config: the derived parallel widths are computed from the leaves - #36790
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6767a69319
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| moe_dp_size=moe_data_model_parallel_size, | ||
| dcp_size=decode_context_parallel_size, | ||
| dcp_enabled=_DCP is not None, | ||
| world_size=world_size, |
There was a problem hiding this comment.
Preserve expanded WORLD size for scale joiners
When an elastic scale joiner starts with a nonzero ep_join_rank_offset, bootstrap.py initializes WORLD with rank_offset + tp_size * pp_size, but initialize_model_parallel() deliberately replaces its local world_size variable with tp_size * pp_size when recovered_rank=True. Stamping that local value makes get_parallel().world_size return only the joining cohort's size instead of the actual expanded WORLD size previously returned by get_world_size(); derive this value from the live world group or pass the original distributed size separately.
Useful? React with 👍 / 👎.
6767a69 to
ad5f8e3
Compare
304a6ac to
a242476
Compare
ad5f8e3 to
c6a9b46
Compare
a242476 to
35d45cb
Compare
c6a9b46 to
8595626
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 859562673c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @property | ||
| def attn_dp_size(self) -> int: | ||
| return self._v("attn_dp_size", _dp().get_attention_dp_size) | ||
| return self._derived_width("attn_dp_size", _dp().get_attention_dp_size) |
There was a problem hiding this comment.
Honor temporary DP-size overrides
When code enters dp_attention.disable_dp_size() for a draft-model scope, that context manager changes the live _ATTN_DP_SIZE to 1, but this property now returns the previously stamped value without consulting the live getter. Consequently, callers using the canonical get_parallel().attn_dp_size inside that scope observe the target model's DP width rather than the disabled width, unlike the pre-change behavior; temporarily override/restamp the derived value together with _ATTN_DP_SIZE.
Useful? React with 👍 / 👎.
8595626 to
90c7cc1
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 90c7cc125e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
90c7cc1 to
ba959c3
Compare
35d45cb to
f4472ae
Compare
ba959c3 to
dcb2776
Compare
f4472ae to
6be2ce6
Compare
dcb2776 to
3e2d584
Compare
… read back
`get_parallel().attn_tp_size` and its five siblings -- `attn_dp_size`,
`moe_ep_size`, `moe_tp_size`, `dcp_enabled`, `attn_dcp_size` -- were read back
off the group coordinators that had just been built from the
configured leaves. That made them the one corner of `get_parallel()` where a
name needs distributed init rather than a publish, and the two spellings are
indistinguishable at the call site: `tp_size` raises `ValueError("'parallel'
not published")`, `attn_tp_size` raised a group getter's bare `AssertionError`.
The arithmetic moves to `runtime_context.derive_parallel_widths`, and
`initialize_model_parallel` builds its groups from that dict and stamps it. One
formula, one place: the width a group was built at and the width a reader gets
cannot drift apart. `initialize_dp_attention` stamps `attn_dp_size`, and
`update_dp_attention_post_scale` restamps it where it already updates the live
one -- which is the case that makes a plain per-read derivation wrong. Elastic
EP rewrites `dp_size` and `ep_size` on the published bag while the coordinators
keep their construction width, so `tp_size == attn_tp_size * attn_dp_size *
attn_cp_size` stops holding after a scale-up; deriving on every read would
answer with a width that shrank, or zero.
`world_size` stays a live read and is deliberately not in the set. It is not a
quotient of the leaves, and its getter is right at every moment: a stamp taken
when the groups are built would answer with the launch count after
`try_admit_scale_ranks` expands WORLD, and with the joining cohort's own width
on a scale-joiner, which lays its groups out at `tp * pp` while WORLD spans
`ep_join_rank_offset + tp * pp`.
A scope that moves a width temporarily moves the derived one with it, the way
`patch_tensor_parallel_group` already does for `tp_size`: `disable_dp_size()`
runs a draft worker without DP attention, and overrides `attn_dp_size` for its
duration so the legacy getter and `get_parallel()` cannot disagree inside it.
A read is the override, then the stamp, then the live group. The fallback keeps
a process that installed groups without `initialize_model_parallel` working;
when neither is present the error says which of the two is missing.
`TestDerivedWidths` pins the quotients, the stamp winning over the group, the
override winning over the stamp, the restamp, the named failure, and that
`parallel_state` no longer carries its own copy of the arithmetic.
The 265 registered tests that mention the record report the same failure set as
the base commit.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…k helper `compute_dp_attention_world_info` recomputed `attn_dp_size` and `attn_tp_size` with the same two lines the stamped derivation uses. It cannot go away -- four of its six callers want only the ranks it returns, and ranks are per-process, so they are not part of the stamped set -- but the widths in it were a second copy that could drift from the one the groups are built at. `derive_attention_widths` is now the single home for that pair; `derive_parallel_widths` and the rank helper both call it. `initialize_dp_attention` had a third copy -- it called the rank helper, dropped the width it returns and recomputed `dp_size if enable_dp_attention else 1` three lines below -- and now takes what the helper already gave it. The guard that already refused a duplicate quotient in `parallel_state.py` covers `dp_attention.py` too, and refuses that `else 1` spelling as well, which the quotient patterns did not match. A new case checks the helper and the stamp agree. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
3e2d584 to
4b7ea9b
Compare
…-project#36790) Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
…-project#36790) Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
attn_dcp_size is now a derived parallel width resolved at publish time (sgl-project#36790, sgl-project#38113), so it no longer follows an overridden dcp_enabled and a draft forward kept the target's DCP width. State the widths the guard means, the way the other production override sites do.
Motivation
get_parallel()answers a name one way: what the operator configures comes from the publishedparallelbag, what only exists once the process groups are up comes from the groups. Sixnames broke that rule —
attn_tp_size,attn_dp_size,moe_ep_size,moe_tp_size,dcp_enabledandattn_dcp_size. None is aparallelleaf: each is aquotient of leaves that no flag sets, and each was read back off the group coordinator that had
just been built from those leaves.
That made them the one corner of
get_parallel()where a name needs distributed init ratherthan a publish, and the difference is invisible at the call site:
Modifications
The arithmetic moves to
runtime_context.derive_parallel_widths.initialize_model_parallelbuilds its groups from that dict and stamps it, so the width agroup was built at and the width a reader gets cannot drift apart.
initialize_dp_attentionstampsattn_dp_size;update_dp_attention_post_scalerestamps itwhere it already updates the live one;
destroy_model_parallelclears the stamp.That last case is why a plain per-read derivation is not enough. Elastic EP rewrites
dp_sizeand
ep_sizeon the published bag while the coordinators keep their construction width, sotp_size == attn_tp_size * attn_dp_size * attn_cp_sizestops holding after a scale-up: aformula evaluated on every read would answer with a width that shrank, or with zero.
world_sizestays a live read and is deliberately not in the set. It is not a quotient of theleaves, and its getter is right at every moment: a stamp taken when the groups are built would
answer with the launch count after
try_admit_scale_ranksexpands WORLD, and with the joiningcohort's own width on a scale-joiner, which lays its groups out at
tp * ppwhile WORLD spansep_join_rank_offset + tp * pp.A read is the override, then the stamp, then the live group. The fallback keeps a process that
installed groups without going through
initialize_model_parallelworking; when neither ispresent, the error names which of the two is missing instead of surfacing a group getter's
assertion.
The second commit removes the last duplicate of the arithmetic.
compute_dp_attention_world_inforecomputedattn_dp_sizeandattn_tp_sizewith the sametwo lines. It stays — four of its six callers want only the ranks it returns, and ranks are
per-process, so they are not part of the stamped set — but the widths in it now come from
derive_attention_widths, whichderive_parallel_widthscalls as well.Accuracy Tests
TestDerivedWidthsis new. It pins the quotients, the stamp winning over the group, anoverride winning over the stamp, the restamp, the named failure when neither is present, that
neither
parallel_state.pynordp_attention.pycarries a second copy of a quotient, and thatthe rank helper and the stamp agree across three topologies.
The 159 registered tests that mention the record report the same failure set as the base, and
the 24-shape resolution dump is field-for-field identical.
Speed Tests and Profiling
None. A stamped width is a dict lookup instead of an attribute read on a group coordinator.
Checklist
Review and Merge Process
Small enough to read in full, and worth reading in full: the elastic-EP case is the whole
argument for stamping rather than deriving per read.
CI States
Latest PR Test (Base): ❌ Run #33194068745
Latest PR Test (Extra): ❌ Run #33194068628
Latest PR Test (AMD ROCm 7.2): 🚫 Run #33194068935