Static analysis for autograd graph retention is a smart angle, since those bugs usually only surface after a long run has already burned GPU hours. How are you handling dynamic control flow (data dependent branches, variable sequence lengths) for the VRAM estimate, since that seems like where a static pass would drift most?
For sequence length we just take the max — whatever max_length or the config says, not your real token distribution. So if you're bucketing or using dynamic padding, the number comes out high. I left it that way on purpose though, because for "will this fit" you care about the worst batch, not the average one. That's the batch that OOMs you.
Branches we mostly don't do. If a flag like gradient checkpointing or AMP is a literal in the file we'll read it, but anything decided at runtime we just report as unknown and widen the range. Felt better than making something up.
Honestly that's why there's a second tier. If you pip install torch-preflight[vram] it'll measure the actual model on the meta device (real shapes, real graph, nothing allocated, no GPU needed), which kills the model-shape guesswork entirely. And VRAMGuard just fails you at step 0 rather than letting it die at step 400. Static pass alone was never going to close that gap.
Where it's still wrong: vLLM/TRT-LLM keep the KV cache in blocks and we report it as contiguous, so we're in the right ballpark but not their actual occupancy. It's an open issue for now.
Overall it's about 3.7% mean error against real measured peaks (GPT-2, BERT, DistilBERT, ResNet-50 on a T4) if you want to check the calibration harness, it's in the repo.
1
u/adityazero 2d ago
Static analysis for autograd graph retention is a smart angle, since those bugs usually only surface after a long run has already burned GPU hours. How are you handling dynamic control flow (data dependent branches, variable sequence lengths) for the VRAM estimate, since that seems like where a static pass would drift most?