Ssameni/puzzletron bypass 3 integration

[Separius]

Summary

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

This PR wires the bypass distillation core into the full Puzzletron pipeline and adds runnable configs, docs, and end-to-end GPU
coverage.

Stack:

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

What Changed

• Added bypass as an optional Puzzletron pipeline stage after pruning and before replacement-library construction.
• Added dynamic Puzzletron progress numbering so runs report the correct total with or without bypass.
• Added config-specific bypass skip-if-complete detection.
• Added replacement-library stale detection so newly realized bypass checkpoints are picked up.
• Updated replacement-library extraction to infer which subblocks to extract from bypass checkpoint metadata.
• Prioritized bypass-trained checkpoints over plain pruned checkpoints when duplicate architectural candidates exist.
• Added MIP support for target_num_kv_heads.
• Added Llama bypass config coverage.
• Added Nemotron-3 Nano configs for pruning, validation, bypass, and full Puzzletron runs.
• Added Nemotron-3 example documentation.
• Added GPU/integration coverage for bypass training, checkpointing, resume, and full Puzzletron behavior.

Why

The previous PR adds the bypass training engine, but Puzzletron still needs pipeline wiring so bypass-trained blocks become usable
replacement-library candidates.

This PR keeps that integration separate from the core engine so reviewers can focus on pipeline behavior, config surface,
replacement-library semantics, and end-to-end coverage.

Tests

Added/updated coverage for:

• Puzzletron progress numbering
• bypass replacement-library discovery/extraction/priority
• bypass GPU smoke/integration behavior
• bypass checkpoint GPU behavior
• bypass resume behavior
• full Puzzletron pipeline behavior with bypass

[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.

[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: 298107c6-6ccd-4eaf-a4d8-7e74eed430c8

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-3-integration

---

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

[github-actions]

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