ZeRO-3 OPSD rollout deadlocks when data-parallel ranks finish generation at different lengths #8262

Summary

OPSD rollout can deadlock when Hugging Face generate() runs on a ZeRO-3-partitioned student and data-parallel ranks finish at different token positions.

One rank may emit EOS and leave generation while other ranks enter the next decode forward. The latter ranks then wait in the next ZeRO-3 parameter all-gather, while the completed rank has already entered teacher forward and waits in a different collective.

This differs from #7844: that report had all ranks hanging in engine.step(). Here the ranks diverge before backward/step, during student rollout and teacher scoring.

Environment

Reproduced independently on:

• 8 x NVIDIA H200, CUDA 12.6, PyTorch 2.7.0, DeepSpeed cf443004
• 8 x AMD MI250 ROCm devices, ROCm 6.4.1, PyTorch 2.9.1, DeepSpeed 22969fa7
• DeepSpeedExamples OPSD trainer with ZeRO stage 3
• Qwen2.5-0.5B-Instruct student
• Qwen2.5-Math-7B-Instruct teacher
• Different prompts on each data-parallel rank

Steps to reproduce

1. Use the DeepSpeedExamples OPSD Hybrid Engine configuration with ZeRO stage 3.
2. Run with multiple data-parallel ranks and different prompts per rank.
3. Let module.generate() use the tokenizer's default EOS stopping criterion.
4. Run at least two rollout/training steps.

Observed behavior

After a completed training step, the job permanently stops making progress. Per-rank stack traces show:

• one rank in teacher forward: reset_step -> broadcast;
• the other ranks in student decode: fetch_sub_module -> all_gather_into_tensor.

The divergent collective sequence is:

rank 1:   student decode N -> EOS -> teacher forward broadcast
others:   student decode N -> student decode N+1 parameter all-gather

Padding generated sequences afterward cannot help because the mismatch occurs inside autoregressive generation.

Root cause

Under ZeRO-3, each decode forward gathers partitioned parameters. Data-parallel ranks must execute the same number and order of forwards. Default per-sequence EOS termination violates that requirement when prompts produce different response lengths.

Verified workaround

Force every rank to execute exactly max_new_tokens decode iterations:

output_ids = module.generate(
    ...,
    eos_token_id=None,
)

The two-step eight-GPU reproducer completed on both H200 and MI250 after this change. A full 8-device MI250 prompt epoch and a subsequent 200-step run also completed.

Expected behavior

The Hybrid Engine rollout should preserve an identical ZeRO-3 collective sequence across data-parallel ranks. At minimum, the HF generation path should disable rank-local early termination. A future optimized implementation could mask finished sequences while keeping all ranks in the decode loop.