easy - add compile flag to configs #634
Conversation
meta-cla
bot
added
the
CLA Signed
felipemello1
changed the title
| model: "Qwen/Qwen3-4B" | ||
| off_by_n: 1 # Off by one by default | ||
| launcher: mast | ||
| compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM |
There was a problem hiding this comment.
This file should have been removed?
There was a problem hiding this comment.
its there: https://github.com/meta-pytorch/torchforge/blob/main/.meta/mast/qwen3_4b_mast.yaml
But i can delete it in this PR if you want
There was a problem hiding this comment.
OOf, why are there so many configs.
Yes, i missed it in https://github.com/meta-pytorch/torchforge/pull/632/files . Please just remove it.
| max_res_tokens: 2048 | ||
| model: "Qwen/Qwen3-8B" | ||
| off_by_n: 1 # Off by one by default | ||
| compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM |
There was a problem hiding this comment.
Why not enabling it by default if you're updating all of the configs?
There was a problem hiding this comment.
wdym by "enabling it by default"? We still need to expose the flag because compile can be tricky in some setups. It also add a bit of warmup time, so if someone is just quickly testing something, they may want to set it to false
There was a problem hiding this comment.
I see. I was suggesting that to reduce the number of hyper parameters in the yaml config because
- We seem to want it to be enable for production runs
- This is a niche config (many things can slow down warmup time) that I don't expect people to remember to toggle in practice. We can set the default to be false when launching the job in local mode or ONLY set them to be true for large models.
Not a big deal.
|
Will enabling |
Yes, the larger the model, the longer it takes. Its between a couple of seconds to 60s. But it decreases peak memory by 40% and increase Tok/s > 20%, in my experience |
2 participants
We currently have 3 places for compile
This PR creates a single flag at the top and sets it to true as default, which helps with memory / tok/s
For the current bsz / seq len, its not a huge difference in speed/memory, but as the models/sequence grow, it becomes more relevant.