Repository navigation
[Umbrella] Continuous batching fixes for training-while-generating on one weight copy - #48294
Closed
qgallouedec wants to merge 156 commits into
Closed
qgallouedec wants to merge 156 commits into
qgallouedec wants to merge 156 commits into
Conversation
…ite models PagedAttentionCache and the CB config resolution read decoder attributes (num_key_value_heads, num_hidden_layers, head_dim, ...) from the top-level config, which does not carry them for composite configs (vision-language models keep them on text_config). Continuous batching on any ForConditionalGeneration model crashed with 'num_key_value_heads or num_attention_heads could not be found in the config'. Resolve against config.get_text_config(), a no-op for text-only configs.
When the generation loop dies on a fatal error, add_request kept accepting requests and get_result returned None until the caller's timeout, so a generate_batch user saw a silent timeout with no trace of the underlying error. Raise from fatal_error in add_request, and in get_result once the remaining results are drained.
… forward has no cache A standard forward on a model switched to a paged attention implementation (as done by ContinuousBatchingManager) silently computed non-causal attention: sdpa_attention_paged_forward treats cache as optional and mask creation early-exits to None for paged implementations. Now, when the continuous batching kwargs are absent, the paged sdpa/flash forwards delegate to their base implementation, the eager paged forward materializes a causal mask, and mask creation resolves 2D padding masks with the base implementation's mask function.
DistributedHelper creates the TP CPU (gloo) group with a plain new_group, which requires every rank of the default group to participate in the call. When the TP group is a subset of the world (one continuous batching manager per node, data parallelism across nodes), the concurrent per-subset calls deadlock in gloo's connectFullMesh until timeout. use_local_synchronization restricts the creation to the group members, which is exactly the situation here; for the world-spanning TP case it is a no-op behavior-wise.
… instead of crashing The async-batching path bootstraps the first batch unconditionally at background-thread start and raises 'Failed to bootstrap the first batch' when the request queue is still empty. Whether a manager works therefore depends on a race between start() and the first add_request; user code that starts the manager and submits from another thread wins or loses it nondeterministically. Wait for the first request with the same stop conditions as the main loop instead.
qgallouedec
marked this pull request as draft
August 25, 2026 15:14
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
…ite models PagedAttentionCache and the CB config resolution read decoder attributes (num_key_value_heads, num_hidden_layers, head_dim, ...) from the top-level config, which does not carry them for composite configs (vision-language models keep them on text_config). Continuous batching on any ForConditionalGeneration model crashed with 'num_key_value_heads or num_attention_heads could not be found in the config'. Resolve against config.get_text_config(), a no-op for text-only configs.
…e/transformers into fix-cb-bootstrap-wait
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: qwen3_5, qwen3_5_moe, qwen3_next |
Contributor
CI recapDashboard: View test results in Grafana
|
`_get_optimizer_file` answers for the tensor-parallel and the plain case, so the caller that loads it no longer has to pick between them.
# Conflicts: # src/transformers/trainer.py
`init_continuous_batching` put `paged|<impl>` on the model and only `stop()` took it off, so any ordinary forward while the engine ran landed in the paged kernel with no paged cache and raised. The engine already passes its cache and the packed-input metadata through the forward kwargs, so `flash_attention_forward` and `sdpa_attention_forward` can route to their paged counterpart when a call carries a `PagedAttentionCache` and stay ordinary otherwise. The model is then left on the implementation it was loaded with and serves both. `eager` keeps the `paged|` prefix: every model defines its own eager forward, so there is no single one to route. `attn_mask_is_needed` now keys off the implementation instead of the `paged|` prefix, which only eager still carries.
# Conflicts: # src/transformers/trainer.py
# Conflicts: # src/transformers/distributed/utils.py # src/transformers/generation/continuous_batching/distributed.py
# Conflicts: # src/transformers/generation/continuous_batching/cache.py
# Conflicts: # src/transformers/distributed/utils.py
# Conflicts: # src/transformers/distributed/tensor_parallel.py # src/transformers/trainer.py
…ansformers into cb-train-umbrella
Member
Author
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.
Draft / umbrella PR: this branch is the sum of six independent fixes, each also opened as its own PR (the mergeable units). It exists so the whole set can be run and reviewed together.
Context
We explored running a training loop and continuous batching generation simultaneously on ONE copy of the weights (rollouts pick up optimizer updates with no weight sync). It works: dense, MoE, and VLM text-only models, greedy and sampled decoding, single GPU up to Qwen3-30B-A3B on TP4 and a 2-node TP-within-node / gradients-only-across-nodes layout, including a GRPO loop learning through it (reward 0.09 -> 1.0). Along the way we hit and fixed six CB bugs; this branch carries all of them.
Measured on this branch: 300 training steps with a CB manager attached give bitwise-identical final weights to training alone (same sha256); paged|sdpa == sdpa and paged|eager == eager exactly (max logit diff 0.0);
tests/generation/test_continuous_batching.py74 passed, 0 failed on the full set (H100, torch 2.13.0+cu130).The six fixes
paged|sdpasilently computes non-causal attention on a standard forward #48296). A standard forward on a CB-switched model silently computed non-causal attention (loss 2.84 vs 4.03, trains on future-token leakage).Issue reports with self-contained repros accompany 1, 2, 3 (secondary), 5, and 6.
@remi-or