Repository navigation
chore(trainer): unify context-parallel setup across models - #3533
Merged
Merged
Conversation
garrett361
marked this pull request as ready for review
September 11, 2026 18:11
S1ro1
reviewed
Sep 11, 2026
samsja
reviewed
Sep 11, 2026
samsja
previously approved these changes
Sep 12, 2026
CP wiring had grown two trainer-side helpers calling the same method name at different depths, a duplicated setup block in both trainers, and three attribute spellings for the same three values. One trainer entry point: `setup_context_parallel(model, config, parallel_dims)` in `utils/cp.py` replaces `setup_model_cp` and `setup_sparse_mla_cp`, and both trainers call it once behind their existing `cp_enabled` guard. It walks `model.modules()` and calls the method wherever it is defined, depth-first pre-order, so a module is always configured before its descendants. An unknown `cp_style` now raises instead of silently taking the ulysses branch. One model-side method: `PreTrainedModelPrimeRL.setup_context_parallel` stores the four values for its own module and nothing else. Each family therefore keeps exactly one such method, on the module that consumes the state, and the delegating overrides are gone: deepseek_v4 and glm_moe_dsa on their attention, nemotron_h on its Mamba mixer, qwen3_5 on its GatedDeltaNet. The walk reaches through activation-checkpoint wrappers and VLM bodies, which also fixes `setup_sparse_mla_cp` silently no-oping on a VLM. One spelling: `cp_group` / `cp_rank` / `cp_world_size` / `cp_style`, unprefixed everywhere, with class-level defaults on the base so read-back needs no getattr fallbacks. Carrying `cp_style` removes the hardcoded `cp_style="ulysses"` from the qwen3_5 VLM forward. Also drops `gpt_oss`'s `cp_support` override, which returned the base default. The CP-support matrix is otherwise unchanged. No behavioural change intended. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
garrett361
force-pushed
the
feat/uniform-cp-setup
branch
from
September 14, 2026 13:33
ec7dd2f to
c80b873
Compare
The four CP values were homogeneous by convention: every consuming module declared `cp_group`, `cp_rank`, `cp_world_size` and `cp_style` and implemented a `setup_context_parallel` method assigning exactly those names. Nothing enforced it, so a typo in one name would have miswired CP silently. `CPContext` in `utils/cp.py` now carries the four values and `cp_enabled`, and a module that consumes CP state declares one field, `self.cp_context`, defaulting to `CPContext()`, which is the non-CP case. `PreTrainedModelPrimeRL` supplies that default for every model. `setup_context_parallel` builds one frozen instance and assigns it to every module that has the field, so all consumers share a single object, and raises if that attribute is not a `CPContext`, catching a name collision at setup rather than at the first forward. This deletes five identical `setup_context_parallel` methods, three duplicate `cp_enabled` properties and sixteen default assignments. `Qwen3_5VLMModel` holds its own context instead of reading its child's, leaving the `eager_reference` helper, which takes the attention module as an argument, as the only place that reads CP state off another module. No behavioural change intended. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
garrett361
commented
Sep 14, 2026
Comment on lines
+59
to
+68
| cp_context = CPContext(cp_group, cp_rank, parallel_dims.cp, config.cp_style) | ||
| for module in model.modules(): | ||
| if not hasattr(module, "cp_context"): | ||
| continue | ||
| if not isinstance(module.cp_context, CPContext): | ||
| raise TypeError( | ||
| f"{type(module).__name__}.cp_context is {type(module.cp_context).__name__}, not a CPContext; " | ||
| "context-parallel setup claims that attribute name, so rename it" | ||
| ) | ||
| module.cp_context = cp_context |
Contributor
Author
There was a problem hiding this comment.
I don't love this keying off cp_context attr existence, but it is better than threading a cp_context attr through all __init__ methods. Open to suggestions.
Contributor
There was a problem hiding this comment.
agreed but can't think of better one
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.
Unifies all context parallel setup so that a single frozen
CPContextdataclass is shared by all cp-aware modules, stored under aself.cp_contex: CPContextattr.