[deepakn94]

What

evaluate() all-reduced its loss over pg_collection.dp_cp, which is the replicate data-parallel group and excludes the GTP weight-remat axis (use_mpu_process_groups() builds it with with_gtp_remat=False, megatron/core/process_groups_config.py:261). But gtp_remat peers hold distinct micro-batches: the eval sampler uses mpu.get_data_parallel_rank() / get_data_parallel_world_size(), both of which default to with_gtp_remat=True.

See figure demonstrating the issue:

So the full eval batch was read and forwarded, and then only a 1/gtp_remat slice of it was averaged into the reported number. On a run with dp=6, gtp_remat=64, eval_global_batch_size=768: all 384 distinct-data ranks read their micro-batches (768 sequences per eval iteration), the all-reduce covered 6 of those ranks (12 sequences), and consumed_valid_samples still incremented by the full 768. Over --eval-iters 14 that is ~168 sequences behind each reported validation loss instead of 10,752. The estimate stays unbiased, but its standard error is ~8x larger, which shows up as visible scatter between evaluations.

train_step already handles this correctly and picks dp_cp_gtp_remat with a comment explaining exactly this hazard (megatron/training/training.py:3178). evaluate() was missed in that migration. This change makes it consistent.

Also

validate_args asserted eval_global_batch_size % (eval_micro_batch_size * data_parallel_size) == 0, but evaluate() divides by eval_micro_batch_size * data_parallel_size * gtp_weight_remat_size to derive its micro-batch count. With GTP active the assert can pass while the division floors to a smaller count, or to zero micro-batches, silently evaluating nothing. The assert now includes gtp_weight_remat_size.

megatron/training/yaml_arguments.py carries the same assert, but that path never factors GTP into data_parallel_size at all, so it needs a broader change and is left alone here.

Impact

No behavior change when gtp_weight_remat_size == 1: dp_cp_gtp_remat aliases dp_cp in that case, and the added assert factor is 1.

🤖 Generated with Claude Code

[copy-pr-bot]

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.

[deepakn94]

/claude review

[deepakn94]

/ok to test b3ecb9b

[claude]

Light review — the fix itself looks correct. Verified that args.gtp_weight_remat_size is derived at arguments.py:415, well before the new assert at :713; that the assert now matches the division evaluate() actually performs at training.py:5081; and that the getattr(..., "dp_cp_gtp_remat", None) or eval_pgc.dp_cp fallback mirrors train_step (training.py:3178). Both changed files are under megatron/training, so the process-group migration guidance does not apply.

Two notes:

1. The existing divisibility unit test was not updated for the new factor.

tests/unit_tests/test_eval_batch_size.py::TestEvalBatchSizeDivisibility re-implements the assert expression inline rather than calling validate_args, so it silently keeps testing the old two-factor formula and will pass unchanged against this PR. The regression this PR fixes — the assert passing while evaluate() floors its microbatch count to zero — therefore still has no test guarding it. Worth extending that class, e.g.:

    def test_gtp_remat_factored_into_divisibility(self):
        """Regression: 32 %% (4*2) == 0 passes the old assert, but
        32 // (4*2*8) == 0 microbatches, so evaluate() evaluates nothing.
        """
        args = create_test_args(
            eval_global_batch_size=32,
            eval_micro_batch_size=4,
            data_parallel_size=2,
            gtp_weight_remat_size=8,
        )
        assert args.eval_global_batch_size %% (
            args.eval_micro_batch_size * args.data_parallel_size
        ) == 0
        assert args.eval_global_batch_size %% (
            args.eval_micro_batch_size
            * args.data_parallel_size
            * args.gtp_weight_remat_size
        ) != 0

create_test_args in that file will also need gtp_weight_remat_size = 1 added to its defaults so the existing cases keep working.

2. Minor: error-message misdirection when --global-batch-size is omitted under GTP.

arguments.py:701 defaults global_batch_size = micro_batch_size * data_parallel_size without the GTP-remat factor, and eval_global_batch_size then inherits it. With gtp_weight_remat_size > 1 the new assert therefore always trips, pointing the user at eval_global_batch_size when the real problem is the global_batch_size default. Not a regression — that config already died later in init_num_microbatches_calculator (num_microbatches_calculator.py:381, or the num_micro_batches >= 1 assert), so nothing that previously worked breaks. Only the message the user sees changes, and it now names the wrong knob. Out of scope here, but perhaps a follow-up alongside the yaml_arguments.py:132 gap already called out in the description.

[claude]

Correction to the snippet above: the %% in the two assert lines should be single % (modulo). Shell-escaping artifact, not part of the suggestion.

[deepakn94]

/claude review

[deepakn94]

/ok to test 873e7c0

[claude]

Light review — the core change looks correct. eval_dp_cp_group matches the existing pattern at megatron/training/training.py:3178 and megatron/core/distributed/finalize_model_grads.py:603, the getattr(..., None) or fallback correctly handles pg_collections that do not populate dp_cp_gtp_remat (e.g. the MIMO namespaces in tests/unit_tests/models/mimo/test_mimo_mock_data.py:52), and the sft branch's get_world_size divisor now uses the same group as its all_reduce. Both edits are no-ops at gtp_weight_remat_size == 1.

Two items:

1. Stale user-facing docstring for the divisibility rule.

megatron/training/config/training_config.py:132 still documents the two-factor rule. This docstring is what ArgumentGroupFactory turns into --eval-global-batch-size CLI help, so users now get help text that understates the constraint the assert enforces:

    eval_global_batch_size: int | None = None
    """Global batch size to use during evaluation. If not set, defaults to global_batch_size.
    Must be divisible by (eval_micro_batch_size * data_parallel_size * gtp_weight_remat_size),
    where data_parallel_size is the replicate degree and gtp_weight_remat_size is the GTP
    weight-remat axis (both are axes over which distinct micro-batches are distributed).
    """

2. No regression test for the group-selection half of the fix.

The validate_args half got TestEvalBatchSizeDivisibilityWithGTP, but the dp_cp -> dp_cp_gtp_remat swap in evaluate() — the actual reported-loss bug — has none, so a future refactor could silently revert it. A cheap non-distributed guard in the same spirit as tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py:762 (which asserts group selection rather than numerics) would cover it: build a SimpleNamespace pg_collection whose dp_cp and dp_cp_gtp_remat are distinguishable stubs, monkeypatch torch.distributed.all_reduce to record the group= it receives, and assert the recorded group is the gtp_remat-inclusive one — plus a second case with dp_cp_gtp_remat=None asserting the dp_cp fallback still works.

That does require driving evaluate(), so if that is more scaffolding than it is worth, an assertion in an existing GTP test that the eval path resolves to the gtp_remat-inclusive group would at least pin the intent.

Also, thanks for calling out megatron/training/yaml_arguments.py:132 explicitly in the PR description rather than leaving it as a silent divergence — agreed that path needs a broader change since it never folds GTP into data_parallel_size at all.

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32454847897