Bug: torch.compile + DDP causes _orig_mod. prefix leakage in checkpoint keys

Bug Report: _orig_mod. Prefix Leakage in Checkpoint Loading

Problem Description

When torch.compile is enabled during pre-training, saved checkpoints contain an _orig_mod. prefix in their state dict keys. The existing prefix-stripping logic only handled the module. prefix from DDP, causing a complete key mismatch during post-training (SFT/DPO) and inference.

Observed Behavior:

• Checkpoint keys: _orig_mod.attn_blocks.0.attn.heads.0.key.weight
• Expected keys: attn_blocks.0.attn.heads.0.key.weight
• Result: 163 missing keys during loading, causing the model to load with random initialization instead of pre-trained weights.

Root Cause: Two separate code paths had flawed prefix-stripping logic:

1. src/post_training/utils.py (_strip_ddp_prefix): Used a single if check that only stripped module., ignoring _orig_mod. and nested prefixes.
2. src/post_training/inference.py (load_model_from_ckpt): Had independent logic that only stripped module. and transformer. once, also ignoring _orig_mod. and nested prefixes.

Fix

1. Unified Prefix Stripping in utils.py Replaced the single if check with a while loop to iteratively strip both module. and _orig_mod. prefixes, handling arbitrary nesting:

# src/post_training/utils.py
def _strip_ddp_prefix(state_dict: dict) -> dict:
    """Remove leading module. (DDP) and _orig_mod. (torch.compile) prefixes."""
    out = {}
    for k, v in state_dict.items():
        while k.startswith("module.") or k.startswith("_orig_mod."):
            k = k.removeprefix("module.").removeprefix("_orig_mod.")
        out[k] = v
    return out

2. Deduplicated Logic in inference.py Removed the independent prefix-stripping logic and reused the unified _strip_ddp_prefix function:

# src/post_training/inference.py
from src.post_training.utils import _strip_ddp_prefix

def load_model_from_ckpt(ckpt_path, device, overrides=None):
    # ... build model ...
    state = ck["model_state_dict"] if "model_state_dict" in ck else ck
    state = _strip_ddp_prefix(state)  # Reuse unified function
    keys = set(model.state_dict().keys())
    model.load_state_dict({k: v for k, v in state.items() if k in keys}, strict=False)
    return model.to(device).eval()

Verification

After applying the fix, checkpoint loading produces zero missing keys:

ephemeral/ckpts/base_pretrained.pt:
  total ckpt keys: 511, matched: 163, missing: 0
ephemeral/ckpts/sft.pt:
  total ckpt keys: 511, matched: 163, missing: 0

All model loading paths (training, inference, and evaluation) now consistently handle torch.compile and DDP prefixes.