Ssameni/puzzletron bypass 2 core

[Separius]

Summary

This is PR 2 of 3 in the Puzzletron bypass/local-distillation stack.

This PR adds the bypass distillation core engine. It builds on PR 1’s shared infrastructure, but it does not yet wire bypass into
the full Puzzletron pipeline.

Stack:

1. ssameni/puzzletron-bypass-1-prereqs: shared prerequisites
2. This PR: bypass distillation core
3. ssameni/puzzletron-bypass-3-integration: Puzzletron integration, configs, docs, GPU coverage

What Changed

• Added modelopt.torch.puzzletron.bypass_distillation.
• Added bypass run identity/fingerprinting, experiment naming, state manifests, and completion tracking.
• Added stitched teacher/student model construction for local blockwise distillation.
• Added bypass training loop with:
  • pipeline-parallel teacher activation stitching
  • per-block student losses
  • gradient accumulation
  • checkpoint/resume support
  • validation hooks
  • best/latest checkpoint realization
• Added bypass checkpoint helpers for saving optimizer/scaler state and HF-format model checkpoints.
• Added scalable distributed checkpoint saving that gathers tensors per safetensors file instead of materializing all shards on
  rank 0.
• Restricted v1 keys_to_learn to subblock-level targets only:
  • entire_block
  • subblock_attention
  • subblock_ffn
  • subblock_mamba
  • lists of those keys
• Added pipeline ownership helper for deriving owned blocks and neighboring PP ranks.
• Added robust trust_remote_code / auto_map handling for checkpoint config saves.

Why

Bypass distillation is the local-distillation stage used to train pruned/reconfigured blocks before they are added to the
replacement library.

Keeping the core engine separate from Puzzletron pipeline wiring makes this PR reviewable on its own: reviewers can focus on
distributed training, checkpoint/resume semantics, and Sewing Kit stitching behavior without also reviewing configs/docs/pipeline
integration.

Tests

Added focused unit coverage for:

• bypass checkpoint utilities
• bypass run identity/fingerprinting and completion state
• keys_to_learn subblock selection
• LR scheduler behavior
• launch/sweep dispatch behavior
• HF checkpoint utility behavior
• stitched model factory buffer ownership

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1469/
Built to branch gh-pages at 2026-05-12 14:43 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 10cc3f86-0246-4743-b755-8ea7b4570124

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch ssameni/puzzletron-bypass-2-core

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}