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.

[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.

[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.

[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.