From a68c50e34f30196f526c0f8b418965c63d86b710 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 6 May 2026 00:13:08 +0000
Subject: [PATCH 01/19] Add per-block torch.compile support for SFT training

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).
---
 src/mini_trainer/api_train.py                |  3 ++
 src/mini_trainer/setup_model_for_training.py | 29 ++++++++++++++------
 src/mini_trainer/train.py                    |  8 ++++++
 src/mini_trainer/training_types.py           |  9 ++++++
 4 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/src/mini_trainer/api_train.py b/src/mini_trainer/api_train.py
index a50c5071..228ceb1e 100644
--- a/src/mini_trainer/api_train.py
+++ b/src/mini_trainer/api_train.py
@@ -167,6 +167,9 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
         command.append(f"--min-samples-per-checkpoint={train_args.min_samples_per_checkpoint}")
 
     # Add optional boolean flags
+    if train_args.compile_model:
+        command.append("--compile-model")
+
     if train_args.use_liger_kernels:
         command.append("--use-liger-kernels")
 
diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index 39f0f67b..6b9bb77e 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -415,7 +415,7 @@ def prepare_model_for_fsdp2(model: torch.nn.Module) -> ModelInitializationContex
     return context
 
 
-def wrap_fsdp2(model: torch.nn.Module) -> torch.nn.Module:
+def wrap_fsdp2(model: torch.nn.Module, compile_model: bool = False) -> torch.nn.Module:
     """
     Phase 2: Pure FSDP2 wrapping with activation checkpointing.
 
@@ -482,6 +482,20 @@ def wrap_fsdp2(model: torch.nn.Module) -> torch.nn.Module:
             # preserve_rng_state needs to be true so that the backward pass can be accurate
             layers[idx] = ptd_checkpoint_wrapper(block, preserve_rng_state=True)
 
+    # Apply torch.compile to each block (after AC, before FSDP2)
+    if compile_model:
+        class_name = model.__class__.__name__
+        moe_classes = ("MixtralForCausalLM", "GraniteMoeHybridForCausalLM")
+        if class_name in moe_classes:
+            raise ValueError(
+                f"--compile-model is not compatible with MoE architecture {class_name}. "
+                "MoE router logic causes graph breaks with fullgraph=True."
+            )
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+        log_rank_0(f"🔄 [Phase 2] Compiling {len(layers)} transformer blocks with torch.compile")
+        for idx, block in enumerate(layers):
+            layers[idx] = torch.compile(block, backend="inductor", fullgraph=True)
+
     # Build 1D device mesh over all ranks
     world_size = dist.get_world_size()
     mesh = init_device_mesh("cuda", [world_size], mesh_dim_names=["fsdp"])
@@ -1262,11 +1276,6 @@ def load_osft_model():
             to_print=True,
         )
 
-    # NOTE: Don't enable HuggingFace gradient checkpointing with FSDP2
-    # It causes conflicts. TorchTitan applies PyTorch's checkpoint wrapper
-    # BEFORE FSDP2 wrapping if needed.
-    # model.gradient_checkpointing_enable()
-    # torch.compile(model)
     return model
 
 
@@ -1283,13 +1292,14 @@ def setup_training_components(
     eps: float = 1e-8,
     weight_decay: float = 0.0,
     resume_from_checkpoint: str | None = None,
+    compile_model: bool = False,
 ) -> tuple[torch.nn.Module, torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]:
     """
     Set up training components including model wrapping, optimizer, and learning rate scheduler.
 
     This function orchestrates the three-phase model initialization pipeline:
     1. Phase 1: Prepare model for FSDP2 (extract state dicts, materialize buffers)
-    2. Phase 2: Pure FSDP2 wrapping (activation checkpointing + sharding)
+    2. Phase 2: Pure FSDP2 wrapping (activation checkpointing + compilation + sharding)
     3. Phase 3: Finalize initialization (distribute weights, compute SVD for OSFT)
 
     Args:
@@ -1299,6 +1309,7 @@ def setup_training_components(
         lr_scheduler: Type of learning rate scheduler to use
         num_training_steps: Total number of training steps (required for some schedulers)
         scheduler_kwargs: Additional scheduler-specific keyword arguments
+        compile_model: Whether to compile transformer blocks with torch.compile
 
     Returns:
         Tuple of (wrapped_model, optimizer, lr_scheduler)
@@ -1313,8 +1324,8 @@ def setup_training_components(
     init_context = prepare_model_for_fsdp2(model)
     init_context.resume_from_checkpoint = resume_from_checkpoint
 
-    # Phase 2: Pure FSDP2 wrapping
-    model = wrap_fsdp2(model)
+    # Phase 2: Pure FSDP2 wrapping (+ optional compilation)
+    model = wrap_fsdp2(model, compile_model=compile_model)
 
     # Phase 3: Finalize model initialization (distribute weights)
     model = finalize_model_initialization(model, init_context)
diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index 42fbfa76..d244efd8 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -1283,6 +1283,7 @@ def main(
     beta2: Annotated[float, Option(help="AdamW beta2 parameter (RMSprop coefficient)")] = 0.95,
     eps: Annotated[float, Option(help="AdamW epsilon for numerical stability")] = 1e-8,
     weight_decay: Annotated[float, Option(help="AdamW weight decay (L2 penalty)")] = 0.0,
+    compile_model: Annotated[bool, Option(help="Compile transformer blocks with torch.compile")] = False,
     use_liger_kernels: Annotated[bool, Option(help="Whether to use Liger kernels")] = False,
     osft: Annotated[bool, Option(help="Enable OSFT (Orthogonal Subspace Fine-Tuning)")] = False,
     osft_unfreeze_rank_ratio: Annotated[
@@ -1406,6 +1407,12 @@ def main(
 
     # validation, do this before continuing execution flow so we don't log experiments that are invalid from
     # the get-go
+    if compile_model and osft:
+        raise ValueError(
+            "--compile-model is not compatible with --osft. "
+            "OSFT uses dynamic forward methods that cannot be traced by torch.compile."
+        )
+
     if osft:
         if osft_unfreeze_rank_ratio is None:
             raise ValueError("osft_unfreeze_rank_ratio is required when osft is True")
@@ -1613,6 +1620,7 @@ def main(
         eps=eps,
         weight_decay=weight_decay,
         resume_from_checkpoint=resume_from_full_state_checkpoint,
+        compile_model=compile_model,
     )
 
     # Reconstruct callbacks from serialized CLI arg
diff --git a/src/mini_trainer/training_types.py b/src/mini_trainer/training_types.py
index 67736b6d..11e46b17 100644
--- a/src/mini_trainer/training_types.py
+++ b/src/mini_trainer/training_types.py
@@ -96,6 +96,15 @@ class TrainingArgs:
     )
     weight_decay: float = field(default=0.0, metadata={"help": "Weight decay (L2 penalty) for AdamW optimizer."})
 
+    # Compilation
+    compile_model: bool = field(
+        default=False,
+        metadata={
+            "help": "Compile transformer blocks with torch.compile for improved throughput. "
+            "Not compatible with OSFT or MoE architectures."
+        },
+    )
+
     # Model configuration
     use_liger_kernels: bool = field(default=False, metadata={"help": "Whether to use Liger kernels."})
     osft: bool = field(

From 83b7b2ecc01835c50a9ebdcff68afbceca11471f Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 6 May 2026 21:21:47 +0000
Subject: [PATCH 02/19] Move validation and global config out of wrap_fsdp2

- 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
---
 src/mini_trainer/setup_model_for_training.py |  8 --------
 src/mini_trainer/train.py                    | 15 +++++++++++++++
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index 6b9bb77e..ed01eb6d 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -484,14 +484,6 @@ def wrap_fsdp2(model: torch.nn.Module, compile_model: bool = False) -> torch.nn.
 
     # Apply torch.compile to each block (after AC, before FSDP2)
     if compile_model:
-        class_name = model.__class__.__name__
-        moe_classes = ("MixtralForCausalLM", "GraniteMoeHybridForCausalLM")
-        if class_name in moe_classes:
-            raise ValueError(
-                f"--compile-model is not compatible with MoE architecture {class_name}. "
-                "MoE router logic causes graph breaks with fullgraph=True."
-            )
-        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
         log_rank_0(f"🔄 [Phase 2] Compiling {len(layers)} transformer blocks with torch.compile")
         for idx, block in enumerate(layers):
             layers[idx] = torch.compile(block, backend="inductor", fullgraph=True)
diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index d244efd8..cd0ac48e 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -1413,6 +1413,12 @@ def main(
             "OSFT uses dynamic forward methods that cannot be traced by torch.compile."
         )
 
+    if compile_model and use_liger_kernels:
+        raise ValueError(
+            "--compile-model is not compatible with --use-liger-kernels. "
+            "Both replace the same memory-bound ops; the interaction is untested."
+        )
+
     if osft:
         if osft_unfreeze_rank_ratio is None:
             raise ValueError("osft_unfreeze_rank_ratio is required when osft is True")
@@ -1571,6 +1577,15 @@ def main(
                             model.resize_token_embeddings(ckpt_embed_size)
                     break
 
+    if compile_model:
+        moe_classes = ("MixtralForCausalLM", "GraniteMoeHybridForCausalLM")
+        if model.__class__.__name__ in moe_classes:
+            raise ValueError(
+                f"--compile-model is not compatible with MoE architecture {model.__class__.__name__}. "
+                "MoE router logic causes graph breaks with fullgraph=True."
+            )
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+
     # Create PretrainingConfig if block_size is provided
     pretraining_config = None
     if block_size is not None:

From 4fa6054196379fbd88c8f01355d71a622ce46345 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 7 May 2026 04:25:49 +0000
Subject: [PATCH 03/19] Refine torch.compile call: drop redundant backend, add
 dynamic=True

- 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
---
 src/mini_trainer/setup_model_for_training.py | 2 +-
 src/mini_trainer/train.py                    | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index ed01eb6d..ea84d1d5 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -486,7 +486,7 @@ def wrap_fsdp2(model: torch.nn.Module, compile_model: bool = False) -> torch.nn.
     if compile_model:
         log_rank_0(f"🔄 [Phase 2] Compiling {len(layers)} transformer blocks with torch.compile")
         for idx, block in enumerate(layers):
-            layers[idx] = torch.compile(block, backend="inductor", fullgraph=True)
+            layers[idx] = torch.compile(block, fullgraph=True, dynamic=True)
 
     # Build 1D device mesh over all ranks
     world_size = dist.get_world_size()
diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index cd0ac48e..cd3ea92a 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -1584,6 +1584,7 @@ def main(
                 f"--compile-model is not compatible with MoE architecture {model.__class__.__name__}. "
                 "MoE router logic causes graph breaks with fullgraph=True."
             )
+        # Without this, AC's RNG side effects cause a graph break under compile.
         torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
 
     # Create PretrainingConfig if block_size is provided

From 578a60bea8bca67bd0efe507088d7cb3ff120dd4 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 7 May 2026 04:26:24 +0000
Subject: [PATCH 04/19] Add tests for per-block torch.compile

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)
---
 tests/gpu_tests/test_compile.py | 279 ++++++++++++++++++++++++++++++++
 tests/test_api_train.py         |   4 +
 tests/test_compile_guards.py    |  51 ++++++
 3 files changed, 334 insertions(+)
 create mode 100644 tests/gpu_tests/test_compile.py
 create mode 100644 tests/test_compile_guards.py

diff --git a/tests/gpu_tests/test_compile.py b/tests/gpu_tests/test_compile.py
new file mode 100644
index 00000000..5c932613
--- /dev/null
+++ b/tests/gpu_tests/test_compile.py
@@ -0,0 +1,279 @@
+"""GPU tests for per-block torch.compile integration with FSDP2."""
+
+import os
+
+os.environ["TESTING"] = "true"
+
+import gc
+
+import pytest
+import torch
+import torch.distributed as dist
+from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
+
+from mini_trainer.setup_model_for_training import setup_model, setup_training_components
+from mini_trainer.utils import patch_target_module
+
+
+def create_tiny_llama_model():
+    config = LlamaConfig(
+        vocab_size=1000,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=128,
+        rope_theta=10000.0,
+        hidden_act="silu",
+    )
+    return LlamaForCausalLM(config), config
+
+
+def _run_steps(model_path, compile_model, input_ids, labels, num_steps=3):
+    """Load model, wrap with FSDP2 (± compile), run steps, return losses."""
+    model = setup_model(
+        model_name_or_path=str(model_path),
+        use_liger_kernels=False,
+        osft=False,
+        local_rank=0,
+    )
+    model, optimizer, lr_scheduler = setup_training_components(
+        model,
+        learning_rate=1e-3,
+        num_warmup_steps=0,
+        lr_scheduler="constant",
+        compile_model=compile_model,
+    )
+
+    losses = []
+    for _ in range(num_steps):
+        optimizer.zero_grad()
+        output = model(input_ids=input_ids, labels=labels)
+        loss = output.loss.float().sum() / input_ids.shape[0]
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        lr_scheduler.step()
+        losses.append(loss.item())
+
+    return losses, model
+
+
+@pytest.mark.gpu
+class TestCompile:
+    @pytest.fixture(autouse=True, scope="class")
+    def dist_env(self):
+        """Single process group for the entire test class."""
+        os.environ["RANK"] = "0"
+        os.environ["WORLD_SIZE"] = "1"
+        os.environ["LOCAL_RANK"] = "0"
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = "12356"
+        dist.init_process_group(backend="nccl", rank=0, world_size=1)
+
+        from mini_trainer.none_reduction_losses import (
+            hf_fixed_cross_entropy_none_reduction,
+        )
+
+        patch_target_module(
+            "transformers.loss.loss_utils.fixed_cross_entropy",
+            hf_fixed_cross_entropy_none_reduction,
+        )
+
+        yield
+
+        dist.destroy_process_group()
+
+    @pytest.fixture(autouse=True)
+    def reset_dynamo(self):
+        torch._dynamo.reset()
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = False
+        yield
+        torch._dynamo.reset()
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = False
+
+    @pytest.fixture
+    def saved_model(self, tmp_path):
+        """Create and save a tiny Llama model + tokenizer to disk."""
+        torch.manual_seed(42)
+        model, config = create_tiny_llama_model()
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        tokenizer.pad_token = tokenizer.eos_token
+        model_path = tmp_path / "tiny_llama"
+        model.save_pretrained(model_path)
+        tokenizer.save_pretrained(model_path)
+        return model_path, config
+
+    def test_compiled_matches_eager(self, saved_model, single_gpu_device):
+        """Compiled and eager produce approximately equal losses.
+
+        Not exact equality: bf16 mixed precision + inductor kernel fusion
+        reorder floating-point ops, so small divergence is expected.
+        Observed gap is O(1e-3) on loss values O(1e+2), growing across
+        steps as differences accumulate. Tolerance is set at 10x the
+        observed per-step gap.
+        """
+        model_path, config = saved_model
+
+        torch.manual_seed(99)
+        input_ids = torch.randint(0, config.vocab_size, (2, 32), device=single_gpu_device)
+        labels = input_ids.clone()
+
+        # Eager run
+        torch.manual_seed(7)
+        torch.cuda.manual_seed(7)
+        eager_losses, eager_model = _run_steps(model_path, compile_model=False, input_ids=input_ids, labels=labels)
+
+        del eager_model
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch._dynamo.reset()
+
+        # Compiled run (same model weights from disk, same seed)
+        torch.manual_seed(7)
+        torch.cuda.manual_seed(7)
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+        compiled_losses, _ = _run_steps(model_path, compile_model=True, input_ids=input_ids, labels=labels)
+
+        for step, (e, c) in enumerate(zip(eager_losses, compiled_losses)):
+            assert abs(e - c) < 0.1, (
+                f"Step {step}: eager loss {e:.6f} vs compiled loss {c:.6f}, diff {abs(e - c):.2e} exceeds tolerance 0.1"
+            )
+
+    def test_no_graph_breaks_and_dynamic_shapes(self, saved_model, single_gpu_device):
+        """Forward/backward completes under fullgraph=True with varied seq_len.
+
+        Two forward/backward passes with different sequence lengths verify:
+        1. No graph breaks (fullgraph=True contract)
+        2. dynamic=True reuses the same compiled graph (no recompilation)
+        """
+        model_path, config = saved_model
+
+        torch._dynamo.utils.counters.clear()
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+
+        model = setup_model(
+            model_name_or_path=str(model_path),
+            use_liger_kernels=False,
+            osft=False,
+            local_rank=0,
+        )
+        model, optimizer, lr_scheduler = setup_training_components(
+            model,
+            learning_rate=1e-3,
+            num_warmup_steps=0,
+            lr_scheduler="constant",
+            compile_model=True,
+        )
+
+        # Step 1: seq_len=32
+        input_ids_1 = torch.randint(0, config.vocab_size, (2, 32), device=single_gpu_device)
+        optimizer.zero_grad()
+        loss = model(input_ids=input_ids_1, labels=input_ids_1.clone()).loss.float().sum()
+        loss.backward()
+        optimizer.step()
+
+        compilations_after_first = torch._dynamo.utils.counters["stats"]["ok"]
+
+        # Step 2: seq_len=48 (different shape — should reuse graph via dynamic=True)
+        input_ids_2 = torch.randint(0, config.vocab_size, (2, 48), device=single_gpu_device)
+        optimizer.zero_grad()
+        loss = model(input_ids=input_ids_2, labels=input_ids_2.clone()).loss.float().sum()
+        loss.backward()
+        optimizer.step()
+
+        compilations_after_second = torch._dynamo.utils.counters["stats"]["ok"]
+
+        graph_breaks = dict(torch._dynamo.utils.counters["graph_break"])
+        assert len(graph_breaks) == 0, f"Graph breaks detected: {graph_breaks}"
+
+        assert compilations_after_second == compilations_after_first, (
+            f"dynamic=True should prevent recompilation on shape change, "
+            f"but compilations went from {compilations_after_first} to {compilations_after_second}"
+        )
+
+    def test_optimized_module_wrappers(self, saved_model, single_gpu_device):
+        """Compiled blocks are OptimizedModule; uncompiled blocks are not."""
+        model_path, _ = saved_model
+
+        # Compiled path
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+        model_c = setup_model(
+            model_name_or_path=str(model_path),
+            use_liger_kernels=False,
+            osft=False,
+            local_rank=0,
+        )
+        model_c, _, _ = setup_training_components(
+            model_c,
+            learning_rate=1e-3,
+            num_warmup_steps=0,
+            lr_scheduler="constant",
+            compile_model=True,
+        )
+
+        from torch._dynamo.eval_frame import OptimizedModule
+
+        layers_c = model_c.model.layers
+        for idx, block in enumerate(layers_c):
+            assert isinstance(block, OptimizedModule), f"Block {idx} should be OptimizedModule, got {type(block)}"
+
+        del model_c
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        # Eager path
+        model_e = setup_model(
+            model_name_or_path=str(model_path),
+            use_liger_kernels=False,
+            osft=False,
+            local_rank=0,
+        )
+        model_e, _, _ = setup_training_components(
+            model_e,
+            learning_rate=1e-3,
+            num_warmup_steps=0,
+            lr_scheduler="constant",
+            compile_model=False,
+        )
+
+        layers_e = model_e.model.layers
+        for idx, block in enumerate(layers_e):
+            assert not isinstance(block, OptimizedModule), (
+                f"Block {idx} should NOT be OptimizedModule, got {type(block)}"
+            )
+
+    def test_compile_works_without_dynamo_config_flag(self, saved_model, single_gpu_device):
+        """AC + compile works without skip_fwd_side_effects_in_bwd_under_checkpoint.
+
+        The flag is set defensively in train.py but is not required on current
+        PyTorch. If this test starts failing on a future version, the flag
+        becomes load-bearing and the comment in train.py should be updated.
+        """
+        model_path, config = saved_model
+
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = False
+
+        model = setup_model(
+            model_name_or_path=str(model_path),
+            use_liger_kernels=False,
+            osft=False,
+            local_rank=0,
+        )
+        model, optimizer, _ = setup_training_components(
+            model,
+            learning_rate=1e-3,
+            num_warmup_steps=0,
+            lr_scheduler="constant",
+            compile_model=True,
+        )
+
+        input_ids = torch.randint(0, config.vocab_size, (2, 32), device=single_gpu_device)
+        labels = input_ids.clone()
+
+        optimizer.zero_grad()
+        output = model(input_ids=input_ids, labels=labels)
+        loss = output.loss.float().sum()
+        loss.backward()
+        optimizer.step()
diff --git a/tests/test_api_train.py b/tests/test_api_train.py
index e40c8832..158073cb 100644
--- a/tests/test_api_train.py
+++ b/tests/test_api_train.py
@@ -748,6 +748,7 @@ def test_all_boolean_flags_passed(self):
                 max_tokens_per_gpu=1000,
                 learning_rate=1e-5,
                 output_dir=tmpdir,
+                compile_model=True,
                 use_liger_kernels=True,
                 checkpoint_at_epoch=True,
                 save_final_checkpoint=True,
@@ -767,6 +768,7 @@ def test_all_boolean_flags_passed(self):
                 _, command = call_args[0]
 
                 # Verify all boolean flags are present
+                assert "--compile-model" in command
                 assert "--use-liger-kernels" in command
                 assert "--osft" in command
                 assert "--checkpoint-at-epoch" in command
@@ -783,6 +785,7 @@ def test_boolean_flags_not_passed_when_false(self):
                 max_tokens_per_gpu=1000,
                 learning_rate=1e-5,
                 output_dir=tmpdir,
+                compile_model=False,
                 use_liger_kernels=False,
                 osft=False,
                 checkpoint_at_epoch=False,
@@ -800,6 +803,7 @@ def test_boolean_flags_not_passed_when_false(self):
                 _, command = call_args[0]
 
                 # Verify boolean flags are NOT present when False
+                assert "--compile-model" not in command
                 assert "--use-liger-kernels" not in command
                 assert "--osft" not in command
                 assert "--checkpoint-at-epoch" not in command
diff --git a/tests/test_compile_guards.py b/tests/test_compile_guards.py
new file mode 100644
index 00000000..4e9954ea
--- /dev/null
+++ b/tests/test_compile_guards.py
@@ -0,0 +1,51 @@
+"""Unit tests for torch.compile validation guard contracts.
+
+These tests document the expected guard behavior from train.py:main().
+They test the guard conditions directly (not via main()) because main()
+requires a full distributed environment. The GPU tests in
+gpu_tests/test_compile.py exercise the real code path end-to-end.
+"""
+
+import pytest
+
+
+class TestCompileValidationGuards:
+    def test_compile_osft_incompatible(self):
+        compile_model = True
+        osft = True
+        with pytest.raises(ValueError, match="not compatible with --osft"):
+            if compile_model and osft:
+                raise ValueError(
+                    "--compile-model is not compatible with --osft. "
+                    "OSFT uses dynamic forward methods that cannot be traced by torch.compile."
+                )
+
+    def test_compile_liger_incompatible(self):
+        compile_model = True
+        use_liger_kernels = True
+        with pytest.raises(ValueError, match="not compatible with --use-liger-kernels"):
+            if compile_model and use_liger_kernels:
+                raise ValueError(
+                    "--compile-model is not compatible with --use-liger-kernels. "
+                    "Both replace the same memory-bound ops; the interaction is untested."
+                )
+
+    def test_compile_moe_incompatible(self):
+        moe_classes = ("MixtralForCausalLM", "GraniteMoeHybridForCausalLM")
+        for cls_name in moe_classes:
+            with pytest.raises(ValueError, match="not compatible with MoE"):
+                if cls_name in moe_classes:
+                    raise ValueError(
+                        f"--compile-model is not compatible with MoE architecture {cls_name}. "
+                        "MoE router logic causes graph breaks with fullgraph=True."
+                    )
+
+    def test_compile_guards_do_not_fire_when_disabled(self):
+        compile_model = False
+        osft = True
+        use_liger_kernels = True
+
+        if compile_model and osft:
+            raise ValueError("should not reach")
+        if compile_model and use_liger_kernels:
+            raise ValueError("should not reach")

From be569cb6684f169430b5d88c5ef9f197862bd53c Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 7 May 2026 05:14:08 +0000
Subject: [PATCH 05/19] Fix dynamo config comment and add compile_model to
 docstring

- 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.
---
 src/mini_trainer/setup_model_for_training.py | 2 ++
 src/mini_trainer/train.py                    | 4 +++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index ea84d1d5..5d7cdb81 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -427,6 +427,8 @@ def wrap_fsdp2(model: torch.nn.Module, compile_model: bool = False) -> torch.nn.
 
     Args:
         model: Model to wrap with FSDP2 (should already have buffers materialized)
+        compile_model: If True, compile each transformer block with torch.compile
+            (fullgraph=True, dynamic=True) between AC wrapping and FSDP2 sharding.
 
     Returns:
         FSDP2-wrapped model
diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index cd3ea92a..77257120 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -1584,7 +1584,9 @@ def main(
                 f"--compile-model is not compatible with MoE architecture {model.__class__.__name__}. "
                 "MoE router logic causes graph breaks with fullgraph=True."
             )
-        # Without this, AC's RNG side effects cause a graph break under compile.
+        # Defensive: not required on current PyTorch but may be needed on
+        # future versions where AC's RNG side effects cause graph breaks.
+        # See test_compile_works_without_dynamo_config_flag.
         torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
 
     # Create PretrainingConfig if block_size is provided

From c8bec87c05ba6e990d93a3dcefed52052277c7df Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Mon, 15 Jun 2026 19:38:11 +0000
Subject: [PATCH 06/19] Remove compile-hostile patterns from training loop
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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).
---
 src/mini_trainer/sampler.py        |  1 +
 src/mini_trainer/train.py          | 20 +++++++++++---------
 tests/test_model_initialization.py |  2 +-
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/mini_trainer/sampler.py b/src/mini_trainer/sampler.py
index d2702f60..3aa0d757 100644
--- a/src/mini_trainer/sampler.py
+++ b/src/mini_trainer/sampler.py
@@ -503,6 +503,7 @@ def padded_mb_collate_fn(minibatch: list[dict], batch_num_loss_counted_tokens: i
         }
 
     max_len = max(len(item["input_ids"]) for item in minibatch)
+    max_len = (max_len + 7) & ~7  # round up to multiple of 8 for torch.compile
 
     padded_input_ids = []
     padded_labels = []
diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index 77257120..d34e32e4 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -363,9 +363,6 @@ def compute_validation_loss(model, val_data_loader, device):
                 loss = output.loss.float().sum()
                 loss_metrics = loss.detach().item()
 
-                # Clear cache after each minibatch to prevent OOM
-                torch.cuda.empty_cache()
-
                 val_batch_totals.accumulate_minibatch_metrics(
                     num_loss_counted_tokens=mb_num_loss_counted_tokens,
                     num_total_tokens=mb["input_ids"].numel(),
@@ -943,23 +940,19 @@ def train(
 
                 # Ensure scalar loss even if model returns per-token loss
                 loss = (loss / batch_num_loss_counted_tokens) * world_size
-                loss_metrics = loss.detach().cpu().item()
+                loss_metrics = loss.detach().item()
                 loss.backward()
 
                 if callback_manager and callback_manager.has_callbacks("on_after_backward"):
                     callback_manager.context.loss = loss_metrics
                     callback_manager.fire("on_after_backward")
 
-                torch.cuda.empty_cache()
-
                 batch_totals.accumulate_minibatch_metrics(
                     num_loss_counted_tokens=mb_num_loss_counted_tokens,
                     num_total_tokens=mb["input_ids"].shape[1],
                     num_samples=mb_num_samples,
                     loss=loss_metrics,
-                    # since FSDP2 automatically averages gradients by the world-size,
-                    # each rank's gradient contributes 1/8 to the backward
-                    loss_backward=loss.detach().item() / world_size,
+                    loss_backward=loss_metrics / world_size,
                     time_per_minibatch=time.time() - mb_start_time,
                 )
 
@@ -1076,6 +1069,15 @@ def train(
 
             dist.barrier()
 
+            if step == 1 and is_local_main_process:
+                peak_gb = batch_metrics["peak_memory_usage_GB"]
+                gpu_total_gb = torch.cuda.get_device_properties(device).total_memory / 1e9
+                utilization = peak_gb / gpu_total_gb
+                log_rank_0(
+                    f"Memory after step 1: {peak_gb:.1f}GB / {gpu_total_gb:.1f}GB "
+                    f"({utilization:.0%} utilization)"
+                )
+
             # On-demand full-state checkpoint check
             if full_state_checkpointer is not None and full_state_checkpointer.should_save(device):
                 log_rank_0("Signal received — saving full-state checkpoint and exiting")
diff --git a/tests/test_model_initialization.py b/tests/test_model_initialization.py
index 7d4558ad..681e005d 100644
--- a/tests/test_model_initialization.py
+++ b/tests/test_model_initialization.py
@@ -297,7 +297,7 @@ def test_setup_training_components_basic(
         assert lr_scheduler == mock_lr_scheduler
 
         # Check FSDP2 wrapping
-        mock_wrap.assert_called_once_with(mock_model)
+        mock_wrap.assert_called_once_with(mock_model, compile_model=False)
 
         # Check optimizer creation
         mock_adamw.assert_called_once_with(

From a89c3e6e99dd95b35e16f934964c40e5a0f590a2 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Mon, 15 Jun 2026 19:38:11 +0000
Subject: [PATCH 07/19] Fix compile recompiles: mod-8 padding and skip cache
 shape guards
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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
---
 src/mini_trainer/sampler.py | 6 ++++++
 src/mini_trainer/train.py   | 1 +
 2 files changed, 7 insertions(+)

diff --git a/src/mini_trainer/sampler.py b/src/mini_trainer/sampler.py
index 3aa0d757..3d1ee193 100644
--- a/src/mini_trainer/sampler.py
+++ b/src/mini_trainer/sampler.py
@@ -455,6 +455,12 @@ def mb_collate_fn(minibatch, batch_num_loss_counted_tokens):
     #     f"num_loss_counted_tokens: {num_loss_counted_tokens}\033[0m"
     # )
 
+    pad_len = (8 - total_len % 8) % 8
+    if pad_len > 0:
+        input_ids.extend([0] * pad_len)
+        labels.extend([-100] * pad_len)
+        position_ids.extend(range(pad_len))
+
     return {
         "input_ids": torch.tensor([input_ids], dtype=torch.long),
         "labels": torch.tensor([labels], dtype=torch.long),
diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index d34e32e4..6f514e03 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -1590,6 +1590,7 @@ def main(
         # future versions where AC's RNG side effects cause graph breaks.
         # See test_compile_works_without_dynamo_config_flag.
         torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+        torch._inductor.config.unsafe_skip_cache_dynamic_shape_guards = True
 
     # Create PretrainingConfig if block_size is provided
     pretraining_config = None

From 9bf5209390392dbe1474f53e67e105137ef1a5d1 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Mon, 15 Jun 2026 19:38:11 +0000
Subject: [PATCH 08/19] Fix CI: guard CUDA call, update tests for mod-8 padding

- 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)
---
 src/mini_trainer/train.py   |  5 ++---
 tests/test_data_loader.py   | 20 ++++++++++----------
 tests/test_training_loop.py |  1 -
 3 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index 6f514e03..d6d81973 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -1069,13 +1069,12 @@ def train(
 
             dist.barrier()
 
-            if step == 1 and is_local_main_process:
+            if step == 1 and is_local_main_process and torch.cuda.is_available():
                 peak_gb = batch_metrics["peak_memory_usage_GB"]
                 gpu_total_gb = torch.cuda.get_device_properties(device).total_memory / 1e9
                 utilization = peak_gb / gpu_total_gb
                 log_rank_0(
-                    f"Memory after step 1: {peak_gb:.1f}GB / {gpu_total_gb:.1f}GB "
-                    f"({utilization:.0%} utilization)"
+                    f"Memory after step 1: {peak_gb:.1f}GB / {gpu_total_gb:.1f}GB ({utilization:.0%} utilization)"
                 )
 
             # On-demand full-state checkpoint check
diff --git a/tests/test_data_loader.py b/tests/test_data_loader.py
index 263c0d1f..6f7451b6 100644
--- a/tests/test_data_loader.py
+++ b/tests/test_data_loader.py
@@ -132,10 +132,10 @@ def test_collate_single_sample(self):
 
         result = mb_collate_fn(minibatch, batch_num_loss_counted_tokens=4)
 
-        assert result["input_ids"].shape == (1, 5)
-        assert result["labels"].shape == (1, 5)
-        assert result["position_ids"].shape == (1, 5)
-        assert result["position_ids"].tolist() == [[0, 1, 2, 3, 4]]
+        assert result["input_ids"].shape == (1, 8)
+        assert result["labels"].shape == (1, 8)
+        assert result["position_ids"].shape == (1, 8)
+        assert result["position_ids"].tolist() == [[0, 1, 2, 3, 4, 0, 1, 2]]
         assert result["num_loss_counted_tokens"] == 4
         assert result["num_samples"] == 1
         assert result["batch_num_loss_counted_tokens"] == 4
@@ -157,13 +157,13 @@ def test_collate_multiple_samples(self):
 
         result = mb_collate_fn(minibatch, batch_num_loss_counted_tokens=6)
 
-        # Check concatenation
-        assert result["input_ids"].shape == (1, 7)
-        assert result["input_ids"].tolist() == [[1, 2, 3, 4, 5, 6, 7]]
-        assert result["labels"].tolist() == [[10, 20, 30, 40, -100, 60, 70]]
+        # Check concatenation (padded to mod-8)
+        assert result["input_ids"].shape == (1, 8)
+        assert result["input_ids"].tolist() == [[1, 2, 3, 4, 5, 6, 7, 0]]
+        assert result["labels"].tolist() == [[10, 20, 30, 40, -100, 60, 70, -100]]
 
-        # Check position_ids reset for each sequence
-        assert result["position_ids"].tolist() == [[0, 1, 2, 0, 1, 2, 3]]
+        # Check position_ids reset for each sequence (pad token starts new mini-sequence)
+        assert result["position_ids"].tolist() == [[0, 1, 2, 0, 1, 2, 3, 0]]
 
         assert result["num_loss_counted_tokens"] == 6
         assert result["num_samples"] == 2
diff --git a/tests/test_training_loop.py b/tests/test_training_loop.py
index 315b3de4..78269f43 100644
--- a/tests/test_training_loop.py
+++ b/tests/test_training_loop.py
@@ -949,7 +949,6 @@ def test_memory_tracking(
 
         # Verify memory management calls
         mock_reset_stats.assert_called()
-        mock_empty_cache.assert_called()
 
         # Verify memory was tracked in metrics
         logged_metrics = mock_logger.log_sync.call_args[0][0]

From c41645cf9ff9ede2c268fc5ab9e35adc68d5a331 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 24 Jun 2026 14:42:11 +0000
Subject: [PATCH 09/19] Enable torch.compile for OSFT via OSFTLinear module
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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
---
 benchmarks/bench_osft_compile.py | 219 ++++++++++++++++++++++
 src/mini_trainer/osft_utils.py   | 301 ++++++++++++-------------------
 src/mini_trainer/train.py        |   6 -
 tests/gpu_tests/test_compile.py  | 152 +++++++++++++++-
 tests/test_compile_guards.py     |  13 --
 tests/test_osft.py               |   8 +-
 tests/test_osft_fidelity.py      |  55 ++----
 7 files changed, 505 insertions(+), 249 deletions(-)
 create mode 100644 benchmarks/bench_osft_compile.py

diff --git a/benchmarks/bench_osft_compile.py b/benchmarks/bench_osft_compile.py
new file mode 100644
index 00000000..b989a0c6
--- /dev/null
+++ b/benchmarks/bench_osft_compile.py
@@ -0,0 +1,219 @@
+"""Benchmark: OSFT eager vs compiled forward/backward.
+
+Profiles OSFT training steps with torch.profiler to identify graph breaks
+and measure compile speedup.
+
+Usage:
+  # Eager baseline
+  python bench_osft_compile.py
+
+  # Compiled
+  python bench_osft_compile.py --compile
+
+  # With Chrome trace export
+  python bench_osft_compile.py --compile --trace-dir benchmarks/traces
+
+  # dynamo.explain report (graph break analysis)
+  python bench_osft_compile.py --explain
+"""
+
+import argparse
+import os
+import time
+
+os.environ["TESTING"] = "true"
+
+import torch
+import torch.distributed as dist
+from torch.profiler import ProfilerActivity, profile, schedule
+from transformers import LlamaConfig, LlamaForCausalLM
+
+from mini_trainer.none_reduction_losses import hf_fixed_cross_entropy_none_reduction
+from mini_trainer.setup_model_for_training import setup_model, setup_training_components
+from mini_trainer.utils import patch_target_module
+
+
+def create_tiny_llama(tmp_dir):
+    config = LlamaConfig(
+        vocab_size=1000,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=128,
+        rope_theta=10000.0,
+        hidden_act="silu",
+    )
+    model = LlamaForCausalLM(config)
+    model.save_pretrained(tmp_dir)
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained("gpt2")
+    tok.pad_token = tok.eos_token
+    tok.save_pretrained(tmp_dir)
+    return tmp_dir
+
+
+def setup_dist():
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "12399")
+    dist.init_process_group(backend="nccl", rank=0, world_size=1)
+    patch_target_module(
+        "transformers.loss.loss_utils.fixed_cross_entropy",
+        hf_fixed_cross_entropy_none_reduction,
+    )
+
+
+def build_model(model_path, compile_model, osft=True, osft_rank_ratio=0.25):
+    model = setup_model(
+        model_name_or_path=model_path,
+        use_liger_kernels=False,
+        osft=osft,
+        osft_rank_ratio=osft_rank_ratio if osft else None,
+        local_rank=0,
+    )
+    model, optimizer, lr_scheduler = setup_training_components(
+        model,
+        learning_rate=1e-3,
+        num_warmup_steps=0,
+        lr_scheduler="constant",
+        compile_model=compile_model,
+    )
+    return model, optimizer, lr_scheduler
+
+
+def run_steps(model, optimizer, lr_scheduler, input_ids, labels, num_steps):
+    losses = []
+    for _ in range(num_steps):
+        optimizer.zero_grad()
+        output = model(input_ids=input_ids, labels=labels)
+        loss = output.loss.float().sum() / input_ids.shape[0]
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        lr_scheduler.step()
+        losses.append(loss.item())
+    return losses
+
+
+def run_explain(model_path):
+    """Run torch._dynamo.explain on an OSFT model to report graph breaks."""
+    model, optimizer, lr_scheduler = build_model(model_path, compile_model=False, osft=True)
+
+    input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+    labels = input_ids.clone()
+
+    explanation = torch._dynamo.explain(model)(input_ids=input_ids, labels=labels)
+    print("\n" + "=" * 80)
+    print("torch._dynamo.explain report")
+    print("=" * 80)
+    print(explanation)
+    print("=" * 80)
+
+
+def run_profile(model_path, compile_model, trace_dir, num_warmup=3, num_active=5):
+    mode_str = "compiled" if compile_model else "eager"
+    osft_str = "osft"
+    print(f"\nProfiling {osft_str} {mode_str}...")
+
+    model, optimizer, lr_scheduler = build_model(model_path, compile_model=compile_model, osft=True)
+
+    input_ids = torch.randint(0, 1000, (2, 32), device="cuda")
+    labels = input_ids.clone()
+
+    # Warmup (includes compilation for compiled mode)
+    print(f"  Warmup: {num_warmup} steps...")
+    t0 = time.perf_counter()
+    run_steps(model, optimizer, lr_scheduler, input_ids, labels, num_warmup)
+    torch.cuda.synchronize()
+    warmup_time = time.perf_counter() - t0
+    print(f"  Warmup done in {warmup_time:.2f}s")
+
+    # Profiled steps
+    print(f"  Profiling: {num_active} steps...")
+    trace_path = None
+    if trace_dir:
+        os.makedirs(trace_dir, exist_ok=True)
+        trace_path = os.path.join(trace_dir, f"osft_{mode_str}.json")
+
+    with profile(
+        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+        record_shapes=True,
+        with_stack=True,
+        schedule=schedule(wait=0, warmup=1, active=num_active - 1, repeat=1),
+    ) as prof:
+        for _ in range(num_active):
+            optimizer.zero_grad()
+            output = model(input_ids=input_ids, labels=labels)
+            loss = output.loss.float().sum() / input_ids.shape[0]
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            lr_scheduler.step()
+            prof.step()
+
+    torch.cuda.synchronize()
+
+    if trace_path:
+        prof.export_chrome_trace(trace_path)
+        print(f"  Chrome trace: {trace_path}")
+
+    print(f"\n  === {osft_str} {mode_str} — Top CUDA ops ===")
+    print(
+        prof.key_averages().table(
+            sort_by="cuda_time_total", row_limit=20
+        )
+    )
+
+    # Wall-clock timing (separate from profiler)
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    run_steps(model, optimizer, lr_scheduler, input_ids, labels, 10)
+    torch.cuda.synchronize()
+    wall = time.perf_counter() - t0
+    print(f"\n  Wall-clock: 10 steps in {wall:.3f}s ({wall / 10 * 1000:.1f} ms/step)")
+
+    return prof
+
+
+def main():
+    parser = argparse.ArgumentParser(description="OSFT compile benchmark")
+    parser.add_argument("--compile", action="store_true", help="Enable torch.compile")
+    parser.add_argument("--explain", action="store_true", help="Run torch._dynamo.explain")
+    parser.add_argument("--trace-dir", type=str, default=None, help="Directory for Chrome traces")
+    parser.add_argument("--both", action="store_true", help="Run both eager and compiled for comparison")
+    args = parser.parse_args()
+
+    import tempfile
+
+    tmpdir = tempfile.mkdtemp()
+
+    torch.manual_seed(42)
+    setup_dist()
+    model_path = create_tiny_llama(tmpdir)
+
+    try:
+        if args.explain:
+            run_explain(model_path)
+        elif args.both:
+            run_profile(model_path, compile_model=False, trace_dir=args.trace_dir)
+            torch._dynamo.reset()
+            dist.destroy_process_group()
+            setup_dist()
+            run_profile(model_path, compile_model=True, trace_dir=args.trace_dir)
+        else:
+            run_profile(model_path, compile_model=args.compile, trace_dir=args.trace_dir)
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        import shutil
+
+        shutil.rmtree(tmpdir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/mini_trainer/osft_utils.py b/src/mini_trainer/osft_utils.py
index 29c80b96..465190ca 100644
--- a/src/mini_trainer/osft_utils.py
+++ b/src/mini_trainer/osft_utils.py
@@ -5,7 +5,6 @@
 import typing as t
 from dataclasses import dataclass
 from typing import Protocol
-from weakref import ref as weakref
 
 import numpy as np
 import torch
@@ -134,6 +133,73 @@ class SVDDecompositionDict(SVDDictBase, total=False):
     rank_high: int
 
 
+class OSFTLinear(nn.Module):
+    """Factorized linear using SVD components: W = U_high @ diag(S_high) @ V_high + U_low @ diag(S_low) @ V_low.
+
+    Replaces closure-based forward overrides with a proper nn.Module whose
+    forward is pure tensor math — no weakrefs, dict construction, getattr,
+    or runtime .to() calls — so torch.compile can trace it without graph breaks.
+
+    Attribute names (osft_U_high, osft_S_high, osft_V_high, osft_params.U_low,
+    osft_params.S_low, osft_params.V_low) match the existing convention so that
+    project_gradients, project_parameters, get_svd_dict_for_module, and
+    prepare_state_dict_for_save continue to work unchanged.
+    """
+
+    def __init__(self, U_high, S_high, V_high, U_low, S_low, V_low, bias=None, rank_high=None):
+        super().__init__()
+        self.osft_U_high = nn.Parameter(U_high, requires_grad=False)
+        self.osft_S_high = nn.Parameter(S_high, requires_grad=False)
+        self.osft_V_high = nn.Parameter(V_high, requires_grad=False)
+        osft_params = nn.Module()
+        osft_params.U_low = U_low
+        osft_params.S_low = S_low
+        osft_params.V_low = V_low
+        self.osft_params = osft_params
+        # rank_high: full (unsharded) k_high dimension, needed by V all-gather
+        # in project_gradients. Stored as a persistent int buffer so it survives
+        # FSDP2 sharding (S_high.shape[0] becomes the local shard size, not full).
+        if rank_high is not None:
+            # Infer device from an existing parameter so the buffer lives on the
+            # same device (meta during lazy init, CPU/CUDA during standard init).
+            buf_device = U_high.device
+            self.register_buffer(
+                "_rank_high_buf",
+                torch.tensor(rank_high, dtype=torch.long, device=buf_device),
+                persistent=True,
+            )
+            self._rank_high_cached: int | None = rank_high
+        else:
+            self._rank_high_buf = None
+            self._rank_high_cached = None
+        if bias is not None:
+            self.bias = nn.Parameter(bias.data, requires_grad=bias.requires_grad)
+        else:
+            self.bias = None
+
+    @property
+    def rank_high(self) -> int:
+        if self._rank_high_cached is not None:
+            return self._rank_high_cached
+        # Lazy extract after set_model_state_dict materializes the buffer.
+        # Single .item() call, then cached for all subsequent accesses.
+        buf = self._rank_high_buf
+        if buf is not None and buf.device.type != "meta":
+            self._rank_high_cached = buf.item()
+            return self._rank_high_cached
+        return self.osft_S_high.shape[0]
+
+    def forward(self, x):
+        x_V_high = x @ self.osft_V_high.transpose(0, 1)
+        result_high = (x_V_high * self.osft_S_high) @ self.osft_U_high.transpose(0, 1)
+        x_V_low = x @ self.osft_params.V_low.transpose(0, 1)
+        result_low = (x_V_low * self.osft_params.S_low) @ self.osft_params.U_low.transpose(0, 1)
+        result = result_high + result_low
+        if self.bias is not None:
+            result = result + self.bias
+        return result
+
+
 class OSFTModelProtocol(Protocol):
     """
     Protocol defining the interface for models with OSFT capabilities.
@@ -1166,7 +1232,6 @@ def __init__(
             # create a set of logical keys
             self.lazy_init_param_keys = lazy_init_param_keys if lazy_init_param_keys else []
             self.lazy_init_buffer_dict = lazy_init_buffer_dict if lazy_init_buffer_dict else {}
-            self._osft_handles: dict[str, tuple[weakref[nn.Module], str]] = {}
             self.logical_osft_keys = []
             self.orig_param_registry: dict[str, ParamSpec] = {}  # stores all of the original params
             self.osft_paramspec_registry: dict[str, OSFTFactorSpec] = {}
@@ -1196,8 +1261,6 @@ def _reset_osft_metadata(self):
             self.logical_osft_keys = []
             self.orig_param_registry = {}
             self.osft_paramspec_registry = {}
-            self._osft_handles = {}
-            self.osft_params = {}
             # Clear any cached all-gathered V_high — V_high changes on reinit.
             for module in self.modules():
                 if hasattr(module, "_osft_v_high_full"):
@@ -1589,7 +1652,7 @@ def ensure_dtype(tensor, expected_dtype):
                             osft_spec.U_high: ensure_dtype(svd_dict["U_high"], expected_dtype),
                             osft_spec.S_high: ensure_dtype(svd_dict["S_high"], expected_dtype),
                             osft_spec.V_high: ensure_dtype(svd_dict["V_high"], expected_dtype),
-                            osft_spec.rank_high: svd_dict["rank_high"],
+                            osft_spec.rank_high: torch.tensor(svd_dict["rank_high"], dtype=torch.long),
                         }
                     )
 
@@ -1733,12 +1796,7 @@ def _pre_fsdp2_wrap_initialize_lazy_osft(self):
                 if is_osft_param:
                     self.logical_osft_keys.append(pk)
 
-            # next, we actually register these keys
-            # and replace them with the OSFT equivalents
             for key in self.logical_osft_keys:
-                # here we build the association to the original key
-                mod, attr = self._get_module_by_name(key)
-                self._register_osft_target(key, mod, attr)
                 self._prepare_osft_param(key)
 
         def reinitialize_osft(self, decompose_existing_weights: bool, assigned_params=None):
@@ -1795,23 +1853,6 @@ def mark_fsdp2_initialized(self):
             self._lazy_init_pending = False
             set_fsdp2_lazy_init_mode(self, None)
 
-        def _get_module_by_logical_key(self, logical_key: str):
-            """Return (module, attr) using the stable handle; independent of FQNs/wrappers."""
-            try:
-                wr, attr = self._osft_handles[logical_key]
-            except KeyError:
-                return None, None
-            mod = wr()
-            return (mod, attr) if mod is not None else (None, None)
-
-        # call this BEFORE activation checkpointing / fully_shard
-        def _register_osft_target(self, logical_key: str, module: nn.Module, attr: str):
-            """Record a stable handle to the parent module + attribute name for a target param."""
-            # Optional: tag the Parameter for debugging/validation
-            p = getattr(module, attr)
-            p._osft_key = logical_key
-            self._osft_handles[logical_key] = (weakref(module), attr)
-
         def _record_osft_factor_spec(self, logical_key: str, attr: str) -> OSFTFactorSpec:
             """Create and store the factor spec describing where OSFT tensors live."""
             parent_logical_key = logical_key.rsplit(".", 1)[0] if "." in logical_key else ""
@@ -1827,7 +1868,7 @@ def _compose(parent: str, suffix: str) -> str:
                 U_low=_compose(parent_logical_key, "osft_params.U_low"),
                 S_low=_compose(parent_logical_key, "osft_params.S_low"),
                 V_low=_compose(parent_logical_key, "osft_params.V_low"),
-                rank_high=_compose(parent_logical_key, "osft_params.rank_high"),
+                rank_high=_compose(parent_logical_key, "_rank_high_buf"),
             )
             self.osft_paramspec_registry[logical_key] = spec
             return spec
@@ -1846,15 +1887,13 @@ def eject_og_state_dict(self):
 
         def _prepare_osft_param(self, logical_key: str):
             """
-            Prepares an OSFT parameter by initializing an OSFT module at the given
-            key and removes the actual weight.
+            Prepares an OSFT parameter by replacing the Linear module with an
+            OSFTLinear at the same position in the module tree.
             """
-            mod_ref, attr = self._osft_handles[logical_key]
-            mod = mod_ref()
+            mod, attr = self._get_module_by_name(logical_key)
             if mod is None:
-                raise ValueError(f"requested module {logical_key} but ref is None")
+                raise ValueError(f"requested module {logical_key} but could not be found")
 
-            # next we register the parameter and remove the attribute
             meta_weight = getattr(mod, attr)
             top_K = self.osft_config[logical_key]
             svd_dict = create_svd_dict(
@@ -1866,51 +1905,26 @@ def _prepare_osft_param(self, logical_key: str):
                 output_dtype=meta_weight.dtype,
             )
 
-            # next, we create a new module and register the parameters
-            # TODO: move these no-grad params onto the SVD module
-            mod.register_parameter("osft_U_high", nn.Parameter(svd_dict["U_high"], requires_grad=False))
-            mod.register_parameter("osft_S_high", nn.Parameter(svd_dict["S_high"], requires_grad=False))
-            mod.register_parameter("osft_V_high", nn.Parameter(svd_dict["V_high"], requires_grad=False))
-
-            # Trainable low-rank components
-            module_svd = nn.Module()
-            module_svd.U_low = svd_dict["U_low"]
-            module_svd.S_low = svd_dict["S_low"]
-            module_svd.V_low = svd_dict["V_low"]
-            module_svd.rank_high = svd_dict["rank_high"]
-
-            # this we want to improve
-            mod.add_module("osft_params", module_svd)
-
-            # Override linear projection to use module-local OSFT params
-            # Note: we use the logical key to look up the module dynamically via the handle registry
-            # to ensure the reference survives FSDP2 wrapping and activation checkpointing
-            def make_forward(lkey):
-                def forward(x):
-                    owner_mod, _ = self._get_module_by_logical_key(lkey)
-                    if owner_mod is None:
-                        raise RuntimeError(f"Module for logical key '{lkey}' not found in handle registry")
-                    svd_dict = {
-                        "U_high": owner_mod.osft_U_high,
-                        "S_high": owner_mod.osft_S_high,
-                        "V_high": owner_mod.osft_V_high,
-                        "U_low": owner_mod.osft_params.U_low,
-                        "S_low": owner_mod.osft_params.S_low,
-                        "V_low": owner_mod.osft_params.V_low,
-                        "rank_high": owner_mod.osft_params.rank_high,
-                    }
-                    # retrieve bias dynamically to avoid meta tensor issues
-                    bias = getattr(owner_mod, "bias", None)
-                    return self._factorized_linear(x, svd_dict, bias)
-
-                return forward
-
-            # update the forward
-            mod.forward = make_forward(logical_key)
-            meta_weight.requires_grad = False
-            self._record_osft_factor_spec(logical_key, attr)
+            bias = getattr(mod, "bias", None)
+            osft_linear = OSFTLinear(
+                U_high=svd_dict["U_high"],
+                S_high=svd_dict["S_high"],
+                V_high=svd_dict["V_high"],
+                U_low=svd_dict["U_low"],
+                S_low=svd_dict["S_low"],
+                V_low=svd_dict["V_low"],
+                bias=bias,
+                rank_high=svd_dict["rank_high"],
+            )
+
+            safe_name = logical_key.replace(".", "_")
+            self.name_mapping[logical_key] = safe_name
+            osft_linear.osft_params.safe_name = safe_name
 
-            mod._parameters.pop(attr)
+            mod_path = logical_key.rsplit(".", 1)[0]
+            parent, child_name = self._get_module_by_name(mod_path)
+            setattr(parent, child_name, osft_linear)
+            self._record_osft_factor_spec(logical_key, attr)
 
         @torch.no_grad()
         def process_param_into_svd_dict(self, param: torch.Tensor, name: str) -> SVDDecompositionDict:
@@ -2029,10 +2043,17 @@ def _initialize_osft_parameters(self, decompose_existing_weights: bool, assigned
                         output_dtype=self.output_dtype,
                     )
 
-                    # Move SVD components to target device and clear GPU cache
+                    # Move SVD results back to the original parameter's device.
+                    # SVD runs on GPU for speed, but the module should live where
+                    # its predecessor lived.  FSDP2 handles final placement.
+                    # Tensor.to() across devices strips nn.Parameter, so re-wrap.
+                    orig_device = param.device
                     for key in svd_dict:
                         if isinstance(svd_dict[key], torch.Tensor):
-                            svd_dict[key] = svd_dict[key].to(target_device)
+                            moved = svd_dict[key].to(orig_device)
+                            if isinstance(svd_dict[key], nn.Parameter) and not isinstance(moved, nn.Parameter):
+                                moved = nn.Parameter(moved, requires_grad=svd_dict[key].requires_grad)
+                            svd_dict[key] = moved
 
                     # Clear the temporary GPU + CPU tensor
                     del param_gpu
@@ -2047,63 +2068,25 @@ def _initialize_osft_parameters(self, decompose_existing_weights: bool, assigned
                     safe_name = name.replace(".", "_")
                     self.name_mapping[name] = safe_name
 
-                    # Attach OSFT components to the owning module so only block-local params materialize
                     mod, attr = self._get_module_by_name(name)
-
-                    # Register this target in the handle registry for stable lookups
-                    self._register_osft_target(name, mod, attr)
-
-                    # High-rank frozen components
-                    mod.register_parameter(
-                        "osft_U_high",
-                        nn.Parameter(svd_dict["U_high"], requires_grad=False),
-                    )
-                    mod.register_parameter(
-                        "osft_S_high",
-                        nn.Parameter(svd_dict["S_high"], requires_grad=False),
-                    )
-                    mod.register_parameter(
-                        "osft_V_high",
-                        nn.Parameter(svd_dict["V_high"], requires_grad=False),
+                    bias = getattr(mod, "bias", None)
+                    osft_linear = OSFTLinear(
+                        U_high=svd_dict["U_high"],
+                        S_high=svd_dict["S_high"],
+                        V_high=svd_dict["V_high"],
+                        U_low=svd_dict["U_low"],
+                        S_low=svd_dict["S_low"],
+                        V_low=svd_dict["V_low"],
+                        bias=bias,
+                        rank_high=svd_dict["rank_high"],
                     )
-                    # Trainable low-rank components
-                    module_svd = nn.Module()
-                    module_svd.U_low = svd_dict["U_low"]
-                    module_svd.S_low = svd_dict["S_low"]
-                    module_svd.V_low = svd_dict["V_low"]
-                    module_svd.rank_high = svd_dict["rank_high"]
-                    module_svd.safe_name = safe_name
-                    mod.add_module("osft_params", module_svd)
-
-                    # Override linear projection to use module-local OSFT params
-                    # Note: we use the logical key to look up the module dynamically via the handle registry
-                    # to ensure the reference survives FSDP2 wrapping and activation checkpointing
-                    def make_forward(lkey):
-                        def forward(x):
-                            owner_mod, _ = self._get_module_by_logical_key(lkey)
-                            if owner_mod is None:
-                                raise RuntimeError(f"Module for logical key '{lkey}' not found in handle registry")
-                            svd_dict = {
-                                "U_high": owner_mod.osft_U_high,
-                                "S_high": owner_mod.osft_S_high,
-                                "V_high": owner_mod.osft_V_high,
-                                "U_low": owner_mod.osft_params.U_low,
-                                "S_low": owner_mod.osft_params.S_low,
-                                "V_low": owner_mod.osft_params.V_low,
-                                "rank_high": owner_mod.osft_params.rank_high,
-                            }
-                            # retrieve bias dynamically to avoid meta tensor issues
-                            bias = getattr(owner_mod, "bias", None)
-                            return self._factorized_linear(x, svd_dict, bias)
-
-                        return forward
-
-                    mod.forward = make_forward(name)
-                    param.requires_grad = False
-                    # Remove original parameter so it doesn't get updated
-                    mod._parameters.pop(attr, None)
+                    osft_linear.osft_params.safe_name = safe_name
+
+                    # Replace Linear with OSFTLinear in the parent module
+                    mod_path = name.rsplit(".", 1)[0]
+                    parent, child_name = self._get_module_by_name(mod_path)
+                    setattr(parent, child_name, osft_linear)
                     self._record_osft_factor_spec(name, attr)
-                    torch.cuda.empty_cache()
 
                     osft_params_processed += 1
 
@@ -2147,49 +2130,6 @@ def _reconstruct_weight(
             svd_dict = self.get_svd_dict_for_module(mod)
             return reconstruct_weight_matrix(svd_dict, upcast_dtype=upcast_dtype, output_dtype=output_dtype)
 
-        def _factorized_linear(self, x, svd_dict, bias=None):
-            """
-            Efficient factorized linear operation using SVD components.
-
-            Computes: x @ (U_high @ S_high @ V_high + U_low @ S_low @ V_low).T + bias
-            As: (x @ V_high.T) @ (S_high * U_high).T + (x @ V_low.T) @ (S_low * U_low).T
-            """
-            # Extract components
-            U_high = svd_dict["U_high"]
-            S_high = svd_dict["S_high"]
-            V_high = svd_dict["V_high"]
-            U_low = svd_dict["U_low"]
-            S_low = svd_dict["S_low"]
-            V_low = svd_dict["V_low"]
-
-            device = x.device
-            dtype = x.dtype
-
-            # Move to correct device (keep native dtype)
-            U_high = U_high.to(device=device)
-            S_high = S_high.to(device=device)
-            V_high = V_high.to(device=device)
-            U_low = U_low.to(device=device)
-            S_low = S_low.to(device=device)
-            V_low = V_low.to(device=device)
-
-            # High-rank path (frozen): x @ V_high.T -> (batch, seq, rank_high)
-            x_V_high = x @ V_high.transpose(0, 1)
-            result_high = (x_V_high * S_high) @ U_high.transpose(0, 1)
-
-            # Low-rank path (trainable): x @ V_low.T -> (batch, seq, rank_low)
-            x_V_low = x @ V_low.transpose(0, 1)
-            result_low = (x_V_low * S_low) @ U_low.transpose(0, 1)
-
-            # Combine both paths
-            result = result_high + result_low
-
-            # Add bias if present
-            if bias is not None:
-                result = result + bias.to(device=device, dtype=dtype)
-
-            return result
-
         def get_svd_dict_for_module(self, module) -> SVDDecompositionDict:
             if not hasattr(module, "osft_params"):
                 raise ValueError("Module does not have OSFT parameters attached")
@@ -2199,16 +2139,14 @@ def get_svd_dict_for_module(self, module) -> SVDDecompositionDict:
             ):
                 raise ValueError("Module is missing OSFT high-rank tensors (U/S/V_high)")
             module_svd = module.osft_params
-            S_high = module.osft_S_high
-            rank_high = S_high.shape[0]
             svd_dict: SVDDecompositionDict = {
                 "U_high": module.osft_U_high,
-                "S_high": S_high,
+                "S_high": module.osft_S_high,
                 "V_high": module.osft_V_high,
                 "U_low": module_svd.U_low,
                 "S_low": module_svd.S_low,
                 "V_low": module_svd.V_low,
-                "rank_high": rank_high,
+                "rank_high": module.rank_high,
             }
             return svd_dict
 
@@ -2451,6 +2389,7 @@ def prepare_state_dict_for_save(self, state_dict):
                 U_low = state_dict.pop(osft_factors.U_low)
                 S_low = state_dict.pop(osft_factors.S_low)
                 V_low = state_dict.pop(osft_factors.V_low)
+                state_dict.pop(osft_factors.rank_high, None)
                 W = reconstruct_weight_matrix(
                     {
                         "U_high": U_high,
diff --git a/src/mini_trainer/train.py b/src/mini_trainer/train.py
index d6d81973..a8daa3bb 100644
--- a/src/mini_trainer/train.py
+++ b/src/mini_trainer/train.py
@@ -1408,12 +1408,6 @@ def main(
 
     # validation, do this before continuing execution flow so we don't log experiments that are invalid from
     # the get-go
-    if compile_model and osft:
-        raise ValueError(
-            "--compile-model is not compatible with --osft. "
-            "OSFT uses dynamic forward methods that cannot be traced by torch.compile."
-        )
-
     if compile_model and use_liger_kernels:
         raise ValueError(
             "--compile-model is not compatible with --use-liger-kernels. "
diff --git a/tests/gpu_tests/test_compile.py b/tests/gpu_tests/test_compile.py
index 5c932613..0fd3dae1 100644
--- a/tests/gpu_tests/test_compile.py
+++ b/tests/gpu_tests/test_compile.py
@@ -30,12 +30,13 @@ def create_tiny_llama_model():
     return LlamaForCausalLM(config), config
 
 
-def _run_steps(model_path, compile_model, input_ids, labels, num_steps=3):
+def _run_steps(model_path, compile_model, input_ids, labels, num_steps=3, osft=False, osft_rank_ratio=0.25):
     """Load model, wrap with FSDP2 (± compile), run steps, return losses."""
     model = setup_model(
         model_name_or_path=str(model_path),
         use_liger_kernels=False,
-        osft=False,
+        osft=osft,
+        osft_rank_ratio=osft_rank_ratio if osft else None,
         local_rank=0,
     )
     model, optimizer, lr_scheduler = setup_training_components(
@@ -277,3 +278,150 @@ def test_compile_works_without_dynamo_config_flag(self, saved_model, single_gpu_
         loss = output.loss.float().sum()
         loss.backward()
         optimizer.step()
+
+
+@pytest.mark.gpu
+class TestOSFTCompile:
+    @pytest.fixture(autouse=True, scope="class")
+    def dist_env(self):
+        os.environ["RANK"] = "0"
+        os.environ["WORLD_SIZE"] = "1"
+        os.environ["LOCAL_RANK"] = "0"
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = "12357"
+        dist.init_process_group(backend="nccl", rank=0, world_size=1)
+
+        from mini_trainer.none_reduction_losses import (
+            hf_fixed_cross_entropy_none_reduction,
+        )
+
+        patch_target_module(
+            "transformers.loss.loss_utils.fixed_cross_entropy",
+            hf_fixed_cross_entropy_none_reduction,
+        )
+
+        yield
+
+        dist.destroy_process_group()
+
+    @pytest.fixture(autouse=True)
+    def reset_dynamo(self):
+        torch._dynamo.reset()
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = False
+        yield
+        torch._dynamo.reset()
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = False
+
+    @pytest.fixture
+    def saved_model(self, tmp_path):
+        torch.manual_seed(42)
+        model, config = create_tiny_llama_model()
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        tokenizer.pad_token = tokenizer.eos_token
+        model_path = tmp_path / "tiny_llama"
+        model.save_pretrained(model_path)
+        tokenizer.save_pretrained(model_path)
+        return model_path, config
+
+    def test_osft_compiled_matches_eager(self, saved_model, single_gpu_device):
+        """Compiled and eager OSFT produce approximately equal losses."""
+        model_path, config = saved_model
+
+        torch.manual_seed(99)
+        input_ids = torch.randint(0, config.vocab_size, (2, 32), device=single_gpu_device)
+        labels = input_ids.clone()
+
+        torch.manual_seed(7)
+        torch.cuda.manual_seed(7)
+        eager_losses, eager_model = _run_steps(
+            model_path, compile_model=False, input_ids=input_ids, labels=labels, osft=True,
+        )
+
+        del eager_model
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch._dynamo.reset()
+
+        torch.manual_seed(7)
+        torch.cuda.manual_seed(7)
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+        compiled_losses, _ = _run_steps(
+            model_path, compile_model=True, input_ids=input_ids, labels=labels, osft=True,
+        )
+
+        for step, (e, c) in enumerate(zip(eager_losses, compiled_losses)):
+            assert abs(e - c) < 0.1, (
+                f"Step {step}: eager loss {e:.6f} vs compiled loss {c:.6f}, diff {abs(e - c):.2e} exceeds tolerance 0.1"
+            )
+
+    def test_osft_no_graph_breaks(self, saved_model, single_gpu_device):
+        """OSFT forward/backward completes under fullgraph=True with varied seq_len."""
+        model_path, config = saved_model
+
+        torch._dynamo.utils.counters.clear()
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+
+        model = setup_model(
+            model_name_or_path=str(model_path),
+            use_liger_kernels=False,
+            osft=True,
+            osft_rank_ratio=0.25,
+            local_rank=0,
+        )
+        model, optimizer, lr_scheduler = setup_training_components(
+            model,
+            learning_rate=1e-3,
+            num_warmup_steps=0,
+            lr_scheduler="constant",
+            compile_model=True,
+        )
+
+        input_ids_1 = torch.randint(0, config.vocab_size, (2, 32), device=single_gpu_device)
+        optimizer.zero_grad()
+        loss = model(input_ids=input_ids_1, labels=input_ids_1.clone()).loss.float().sum()
+        loss.backward()
+        optimizer.step()
+
+        compilations_after_first = torch._dynamo.utils.counters["stats"]["ok"]
+
+        input_ids_2 = torch.randint(0, config.vocab_size, (2, 48), device=single_gpu_device)
+        optimizer.zero_grad()
+        loss = model(input_ids=input_ids_2, labels=input_ids_2.clone()).loss.float().sum()
+        loss.backward()
+        optimizer.step()
+
+        compilations_after_second = torch._dynamo.utils.counters["stats"]["ok"]
+
+        graph_breaks = dict(torch._dynamo.utils.counters["graph_break"])
+        assert len(graph_breaks) == 0, f"Graph breaks detected: {graph_breaks}"
+
+        assert compilations_after_second == compilations_after_first, (
+            f"dynamic=True should prevent recompilation on shape change, "
+            f"but compilations went from {compilations_after_first} to {compilations_after_second}"
+        )
+
+    def test_osft_optimized_module_wrappers(self, saved_model, single_gpu_device):
+        """Compiled OSFT blocks are OptimizedModule."""
+        model_path, _ = saved_model
+
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+        model = setup_model(
+            model_name_or_path=str(model_path),
+            use_liger_kernels=False,
+            osft=True,
+            osft_rank_ratio=0.25,
+            local_rank=0,
+        )
+        model, _, _ = setup_training_components(
+            model,
+            learning_rate=1e-3,
+            num_warmup_steps=0,
+            lr_scheduler="constant",
+            compile_model=True,
+        )
+
+        from torch._dynamo.eval_frame import OptimizedModule
+
+        layers = model.model.layers
+        for idx, block in enumerate(layers):
+            assert isinstance(block, OptimizedModule), f"Block {idx} should be OptimizedModule, got {type(block)}"
diff --git a/tests/test_compile_guards.py b/tests/test_compile_guards.py
index 4e9954ea..70f9f7a8 100644
--- a/tests/test_compile_guards.py
+++ b/tests/test_compile_guards.py
@@ -10,16 +10,6 @@
 
 
 class TestCompileValidationGuards:
-    def test_compile_osft_incompatible(self):
-        compile_model = True
-        osft = True
-        with pytest.raises(ValueError, match="not compatible with --osft"):
-            if compile_model and osft:
-                raise ValueError(
-                    "--compile-model is not compatible with --osft. "
-                    "OSFT uses dynamic forward methods that cannot be traced by torch.compile."
-                )
-
     def test_compile_liger_incompatible(self):
         compile_model = True
         use_liger_kernels = True
@@ -42,10 +32,7 @@ def test_compile_moe_incompatible(self):
 
     def test_compile_guards_do_not_fire_when_disabled(self):
         compile_model = False
-        osft = True
         use_liger_kernels = True
 
-        if compile_model and osft:
-            raise ValueError("should not reach")
         if compile_model and use_liger_kernels:
             raise ValueError("should not reach")
diff --git a/tests/test_osft.py b/tests/test_osft.py
index 7723d712..7d25b57b 100644
--- a/tests/test_osft.py
+++ b/tests/test_osft.py
@@ -752,8 +752,6 @@ def __init__(self, config, **kwargs):
         model = OSFTModelClass(config, osft_config={}, initialize_osft=False)
 
         assert model.osft_config == {}
-        assert hasattr(model, "osft_params")
-        assert len(model.osft_params) == 0
 
 
 class TestSetupModelIntegration:
@@ -2290,10 +2288,10 @@ def test_project_parameters_basic(self):
                 # Add a component in the frozen subspace direction
                 with torch.no_grad():
                     module.osft_params.U_low.data += U_high @ torch.randn(
-                        U_high.shape[1], module.osft_params.U_low.shape[1]
+                        U_high.shape[1], module.osft_params.U_low.shape[1], device=U_high.device
                     )
                     module.osft_params.V_low.data += (
-                        torch.randn(module.osft_params.V_low.shape[0], V_high.shape[0]) @ V_high
+                        torch.randn(module.osft_params.V_low.shape[0], V_high.shape[0], device=V_high.device) @ V_high
                     )
 
         # Now project parameters
@@ -2471,7 +2469,7 @@ def test_project_parameters_idempotent(self):
             if hasattr(module, "osft_params") and hasattr(module, "osft_U_high"):
                 with torch.no_grad():
                     module.osft_params.U_low.data += module.osft_U_high @ torch.randn(
-                        module.osft_U_high.shape[1], module.osft_params.U_low.shape[1]
+                        module.osft_U_high.shape[1], module.osft_params.U_low.shape[1], device=module.osft_U_high.device
                     )
 
         # First projection
diff --git a/tests/test_osft_fidelity.py b/tests/test_osft_fidelity.py
index e76447b5..aca7c6f4 100644
--- a/tests/test_osft_fidelity.py
+++ b/tests/test_osft_fidelity.py
@@ -424,85 +424,56 @@ def __init__(self, config, **kwargs):
         }
 
     def test_factorized_linear_2d_input(self, simple_model_with_osft):
-        """Test factorized linear with 2D input matches standard linear operation."""
+        """Test OSFTLinear forward with 2D input matches standard linear operation."""
         model = simple_model_with_osft["model"]
         original_weight = simple_model_with_osft["original_weight"]
         original_bias = simple_model_with_osft["original_bias"]
 
-        # Create test input (batch_size=8, input_dim=64)
         test_input = torch.randn(8, 64, dtype=torch.float32)
-
-        # Get expected result using standard linear operation
         expected_output = F.linear(test_input, original_weight, original_bias)
 
-        # Get SVD dict for the linear layer using the proper API
-        linear_module, _ = model._get_module_by_name("linear.weight")
-        svd_dict = model.get_svd_dict_for_module(linear_module)
-
-        # Get actual result using factorized linear
-        actual_output = model._factorized_linear(test_input, svd_dict, original_bias)
+        actual_output = model.linear(test_input)
 
-        # Check shapes match
         assert actual_output.shape == expected_output.shape, (
             f"Shape mismatch: {actual_output.shape} vs {expected_output.shape}"
         )
-
-        # Check outputs are approximately equal (with reasonable tolerance for SVD approximation)
         assert torch.allclose(actual_output, expected_output, rtol=1e-3, atol=1e-4), (
-            f"Factorized linear output differs from standard linear. Max diff: {torch.max(torch.abs(actual_output - expected_output))}"
+            f"OSFTLinear output differs from standard linear. Max diff: {torch.max(torch.abs(actual_output - expected_output))}"
         )
 
     def test_factorized_linear_3d_input(self, simple_model_with_osft):
-        """Test factorized linear with 3D input matches standard linear operation."""
+        """Test OSFTLinear forward with 3D input matches standard linear operation."""
         model = simple_model_with_osft["model"]
         original_weight = simple_model_with_osft["original_weight"]
         original_bias = simple_model_with_osft["original_bias"]
 
-        # Create test input (batch_size=4, seq_len=16, input_dim=64)
         test_input = torch.randn(4, 16, 64, dtype=torch.float32)
-
-        # Get expected result using standard linear operation
         expected_output = F.linear(test_input, original_weight, original_bias)
 
-        # Get SVD dict for the linear layer using the proper API
-        linear_module, _ = model._get_module_by_name("linear.weight")
-        svd_dict = model.get_svd_dict_for_module(linear_module)
-
-        # Get actual result using factorized linear
-        actual_output = model._factorized_linear(test_input, svd_dict, original_bias)
+        actual_output = model.linear(test_input)
 
-        # Check shapes match
         assert actual_output.shape == expected_output.shape, (
             f"Shape mismatch: {actual_output.shape} vs {expected_output.shape}"
         )
-
-        # Check outputs are approximately equal (with reasonable tolerance for SVD approximation)
         assert torch.allclose(actual_output, expected_output, rtol=1e-3, atol=1e-4), (
-            f"Factorized linear output differs from standard linear. Max diff: {torch.max(torch.abs(actual_output - expected_output))}"
+            f"OSFTLinear output differs from standard linear. Max diff: {torch.max(torch.abs(actual_output - expected_output))}"
         )
 
     def test_factorized_linear_without_bias(self, simple_model_with_osft):
-        """Test factorized linear without bias term."""
+        """Test OSFTLinear forward without bias term."""
         model = simple_model_with_osft["model"]
         original_weight = simple_model_with_osft["original_weight"]
 
-        # Create test input
         test_input = torch.randn(6, 64, dtype=torch.float32)
-
-        # Get expected result using standard linear operation (no bias)
         expected_output = F.linear(test_input, original_weight, None)
 
-        # Get SVD dict for the linear layer using the proper API
-        linear_module, _ = model._get_module_by_name("linear.weight")
-        svd_dict = model.get_svd_dict_for_module(linear_module)
+        # OSFTLinear includes bias from original — subtract it for comparison
+        osft_module = model.linear
+        actual_output = osft_module(test_input)
+        if osft_module.bias is not None:
+            actual_output = actual_output - osft_module.bias
 
-        # Get actual result using factorized linear (no bias)
-        actual_output = model._factorized_linear(test_input, svd_dict, None)
-
-        # Check shapes match
         assert actual_output.shape == expected_output.shape
-
-        # Check outputs are approximately equal
         assert torch.allclose(actual_output, expected_output, rtol=1e-3, atol=1e-4), (
-            f"Factorized linear without bias differs from standard linear. Max diff: {torch.max(torch.abs(actual_output - expected_output))}"
+            f"OSFTLinear without bias differs from standard linear. Max diff: {torch.max(torch.abs(actual_output - expected_output))}"
         )

From f7c648826d24817eb9cc8397001e05567c47df26 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 24 Jun 2026 16:55:18 +0000
Subject: [PATCH 10/19] Fix reinitialize_osft when called on
 already-initialized model

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.
---
 src/mini_trainer/osft_utils.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/src/mini_trainer/osft_utils.py b/src/mini_trainer/osft_utils.py
index 465190ca..64f964ef 100644
--- a/src/mini_trainer/osft_utils.py
+++ b/src/mini_trainer/osft_utils.py
@@ -1799,6 +1799,33 @@ def _pre_fsdp2_wrap_initialize_lazy_osft(self):
             for key in self.logical_osft_keys:
                 self._prepare_osft_param(key)
 
+        @torch.no_grad()
+        def _restore_dense_linears(self):
+            """Replace OSFTLinear modules with standard nn.Linear so that
+            the original parameter FQNs (e.g. ``q_proj.weight``) are restored.
+
+            Must be called before ``_reset_osft_metadata`` clears the registry.
+            """
+            for orig_key, spec in self.osft_paramspec_registry.items():
+                mod_path = orig_key.rsplit(".", 1)[0]
+                parent, child_name = self._get_module_by_name(mod_path)
+                osft_mod = getattr(parent, child_name, None)
+                if not isinstance(osft_mod, OSFTLinear):
+                    continue
+                svd_dict = self.get_svd_dict_for_module(osft_mod)
+                W = reconstruct_weight_matrix(
+                    svd_dict,
+                    upcast_dtype=self.upcast_dtype,
+                    output_dtype=self.output_dtype,
+                )
+                out_features, in_features = W.shape
+                has_bias = osft_mod.bias is not None
+                linear = nn.Linear(in_features, out_features, bias=has_bias, device=W.device, dtype=W.dtype)
+                linear.weight = nn.Parameter(W, requires_grad=True)
+                if has_bias:
+                    linear.bias = nn.Parameter(osft_mod.bias.data, requires_grad=osft_mod.bias.requires_grad)
+                setattr(parent, child_name, linear)
+
         def reinitialize_osft(self, decompose_existing_weights: bool, assigned_params=None):
             """
             Reinitializes the OSFT decomposition (e.g., when learning a new task in continual learning).
@@ -1814,6 +1841,9 @@ def reinitialize_osft(self, decompose_existing_weights: bool, assigned_params=No
             log_rank_0(f"   • decompose_existing_weights: {decompose_existing_weights}")
             log_rank_0(f"   • assigned_params: {len(assigned_params) if assigned_params else 'None (all params)'}")
 
+            if self.osft_paramspec_registry:
+                self._restore_dense_linears()
+
             self._reset_osft_metadata()
 
             log_rank_0("🚀 [reinitialize_osft] Calling _initialize_osft_parameters")

From be9301c950d36074211d4355729d107e2647489e Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 24 Jun 2026 21:06:13 +0000
Subject: [PATCH 11/19] Add regression tests for reinitialize_osft double-init

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.
---
 tests/test_osft.py | 69 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/tests/test_osft.py b/tests/test_osft.py
index 7d25b57b..16734f58 100644
--- a/tests/test_osft.py
+++ b/tests/test_osft.py
@@ -1052,6 +1052,75 @@ def __init__(self, config=None, **kwargs):
         assert reconstructed["linear.weight"].dtype == torch.float32
 
 
+class TestOSFTReinitialize:
+    """Test reinitialize_osft handles the double-init case correctly."""
+
+    def test_reinitialize_after_initialize(self):
+        """Regression: reinitialize_osft must work when the model was already initialized.
+
+        After the first OSFT init, parameter FQNs change (e.g. q_proj.weight becomes
+        q_proj.osft_U_high). Without restoring dense linears first, the second init
+        finds zero matching targets and silently produces a broken model.
+        """
+
+        class SimpleModel(nn.Module):
+            def __init__(self, config=None, **kwargs):
+                super().__init__()
+                self.linear = nn.Linear(8, 8, bias=False)
+                self.config = config or MagicMock()
+                self.dtype = torch.float32
+
+        OSFTModel = create_osft_model_class(SimpleModel)
+        osft_config = {"linear.weight": 4}
+
+        model = OSFTModel(MagicMock(), osft_config=osft_config, initialize_osft=True)
+        assert len(model.osft_paramspec_registry) == 1
+
+        model.reinitialize_osft(decompose_existing_weights=True)
+        assert len(model.osft_paramspec_registry) == 1, (
+            f"Expected 1 OSFT param after reinit, got {len(model.osft_paramspec_registry)}"
+        )
+
+        from mini_trainer.osft_utils import OSFTLinear
+
+        assert isinstance(model.linear, OSFTLinear)
+
+        x = torch.randn(2, 8)
+        out = model.linear(x)
+        assert out.shape == (2, 8)
+
+        sd = model.state_dict()
+        osft_keys = [k for k in sd if "osft" in k or "_rank_high" in k]
+        assert len(osft_keys) > 0, "No OSFT keys in state dict after reinit"
+
+    def test_reinitialize_preserves_weight_reconstruction(self):
+        """The reconstructed weight after reinit should approximate the original."""
+
+        class SimpleModel(nn.Module):
+            def __init__(self, config=None, **kwargs):
+                super().__init__()
+                self.linear = nn.Linear(16, 16, bias=False)
+                self.config = config or MagicMock()
+                self.dtype = torch.float32
+
+        OSFTModel = create_osft_model_class(SimpleModel)
+        osft_config = {"linear.weight": 12}
+
+        model = OSFTModel(MagicMock(), osft_config=osft_config, initialize_osft=True)
+
+        sd_before = model.prepare_state_dict_for_save(model.state_dict().copy())
+        W_before = sd_before["linear.weight"].clone()
+
+        model.reinitialize_osft(decompose_existing_weights=True)
+
+        sd_after = model.prepare_state_dict_for_save(model.state_dict().copy())
+        W_after = sd_after["linear.weight"]
+
+        assert torch.allclose(W_before, W_after, atol=1e-5), (
+            f"Weight changed after reinit: max diff = {(W_before - W_after).abs().max().item()}"
+        )
+
+
 class TestOSFTOrthogonality:
     """Test OSFT orthogonality constraints during training."""
 

From 74893cae5d7d15989723a89f9bdf831f3a600eff Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 24 Jun 2026 21:07:15 +0000
Subject: [PATCH 12/19] Add distributed benchmark for OSFT eager vs compiled

torchrun-based benchmark measuring per-step throughput with CUDA events.
Results from Granite-3.3-8B on 6x H200: 10-17% speedup with compile.
---
 benchmarks/bench_osft_compile_distributed.py | 254 +++++++++++++++++++
 benchmarks/results_osft_compile.json         | 196 ++++++++++++++
 2 files changed, 450 insertions(+)
 create mode 100644 benchmarks/bench_osft_compile_distributed.py
 create mode 100644 benchmarks/results_osft_compile.json

diff --git a/benchmarks/bench_osft_compile_distributed.py b/benchmarks/bench_osft_compile_distributed.py
new file mode 100644
index 00000000..dc4d0e2b
--- /dev/null
+++ b/benchmarks/bench_osft_compile_distributed.py
@@ -0,0 +1,254 @@
+"""Benchmark: OSFT eager vs compiled training step throughput.
+
+Measures steady-state step time for OSFT training with and without
+torch.compile on a real model (Llama-8B) under FSDP2.
+
+Usage:
+  # Run on 6 GPUs (adjust --nproc-per-node as needed)
+  CUDA_VISIBLE_DEVICES=2,3,4,5,6,7 torchrun --nnodes=1 --nproc-per-node=6 \
+      benchmarks/bench_osft_compile_distributed.py
+
+  # Eager only
+  CUDA_VISIBLE_DEVICES=2,3,4,5,6,7 torchrun --nnodes=1 --nproc-per-node=6 \
+      benchmarks/bench_osft_compile_distributed.py --mode eager
+
+  # Compiled only
+  CUDA_VISIBLE_DEVICES=2,3,4,5,6,7 torchrun --nnodes=1 --nproc-per-node=6 \
+      benchmarks/bench_osft_compile_distributed.py --mode compiled
+"""
+
+import argparse
+import gc
+import json
+import os
+import statistics
+import time
+
+os.environ["TESTING"] = "true"
+
+import torch
+import torch.distributed as dist
+
+from mini_trainer.none_reduction_losses import hf_fixed_cross_entropy_none_reduction
+from mini_trainer.setup_model_for_training import setup_model, setup_training_components
+from mini_trainer.utils import log_rank_0, patch_target_module
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="OSFT compile benchmark (distributed)")
+    parser.add_argument(
+        "--mode",
+        choices=["both", "eager", "compiled"],
+        default="both",
+        help="Which mode(s) to benchmark",
+    )
+    parser.add_argument("--model", type=str, default="meta-llama/Llama-3.1-8B", help="Model name or path")
+    parser.add_argument("--seq-lens", type=str, default="512,1024,2048", help="Comma-separated sequence lengths")
+    parser.add_argument("--batch-size", type=int, default=1, help="Per-GPU batch size")
+    parser.add_argument("--warmup-steps", type=int, default=5, help="Warmup steps (excluded from timing)")
+    parser.add_argument("--measure-steps", type=int, default=20, help="Steps to measure")
+    parser.add_argument("--osft-rank-ratio", type=float, default=0.25, help="OSFT rank ratio")
+    parser.add_argument("--output-json", type=str, default=None, help="Path to write results JSON")
+    return parser.parse_args()
+
+
+def run_arm(
+    model_name: str,
+    compile_model: bool,
+    seq_lens: list[int],
+    batch_size: int,
+    warmup_steps: int,
+    measure_steps: int,
+    osft_rank_ratio: float,
+    local_rank: int,
+):
+    label = "compiled" if compile_model else "eager"
+    log_rank_0(f"\n{'='*70}")
+    log_rank_0(f"  {label.upper()} ARM")
+    log_rank_0(f"{'='*70}")
+
+    if compile_model:
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+
+    model = setup_model(
+        model_name_or_path=model_name,
+        use_liger_kernels=False,
+        osft=True,
+        osft_rank_ratio=osft_rank_ratio,
+        local_rank=local_rank,
+    )
+    model, optimizer, lr_scheduler = setup_training_components(
+        model,
+        learning_rate=1e-5,
+        num_warmup_steps=0,
+        lr_scheduler="constant",
+        compile_model=compile_model,
+    )
+
+    results_by_seqlen = {}
+
+    for seq_len in seq_lens:
+        log_rank_0(f"\n--- {label} | seq_len={seq_len} ---")
+
+        input_ids = torch.randint(0, 32000, (batch_size, seq_len), device=f"cuda:{local_rank}")
+        labels = input_ids.clone()
+
+        # Warmup — extra steps for compiled to amortize torch.compile
+        actual_warmup = warmup_steps + (5 if compile_model else 0)
+        log_rank_0(f"  Warmup: {actual_warmup} steps...")
+        for _ in range(actual_warmup):
+            optimizer.zero_grad()
+            output = model(input_ids=input_ids, labels=labels)
+            loss = output.loss.float().sum() / batch_size
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            lr_scheduler.step()
+
+        torch.cuda.synchronize()
+        dist.barrier()
+
+        # Measure with CUDA events — no host sync between steps
+        log_rank_0(f"  Measuring: {measure_steps} steps...")
+        start_events = [torch.cuda.Event(enable_timing=True) for _ in range(measure_steps)]
+        end_events = [torch.cuda.Event(enable_timing=True) for _ in range(measure_steps)]
+
+        for i in range(measure_steps):
+            start_events[i].record()
+
+            optimizer.zero_grad()
+            output = model(input_ids=input_ids, labels=labels)
+            loss = output.loss.float().sum() / batch_size
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            lr_scheduler.step()
+
+            end_events[i].record()
+
+        torch.cuda.synchronize()
+        step_times_ms = [s.elapsed_time(e) for s, e in zip(start_events, end_events)]
+
+        dist.barrier()
+
+        median_ms = statistics.median(step_times_ms)
+        mean_ms = statistics.mean(step_times_ms)
+        stdev_ms = statistics.stdev(step_times_ms) if len(step_times_ms) > 1 else 0.0
+        p10 = sorted(step_times_ms)[max(0, len(step_times_ms) // 10)]
+        p90 = sorted(step_times_ms)[min(len(step_times_ms) - 1, 9 * len(step_times_ms) // 10)]
+
+        world_size = dist.get_world_size()
+        tokens_per_step = batch_size * seq_len * world_size
+        tokens_per_sec = tokens_per_step / (median_ms / 1000)
+
+        results_by_seqlen[seq_len] = {
+            "median_ms": round(median_ms, 2),
+            "mean_ms": round(mean_ms, 2),
+            "stdev_ms": round(stdev_ms, 2),
+            "p10_ms": round(p10, 2),
+            "p90_ms": round(p90, 2),
+            "tokens_per_sec": round(tokens_per_sec, 0),
+            "all_steps_ms": [round(t, 2) for t in step_times_ms],
+        }
+
+        log_rank_0(
+            f"  Result: median={median_ms:.1f}ms  mean={mean_ms:.1f}ms  "
+            f"stdev={stdev_ms:.1f}ms  p10={p10:.1f}ms  p90={p90:.1f}ms  "
+            f"tok/s={tokens_per_sec:.0f}"
+        )
+
+    # Cleanup
+    del model, optimizer, lr_scheduler
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch._dynamo.reset()
+
+    return results_by_seqlen
+
+
+def main():
+    args = parse_args()
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    torch.cuda.set_device(local_rank)
+
+    dist.init_process_group(backend="nccl")
+    patch_target_module(
+        "transformers.loss.loss_utils.fixed_cross_entropy",
+        hf_fixed_cross_entropy_none_reduction,
+    )
+
+    seq_lens = [int(s) for s in args.seq_lens.split(",")]
+    world_size = dist.get_world_size()
+
+    log_rank_0(f"\nOSFT Compile Benchmark")
+    log_rank_0(f"  Model: {args.model}")
+    log_rank_0(f"  GPUs: {world_size}x {torch.cuda.get_device_name(local_rank)}")
+    log_rank_0(f"  Seq lengths: {seq_lens}")
+    log_rank_0(f"  Batch size: {args.batch_size}/GPU")
+    log_rank_0(f"  OSFT rank ratio: {args.osft_rank_ratio}")
+    log_rank_0(f"  Warmup: {args.warmup_steps} steps, Measure: {args.measure_steps} steps")
+
+    all_results = {
+        "config": {
+            "model": args.model,
+            "world_size": world_size,
+            "gpu": torch.cuda.get_device_name(local_rank),
+            "batch_size_per_gpu": args.batch_size,
+            "osft_rank_ratio": args.osft_rank_ratio,
+            "warmup_steps": args.warmup_steps,
+            "measure_steps": args.measure_steps,
+            "pytorch_version": torch.__version__,
+        }
+    }
+
+    modes = []
+    if args.mode == "both":
+        modes = ["eager", "compiled"]
+    elif args.mode == "eager":
+        modes = ["eager"]
+    else:
+        modes = ["compiled"]
+
+    for mode in modes:
+        compile_model = mode == "compiled"
+        results = run_arm(
+            model_name=args.model,
+            compile_model=compile_model,
+            seq_lens=seq_lens,
+            batch_size=args.batch_size,
+            warmup_steps=args.warmup_steps,
+            measure_steps=args.measure_steps,
+            osft_rank_ratio=args.osft_rank_ratio,
+            local_rank=local_rank,
+        )
+        all_results[mode] = results
+
+    # Summary table
+    if len(modes) == 2 and dist.get_rank() == 0:
+        print(f"\n{'='*70}")
+        print(f"  SUMMARY: OSFT eager vs compiled")
+        print(f"  {args.model} | {world_size} GPUs | batch={args.batch_size}/GPU | rank_ratio={args.osft_rank_ratio}")
+        print(f"{'='*70}")
+        print(f"  {'seq_len':>8}  {'eager_ms':>10}  {'compiled_ms':>12}  {'speedup':>8}  {'tok/s eager':>12}  {'tok/s compiled':>15}")
+        print(f"  {'-'*8}  {'-'*10}  {'-'*12}  {'-'*8}  {'-'*12}  {'-'*15}")
+        for sl in seq_lens:
+            e = all_results["eager"][sl]
+            c = all_results["compiled"][sl]
+            speedup = e["median_ms"] / c["median_ms"]
+            print(
+                f"  {sl:>8}  {e['median_ms']:>10.1f}  {c['median_ms']:>12.1f}  {speedup:>7.2f}x  "
+                f"{e['tokens_per_sec']:>12.0f}  {c['tokens_per_sec']:>15.0f}"
+            )
+        print(f"{'='*70}")
+
+    if args.output_json and dist.get_rank() == 0:
+        with open(args.output_json, "w") as f:
+            json.dump(all_results, f, indent=2)
+        print(f"\nResults written to {args.output_json}")
+
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/results_osft_compile.json b/benchmarks/results_osft_compile.json
new file mode 100644
index 00000000..2b9c849a
--- /dev/null
+++ b/benchmarks/results_osft_compile.json
@@ -0,0 +1,196 @@
+{
+  "config": {
+    "model": "ibm-granite/granite-3.3-8b-base",
+    "world_size": 6,
+    "gpu": "NVIDIA H200",
+    "batch_size_per_gpu": 1,
+    "osft_rank_ratio": 0.25,
+    "warmup_steps": 5,
+    "measure_steps": 20,
+    "pytorch_version": "2.14.0a0+gitd16239c"
+  },
+  "eager": {
+    "512": {
+      "median_ms": 1002.25,
+      "mean_ms": 1059.1,
+      "stdev_ms": 142.45,
+      "p10_ms": 987.83,
+      "p90_ms": 1234.55,
+      "tokens_per_sec": 3065.0,
+      "all_steps_ms": [
+        1576.92,
+        987.83,
+        984.05,
+        987.34,
+        1005.69,
+        1004.63,
+        995.82,
+        1004.56,
+        1013.08,
+        1234.55,
+        1044.75,
+        1002.54,
+        1001.95,
+        995.49,
+        991.3,
+        1000.24,
+        997.58,
+        995.77,
+        1195.74,
+        1162.15
+      ]
+    },
+    "1024": {
+      "median_ms": 1009.29,
+      "mean_ms": 1052.26,
+      "stdev_ms": 106.69,
+      "p10_ms": 1003.68,
+      "p90_ms": 1239.38,
+      "tokens_per_sec": 6087.0,
+      "all_steps_ms": [
+        1002.37,
+        999.78,
+        1013.86,
+        1393.39,
+        1011.09,
+        1009.24,
+        1008.24,
+        1011.72,
+        1017.98,
+        1006.02,
+        1048.33,
+        1006.8,
+        1234.82,
+        1239.38,
+        1007.78,
+        1003.68,
+        1003.97,
+        1009.34,
+        1009.52,
+        1007.97
+      ]
+    },
+    "2048": {
+      "median_ms": 1211.05,
+      "mean_ms": 1248.17,
+      "stdev_ms": 78.31,
+      "p10_ms": 1208.22,
+      "p90_ms": 1416.38,
+      "tokens_per_sec": 10147.0,
+      "all_steps_ms": [
+        1212.38,
+        1208.22,
+        1206.95,
+        1210.73,
+        1213.13,
+        1209.74,
+        1416.38,
+        1393.7,
+        1210.86,
+        1210.43,
+        1209.37,
+        1210.75,
+        1213.68,
+        1211.76,
+        1211.24,
+        1209.94,
+        1427.52,
+        1358.1,
+        1211.39,
+        1207.16
+      ]
+    }
+  },
+  "compiled": {
+    "512": {
+      "median_ms": 855.22,
+      "mean_ms": 858.84,
+      "stdev_ms": 8.8,
+      "p10_ms": 853.62,
+      "p90_ms": 867.49,
+      "tokens_per_sec": 3592.0,
+      "all_steps_ms": [
+        866.06,
+        858.54,
+        854.32,
+        857.26,
+        854.76,
+        854.42,
+        855.63,
+        852.68,
+        854.81,
+        859.19,
+        854.55,
+        854.8,
+        853.62,
+        861.89,
+        891.61,
+        867.49,
+        853.7,
+        853.1,
+        861.95,
+        856.39
+      ]
+    },
+    "1024": {
+      "median_ms": 918.75,
+      "mean_ms": 918.99,
+      "stdev_ms": 2.18,
+      "p10_ms": 916.79,
+      "p90_ms": 922.72,
+      "tokens_per_sec": 6687.0,
+      "all_steps_ms": [
+        920.61,
+        917.22,
+        916.64,
+        919.6,
+        918.94,
+        915.2,
+        924.01,
+        917.83,
+        917.34,
+        917.62,
+        919.58,
+        918.56,
+        920.8,
+        920.68,
+        921.04,
+        922.72,
+        916.79,
+        919.35,
+        917.47,
+        917.85
+      ]
+    },
+    "2048": {
+      "median_ms": 1089.72,
+      "mean_ms": 1089.9,
+      "stdev_ms": 1.14,
+      "p10_ms": 1088.57,
+      "p90_ms": 1092.0,
+      "tokens_per_sec": 11276.0,
+      "all_steps_ms": [
+        1092.21,
+        1089.82,
+        1089.3,
+        1091.16,
+        1088.57,
+        1088.09,
+        1089.46,
+        1089.49,
+        1091.83,
+        1090.27,
+        1090.21,
+        1088.42,
+        1089.37,
+        1089.87,
+        1092.0,
+        1089.76,
+        1089.68,
+        1089.29,
+        1089.09,
+        1090.11
+      ]
+    }
+  }
+}
\ No newline at end of file

From 381353da4c67680e2c6b678debe827a9da00683e Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 24 Jun 2026 22:26:25 +0000
Subject: [PATCH 13/19] Fix lint: remove unused import, f-string prefixes;
 reformat

---
 benchmarks/bench_osft_compile_distributed.py | 21 ++++++++++----------
 tests/gpu_tests/test_compile.py              | 12 +++++++++--
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/benchmarks/bench_osft_compile_distributed.py b/benchmarks/bench_osft_compile_distributed.py
index dc4d0e2b..4db35a08 100644
--- a/benchmarks/bench_osft_compile_distributed.py
+++ b/benchmarks/bench_osft_compile_distributed.py
@@ -22,7 +22,6 @@
 import json
 import os
 import statistics
-import time
 
 os.environ["TESTING"] = "true"
 
@@ -63,9 +62,9 @@ def run_arm(
     local_rank: int,
 ):
     label = "compiled" if compile_model else "eager"
-    log_rank_0(f"\n{'='*70}")
+    log_rank_0(f"\n{'=' * 70}")
     log_rank_0(f"  {label.upper()} ARM")
-    log_rank_0(f"{'='*70}")
+    log_rank_0(f"{'=' * 70}")
 
     if compile_model:
         torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
@@ -181,7 +180,7 @@ def main():
     seq_lens = [int(s) for s in args.seq_lens.split(",")]
     world_size = dist.get_world_size()
 
-    log_rank_0(f"\nOSFT Compile Benchmark")
+    log_rank_0("\nOSFT Compile Benchmark")
     log_rank_0(f"  Model: {args.model}")
     log_rank_0(f"  GPUs: {world_size}x {torch.cuda.get_device_name(local_rank)}")
     log_rank_0(f"  Seq lengths: {seq_lens}")
@@ -226,12 +225,14 @@ def main():
 
     # Summary table
     if len(modes) == 2 and dist.get_rank() == 0:
-        print(f"\n{'='*70}")
-        print(f"  SUMMARY: OSFT eager vs compiled")
+        print(f"\n{'=' * 70}")
+        print("  SUMMARY: OSFT eager vs compiled")
         print(f"  {args.model} | {world_size} GPUs | batch={args.batch_size}/GPU | rank_ratio={args.osft_rank_ratio}")
-        print(f"{'='*70}")
-        print(f"  {'seq_len':>8}  {'eager_ms':>10}  {'compiled_ms':>12}  {'speedup':>8}  {'tok/s eager':>12}  {'tok/s compiled':>15}")
-        print(f"  {'-'*8}  {'-'*10}  {'-'*12}  {'-'*8}  {'-'*12}  {'-'*15}")
+        print(f"{'=' * 70}")
+        print(
+            f"  {'seq_len':>8}  {'eager_ms':>10}  {'compiled_ms':>12}  {'speedup':>8}  {'tok/s eager':>12}  {'tok/s compiled':>15}"
+        )
+        print(f"  {'-' * 8}  {'-' * 10}  {'-' * 12}  {'-' * 8}  {'-' * 12}  {'-' * 15}")
         for sl in seq_lens:
             e = all_results["eager"][sl]
             c = all_results["compiled"][sl]
@@ -240,7 +241,7 @@ def main():
                 f"  {sl:>8}  {e['median_ms']:>10.1f}  {c['median_ms']:>12.1f}  {speedup:>7.2f}x  "
                 f"{e['tokens_per_sec']:>12.0f}  {c['tokens_per_sec']:>15.0f}"
             )
-        print(f"{'='*70}")
+        print(f"{'=' * 70}")
 
     if args.output_json and dist.get_rank() == 0:
         with open(args.output_json, "w") as f:
diff --git a/tests/gpu_tests/test_compile.py b/tests/gpu_tests/test_compile.py
index 0fd3dae1..3a5de71a 100644
--- a/tests/gpu_tests/test_compile.py
+++ b/tests/gpu_tests/test_compile.py
@@ -334,7 +334,11 @@ def test_osft_compiled_matches_eager(self, saved_model, single_gpu_device):
         torch.manual_seed(7)
         torch.cuda.manual_seed(7)
         eager_losses, eager_model = _run_steps(
-            model_path, compile_model=False, input_ids=input_ids, labels=labels, osft=True,
+            model_path,
+            compile_model=False,
+            input_ids=input_ids,
+            labels=labels,
+            osft=True,
         )
 
         del eager_model
@@ -346,7 +350,11 @@ def test_osft_compiled_matches_eager(self, saved_model, single_gpu_device):
         torch.cuda.manual_seed(7)
         torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
         compiled_losses, _ = _run_steps(
-            model_path, compile_model=True, input_ids=input_ids, labels=labels, osft=True,
+            model_path,
+            compile_model=True,
+            input_ids=input_ids,
+            labels=labels,
+            osft=True,
         )
 
         for step, (e, c) in enumerate(zip(eager_losses, compiled_losses)):

From 64cb60961e68a4900fd422f3c8b717bd5ccfa864 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Wed, 24 Jun 2026 22:55:03 +0000
Subject: [PATCH 14/19] Fix review findings: meta device in restore, stale
 Protocol, docstring

- _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
---
 src/mini_trainer/osft_utils.py | 4 ++--
 tests/test_osft_fidelity.py    | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/mini_trainer/osft_utils.py b/src/mini_trainer/osft_utils.py
index 64f964ef..ae2b3e6d 100644
--- a/src/mini_trainer/osft_utils.py
+++ b/src/mini_trainer/osft_utils.py
@@ -210,7 +210,7 @@ class OSFTModelProtocol(Protocol):
 
     osft_config: dict[str, int]
     name_mapping: dict[str, str]
-    osft_params: nn.ModuleDict
+    osft_paramspec_registry: dict
     upcast_dtype: torch.dtype
     output_dtype: torch.dtype
 
@@ -1820,7 +1820,7 @@ def _restore_dense_linears(self):
                 )
                 out_features, in_features = W.shape
                 has_bias = osft_mod.bias is not None
-                linear = nn.Linear(in_features, out_features, bias=has_bias, device=W.device, dtype=W.dtype)
+                linear = nn.Linear(in_features, out_features, bias=has_bias, device="meta", dtype=W.dtype)
                 linear.weight = nn.Parameter(W, requires_grad=True)
                 if has_bias:
                     linear.bias = nn.Parameter(osft_mod.bias.data, requires_grad=osft_mod.bias.requires_grad)
diff --git a/tests/test_osft_fidelity.py b/tests/test_osft_fidelity.py
index aca7c6f4..a10b9ae5 100644
--- a/tests/test_osft_fidelity.py
+++ b/tests/test_osft_fidelity.py
@@ -372,7 +372,7 @@ def test_svd_gradient_flow(self):
 
 
 class TestFactorizedLinearAccuracy:
-    """Test that _factorized_linear produces the same results as standard linear operations."""
+    """Test that OSFTLinear forward produces the same results as standard linear operations."""
 
     @pytest.fixture
     def simple_model_with_osft(self):

From 182a144385cefa0beea3130d6819b1648189ef6c Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 25 Jun 2026 18:16:39 +0000
Subject: [PATCH 15/19] Use SDPA as default attention implementation

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.
---
 src/mini_trainer/setup_model_for_training.py | 44 +++++++-------------
 1 file changed, 16 insertions(+), 28 deletions(-)

diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index 5d7cdb81..cb641d70 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -31,7 +31,6 @@
     is_vlm_for_direct_loading,
     is_vlm_with_causal_lm,
     load_vlm_for_text_training,
-    needs_sdpa,
 )
 
 
@@ -1043,47 +1042,36 @@ def setup_model(
         except ImportError:
             log_rank_0("⚠️ GPT-OSS model detected but Mxfp4Config not available - using default config")
 
-    # Check if model requires SDPA instead of Flash Attention 2.
-    # This covers M-RoPE models (3D position_ids) and models with timm vision
-    # towers (TimmWrapperModel rejects flash_attention_2).
-    _needs_sdpa = needs_sdpa(model_config)
-
-    # Handle models that need SDPA (doesn't require flash_attn)
-    if _needs_sdpa:
-        base_model_args["attn_implementation"] = "sdpa"
-        log_rank_0(f"Using SDPA for {model_name_or_path} (model incompatible with Flash Attention 2)")
-    else:
-        # Check if flash_attn is available for non-SDPA models
+    # GPT-OSS models require vllm-flash-attn3 (Hopper+) or eager.
+    # All other models use SDPA, which dispatches to FlashAttention-2
+    # kernels via PyTorch's backend and is compatible with torch.compile.
+    if is_gpt_oss:
         try:
             import flash_attn as _
 
-            if is_gpt_oss:
-                # vllm-flash-attn3 requires Hopper (SM 9.0+) GPUs;
-                # GPT-OSS only supports flash-attn3 or eager
-                major, _ = torch.cuda.get_device_capability(0)
-                if major >= 9:
-                    base_model_args["attn_implementation"] = "kernels-community/vllm-flash-attn3"
-                    log_rank_0("Set attention implementation to vllm-flash-attn3 for GPT-OSS")
-                else:
-                    base_model_args["attn_implementation"] = "eager"
-                    log_rank_0(
-                        f"GPT-OSS: flash-attn3 requires Hopper (SM 9.0+) GPUs, "
-                        f"but found SM {major}.x. Using eager attention instead."
-                    )
+            major, _ = torch.cuda.get_device_capability(0)
+            if major >= 9:
+                base_model_args["attn_implementation"] = "kernels-community/vllm-flash-attn3"
+                log_rank_0("Set attention implementation to vllm-flash-attn3 for GPT-OSS")
             else:
-                base_model_args["attn_implementation"] = "flash_attention_2"
-
+                base_model_args["attn_implementation"] = "eager"
+                log_rank_0(
+                    f"GPT-OSS: flash-attn3 requires Hopper (SM 9.0+) GPUs, "
+                    f"but found SM {major}.x. Using eager attention instead."
+                )
         except ImportError as e:
             if os.environ.get("TESTING", "false").lower() == "true":
                 base_model_args["attn_implementation"] = "sdpa"
             else:
                 raise e
+    else:
+        base_model_args["attn_implementation"] = "sdpa"
 
     # For models with timm vision towers: set vision config to eager
     # while keeping the text model's attention implementation.
     # timm's TimmWrapperModel rejects both FA2 and SDPA.
     if has_timm_vision_tower(model_config):
-        attn_impl = base_model_args.get("attn_implementation", "flash_attention_2")
+        attn_impl = base_model_args.get("attn_implementation", "sdpa")
         base_model_args["attn_implementation"] = {
             "text_config": attn_impl,
             "vision_config": "eager",

From 10310dfa1cc86cc35b1aea2a65b69786695d4609 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 25 Jun 2026 18:27:18 +0000
Subject: [PATCH 16/19] Add orthogonality-under-compile test for OSFT

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.
---
 tests/gpu_tests/test_compile.py | 59 +++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/tests/gpu_tests/test_compile.py b/tests/gpu_tests/test_compile.py
index 3a5de71a..26f5dc15 100644
--- a/tests/gpu_tests/test_compile.py
+++ b/tests/gpu_tests/test_compile.py
@@ -11,8 +11,14 @@
 import torch.distributed as dist
 from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
 
+from mini_trainer.osft_utils import OSFTLinear
 from mini_trainer.setup_model_for_training import setup_model, setup_training_components
 from mini_trainer.utils import patch_target_module
+from tests.test_utils.orthogonality import (
+    OrthogonalityTracker,
+    check_gradient_orthogonality,
+    check_parameter_orthogonality,
+)
 
 
 def create_tiny_llama_model():
@@ -433,3 +439,56 @@ def test_osft_optimized_module_wrappers(self, saved_model, single_gpu_device):
         layers = model.model.layers
         for idx, block in enumerate(layers):
             assert isinstance(block, OptimizedModule), f"Block {idx} should be OptimizedModule, got {type(block)}"
+
+    def test_osft_orthogonality_under_compile(self, saved_model, single_gpu_device):
+        """OSFT subspace orthogonality is preserved under torch.compile.
+
+        Runs 10 training steps with compiled OSFT and checks that gradient
+        and parameter orthogonality are maintained within 1 degree at every
+        step. The optim_wrapper monkey-patch calls project_gradients() and
+        project_parameters() inside optimizer.step(), so checks run after
+        step returns.
+        """
+        model_path, config = saved_model
+
+        torch.manual_seed(7)
+        torch.cuda.manual_seed(7)
+        torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True
+
+        model = setup_model(
+            model_name_or_path=str(model_path),
+            use_liger_kernels=False,
+            osft=True,
+            osft_rank_ratio=0.25,
+            local_rank=0,
+        )
+        model, optimizer, lr_scheduler = setup_training_components(
+            model,
+            learning_rate=1e-3,
+            num_warmup_steps=0,
+            lr_scheduler="constant",
+            compile_model=True,
+        )
+
+        tracker = OrthogonalityTracker(margin_deg=1.0)
+        num_steps = 10
+
+        for step in range(1, num_steps + 1):
+            input_ids = torch.randint(0, config.vocab_size, (2, 32), device=single_gpu_device)
+            labels = input_ids.clone()
+
+            optimizer.zero_grad()
+            output = model(input_ids=input_ids, labels=labels)
+            loss = output.loss.float().sum() / input_ids.shape[0]
+            loss.backward()
+
+            # optim_wrapper runs project_gradients → step → project_parameters
+            optimizer.step()
+            lr_scheduler.step()
+
+            for module in model.modules():
+                if isinstance(module, OSFTLinear):
+                    check_gradient_orthogonality(model, module, step, tracker)
+                    check_parameter_orthogonality(model, module, step, tracker)
+
+        assert tracker.is_successful(), f"Orthogonality violated under compile:\n{tracker.get_summary()}"

From 3735236e2ed9b316ce20db4de9dd5f35a5910425 Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 25 Jun 2026 20:14:12 +0000
Subject: [PATCH 17/19] Fix lint: use noqa for flash_attn availability check
 import

---
 src/mini_trainer/setup_model_for_training.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index cb641d70..00f08088 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -1047,7 +1047,7 @@ def setup_model(
     # kernels via PyTorch's backend and is compatible with torch.compile.
     if is_gpt_oss:
         try:
-            import flash_attn as _
+            import flash_attn  # noqa: F401
 
             major, _ = torch.cuda.get_device_capability(0)
             if major >= 9:

From dcce1f65f963470ad6e3675f7ee6ffcd83a4df9c Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 25 Jun 2026 21:16:12 +0000
Subject: [PATCH 18/19] Tighten comments: fix imprecise SDPA claim, trim
 docstrings

---
 src/mini_trainer/osft_utils.py               | 38 +++++---------------
 src/mini_trainer/setup_model_for_training.py |  4 +--
 tests/gpu_tests/test_compile.py              | 15 ++------
 3 files changed, 14 insertions(+), 43 deletions(-)

diff --git a/src/mini_trainer/osft_utils.py b/src/mini_trainer/osft_utils.py
index ae2b3e6d..96ce81db 100644
--- a/src/mini_trainer/osft_utils.py
+++ b/src/mini_trainer/osft_utils.py
@@ -134,16 +134,11 @@ class SVDDecompositionDict(SVDDictBase, total=False):
 
 
 class OSFTLinear(nn.Module):
-    """Factorized linear using SVD components: W = U_high @ diag(S_high) @ V_high + U_low @ diag(S_low) @ V_low.
+    """Factorized linear: W = U_high @ diag(S_high) @ V_high + U_low @ diag(S_low) @ V_low.
 
-    Replaces closure-based forward overrides with a proper nn.Module whose
-    forward is pure tensor math — no weakrefs, dict construction, getattr,
-    or runtime .to() calls — so torch.compile can trace it without graph breaks.
-
-    Attribute names (osft_U_high, osft_S_high, osft_V_high, osft_params.U_low,
-    osft_params.S_low, osft_params.V_low) match the existing convention so that
-    project_gradients, project_parameters, get_svd_dict_for_module, and
-    prepare_state_dict_for_save continue to work unchanged.
+    forward() is pure tensor math — torch.compile traces it without graph breaks.
+    Attribute names match the existing convention used by project_gradients,
+    project_parameters, get_svd_dict_for_module, and prepare_state_dict_for_save.
     """
 
     def __init__(self, U_high, S_high, V_high, U_low, S_low, V_low, bias=None, rank_high=None):
@@ -156,12 +151,8 @@ def __init__(self, U_high, S_high, V_high, U_low, S_low, V_low, bias=None, rank_
         osft_params.S_low = S_low
         osft_params.V_low = V_low
         self.osft_params = osft_params
-        # rank_high: full (unsharded) k_high dimension, needed by V all-gather
-        # in project_gradients. Stored as a persistent int buffer so it survives
-        # FSDP2 sharding (S_high.shape[0] becomes the local shard size, not full).
+        # Persistent buffer: FSDP2 shards S_high, so S_high.shape[0] != k_high.
         if rank_high is not None:
-            # Infer device from an existing parameter so the buffer lives on the
-            # same device (meta during lazy init, CPU/CUDA during standard init).
             buf_device = U_high.device
             self.register_buffer(
                 "_rank_high_buf",
@@ -179,10 +170,9 @@ def __init__(self, U_high, S_high, V_high, U_low, S_low, V_low, bias=None, rank_
 
     @property
     def rank_high(self) -> int:
+        """Full (unsharded) k_high. Not safe inside a compiled region (.item() graph break)."""
         if self._rank_high_cached is not None:
             return self._rank_high_cached
-        # Lazy extract after set_model_state_dict materializes the buffer.
-        # Single .item() call, then cached for all subsequent accesses.
         buf = self._rank_high_buf
         if buf is not None and buf.device.type != "meta":
             self._rank_high_cached = buf.item()
@@ -1801,11 +1791,7 @@ def _pre_fsdp2_wrap_initialize_lazy_osft(self):
 
         @torch.no_grad()
         def _restore_dense_linears(self):
-            """Replace OSFTLinear modules with standard nn.Linear so that
-            the original parameter FQNs (e.g. ``q_proj.weight``) are restored.
-
-            Must be called before ``_reset_osft_metadata`` clears the registry.
-            """
+            """Replace OSFTLinear → nn.Linear before re-decomposition. Must run before _reset_osft_metadata."""
             for orig_key, spec in self.osft_paramspec_registry.items():
                 mod_path = orig_key.rsplit(".", 1)[0]
                 parent, child_name = self._get_module_by_name(mod_path)
@@ -1916,10 +1902,7 @@ def eject_og_state_dict(self):
             return sd
 
         def _prepare_osft_param(self, logical_key: str):
-            """
-            Prepares an OSFT parameter by replacing the Linear module with an
-            OSFTLinear at the same position in the module tree.
-            """
+            """Replace nn.Linear at logical_key with an OSFTLinear."""
             mod, attr = self._get_module_by_name(logical_key)
             if mod is None:
                 raise ValueError(f"requested module {logical_key} but could not be found")
@@ -2073,10 +2056,7 @@ def _initialize_osft_parameters(self, decompose_existing_weights: bool, assigned
                         output_dtype=self.output_dtype,
                     )
 
-                    # Move SVD results back to the original parameter's device.
-                    # SVD runs on GPU for speed, but the module should live where
-                    # its predecessor lived.  FSDP2 handles final placement.
-                    # Tensor.to() across devices strips nn.Parameter, so re-wrap.
+                    # .to() across devices strips nn.Parameter — re-wrap.
                     orig_device = param.device
                     for key in svd_dict:
                         if isinstance(svd_dict[key], torch.Tensor):
diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index 00f08088..a6d26278 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -1043,8 +1043,8 @@ def setup_model(
             log_rank_0("⚠️ GPT-OSS model detected but Mxfp4Config not available - using default config")
 
     # GPT-OSS models require vllm-flash-attn3 (Hopper+) or eager.
-    # All other models use SDPA, which dispatches to FlashAttention-2
-    # kernels via PyTorch's backend and is compatible with torch.compile.
+    # All other models use SDPA, which is compatible with torch.compile
+    # (HF's flash_attention path has a data-dependent graph break).
     if is_gpt_oss:
         try:
             import flash_attn  # noqa: F401
diff --git a/tests/gpu_tests/test_compile.py b/tests/gpu_tests/test_compile.py
index 26f5dc15..b7f4c00d 100644
--- a/tests/gpu_tests/test_compile.py
+++ b/tests/gpu_tests/test_compile.py
@@ -441,14 +441,7 @@ def test_osft_optimized_module_wrappers(self, saved_model, single_gpu_device):
             assert isinstance(block, OptimizedModule), f"Block {idx} should be OptimizedModule, got {type(block)}"
 
     def test_osft_orthogonality_under_compile(self, saved_model, single_gpu_device):
-        """OSFT subspace orthogonality is preserved under torch.compile.
-
-        Runs 10 training steps with compiled OSFT and checks that gradient
-        and parameter orthogonality are maintained within 1 degree at every
-        step. The optim_wrapper monkey-patch calls project_gradients() and
-        project_parameters() inside optimizer.step(), so checks run after
-        step returns.
-        """
+        """Gradient and parameter orthogonality hold under torch.compile."""
         model_path, config = saved_model
 
         torch.manual_seed(7)
@@ -471,9 +464,8 @@ def test_osft_orthogonality_under_compile(self, saved_model, single_gpu_device):
         )
 
         tracker = OrthogonalityTracker(margin_deg=1.0)
-        num_steps = 10
 
-        for step in range(1, num_steps + 1):
+        for step in range(1, 11):
             input_ids = torch.randint(0, config.vocab_size, (2, 32), device=single_gpu_device)
             labels = input_ids.clone()
 
@@ -481,11 +473,10 @@ def test_osft_orthogonality_under_compile(self, saved_model, single_gpu_device):
             output = model(input_ids=input_ids, labels=labels)
             loss = output.loss.float().sum() / input_ids.shape[0]
             loss.backward()
-
-            # optim_wrapper runs project_gradients → step → project_parameters
             optimizer.step()
             lr_scheduler.step()
 
+            # .grad still holds projected gradients (AdamW doesn't zero in step)
             for module in model.modules():
                 if isinstance(module, OSFTLinear):
                     check_gradient_orthogonality(model, module, step, tracker)

From d1b05cf9f55d70f0673e82b2ea494b1d25ac999c Mon Sep 17 00:00:00 2001
From: Sean McGovern <smcgover@redhat.com>
Date: Thu, 25 Jun 2026 21:36:24 +0000
Subject: [PATCH 19/19] Fix stale attribute names in _is_osft_owned_attribute

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.
---
 src/mini_trainer/setup_model_for_training.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/src/mini_trainer/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index a6d26278..be544042 100644
--- a/src/mini_trainer/setup_model_for_training.py
+++ b/src/mini_trainer/setup_model_for_training.py
@@ -135,14 +135,11 @@ def _sanitize_meta_attribute_aliases(model: torch.nn.Module) -> int:
     default_device = torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu")
 
     def _is_osft_owned_attribute(module: torch.nn.Module, name: str) -> bool:
-        if name.startswith("osft_") or name in {"U_low", "S_low", "V_low", "rank_high"}:
+        if name.startswith("osft_") or name.startswith("_rank_high"):
             return True
-        return hasattr(module, "osft_params") and name in {
-            "U_low",
-            "S_low",
-            "V_low",
-            "rank_high",
-        }
+        if hasattr(module, "osft_params") and name in {"U_low", "S_low", "V_low"}:
+            return True
+        return False
 
     for module in model.modules():
         # collect available candidates from this module