ZeroSyncGRPOTrainer - #6949
Draft
qgallouedec wants to merge 63 commits into
Draft
Conversation
… never-idle generation
…ompletion, batches from the ready pool
…sed_columns default False
…drift reconciliation with forking
…fault leaves room for training
…urrent batch's prompts
… tells you which lever applies
…ts own layers Under tp_size > 1 both generation and training issue collectives, and NCCL matches them by position across ranks, so a background generation thread and the training step pair a decode collective on one rank with a training collective on the other and deadlock. Stepping the engine at fixed points in the trainer's own forward and backward gives every rank the same order by construction, and generation keeps advancing through the step rather than waiting for it: generation/decode_steps reports two per layer. The engine decodes through a second view of the model, sharing every parameter, because continuous batching switches a model to a paged attention implementation that cannot serve the training forward. The vocabulary projection is replicated rather than split, which is what lets the chunked lm_head multiply by it, and it stays a DTensor so gradient clipping sees one kind of parameter. Measured on Qwen3-0.6B, GSM8K: tp_size 2 and 4 both train, 56 decode steps inside every training step, and the same run with the engine left in its background thread deadlocks.
The engine's thread is not a daemon, and Python joins those before it runs any atexit hook, so the hook registered for it never fired and a script that finished training never exited. Stopping it in `train` gets there first. The stop is a hard one: a flush waits for every rollout still in flight, and there are always some, since the trainer keeps `generation_ahead` batches queued, and none of them will be trained on now.
Data parallel was broken on more than one process: create_model_from_path defaults device_map="auto" on GPU, and this trainer only cleared it under tensor parallelism, so the documented `accelerate launch --num_processes N` failed with "You can't train a model that has been loaded with device_map='auto' in any distributed mode". Copy GRPO's guard verbatim. Also copy the model_init_kwargs dict rather than mutating the user's config, as GRPO does. The explanation of why the trainer steps the engine itself was wrong. It is not that NCCL matches collectives by position across ranks, and it is not unfixable from outside. NCCL requires every rank to issue the operations on its communicators in the same host-side order, and recommends "a deterministic order issued from a single host thread per-device"; the interleave is exactly that. Measured: two threads deadlock 4/4, one thread completes 4/4, and NCCL_LAUNCH_ORDER_IMPLICIT does not help because it derives the order from host issue order.
Capturing the cuda graphs submits the engine's own warmup requests, whose results come back through
the same queue as rollouts. `_drain` popped every result id unconditionally, so under tensor
parallelism, where the trainer owns the engine and drains it itself, the run died with
KeyError: '__warmup_DECODING_0__'. Cuda graphs are on by default here (the attention implementation
needs no mask, and transformers' heuristic enables them in that case), so every tp_size > 1 run so
far had to pass use_cuda_graph=False, and every number measured that way had graphs off.
Data parallel is unaffected: there the engine's own background thread consumes those results.
Measured on 2 x H100 with Qwen3-0.6B, GSM8K, 128-token completions, after the fix:
tp_size=2 5.13 s/step -> 2.64 s/step
data parallel (2 processes) 3.56 s/step -> 2.20 s/step
ZS_STEPS_PER_BOUNDARY and its inline import were scaffolding left in the working tree and swept into the previous commit. One decode step per layer boundary is the intended behaviour.
Two changes, one following from the other.
The trainer no longer interleaves generation through the training step. It drove a decode step at
every layer boundary, from forward and backward hooks; it now runs the engine between its own steps.
Measured on 2 x H100 with cuda graphs on, interleaving is worth nothing on short completions (1.55
against 1.54 s/step) and about 7% on long ones (2.28 against 2.44 with 512-token completions and 8
rollouts per prompt), and it makes the memory below impossible, since generation is then never idle.
With generation quiescent while the forward and backward run, its KV cache does not have to stay
resident. `release_kv_cache_during_step` frees it for the step and takes it back before generating
again, using the new release_memory/restore_memory in transformers. Live rollouts are copied to host
memory and back, so the cost follows the live KV rather than the pool size.
Measured, tp_size=2, Qwen3-0.6B, GSM8K, 128-token completions:
release off 1.44 s/step
release on 1.43 s/step, handing 5.36 GiB per step to the training step
Requires use_async_batching=False, which costs nothing here (1.52 against 1.53 s/step), and raises if
set with tp_size=1, where generation runs in its own thread and is never quiescent.
A batch still in flight cannot have its cache taken from under it, and transformers refuses that combination. Rather than make the user set the flag and hit the error mid-training, set it here. It costs nothing anyway: 1.52 against 1.53 s/step with and without async batching.
Above roughly 128-token completions a run with it enabled dies with 'probability tensor contains inf, nan or element < 0': the generation state does not survive the release and restore. Isolated to the completion length (128 tokens finishes with 4 or 8 rollouts, 256 fails with 4), not the number of rollouts, and not the cuda graphs, since it fails with them off. The option is off by default.
The corruption above 128-token completions was in transformers: releasing the cache destroyed what the blocks contained while the block manager went on advertising them by content hash. Verified by comparing greedy completions token for token against a run that never releases, up to 512-token completions releasing every step. At 512-token completions with 8 rollouts per prompt it is free: 2.52 and 2.40 s/step with it, against 2.51 and 2.45 without, all runs completing.
Data parallel attached the manager to the training model itself, so the training forward went through the paged attention implementation with no cache. That only worked because of a redirect in transformers that #48297 replaced with a raise, so it breaks as soon as this trainer runs against a current transformers. Tensor parallelism already avoided this by giving the engine a second view over the same parameters. Use it in both cases: the attention implementation is a setting on the config, shared by every module and every thread, so it cannot be flipped around each forward, and the view means it never has to be. The tensor-parallel transforms in the view are guarded by DTensor checks and do not fire without one. Verified: data parallel on two processes 2.40 s/step, tensor parallel with the KV cache released each step 2.44 s/step, both with the parameters changing.
Two places asked whether a parameter was a DTensor by looking for to_local or placements. Both are type checks in disguise; isinstance says what is meant and matches the repository's guidance against hasattr and getattr.
The backward multiplied fp32 grad_logits against .float() copies of the weight and hidden states, which lands on cutlass SIMT sgemm kernels (~55 TFLOP/s on H100 against 989 for bf16 tensor cores). The fp32 was there for accumulation, which the add_ into the pre-allocated fp32 buffers already provides, so the GEMMs now run in the model dtype and only the accumulation stays fp32. Measured on Qwen3-4B tp=4, B=16 x L=330: backward 550 -> 297 ms. The 98 chunked lm_head tests pass, including the bf16 ones.
With `packed_training=True` the training batch packs samples first-fit by decreasing length into rows of at least 4096 tokens. Position ids restart at every sample and no attention mask is passed, so the model builds the block-diagonal mask itself (transformers' packed-sequence detection); with `attn_implementation="flex_attention"` the attention is block-sparse and the pad-token compute disappears. Advantages and behavior logprobs become per-token, since one row holds several samples; the loss's global shift never crosses a seam because a sample's first token is always context. Verified on a tiny model that packed per-token logprobs match the padded layout exactly (5e-4, fp32). On the Qwen3-4B tp=4 GSM8K bench at batch 128: 3.73 to 3.42 s/step, MFU 8.3 to 9.2%.
…ong batches into OOMs Length-matched batch selection helps the padded path (padding follows the longest row), but under packed training it concentrated the near-cap samples into single batches: 256 samples all around 650 tokens is over twice the activation memory of a mixed draw, and whether the first primed sample was long decided whether step one fit, which showed up as the same config OOMing or passing run to run. Packing already removes the padding, so it now draws in arrival order and keeps the length mix.
…ce is written once
- packed_training defaults attn to flex_attention (flash attention for hybrid models, whose kernels read sample boundaries from the varlen kwargs) - hybrid models pack all samples in one flat unpadded row and pass cu_seq_lens_q/k + max_length_q/k through compute_loss - the generation view maps each submodule to its own copied config via the deepcopy memo (the top-level config clobbered text_config on composite models) - _tokenize_conversation uses the batched apply_chat_template call and takes row 0, like GRPO (processors return nested input_ids otherwise) Smoked on Qwen/Qwen3.5-2B, 1 GPU: generation through the CB linear-attention path, packed fla varlen training forward, clip_ratio ~0.001-0.003 (training and engine logprobs agree across the two code paths), grad_norm 11.9/13.4 with a varying reward. Requires flash-linear-attention.
…allel plan Qwen3.5 keeps lm_head out of its tp_plan, so the weight is already a full tensor and full_tensor() raised. Take the local data in that case, and use the model's device mesh for both embeddings.
HotpotQA with a keyword-search tool over the dataset's own paragraphs. The trainer's tool loop resubmits each turn, so the continuous batching engine serves the repeated prefix from cache, and Qwen3.5's hybrid attention keeps 5x less KV per token than a dense model of similar size, which is what pays for concurrent rollouts when generation and training share the GPU. Smoked on Qwen3.5-2B, 1 GPU, 6 steps: 2.47 search calls per question, 0.28 answer F1.
Two things kept the search example single-GPU, neither of them the tool loop: - Its search tool ranked articles out of a set, whose iteration order over strings differs between processes, so ranks retrieved different articles for the same query and their requests diverged. Sorted postings and an explicit tie-break make it deterministic. - Parameters the tp_plan leaves out (norms, and the gated delta net's convolution and gates) load as plain tensors, which fused optimizers and gradient clipping refuse to mix with DTensors. They are now replicated like the embeddings, and their modules compute on the local view of their whole subtree: the gated delta net reads self.conv1d.weight itself, so unwrapping only the convolution never fires, and gradient checkpointing replays these forwards in the backward pass. Verified on Qwen/Qwen3.5-9B, 4 GPUs at tp=4: 1.89 search calls per question, 0.76 answer F1.
… packing flash-linear-attention's kernels are the only ones that take sequence boundaries for the delta rule, and setting FLA_CACHE_MODE to skip its broken Triton autotuning makes them ignore those boundaries: measured on Qwen3.5-2B, a packed forward then differs from the per-sample forward by 3.30 against 4.3e-4 with autotuning left on. Padded batches are exact either way, so the example uses them, and packed training on a hybrid model without flash-linear-attention now raises instead of quietly running the recurrent state through the sample boundaries. Verified on Qwen/Qwen3.5-9B, 4 GPUs at tp=4: 1.74 search calls per question, 0.76 answer F1.
The sweep already replicates any plain parameter, which is what an untied input embedding is, so its own branch went away; and the two local-view helpers became one, since covering a subtree covers a lone module too. Verified on Qwen/Qwen3.5-9B, 4 GPUs at tp=4 with packed training: 2.09 search calls per question, 0.73 answer F1, 325 s for 3 steps against 372 s for 4 padded ones.
The world had to be exactly as wide as the model was split: tp_size == WORLD_SIZE. Ranks beyond that now form replicas that generate from their own prompts and sum their gradients over a group holding one rank per replica, so parameters stay whole and generation keeps reading them in place. Verified on Qwen3-0.6B, 4 GPUs at tp=2: the two replicas draw different prompts, and after training the ranks carrying the same shard agree exactly (4.638126 and -63.683460 on both replicas), where before they drifted apart (4.938084 against 7.568392). Needs two transformers fixes to build the mesh at all (tp over a subset of ranks) and, for now, a relaxed accelerate check: DistributedDataParallel cannot wrap a tensor parallel model, since it broadcasts parameters at construction and DTensor has no sharding strategy for c10d.broadcast_.
A replica whose completions come back short filled its batch first, trained, and then waited at the gradient sum for a replica still decoding, and it carried a lighter forward while that one carried the long tail alone: the wait and the activation peak both landed unevenly. The replicas now agree on when enough samples exist between them, and hand them out longest-first to whichever replica carries the least, by sum of squared lengths, since attention is quadratic and that is what predicts the step time and the peak. Nothing ties a sample to the replica that generated it, so ownership moves with the assignment. Every rank runs the same assignment over the same pool, so none of them has to be told the result. Measured on Qwen3-1.7B, 4 GPUs at tp=2: load imbalance 1.000-1.007 across replicas, with up to 100% of a replica's batch coming from its neighbour, and the replicas still agree on the weights.
Adam carries two fp32 moments per parameter, four times the weights themselves, and it is the term that decides whether a model fits: for a 14B at tp=4 it is 27.6 GiB of the 41.5 GiB a rank holds. The replicas divide the parameters between them, each updates its own share, and passes the result back to the others, so the parameters stay whole everywhere, which is what generation reads. The optimizer only knows its own share, so its zero_grad would leave the rest of the model holding gradients for the next step to accumulate onto: the whole model is cleared after the step instead. exp/check_shard_optim.py runs both paths in one process from one snapshot on the same gradients. Adam is elementwise, so which replica holds the moments cannot change the arithmetic, and the weights come out bit-identical after 3 steps.
A replica is a replica whether or not the model inside it is split across devices: the pool that balances the batches and the optimizer state that no replica needs twice apply either way. Only the gradient reduction is specific, because DDP already does it when it can wrap the model, and it cannot wrap a tensor-parallel one. Checked on 4 GPUs at tp=2/dp=2, tp=1/dp=4 and tp=4/dp=1, packed and padded, with and without accumulation: the replicas end every run holding the same weights.
Which samples to train on was decided by length, since the same longest-first pass both chose them and balanced them. That left the short ones sitting in the ready lists, and a rollout is either in flight or waiting there: a growing backlog is fewer sequences decoding at once, which is what makes the engine efficient. Measured on Qwen3-14B at 8 GPUs, concurrent sequences fell from about 575 to 368 and decoding lost 9%. Age now decides what is taken and length only decides where it goes, so the balance is the same: 1.009 mean imbalance on the skewed-length check, unchanged, with the replicas still starting their forward within a millisecond of each other.
The readiness gate ran a collective before every decode step, on the same devices the engine is issuing its own collectives on through a different communicator. Decoding lost 10% to it: 4364 trained tokens/s without replicas against 3943 with, on Qwen3-14B at 8 GPUs, with the same tokens generated in a longer wall time. Asked once every eight decode steps instead, on a schedule identical everywhere: a replica that skipped a collective the others were waiting on would hang them all. Overshooting the count by a few decode steps only leaves more samples ready.
The split of the world into replicas follows from tp_size, so there was nothing for a caller to decide and every one of them had to build the same ParallelismConfig by hand. Accelerate has to be told it, or it wraps the model for data parallelism itself, which it cannot do to a tensor-parallel one. The recipe pinned tp_size to the whole world, so it could never hold more than one replica. It now splits a 14B four ways and puts the rest of the ranks into replicas, which is what the measurements say: 4.1k trained tokens/s on 8 GPUs and 9.2k on 16, against 2.8k for the widest split on 8.
qgallouedec
changed the title
…es at once The engine used to be handed a whole batch of rollouts per optimizer step and nothing in between, so its decode batch drained as rollouts finished and refilled in a burst at the next step. A decode step costs about the same whatever it carries, since it reads the whole model and pays the per-layer all-reduces either way, so a batch that is half empty is a decode step half wasted. Rollouts are queued instead of started, and a rollout starts the moment one finishes. The members of a group therefore start at different times, which is no change in kind: zero-sync updates the weights while generation runs, so a group already sees more than one version of the policy. `generation_ahead` is gone, replaced by `rollouts_in_flight`. How many rollouts decode at once is a property of the engine and the KV cache, not of the batch being trained on, and the old name counted batches of a quantity it did not depend on. Measured on Qwen3-1.7B at tp=2/dp=2, steps 10 to 18, one replica: 3.22 s/step against 3.51, 4143 generated tokens/s against 3810, and the number of rollouts decoding at once holds between 34 and 40 where it used to swing between 24 and 44.
…lready walking Keeping the engine full means the queue runs ahead of the training loop, and the loop only hands over one batch a step. Reading ahead through a second iterator over the same dataset gave both of them the same prompts, so the opening batches of an epoch were generated and trained on twice. The dataloader now hands back an iterator the trainer keeps a reference to, and the queue draws from that one. The loop carries on from wherever the queue left off, so each prompt is read exactly once and none is skipped. Checked by hashing every prompt queued over a run: 22 and 24 queued on the two replicas, 22 and 24 distinct, none repeated. Reading one batch beyond what the engine holds keeps the queue from emptying when a burst of rollouts finishes together, which is what pins the number generating rather than only capping it. Qwen3-1.7B at tp=2/dp=2, steps 10 to 18, one replica: 2.64 s/step against 3.51 before, with the rollouts decoding at once held at exactly 40 where they used to swing between 24 and 44. Not generating the same prompt twice is most of the difference against the 2.99 s/step the second iterator gave.
packed_training gated a path that is strictly worse: padding rows to the longest in the batch wastes up to a third of the forward on pad tokens, measured 6.98 against 5.82 s/step on Qwen3-4B. Nothing chose it on purpose. Removing it takes the flag, the padded collator, and the length-anchored batch selection that only existed to make padding cheaper. The attention implementation still defaults to flex_attention, or flash attention for models with linear attention layers.
17b573c removed the flag but left the test that passes it, so the suite fails on that branch. The remaining test_train covers the same path, since every row is packed now.
The flag was removed in 17b573c, so both example scripts raise TypeError on the config before they reach a GPU.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.