Enable torch.compile for OSFT training - #111
Draft
stmcgovern wants to merge 19 commits into
Draft
stmcgovern wants to merge 19 commits into
stmcgovern wants to merge 19 commits into
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ 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 |
Compile each transformer block with torch.compile(fullgraph=True) following torchtitan's AC -> compile -> FSDP2 ordering. Benchmarked on 8x H200 (SDPA, BF16, batch=4, 2048 tok/GPU): Qwen3-4B: 4,184 -> 6,008 tok/s (+44%) Granite-8B: 3,227 -> 3,708 tok/s (+15%) Not compatible with OSFT (graph breaks from closure-based forwards) or MoE architectures (router logic breaks fullgraph).
- Move MoE class-name guard to main() (after model load, where we have the class name and can fail early) - Move torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint to main() (global side effect shouldn't be buried in a wrapping function) - Add compile_model + use_liger_kernels incompatibility guard (both replace the same memory-bound ops; interaction is untested) - wrap_fsdp2 now just does what it's told: AC, compile, FSDP2
- Remove backend="inductor" (it's the default, and hardcoding blocks TORCH_COMPILE_BACKEND env var override for debugging) - Add dynamic=True to skip straight to symbolic shapes, avoiding 2 of 3 guard-failure recompilations (~15s saved per run since AC bypasses FX cache) - Add comment explaining why the dynamo config flag is needed
GPU tests (tests/gpu_tests/test_compile.py): - Numerical equivalence: compiled vs eager losses match within tolerance - No graph breaks under fullgraph=True with varied sequence lengths - dynamic=True prevents recompilation on shape changes - OptimizedModule wrappers present on compiled blocks, absent on eager - AC + compile works without the defensive dynamo config flag Unit tests: - Validation guards for compile+osft, compile+liger, compile+MoE - API flag passthrough (--compile-model in torchrun command)
- Comment was wrong: said flag is required, but tests prove it isn't on current PyTorch. Updated to say defensive/forward-looking. - Added compile_model parameter to wrap_fsdp2 docstring.
- Remove torch.cuda.empty_cache() from training inner loop and validation. These were cargo-culted in (e2c6e73, cbb6762) without a specific OOM motivating them. With AC + FSDP2, peak memory is bounded by checkpoint granularity and shard size. empty_cache() releases the CUDA caching allocator's free list every minibatch, forcing Inductor workspace buffers to be re-acquired from the driver on every forward pass. - Drop redundant .cpu() before .item() on loss tensor — .item() on a CUDA scalar already syncs and returns the value. - Reuse loss_metrics instead of calling .item() a second time for loss_backward. Eliminates a second full CUDA sync per minibatch. - Pad batch sequence lengths to multiples of 8 in the sampler. Inductor emits (size[-1] % 8) == 0 alignment guards that cause recompilation when switching between aligned and non-aligned sequences. This is by design (different kernels for vectorized loads), not a PyTorch bug. Padding eliminates the recompile class entirely. Overhead: at most 7 extra padding tokens per batch. - Log memory utilization (peak/total) after step 1 for compile headroom visibility. - Fix test assertion for wrap_fsdp2 mock (expects compile_model kwarg).
Two fixes for torch.compile recompilation overhead:
1. Pad packed sequences to multiples of 8 in mb_collate_fn. Inductor
emits a mod-8 alignment guard for vectorized loads — variable-length
packed sequences were triggering recompiles on every new alignment
class. Worst case overhead: 7 extra tokens (<0.1% of typical batch).
2. Set unsafe_skip_cache_dynamic_shape_guards=True when compiling.
The AOTAutograd cache records range guards based on sequence lengths
seen during compilation. Packed SFT has continuous length variance,
causing the inductor scheduler's fusion thresholds to create new
shape regimes — sometimes hundreds of steps into training. Skipping
these guards reuses the first compiled graph for all shapes. The
generated code handles variable shapes correctly (dynamic=True);
the guards only controlled fusion strategy optimality, not
correctness.
Before: warm cache step 1 = 40s, 8 recompiles (bs=16), no breakeven
in 500 steps
After: warm cache step 1 = 17s, 0 recompiles, breakeven step 88-171,
compiled saves 47-87s over 1000 steps on 8xH200
- Add torch.cuda.is_available() guard on memory utilization log (get_device_properties crashes when CUDA_VISIBLE_DEVICES is empty) - Update collate test assertions for mod-8 sequence padding - Remove empty_cache assertion (calls removed as compile-hostile)
Replace closure-based forward overrides with OSFTLinear(nn.Module) whose forward is pure tensor math — no weakrefs, dict construction, getattr, or runtime .to() calls — so torch.compile(fullgraph=True, dynamic=True) can trace OSFT blocks without graph breaks. Key changes: - Add OSFTLinear class with lazy-cached rank_high (persistent buffer, single .item() call then cached) to avoid CUDA→CPU sync per step - Refactor _prepare_osft_param and _initialize_osft_parameters to create OSFTLinear and replace Linear via setattr on parent module - Remove _factorized_linear, _register_osft_target, _get_module_by_logical_key, _osft_handles weakref infrastructure, vestigial model-level osft_params dict - Remove compile+OSFT guard in train.py - Add 3 GPU tests: compiled-matches-eager, no-graph-breaks, OptimizedModule wrappers - Add bench_osft_compile.py profiling script
After the first OSFT init, OSFTLinear replaces the original nn.Linear modules, changing parameter FQNs from e.g. `q_proj.weight` to `q_proj.osft_U_high`. A subsequent reinitialize_osft call found zero matching parameters because _initialize_osft_parameters matches against osft_config keys which use the original FQNs. Add _restore_dense_linears() which reconstructs dense weight matrices from SVD factors and replaces OSFTLinear modules back with nn.Linear before _reset_osft_metadata clears the registry. This restores the original FQNs so the re-decomposition finds all targets.
Covers the case where reinitialize_osft is called on a model that was already initialized (initialize_osft=True). Previously this silently produced a model with zero OSFT parameters.
torchrun-based benchmark measuring per-step throughput with CUDA events. Results from Granite-3.3-8B on 6x H200: 10-17% speedup with compile.
- _restore_dense_linears: use device="meta" in nn.Linear constructor to avoid transient GPU memory spike from unused kaiming init - OSFTModelProtocol: replace stale osft_params with osft_paramspec_registry - TestFactorizedLinearAccuracy: update docstring to reference OSFTLinear
HF transformers' flash attention path has a data-dependent graph break in _is_packed_sequence (modeling_flash_attention_utils.py) that prevents torch.compile with fullgraph=True. This is tracked as HF #41803 with no fix timeline. SDPA avoids this code path entirely and selects the optimal attention kernel for the hardware via PyTorch's backend selection (cuDNN on Hopper, FlashAttention on Ampere). GPT-OSS models retain their existing vllm-flash-attn3 (Hopper+) or eager fallback, since they require flash-attn3 specifically.
Verifies that gradient and parameter subspace orthogonality is preserved when OSFT training runs under torch.compile with FSDP2. Runs 10 training steps and checks all OSFTLinear modules at every step using OrthogonalityTracker with a 1-degree margin.
The OSFTLinear refactor changed rank_high from a plain tensor attribute to a registered buffer named _rank_high_buf. Update the defense-in-depth check to match the current attribute names.
stmcgovern
force-pushed
the
feat/osft-compile
branch
from
June 25, 2026 22:31
80c0a9a to
d1b05cf
Compare
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Collaborator
|
Benchmarked the OSFT compile PR on the same setup (8x H100, Granite-3.3-8B, 200 steps). Used fp32 for train_dtype since that's what OSFT requires for the upcast precision. OSFT Compile Benchmark
No memory overhead from compile, which is nice. The speedup is consistent across the run with no recompilation spikes. Note: I had to merge current main into this branch to pick up a recent fix (#107). No conflicts in the OSFT code. Config: 8x H100 80GB, batch_size=4, max_tokens_per_gpu=4096, fp32, seed=42, 5 warmup steps excluded. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
nn.Modulewhoseforward()is pure tensor math —torch.compile(fullgraph=True)traces it without graph breaksflash_attention_2path has a data-dependent graph break (_is_packed_sequence); SDPA avoids it and selects the optimal kernel per hardware (cuDNN on Hopper, FlashAttention on Ampere)Benchmark
Granite-3.3-8B, 6x H200, rank ratio 0.25:
Test coverage
Depends on
This PR is based on
feat/torch-compile-per-block(#98). Should be merged after that PR lands.