Save hf checkpoint at every valitation iteration during distillation.#1897
Save hf checkpoint at every valitation iteration during distillation.#1897danielkorzekwa wants to merge 25 commits into
Conversation
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
|
/claude review |
|
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:
📝 WalkthroughWalkthroughAdds per-validation HuggingFace exports, token-budgeted preprocessing and blend generation, plus tests and documentation for the new iterative workflows. ChangesIterative distillation tooling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant distill.py CLI
participant pretrain
participant _HFValidationExportCallback
participant AutoConfig
participant Filesystem
distill.py CLI->>pretrain: pretrain(config, callbacks=[callback])
pretrain->>_HFValidationExportCallback: on_eval_end(iteration)
_HFValidationExportCallback->>Filesystem: create iter_<iteration>/ export path
_HFValidationExportCallback->>AutoConfig: save_pretrained(...)
AutoConfig->>Filesystem: write config.json
_HFValidationExportCallback->>pretrain: torch.distributed.barrier()
sequenceDiagram
participant CLI
participant megatron_preprocess_data
participant process_hf_split
participant process_json_file
participant _encode_docs
CLI->>megatron_preprocess_data: --max_tokens
megatron_preprocess_data->>process_hf_split: remaining_tokens
megatron_preprocess_data->>process_json_file: remaining_tokens
process_hf_split->>_encode_docs: may_stop_early=True
process_json_file->>_encode_docs: may_stop_early=True
process_hf_split-->>megatron_preprocess_data: stop at token budget
process_json_file-->>megatron_preprocess_data: stop at token budget
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| ) | ||
| # TODO: Use distill(..., callbacks=[callback]) once Megatron-Bridge supports callbacks. | ||
| pretrain(config, forward_step_modelopt, callbacks=[callback]) | ||
| else: |
There was a problem hiding this comment.
[SUGGESTION] When --hf_validation_export_path is set, the training entrypoint switches from distill(config) to pretrain(config, forward_step_modelopt, callbacks=[callback]). This creates two distinct training code paths that must stay behaviorally identical — otherwise enabling validation export silently changes how the model is trained, not just whether checkpoints are dumped.
If distill() does any distillation-specific setup beyond pretrain + forward_step_modelopt (e.g. loss-balancer wiring, KD config injection, provider hooks), that setup would be skipped on the export path. The assert isinstance(config.model, DistillationProvider) guard suggests you've considered this, but it would be worth (1) confirming distill() is genuinely just pretrain(config, forward_step_modelopt, ...) under the hood, and (2) leaving a one-line comment near the fork noting that equivalence, so a future change to distill() doesn't quietly diverge the two paths. The existing TODO about distill(..., callbacks=...) partially covers this, but the equivalence risk is the part worth calling out explicitly.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
🧹 Nitpick comments (1)
examples/megatron_bridge/distill.py (1)
209-214: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winReload of
AutoConfigon every validation export; duplicated with the final-export path.
AutoConfig.from_pretrained(self.student_hf_path, ...)is re-executed on everyon_eval_endcall even thoughstudent_hf_path/trust_remote_codenever change across the run — this repeats file/Hub I/O unnecessarily. The same load-then-save_pretrainedpattern (with an identical comment) is also duplicated inmain()at Lines 557-560.Consider loading the config once in
__init__and reusing it inon_eval_end, and/or extracting a small shared helper used by both this callback and the final--hf_export_pathblock to avoid drift between the two copies.♻️ Proposed fix to cache the config once
def __init__( self, export_dir: str, student_hf_model: str, student_hf_path: str, trust_remote_code: bool, ) -> None: self.export_dir = Path(export_dir) - self.student_hf_path = student_hf_path - self.trust_remote_code = trust_remote_code self._last_exported_iteration: int | None = None self.bridge = AutoBridge.from_hf_pretrained( student_hf_model, trust_remote_code=trust_remote_code ) + self._student_config = AutoConfig.from_pretrained( + student_hf_path, trust_remote_code=trust_remote_code + )if dist.rank() == 0: - # Preserve the student architecture from student_hf_path, including heterogeneous - # layer changes; AutoConfig supports both local paths and Hugging Face model IDs. - AutoConfig.from_pretrained( - self.student_hf_path, trust_remote_code=self.trust_remote_code - ).save_pretrained(output_path) + self._student_config.save_pretrained(output_path)🤖 Prompt for 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. In `@examples/megatron_bridge/distill.py` around lines 209 - 214, `DistillCallback.on_eval_end` is reloading the student config on every validation export even though `student_hf_path` and `trust_remote_code` are ثابت, and the same `AutoConfig.from_pretrained(...).save_pretrained(...)` logic is duplicated in `main()`. Cache the loaded config once in `DistillCallback.__init__` (or a shared helper) and reuse it in `on_eval_end`, then have the final `--hf_export_path` export path call the same helper to keep the behavior in sync and avoid repeated I/O.
🤖 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 `@examples/megatron_bridge/distill.py`:
- Around line 167-218: The _HFValidationExportCallback currently saves a new
Hugging Face export on every validation with no cleanup, which can cause
unbounded disk growth. Add a retention setting such as keep_last_n to the
callback, and in on_eval_end prune older iter_* export directories after a
successful save, keeping only the most recent exports. Use the existing export
flow in _HFValidationExportCallback and mirror the retention approach already
used for main checkpointing (most_recent_k) so the logic is easy to locate and
consistent.
---
Nitpick comments:
In `@examples/megatron_bridge/distill.py`:
- Around line 209-214: `DistillCallback.on_eval_end` is reloading the student
config on every validation export even though `student_hf_path` and
`trust_remote_code` are ثابت, and the same
`AutoConfig.from_pretrained(...).save_pretrained(...)` logic is duplicated in
`main()`. Cache the loaded config once in `DistillCallback.__init__` (or a
shared helper) and reuse it in `on_eval_end`, then have the final
`--hf_export_path` export path call the same helper to keep the behavior in sync
and avoid repeated I/O.
🪄 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: c41374ad-ecce-4dc2-b44d-66aedfb9cdc0
📒 Files selected for processing (3)
examples/megatron_bridge/README.mdexamples/megatron_bridge/distill.pytests/examples/megatron_bridge/test_distill.py
| class _HFValidationExportCallback(Callback): | ||
| """Export the live student to Hugging Face format after each validation stage.""" | ||
|
|
||
| def __init__( | ||
| self, | ||
| export_dir: str, | ||
| student_hf_model: str, | ||
| student_hf_path: str, | ||
| trust_remote_code: bool, | ||
| ) -> None: | ||
| self.export_dir = Path(export_dir) | ||
| self.student_hf_path = student_hf_path | ||
| self.trust_remote_code = trust_remote_code | ||
| self._last_exported_iteration: int | None = None | ||
| self.bridge = AutoBridge.from_hf_pretrained( | ||
| student_hf_model, trust_remote_code=trust_remote_code | ||
| ) | ||
|
|
||
| def on_eval_end(self, context) -> None: | ||
| """Export the student at the iteration that was just validated.""" | ||
| iteration = context.state.train_state.step | ||
| # The final iteration can be validated both on its regular interval and after training. | ||
| # Avoid exporting and overwriting the same Hugging Face checkpoint twice. | ||
| if iteration == self._last_exported_iteration: | ||
| return | ||
| output_path = self.export_dir / f"iter_{iteration:07d}" | ||
| print_rank_0(f"Exporting validation checkpoint {iteration} to {output_path}") | ||
|
|
||
| # DistillationModel is the student with teacher and KD-loss modules attached. Hide the | ||
| # auxiliary modules temporarily so the Hugging Face export contains only student weights. | ||
| with contextlib.ExitStack() as stack: | ||
| for model_chunk in unwrap_model(context.model): | ||
| if isinstance(model_chunk, mtd.DistillationModel): | ||
| stack.enter_context(model_chunk.hide_teacher_model()) | ||
| stack.enter_context(model_chunk.hide_loss_modules()) | ||
| self.bridge.save_hf_pretrained( | ||
| context.model, | ||
| output_path, | ||
| show_progress=True, | ||
| strict=True, | ||
| ) | ||
|
|
||
| if dist.rank() == 0: | ||
| # Preserve the student architecture from student_hf_path, including heterogeneous | ||
| # layer changes; AutoConfig supports both local paths and Hugging Face model IDs. | ||
| AutoConfig.from_pretrained( | ||
| self.student_hf_path, trust_remote_code=self.trust_remote_code | ||
| ).save_pretrained(output_path) | ||
| torch.distributed.barrier() | ||
| self._last_exported_iteration = iteration | ||
|
|
||
|
|
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
No retention policy for per-validation HF exports — unbounded disk growth risk.
Every validation stage writes a full HuggingFace checkpoint to iter_<n>/ with no cap on how many are kept. For a realistic run (e.g. the README's --train_iters 15000 --eval_interval 100 example), that's ~150 full model copies with no cleanup — unlike the main Megatron checkpoint config in this same file, which bounds retention via most_recent_k=5 (Line 505). For multi-billion-parameter students this can exhaust disk and fail the job mid-training.
Consider adding a retention parameter (e.g. keep_last_n) to _HFValidationExportCallback that prunes older iter_* export directories, mirroring the most_recent_k pattern already used for the main checkpoint.
🤖 Prompt for 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.
In `@examples/megatron_bridge/distill.py` around lines 167 - 218, The
_HFValidationExportCallback currently saves a new Hugging Face export on every
validation with no cleanup, which can cause unbounded disk growth. Add a
retention setting such as keep_last_n to the callback, and in on_eval_end prune
older iter_* export directories after a successful save, keeping only the most
recent exports. Use the existing export flow in _HFValidationExportCallback and
mirror the retention approach already used for main checkpointing
(most_recent_k) so the logic is easy to locate and consistent.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Examples-only PR (3 files: examples/megatron_bridge/distill.py, its README, and tests/examples/megatron_bridge/test_distill.py). megatron.bridge is an external dependency not installed in this environment, so distill()/pretrain() internals could not be introspected; review is based on the diff plus ModelOpt's DistillationModel API.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
What I verified
- The export path correctly strips auxiliary modules before saving:
hide_teacher_model()andhide_loss_modules()exist inmodelopt/torch/distill/distillation_model.pyand are the same mechanism used by the minimal state-dict export path. - Distributed ordering is sound: all ranks run the collective
save_hf_pretrained, only rank 0 writesconfig.json, followed by abarrier(). - The
_last_exported_iterationguard sensibly prevents double-export/overwrite when the final iteration is validated both on its interval and post-training. - Argument validation was correctly widened to require
--student_hf_modelwhen either HF export flag is set.
Suggestion (non-blocking)
- Enabling
--hf_validation_export_pathswitches training fromdistill(config)topretrain(config, forward_step_modelopt, callbacks=[...]). Worth confirming these two paths are behaviorally identical and leaving a comment noting the equivalence, so a future change todistill()doesn't silently diverge the export path from the normal one.
Overall risk: low — additive, opt-in feature gated behind a new flag, backward compatible, with test coverage added.
Codecov Report✅ All modified and coverable lines are covered by tests.
Additional details and impacted files@@ Coverage Diff @@
## main #1897 +/- ##
===========================================
- Coverage 78.38% 64.95% -13.43%
===========================================
Files 518 525 +7
Lines 58269 59363 +1094
===========================================
- Hits 45674 38561 -7113
- Misses 12595 20802 +8207
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: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/utils/plugins/megatron_preprocess_data.py (1)
308-368: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAdd
max_tokensto file-backed dataset prefixes
process_json_file()returns the same.bin/.idxprefix regardless of token budget, so reruns with a differenttarget_tokenscan silently reuse stale files for"files"sources. This also leaves the file-backed path out of sync withprocess_hf_split()and the blend test’s_tokensexpectation. Consider appending the sametoken_tagthere.🤖 Prompt for 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. In `@modelopt/torch/utils/plugins/megatron_preprocess_data.py` around lines 308 - 368, `process_json_file()` currently builds output prefixes only from the input stem, so file-backed runs with different token budgets can reuse the same `.bin/.idx` artifacts. Update `process_json_file` in `megatron_preprocess_data.py` to append the same `token_tag` used by `process_hf_split()` when constructing `output_prefix`/`prefixes`, so the `process_json_file` and `process_hf_split` paths produce consistent, budget-specific dataset names. Ensure the new naming is derived from the existing `max_tokens`/token budget logic and is applied before checking for existing builders or returning skipped prefixes.
🤖 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 `@examples/dataset/prepare_data_blend.py`:
- Around line 88-141: The weight-to-token allocation in prepare_data_blend can
produce zero or negative max_tokens for the last source when YAML weights are
misconfigured, leading to an empty prefixes list and a later ZeroDivisionError.
Add upfront validation in prepare_data_blend/load_config that source weights sum
to about 100 (or otherwise ensure source_tokens stays positive), and also guard
the prefix_weight calculation after megatron_preprocess_data so an empty
prefixes result raises a clear configuration error instead of dividing by zero.
---
Outside diff comments:
In `@modelopt/torch/utils/plugins/megatron_preprocess_data.py`:
- Around line 308-368: `process_json_file()` currently builds output prefixes
only from the input stem, so file-backed runs with different token budgets can
reuse the same `.bin/.idx` artifacts. Update `process_json_file` in
`megatron_preprocess_data.py` to append the same `token_tag` used by
`process_hf_split()` when constructing `output_prefix`/`prefixes`, so the
`process_json_file` and `process_hf_split` paths produce consistent,
budget-specific dataset names. Ensure the new naming is derived from the
existing `max_tokens`/token budget logic and is applied before checking for
existing builders or returning skipped prefixes.
🪄 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: 9f8a8b4f-8080-4109-952c-6c7d31f8c071
📒 Files selected for processing (8)
examples/dataset/MEGATRON_DATA_PREP.mdexamples/dataset/prepare_data_blend.pyexamples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.mdexamples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.mdexamples/researcher_guide/README.mdmodelopt/torch/utils/plugins/megatron_preprocess_data.pytests/examples/dataset/test_prepare_data_blend.pytests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py
✅ Files skipped from review due to trivial changes (3)
- examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md
- examples/researcher_guide/README.md
- examples/dataset/MEGATRON_DATA_PREP.md
| for index, source in enumerate(sources): | ||
| weight = float(source["weight"]) | ||
| if total_tokens is None: | ||
| source_tokens = None | ||
| elif index == len(sources) - 1: | ||
| source_tokens = total_tokens - allocated_tokens | ||
| else: | ||
| source_tokens = round(total_tokens * weight / 100) | ||
| allocated_tokens += source_tokens | ||
|
|
||
| dataset = source["hf_dataset"] | ||
| source_dir = output_dir / f"{index:02d}_{dataset.replace('/', '--')}" | ||
| content_field = source["content_field"] | ||
| input_args: dict[str, Any] | ||
| if "files" in source: | ||
| raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") | ||
| paths = [ | ||
| hf_hub_download( | ||
| repo_id=dataset, | ||
| filename=file, | ||
| repo_type="dataset", | ||
| local_dir=raw_dir, | ||
| ) | ||
| for file in source["files"] | ||
| ] | ||
| input_args = {"jsonl_paths": paths} | ||
| else: | ||
| input_args = { | ||
| "hf_dataset": dataset, | ||
| "hf_name": source.get("config"), | ||
| "hf_split": source["split"], | ||
| "hf_max_samples_per_split": source.get("max_samples"), | ||
| "hf_streaming": True, | ||
| } | ||
|
|
||
| # Each prefix is the path shared by a tokenized Megatron .bin/.idx file pair. | ||
| prefixes = megatron_preprocess_data( | ||
| **input_args, | ||
| output_dir=source_dir, | ||
| tokenizer_name_or_path=tokenizer, | ||
| json_keys=content_field, | ||
| # Plain text lacks chat-template boundary tokens, so terminate each document with EOS. | ||
| append_eod=content_field == "text", | ||
| # Join lines in text documents by replacing each newline with a space. | ||
| strip_newlines=content_field == "text", | ||
| reasoning_content="inline" if content_field == "messages" else "strip", | ||
| # Guard against pathological records by capping each tokenized document at 256K tokens. | ||
| max_sequence_length=256_000, | ||
| max_tokens=source_tokens, | ||
| workers=workers, | ||
| ) | ||
| prefix_weight = weight / len(prefixes) | ||
| blend.extend((prefix_weight, prefix) for prefix in prefixes) | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Unvalidated weight allocation can crash with ZeroDivisionError on misconfigured YAML.
source_tokens for non-last sources is round(total_tokens * weight / 100) (Line 95), with the last source getting total_tokens - allocated_tokens (Line 93). If the configured weights don't sum to ~100 (or accumulate enough rounding error), the last source's source_tokens can end up <= 0. Inside megatron_preprocess_data, a non-positive max_tokens on the very first split/file causes an immediate break with zero prefixes returned (Lines 600-601, 628-629 of megatron_preprocess_data.py). Back here, prefix_weight = weight / len(prefixes) (Line 139) then raises ZeroDivisionError.
Since load_config/prepare_data_blend is the interface boundary for this user-authored YAML, consider validating that sources weights sum to 100 (within tolerance) up front, and/or guarding against an empty prefixes result with a clear error message instead of a raw ZeroDivisionError.
🛡️ Proposed validation
def prepare_data_blend(config_path: Path) -> list[tuple[float, str]]:
"""Download and tokenize the configured weighted data sources."""
config = load_config(config_path)
output_dir = Path(config["output_dir"])
output_dir.mkdir(parents=True, exist_ok=True)
target_tokens = config.get("target_tokens")
total_tokens = None if target_tokens is None else int(target_tokens)
tokenizer = str(config["tokenizer"])
+
+ total_weight = sum(float(source["weight"]) for source in config["sources"])
+ if not math.isclose(total_weight, 100, abs_tol=0.5):
+ raise ValueError(f"Source weights must sum to 100, got {total_weight}")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for index, source in enumerate(sources): | |
| weight = float(source["weight"]) | |
| if total_tokens is None: | |
| source_tokens = None | |
| elif index == len(sources) - 1: | |
| source_tokens = total_tokens - allocated_tokens | |
| else: | |
| source_tokens = round(total_tokens * weight / 100) | |
| allocated_tokens += source_tokens | |
| dataset = source["hf_dataset"] | |
| source_dir = output_dir / f"{index:02d}_{dataset.replace('/', '--')}" | |
| content_field = source["content_field"] | |
| input_args: dict[str, Any] | |
| if "files" in source: | |
| raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") | |
| paths = [ | |
| hf_hub_download( | |
| repo_id=dataset, | |
| filename=file, | |
| repo_type="dataset", | |
| local_dir=raw_dir, | |
| ) | |
| for file in source["files"] | |
| ] | |
| input_args = {"jsonl_paths": paths} | |
| else: | |
| input_args = { | |
| "hf_dataset": dataset, | |
| "hf_name": source.get("config"), | |
| "hf_split": source["split"], | |
| "hf_max_samples_per_split": source.get("max_samples"), | |
| "hf_streaming": True, | |
| } | |
| # Each prefix is the path shared by a tokenized Megatron .bin/.idx file pair. | |
| prefixes = megatron_preprocess_data( | |
| **input_args, | |
| output_dir=source_dir, | |
| tokenizer_name_or_path=tokenizer, | |
| json_keys=content_field, | |
| # Plain text lacks chat-template boundary tokens, so terminate each document with EOS. | |
| append_eod=content_field == "text", | |
| # Join lines in text documents by replacing each newline with a space. | |
| strip_newlines=content_field == "text", | |
| reasoning_content="inline" if content_field == "messages" else "strip", | |
| # Guard against pathological records by capping each tokenized document at 256K tokens. | |
| max_sequence_length=256_000, | |
| max_tokens=source_tokens, | |
| workers=workers, | |
| ) | |
| prefix_weight = weight / len(prefixes) | |
| blend.extend((prefix_weight, prefix) for prefix in prefixes) | |
| tokenizer = str(config["tokenizer"]) | |
| total_weight = sum(float(source["weight"]) for source in config["sources"]) | |
| if not math.isclose(total_weight, 100, abs_tol=0.5): | |
| raise ValueError(f"Source weights must sum to 100, got {total_weight}") |
🤖 Prompt for 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.
In `@examples/dataset/prepare_data_blend.py` around lines 88 - 141, The
weight-to-token allocation in prepare_data_blend can produce zero or negative
max_tokens for the last source when YAML weights are misconfigured, leading to
an empty prefixes list and a later ZeroDivisionError. Add upfront validation in
prepare_data_blend/load_config that source weights sum to about 100 (or
otherwise ensure source_tokens stays positive), and also guard the prefix_weight
calculation after megatron_preprocess_data so an empty prefixes result raises a
clear configuration error instead of dividing by zero.
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
…al (data up to 100m tokens) Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
tutorial Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
…ter_during_distill Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
| parser.add_argument("--lr", type=float, default=1e-4, help="Peak learning rate") | ||
| parser.add_argument("--min_lr", type=float, default=1e-5, help="Minimum learning rate") | ||
| parser.add_argument("--lr_warmup_iters", type=int, default=50, help="Number of LR warmup steps") | ||
| parser.add_argument( |
There was a problem hiding this comment.
when is reset optimizer used?
There was a problem hiding this comment.
reset optimizer was removed from this MR
| " in megatron distributed checkpoint format.\n" | ||
| ) | ||
|
|
||
| # TODO: Extend _HFValidationExportCallback to export --hf_export_path from the in-memory |
There was a problem hiding this comment.
would it be possible to export from the in memory student in this PR? loading an extra model is very expensive memory wise. you can probably do some async save of the in-memory student
There was a problem hiding this comment.
@AAnoosheh , @kevalmorabia97 any idea if we could do it?
There was a problem hiding this comment.
We are exporting now offline. Once distillation is finished, the export script is called to export multiple mbridge checkpoints
| | 0 | 0.5187 | 1.4739 | 48.69% (full) | 23.09% (full) | | ||
| | 0.7M | 0.1919 | 1.0931 | 58.74% | 13.14% | | ||
| | 3.3M | 0.1342 | 1.0550 | 60.56% | 14.29% | | ||
| | 39.3M | 0.0675 | 1.0296 | 64.63% | 6.14% | |
There was a problem hiding this comment.
why does MMLU pro drop so low here?
There was a problem hiding this comment.
@jenchen13 The catastrophic drop in MMLU to ~6% happens because the Multi-Layer Perceptron (MLP) is losing its non-linearity and collapsing into a purely linear projection during export.
@danielkorzekwa I did a deep dive into your latest commit (0fee306). While it successfully protects the live teacher by detaching it, the student's config is still forced to activation_func = None to remain pickleable.
When this sanitized config hits the Model-Optimizer export pipeline (modelopt/torch/export/layer_utils.py), a silent failure occurs:
-
_get_hidden_actrelies onhasattr(module, "activation_func"). In Python, this evaluates toTrueeven if the attribute is explicitlyNone. - Because it bypasses this check, Hugging Face receives the un-mappable
Nonetype and gracefully initializes the MLP without an activation function instead of crashing. - Without the activation function,
$W_2(\sigma(W_1(x)))$ becomes$W_2(W_1(x))$ , lobotomizing the model's reasoning.
To permanently fix this, we need to enforce a fail-fast boundary in _get_hidden_act: trapping the explicit None state with getattr, unwrapping functools.partial objects, and explicitly validating against Hugging Face's ACT2FN dictionary.
I have the exact structural code fix and the unittest.mock regression tests ready locally. Let me know if you'd like me to open a PR against your branch to inject it!
There was a problem hiding this comment.
yes, please open MR into this dkorzekwa/save_hf_checkpoint_for_every_val_iter_during_distill, thanks
Add @AAnoosheh , @kevalmorabia97 as reviewers
There was a problem hiding this comment.
yes, please open MR into this dkorzekwa/save_hf_checkpoint_for_every_val_iter_during_distill, thanks
Add @AAnoosheh , @kevalmorabia97 as reviewers
Done! I just opened PR #1997 against your branch with the strict validation fixes and regression tests. I've also tagged @AAnoosheh and @kevalmorabia97 for review over there as requested!
There was a problem hiding this comment.
lets abandon #1997. Now we export to HF after distillation process is completed. See the latest commit in this branch.
There was a problem hiding this comment.
@jenchen13 regarding, why does MMLU pro drop so low here? - it is due to overfitting to nemotron v2 style, added to the README.md:
Nemotron v2 KD continues to decrease, but its MMLU-Pro score remains below the pruned baseline. The severe
MMLU-Pro regression appears to come from overfitting to Nemotron v2-style responses, even though MMLU-Pro prompts
ask the model to answer in a specific multiple-choice style.
There was a problem hiding this comment.
so the MMLU pro drop is not because you are discarding non-linearities during export? does @Surya-5555 comment still apply?
There was a problem hiding this comment.
so the MMLU pro drop is not because you are discarding non-linearities during export? does @Surya-5555 comment still apply?
@jenchen13 Yes, my previous comment applied to the architecture prior to commit 5435743.
As Daniel mentioned, we changed the approach and abandoned my fix (PR #1997) because the workflow has pivoted to exporting all HF checkpoints after the distillation process is finished, rather than live from memory.
In the previous live-export setup, the activation_func = None hack was absolutely causing a catastrophic MLP linear collapse, which resulted in those ~6% MMLU scores. However, because the new offline approach no longer needs to sanitize the in-memory config, this structural graph corruption is bypassed entirely.
So, if there are still MMLU regressions occurring in the newly exported models, they are indeed due to training dynamics (like the overfitting to the Nemotron v2 style that Daniel noted in the README), not an export bug!
…ntermediate HF checkpoints during distillation. Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
…ter_during_distill Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
…checkpoints, once distillation is finished. Keep all mcore distilaltion on disk. Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
jenchen13
left a comment
jenchen13
left a comment
There was a problem hiding this comment.
is it necessary to add the export script to exporting multiple checkpoints? it seems like we could just have one export script, and if you want to export multiple checkpoints that's something done in a bash script calling export.py
| --megatron_path /tmp/distill-out/checkpoints/iter_0000500 \ | ||
| --hf_export_path /tmp/distilled-vlm-hf_iter_0000500 | ||
|
|
||
| Example selected validation iterations: |
There was a problem hiding this comment.
why is there also a export_quantized_megatron_to_hf.py file? what are the differences between exporting a quantized model vs distilled model? this could probably be consolidated
There was a problem hiding this comment.
That one handles quantized megatron to unified export format. It will be dropped once we natively support quantized export in MBridge (any parallelism) and fold it in quantize.py
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
| To use Weights & Biases for logging, set the `WANDB_API_KEY` environment variable and pass the `--wandb_project` argument. | ||
| Optionally, you can also pass `--wandb_entity` and `--wandb_exp_name` arguments to group runs under a project and experiment name. | ||
|
|
||
| To measure the initial student's CE and distillation losses, add `--validate_only` to the command. |
There was a problem hiding this comment.
this seems outdated. I dont see --validate_only anywhere
| ```text | ||
| hf_validation/ | ||
| ├── iter_0000100/ | ||
| ├── iter_0000200/ | ||
| └── iter_0000300/ |
There was a problem hiding this comment.
suggestion: What if we save in same checkpoints directory as hf_iter_... name? For example:
checkpoints/
├── hf_iter_0000100/
├── iter_0000200/
├── hf_iter_0000100/
├── iter_0000200/
Feel free to disregard if you prefer your current approach
| ] | ||
|
|
||
|
|
||
| def test_export_distilled_megatron_iterations(tmp_path: Path, num_gpus): |
There was a problem hiding this comment.
what if we just update existing test_distill.py::test_distill_vlm test? It already runs distill + export_distilled_megatron_to_hf.py so we just need to update flags there.
The other test_distill_llm covers the distill.py --hf_export_path path anyways
Saves ~3-4mins CI/CD time
5 participants
What does this PR do?
Save hf checkpoint at every valitation iteration during distillation.
Addition functionality:
--validate_onlyinexamples/megatron_bridge/distill.pyto enable computing validation losses for iter 0--reset_optimizerinexamples/megatron_bridge/distill.pyto enable not using presaved optimizer, e.g., when changing the number of train iters.Usage
Testing
Before your PR is "Ready for review"
Summary by CodeRabbit
--validate_onlyfor student-only validation at iteration 0.--hf_validation_export_pathwith--hf_validation_export_intervalto export validated student HuggingFace artifacts during distillation.prepare_data_blend.pyfor YAML-driven token-budgeted data blends.--max_tokensto stop Megatron preprocessing after a token budget.