feat(training): support multiple validation sets and eval at start

[MoPl90]

What does this PR do ?

Implements #3996 (eval at start; approved by @yaoyu-33) and #618 (multiple validation sets + separate loss reporting), bringing Bridge's training loop to parity with Megatron-LM's multiple_validation_sets / validation_set_names convention. Bridge inherits both fields in ValidationConfig, and BlendedMegatronDatasetBuilder already returns a list of validation datasets when the flag is set, but Bridge's loaders and eval cannot consume it today (likely related to the crash reported in #2101).

Changelog

• training/config.py: add Bridge ValidationConfig subclass (same pattern as TrainingConfig) with eval_at_start: bool = False
• training/train.py: extract the in-loop eval into _run_validation() (all existing logic preserved one-to-one); called once before the loop when eval_at_start and do_valid are set
• training/eval.py: per-set evaluation and suffixed logging in evaluate_and_print_results() for the supported loggers (TensorBoard / W&B / MLflow / Comet); cross-rank dataset-count guard; raises ValueError on validation_set_names length mismatch
• data/loaders.py: folds the three duplicated valid/test dataloader build sites into _build_eval_test_dataloaders(), mapping over dataset lists; errors when a provider returns multiple validation datasets without the flag
• Tests: see Testing section below

Design notes

1. No combined loss in multi-val mode. Megatron-LM's convention evaluates each set separately and logs no aggregate; we match it here
2. Collective balance under PP. Every rank runs the same number of fully-flushed evaluate() calls per eval round, so the per-set loop cannot desync pipeline collectives
3. Fail loud instead of hanging. A min/max all-reduce on the set count raises on every rank if ranks ever disagree (e.g. a provider deriving sets from per-rank file reads), as otherwise the mismatch surfaces as a silent NCCL hang
4. Why one PR and not two: both features reshape evaluate_and_print_results and the eval plumbing in train()/loaders.py, so most of the diff (the _run_validation() extraction, the per-set eval restructuring, the test scaffolding) would be needed for either feature alone. They are also used together in practice: the motivating use case is a step-0 baseline per validation set, and the end-to-end evidence below exercises exactly that combination

Testing

Unit tests

• test_eval.py: indexed and named suffixes; exact logger calls (metric names, ppl, step) on TensorBoard/W&B; validation_set_names length mismatch raises before any evaluation; timelimit mid-loop aborts remaining sets; a list without the flag is evaluated once (VPP); rank set-count agreement passes / mismatch raises; is_test never splits a list.
• test_loaders.py: one dataloader per set with correct per-set contents; iterator layer yields one RerunDataIterator per set in blend order; guard error without the flag.
• test_config.py: eval_at_start default/settable; multiple_validation_sets validates under PP 1 and 2.

End-to-end training

Ran 6 short training jobs (train_iters=4, eval_interval=2, eval_iters=2) with eval_at_start=true and multiple_validation_sets=true on a blend of 5 validation datasets, using dense Qwen3.5-VL models (2B and 9B sizes). All 6 finished and logged one loss series per validation set, starting at iteration 0. The matrix covers:

• Model size. 2B on 4 GPUs and 9B on 8 GPUs.
• Set naming. validation_set_names=["set-A", ..., "set-E"] vs unset (falls back to the index suffix -0 ... -4).
• Parallelism. Pure DP (TP1/PP1/DP4), TP-only (TP4/DP2), and pipeline parallel (TP2/PP2/DP2), confirming per-set loss reduction and logging work when the loss lives on the last PP stage.

Validation on start. With eval_at_start=true, all validation sets are evaluated before the first training step and logged at iteration 0. From the 2B run with named sets:

Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
--------------------------------------------------------------------------------------------------------------------------
 validation-set-A loss at iteration 0 (pre-train validation) | lm loss value: 8.543873E-01 | lm loss PPL: 2.349934E+00 |
--------------------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
--------------------------------------------------------------------------------------------------------------------------
 validation-set-B loss at iteration 0 (pre-train validation) | lm loss value: 1.473683E+00 | lm loss PPL: 4.365282E+00 |
--------------------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
--------------------------------------------------------------------------------------------------------------------------
 validation-set-C loss at iteration 0 (pre-train validation) | lm loss value: 9.426797E-01 | lm loss PPL: 2.566851E+00 |
--------------------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
--------------------------------------------------------------------------------------------------------------------------
 validation-set-D loss at iteration 0 (pre-train validation) | lm loss value: 8.461825E-01 | lm loss PPL: 2.330732E+00 |
--------------------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
--------------------------------------------------------------------------------------------------------------------------
 validation-set-E loss at iteration 0 (pre-train validation) | lm loss value: 9.909064E-01 | lm loss PPL: 2.693675E+00 |
--------------------------------------------------------------------------------------------------------------------------

Multi-dataset validation. Each set is evaluated independently at every eval_interval and at the end of training, printed as its own block with the set suffix. From the 9B TP2/PP2 run (default index naming, final validation):

Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
-----------------------------------------------------------------------------------------------------------------
 validation-0 loss at iteration 4 on validation set | lm loss value: 9.499896E-01 | lm loss PPL: 2.585683E+00 |
-----------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
-----------------------------------------------------------------------------------------------------------------
 validation-1 loss at iteration 4 on validation set | lm loss value: 1.331964E+00 | lm loss PPL: 3.788476E+00 |
-----------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
-----------------------------------------------------------------------------------------------------------------
 validation-2 loss at iteration 4 on validation set | lm loss value: 8.635647E-01 | lm loss PPL: 2.371600E+00 |
-----------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
-----------------------------------------------------------------------------------------------------------------
 validation-3 loss at iteration 4 on validation set | lm loss value: 7.436793E-01 | lm loss PPL: 2.103661E+00 |
-----------------------------------------------------------------------------------------------------------------
Evaluating on 64 samples
Evaluating iter 1/2
Evaluating iter 2/2
-----------------------------------------------------------------------------------------------------------------
 validation-4 loss at iteration 4 on validation set | lm loss value: 7.696714E-01 | lm loss PPL: 2.159057E+00 |
-----------------------------------------------------------------------------------------------------------------

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

• Closes [feature] Eval at start (step-0 eval) #3996
• Closes Feature Request: Support for multiple validation sets + separate loss validation reporting #618
• Related to [BUG/CLARIFICATION] enabling multiple_validation_sets in GPTDatasetConfig results in a ZeroDivisionError when entering in a validation cycle #2101
• Suggested labels: feature, area:training, needs-review

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

Thanks for adding these two features. I found four correctness issues around validation iterator shape, checkpoint resume, and step-zero evaluation bookkeeping; details are inline.

[yaoyu-33]

This applies the single aggregate consumed_valid_samples offset to every validation set. evaluate() increments that scalar for every set, so after N sets each individual loader has consumed only 1/N of the recorded total. After checkpoint resume, every loader skips the full aggregate offset (for example, 5 sets x 2 eval iterations makes each set resume at 10 iterations rather than 2), silently changing the evaluated samples or repeating cyclic data. MCore currently refuses multiple validation sets without full validation for this exact reason. Please either persist per-set offsets, or require and fully implement full_validation, and add save/resume coverage.

[MoPl90]

As I outline below, I opted for the full bookkeeping approach instead of full_validation. See my comment below for details.

[yaoyu-33]

When multiple_validation_sets=True but the provider returns a single validation dataset, this branch produces one RerunDataIterator. evaluate_and_print_results() then assumes the flag means a sequence of sets and calls len() on it, but RerunDataIterator has no __len__ or __iter__, so the first validation fails. The built-in blended builder does return a single dataset when the validation blend has one prefix. Please normalize it to a one-element list whenever the flag is enabled, or reject that configuration early, and cover the singleton case.

[MoPl90]

Addressed in 5d93fab

[yaoyu-33]

This extra validation pass is not included in get_train_valid_test_num_samples(): the existing formula reserves samples for interval evaluations plus the final evaluation only. For train_iters=10, eval_interval=5, and eval_iters=2, it builds 6 validation iterations but this change executes 8. A cyclic loader hides the shortage by reusing early samples, while a single-pass loader can exhaust early. Please add one eval_iters * eval_global_batch_size allocation when the pre-train pass can run.

[MoPl90]

With cb5601f this is now correctly calculated.

[yaoyu-33]

This condition also runs at a nonzero resumed step and on every NVRx in-process retry, because recovery re-enters train() with the restored step. That duplicates validation/logging and advances validation state during recovery. Issue #3996 proposed guarding this with train_state.step == 0. If ordinary checkpoint resume should intentionally evaluate, please distinguish it from an automatic in-process retry; otherwise add the step-zero guard and a nonzero-step test.

