Add Megatron-Bridge prune & quantize launcher pipelines - #2031
Add Megatron-Bridge prune & quantize launcher pipelines#2031kevalmorabia97 wants to merge 3 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
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 adds post-run accuracy and pruning-score gates. It also adds inline launcher commands, per-task requirements, configurable Docker users, cache environment overrides, workflow examples, documentation, and validation tests. ChangesEvaluation and pruning gates
Launcher task workflows
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant EvaluationCLI
participant EvalRunner
participant ResultsJSON
EvaluationCLI->>EvaluationCLI: validate threshold, output path, and one task
EvaluationCLI->>EvalRunner: execute evaluation
EvalRunner->>ResultsJSON: write results JSON
EvaluationCLI->>ResultsJSON: read task accuracy
ResultsJSON-->>EvaluationCLI: return accuracy
EvaluationCLI-->>EvaluationCLI: report PASS or exit status 1
sequenceDiagram
participant SandboxPipeline
participant DockerExecutor
participant TaskShell
SandboxPipeline->>SandboxPipeline: resolve task fields
SandboxPipeline->>DockerExecutor: configure Docker user
DockerExecutor->>TaskShell: run task
TaskShell->>TaskShell: install requirements
TaskShell-->>DockerExecutor: execute inline or script command
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
SCOPE: Reviewed all 9 files in the PR authoritative changed-file list: tools/launcher/{core.py, slurm_config.py, docs/configuration.md}, the two new example YAMLs, three test files, and examples/llm_eval/lm_eval_hf.py. Pure tooling/examples change — no modelopt/ source, mode registration, config schema, or export code touched, so ModeState/Export/Algorithm categories do not apply. (A raw two-dot diff against the shallow base also surfaced unrelated MiniMax deletions and hf_ptq README edits; those are NOT in the PR file list and were treated as base-mismatch noise, not reviewed.)
FINDINGS: CRITICAL 0, IMPORTANT 0, SUGGESTION 1.
- SUGGESTION: lm_eval_hf.py accuracy gate uses sorted(glob(...))[-1] to pick the newest results file, which orders by path (directory name first), not timestamp. Robust only when --output_path is not reused across runs; shipped examples use ephemeral per-run paths so non-blocking. Suggested max(files, key=os.path.getmtime).
ASSESSMENT: Low risk. Two substantive paths check out:
- core.py inline/reqs/docker_user — new SandboxTask fields (inline, reqs, reqs_file) and SlurmConfig.docker_user default to None, slurm_factory gains an optional kwarg — backward compatible. shlex imported; reqs tokens shlex-quoted so specifiers like transformers<5 stay literal; args+inline correctly rejected; docker_user falls back to host uid:gid.
- lm_eval_hf.py gate + backend handling — for non-hf backends ModelOpt-only kwargs are dropped from the namespace (not injected into model_args), so --model vllm on a quantized checkpoint will not break the constructor. Metric extraction matches acc / acc, while excluding acc_stderr / acc_norm; single-task validation fails fast before the expensive eval.
Tests cover the new fields, the guard, and docker_user override. LGTM.
Codecov Report✅ All modified and coverable lines are covered by tests.
Additional details and impacted files@@ Coverage Diff @@
## main #2031 +/- ##
==========================================
- Coverage 75.79% 69.24% -6.55%
==========================================
Files 518 519 +1
Lines 58658 61015 +2357
==========================================
- Hits 44459 42251 -2208
- Misses 14199 18764 +4565
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:
|
ac86810 to
e8b8914
Compare
|
6c137be to
9d4ee18
Compare
9d4ee18 to
4aef972
Compare
![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: 8
🤖 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/llm_eval/lm_eval_hf.py`:
- Around line 339-351: Update the lower_bound validation branch around
_pop_gate_bound and _enforce_accuracy_gate to require a non-empty output_path
before cli.execute(args) runs. Raise the existing appropriate validation error
immediately when --accuracy_lower_bound is used without --output_path, while
preserving the single-task validation and normal execution path.
In `@tests/examples/megatron_bridge/test_prune_minitron.py`:
- Line 73: Strengthen the score-gate coverage in the real-script tests around
the existing score_lower_bound cases by adding an invocation with an impossible
MMLU bound such as 1.1. Assert that this command exits non-zero, while
preserving the existing successful-path checks for valid bounds.
In `@tools/launcher/core.py`:
- Around line 788-791: Enforce the inline-task command-source contract before
command construction in tools/launcher/core.py around the existing task.inline
validation: reject tasks that define both script and inline, and reject tasks
that define neither, while preserving the existing inline args restriction.
Update tools/launcher/tests/test_examples_resolve.py in the example validation
logic to reject args when inline is configured, matching runtime behavior.
- Around line 848-852: Update the marker construction in the launcher command
that defines reqs_prefix to remove the # nosec suppression and derive a task-
and node-unique synchronization path using the available job/task and node
identifiers. Ensure the resulting path remains safely rooted in the temporary
directory and is shared by the intended local ranks, while preserving the
existing wait-and-touch synchronization behavior.
- Around line 855-857: Update the reqs_prefix branch around script_cmd
construction to shell-quote task.script and every value in task_args before
concatenating them into the inline bash command. Preserve the existing argument
order and run.Script invocation while ensuring YAML/global-variable contents are
passed literally rather than interpreted by the shell.
In `@tools/launcher/docs/configuration.md`:
- Around line 63-66: Update the inline task documentation to state that args
must be omitted when inline is used, because run_jobs() raises ValueError if
inline is combined with non-empty args; remove the claim that args are ignored
while preserving the existing inline command behavior description.
In
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml`:
- Around line 19-20: Add the required environment variables to both model
configurations: in
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml
lines 19-20, add MLM_MODEL_CFG using the Nemotron HuggingFace repository ID; in
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml
lines 16-17, add MLM_MODEL_CFG with the same repository ID and QUANT_CFG
matching the CLI quantization configuration.
In `@tools/launcher/tests/test_examples_resolve.py`:
- Around line 79-91: Update the task-shape validation in the test around the
script/inline exclusivity assertions to reject any task that combines a
non-empty inline command with args, matching run_jobs() behavior. Add an
assertion covering inline tasks’ args absence so the test exercises this
runtime-invalid shape while preserving existing script/inline and single-line
checks.
🪄 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: 4b2432c3-8b73-47b4-85a4-b0f1bbdc814f
📒 Files selected for processing (12)
examples/llm_eval/lm_eval_hf.pyexamples/megatron_bridge/prune_minitron.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/examples/megatron_bridge/test_prune_minitron.pytools/launcher/core.pytools/launcher/docs/configuration.mdtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yamltools/launcher/slurm_config.pytools/launcher/tests/test_docker_execution.pytools/launcher/tests/test_examples_resolve.pytools/launcher/tests/test_yaml_formats.py
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Mostly a tidy set of additive launcher/example features, and the E2E evidence in the PR body is appreciated. A few substantive issues, though:
Correctness / backward compat
- Adding
"best"toMCoreMinitronSearcher.default_state_dictis not backward compatible for cached score checkpoints:BaseSearcher.load_search_checkpoint()is called withstrict=Trueand assertsstate_dict.keys() == self.state_dict().keys(), so any pre-existing--prune_intermediate_ckpt/modelopt_pruning_scoresdir written by an older version now fails withKeys in checkpoint don't match!instead of resuming. The PR body claims "new state-dict key is additive" — it is additive for writers, not readers. self.bestshadowsBaseSearcher.best, which is the documented search-result state dict returned byBaseSearcher.search() -> SearchStateDict(reset_searchsetsself.best: SearchStateDict = {}).MCoreMinitronSearcher.search()now returns aCandidateSubnet | None, breaking that annotation/contract (and likely tripping mypy on the attribute override). Please rename (e.g.best_candidate).lm_eval_hf.pynow silently drops--quant_cfg/--sparse_cfg/--auto_quantize_bitsfor non-hfbackends. A user running--model vllm --quant_cfg FP8_DEFAULT_CFGgets an eval of the unquantized model with no warning — and with--accuracy_lower_boundthat even reports PASS. This should raise instead of silently ignoring. Alsois_hfmisses thehuggingfacealias andhf-multimodal, both of which are HFLM-based.
Tests
4. The new gate helpers (_requested_tasks, _enforce_accuracy_gate) are pure functions with real edge cases (mtime selection among nested results*.json, "acc,none" vs acc_stderr, missing task) and have no tests. Similarly, the new run_jobs logic (reqs_prefix shell string, args+inline ValueError) is untested — the added test_examples_resolve assertion covers YAML-level script/inline exclusivity, not the args+inline guard the PR body claims.
5. get_default_env now reads ambient HF_HOME/TRITON_CACHE_DIR, but the existing tests/test_core.py::TestGetDefaultEnv assertions ("/cicd/hf-cache") weren't updated to monkeypatch.delenv, so they'll fail on any dev box/CI runner that exports those vars. No test covers the new override path either.
Robustness / docs
6. The pip barrier marker /tmp/.modelopt_launcher_reqs_done is not run-unique, and under enroot/pyxis /tmp is commonly the host /tmp; a stale marker from an earlier task on the same node makes non-zero local ranks skip the wait, reintroducing exactly the concurrent-pip race the barrier guards against. Ranks also proceed after the 600 s timeout even when rank 0's install failed.
7. docs/configuration.md says args "is ignored" with inline, but run_jobs raises ValueError.
Design (per the complexity gate): inline/reqs extend the existing launcher config rather than adding a new subsystem, which is good, but they do add a third way to express a task command alongside the documented common/**/*.sh wrapper convention (CLAUDE.md: "Scripts go in common/") and typed _target_ tasks — and common/megatron_bridge/import/import.sh shows the wrapper pattern already covers this flow (a wrapper could also do the pip install + rank-0 barrier in bash, making reqs/reqs_file unnecessary). The PR body describes the new fields but doesn't say why the wrapper-script or typed-task route was rejected; please add that rationale, and update CLAUDE.md's conventions list so contributors know which of the three to reach for.
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: 2
🤖 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/prune_minitron.py`:
- Around line 683-693: The score-gate logic currently runs after checkpoint
export, allowing failed pruned models to remain at the output path and bypass
evaluation on reruns. Move the `args.score_lower_bound` validation block,
including the `pruning_scores["best"]` check and `sys.exit(1)`, to immediately
after `mtp.prune()` and before either artifact-save path.
In `@tools/launcher/core.py`:
- Around line 849-852: Update the reqs_prefix shell flow around marker creation
so nonzero local ranks fail when the marker is absent after the wait loop,
preventing fi && task_cmd from running without installed dependencies. Preserve
successful continuation when {marker} appears and propagate an explicit failure
status when the 600-iteration wait expires.
🪄 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: 2e1eb71d-dedb-4e11-adee-27243a5a5b6e
📒 Files selected for processing (12)
examples/llm_eval/lm_eval_hf.pyexamples/megatron_bridge/prune_minitron.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/examples/megatron_bridge/test_prune_minitron.pytools/launcher/core.pytools/launcher/docs/configuration.mdtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yamltools/launcher/slurm_config.pytools/launcher/tests/test_docker_execution.pytools/launcher/tests/test_examples_resolve.pytools/launcher/tests/test_yaml_formats.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tools/launcher/docs/configuration.md
4aef972 to
d4cf77f
Compare
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: 6
🤖 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/llm_eval/lm_eval_hf.py`:
- Line 267: Update the argparse definitions for calib_size,
auto_quantize_method, and auto_quantize_score_size to default to None, so non-HF
backends only reject values explicitly supplied by the user. In the HF
model_args injection path, apply the existing effective defaults (512,
"gradient", and 128) when these options remain unset, while preserving the
current validation for explicitly provided values.
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Line 645: Update the assignment in the searcher implementation from self.best
to a distinct field such as self.best_candidate when storing asdict(best).
Preserve BaseSearcher.self.best exclusively for its SearchStateDict so the
inherited search() contract remains unchanged.
- Line 645: Update the search completion flow around self.best = asdict(best) to
save the search checkpoint after assigning the selected candidate, ensuring the
final checkpoint’s "best" state is populated. Use the existing
save_search_checkpoint() mechanism and preserve the earlier checkpoint behavior.
- Line 297: The new "best" key added to the checkpoint structure causes strict
validation to fail when loading older checkpoints that lack this field. Either
exclude the "best" field from the strict state validation in
BaseSearcher.load_search_checkpoint, or add migration logic that populates the
"best" key in older checkpoints before strict validation occurs. Ensure backward
compatibility so existing checkpoints can resume without modification.
In
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml`:
- Around line 20-34: Remove the hardcoded --trust_remote_code flags from both
quantization and export commands in the example; remote-code execution must
default to False and require explicit caller opt-in, with the documented
security-exception approval obtained if it remains necessary.
- Around line 17-24: Update the environment section to define MLM_MODEL_CFG as
the Hugging Face repository ID and QUANT_CFG as MAMBA_MOE_FP8_CONSERVATIVE_CFG,
then replace the hardcoded model path and quantization config in the inline
quantize.py command with these variables.
🪄 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: d767c065-f193-45ba-8fca-60d5f62bb30f
📒 Files selected for processing (13)
examples/llm_eval/lm_eval_hf.pyexamples/megatron_bridge/prune_minitron.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/examples/megatron_bridge/test_prune_minitron.pytools/launcher/core.pytools/launcher/docs/configuration.mdtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yamltools/launcher/slurm_config.pytools/launcher/tests/test_core.pytools/launcher/tests/test_docker_execution.pytools/launcher/tests/test_examples_resolve.pytools/launcher/tests/test_yaml_formats.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tests/examples/megatron_bridge/test_prune_minitron.py
- examples/megatron_bridge/prune_minitron.py
- tools/launcher/tests/test_examples_resolve.py
- tools/launcher/slurm_config.py
- tools/launcher/tests/test_yaml_formats.py
- tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml
- tools/launcher/docs/configuration.md
End-to-end ModelOpt launcher pipelines for Megatron-Bridge on Nemotron-3-Nano-30B-A3B: - mbridge_prune.yaml: prune -> vLLM gen -> MMLU gate - mbridge_quantize.yaml: FP8 quantize -> unified-HF export -> vLLM gen -> MMLU gate (vLLM backend) Launcher (tools/launcher) additions so these run wrapper-free from YAML: - SandboxTask.inline: run a command directly from YAML (no .sh wrapper) - SandboxTask.reqs / reqs_file: pip-install deps in-container before the command - SlurmConfig.docker_user: run local Docker as a chosen user (e.g. root) - reject args together with inline examples/llm_eval/lm_eval_hf.py: - --accuracy_lower_bound: gate the run on a single task's acc (exit non-zero if below) - drop ModelOpt args for non-hf backends so --model vllm works on quantized checkpoints Docs (tools/launcher/docs/configuration.md) + unit tests updated. Verified end-to-end on Qwen3-0.6B (prune, FP8 quantize/export, vLLM gen, MMLU gates). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
d4cf77f to
b090a7c
Compare
|
Addressed the first round of review comments and resolved those threads (changes in Prune score gate —
Not addressed yet (will follow up): the test-coverage suggestions (failing-gate assertion, |
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.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
🤖 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 `@tests/examples/llm_eval/test_llm_eval.py`:
- Around line 39-41: Add a focused regression test alongside the existing LLM
evaluation accuracy test that configures an accuracy lower bound above the
produced score, invokes the relevant lm_eval_hf.py execution path, and asserts a
non-zero result. Keep the existing passing-bound test unchanged and ensure the
new test specifically exercises the failing accuracy-gate behavior.
🪄 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: d9a9085f-2967-45cb-975a-9cda12259235
📒 Files selected for processing (1)
tests/examples/llm_eval/test_llm_eval.py
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Most of the previous round is genuinely resolved (see below); what's left are owner-judgment calls plus one test gap, so I'd like a human sign-off rather than auto-approving.
Confirmed fixed: self.best = asdict(best) preserves the BaseSearcher.search() -> SearchStateDict contract (base returns self.best; reset_search()'s comment anticipates best in default_state_dict; mtn.search() returns state_dict(), so pruning_scores["best"]["score"] resolves); non-HF backends now raise on quantization/sparsity request args and is_hf covers all four HFLM aliases; results file selected by mtime; --accuracy_lower_bound validates --output_path/single-task before the expensive eval; TestGetDefaultEnv env isolation added; docs say args is rejected (not ignored); the reqs barrier now rm -fs the marker and waiting ranks fail via [ -f {marker} ]; script/inline exclusivity enforced in run_jobs.
Still warranting a look:
-
💬 Author decided to accept the strict-checkpoint-key break (
"best"added todefault_state_dict) because pruning is a one-shot step — reasonable, but flagging for human sign-off because the failure mode for anyone resuming a pre-upgrade--prune_intermediate_ckpt/modelopt_pruning_scoresdir is an opaqueassert state_dict.keys() == self.state_dict().keys(), "Keys in checkpoint don't match!"with no hint to delete the dir, and the mitigation is one line (overrideload_search_checkpoint(strict=False)inMCoreMinitronSearcher, or derive the score from the already-presentall_candidates_per_constraint). Owner's call, but worth a conscious yes/no. -
💬 Author declined new tests for the
run_jobsguard as "a one-line guard" — the guard itself, agreed, but thereqs_prefixconstruction (shlex quoting ofreqs/reqs_file, the rank-0 barrier shell string, and thereqs+script→run.Script(inline=...)branch) is the most intricate new code in this PR and is currently exercised by nothing.tests/test_core_extended.py::TestRunJobsExtendedalready patchescore.run.Scriptand asserts on the call kwargs, so a test asserting theinline=string for areqs/inlinetask is ~10 lines. Both failing-gate paths (--accuracy_lower_boundabove the score,--score_lower_boundabove the score) also remain untested; the prune one is fairly argued as too expensive, but thelm_evalone could be a cheap unit test of_enforce_accuracy_gateagainst a syntheticresults.json. -
💬 Design rationale (raised previously): the PR body describes
inline/reqsbut still doesn't say why the existingcommon/**/*.shwrapper convention (whichcommon/megatron_bridge/import/import.shalready demonstrates for this flow, and which could do the pip install + barrier in bash) or typed_target_tasks were rejected.tools/launcher/CLAUDE.mdstill lists only "Scripts go incommon/" under key conventions, so contributors now have three ways to express a task command with no guidance on which to pick. Please add a sentence of rationale to the PR body and a line toCLAUDE.md(docs/configuration.mdis updated nicely already). -
Multi-node
reqsbarrier (nit): the marker is created in the per-task working dir (/nemo_run/code), which is the packaged code dir mounted from the shared job dir — i.e. shared across nodes, not node-local. On a multi-node task, node B's non-zero local ranks can observe node A's marker and skip the wait before node B's local rank 0 has installed anything. Folding${SLURM_NODEID:-0}into the marker name (or using a node-local path) would make it truly per-node; alternatively documentreqsas single-node only. -
score_lower_bound=0.01intest_prune_minitron(nit): couples the GPU test's pass/fail to a tiny random-init model's MMLU score. Low risk (4-choice MMLU floor is ~25%), but it is a new flake surface for a test whose purpose is the prune path, not accuracy.
3 participants
What does this PR do?
Type of change: new example + small launcher / modelopt-example features (backward compatible)
Adds end-to-end ModelOpt launcher pipelines for the Megatron-Bridge flow on Nemotron-3-Nano-30B-A3B, the minimal launcher features to run them wrapper-free from YAML, and an in-step accuracy gate for Minitron pruning.
New launcher examples (
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/):mbridge_prune.yaml— Minitron prune with an in-step MMLU gate → vLLM sanity gen (2 tasks)mbridge_quantize.yaml— FP8 quantize → unified-HF export → MMLU gate on the vLLM backend, which doubles as the deploy sanity check (3 tasks). Matches the tutorial.Prune accuracy gate (reuses the search's own score — no separate eval step):
modelopt/torch/prune/plugins/mcore_minitron.py—MCoreMinitronSearchernow stores the exported bestCandidateSubnetunderstate_dict["best"](additive; sits beside the existingsorted_layerskey).examples/megatron_bridge/prune_minitron.py— new--score_lower_bound: readspruning_scores["best"].scoreand exits non-zero if the exported model is below the floor. Score-agnostic (any--prune_score_func); rejected with--prune_export_config(manual pruning has no score).Launcher (
tools/launcher) — run single-node Megatron-Bridge one-liners directly from YAML:SandboxTask.inline— a command in the YAML, nocommon/**/*.shwrapper (single-line; folded scalar)SandboxTask.reqs/reqs_file— pip-install deps in the container before the command (shell-safe; on Slurm the install is rank-0-guarded so multi-rank tasks don't race)SlurmConfig.docker_user— local-Docker user (e.g.root); ignored on Slurmget_default_envhonorsHF_HOME/TRITON_CACHE_DIRenv overrides, so a non-CI user can point caches at a writable path (the shared/cicd/hf-cacheis owned by the CI account)argstogether withinlineexamples/llm_eval/lm_eval_hf.py:--accuracy_lower_bound— gate on the single requested task'sacc(used by the quantize MMLU step; exits non-zero if below)hfbackends, so--model vllmworks on a deployable quantized checkpointUsage
Testing
inline,reqs/reqs_file,docker_user, theargs+inlineguard, and example-resolve; ruff / mypy / bandit clean.tests/examples/megatron_bridge/test_prune_minitron.pynow passes--score_lower_bound=0.01to exercise the gate path on the tiny models.[score_gate] mmlu_10pct = 0.5196 >= 0.45 PASS; vLLM gen coherent.Detected ModelOpt fp8 checkpoint) → MMLU on the vLLM backendacc = 0.7077 >= 0.60 PASS.nemo:26.06through the same flow.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
transformers<5, sombridge_prune.yamlruns onnemo:26.04(26.06 drops it); quantize/export run onnemo:26.06.docker_user: rootis set on all example tasks — local-Docker only (ignored on Slurm), needed so downstream tasks can read task_0's root-owned checkpoints and to read the image's root-only/opt/Megatron-Bridge.enforce_eager=Trueto vLLM — for a run-once eval this skips ~17 min of CUDA-graph capture /torch.compilewith no accuracy change.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Bug Fixes