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/benchmarks/bench_osft_compile_distributed.py b/benchmarks/bench_osft_compile_distributed.py
new file mode 100644
index 00000000..4db35a08
--- /dev/null
+++ b/benchmarks/bench_osft_compile_distributed.py
@@ -0,0 +1,255 @@
+"""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
+
+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("\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("  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
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/osft_utils.py b/src/mini_trainer/osft_utils.py
index 29c80b96..96ce81db 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,63 @@ class SVDDecompositionDict(SVDDictBase, total=False):
     rank_high: int
 
 
+class OSFTLinear(nn.Module):
+    """Factorized linear: W = U_high @ diag(S_high) @ V_high + U_low @ diag(S_low) @ V_low.
+
+    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):
+        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
+        # Persistent buffer: FSDP2 shards S_high, so S_high.shape[0] != k_high.
+        if rank_high is not None:
+            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:
+        """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
+        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.
@@ -144,7 +200,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
 
@@ -1166,7 +1222,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 +1251,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 +1642,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,14 +1786,32 @@ 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)
 
+        @torch.no_grad()
+        def _restore_dense_linears(self):
+            """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)
+                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="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)
+                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).
@@ -1756,6 +1827,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")
@@ -1795,23 +1869,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 +1884,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
@@ -1845,16 +1902,11 @@ def eject_og_state_dict(self):
             return sd
 
         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.
-            """
-            mod_ref, attr = self._osft_handles[logical_key]
-            mod = mod_ref()
+            """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 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 +1918,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"],
+            )
 
-            mod._parameters.pop(attr)
+            safe_name = logical_key.replace(".", "_")
+            self.name_mapping[logical_key] = safe_name
+            osft_linear.osft_params.safe_name = safe_name
+
+            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 +2056,14 @@ 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
+                    # .to() across devices strips nn.Parameter — 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 +2078,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),
+                    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"],
                     )
-                    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"]
-                    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 +2140,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 +2149,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 +2399,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/sampler.py b/src/mini_trainer/sampler.py
index d2702f60..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),
@@ -503,6 +509,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/setup_model_for_training.py b/src/mini_trainer/setup_model_for_training.py
index 39f0f67b..be544042 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,
 )
 
 
@@ -136,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
@@ -415,7 +411,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.
 
@@ -427,6 +423,8 @@ def wrap_fsdp2(model: torch.nn.Module) -> torch.nn.Module:
 
     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
@@ -482,6 +480,12 @@ 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:
+        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, fullgraph=True, dynamic=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"])
@@ -1035,47 +1039,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 is compatible with torch.compile
+    # (HF's flash_attention path has a data-dependent graph break).
+    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."
-                    )
-            else:
-                base_model_args["attn_implementation"] = "flash_attention_2"
+            import flash_attn  # noqa: F401
 
+            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."
+                )
         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",
@@ -1262,11 +1255,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 +1271,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 +1288,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 +1303,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..a8daa3bb 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,14 @@ def train(
 
             dist.barrier()
 
+            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 ({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")
@@ -1283,6 +1284,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 +1408,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 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")
@@ -1564,6 +1572,19 @@ 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."
+            )
+        # 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
+        torch._inductor.config.unsafe_skip_cache_dynamic_shape_guards = True
+
     # Create PretrainingConfig if block_size is provided
     pretraining_config = None
     if block_size is not None:
@@ -1613,6 +1634,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(
diff --git a/tests/gpu_tests/test_compile.py b/tests/gpu_tests/test_compile.py
new file mode 100644
index 00000000..b7f4c00d
--- /dev/null
+++ b/tests/gpu_tests/test_compile.py
@@ -0,0 +1,485 @@
+"""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.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():
+    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, 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=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,
+    )
+
+    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()
+
+
+@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)}"
+
+    def test_osft_orthogonality_under_compile(self, saved_model, single_gpu_device):
+        """Gradient and parameter orthogonality hold under torch.compile."""
+        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)
+
+        for step in range(1, 11):
+            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()
+            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)
+                    check_parameter_orthogonality(model, module, step, tracker)
+
+        assert tracker.is_successful(), f"Orthogonality violated under compile:\n{tracker.get_summary()}"
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..70f9f7a8
--- /dev/null
+++ b/tests/test_compile_guards.py
@@ -0,0 +1,38 @@
+"""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_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
+        use_liger_kernels = True
+
+        if compile_model and use_liger_kernels:
+            raise ValueError("should not reach")
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_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(
diff --git a/tests/test_osft.py b/tests/test_osft.py
index 7723d712..16734f58 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:
@@ -1054,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."""
 
@@ -2290,10 +2357,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 +2538,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..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):
@@ -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))}"
         )
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]