[MoPl90]

Addressed in b045fe2

[MoPl90]

Thank you for the review and spotting those correctness issues. I will address your inline comments directly.

Regarding the consumed_valid_samples offset, I considered following the MCore logic (full_validation), but it derives per-set eval iterations from len(dataloader), which doesn't exist for external dataloaders. This is therefore orthogonal to the offset book keeping approach and doesn't fit our use case. Happy to address as a follow-up. Instead I:

• added TrainState.consumed_valid_samples_per_set: one offset counter per validation set, persisted in the checkpoint train state. Checkpoints written before this field existed load as an empty list and the counters initialize to zeros, so checkpoints resume cleanly
• enabled per-set resume in the loader builder. build_train_valid_test_data_loaders sizes the per-set list once per build and resumes each validation loader from its own offset. If the number of sets changed since the checkpoint, the counters reset to zero with a warning
• advance the aggregate: evaluate() advances consumed_valid_samples for every set, and additionally the indexed per-set counter when multiple_validation_sets is on. Rationale: preserve the field's existing semantics (total across sets); consumed_valid_samples == sum(consumed_valid_samples_per_set) holds, and a multi-set checkpoint resumed with the flag disabled falls back to the aggregate offset with a warning, i.e. the pre-PR behavior for that configuration
• note: External dataloaders (dataloader_type="external") pass offsets through as before; for them the per-set counters are bookkeeping only, since sampler offsets never applied to external loaders in the first place

Moreover, I added unit tests and test_pretrain_resume_multi_val.py, a functional end-to-end test (L1) that checkpoints mid-run, resumes, and verifies the per-set counters persist and continue from their restored values (verified on a 2-GPU machine). Two notes:

1. It reads the tracker with raw torch.load on all ranks rather than the read_train_state helper, which is lru_cached (returns stale state on a second read) and broadcasts its result across ranks
2. It sets async_save=False because it reads latest_train_state.pt immediately after pretrain() returns, which otherwise races the async finalization write

Happy to address any remaining concerns!

[yaoyu-33]

Thanks for the thorough revision — all four items from my previous round are genuinely addressed, and the approach you picked for the per-set offsets (real bookkeeping in TrainState rather than requiring full_validation) is the right call. The save/resume functional test with two sets and dataloader_type="single" is exactly the coverage that was missing.

I also checked two things that are easy to get wrong here, and you got both right, so no action needed:

• PP collective balance holds. evaluate() returns {} rather than None off the last pipeline stage, and the only return None path is the timelimit one, which is gated on an all_reduce(MAX). So the early return None inside the per-set loop cannot leave ranks with unequal evaluate() counts.
• _run_validation(toggle_pre_hook=False) is correct for the pre-train pass. The forward pre-hook is already disabled at train.py:338 before your insertion point, and interval-time is started at train.py:189, so the timers("interval-time").stop() inside the helper will not trip MCore's assert self._started.

One blocking issue and a few smaller ones inline. The PR is also currently conflicting with main and will need a rebase.

[yaoyu-33]

Blocking: the MegatronMIMO path bypasses all of this.

build_train_valid_test_data_loaders() returns early for MegatronMIMOProvider at line 259, which is above this normalization block and above the consumed_valid_samples_per_set sizing at line 400. If a user sets multiple_validation_sets=True on a MIMO model, three things go wrong:

1. train_state.consumed_valid_samples_per_set is never sized, so it stays []. evaluate() then does state.train_state.consumed_valid_samples_per_set[0] += eval_batch_size → IndexError.
2. train_megatron_mimo.py:390-393 gates evaluation on a per-rank condition (valid_data_iterator is not None), because heterogeneous MIMO legitimately leaves some ranks without a validation loader. The new set-count all_reduce in evaluate_and_print_results() is unconditional, so only a subset of ranks enters it → collective hang. That is the exact failure mode the all_reduce was added to prevent.
3. len(data_iterators) is called on a single RerunDataIterator, which has no __len__ → TypeError.

Given design note 3 in the PR description ("fail loud instead of hanging"), the consistent fix is to reject the combination up front — e.g. in ConfigContainer.validate(), raise when validation.multiple_validation_sets is set and cfg.model is a MegatronMIMOProvider. Supporting it properly is fine too, but out of scope for this PR; an explicit error is enough.

Please add a unit test for whichever guard you choose.

