[2/2] Track every Megatron-Bridge script with MLflow - #2514
kevalmorabia97 wants to merge 2 commits into
Conversation
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. 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 change adds shared MLflow configuration and checkpoint-provenance utilities. Megatron-Bridge pruning, quantization, distillation, and export workflows use these utilities. Hugging Face PTQ and vLLM example utilities also adopt shared tracking helpers. Tests and documentation cover the updated workflows. ChangesMLflow tracking workflows
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~60 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant Distill as distill.py
participant RunContext as distill_run
participant BridgeLogger as Megatron-Bridge logger
participant MLflow
participant Checkpoint
Distill->>RunContext: Enter the distillation run context
RunContext->>MLflow: Open the run on the last process
Distill->>BridgeLogger: Supply logger_kwargs
BridgeLogger->>MLflow: Log training metrics and resolved config
Distill->>RunContext: Record provenance after the checkpoint marker changes
RunContext->>Checkpoint: Write the .experiment.json pointer
Suggested reviewers: Merge Risk: 🟡 Moderate · up to If another component leaves a different MLflow run active, a tracked script can write its final metrics and artifacts to that run and then close it, which corrupts experiment records. Separately, a tracking URI that contains only a token (or only a username) can end up stored in Megatron-Bridge run parameters and checkpoint configuration. Fix the run-ID check before merging; the credential exposure is narrow but simple to fix. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## kmorabia/mlflow-tool-core #2514 +/- ##
=============================================================
- Coverage 78.44% 76.46% -1.99%
=============================================================
Files 607 605 -2
Lines 68843 67176 -1667
=============================================================
- Hits 54006 51368 -2638
- Misses 14837 15808 +971
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:
|
|
/claude review |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Nicely scoped follow-up (~300 core-logic lines, well tested), but log_active_run_experiment_json uses the fluent mlflow.log_text, which starts a brand-new run on the already-closed path this PR explicitly supports.
Needs action:
- Log the artifact through
MlflowClient().log_text(run.info.run_id, ...)inmodelopt/torch/utils/mlflow.py—mlflow.log_textcalls_get_or_start_run(), so thelast_active_run()branch silently opens a second run. See inline comment. - Extend the
ForeignMlflowfake intests/unit/torch/utils/test_mlflow.pyto fail if a run is started, so the closed-run path is actually covered. - Tighten
record_checkpoint_provenanceinexamples/megatron_bridge/mlflow_utils.py:checkpoint_dir.is_dir()is true on any resumed run, so a run that dies before its first save stamps a previous run's checkpoint. See inline comment. - Replace the
if tool is DISTILLidentity check inadd_mlflow_argswith aToolfield (e.g.log_checkpoints_flag: bool = False), so the per-script differences all live in one place.
No action needed:
- Existing test updates are all signature/refactor follow-through (
_run_tags/_run_outputsfolded intoTool), coverage intact.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — 1 CRITICAL, 2 IMPORTANT, 2 SUGGESTION
Full review (scope: /claude review with no extra instructions). 9 files changed (+517/−86); reviewed all of modelopt/, examples/ and both test files, plus the README.md and CHANGELOG.rst diffs.
The refactor itself is clean: collapsing three copies of the wiring into one Tool record per script is the right call, and the decision not to open a competing run in distill.py — letting Megatron-Bridge’s LoggerConfig own it — is well reasoned and well documented, as is keeping mlflow_log_artifacts off by default. The findings are all in the new provenance-pointer path.
Most impactful
1. log_active_run_experiment_json uploads its artifact to a brand-new run when the real run has already closed (modelopt/torch/utils/mlflow.py:779) — CRITICAL. mlflow.log_text is the fluent API, so it resolves its target through _get_or_start_run(); with mlflow.active_run() returning None, that starts a run rather than reusing the one last_active_run() just returned. That is precisely the state the last_active_run() fallback was added for (Megatron-Bridge sys.exit()s from inside train()), so on the documented main path every QAD job leaves a spurious empty run and attaches experiment.json to it, while the on-disk pointer names the real run. Nothing raises, and the stubbed ForeignMlflow.log_text in the new test cannot see it. Fix: go through MlflowClient().log_text(run_id=..., ...).
2. record_checkpoint_provenance never clears a stale pointer (examples/megatron_bridge/mlflow_utils.py:237) — IMPORTANT. quantize.py and the export get this for free from track_run’s untracked branch; this path returns early instead. Since distill.py passes load=checkpoint_dir and a reused --output_dir is the normal Slurm-requeue flow, an untracked re-run — or a tracked one where mlflow is absent or the server unreachable — leaves the earlier run’s .experiment.json claiming the new weights, contradicting the invariant log_experiment_json documents.
3. QAD runs get MLflow’s auto-generated run name, not the documented UTC timestamp (examples/megatron_bridge/mlflow_utils.py:223) — IMPORTANT. The timestamp default lives in MlflowRunLogger.start, which distill.py never reaches, so mlflow_run_name=None is what Megatron-Bridge receives. Both the --mlflow_run_name help text registered on all three parsers and the README line this PR edits promise the UTC start time.
Two SUGGESTIONs are inline: the consumed Megatron checkpoint appears in no tag, so the PTQ→QAD→export chain is joinable only via the on-disk pointer rather than a server-side tag query; and record_checkpoint_provenance’s “no checkpoint to point at” claim does not hold for a resumed --output_dir.
Verified as correct
DISTILL.checkpointmatchesdistill.py’scheckpoint_dir = os.path.join(args.output_dir, "checkpoints").dist.is_last_process()exists and is global-rank based, matching where Megatron-Bridge opens the run; the process group is still alive where thefinallyruns.- The
finallyplacement is right —sys.exit()raisesSystemExit, which propagates through it. log_active_run_experiment_json’s JSON keys matchMlflowRunLogger.run_infoexactly, so consumers see one format.resolved_recipe_texts(getattr(args, "recipe", None))correctly tolerates the export’s missing--recipe;_NON_PARAM_ARGSand theprint_args(masked_args(args))/checkpoint_exportedordering are consistent across both single-pass scripts.logger_kwargsreturning{}when untracked does keep an older Megatron-Bridge working, as the PR body claims.
Minor, not raised inline
distill.py --hf_export_path writes a second, deployable HF checkpoint (export_llm_to_hf / save_vlm_to_hf) that gets neither a pointer nor stale-pointer cleanup. Arguably out of scope given export_quantized_megatron_to_hf.py covers the primary export path — noting it in case it was an oversight.
Risk
Moderate. The behavioural blast radius is examples-only, and the single library change is additive (log_active_run_experiment_json is new; nothing else calls it) with no modelopt_state, mode-registration or config-schema surface touched, so there is no checkpoint or public-API compatibility risk. Finding 1 is worth fixing before merge because it writes to the tracking server on the default path and fails silently.
🤖 Generated with Claude Code
|
On the minor point raised but not filed inline — It is not a two-line addition. That branch does: is_rank_0 = dist.rank() == 0
dist.cleanup() # process group destroyed; export_ckpt makes its own
if is_rank_0:
export_llm_to_hf(...)So the rank that writes the checkpoint is rank 0, while Megatron-Bridge owns the MLflow run on the last rank — and after Happy to do it as a follow-up if you'd rather it not wait. For the record, everything else from both review rounds is now in 🤖 Generated with Claude Code |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Re-review: the spurious-run bug is gone and the other five threads are properly closed, but the chosen fix uses a log_text(..., run_id=) kwarg that does not exist at the declared mlflow-skinny>=2.9 floor.
Needs action:
- 💬 Author replied "verified against the installed client" — the floor in
pyproject.tomlismlflow-skinny>=2.9, where fluentlog_text(text, artifact_file)takes norun_id. UseMlflowClient().log_text(run_id, text, artifact_file)(works on every version) or raise the floor. See inline onmodelopt/torch/utils/mlflow.py:808. - Add a regression that a
TypeErrorfromlog_textis visible, or make theexcept Exceptioninlog_active_run_experiment_jsonnarrower — today an unsupported signature degrades to a warning with no pointer on disk and a stale pointer left in place. - Confirm the new cross-stage join is robust to relative paths:
checkpoint_run_tagsresolvescheckpoint_pathbut leavessource_checkpoint_pathas given, so a relative--student_megatron_pathbreaks the PTQ→QAD→export tag querytest_the_three_stages_chain_on_the_serverpins.
No action needed:
- ✔️ Resolved since the last review: the closed-run double-open, the
is_dir()provenance guard (nowcheckpoint_marker), stale-pointer clearing, thetool is DISTILLcheck, the run-name default, and the source-checkpoint tag.
|
/claude review |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Approving: all three open items from the last round are fixed in the diff, and nothing new surfaced.
No action needed:
- ✔️ Resolved since the last review:
log_active_run_experiment_jsonnow usesMlflowClient().log_text(info["run_id"], text, ...)(positional, valid at themlflow-skinny>=2.9floor); the pointer is written before the upload withtest_a_failed_upload_still_leaves_the_pointer_on_diskpinning that aTypeErrorno longer costs the on-disk record; andcheckpoint_run_tagsnow resolvessource_checkpoint_pathon the same terms ascheckpoint_path, with relative/absolute and Hub-id tests. - Existing test edits are signature follow-through from folding
_run_tags/_run_outputsintoTool; coverage is intact. - Minor, not blocking: if
mlflow.get_experimentraises,log_active_run_experiment_jsonreturns without dropping a stale pointer, unlike therun is Nonepath — narrow edge, worth a follow-up at most.
Complex PR: 1 existing test file modified or removed. Looping in a human for approval.
There was a problem hiding this comment.
Claude review — 0 CRITICAL, 2 IMPORTANT, 3 SUGGESTION
Full review (scope: /claude review with no extra instructions). 9 files changed (+837/−102); reviewed all of modelopt/ and examples/, both test files, and the README.md/CHANGELOG.rst diffs — nothing skipped.
Both items the previous round flagged are genuinely fixed: MlflowClient().log_text(run_id, text, artifact_file) is positional and works at the mlflow-skinny>=2.9 floor, and source_checkpoint_path is now resolved on the same terms as the upstream checkpoint_path while leaving a Hub org/name alone — test_a_source_path_joins_whatever_the_caller_typed pins exactly the case that was broken. ForeignMlflow.start_run raising is a good way to make the closed-run path assert itself.
The two IMPORTANT findings are both new.
Most impactful
1. logger_kwargs hands Megatron-Bridge the unmasked tracking URI (examples/megatron_bridge/mlflow_utils.py:236) — IMPORTANT. args.mlflow keeps any user:token@ the user passed, and that is the one thing every other path in this module redacts: _redact_argv for command.txt, the _SECRET_NAME/_redact params filter at modelopt/torch/utils/mlflow.py:549, run_info["tracking_uri"], print_args(masked_args(args)) in both wired scripts. By this PR's own README wording LoggerConfig records "the full resolved config as params", and the resolved config is also serialised to checkpoints/iter_*/run_config.yaml (asserted at tests/examples/megatron_bridge/test_distill.py:274) — so on the QAD path a URI-embedded token becomes a searchable MLflow param and ships inside the distilled checkpoint. The other two scripts never hand the URI to anything that serialises it, so this is specific to the new path. Either strip the credentials into MLFLOW_TRACKING_USERNAME/MLFLOW_TRACKING_PASSWORD before passing the URI, or document in the DISTILL help text and the README that distill.py needs them from the environment.
2. log_active_run_experiment_json leaves a stale pointer when it cannot read the run (modelopt/torch/utils/mlflow.py:806) — IMPORTANT. The run is None branch four lines up calls drop_experiment_json; this one returns after a warning. Both this docstring and log_experiment_json's promise the pointer beside a saved checkpoint is "this run's or absent -- never a previous run's". The reachable case is the ordinary requeue flow: distill.py points load at the same <output_dir>/checkpoints, so a resumed run starts with the previous run's .experiment.json there, and record_checkpoint_provenance has already proven the marker moved. A server that goes away between the last save and this call therefore leaves a pointer misnaming the author of the new weights — a wrong answer rather than no answer. test_recording_the_active_run_never_fails_the_job covers this path but starts from an empty directory, so it cannot see it.
Three SUGGESTIONs are inline: default_run_name() evaluated per-rank so run_config.yaml can name a run name that does not exist; --mlflow_log_checkpoints registered in the underscored spelling only, against the dual-spelling convention the sibling flags document; and two _run_inputs monkeypatch stubs left at the old one-argument signature.
Verified as correct
checkpoint_markerreadslatest_checkpointed_iteration.txt, which Megatron-Bridge really does write after distillation —tests/examples/megatron_bridge/test_qad.py:110asserts it on the distilled checkpoint, so the marker-moved gate is not silently always-false.- The marker gate itself is sound: read before
distill(), compared after, so a resumed run that died before its own first save leaves the inherited checkpoint and its pointer alone.test_no_pointer_when_this_run_saved_nothingwalks all three states. dist.is_last_process()is global-rank based (rank() == size() - 1) and the default process group is still alive where thefinallyruns — confirmed bydistill.py's own "Save rank before destroying process group" comment in the--hf_export_pathbranch, which is the first place the group goes away. So exactly one rank writes the pointer; there is no all-ranks-look-like-rank-0 fanout where a non-owning rank woulddrop_experiment_jsonover the owner's write.finallyplacement is right: Megatron-Bridge'ssys.exit()from insidetrain()raisesSystemExit, which propagates through it, and thenot args.validate_onlyguard matchesskip_train=args.validate_only.Tool.checkpointforDISTILLmatchesdistill.py:384'sos.path.join(args.output_dir, "checkpoints"); every attribute the threeToolrecords read (student_hf_path,student_megatron_path,output_dir,hf_model_name_or_path,megatron_path,export_unified_hf_path) is a flag its script registers beforeresolve_mlflow_argsruns.resolve_mlflow_argsassignsargs.mlflowfrom$MLFLOW_TRACKING_URI, sologger_kwargs'sif not args.mlflowearly return does not quietly drop env-configured tracking fordistill.py— and it guaranteesmlflow_experimentis non-Nonewhenever the dict is non-empty._experiment_jsonunifies the two writers on one schema, andrun_urlno longer double-slashes;MlflowRunLogger.run_urlreading throughrun_infois not circular now thatrun_infobuilds the URL itself, and still returns""before the run opens.osis imported at module top, socheckpoint_run_tags' newos.path.existsguard is fine;from mlflow.tracking import MlflowClientand positionallog_textboth predate the 2.9 floor.args.checkpoint_exported = Falseis set afterprint_args, and_NON_PARAM_ARGSexcludes it plusmlflow_log_checkpoints, so neither reaches the params.
Risk
Moderate, and confined. The single library change is additive — log_active_run_experiment_json and default_run_name are new, and the checkpoint_run_tags(..., model=) parameter is keyword-with-default, so the existing hf_ptq caller is unaffected. No modelopt_state, mode registration, config schema or export-format surface is touched, so there is no checkpoint or public-API compatibility exposure. Finding 1 is the one I would fix before merge, because it writes a secret to durable storage on a documented invocation form and nothing surfaces it.
🤖 Generated with Claude Code
|
/claude review |
There was a problem hiding this comment.
Claude review — 1 IMPORTANT, 2 SUGGESTION
Scope: the trigger was a bare /claude review, so this is a full review. All 9 changed files opened (small PR, +946/-104): modelopt/torch/utils/mlflow.py, the four examples/megatron_bridge/ files, both test files, README, CHANGELOG. This is the first Claude pass on the PR — no prior findings to reconcile. Reviewed against the stacked base kmorabia/mbridge-quantize-mlflow (two-dot diff, since the checkout is shallow).
Findings
| No. | Severity | Where | What |
|---|---|---|---|
| 1 | IMPORTANT Compatibility | examples/megatron_bridge/mlflow_utils.py:280-287 |
A tracked QAD run that cannot find Megatron-Bridge’s run writes no pointer, prints nothing, and deletes the pointer it inherited |
| 2 | SUGGESTION | modelopt/torch/utils/mlflow.py:826-836 |
A transient get_experiment() failure drops a pointer whose run_id is already in memory |
| 3 | SUGGESTION | examples/megatron_bridge/README.md:129 |
"All three scripts that write a checkpoint" — four do; distill.py --hf_export_path and export_distilled_megatron_to_hf.py are uncovered |
Most impactful
Finding 1 is the one worth acting on. log_active_run_experiment_json() collapses "this job is untracked" and "tracking was requested but no run is visible on this rank" into the same quiet drop_experiment_json() + return. The second case is reachable — a Megatron-Bridge version that opens the run somewhere other than rank == world_size - 1, an mlflow client absent from the image, MLflow logging disabled inside LoggerConfig — and in it the user passed --mlflow, got no error, got no pointer, and lost the one a previous run left. For a feature whose entire purpose is "the checkpoint names the run that produced it", that failure is invisible in the job log. record_checkpoint_provenance() already has args.mlflow in hand, so telling the two cases apart is a returned bool plus a one-line warning.
The design leans on "Megatron-Bridge opens the run on the last rank". That is consistent with Megatron’s own is_last_rank() and with dist.is_last_process(), and I could not verify it from this repo (megatron.bridge is not installed in this environment), so it is not a separate finding — but it is the assumption that finding 1 would make safe to be wrong about.
What I traced and found correct
- The three-stage join keys line up.
quantize.checkpoint_path(resolved--export_megatron_path) equalsdistill.source_checkpoint_path(resolved--student_megatron_path);distill.checkpoint_path(resolved<output_dir>/checkpoints) equalsexport.source_checkpoint_path(resolved--megatron_path, which README:328 points at<output>/checkpoints). The newos.path.exists()guard incheckpoint_run_tagscorrectly leaves a Hub id likeorg/nameunresolved, andmodel=keeps themodeltag naming the model rather than the checkpoint directory for the two stages whose source is a checkpoint. dist.broadcast(default_run_name())is safe where it is called.distill.py:680runsdist.setup()beforemain(), so the default process group is initialized andtorch.cuda.set_device(local_rank())has run — the.cuda()insidebroadcastwill not collide ranks on device 0. All ranks parse identical argv, soargs.mlflow/args.mlflow_run_nameare identical and every rank takes the same branch of theor; no rank-divergent collective.- The
finallyarounddistill(config)is right, andSystemExitsurvives it.dist.abort()re-raisesSystemExitrather than swallowing it, so the--exit_intervalpath still exits cleanly after the pointer is written.record_checkpoint_provenancedoes no collective, so the last rank doing MLflow HTTP while peers unwind cannot deadlock. - The
saved_beforemarker guard holds. Fresh save, resume-then-save, and resume-then-crash all behave as documented;checkpoint_markerreads<output_dir>/checkpoints/latest_checkpointed_iteration.txt, which matchescheckpoint_diratdistill.py:384andsave=at:613. Tool.outputs = field(default=lambda args: {})does not bind as a method. The generated__init__always assigns the instance attribute (viaobject.__setattr__underfrozen=True), so the instance dict shadows the class-level function;test_the_export_tags_point_at_the_deployable_checkpointexercises this through_describe.- The refactor is behaviour-preserving.
_experiment_jsonadds a.rstrip("/")thatrun_urlpreviously lacked (strictly better — no doubled slash before#/),run_nameresolution is unchanged, and no existing test needed editing (0 deletions intest_mlflow.py). split_tracking_credentialsis the right call for the distill path. Megatron-Bridge logs its resolved config as params and serialises it intorun_config.yamlinside the checkpoint, so masking would not work and leaving the credential would make it durable in two places;os.environ.setdefaultcorrectly lets a deliberately-exported variable win.- Plugin laziness respected:
mlflowandmlflow.trackingare both imported inside the function, andmlflow_utils.pystill imports no Megatron.
No mode registration, config schema, modelopt_state, or public modelopt/torch/*/__init__.py surface is touched, so nothing here affects checkpoint restore. The only public-API change is additive — three new __all__ entries — plus a new model= parameter on checkpoint_run_tags with a backward-compatible default.
Risk: low. Confined to example scripts and one opt-in tracking utility; the optimization and export paths are untouched, and every failure mode in the new library helper is caught and warned rather than raised. Finding 1 is an observability hole in a provenance feature, not a correctness bug in anything that produces weights.
🤖 Generated with Claude Code
There was a problem hiding this comment.
Claude review — CRITICAL: 1 · IMPORTANT: 0 · SUGGESTION: 3
Scope: the trigger was a bare /claude review, so this is a full review. 12 files changed (+1400/−53); I opened all 8 source files (modelopt/torch/utils/mlflow.py, the six examples/megatron_bridge/ scripts, the README) plus surrounding context in distributed.py, validate_tracking_uri/resolve_tracking_uri/_closing_run/tracked_run, and the logger_kwargs and credential tests in tests/examples/megatron_bridge/test_mlflow_utils.py. Two-dot diff against origin/kmorabia/mlflow-tool-core, since the checkout is shallow.
Prior rounds reconciled. The findings from the last two passes are genuinely fixed:
- The stray-run bug is gone:
log_textnow gates on_reattach(), solog_experiment_json(None)intracked_run.close()can no longer open a second run ahead of the re-attach. _reattachcomparesactive.info.run_idagainstself._run.info.run_idrather than assuming any active run is this one — the guard the last round asked for, and it is what makes theKILLEDcarry-over infinish()land on the right run.split_tracking_credentialspercent-decodes both halves, sohttps://alice:tok%2Fen@server/authenticates the same fromdistill.pyas fromquantize.py._checkpoint_rootnormalizes aniter_*directory up to the checkpoints root, so thedistill→distill_exporthop joins on the same key the distillation tagged.args.checkpoint_exportedis no longer dead state inexport_distilled_megatron_to_hf.py; thesettles_pointer=Falsetools getexported=lambda: Falseand never read the attribute.
The blocking one
[CRITICAL] logger_kwargs takes its "hand Megatron-Bridge nothing" decision only on rank 0 (examples/megatron_bridge/mlflow_utils.py:292). if recordable is None and dist.is_master(): print(...); return {} — the rank guard belongs to the print, but it also gates the return, so every non-master rank falls through and receives mlflow_tracking_uri=None with mlflow_experiment set. That is the rank that matters: distill_run opens the run on dist.is_last_process() precisely because that is where Megatron-Bridge looks, and for any world_size > 1 that rank is not master. It therefore enters Megatron-Bridge's MLflow path (which this module documents as keyed on mlflow_experiment) with no URI, calls set_experiment unguarded from inside the training loop against ./mlruns, and repoints the process-global tracking URI out from under MlflowRunLogger — the same class of failure the sibling logger.enabled broadcast in distill_run was added to prevent. It also makes LoggerConfig, and so run_config.yaml, differ by rank. test_an_unmovable_credential_is_not_recorded_at_all passes only because a single-process test has dist.is_master() == True. Reachable with any half-credential URI, e.g. a bearer token pasted as userinfo (--mlflow https://token@server/).
Non-blocking
finish()leaves the run with no terminal status on thenot ourspath — when a different run holds the fluent slot our run staysRUNNINGforever, and a_reattach()exception now skipsend_runwhere the old code fell through to it.MlflowClient().set_terminated(run_id, status)closes ours without needing the active slot.log_active_run_experiment_jsonreturnsFalsefor two different reasons — theOSErrorbranch reportswritten=Falseafter finding the run and uploading the artifact, and both callers then print "no run was found", pointing the reader at rank placement instead of the filesystem.- A pruning run given both output paths writes two checkpoints and one pointer —
record_exported_checkpoint, which this PR introduces, would cover--output_hf_paththe way it covers the per-iteration exports.
One smaller note not worth its own thread: split_tracking_credentials uses os.environ.setdefault, so a user who has already exported MLFLOW_TRACKING_USERNAME/PASSWORD and passes a differently-credentialed --mlflow gets our logger authenticating off the URI userinfo (where requests wins) and Megatron-Bridge authenticating off the pre-existing variables — a 401 raised from inside the training loop. The docstring states the precedence but not that the mismatch is a hazard.
What I traced and found correct
- The
_reattachstate machine. On the preempted pathclose()→log_experiment_json(None)→log_text→_reattachcaptures_closed_as = "KILLED"and re-opens;finish()'s second_reattachthen sees its own run, returnsTrue, and the("FAILED", "KILLED")filter restoresKILLED— a non-terminalRUNNINGnever reachesend_run, as the comment claims. - No stray run from the provenance writer.
log_active_run_experiment_jsongoes throughMlflowClient().log_text(run_id, ...)and otherwise touches onlyget_experiment/get_tracking_uri, neither of which resolves through_get_or_start_run(). - The pointer invariant holds on every path I walked, including the requeue flow:
checkpoint_markerreadslatest_checkpointed_iteration.txtrather than trusting the directory, so a run that dies before its own first save leaves the resumed checkpoint's pointer alone, andrecord_checkpoint_provenanceis called unconditionally so an untracked re-run into a reused--output_dirclears what it inherited.DISTILL_EXPORT/DISTILLcorrectly opt out oftracked_run's own pointer withsettles_pointer=False, and the untracked-cleanup branch atmlflow.py:1072is unreachable for them precisely becauserecord_exported_checkpoint/record_checkpoint_provenanceown it instead. - Collectives are entered by all ranks under identical predicates. Both
dist.broadcastcalls indistill_runsit underif args.mlflow, which every rank parsed identically;src=dist.size() - 1matches the ranktracked_runwas told is main;size() == 1short-circuits to the local object, which isNonewhen the logger disabled itself, so a single-process run degrades correctly. - Rank capture around
dist.cleanup().export_distilled_megatron_to_hf.py:300savesis_rank_0before the teardown andrecord_exported_checkpointrefuses on anything else, so the post-cleanup "every rank is rank 0" state cannot produce N writers; it is also the same rankmlflow_runenabled the logger on, so the fluentactive_run()it reads is really there. split_tracking_credentials' fail-closed scheme check is consistent withvalidate_tracking_uri(mlflow.py:166), which rejects the same non-http(s)URIs — so theNonereturn cannot be reached by a legitimatefile:/sqlite:backend, andrpartition/partitionhandle an@or:inside the password correctly.- Exit paths.
_closing_runmapsSystemExit(0)toFINISHEDand non-zero toFAILED, sotrain()'s--exit_intervalexit andprune_minitron.py's accuracy gate are both recorded as the PR describes, andrecord_checkpoint_provenancein thefinallystill runs on the exit path. - Every
Toolcallable resolves against an argument its script defines, andprune_score/checkpoint_exported/mlflow_log_checkpointsare all in_BOOKKEEPING, so none of them leaks into the params.
Risk
Moderate, and confined to examples/megatron_bridge/ plus additive surface in modelopt/torch/utils/mlflow.py. The library changes are behind enabled/self._run is None guards and everything after the run opens is best-effort, so nothing here can fail a job that would otherwise have worked — with the one exception of the CRITICAL, which hands a training loop a tracking configuration it will trip over. Fix that and the rest is documentation-grade.
🤖 Generated with Claude Code
4a484f8 to
82c1636
Compare
|
/claude review |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5-5) — DM the bot to share feedback.
Nudge: the two open code items from the last round are fixed, but the PR is still 566 lines of core logic, over the 500-line budget, and none of the size exceptions apply.
Needs action:
- ✂️ Split this PR into stacked
[x/N]PRs, merged in order after #2544, each passing CI alone and linking the others.[2/3] Add shared-run MLflow helpers:split_tracking_credentials,log_active_run_experiment_json, and_reattach/_closed_as/set_terminatedinmodelopt/torch/utils/mlflow.py, plus theirtest_mlflow.pycases and theFakeMlflowresume support (about 151 lines).[3/3] Track every Megatron-Bridge script with MLflow:examples/megatron_bridge/*plustest_mlflow_utils.py(about 415 lines). If the owner accepts the overage instead, post that waiver on the PR.
No action needed:
- ✔️ Resolved since the last review: the README now says the uploaded
distill.pylog comes from the last rank and does not include theprint_rank_0output.mlflow_runnow clearsargs.mlflowwhen the logger disables itself, and a test pins thatrecord_exported_checkpointthen stays quiet. Also fixed: every rank now declines the half-credential config, the owning rank is captured beforedist.cleanup(), a foreign-runfinish()closes this run by id, andprune_scoreis stashed right after the search. - The edits to existing tests are justified. The naming test is now parametrised over all five tools, pointer gating covers the prune and export scripts, and
run_tags/add_mlflow_argsare now reached throughmlflow_utils. No assertion was loosened. - The design is settled. This PR extends the existing
Tool/tracked_runcore from #2544 rather than adding a second system.
| _name(args.student_megatron_path) if args.student_megatron_path else "bf16" | ||
| ), | ||
| model=lambda args: args.student_hf_path, | ||
| checkpoint=lambda args: str(Path(args.output_dir) / "checkpoints"), |
There was a problem hiding this comment.
[SUGGESTION] DISTILL.checkpoint is unconditional, so a --validate_only run still tags checkpoint_path=<output_dir>/checkpoints even though it writes nothing.
run_tags states the invariant explicitly — "Omitted rather than empty when the run writes no checkpoint: a search for runs that produced one should not match it." With --validate_only against a reused --output_dir, a tag query for "which run produced this checkpoint" now returns the validation run alongside the training run that actually wrote it. The on-disk pointer is correct (distill.py guards record_checkpoint_provenance with if not args.validate_only), so this is server-side only.
Folding the condition into the Tool also removes the need for that guard, and makes checkpoint_marker return None for the same reason:
# None for --validate_only: the run writes no checkpoint, so nothing to tag or point at.
checkpoint=lambda args: (
None if args.validate_only else str(Path(args.output_dir) / "checkpoints")
),
| def resolve_mlflow_args( | ||
| args: argparse.Namespace, parser: argparse.ArgumentParser, tool: Tool | ||
| ) -> None: | ||
| """Settle where tracking is configured from, and name the experiment.""" | ||
| _resolve_mlflow_args(args, parser, tool) | ||
|
|
There was a problem hiding this comment.
[SUGGESTION] This wrapper adds nothing — it forwards all three arguments to _resolve_mlflow_args unchanged. add_mlflow_args below genuinely needs to wrap (it appends --mlflow_log_checkpoints), but this one only exists so the scripts can import both names from mlflow_utils. A plain alias says that, and keeps _resolve_mlflow_args's docstring as the single source of truth instead of a second, shorter paraphrase that can drift:
| def resolve_mlflow_args( | |
| args: argparse.Namespace, parser: argparse.ArgumentParser, tool: Tool | |
| ) -> None: | |
| """Settle where tracking is configured from, and name the experiment.""" | |
| _resolve_mlflow_args(args, parser, tool) | |
| # Re-exported so a script imports the whole tracking CLI from one module; add_mlflow_args | |
| # below wraps because it has something to add. | |
| resolve_mlflow_args = _resolve_mlflow_args |
| # The Tool names the Megatron checkpoint, which is what a chain joins on, so the HF | ||
| # one is pointed at the run here -- both are written when both flags are given. | ||
| record_exported_checkpoint(args, args.output_hf_path, dist.is_master()) |
There was a problem hiding this comment.
[SUGGESTION] The comment holds only when --output_megatron_path is also given. PRUNE.checkpoint is output_megatron_path or output_hf_path, so with --output_hf_path alone the Tool names the HF checkpoint, tracked_run writes .experiment.json there at close, and this call writes the same JSON to the same directory — a harmless duplicate, but the comment's premise ("the Tool names the Megatron checkpoint") is false in that configuration.
Worth narrowing both the comment and the call, e.g.:
# The Tool points at the Megatron checkpoint when there is one, since that is what a
# chain joins on, so the HF one needs its own pointer only then.
if args.output_megatron_path:
record_exported_checkpoint(args, args.output_hf_path, dist.is_master())
|
Claude review — 0 CRITICAL, 0 IMPORTANT, 3 SUGGESTION Full review (scope: Prior round's findings are resolved. The CRITICAL I raised last time — What I traced this round
Findings — all three are non-blocking polish in the provenance-metadata layer, none affects a checkpoint's own pointer:
Risk: low. Everything lands in 🤖 Generated with Claude Code |
[1/2] Merges with #2514 after this. Three example scripts had each reimplemented the same tracking wiring: the flags, the $USER/<tool>/<model>-<variant> experiment convention, the params and tags and artifacts a run uploads, and the open/close dance with its status. The copies had already drifted -- only hf_ptq wrote a provenance pointer, only vllm_serve republished the resolved URI -- and every new script meant another copy. What a script records is now one declarative Tool record, declared in the script itself beside the flags it reads: which arguments name its model, its checkpoint and what it consumed, what it uploads, what it measures. tracked_run takes that record and runs the whole thing, so a script adds tracking in three lines. examples/hf_ptq, examples/vllm_serve and examples/megatron_bridge/quantize.py move onto it with no change in behaviour, beyond hf_ptq's source_checkpoint_path tag now resolving to an absolute path so it can join the run that produced its input. The four test suites shared four copies of a stand-in for the mlflow module, which had drifted far enough that one made log_artifact a no-op -- a test asserting on an upload asserted nothing. They now share one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
82c1636 to
783ee5b
Compare
[1/2] Merges with #2514 after this. Three example scripts had each reimplemented the same tracking wiring: the flags, the $USER/<tool>/<model>-<variant> experiment convention, the params and tags and artifacts a run uploads, and the open/close dance with its status. The copies had already drifted -- only hf_ptq wrote a provenance pointer, only vllm_serve republished the resolved URI -- and every new script meant another copy. What a script records is now one declarative Tool record, declared in the script itself beside the flags it reads: which arguments name its model, its checkpoint and what it consumed, what it uploads, what it measures. tracked_run takes that record and runs the whole thing, so a script adds tracking in three lines. examples/hf_ptq, examples/vllm_serve and examples/megatron_bridge/quantize.py move onto it with no change in behaviour, beyond hf_ptq's source_checkpoint_path tag now resolving to an absolute path so it can join the run that produced its input. The four test suites shared four copies of a stand-in for the mlflow module, which had drifted far enough that one made log_artifact a no-op -- a test asserting on an upload asserted nothing. They now share one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
783ee5b to
87ac734
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5-5) — DM the bot to share feedback.
Nudge: the code is ready, but this PR is still 531 lines of core logic, over the 500-line budget, and none of the size exceptions apply.
Needs action:
- ✂️ Split this PR into stacked
[x/N]PRs, merged in order after #2544, each passing CI on its own and linking the others. You said you won't split because it's only slightly over; that's the owner's call, but "slightly over" isn't an exception, so split or have the owner post a waiver on the PR.[2/3] Add shared-run MLflow helpers:split_tracking_credentials,log_active_run_experiment_json,_reattach/_closed_as/set_terminatedinmodelopt/torch/utils/mlflow.py, with theirtest_mlflow.pycases andFakeMlflowresume support (about 151 lines).[3/3] Track every Megatron-Bridge script with MLflow:examples/megatron_bridge/*plustest_mlflow_utils.py(about 380 lines).
No action needed:
- ✔️ Resolved since the last review:
DISTILL.checkpointreturnsNonefor--validate_only,resolve_mlflow_argsis a plain alias, andprune_minitron.pywrites the extra HF pointer only when a Megatron checkpoint also exists. - The edits to existing tests are justified: the naming test is parametrised over all five tools, and pointer gating now covers prune and export. No assertion was loosened.
- The design is settled: this extends the
Tool/tracked_runcore from #2544 rather than adding a second system.
[1/2] Merges with #2514 after this. Three example scripts had each reimplemented the same tracking wiring: the flags, the $USER/<tool>/<model>-<variant> experiment convention, the params and tags and artifacts a run uploads, and the open/close dance with its status. The copies had already drifted -- only hf_ptq wrote a provenance pointer, only vllm_serve republished the resolved URI -- and every new script meant another copy. What a script records is now one declarative Tool record, declared in the script itself beside the flags it reads: which arguments name its model, its checkpoint and what it consumed, what it uploads, what it measures. tracked_run takes that record and runs the whole thing, so a script adds tracking in three lines. examples/hf_ptq, examples/vllm_serve and examples/megatron_bridge/quantize.py move onto it with no change in behaviour, beyond hf_ptq's source_checkpoint_path tag now resolving to an absolute path so it can join the run that produced its input. The four test suites shared four copies of a stand-in for the mlflow module, which had drifted far enough that one made log_artifact a no-op -- a test asserting on an upload asserted nothing. They now share one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
87ac734 to
6f76b6a
Compare
|
Size waiver (repo owner). This work is intentionally split as 2 PRs, not 3: #2544 is the shared tracking core, this PR is the This PR is ~563 lines of source against the ~500-line budget. The bulk is five example scripts each declaring their own |
[1/2] Merges with #2514 after this. Three example scripts had each reimplemented the same tracking wiring: the flags, the $USER/<tool>/<model>-<variant> experiment convention, the params and tags and artifacts a run uploads, and the open/close dance with its status. The copies had already drifted -- only hf_ptq wrote a provenance pointer, only vllm_serve republished the resolved URI -- and every new script meant another copy. What a script records is now one declarative Tool record, declared in the script itself beside the flags it reads: which arguments name its model, its checkpoint and what it consumed, what it uploads, what it measures. tracked_run takes that record and runs the whole thing, so a script adds tracking in three lines. examples/hf_ptq, examples/vllm_serve and examples/megatron_bridge/quantize.py move onto it with no change in behaviour, beyond hf_ptq's source_checkpoint_path tag now resolving to an absolute path so it can join the run that produced its input. The four test suites shared four copies of a stand-in for the mlflow module, which had drifted far enough that one made log_artifact a no-op -- a test asserting on an upload asserted nothing. They now share one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
[2/2] Merges after #2544. of five scripts in that directory that write a checkpoint; the other four recorded nothing, so the provenance chain stopped at the PTQ checkpoint and a deployed model could not be traced back to the run that produced it. prune_minitron.py, distill.py, export_quantized_megatron_to_hf.py and export_distilled_megatron_to_hf.py now take the same flags, each declaring what it records as a Tool beside its own flags. Each writes .experiment.json into the checkpoint it produced and tags what it consumed, so prune -> quantize -> distill -> export is walkable both from disk and by tag query. distill.py opens the run rather than being wrapped by one: Megatron-Bridge's LoggerConfig records per-iteration metrics and the full resolved config, which a wrapper cannot see, and it joins mlflow.active_run() when there is one. So the run is opened on the rank Megatron-Bridge looks at -- the last one -- and the two share it. Its early exit is handled explicitly: train() leaves through sys.exit(0) on --exit_interval, which a blanket handler would record as FAILED. The shared wiring these four need lands here with them: distill_run, logger_kwargs, checkpoint_marker and the two record_* helpers in mlflow_utils, the add_mlflow_args wrapper that adds --mlflow_log_checkpoints, the exported callback's settles_pointer branch and the args.mlflow reset -- each with the caller that earns it -- plus split_tracking_credentials, log_active_run_experiment_json and MlflowRunLogger._reattach in the library, the last because a co-owner can end the run first, as Megatron-Bridge does on SIGTERM. Two of Megatron-Bridge's defaults are deliberately not inherited: checkpoint artifact upload stays off unless --mlflow_log_checkpoints, and an untracked run passes no mlflow_* fields at all, since they landed in Megatron-Bridge 0.6 and sending them unconditionally would break an untracked run on an older one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
prune_minitron.py puts --output_megatron_path and --output_hf_path in a required mutually exclusive group, so exactly one is ever set. The extra record_exported_checkpoint call added last round sat inside the HF save branch behind `if args.output_megatron_path:` -- a condition that can never hold there -- and PRUNE.checkpoint's `or None` tail is dead for the same reason. The review that asked for it said passing both was a normal invocation, and I took that without checking the parser. The test could not catch it because _parse_prune builds its own parser from _SCRIPT_ARGS, so it asserted on two combinations argparse rejects outright: both flags, and neither. It now covers the two legal invocations and checks the pointer lands in whichever output was saved. Found by running the scripts for real against a live MLflow server, where the prune step exited 2 before the run was ever opened. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
7142303 to
442547f
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5-5) — DM the bot to share feedback.
Nudge: the code is ready, but this PR is 552 lines of core logic, over the 500-line budget. The owner's waiver is not one of the size exceptions, so a human makes the merge call.
Needs action:
- ✂️ Split this PR into stacked
[x/N]PRs, merged in order after #2544. Each PR should pass CI on its own and link the others.[2/3] Add shared-run MLflow helpers(about 151 lines):split_tracking_credentials,log_active_run_experiment_json,_reattach/_closed_as/set_terminatedinmodelopt/torch/utils/mlflow.py, with theirtest_mlflow.pycases and theFakeMlflowresume support.[3/3] Track every Megatron-Bridge script with MLflow(about 401 lines):examples/megatron_bridge/*plustest_mlflow_utils.py. If #2544 also splits, renumber the whole stack to[x/4]. If a reviewer accepts the waiver instead, they should say so on the PR.
No action needed:
- ✔️ Resolved since the last review:
DISTILL.checkpointreturnsNonefor--validate_only,resolve_mlflow_argsis a plain alias,mlflow_runclearsargs.mlflowwhen the logger turns itself off, every rank declines a URI with only half auser:token, andEXPORT.sourcenormalisesiter_*paths. - The edits to existing tests are justified: the naming test now runs over all five tools, and the pointer checks now also cover the prune and export scripts. No assertion was loosened.
- The design is settled: this builds on the
Tool/tracked_runcore from #2544 rather than adding a second system.
This branch has not been deployed
2 participants
What does this PR do?
Type of change: new feature
[2/2] of a split. Based on #2544 — merge that first; this PR's diff is only the Megatron-Bridge half.
#2477 added MLflow tracking to
examples/megatron_bridge/quantize.py. It was one of five scripts in that directory that write a checkpoint; the other four recorded nothing, so the provenance chain stopped at the PTQ checkpoint and a deployed model could not be traced back to the run that produced it.All five now take the same
--mlflow/--mlflow_experiment/--mlflow_run_nameflags, and each declares what it records as aToolbeside its own flags — the sharedmlflow_utils.pyknows none of them:prune_minitron.pyprune_scoremetric, pointerquantize.py(#2477, moved onto the sharedToolin #2544)distill.pyexport_quantized_megatron_to_hf.pyexport_distilled_megatron_to_hf.pyEach writes
.experiment.jsoninto the checkpoint it produced, and each tags what it consumed, soprune → quantize → distill → exportis walkable both from disk and by tag query on the server.distill.pyopens the run and Megatron-Bridge joins it. ItsLoggerConfigrecords per-iteration metrics and the full resolved config — which a wrapper aroundmain()cannot see — but nothing ofdistill.py's own arguments and no invocation. Megatron-Bridge takesmlflow.active_run()when one exists, applies the tags and logs into it, sodistill_run()opens the run on the rank Megatron-Bridge looks at (the last one) and the two share it. Its early exit is handled explicitly:train()leaves throughsys.exit(0)on--exit_interval, which a blanket handler would record asFAILED.The library pieces that exist for that shared run land here with their first caller, rather than in [1/2] where they would have none:
split_tracking_credentials, so a URI handed to something which records it carries no credential;log_active_run_experiment_json, for pointing a checkpoint at a run this process did not open; andMlflowRunLogger._reattach, because a co-owner can end the run first — Megatron-Bridge does, asKILLED, when SIGTERM arrives mid-training.Two of Megatron-Bridge's defaults are deliberately not inherited: checkpoint artifact upload stays off unless
--mlflow_log_checkpoints(it pushes the whole checkpoint over HTTP after every save), and an untracked run passes nomlflow_*fields at all, since they landed in Megatron-Bridge 0.6 and sending them unconditionally would break an untracked run on an older one.Usage
Experiments default to
$USER/megatron_bridge_{prune,quantize,distill,export,distill_export}/<model basename>-<variant>.Testing
hf_ptq— prune, quantize, QAD distillation, quantized export, BF16 distillation, distilled export, HF PTQ — each closingFINISHEDwith the invocation, its arguments as params, its log, and a matching.experiment.jsonon disk. The chain tags line up: each stage'ssource_checkpoint_pathis the previous stage'scheckpoint_path.tests/examples/megatron_bridgeinnvcr.io/nvidia/nemo:26.08, the only lane that runs it: 76 passed. Plus the three suites from [1/2] One MLflow tracking core behind a Tool record #2544: 195 pass.pre-commit run --files <changed>: all hooks pass.LoggerConfig, the exit-callback guard, and theiter_*join.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Split from a single ~1150-line PR at review's request; #2544 carries the library consolidation this builds on, and this branch is based on it. Earlier review threads here show as outdated after the rebases — they are all resolved and their fixes are in this branch.
One known gap, stated in the README rather than implied:
distill.py --hf_export_pathwrites a second HuggingFace checkpoint from rank 0, which is not the rank that owns the run, so it carries no pointer yet. For the same reason the uploadedlogs/distill.logholds the last rank's output —print_rank_0keeps the script's own lines on rank 0 — which the README now says outright; carrying rank 0's log into a run owned by another rank needs cross-rank upload and is a follow-up.Two defects found on shared-run paths during review, both verified against the installed Megatron-Bridge 0.6 rather than its docs. Megatron-Bridge ends the run it shares with
distill.pyasKILLEDfrom its SIGTERM handler (train.py:1413) and then leaves throughsys.exit()(train.py:805), i.e. beforedistill_run'sfinally— and MLflow's fluent calls resolve their target by opening a run when none is active, so a preempted distillation's log and metrics went to a second, empty run and itsKILLEDstatus was overwritten. Separately, an unreachable server disabled our logger butlogger_kwargsstill handed Megatron-Bridge the same URI, andstate.pycallsset_experimentunguarded from inside the training loop — so a best-effort$MLFLOW_TRACKING_URIaborted the training instead of degrading to untracked.🤖 Generated with Claude Code