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.

[delock]

Hi @LiRunGuo thanks for the followup. Can you resolve merge conflicts? Thanks!

I'll create an issue for early stopping and assign to you, let me know if it works.

[LiRunGuo]

Thanks for the reminder. I have rebased the PR onto the latest master and resolved the conflicts with the newly added HybridEngineRollout profiling changes. The updated branch preserves both profiling and synchronized fixed-length decoding with post-EOS padding. I kept synchronized early stopping out of this PR and will address it separately in #8321.

Local validation after the rebase:

• pytest -q tests/unit/runtime/rollout/test_hybrid_engine_rollout.py: 18 passed
• targeted pre-commit hooks: passed
• git diff --check: passed

[LiRunGuo]

The PR is now mergeable, approved, and all available checks/DCO have passed. I attempted to merge it, but this repository uses a maintainer-controlled merge queue and my account does not have permission to add it. Could you please add #8264 to the merge queue when convenient? Thanks!

[LiRunGuo]

Hi @delock, I checked the failed merge-queue run (32981748015). The only failures were two unrelated FusedAdam numerical-reference cases:

• test_fused_adam_matches_reference[fp32-adam]
• test_fused_adam_matches_reference[fp16-adam]

All HybridEngineRollout tests passed in that run. The temporary merge commit only changes deepspeed/runtime/rollout/hybrid_engine_rollout.py and its unit-test file; it does not touch FusedAdam or test_adamw.py. The full CPU unit-test job also passed previously on the PR head.

This therefore appears to be an unrelated intermittent CI/test failure rather than a regression from #8264. Could you please rerun the failed merge-group job or add the PR back to the merge queue? Thanks!

[LiRunGuo]

Hi @delock, thank you for requeueing the PR. I checked the second merge-group commit (49aba2d3741bde3bc204a39bd08e8619fb815005), and it appears to be stuck because of GitHub Actions startup/scheduling failures rather than a code or test failure:

• cpu-torch-latest has remained queued without creating any jobs or steps.
• DCO / required and Formatting ended with startup_failure.
• Python was marked failed before its queued jobs started and has no test output.
• NV precompile completed successfully.

Because required workflows for this merge group have already failed, the still-queued CPU workflow cannot make the merge complete. Could you please cancel/ignore this stale merge group and add #8264 to the merge queue again? No code changes appear necessary. Thanks!

[LiRunGuo]

Hi @delock, I have rebased #8264 onto the latest master (32e301ffa) and resolved the overlap with the newly merged shared-prefill implementation from #8296.

The updated code preserves the complete shared-prefill hook lifecycle and places the ZeRO-3-safe eos_token_id=None argument inside the shared generation path. Post-EOS padding is applied after generation/profiling, and both the shared-prefill helper and EOS-padding helper/tests are retained.

Local validation after the rebase:

• pytest -q tests/unit/runtime/rollout/test_hybrid_engine_rollout.py: 23 passed, 1 CUDA-only test skipped
• targeted pre-commit hooks: passed
• git diff --check: passed
• range-diff reviewed: the two original fixes are preserved, with only the expected shared-prefill context changes

The branch has been updated with --force-with-lease. Once the refreshed CI passes, could you please add #8264 back to the merge queue? Thanks!