Fix ZeRO-3 synchronization during OPSD rollout

[LiRunGuo]

Summary

• disable rank-local EOS early exit in the OPSD Hybrid Engine Hugging Face rollout path
• keep ZeRO-3 decode forwards and parameter collectives in the same order on every data-parallel rank
• add a CPU unit test that verifies the collective-safe generation argument

Fixes #8262.

Why

ZeRO-3 gathers partitioned parameters during every autoregressive decode forward. If one data-parallel rank emits EOS and returns from generate() while another rank continues decoding, the continuing rank enters another parameter all-gather while the finished rank moves to a later collective. The job then deadlocks.

Passing eos_token_id=None makes every rank execute exactly max_new_tokens iterations. Output masking remains unchanged, so EOS and padding tokens are still excluded from downstream response loss as applicable.

Validation

• pytest -q tests/unit/runtime/rollout/test_hybrid_engine_rollout.py (10 passed)
• pre-commit run --files deepspeed/runtime/rollout/hybrid_engine_rollout.py tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
• 2-step distributed reproducer completed on 8 x H200 and 8 x MI250 after the fix
• full prompt epoch and a separate 200-step OPSD run completed on 8 x MI250

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[delock]

Hi @LiRunGuo thanks for your issue and fix. I have some comments:

1. The fix by setting eos_token_id=None will have the hybridenginerollout keep on generating tokens instead of padding beyond eos. This will cause a change of behavior. A proper fix would be apply the padding (detect first EOS token and pad the rest) to ensure output identical to before the fix.
2. It will help performance if all ranks can check whether an early stop condition (EOS had been generated on every batch) had been met and all ranks can stop rollout early (not necessary rollout to max length). We do not need to address this in this PR, but it would help if you can share your batch information (how many prompts per training iteration, how many rollouts per prompt), these information will help us prioritize our optimizations on HybridEngineRollout, thanks!

@PKUWZP for this discussion.

[LiRunGuo]

Thanks for pointing this out. I agree that setting eos_token_id=None is only a deadlock workaround and changes the originalgeneration behavior by disabling EOS termination.

I will revise the fix so that EOS detection and post-EOS padding are preserved. The rollout will keep all ranks in lockstep, retain the first EOS token, and pad all subsequent positions, matching the previous output semantics.

For the current OPSD configuration, the main 8-GPU run uses micro_batch_size_per_gpu=1, gradient_accumulation_steps=1, and n_samples_per_prompt=1. With 8 data-parallel ranks, this corresponds to 8 prompts and 8 rollouts per training iteration. The 2-GPU smoke test uses 2 prompts and 2 rollouts per iteration. The verification runs use a maximum response length of 64, while the full reproduction uses 1024.

I agree that synchronized early stopping across ranks would improve performance. I will treat that as a follow-up optimization after restoring the original EOS/padding behavior.

[delock]

Hi @LiRunGuo , I saw you force-pushed after comments but eos_token_id=None stays the same. A gentle remind if you intend to push your fix.

Thanks for your setup information, we definely should support such use case.