Enable torch.compile for OSFT training - #111

Draft

stmcgovern wants to merge 19 commits into

Red-Hat-AI-Innovation-Team:mainfrom

stmcgovern:feat/osft-compile

Commits on Jun 25, 2026

Add per-block torch.compile support for SFT training
stmcgovern
committed
a68c50e
Move validation and global config out of wrap_fsdp2
stmcgovern
committed
83b7b2e
Refine torch.compile call: drop redundant backend, add dynamic=True
stmcgovern
committed
4fa6054
Add tests for per-block torch.compile
stmcgovern
committed
578a60b
Fix dynamo config comment and add compile_model to docstring
stmcgovern
committed
be569cb
Remove compile-hostile patterns from training loop
stmcgovern
committed
c8bec87
Fix compile recompiles: mod-8 padding and skip cache shape guards
stmcgovern
committed
a89c3e6
Fix CI: guard CUDA call, update tests for mod-8 padding
stmcgovern
committed
9bf5209
Enable torch.compile for OSFT via OSFTLinear module
stmcgovern
committed
c41645c
Fix reinitialize_osft when called on already-initialized model
stmcgovern
committed
f7c6488
Add regression tests for reinitialize_osft double-init
stmcgovern
committed
be9301c
Add distributed benchmark for OSFT eager vs compiled
stmcgovern
committed
74893ca
Fix lint: remove unused import, f-string prefixes; reformat
stmcgovern
committed
381353d
Fix review findings: meta device in restore, stale Protocol, docstring
stmcgovern
committed
64cb609
Use SDPA as default attention implementation
stmcgovern
committed
182a144
Add orthogonality-under-compile test for OSFT
stmcgovern
committed
10310df
Fix lint: use noqa for flash_attn availability check import
stmcgovern
committed
3735236
Tighten comments: fix imprecise SDPA claim, trim docstrings
stmcgovern
committed
dcce1f6
Fix stale attribute names in _is_osft_owned_attribute
stmcgovern
committed
d1b05cf

Navigation Menu

Enable torch.compile for OSFT training - #111

Enable torch.compile for OSFT training#111
stmcgovern wants to merge 19 commits into
Red-Hat-AI-Innovation-Team:mainfrom
stmcgovern:feat/osft-compile

Commits on Jun 25, 2026

Add per-block torch.compile support for SFT training

Move validation and global config out of wrap_fsdp2

Refine torch.compile call: drop redundant backend, add dynamic=True

Add tests for per-block torch.compile

Fix dynamo config comment and add compile_model to docstring

Remove compile-hostile patterns from training loop

Fix compile recompiles: mod-8 padding and skip cache shape guards

Fix CI: guard CUDA call, update tests for mod-8 padding

Enable torch.compile for OSFT via OSFTLinear module

Fix reinitialize_osft when called on already-initialized model

Add regression tests for reinitialize_osft double-init

Add distributed benchmark for OSFT eager vs compiled

Fix lint: remove unused import, f-string prefixes; reformat

Fix review findings: meta device in restore, stale Protocol, docstring

Use SDPA as default attention implementation

Add orthogonality-under-compile test for OSFT

Fix lint: use noqa for flash_attn availability check import

Tighten comments: fix imprecise SDPA claim, trim docstrings

Fix stale attribute names in _is_osft_owned_attribute