This is an example of how to hotswap PyTorch training code without unloading your model weights from VRAM.
For large LLMs it can take upwards of 30 seconds to load a model from disk to VRAM. Waiting 30 seconds every time you want to rerun your script slows down development. This is a barebones implementation of a method to keep large models in VRAM even after your training script exits. If a model reload is necessary, it happens in the background after exit ensuring the model will be ready immediately the next time your script is run.
This works by spawning a second process that stays active after your target script exits. The script you change is not run directly. Instead, this background process runs the code on your behalf using Python’s eval().
This can also be used over a VPN for remote code execution. IntelliJ’s remote SSH interpreter is quite buggy and not ideal for seamless remote development. Configure model_server.py to run on a remote machine, and run client.py on your development machine. Debugging with the IntelliJ debugger is supported in this configuration as well, enabling an almost seamless development experience with scripts that run instantly and are easily debuggable.
Some work has also been done to ensure compatibility with the DearImgui Python bindings. UI code can be submitted to the server along with your training script. I personally like to build out UI for my training scripts to monitor progress, loss over time, and enable easy evaluation. Submitting your UI code along with your training code ensures that your app will launch instantly.
Here’s a GUI from an app that
4 Comments
mathibela
[flagged]
mathibela
[flagged]
pizza
Tensor visualizer app itself already looks pretty interesting
NitpickLawyer
We use python notebooks for that functionality in the early stages of script testing. Load a cell up top with the model, then do your stuff below, and once things look good convert it to a normal python script.