[deepakn94]

Summary

get_batch in pretrain_gpt.py and pretrain_hybrid.py rebuilt the transformer config from args on every microbatch via core_transformer_config_from_args. That re-runs TransformerConfig.__post_init__ each time, so its validation warnings repeat in the log every step. For example, with --fp32-residual-connection the rank-0 log gets fp32_residual_connection is enabled, overriding pipeline_dtype from torch.bfloat16 to torch.float ... once per microbatch.

This caches the config per args object (_get_transformer_config(args)), so it is built once per training run and rebuilt only when the global args object is replaced. The fields get_batch reads (pipeline layout, MTP layer count, virtual pipeline size, tensor-parallel size and sequence parallelism, CP layouts) are fixed for the run.

The cache is keyed on the args object rather than being a bare lru_cache, so tests/unit_tests/data/test_get_batch.py, which builds fresh args per test and patches core_transformer_config_from_args, still gets a config built from each test's own args.

Testing

Not yet run; tests/unit_tests/data/test_get_batch.py is the relevant suite. Both changed files pass python -m py_compile.

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

[svcnvidia-nemo-ci]

Add regression coverage for cache reuse and invalidation

The caches added in pretrain_gpt.py and pretrain_hybrid.py have no accompanying regression test. The existing mocked packing case and pretraining case fetch only one batch per args object and do not assert constructor counts. Restoring construction on every microbatch would leave those assertions unchanged, allowing the repeated-warning bug to return undetected. This is a coverage gap, not an observed runtime failure in the proposed implementation.

Add a parameterized regression to test_get_batch.py. Exercise get_batch so the test protects both the helper and its call-site wiring. Repeated calls with the same args must reuse the config, while replacement with an equal-valued but distinct args object must rebuild it.

Suggested addition using the file’s existing imports; source-reasoned and not executed, so verify before adopting:

@pytest.mark.parametrize("script", [pretrain_gpt, pretrain_hybrid], ids=["gpt", "hybrid"])
def test_get_batch_caches_config_until_args_replaced(script, monkeypatch):
    args = SimpleNamespace(
        sequence_packing_scheduler="dp_balanced", dynamic_context_parallel=False
    )
    configs = [
        SimpleNamespace(
            virtual_pipeline_model_parallel_size=None,
            pipeline_model_parallel_layout=None,
            mtp_num_layers=None,
        )
        for _ in range(2)
    ]
    build_config = MagicMock(side_effect=configs)
    result = object()
    fetch_batch = MagicMock(return_value=result)
    monkeypatch.setattr(script, "_transformer_config_cache", None)
    monkeypatch.setattr(script, "get_args", lambda: args)
    monkeypatch.setattr(script, "core_transformer_config_from_args", build_config)
    monkeypatch.setattr(script, "mtp_on_this_rank_func", lambda **kwargs: False)
    monkeypatch.setattr(script, "get_batch_on_this_rank_for_sequence_packing", fetch_batch)
    pgs = ProcessGroupCollection(pp=object())

    for _ in range(2):
        assert script.get_batch(None, pg_collection=pgs) is result
        assert fetch_batch.call_args.kwargs["config"] is configs[0]
    build_config.assert_called_once_with(args)

    previous_args = args
    args = SimpleNamespace(**vars(previous_args))
    assert args == previous_args and args is not previous_args
    for _ in range(2):
        assert script.get_batch(None, pg_collection=pgs) is result
        assert fetch_batch.call_args.kwargs["config"] is configs[1]
    assert build_config.call_count == 2
    assert build_config.call_args.args[0] is args

Acceptance and preservation: for each entrypoint, two batches with the original args construct once; two further batches with replacement args bring the total to exactly two and receive the replacement config. Both phases must return the downstream result unchanged. Retain the existing packed-entrypoint assertions that explicit process groups and config are forwarded by identity, MTP selection respects vp_stage, and only Hybrid requests ContextParallelBatch.

Completed the light review of both changed files and their directly affected args-lifecycle, MTP, packing, and CP paths. No production correctness regression identified. Tests were not executed.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[Phlip79]

/review

[svcnvidia-nemo-ci]

LGTM. The prior coverage finding is resolved by the new parameterized regression. For both entrypoints, it checks one config construction across repeated calls, rebuilding when args is replaced by an equal-valued but distinct object, config forwarding by identity, and unchanged downstream return identity.

Completed the light follow-up against the previous reviewed head, including the new test’s integration and isolation. The existing process-group, virtual-stage MTP, and Hybrid return-interface assertions remain intact. No unresolved prior finding or qualifying introduced regression identified.

Source inspection only; tests were not executed.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a findings-only review. Request a full review with /review reset=true.

[nemo-automation-bot]

🔄 Merge queue validation started!

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

[nemo-automation-bot]

🔄 Merge queue validation started!

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