Add Megatron-Bridge prune & quantize launcher pipelines - #2031
Add Megatron-Bridge prune & quantize launcher pipelines#2031kevalmorabia97 wants to merge 1 commit 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. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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.
| files = sorted(glob.glob(os.path.join(output_path, "**", "results*.json"), recursive=True)) | ||
| if not files: | ||
| raise FileNotFoundError(f"No results*.json under {output_path}") | ||
| with open(files[-1]) as f: # lm-eval timestamps names, so last is newest |
There was a problem hiding this comment.
[SUGGESTION] The "last is newest" assumption only holds within a single results subdirectory. The recursive glob descends into every <model_name>/ subdir under output_path, and sorted() orders by full path — directory name first, then filename. If output_path is reused across runs with different checkpoint names (e.g. the prune and quantize pipelines both point --output_path at /scratchspace/mmlu_results), files[-1] picks the lexicographically-last directory, not the most recent file by timestamp — so the gate can silently read a stale result and pass/fail on the wrong run.
Since the timestamp lives in the filename, sorting by mtime (or by basename) is more robust:
files = glob.glob(os.path.join(output_path, "**", "results*.json"), recursive=True)
if not files:
raise FileNotFoundError(f"No results*.json under {output_path}")
newest = max(files, key=os.path.getmtime)
with open(newest) as f:
scores = json.load(f)["results"].get(task, {})Not blocking — the shipped examples use an ephemeral per-run path, so it only bites if someone reuses output_path.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2031 +/- ##
==========================================
- Coverage 75.79% 75.15% -0.64%
==========================================
Files 518 519 +1
Lines 58658 60339 +1681
==========================================
+ Hits 44459 45347 +888
- Misses 14199 14992 +793
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:
|
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>
e2000f9 to
247cc1c
Compare
1 participant
What does this PR do?
Type of change: new example (+ small launcher/example-script features)
Adds end-to-end ModelOpt launcher pipelines for the Megatron-Bridge flow on Nemotron-3-Nano-30B-A3B, and the minimal launcher features to run them wrapper-free from YAML.
New launcher examples (
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/):mbridge_prune.yaml— Minitron prune → vLLM sanity gen → MMLU gatembridge_quantize.yaml— FP8 quantize → unified-HF export → vLLM sanity gen → MMLU gate (matches the tutorial)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)SlurmConfig.docker_user— local-Docker user (e.g.root); ignored on Slurmargstogether withinlineexamples/llm_eval/lm_eval_hf.py:--accuracy_lower_bound— gate the run on the single requested task'sacc(exit non-zero if below)hfbackends, so--model vllmworks on a deployable quantized checkpointUsage
Testing
inline,reqs/reqs_file,docker_user, and theargs+inlineguard; full launcher suite (142) + ruff pass.nemo:26.06vialaunch.py: prune → gen → MMLU gate; FP8 quantize → unified-HF export (hf_quant_config.json= FP8) → vLLM gen → MMLU gate through the vLLM backend (Detected ModelOpt fp8 checkpoint).Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Scoped to
tools/launcher+examples/llm_eval/lm_eval_hf.py.docker_user: rootis set on all tasks in the example YAMLs — 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.🤖 Generated with Claude Code