Repository navigation
ZeroSyncGRPOTrainer - #6949
Draft
qgallouedec wants to merge 143 commits into
Draft
ZeroSyncGRPOTrainer#6949qgallouedec wants to merge 143 commits into
qgallouedec wants to merge 143 commits into
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.
The engine decodes from the model's parameters, so they stay bf16, and AdamW stepped them directly: one step at lr 1e-6 lands 5% of its size on Qwen3.5-2B (24% of the weights change at 1e-5). The optimizer now updates fp32 copies and writes them back after each step.
# Conflicts: # trl/experimental/zero_sync_grpo/zero_sync_grpo_config.py
Measured with the example as written on 4x B200: 200 steps in 25 minutes, reward 0.21 -> 0.84 by step 60 and 0.87-0.99 through step 200. The previous lr 1e-5 relied on bf16 rounding most of each step away.
Member
Author
|
Ran this branch on the models the continuous batching stack now unlocks (huggingface/transformers#50350 -> huggingface/transformers#50356 -> huggingface/transformers#50359 -> huggingface/transformers#50365 -> huggingface/transformers#50369, plus huggingface/transformers#50366 for TP on VLM checkpoints). GSM8K, 32 samples/step, lr 1e-5:
Replicas scale linearly: Qwen3.5-4B at 1.9 / 3.9 / 7.4 / 15.7 samples/s on 1 / 2 / 4 / 8 B200, same reward. Fixed on the way:
At lr 0 the engine and the trainer agree to bf16 noise on every model above (0.1% of tokens off by more than 0.18 nats on Qwen3.5, flat up to 6k tokens), except the MoE models, where two plain forwards of the same sequences already disagree as much (router ties in bf16). |
They come onto the GPU one parameter at a time during the step, so the GPU holds only the bf16 parameters and their gradients (4 bytes per parameter instead of 16) next to the KV cache and the activations.
The next parameter's fp32 copy and moments come in on one side stream while the previous one goes back on another. Benchmarked on a B200 node with a 4B parameter set: 2.28 -> 1.60 s per step on one GPU, 3.2 -> 2.9 s with four GPUs sharing the host link.
With grad enabled, DTensor.to_local() returns a new tensor, so swapping its storage left the shard on the host and the fused step failed on mixed devices at tp>1. New 2-GPU test runs the trainer at tp=2 and checks the shards of the fp32 copies end up in pinned host memory.
A model sharded at load time cannot save resumable checkpoints. The test now passes at tp=2 and fails on eb24000 with the mixed-device error.
New release_kv_cache_during_step, on by default: the trainer frees the engine's KV cache inside the pause and brings it back before resuming, so the forward and backward get its memory. The rollouts in flight wait in the engine's CPU swap pool, sized like the cache, and resume where they were. Needs transformers' release_memory/restore_memory.
The Trainer collects every micro-batch of a step before it trains on the first, so pausing after each one only stopped the engine for nothing: wait for the whole step's samples and pause once. With the KV cache released during the step, that is one round trip of the cache per step instead of one per micro-batch (Qwen3.5-4B on a B200: 2.9 s per step instead of 12.0 s).
Every rank of a group read the completions and ran the tools and the reward functions on them, so each ran once per rank, and a tool with side effects ran several times. The first rank now runs the rollouts alone: it submits the requests, scores the completions, computes the advantages and picks the batch, then sends it to the other ranks of its group over gloo, which join the pause and train on it. With replicas, the pool is exchanged between first ranks. The request ids now come from the engine, and the drain after the pause, which kept the ranks' reads aligned, goes.
…ples are ready The wait read completions until none arrived for 1 s (0.2 s with replicas) and only then checked the count, so under a steady stream it kept reading past the target: the pause came late, and with replicas the one still reading held the other at the count sum. It now reads one completion at a time and checks after each. The replicas start the sum of their counts and keep scoring while it travels, so a slow reward on one never holds up the other, and both stop on the same sum. Adds a test for two replicas, which had none.
The engine decodes a batch on its fast path only if every sequence in it fits the block table, whose width falls back to 8192 tokens when nothing sets it. With longer completions nearly every decode step held one, so every step took the slow path. The trainer now sizes it for the model's context unless the continuous batching config sets it. Qwen3-4B on MATH-500 levels 4-5, completions up to 16k tokens, one B200: 34.6 s/step against 410 s/step (12x).
lend_kv_cache_during_step allocates the engine's KV cache with sublet and runs each training step inside a loan of its memory: the step takes the cache's free blocks, then free GPU memory, then the blocks of rollouts in flight, which are copied out and back. Needs the engine's lendable_cache and lend_memory (play tree).
…aises sublet 0.2 starts a loan when lend() returns it; the trainer ends it with end(). If ending it raises (LeakError, after the cache is back), the engine still resumes, and train() still stops it: a failure at the end of a loan used to leave the generation thread waiting, and the process hung until the job timed out.
…() and end with end()
The engine creates its KV cache in its own thread as it starts. The trainer wraps the manager's _create_batch_processor to move that cache onto a tensor made by sublet before the engine uses it (the cache allocators rebuild their views), computes the views the engine still needs from the cache pool's free sectors and blocks, and calls sublet.lend itself. No lendable_cache option, live_views or lend_memory in transformers anymore. Same loans as the engine-side version on a Qwen3-4B run: nothing copied over the first 6 steps, live KV checksums unchanged.
…, keep all but its free memory The trainer creates the engine's batch processor itself, as the engine's warmup does, and moves its KV cache onto a lendable tensor there, so nothing wraps the engine's own methods anymore. The views to keep are everything but the free sectors and the free blocks, which also covers the trash sector padding tokens read. Both helpers are inlined.
…7fbbb5 Nothing reads it. The previous simplification kept it, 36 MiB on Qwen3-4B.
…se or lend release_kv_cache_during_step and lend_kv_cache_during_step made four combinations for three behaviors, and lend silently overrode release when both were on. Same three behaviors, same default (release).
…, as AsyncGRPO does (#8159) TurnRecord, _SampleBuilder, _chain_to_sequences and Rollout are now AsyncGRPO's, byte for byte, and so are the reconciler tests.
… trainer's engine ZeroSyncRolloutLoop is AsyncRolloutLoop with the engine in place of the vLLM server: generate_one, _execute_tool_calls, _calculate_rewards, compute_advantages, _score_group, the generate and score loops and their metrics are AsyncGRPO's. It runs on an asyncio event loop the trainer drives between its steps, draws the epoch's prompts from the trainer's dataloader, and hands each request's result to the turn awaiting it. With it come AsyncGRPO's behaviors: groups whose rewards are all equal are dropped, advantages are (r - mean) / (std + 1e-8), completions are parsed with the response parser, a non-prefix-preserving template is swapped for a training one with tools, rewards are computed per group, environment_factory and loop_cls. rollouts_in_flight is renamed max_concurrent_rollouts, log_completions and num_completions_to_print are added, and the metrics are logged under AsyncGRPO's names.
compute_loss is AsyncGRPO's over the completion tokens of the packed row: tokens whose importance ratio to the engine falls outside [importance_sampling_clip_min, importance_sampling_clip_max] are dropped from the loss (AsyncGRPO's vllm_importance_sampling_clip_*, without the vLLM in their name), the loss is normalized by the micro-batch's trained tokens over the replicas and scales itself for gradient accumulation, and the sampling/* metrics are logged. The step metrics (perf/*, batch/*_per_step, num_tokens) come from AsyncGRPO's _log_step_metrics and its callbacks; generation_wait_s becomes perf/rollout_wait_s.
The config follows AsyncGRPO's sections, docstrings and defaults (max_completion_length 2048, logging_steps 1, ignore_data_skip forced on), with trust_remote_code and router_aux_loss_coef; the engine's options get their own section. __init__ runs in AsyncGRPO's order with its checks: no Liger, the MoE auxiliary loss, a frozen vision tower with an AutoTokenizer for vision-language models, a placeholder dataset when an environment owns the data. The docstring is AsyncGRPO's, with the paper.
…leness The docs page gains AsyncGRPO's sections that apply (vision-language models, customizing the rollout loop, design philosophy, logged metrics) and its API reference (config, trainer, rollout loop, Rollout, TurnRecord, RolloutGroup). The tests gain AsyncGRPO's tests of the code the two now share: metric reduction and push, tool execution, the rollout loop, group scoring and the zero-std drop. sample/staleness_* and sample/time_in_queue_s are logged as AsyncGRPO logs them.
…RPO's, port its loop tests log_rollout_traces and its buffering come from AsyncGRPO, with the trainer taking the samples in place of its queue dataset. The cleanup moves to _inner_training_loop as in AsyncGRPO, and the methods they share follow its order. The result loop polls the engine only while a request is in flight. AsyncGRPO's environment wiring and custom loop tests run on ZeroSyncRolloutLoop, and a trainer test trains with a custom loop_cls.
… the vLLM server's A failed request came back as an empty output, and the response parser then failed on it, hiding the engine's error.
This branch has not been deployed
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
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.
GRPO where training and generation share one copy of the weights: no server and no weight sync, the continuous batching engine decodes from the tensors the optimizer updates and pauses for the step.
Needs the pause mechanism (huggingface/transformers#48462) and continuous batching leaving sdpa and flash models on the implementation they were loaded with (huggingface/transformers#49069), transformers 5.19 is the first release with both.
DistributedConfig) transformers#50306.add_fused_lm_headdoes not do yet.