[MoPl90]

Agree that this is out of scope. I added a guard in 1a80d21, which also loudly fails for one more potential silent failure modes: cfg.validation.full_validation can be set, but has no effect. Now raises a ValueError at startup.

[yaoyu-33]

eval_at_start re-runs on every ordinary checkpoint resume, but only one extra pass is budgeted.

is_inprocess_restart_retry only suppresses NVRx in-process retries. An ordinary resume from a checkpoint is a fresh process with restart_attempt == 0, so this branch runs again.

Meanwhile get_train_valid_test_num_samples() (loaders.py:133-135) reserves exactly one extra eval_iters worth of samples. A job that is preempted and resumed K times executes K+1 pre-train validation passes and consumes K * eval_iters * eval_global_batch_size more validation samples than the dataset was sized for.

Concrete failure: train_iters=10000, eval_interval=1000, eval_iters=2, dataloader_type="single" (non-cyclic — the same type your new functional test uses). Five preemptions ⇒ six pre-train passes ⇒ ten eval iterations of validation data beyond the reservation. The single-pass sampler drifts forward and can exhaust the validation set / raise StopIteration.

Issue #3996 proposed guarding on train_state.step == 0 specifically, and I think that matches the actual use cases: the value of this feature is a baseline for the loaded checkpoint at the start of training (SFT from a converted HF checkpoint, sanity-checking a HF↔MCore conversion before burning GPU-hours, anchoring an A/B of data mixes). On resume you already have that number from the last in-loop eval, so re-running it adds cost and puts a duplicate point at an already-logged step in TensorBoard/W&B.

Either add the step == 0 guard, or — if you deliberately want resume to re-evaluate — say so in the field docstring and widen the sizing formula accordingly. Right now it does the widened behaviour with the narrow budget. Please also add a nonzero-step test for whichever semantics you land on.

[yaoyu-33]

Following up on my own comment — I agree the step == 0 semantics is the one that actually carries meaning, and I checked that it covers the motivating use cases cleanly:

load_checkpoint() resets state.train_state.step = 0 when cfg.checkpoint.finetune or release is set (checkpointing.py:2322-2323). So "SFT / continued pretraining from a converted HF checkpoint" — the strongest reason to want this feature — already lands at step 0 and gets its baseline. A plain mid-run resume does not, and that is exactly the case where the extra pass buys nothing (the previous in-loop eval already produced that number) while costing unbudgeted validation samples and putting a duplicate point at an already-logged step. So step == 0 is not a compromise here; it is the semantics that matches the use cases.

Given that, I think the field should be renamed while it is still unreleased and free to change. eval_at_start invites exactly the ambiguity this thread is about — "start" of the run, or of training? Two candidates, in my order of preference:

1. eval_at_step_zero — states the guard itself, so the behaviour on resume is unambiguous from the name alone and the docstring does not have to carry the disambiguation.
2. baseline_eval — intent-revealing: a baseline is by definition taken once, at the beginning, which makes "why doesn't this re-run on resume?" a non-question.

There is an additional argument for renaming: ValidationConfig already inherits MCore's start_eval_at_iter (megatron/training/config/training_config.py:125), which delays the first evaluation. Having eval_at_start and start_eval_at_iter sit next to each other in the same config — near-identical wording, essentially opposite meanings — is a real readability trap for anyone scanning the config.

Happy to defer to your preference on the exact name; the substantive ask is the step == 0 guard plus a name that does not require a docstring to disambiguate.

[MoPl90]

I think I can agree with this argument. The checkpoint-resume functionality was valuable for development, but practically adds little value.

The suggested re-naming is also reasonable given that Megatron-core already has a near-identical config field.

Fixed in 8f8fa1a

[yaoyu-33]

validation_set_names set without multiple_validation_sets is silently ignored — this length check only runs inside the flag branch, and the suffix logic at line 486 is likewise gated on the flag. The user gets no names and no warning.

MLM asserts on this (megatron/training/training.py:3633: assert args.multiple_validation_sets, "--validation-set-names requires --multiple-validation-sets"). Worth matching.

Ideally both this check and the names/count check move into ConfigContainer.validate() — a static config error should fail at startup, not at the first eval_interval after the user has already paid for model init and dataset build.

[MoPl90]

Added a config level check in e176de0

[yaoyu-33]

Two concerns about this all_reduce:

