A linter for PyTorch 'torch-preflight' [P]

Reddit r/MachineLearning Tools

Summary

torch-preflight is a new linter for PyTorch code that catches common GPU-wasting bugs without executing code or requiring a GPU, plus estimates VRAM usage to check if a training run fits before paying for an instance.

Been working on this for the last few months. I've been working on PyTorch for the past few years and I always felt, many a times my work went into dump, because of some mistakes I made in the code. torch-preflight reads your PyTorch code and catches the bugs costing you GPU hours. Things like losses.append(loss), which holds the autograd graph from every step until CUDA dies on you or no zero_grad() in the loop or gradient accumulation without dividing the loss or DDP with no DistributedSampler, so every rank trains on the same batches. I've been able to get 13 rules so far. Your code never gets imported or executed, so you need no GPU and no torch install. There's another part to this that estimates VRAM. Point the tool at a training script and a GPU, and you learn whether the run fits before you pay for the instance. You also get the list of changes to make the run fit, with the GiB each one saves. pip install torch-preflight https://github.com/highwaterlabs/torch-preflight https://pypi.org/project/torch-preflight/ Please try this out, and I would like to get your feedback! It's still a work in progeress. Would like to know what breaks on your code. False positives kill a linter, and my only large test target so far has been the PyTorch source tree. Same for the memory numbers. Mine land within 4% of measured peaks, but from four models on one T4. PS: open to contributions, and issues are already open on the repo. Soon I'm going to add a few "Good first issues" as well. Feel free to ping me if you have any questions!
Original Article

Similar Articles

A hackable compiler to generate efficient fused GPU kernels for AI models [P]

Reddit r/MachineLearning

The author presents a custom, hackable ML compiler written in Python that lowers LLMs to optimized CUDA kernels through a multi-stage IR pipeline, achieving performance competitive with or superior to PyTorch on specific operations. The article details the compiler's optimization passes, lowering rules, and CLI usage for generating efficient fused GPU kernels.