HybridEngineRollout: synchronized early stopping across ZeRO-3 data-parallel ranks

[delock]

Description

Follow-up to #8264 (tracked there in the review discussion, agreed to be handled separately).

To keep ZeRO-3 data-parallel ranks in lockstep during rollout, HybridEngineRollout.generate currently passes eos_token_id=None to the underlying generate call and restores EOS/padding semantics after the fact via _pad_after_eos. This fixes the deadlock, but as a consequence every rank always decodes the full max_new_tokens, even when every sequence in the global batch has already produced its EOS. The wasted compute is significant when typical responses are much shorter than the max response length.

Proposal

All ranks should collectively detect the early-stop condition — every sequence on every rank's batch has produced EOS — and stop the decode loop together, while preserving the ZeRO-3 invariant that all ranks execute the same number of parameter-gather collectives.

Concretely, the stop decision must be global, not local:

• Periodically (every decode step, or every K steps) all-reduce a local "unfinished sequences remain" flag across data-parallel ranks.
• All ranks stop only when the aggregated flag says no rank has unfinished sequences. Any rank exiting the loop based on local knowledge alone would reintroduce the deadlock Fix ZeRO-3 synchronization during OPSD rollout #8264 fixed.

This can be implemented either as a custom decode loop with an explicit sync point, or as a StoppingCriteria that aggregates the decision globally before returning True on all ranks.