1. Process group. It uses the default world group, while the rest of this function routes collectives through pg_collection (and evaluate() takes an explicit p2p_communicator/pg_collection). For multi-module / VLM setups that is inconsistent.
2. Placement. The number of validation sets is a static property of the config and the dataset provider — it cannot change between eval rounds. Validating it costs a full-world collective on every eval round, on the hot path. build_train_valid_test_data_loaders() already knows the count (line 400) and is where the per-set counters get sized; doing the agreement check once there and trusting it afterwards is both cheaper and closer to the source of truth.

The check itself is a good idea — I just think it belongs at dataloader-build time.

[MoPl90]

Excellent point! Moved to dataloader build in 666aefc

[yaoyu-33]

Non-blocking, but worth pinning down: the return value is now the last set's loss dict when multiple_validation_sets is on. No current caller uses the return value (pretrain.py:165/178, train.py:612, train_megatron_mimo.py:396 all discard it), so nothing breaks today — but this is a silent trap for any future early-stopping / save-best-checkpoint logic, which would end up keying off whichever set happens to be last in the blend.

Consider returning a dict keyed by the set suffix (or None in multi-set mode) so the ambiguity is impossible to consume accidentally. The docstring update is good; making it unrepresentable would be better.

Separately: when the timelimit fires mid-loop you return None after sets 0..k have already advanced their per-set counters while k+1..N-1 have not, so those sets resume from different offsets. That is defensible since the sets are independent, but it deserves a sentence in the docstring.

[MoPl90]

Valid concern, addressed in a8b35d3

[MoPl90]

Thank you for reviewing again! It seems like I missed some crucial changes when merging and resolving conflicts. To make sure this is not the case now, I ran the following tests and made sure they are all green: Unit suite (311 tests) and the L0/L1 pretrain functional tests (2 GPUs) pass on the branch tip.

Addressing your inline comments above. Thanks!

[MoPl90]

A few additional hardening commits on top of the requested changes:

• Reject a list of test datasets (56221bd): previously the list would nest into the dataloader machinery and fail far from the cause
• Clear stale per-set counters when multi-val is disabled (c444263)
• Reject duplicate validation_set_names (7a12110) that would collide in logging keys and in the per-set loss dict
• Fail fast on a non-list data_iterator with multiple_validation_sets set (954c40a)

One behavior noted but intentionally unchanged: weights in the validation blend are ignored under multiple_validation_sets, matching Megatron-LM's builder.

The branch is also merged with the latest main. At the current tip, the unit suites for test_pretrain, test_config, test_loaders, and test_eval, the test_pretrain_resume_multi_val functional test, and the L0 pretrain functional group pass on 2 GPUs.

[MoPl90]

Happy to address any further concerns if there are any @yaoyu-33

[MoPl90]

@yaoyu-33 I've merged with main and incorporated your update to the eval_iters counter in #5810. All still green on the GPU tests.

[MoPl90]

@yaoyu-33 should we proceed with this feature?

[MoPl90]

@yaoyu-33 I merged the latest main into the branch and fixed what the merge broke. Changes since the last head:

• Conflicts. There were three, in config.py, eval.py and test_eval.py. In eval.py, the empty_loss_keys tracking from fix(training): guard eval loss reduction and surface zero-token steps #6239 now gets a fresh set for each validation set. An all-masked set leaves a gap only in its own metrics. A new unit test covers this.
• Global-batch packing. The validation size check from feat(training): global-batch online packing with Megatron-Core's sequence-packing scheduler #6127 runs after the multi-val normalization. It compared the length of the list of sets against the sample budget, so every multi-val run with packing failed. It now checks each set (9ff2684).
• Unsupported combinations. ConfigContainer.validate() now rejects multiple_validation_sets with DPO and eval_at_step_zero with MegatronMIMO. The DPO loader from feat(training): add DPO training with offline-scored reference model #6092 returns before the per-set bookkeeping, and the MIMO training loop never runs the step-zero pass (deebf59).
• Test fixtures. Four train() unit tests in test_pg_collection_wiring.py and test_train.py build a SimpleNamespace validation config without eval_at_step_zero. They now set it to False.
• Docs. The Validation section of docs/training/training-loop-settings.md now covers the three new fields and how multiple validation sets behave (3b21583).

Unit tests for test_eval, test_config, test_loaders, test_train, test_pg_collection_wiring and test_state pass

Could you take another look if it looks good?