Ssameni/puzzletron bypass 3 integration#1470
Open
Separius wants to merge 9 commits into
Open
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds end-to-end bypass distillation orchestration to Puzzletron: new Nemotron and Llama bypass configs and tutorial, config-driven pipeline progress, AnyModel checkpoint completeness checks, bypass-run detection, replacement-library metadata parsing with bypass prioritization, and broad unit + GPU tests covering bypass scenarios. ChangesBypass Distillation Integration
Sequence Diagram(s)sequenceDiagram
participant Caller as CLI / Hydra
participant Converter as convert_puzzletron_model
participant Checkpoint as _is_complete_anymodel_checkpoint
participant HF as HF snapshot_download
participant BypassDetector as _find_incomplete_bypass_runs
Caller->>Converter: start conversion pipeline
Converter->>Checkpoint: check teacher checkpoint completeness
alt teacher incomplete
Checkpoint->>HF: snapshot_download(hf_model_id)
end
Converter->>BypassDetector: enumerate expected bypass runs, load_bypass_state
BypassDetector->>Converter: list of incomplete runs -> launch bypass if needed
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Contributor
|
Separius
force-pushed
the
ssameni/puzzletron-bypass-3-integration
branch
4 times, most recently
from
00c7d13 to
a715b55
Compare
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
Separius
force-pushed
the
ssameni/puzzletron-bypass-3-integration
branch
from
a715b55 to
9648a91
Compare
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
Contributor
Author
|
/ok to test 3cb4f61 |
Contributor
Author
|
/claude review |
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/puzzletron/puzzletron_nas_plugin.py`:
- Around line 276-290: The in-function import "from huggingface_hub import
snapshot_download" inside the block that checks Path(input_model_path).exists()
lacks the required brief justification; add a short comment immediately above
that import (e.g., "Guard optional dependency: huggingface_hub is imported here
to avoid requiring it unless HF model auto-download is needed") to explain why
the import is local, keeping it next to the import in the same block where
input_model_path, snapshot_download, mprint, and Path are used.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 65736e30-6d7d-46bb-bd1b-86ffc0900785
📒 Files selected for processing (21)
examples/puzzletron/Nemotron-3-Nano-30B-A3B-Base-BF16.mdexamples/puzzletron/README.mdexamples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/bypass/defaults.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/bypass/defaults.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/nemotron-3-nano-30b-a3b-with-bypass.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/nemotron-3-nano-30b-a3b.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/pruning/kv_heads_pruning.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/pruning/pruning_defaults.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/validate_model_defaults.yamlexamples/puzzletron/configs/nemotron-3-nano-30b-a3b/validate_solutions_defaults.yamlexamples/puzzletron/main.pymodelopt/torch/puzzletron/mip/run_puzzle.pymodelopt/torch/puzzletron/puzzletron_nas_plugin.pymodelopt/torch/puzzletron/replacement_library/build_replacement_library.pytests/gpu/torch/puzzletron/test_bypass.pytests/gpu/torch/puzzletron/test_bypass_checkpoint_utils.pytests/gpu/torch/puzzletron/test_puzzletron.pytests/unit/torch/puzzletron/test_puzzletron_nas_plugin.pytests/unit/torch/puzzletron/test_puzzletron_progress.pytests/unit/torch/puzzletron/test_replacement_library_bypass_config.py
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
Contributor
Author
|
/ok to test e8b73f6 |
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com> # Conflicts: # examples/puzzletron/main.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1470 +/- ##
===========================================
+ Coverage 56.40% 77.30% +20.90%
===========================================
Files 506 507 +1
Lines 55486 55694 +208
===========================================
+ Hits 31295 43056 +11761
+ Misses 24191 12638 -11553
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
Contributor
Author
|
/ok to test fbfa166 |
Contributor
Author
|
/claude review |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
ssameni/puzzletron-bypass-1-prereqs: shared prerequisitesssameni/puzzletron-bypass-2-core: bypass distillation coreWhat Changed
target_num_kv_heads.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:
Summary by CodeRabbit