Workload context (from #8264)

• Main 8-GPU OPSD run: micro_batch_size_per_gpu=1, gradient_accumulation_steps=1, n_samples_per_prompt=1 → 8 prompts / 8 rollouts per training iteration across 8 DP ranks.
• Verification runs use max response length 64; full reproduction uses 1024. Responses are typically far shorter than the max, so decoding to max_new_tokens on every rank wastes a large fraction of rollout compute.

Design consideration

With micro_batch_size_per_gpu=1 there is exactly one sequence per GPU, so per-slot work reclamation (continuous batching) has nothing to reclaim locally — synchronized early stopping is the whole win for this configuration. Continuous batching becomes interesting for larger per-rank batches.

Please implement the stop-check as a reusable synchronized control-flow primitive (e.g. a periodic "global work remaining" rendezvous) rather than a one-off stop hack: continuous batching generalizes the same rendezvous from a "stop/continue" boolean into a "rebuild the batch composition" decision, so an early-stopping fix designed with that shape becomes a stepping stone instead of rework.

@LiRunGuo — you kindly offered to own this follow-up in #8264. GitHub won't let me assign you until you've commented on this issue (non-collaborator limitation), so please drop a comment here and I'll assign it to you.

[FU-max-boop]

I audited this against current master (715965e) and #8264 head (687cd49).

Before adding a custom StoppingCriteria, could we first freeze the supported process-group topology? Transformers already provides the required control flow through explicit synced_gpus=True in the versions DeepSpeed currently supports. In both 4.32.1 and 4.51.3, every peer participates in a control all-reduce; a locally finished peer still executes the next model forward while another peer remains active; and all peers exit only when the global unfinished count reaches zero.

For the reported OPSD workload, if all eight ranks are one pure data-parallel world and all enter the same generate() call, explicitly passing synced_gpus=True appears to be the smallest first implementation to validate. It also preserves native EOS/padding semantics instead of combining eos_token_id=None with a full-max_new_tokens decode. Explicitly setting the flag matters because Transformers' automatic ZeRO-3 detection is mediated by its own DeepSpeed config weak reference and need not be active for a model initialized directly by DeepSpeed.

The important limitation is that Transformers hard-codes the default world group. That is not equivalent to the ZeRO partition group in every DeepSpeed topology. Standard ZeRO-3 receives engine.seq_data_parallel_group; MiCS then replaces optimizer.dp_process_group with its parameter-shard group. A generalized primitive should therefore rendezvous over the actual ZeRO parameter-gather group, not blindly over engine.data_parallel_group or world. This also makes pure-DP support versus DP x TP/PP or independent rollout subgroups an explicit contract rather than an accidental property.

I suggest a deterministic two-rank CPU/Gloo regression with scripted EOS at decode steps 2 and 5 and max_new_tokens=8. Current master should expose divergent forward counts [2, 5]; #8264's fixed-length workaround gives [8, 8]; synchronized early stopping should give [5, 5] while retaining local response lengths [2, 5], the first EOS, and no valid post-EOS tokens. A world=4 / two-DP2-group sibling test should be required before claiming subgroup support.

One additional scope boundary: use_graph_capture=True bypasses HF generate() entirely, so synced_gpus=True would not cover that path. It should either be explicitly out of scope here or use the same group-aware rendezvous.

I will not open a competing PR while @LiRunGuo owns this follow-up, but I am happy to validate an exact head with the Gloo harness and a DP/TP GPU matrix.

Local control-flow result. I have now run that no-download probe with Python 3.11, PyTorch 2.7.0, Transformers 4.51.3, and two Gloo ranks on CPU. It instantiates a tiny GPT-2 config in memory and forces rank 0 / rank 1 to emit EOS at decode steps 2 / 5:

mode	per-rank forward counts	per-rank response lengths
normal EOS, default synchronization	[2, 5]	[2, 5]
normal EOS, synced_gpus=True	[5, 5]	[2, 5]
eos_token_id=None	[8, 8]	[8, 8]

The synchronized case retained response IDs [3, 2] and [3, 3, 3, 3, 2], so the first EOS and local response lengths were preserved while forward counts stayed in lockstep. This validates the Transformers control-flow contract only; it is not yet a ZeRO-3/NCCL result. I can turn the probe into the repository's DistributedTest form when there is an exact implementation head to review.

[FU-max-boop]

I completed the GPU follow-up promised above against exact DeepSpeed master
715965e027894a2e72ac2e27f2daed2c599e99f0. This is a harness-only
validation result, not an upstream implementation or a generalized topology
claim.

The run used one host with 2 x A100, Python 3.13.13,
PyTorch 2.8.0+cu128, CUDA 12.8, NCCL 2.27.3, and Transformers 4.51.3. The
harness instantiated DeepSpeedHybridEngine with ZeRO-3 and ran two pure
data-parallel ranks (TP=1), scripted local EOS at decode steps 2 and 5, and
max_new_tokens=8. A separate two-rank NCCL smoke test passed before the
matrix.

Before generation, it checked that the optimizer ZeRO group and every
nonpersistent parameter's ds_process_group were equivalent to the default
world for this exact two-rank topology. The rank-0 ready record is emitted
only after those checks and a two-rank barrier.

mode	exit	forward counts	response lengths
current master, native local EOS	124 after the 180 s outer timeout	no completed result	no completed result
harness-only synced_gpus=True	0	[5, 5]	[2, 5]
fixed-length eos_token_id=None control	0	[8, 8]	[8, 8]

For current master, the job reached ready and then stopped making progress.
The scheduled 120 s faulthandler dump captured execution in
partitioned_param_coordinator._fetch_sub_module_impl ->
torch.cuda.streams.synchronize; the outer timeout terminated the run at
180 s with exit code 124. In contrast, explicit synced_gpus=True completed
with both ranks executing five forwards while retaining the unequal local
response lengths.

The synchronized result retained response IDs [3, 2] and
[3, 3, 3, 3, 2], with all valid-token attention entries set for lengths 2
and 5. The fixed-length control completed too, but returned length 8 on both
ranks, retained tokens after each rank's first EOS, and marked all eight
positions valid. So the latter is a liveness workaround, not the same output
semantics.

This is direct ZeRO-3/NCCL evidence for the narrow pure-DP case where the
default world is the actual ZeRO parameter-gather cohort. It should not be
generalized to DP x TP/PP, MiCS, independent rollout subgroups, or graph
capture. In those cases, a reusable rendezvous still needs to be defined over
the actual ZeRO gather group rather than assuming the default world. I am
also not making a latency or throughput claim from this small deterministic
probe; the forward counts are the measured control-flow result.

The copied-off receipt bundle is anchored by RUN_MANIFEST.sha256 with
SHA-256
8d649d73328bbb9240cb750720b2da1c5b05b66845f98ac3980d8fbbab2efc58,
and an independent verifier returned POST_RUN_ACCEPTANCE=PASS.

There is still no exact owner implementation head on this issue, so I did not
run the native or native-pad-eos acceptance modes. I will not open a
competing PR; once an owner head is available, I can rerun the same two-A100
matrix against that exact SHA, including the pad_token_id == eos_token_id
first-EOS attention check.

[LiRunGuo]

Thanks for creating the follow-up issue. I would be happy to work on it. I will implement the synchronized early-stop decision as a reusable global-work-remaining rendezvous and preserve the EOS/padding semantics introduced in #8264. Please feel free to assign the issue to me. I will keep the implementation in a separate PR after #8264 is merged.

[LiRunGuo]

Validation update on AMD MI250 (job 392426): the branch fix/hybrid-rollout-global-early-stop was tested at commit c317430e3b0ff7e507946b1a8d2aa909565c8d84, based on the latest origin/master. The job completed successfully (exit code 0:0). Results: (1) the 4-rank Gloo subgroup test passed, confirming that the all-reduce uses the correct ZeRO-3 data-parallel subgroup; (2) the 2-rank MI250 scripted-EOS GPU test passed, with ranks generating EOS at steps 2 and 5 and final forward counts [5, 5], rather than [8, 8]; no rollout deadlock occurred; (3) the 2-GPU Qwen2.5-0.5B/7B 2-step ZeRO-3 OPSD smoke validation completed with finite losses (0.6831 and 0.1576). The Qwen run still reports inference_containers=0 and therefore validates the ZeRO-3 + Hugging Face generate fallback path, not Qwen Hybrid Engine acceleration. The synchronized early-stop fix itself has now been validated on AMD MI250.

[LiRunGuo]

Thanks for the detailed CPU/Gloo and A100 validation. I have opened #8390 with the owner implementation: #8390

The branch is rebased on current master (56de57058) and the exact head is 020a184db69e2f7aead70fb03e0144f09675a6dc. It rendezvous over optimizer.dp_process_group, preserves first-EOS/padding semantics, and keeps graph capture explicitly out of scope. On AMD MI300X, the 4-rank/two-subgroup Gloo test and the 2-rank scripted-EOS ZeRO-3 GPU test passed with forward counts [5, 5]; a 2-GPU Qwen OPSD ZeRO-3 smoke test also completed.

@FU-max-boop, if convenient, your A100 harness validation against this exact head would be very helpful. Thank you!