consolidate mbridge distillation: merge distill_hf.py into distill.py#1220
Open
j-rausch wants to merge 6 commits intofeature/puzzletronfrom
Open
consolidate mbridge distillation: merge distill_hf.py into distill.py#1220j-rausch wants to merge 6 commits intofeature/puzzletronfrom
j-rausch wants to merge 6 commits intofeature/puzzletronfrom
Conversation
…still.py Signed-off-by: jrausch <jrausch@nvidia.com> Signed-off-by: root <root@pool0-00848.cm.cluster>
j-rausch
requested review from
ChenhanYu and
jenchen13
and removed request for
a team
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:
📝 WalkthroughWalkthroughThe PR consolidates distillation documentation into a single megatron_bridge location, adds optional HuggingFace export capabilities to the megatron_bridge distillation script with new CLI flags and rank-conditional export logic, removes the redundant puzzletron-specific distillation script and documentation, extends test coverage to support AnyModel distillation, and simplifies GitHub Actions workflows by removing unused test jobs. Changes
Sequence Diagram(s)sequenceDiagram
participant Script as Distillation Script
participant Train as Distributed Training
participant Export as AutoBridge Export
participant FS as File System
Script->>Train: Run distillation loop
Train->>Train: Training complete on all ranks
Script->>Train: dist.cleanup() all ranks
Train->>Script: Cleanup done
rect rgba(100, 200, 100, 0.5)
Note over Script,FS: HF Export Phase (rank 0 only)
Script->>Export: Call export_ckpt(megatron_ckpt)
Export->>FS: Write HF format checkpoint
FS->>Script: Export complete
Script->>FS: Copy student config.json to HF dir
FS->>Script: Copy complete
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
kevalmorabia97
requested review from
AAnoosheh,
danielkorzekwa and
kevalmorabia97
and removed request for
ChenhanYu and
jenchen13
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/megatron_bridge/distill.py`:
- Line 304: The shutil.copy call copying config.json from args.student_hf_path
will fail for remote HuggingFace model IDs; update the logic in distill.py
(around the shutil.copy usage) to detect if args.student_hf_path is a remote
model ID and in that case use the HuggingFace API (e.g., hf_hub_download or
transformers.AutoConfig.from_pretrained / AutoConfig.save_pretrained) to fetch
the config and write it to args.hf_export_path/config.json, otherwise keep the
local shutil.copy behavior; reference args.student_hf_path and
args.hf_export_path so the code handles both local paths and remote model IDs.
In `@tests/_test_utils/torch/puzzletron/utils.py`:
- Line 89: The hardcoded trust_remote_code=True must be made
caller-configurable: add a boolean parameter trust_remote_code: bool = False to
the function that loads the HF model (the function around lines 66-72 in
tests/_test_utils/torch/puzzletron/utils.py), then pass that parameter into
AutoConfig.from_pretrained(...) and any other pretrained loaders (e.g., the call
at line ~152) instead of the hardcoded True; ensure the new parameter defaults
to False and is threaded through to all transformer loading calls
(AutoConfig.from_pretrained and tokenizer/model loading sites) so callers can
opt in when necessary.
🪄 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: Pro
Run ID: 5e8fb78d-7646-49ef-ad4c-8b49d491c83b
📒 Files selected for processing (8)
examples/megatron_bridge/README.mdexamples/megatron_bridge/distill.pyexamples/puzzletron/README.mdexamples/puzzletron/mbridge_distillation/README.mdexamples/puzzletron/mbridge_distillation/distill_hf.pytests/_test_utils/torch/puzzletron/utils.pytests/examples/megatron_bridge/test_distill.pytests/examples/puzzletron/mbridge_distillation/test_distill_hf.py
💤 Files with no reviewable changes (3)
- examples/puzzletron/mbridge_distillation/README.md
- tests/examples/puzzletron/mbridge_distillation/test_distill_hf.py
- examples/puzzletron/mbridge_distillation/distill_hf.py
| For more details, you can refer to the checkpoint conversion scripts in the [Megatron-Bridge README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion). | ||
| For more details, see the [Megatron-Bridge conversion README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion). | ||
|
|
||
| > **Known limitation:** HF export does not yet work for Puzzletron AnyModel (heterogeneous) checkpoints -- Megatron-Bridge cannot reload heterogeneous configs from saved checkpoints. Standard models export correctly with both methods. |
Collaborator
There was a problem hiding this comment.
Did you test this in nemo:26.02.00 or nemo:26.02.01? It was fixed in 26.02.01. Please give that a try again
|
|
||
| > **Known limitation:** HF export does not yet work for Puzzletron AnyModel (heterogeneous) checkpoints -- Megatron-Bridge cannot reload heterogeneous configs from saved checkpoints. Standard models export correctly with both methods. | ||
|
|
||
| ### Distillation Results |
Collaborator
There was a problem hiding this comment.
Can you create results/puzzletron.md file and move the results there and add a reference to it here? I plan to add minitron distillation results also so this way we can keep this doc clean.
Collaborator
There was a problem hiding this comment.
Alternatively we can keep in examples/puzzletron/README.md also since thats where actual pruning is happening. Either is fine
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
force-pushed
the
jrausch/distillation-consolidation
branch
from
4ed9996 to
8c2fa10
Compare
kevalmorabia97
requested review from
kevalmorabia97
and removed request for
a team
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## feature/puzzletron #1220 +/- ##
======================================================
+ Coverage 59.90% 61.80% +1.89%
======================================================
Files 453 463 +10
Lines 47957 48440 +483
======================================================
+ Hits 28728 29937 +1209
+ Misses 19229 18503 -726
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Collaborator
|
/ok to test 8c2fa10 |
kevalmorabia97
force-pushed
the
feature/puzzletron
branch
from
2669a26 to
977d60a
Compare
kevalmorabia97
force-pushed
the
jrausch/distillation-consolidation
branch
from
3c0fa5a to
0c2a2ee
Compare
Collaborator
|
/ok to test 0c2a2ee |
Collaborator
|
/ok to test 286916a |
Collaborator
|
/ok to test 494cf38 |
…tion Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
force-pushed
the
jrausch/distillation-consolidation
branch
from
494cf38 to
3e39174
Compare
Collaborator
|
/ok to test 3e39174 |
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
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
2 participants
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
examples/puzzletron/mbridge_distillation/distill_hf.py(AnyModel-specific) intoexamples/megatron_bridge/distill.py(general)--hf_export_path/--student_hf_modelargs for inline HF export after distillation.tests/examples/megatron_bridge/test_distill.pyvocab_size=128(instead of default 102) for TP divisibility including 8.megatron_bridge/README.mdLimitation discovered during consolidation:
HF export via
--hf_export_pathseems to currently not work for Puzzletron AnyModel (heterogeneous) checkpoints. Megatron-Bridge'sexport_ckptcannot reload heterogeneous model configs from saved checkpoints (heterogeneous_layers_config_encoded_jsonisNoneduring__post_init__inheterogeneous_config.py). This affects both inline--hf_export_pathand the separateconvert_checkpoints.py exportscript.The original
distill_hf.pyREADME documented this as supported, but I think it might have been broken there too (on the side of Megatron-Bridge). The consolidated README now documents this as a known limitation. HF export for standard models works fine via both methods.Summary by CodeRabbit
Release Notes
New Features
Documentation
Bug Fixes & Improvements