Synchronize ZeRO-3 rollout early stopping

[LiRunGuo]

Summary

• synchronize EOS stopping decisions across the ZeRO-3 parameter-gather process group
• stop at the latest rank-local EOS instead of always decoding to max_new_tokens
• retain each sequence's first EOS and pad later response positions
• keep non-ZeRO-3 generation on the native EOS path

The rendezvous uses optimizer.dp_process_group, which is the group used by ZeRO-3 parameter gathers, rather than assuming the default world group. CUDA/HIP graph-capture rollout remains outside this change because it bypasses Hugging Face generate().

Closes #8321.

Validation

• DeepSpeed pre-commit checks: passed
• 4-rank CPU/Gloo two-subgroup test: passed
• 2-rank AMD MI300X scripted-EOS ZeRO-3 test: ranks reaching EOS at decode steps 2 and 5 both executed 5 forwards ([5, 5]), while preserving local response lengths
• 2-GPU Qwen2.5-0.5B/7B ZeRO-3 OPSD smoke test: completed two steps with finite losses (0.6784, 0.1554)

Exact validation head: 020a184db69e2f7aead70fb03e0144f09675a6dc.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 020a184db6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Preserve compatibility with legacy stopping criteria

On installations using the supported Transformers 4.32.x range (requirements/requirements-inf.txt permits >=4.32.1), the legacy StoppingCriteriaList.__call__ applies Python any() to each criterion result. This batch-shaped tensor therefore raises “Boolean value of Tensor with more than one value is ambiguous” whenever a ZeRO-3 rollout has more than one sequence, so the new synchronized generation path crashes; either provide version-compatible scalar semantics for those releases or raise the Transformers dependency floor.

Useful? React with 👍 / 👎.

[LiRunGuo]

Fixed in d8a5b17. The stopping criterion now returns a scalar Python bool because the decision is already rank-global. This preserves the intended all-local-sequences behavior, is compatible with Transformers 4.32 (which applies Python any() to criterion results), and also broadcasts correctly in newer StoppingCriteriaList implementations. I added a two-sequence regression assertion covering the 4.32-style any() path.

[chatgpt-codex-connector]

Replace the direct torch distributed import

This directly imports PyTorch's distributed module under an alias, bypassing the repository's required deepspeed.comm abstraction and the textual check-torchdist hook. Rework the Gloo setup through the DeepSpeed communication layer rather than adding this unapproved direct dependency.

AGENTS.md reference: AGENTS.md:L10-L11

Useful? React with 👍 / 👎.

[LiRunGuo]

Fixed in d8a5b17. I removed the direct torch.distributed import and initialize/query process groups only through deepspeed.comm. Since this is specifically a CPU/Gloo subgroup topology test, the test temporarily mocks DeepSpeed's accelerator as CPU during deepspeed.init_distributed(); this avoids binding Gloo to a visible GPU on accelerator hosts while retaining the DeepSpeed communication abstraction. The 4-rank subgroup test and the 2-GPU scripted-EOS ZeRO-3 test both pass in Slurm job 400417.

[FU-max-boop]

@LiRunGuo I found one remaining subgroup liveness issue at exact head d8a5b1779eafcbf15ab1a14e04aa83bf04f75c3e, so I am holding the requested A100 run for now.

HybridEngineRollout.generate() installs the new optimizer.dp_process_group-aware stopping criterion but leaves synced_gpus unset. When the Transformers HfDeepSpeedConfig integration can see ZeRO-3, Transformers auto-enables synced_gpus from the default world size in both 4.32.1 and 4.51.3. Its decode loop then performs an ungrouped/default-world all-reduce in 4.32.1 and 4.51.3. If only one DP subgroup enters rollout, that first world collective can wait forever for non-participating ranks; concurrently active subgroups are also unnecessarily coupled through WORLD.

The current 4-rank Gloo test calls GlobalWorkRemaining directly, so it does not enter Transformers generate(). The 2-rank CUDA test uses the full WORLD and initializes with deepspeed.initialize() directly, which ordinarily does not establish the Transformers ZeRO-3 weakref, so it likely does not trigger the implicit synced_gpus=True path either.

The narrow fix is to pass synced_gpus=False if zero3_sync is not None else None, preserving the Transformers default for non-ZeRO-3/FSDP paths. The custom criterion already gives every rank in the parameter-gather group the same stop decision. I would also add a kwargs assertion plus a generate-level regression with HF-visible ZeRO-3 and a proper subgroup/non-participating WORLD ranks. This is source-level and CPU/Gloo-testable; once a fixed head lands, NVIDIA validation can target that exact commit.