个人研究者复现 NBFNet 药物重定位研究，A100 训练成本成难题

Hi, ML people of Reddit! I have come to ask for some human advice. I am an amateur (as in unpaid) ML researcher working on a personal project to find drug repurposing candidates for my wife, who was Dxed with Stage 4 colon cancer in early 2024. That the new AI coding tools have made it even possible for someone like me to do this research in my spare time is just amazing, and I owe every ML person a big "thank you" for enabling that.

That being said, I am starting to accumulate GPU costs, and it got painful enough this morning for me to wonder if I am using my dollar budget in the best possible way. And I'm wondering what other people do, in similar situations. After granting myself unfettered spending to build a platform this year (the value of money changes in strange ways with a cancer diagnosis), it's time to optimize costs to better focus my resources.

My basic setup is a Google Vertex training job using their A100s (I am currently training NBFnets over biomedical graphs, if that matters), and I am trying to repro the results of a few papers published recently.

The main paper I’m reproducing reports 100 epochs of 500 batches at batch size 512. On an A100 80GB, my runs fall back to batch size 64. That’s because of how NBFNet works. It has no per-node embedding table. For every query (drug, relation, ?) in a batch, it runs 6 layers of message passing over the entire graph and keeps a 64-dim state for every node. My graph has ~103K nodes and ~2.5M training edges, so one copy of that state is 103K × 64 floats ≈ 26 MB. Backprop keeps several of those per layer. Over 6 layers, that’s roughly 1 GB of saved activations per query. So 64 queries fit in the A100’s 80 GB, 128 don’t, and 512 would need about half a terabyte. The training script catches the CUDA out-of-memory error, halves the batch and carries on. Every run downgrades to 512 --> 256 --> 128 --> 64 in its first five seconds.

(Note: I am not certain they used 512 either. The paper doesn't specify; 512 is the configured default. For all I know, they were getting downgraded too. This halving loop came in their published code, and was presumably put there because it was needed at some point.)

The bigger effect is on how much data the model sees. An “epoch” here is a fixed 500 steps, so at batch 64 my 20 epochs see 640K queries. That’s about a quarter of one pass over the training edges. One full pass would take ~76 hours, about $412. Matching the paper’s stated 100 × 500 steps would take ~4 days, roughly $500 a run. And since I am experimenting with different things, the costs are adding up.

I'm kind of taking a step back this morning and working on a "toy" implementation that can at least prove the mechanism works against synthetic data on a much smaller graph. I estimate even that will cost $150 each for 3 seed runs.

My questions, if you are in the mood:

How much training is enough? For from-scratch NBFNet/ULTRA on a ~100K-node, 2.5M-edge graph, what budget do people use, and how do you find convergence cheaply? (The ULTRA paper trained Hetionet, 45K nodes and 2M edges, from scratch for 10,000 steps at batch 16, so small batches may not be my real problem.)

Am I overpaying? I pay ~$5.40/hour on demand. My jobs run ~20 hours and can resume from a checkpoint. What other options are there? At what point am I better off buying a GPU? Can you even do that in 2026 without being a $trillion company?

Memory tricks: would bf16 autocast or activation checkpointing get me to 128+ queries per batch, with a custom sparse CUDA kernel in the loop? Or is that a rabbit hole?

(not cost related but) Popularity bias: favorite fixes for link-prediction scores that mostly track node degree?

Thanks. Sorry if it sounds like I'm lost. I'm publishing my Lab Notebook with Quartz here: https://whafa.github.io/mods-notebook/

Everything there 100% generated but intended for human consumption and, if you point your agent at it and ask how things are going, it will tell you I've built a fairly sophisticated system that has not made any useful predictions or findings yet.

submitted by /u/whafa
[link] [留言]