feat(sft): add SFT v2 stage 1 with the Energon multimodal data path - #3917
Conversation
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
…enchmark path that never runs in a real SFT job. Keeping it forced every SFTv2 config to carry a required sft_v2 block whose only production value was loader_only=false. tools/benchmark_loaders.py already measures loader throughput standalone, and does so against the Megatron reference loader." Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
_run_train_step computed train_results but discarded it, so _policy_metrics was dead code and loss/grad_norm/lr never reached the logger. Merge the flattened policy metrics into the returned metrics dict. Signed-off-by: rohitrango <rohit.rango@gmail.com>
… local adapter Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
|
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. |
github-actions
Bot
added
the
Documentation
There was a problem hiding this comment.
Self-review: findings still outstanding
This PR is a draft. Posting my own review up front so the outstanding work is visible rather than discovered in CI. Reviewed by a coordinated team of five agents
(RL/algorithms, Energon, data-plane bug hunt, tests, config/deps), with every finding independently re-verified against the source before posting. Upstream claims are
checked against megatron-energon tag 7.4.0 (what uv.lock actually resolves for the ~=7.0 pin) and spot-checked as identical at 7.4.1.
This is careful work. A few things worth calling out as done well, because they are the reason this review is about plumbing rather than architecture:
- The optional-dependency boundary is drawn deliberately and correctly.
nemo_rl/data/energon/__init__.pydefers thesft_dataloaderimport into function bodies and
multimodal/__init__.pyis intentionally empty with a comment explaining why, soimport nemo_rl.datanever pullsmegatron.energon. - The checkpoint identity fingerprint validates backend, format version and identity before
torch.load(..., weights_only=False), and reports a named per-field diff
instead of a bare hash mismatch. The restore-before-iteration guard is a strictly tighter, better-worded version of upstream's bareassert. - Every Energon API call checks out:
get_savable_loaderkwargs, all sixWorkerConfigfields including both error handlers,save_state_rank/restore_state_rank
(the right pair, not the deprecated or global variants),FileStoreCachePool,Cooker,SampleDecoder, and all fiveTaskEncoderhooks are really dispatched. No
fabricated kwargs, no dead hooks. - The loss/label masking chain is correct end to end: the encoder builds no labels, mask, padding or ignore-index of its own and delegates entirely to the shared code
the HF path uses, so there is no second implementation to drift. The shift is applied exactly once. _max_train_stepstreatingvirtual_epoch_lengthas optimizer steps is correct - upstream appliesEpochizeDatasetafterbuild_batch, so it counts batches
per rank. We chased this one hard expecting a factor-of-batch-size bug and it is not there; the code, the docs and the4375recipe value all agree.
What blocks merge today
Four items make CI red immediately: a unit test asserting against a method the packing refactor deleted; four new test files that fail collection in the non-mcore
Data shard; two unregistered recipes failing test_recipes_and_test_suites.py twice; and a stale uv.lock failing uv lock --check.
A theme worth naming
Several findings share one root cause: the "strip Energon packing from stage 1" refactor looks half-finished. The input_is_prepacked branch has no producer (its
producer was deleted in 7a541129), select_samples_to_pack's test outlived the method, and packing_cost, pending_sample, sample_schema and batch_grouping
are all set-but-never-read. Individually trivial; together they make it hard to tell which fields are load-bearing. A sweep for leftovers from that refactor would
clear a good fraction of the list.
Performance and convergence evidence
This adds a second independent SFT training loop, and there are no numbers anywhere in the PR. The most useful thing would be a v1-vs-v2 loss curve on the same config
and seed - vlm_sft-qwen2.5-vl-3b-instruct-clevr-1n2g-megatrontp1-energon.v1.yaml for ~200 steps at sft.seed=42, run once through run_sft.py (v1 loop, Energon
backend) and once through run_sft_v2.py, plotted on one axis. Same data, same NLLLossFn; if v2 does not track v1 within noise the placed-metadata normalization is
wrong. Note the metric-prefix mismatch has to be fixed first or the two runs will not share keys. A throughput/peak-memory number and one kill-and-resume round trip
would round it out, since the colocated-loader design is the whole motivation and nothing currently tests resume.
Also worth a look
docs/guides/sft.md's config snippet is missing two required keys and will not load as written, and the av override change re-admits a dependency that was
deliberately excluded from the shipped container for CVE and royalty reasons - worth making that an explicit, recorded decision rather than a commented-out line,
especially since the comment above it still says the opposite.
Findings are inline on this PR, ordered roughly by severity. Several are one-line fixes.
Generated by Claude Code
… in-worker Review fixes: - sft: mask degenerate rows (active but with no trainable token) out of the batch with a warning instead of aborting the run, and carry source_ids through the multimodal prepare path so the data plane keeps per-sample provenance. - sft_v2: validate virtual_epoch_length on the driver before the cluster is allocated, reject the SFTConfig validation knobs that this loop has no path for, index policy config strictly, and drop the duplicated loader-latency metrics. - cookers: try the bare media extension last so two members sharing a suffix no longer collide. - generic_sft: empty the media tensors when a conversation is truncated to a zero-weighted row, and raise the failure tolerance off 0, which Energon reads as "disable tolerance checking". - sft_dataloader: start the first-sample viewer hint enabled, pin batch_size into the V2 loader identity, and set shuffle_over_epochs_multiplier so data.shuffle=false gives a repeatable order. - data_plane: drop the generation of a cleared partition, and roll the previous checkpoint back if the replacing rename fails. - tq_policy: mint a fresh forward padding target per step rather than mutating the caller's metadata and ratcheting the target upward. CI fixes: - importorskip megatron.energon/megatron.core and apply the mcore mark at module scope in the Energon tests; the collection-time mark hook runs too late to stop an import error. - Register the Energon and local data-plane modules in pyrefly.toml, map vlm_sft to its base YAML, and apply ruff formatting. Processor rerouting: - Build the processor inside SFTMegatronPolicyWorker from the tokenizer config instead of shipping it as a constructor argument. A trust_remote_code processor's class lives in transformers_modules, which Ray worker interpreters cannot import while deserializing arguments. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: rohitrango <rohit.rango@gmail.com>
Companion lock update for a6ac194, which moved megatron-energon out of the energon extra and into mcore. Adds the av package entry, drops the sys_platform=='never' override that had been pinning it out, and records the av-decode extra on megatron-energon for both mcore and dev. The rest of the diff is metadata for the Gym and Megatron-Bridge workspace members, picked up because the lock was regenerated against newer submodule revisions than the ones this commit points at. Those pointer bumps land separately. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: rohitrango <rohit.rango@gmail.com>
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 3b41b28 |
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
rohitrango
added
CI:L1
Signed-off-by: rohitrango <rohit.rango@gmail.com>
|
/ok to test df685d7 |
Signed-off-by: rohitrango <rohit.rango@gmail.com>
|
/ok to test c916435 |
|
/ok to test c1ea6e3 |
Signed-off-by: rohitrango <rohit.rango@gmail.com>
|
/ok to commit 6110d9a |
|
/ok to test 6110d9a |
rohitrango
added
CI:Lfast
…age1 Signed-off-by: rohitrango <rohit.rango@gmail.com> # Conflicts: # uv.lock
|
/ok to test a0d0826 |
…age1 Signed-off-by: rohitrango <rohit.rango@gmail.com>
|
/ok to test b8c8b63 |
#3917 added nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker to ray_actor_environment_registry.py with the mcore extra. This branch builds that registry from ACTOR_ENVIRONMENTS instead of a literal dict, so the new entry did not carry over on rebase and the actor would have raised from get_actor_python_env at launch. Signed-off-by: Terry Kong <terryk@nvidia.com>
#3917 added nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker to ray_actor_environment_registry.py with the mcore extra. This branch builds that registry from ACTOR_ENVIRONMENTS instead of a literal dict, so the new entry did not carry over on rebase and the actor would have raised from get_actor_python_env at launch. Signed-off-by: Terry Kong <terryk@nvidia.com>
#3917 added nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker to ray_actor_environment_registry.py with the mcore extra. This branch builds that registry from ACTOR_ENVIRONMENTS instead of a literal dict, so the new entry did not carry over on rebase and the actor would have raised from get_actor_python_env at launch. Signed-off-by: Terry Kong <terryk@nvidia.com>
3 participants
What does this PR do ?
Adds SFT v2 (stage 1): a single-controller SFT training path with colocated Megatron-Energon loaders, one per data-parallel replica.
Three pieces ship here:
nemo_rl/algorithms/sft_v2.py— a single-controller SFT loop (SFTSingleControllerActor) drivingTQPolicydirectly, with its own step accounting, checkpointing and loader placement, plus theexamples/run_sft_v2.pyentry point.nemo_rl/data/energon/— a multimodal Energon SFT data path: cookers, task encoders, a model-family registry, a rank-aware savable loader with a checkpoint identity fingerprint, and a topology mapper from policy parallel ranks to logical loader shards.nemo_rl/data_plane/adapters/local.py— a process-local data-plane adapter for colocated batches, with checkpoint save/load andlist_sample_ids.The existing v1 Hugging Face SFT path is unchanged by default:
data.backenddefaults tohf, and the Energon imports are deferred soimport nemo_rl.datadoes not requiremegatron.energon.Stage 1 boundaries (deliberately narrow):
data.backend=energon)packing_buffer_size/max_samples_per_sequencemust stay null)Issues
None linked.
Usage
Two recipes are included: Qwen2.5-VL-3B on CLEVR (1n2g, TP1) and Nemotron-Omni-30BA3B on CLEVR (1n8g, TP4/ETP4/EP2).
Validation on Qwen2.5-VL-3B, CLEVR (expanded for image payloads to be 9x larger), DP=8, TP=2, seqlen=8192
Blue curve shows V1 and orange curve shows V2.
Before your PR is "Ready for review"
Pre checks: