Add torch_compile flag for training networks#28
Conversation
|
@juliusberner Could you please take a look at my PR? Thanks! |
wenxin0319
force-pushed
the
main
branch
3 times, most recently
from
eb763ac to
78e9f0d
Compare
juliusberner
left a comment
juliusberner
left a comment
There was a problem hiding this comment.
Thanks for the edits! I posted a couple of follow-up comments and also kicked off the CI!
Also, note that per our contribution guidelines "Commits must have verified signatures.".
| if hasattr(self.net, "init_preprocessors") and self.config.enable_preprocessors: | ||
| self.net.init_preprocessors() | ||
|
|
||
| # Preprocessor sub-objects of ``net`` to also compile (see compile_dict). |
There was a problem hiding this comment.
Can we define this at the top and also use _PREPROCESSOR_ATTRS in on_train_begin, so that this tuple will be the only source-of-truth for the preproc. attributes?
| _PREPROCESSOR_ATTRS = ("vae", "text_encoder", "image_encoder") | ||
|
|
||
| @property | ||
| def compile_dict(self) -> dict: |
There was a problem hiding this comment.
Do you think we could take the model_dict as base dict (which already contains, e.g., fake score and discriminator for DMD2) and
- remove the EMA from the
ema_dict? - add the teacher to the base compile dict, guarded by an
getattr(self, "teacher") is not None(similar to thefsdp_dict) - add the preprocessors as done below
| return compile_dict | ||
|
|
||
| @staticmethod | ||
| def _object_compile_targets(name: str, obj) -> dict: |
There was a problem hiding this comment.
Can we just inline this in the compile_dict method, since it's not used anywhere else?
wenxin0319
force-pushed
the
main
branch
4 times, most recently
from
5100a24 to
b7226d8
Compare
|
@juliusberner Thank you for reviewing it; in the latest commit, i am trying to make the code as concise as i can. let me know if you have any comments. |
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@wenxin0319 Thanks a lot, I've just started the final coderabbit review.
After that, we can merge :) |
WalkthroughThis PR introduces optional torch.compile support for FastGen models. A new configuration field enables torch compilation, a central preprocessor attribute list unifies device placement logic, and a new apply_torch_compile() method compiles trainable modules after distributed wrapping. The trainer invokes compilation post-DDP/FSDP, and comprehensive tests validate behavior across SFT and DMD2 models including edge cases and training integration. Changestorch.compile Feature
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_torch_compile.py`:
- Line 174: The test unpacks two values from model.single_train_step into
loss_map and outputs but never uses outputs; to silence RUF059, change the
unpacking in tests/test_torch_compile.py where single_train_step is called
(references: model.single_train_step) to either discard the second value by
assigning it to _ (e.g., loss_map, _ = ...) or only capture the first element
(e.g., loss_map = model.single_train_step(...)[0]); apply the same change for
both occurrences mentioned (the calls at the two test locations).
- Around line 120-124: The test_dmd2_compile_disabled test is missing an
assertion for the discriminator; add an assertion that _is_compiled returns
False for dmd2_model_not_compiled.discriminator (mirroring the enabled-path test
coverage), i.e., extend test_dmd2_compile_disabled to include assert not
_is_compiled(dmd2_model_not_compiled.discriminator) so the disabled-path
validates the same component as the enabled-path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Enterprise
Run ID: 41b1322a-6dbf-4043-809f-ca12de965606
📒 Files selected for processing (4)
fastgen/configs/config.pyfastgen/methods/model.pyfastgen/trainer.pytests/test_torch_compile.py
Add an optional torch_compile_mode config flag ("default",
"reduce-overhead", "max-autotune"; None disables). Networks to compile
are collected in a compile_dict (derived from model_dict minus EMA, plus
the teacher and the net's preprocessors) and compiled in place via
nn.Module.compile by the trainer, after DDP/FSDP wrapping.
hi, I have added it to image/video inference.py; and the commit is verified right now |
Thanks a lot! Could you address the two CodeRabbit comments above and my remaining comment? |
Greptile SummaryThis PR adds an opt-in
Confidence Score: 4/5Safe to merge for training use cases; the inference compile path silently skips the EMA student network when use_ema=True, so users who rely on compiled inference with EMA won't see the expected speedup. The training path is correct and well-structured. The inference path introduces a mismatch: apply_torch_compile excludes EMA networks (correct for training), but in setup_inference_modules the active student is the EMA network when use_ema=True, so the compiled path is never actually exercised during inference. This is a silent correctness issue rather than a crash, but it defeats the purpose of calling apply_torch_compile in the inference path for the majority of production configurations. scripts/inference/inference_utils.py and fastgen/methods/model.py — the interaction between EMA exclusion in apply_torch_compile and EMA selection in setup_inference_modules needs attention. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Trainer
participant Model as FastGenModel
participant Net as model.net
participant EMA as model.ema
Note over Trainer,EMA: Training path
Trainer->>Model: on_train_begin()
Model->>Net: .to(device, dtype)
Trainer->>Model: apply_torch_compile()
Model->>Net: "module.compile(mode=mode)"
Note over EMA: EMA excluded (not run during training)
Note over Trainer,EMA: Inference path (setup_inference_modules)
Trainer->>Model: cleanup_unused_modules()
Note over Model: fake_score = None, discriminator = None
Trainer->>Model: init_preprocessors()
Trainer->>EMA: "student = model.ema (if use_ema)"
Trainer->>Model: apply_torch_compile()
Model->>Net: "module.compile(mode=mode)"
Note over EMA: EMA excluded again — but EMA IS the inference student!
Note over Net: model.net compiled but NOT called during student sampling
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Trainer
participant Model as FastGenModel
participant Net as model.net
participant EMA as model.ema
Note over Trainer,EMA: Training path
Trainer->>Model: on_train_begin()
Model->>Net: .to(device, dtype)
Trainer->>Model: apply_torch_compile()
Model->>Net: "module.compile(mode=mode)"
Note over EMA: EMA excluded (not run during training)
Note over Trainer,EMA: Inference path (setup_inference_modules)
Trainer->>Model: cleanup_unused_modules()
Note over Model: fake_score = None, discriminator = None
Trainer->>Model: init_preprocessors()
Trainer->>EMA: "student = model.ema (if use_ema)"
Trainer->>Model: apply_torch_compile()
Model->>Net: "module.compile(mode=mode)"
Note over EMA: EMA excluded again — but EMA IS the inference student!
Note over Net: model.net compiled but NOT called during student sampling
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # Compilation is applied by the trainer after DDP/FSDP wrapping; emulate that here. | ||
| model.apply_torch_compile() |
There was a problem hiding this comment.
The fixtures call model.apply_torch_compile() before model.on_train_begin(), but in production the trainer runs the opposite order: on_train_begin() at line 97 of trainer.py, then apply_torch_compile() at line 124. Because on_train_begin() is what calls init_preprocessors() (creating net.vae, net.text_encoder, etc.) and moves parameters to the target device/dtype, calling apply_torch_compile() first means preprocessors don't exist yet and are never compiled — even when torch_compile_mode is set. Any model that relies on compiled preprocessors would appear to work in these tests while silently skipping preprocessor compilation in production. The same wrong order appears in the dmd2_model_compiled, sft_model_not_compiled, and dmd2_model_not_compiled fixtures.
| Returns: | ||
| Tuple of (teacher, student, vae) - any may be None | ||
| """ | ||
| model.apply_torch_compile() # no-op if torch_compile_mode is None |
There was a problem hiding this comment.
apply_torch_compile() is called here before init_preprocessors() is called a few lines later (line 181). Since preprocessors (net.vae, net.text_encoder, net.image_encoder) only exist after init_preprocessors() runs, apply_torch_compile() won't find them and will skip compiling them. In training this is handled correctly — on_train_begin() initialises preprocessors before apply_torch_compile() is invoked by the trainer. For inference with torch_compile_mode set, users expecting compiled preprocessors will silently get uncompiled versions.
| def apply_torch_compile(self): | ||
| """Compile the training networks in place with torch.compile. | ||
|
|
||
| Called by the trainer after DDP/FSDP wrapping (and after the networks | ||
| have been moved to their device in ``on_train_begin``) so torch.compile | ||
| composes with the distributed wrappers. No-op when | ||
| ``config.torch_compile_mode`` is None. | ||
|
|
||
| The modules compiled are those in model_dict (e.g. net, plus | ||
| fake_score/discriminator for DMD2) minus the EMA networks, plus the | ||
| teacher (if any, cf. fsdp_dict) and the net's preprocessors. | ||
| """ | ||
| mode = self.config.torch_compile_mode | ||
| if mode is None: | ||
| return | ||
|
|
||
| # model_dict contains the trainable networks (incl. EMA); EMA networks are | ||
| # weight-averaged copies that aren't run during training, so drop them. | ||
| # None entries arise when cleanup_unused_modules has already been called. | ||
| modules = {name: net for name, net in self.model_dict.items() if name not in self.ema_dict and net is not None} | ||
| # The teacher is not part of model_dict; add it when present (cf. fsdp_dict). | ||
| if getattr(self, "teacher", None) is not None: | ||
| modules["teacher"] = self.teacher | ||
| # Preprocessors (VAE, text/image encoders) are often lightweight wrappers | ||
| # (e.g. WanVideoEncoder, SDVAE) that are not nn.Modules themselves but hold | ||
| # the actual nn.Module under an attribute; compile those submodules too. | ||
| for name in self._PREPROCESSOR_ATTRS: | ||
| obj = getattr(self.net, name, None) | ||
| if obj is None: | ||
| continue | ||
| if isinstance(obj, torch.nn.Module): | ||
| modules[name] = obj | ||
| else: | ||
| for attr, submodule in getattr(obj, "__dict__", {}).items(): | ||
| if isinstance(submodule, torch.nn.Module): | ||
| modules[f"{name}.{attr}"] = submodule | ||
|
|
||
| for name, module in modules.items(): | ||
| logger.info(f"Applying torch.compile (mode={mode}) to {name}") | ||
| module.compile(mode=mode) |
There was a problem hiding this comment.
apply_torch_compile() has no guard against being called more than once on the same model instance. nn.Module.compile() applied a second time wraps the already-compiled _call_impl, compiling a compiled callable. While there is no code path in the current repo that calls this twice in one session, any future callback or utility bridging the training and inference paths could double-compile. A simple check like if getattr(module, "_compiled_call_impl", None) is not None: continue inside the loop would prevent this.
There was a problem hiding this comment.
@juliusberner thank you for pointing it out, i have fixed both comments
|
@wenxin0319 it looks great! Could you fix the CI, such that we can merge it? |
2 participants
FastGen currently relies on diffusers-based model execution, which leaves performance on the table during training.
This PR adds an opt-in torch_compile flag that wraps training networks with torch.compile, enabling PyTorch's compiler optimizations (operator fusion, memory planning, kernel autotuning) for significant speedups on common models.
Benchmark (QwenImage, 20.43B params, NVIDIA H100, bfloat16, 512x512):
Setting │ Time/iter │ Std
Baseline (no compile) │ 0.694s │ 0.094s
torch.compile (max-autotune) │ 0.447s │ 0.014s
which is Speedup 1.55x (55% faster)
Compiled iterations also show much lower variance (0.014s vs 0.094s), meaning more consistent training throughput. The one-time compilation overhead (~5-10 min with max-autotune) is amortized over the full training run.
Changes:
Usage:
Set torch_compile=True in model config to enable.
Summary by CodeRabbit
New Features
Tests