[linmuchuiyang]

What does this PR do ?

Adds policy.automodel_cfg.distributed_timeout_minutes (and value.automodel_cfg.… for PPO), the Automodel-path counterpart of the megatron_cfg.distributed_timeout_minutes knob that #4481 added for the Megatron workers. setup_distributed passes timeout=datetime.timedelta(minutes=…) to torch.distributed.init_process_group, so the default group and every FSDP2 / expert-parallel sub-group derived from it use the configured collective timeout instead of torch's default (10 minutes for NCCL). Unset / null keeps the default; non-positive values are rejected with the same message pattern as #4481.

Why: on the Automodel/DTensor path the first training step JIT-compiles kernels (HybridEP / DeepEP dispatch, TileLang attention) and loads large MoE checkpoints; on cold nodes those phases can keep ranks apart for more than 10 minutes, and the only way to survive it today is an environment-variable patch of setup.py. The knob was needed to run the DeepSeek V4.1 Flash recipe (deepseek-v4.1-support) on H100 (NeMo RL #4246, data-point comment) and GLM-5.2 / DeepSeek-V4 Automodel fine-tuning on H100 and GB200.

Changelog

• nemo_rl/models/automodel/setup.py: _distributed_timeout_minutes() (validation) and _init_process_group() (backend + optional timeout); setup_distributed calls it.
• nemo_rl/models/policy/__init__.py: AutomodelConfig.distributed_timeout_minutes: NotRequired[float | None].
• examples/configs/{dpo,grpo_math_1B,grpo_math_1B_megatron,ppo_math_1B,rm,sft,vlm_mpo}.yaml: distributed_timeout_minutes: null with a comment under every automodel_cfg block (same files fix(megatron): make the NCCL collective timeout configurable #4481 annotated for megatron_cfg, plus rm.yaml).
• tests/unit/models/automodel/test_automodel_setup.py: TestDistributedTimeoutMinutes (forwarded with and without CPU offload backend, unset keeps the default, non-positive rejected).

Issues

Related: #4481 (Megatron counterpart), #2229 / #2390 (exposing hard-coded infrastructure parameters), #4246.

Usage

policy:
  automodel_cfg:
    distributed_timeout_minutes: 45   # null = torch default (10 min for NCCL)

Before your PR is "Ready for review"

• Unit tests added (tests/unit/models/automodel/test_automodel_setup.py::TestDistributedTimeoutMinutes). Locally the two helpers were exercised with the same four cases in a standalone harness (7 passed); the module-level test needs main's full environment (our container is the older deepseek-v4.1-support stack), so please run it in CI
• Signed-off commits (DCO)

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

[sharonyu-115]

Thanks for contributing this configurable timeout for Automodel training.

This PR adds automodel_cfg.distributed_timeout_minutes, including support for PPO value workers, to accommodate long initialization, kernel compilation, and checkpoint-loading delays. Unset or null preserves PyTorch’s default; non-positive values are rejected.

Three P2 issues need attention:

1. C1: Forward the timeout to newly created training subgroups and add a regression test for that forwarding.
2. C2: Update the five reference YAML snapshots affected by the new default.
3. C3: Mark the new timeout tests so the Automodel CI lane selects them.

Validation used isolated CPU harnesses and pinned upstream source. The config failures and test-selection gap were reproduced; full repository tests and GPU/NCCL execution were not run.

Generated by Codex

[sharonyu-115]

C1 · 2 action items.
TL;DR — newly created training subgroups retain the backend timeout despite the new setting.

The default-group initialization receives the timeout, but the mesh construction call omits it. Pinned Automodel documents that mesh subgroups do not inherit this value, so those collectives can still time out during the long first-step delays this option targets.

AI-1

Pass timeout_minutes= with the result of _distributed_timeout_minutes(config) to the mesh construction call. Automodel already forwards this parameter. The call is outside this diff, so an in-place suggestion cannot cover it.

AI-2

Extend the existing distributed-setup test to check that an explicit timeout reaches both default-group initialization and mesh construction, and that absent/null settings preserve default behavior. Require the explicit-timeout case to fail before the fix and pass afterward. The new helper-only tests do not exercise mesh construction. The existing test is outside this diff, so an in-place suggestion cannot cover the edit.

[sharonyu-115]

C2 · 1 action item.
TL;DR — the new YAML defaults break five reference-config checks.

The added key is absent from the reference snapshots. An isolated run of the existing parity check, using the real config loader, reproduced five failures for missing timeout keys.

AI-1

Add the policy Automodel timeout default (null) to dpo.yaml, grpo_math_1B.yaml, rm.yaml, sft.yaml, and ppo_math_1B_megatron.yaml. Add the same default to PPO's value Automodel configuration. These snapshots are outside the diff, so an in-place suggestion cannot cover this change.

[sharonyu-115]

C3 · 1 action item.
TL;DR — the Automodel CI lane filters out all seven new timeout cases.

The new test class lacks the marker required by the Automodel lane and its collection filter. An isolated run with that filter selected none; adding the marker ran all seven successfully.

AI-1

Add the class marker:

Suggested change

	class TestDistributedTimeoutMinutes:
	@pytest.mark.automodel
	class TestDistributedTimeoutMinutes:

[linmuchuiyang]

Thanks for the careful review — all three points were right, and they are addressed in bd9d948:

C1 — timeout now reaches the mesh sub-groups. setup_distributed passes timeout_minutes=_distributed_timeout_minutes(config) to MeshContext.build, so the per-axis NCCL groups get the configured timeout through Automodel's backend override instead of keeping torch's default. The _init_process_group docstring that claimed the sub-groups inherit the default group's timeout was wrong and is corrected. TestSetupDistributed gained test_setup_distributed_forwards_timeout_to_process_group_and_mesh (parametrized explicit / null / absent), which asserts that an explicit value reaches both init_process_group and MeshContext.build and that null / absent keep torch's default for both. Against the previous commit the explicit case fails (assert None == 45) and the other two pass; with the fix all three pass.

C2 — reference snapshots updated. automodel_cfg.distributed_timeout_minutes: null is added to dpo.yaml, grpo_math_1B.yaml, rm.yaml, sft.yaml, and both policy and value in ppo_math_1B_megatron.yaml. test_reference_configs_up_to_date now passes for all seven reference files (it failed for exactly those five before).

C3 — tests are in the Automodel lane. TestDistributedTimeoutMinutes is marked @pytest.mark.automodel; with --automodel-only the seven cases are collected (previously 0). The import block of the test module is also re-sorted, which ruff's I001 had flagged on the earlier commit.

Validation (CPU, Python 3.13 venv with torch 2.13 and the pinned Automodel submodule, init_process_group / MeshContext mocked as in the existing tests):

• pytest tests/unit/models/automodel/test_automodel_setup.py -k "TestSetupDistributed or TestDistributedTimeoutMinutes" --automodel-only → 15 passed
• pytest tests/unit/test_config_v2.py -k test_reference_configs_up_to_date → 7 passed (5 failed on the previous commit)
• full tests/unit/models/automodel/test_automodel_setup.py --automodel-only → 113 passed
• ruff format --check, ruff check, ruff check --select I on the two changed Python files → clean

GPU / NCCL execution was not run here; the sub-group timeout path is exercised by Automodel's own _nccl_backend_override.

[sharonyu-115]

@linmuchuiyang Thanks for addressing the review comments. A further review exposed another issue:

Blocked pending Automodel #4601. The new mesh-timeout forwarding exposes a reproducible CPU-offload failure during GRPO initialization. Once the upstream fix lands, this PR can proceed with an updated Automodel pin and regression coverage.

[sharonyu-115]

The mesh-timeout forwarding depends on fixing Automodel #4601. With CPU offload enabled, an explicit timeout currently creates NCCL-only mesh groups, causing prepare_refit_info() → DTensor.full_tensor() to fail with No backend type associated with device type cpu.

Reproduced with Qwen2.5-1.5B on two H100s: timeout 45 fails before training; timeout null completes one step.

Please hold this PR until the upstream fix lands, then update the Automodel pin and add a two-GPU regression test covering CPU offload with an explicit timeout.

[linmuchuiyang]

Agreed, thanks for catching it — forwarding the timeout to the mesh exposes _nccl_backend_override dropping the Gloo half of cuda:nccl,cpu:gloo, which breaks the CPU-offload path that this very knob is meant to help. Holding this PR until Automodel #4602 (the fix for #4601) lands. Once it is in, I will bump the Automodel pin here and add the two-GPU regression test covering CPU offload with an explicit timeout.

[linmuchuiyang]

Automodel #4602 is merged (d1af15d), thanks for pushing that through. On this PR:

Two-GPU regression test — added tests/unit/models/automodel/test_setup_distributed_gpu.py in 3fd4203. It spawns two ranks, runs the real setup_distributed with cpu_offload=True and distributed_timeout_minutes=45 (default group cuda:nccl,cpu:gloo), then checks that the default group and every mesh axis carry the 45-minute timeout on both the CPU and CUDA backends, that CPU and CUDA all-reduces work on every axis, and that a CPU-offloaded FSDP parameter sharded on device_mesh["dp"] reconstructs through full_tensor() — the prepare_refit_info() path from #4601. On two H100s with the current pin (b916107) the rank bodies fail with No backend type associated with device type cpu; with the #4602 change applied on top of the pin the test passes (CPU and CUDA all-reduces on every axis, both backends at 45 min, full_tensor() matches). The test is in the Automodel lane (@pytest.mark.automodel, under tests/unit/models/automodel/).

Automodel pin — I did not bump it in this PR. The pin would move from Automodel 0.6.1 to 0.8.0 (68 commits), and Automodel's own requirements moved with it: transformers==5.18.0 (NeMo RL caps at <=5.16.1), transformer-engine>=2.19.0 (this repo pins v2.14.1), quack-kernels==0.6.4, a new deep_ep revision and Python >= 3.11 — i.e. a relock of the automodel extra, not a submodule-only change. That seems better done by the maintainers in a dedicated bump, where its CI lanes can be watched. To keep this PR green until then, the new test skips with the reason "Automodel pin predates NVIDIA-NeMo/Automodel#4602" when _nccl_backend_override still hard-codes NCCL, and becomes active automatically on the bump; the guard can be deleted at that point. If you would rather have the bump in this PR anyway, tell me and I will do it, but I wanted to flag the scope first.