config: spell the parallel config tier at the call site (#36250)

This commit is contained in:
Cheng Wan
2026-08-26 03:00:28 -07:00
committed by GitHub
parent 689ade69d1
commit 8005df61d3
136 changed files with 1033 additions and 795 deletions
+60 -90
View File
@@ -13,13 +13,16 @@
# ==============================================================================
"""A single structured accessor for process-static runtime state.
``get_parallel()`` returns a ``ParallelContext`` whose attributes — tp / dcp / pp /
moe / attn size and rank, plus the process-group handles — each delegate live to
the canonical getter in ``distributed.parallel_state`` / ``layers.dp_attention``.
Returned values are exactly what those getters return; this is a read-through
wrapper, not a cache. It gives call-sites one import and one naming scheme in
place of a dozen free functions, plus a test-only ``override()`` hook to force a
topology without monkeypatching the underlying getters.
``get_parallel()`` returns a ``ParallelContext`` whose bare attributes — tp / dcp
/ pp / moe / attn size and rank, plus the process-group handles — each delegate
live to the canonical getter in ``distributed.parallel_state`` /
``layers.dp_attention``. Returned values are exactly what those getters return;
this is a read-through wrapper, not a cache. It gives call-sites one import and
one naming scheme in place of a dozen free functions, plus a test-only
``override()`` hook to force a topology without monkeypatching the underlying
getters. The resolved parallel **configuration** is the same object's ``config``
hop (``get_parallel().config.tp_size``), which reads the published ``parallel``
bag: bare is the live group, ``config`` is what was configured.
``get_server_args()`` returns the process-wide ``ServerArgs``. This is the pristine / resolved-at-startup **read-only** record kept
for debug and reproduction; business code reads resolved config from the
@@ -114,14 +117,21 @@ _PARALLEL_FIELDS = frozenset(
class ParallelContext:
"""Parallel-topology namespace.
"""Parallel-topology namespace: the live groups bare, configuration under
``config``.
Live topology (size / rank / group) is read-through via ``@property`` (the
canonical getters). Parallel **config** leaves (``nccl_port``,
``pp_max_micro_batch_size``, ``enable_dp_attention``, …) come from the
published ``parallel`` config bag via ``__getattr__``. Where a config leaf
shares a name with a live property (``tp_size`` …), the property (the live
fact) wins; the same-name==same-value invariant holds once dist is up.
``get_parallel().tp_size`` and its size / rank / group siblings are
read-through ``@property`` over the canonical getters, so they answer with
the **live** process groups and raise before distributed init. The resolved
parallel **configuration** is one hop away, on the published bag:
``get_parallel().config.tp_size``, ``.config.nccl_port``. It answers in any
process at any point after publish, and follows a post-publish ``override``.
The two disagree by design, so which one a call site wants is spelled at the
call site — no ``config`` means live. Elastic EP scales the live world away
from the configured one, and ``initialize_model_parallel`` aliases
``_MOE_DP`` to ``_ATTN_CP`` when ``attn_cp_size > moe_dp_size``, which makes
a live comparison of that pair degenerate.
"""
__slots__ = ("_overrides", "_config")
@@ -130,28 +140,37 @@ class ParallelContext:
self._overrides = {}
self._config = None # parallel config bag, wired at publish
@property
def config(self) -> _ConfigBag:
"""The published ``parallel`` config bag.
Reads the slot directly: ``parallel`` sits outside the per-role
namespace table (every process reads topology config), so no role check
applies here. The body stays
dynamo-traceable — ``get_parallel().config.moe_dense_tp_size`` and the
gate helpers over it run inside compiled model forwards.
"""
config = self._config
if config is None:
raise ValueError("config namespace 'parallel' not published")
return config
def __getattr__(self, name):
# Reached only for names that are neither a live @property nor a slot:
# serve parallel config leaves from the published bag. The body must
# stay dynamo-traceable — config-leaf reads such as
# ``get_parallel().moe_dense_tp_size`` run inside compiled model
# forwards, and ``object.__getattribute__`` graph-breaks.
# Reached only for names that are neither a live @property nor a slot.
# Config leaves are read under ``config``, so naming one here is a
# call-site mistake and this only builds the error that says so.
if name.startswith("_"):
# No config leaf is underscored; this also breaks the recursion
# when the ``_config`` slot itself is still unset (pickle/copy
# protocols probe attributes before __init__ runs).
# This also breaks the recursion when the ``_config`` slot itself is
# still unset (pickle/copy protocols probe attributes before
# __init__ runs).
raise AttributeError(name)
config = self._config
# ``_fields`` is a plain ``__dict__`` entry on the bag; ``in`` on the
# dict avoids ``_ConfigBag.__contains__`` (not traceable).
if config is not None and name in config._fields:
return getattr(config, name)
detail = (
"not a published parallel config leaf"
if config is not None
else "config not published"
)
raise AttributeError(f"ParallelContext has no {name!r} ({detail})")
raise AttributeError(
f"{name!r} is a parallel config leaf, not live topology; read it "
f"as get_parallel().config.{name}"
)
raise AttributeError(f"ParallelContext has no {name!r}")
def _v(self, name, getter):
overrides = self._overrides
@@ -817,9 +836,9 @@ class RuntimeContext:
# truth for config reads). Driven by NS(...) metadata; a mock/partial
# config with no NS markers yields an empty tree (no bags projected).
self._config_bags = _build_config_bags(server_args)
# Wire the parallel config leaves onto the live wrapper (config-only
# leaves like pp_max_micro_batch_size are served via ParallelContext
# __getattr__; live topology properties still win by name).
# Wire the published `parallel` bag onto the live wrapper: it is the slot
# the `config` property reads, which is how config-only leaves like
# pp_max_micro_batch_size are spelled.
self.parallel._config = self._config_bags.get("parallel")
# A direct install is roleless; ``publish`` assigns the role afterwards.
self._overrides_log = []
@@ -1121,8 +1140,8 @@ def get_forward() -> ForwardFlags:
# --- Resolved config namespaces -------------------------
# Each returns the top-level snapshot bag; reads are `get_exec().moe.field` etc.
# All fail with ValueError("... not published") until publish has projected them.
# ``parallel`` config leaves are served by ``get_parallel()`` (live wrapper);
# their config-bag wiring is a scoped follow-up.
# ``parallel`` has no getter of its own: its bag is reached as
# ``get_parallel().config``, alongside the live topology it belongs to.
def get_device() -> _ConfigBag:
return _CONTEXT.config_bag("device")
@@ -1173,8 +1192,8 @@ def get_observability() -> _ConfigBag:
# table declares which top-level config namespaces each role reads. ``None``
# means the full tree — either the role genuinely needs everything (scheduler)
# or its deployment shape has not been audited yet (restrict only what smoke
# coverage can verify). ``parallel`` is served by ``get_parallel()`` and every
# process legitimately reads topology config, so it is not part of this table.
# coverage can verify). ``parallel`` is served by ``get_parallel().config`` and
# every process legitimately reads topology config, so it is not in this table.
#
# ``SGLANG_ROLE_NAMESPACES`` selects the mode (read once at import):
# off (default) no bookkeeping, zero overhead;
@@ -1559,7 +1578,7 @@ def max_prefill_buffer_tokens() -> int:
tokens = chunked
if (
schedule.enable_dynamic_chunking
and _configured_parallel("pp_size") > 1
and get_parallel().config.pp_size > 1
and chunked
):
tokens = max(
@@ -1591,10 +1610,9 @@ def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float:
activation_tokens = max(schedule.chunked_prefill_size, 2048)
else:
activation_tokens = max(schedule.max_prefill_tokens, 2048)
parallel = get_parallel().config
reserved_mem = (
512
+ activation_tokens * 1.5
+ _configured_parallel("tp_size") * _configured_parallel("pp_size") / 8 * 1024
512 + activation_tokens * 1.5 + parallel.tp_size * parallel.pp_size / 8 * 1024
)
if gpu_mem is not None and gpu_mem > 60 * 1024:
reserved_mem = max(reserved_mem, 10 * 1024)
@@ -1714,54 +1732,6 @@ def cutedsl_moe_max_num_tokens() -> int:
return max(prefill_tokens, decode_max_bs * num_tokens_per_req)
# --- Configured (not live) parallel sizes ------------------------------------
#
# ``get_parallel()`` shadows these names with the LIVE topology, which is the
# right answer almost everywhere. A handful of call sites need what was
# *configured* instead — before the groups exist, in a process that has none,
# or where the live value is deliberately aliased to another dimension. Each
# accessor below names that intent so no business call site has to reach for
# the startup record; the per-site reasons live in the read ratchet.
#
# They read the published leaf rather than the record: the bag is what
# ``override`` writes, and once the instance holds only the user's raw input
# the record would answer with what was *typed* instead of what resolution
# produced. Going through the bag directly is what gets past the live property
# that shadows these four names on ``get_parallel()``.
def _configured_parallel(name: str):
# The bag itself, not ParallelContext, whose live property shadows these
# four names. Read through the parallel slot the way the leaf accessor
# does — ``parallel`` is deliberately outside the per-role namespace table
# (every process reads topology config), so this must not route through
# ``config_bag()``'s role check, which would record or reject the read.
config = _CONTEXT.parallel._config
if config is None:
raise ValueError("config namespace 'parallel' not published")
return getattr(config, name)
def configured_tp_size() -> int:
return _configured_parallel("tp_size")
def configured_pp_size() -> int:
return _configured_parallel("pp_size")
def configured_moe_dp_size() -> int:
return _configured_parallel("moe_dp_size")
def configured_attn_cp_size() -> int:
return _configured_parallel("attn_cp_size")
def configured_dcp_size() -> int:
return _configured_parallel("dcp_size")
def is_ep_joiner() -> bool:
"""True in a process launched as an elastic-EP joiner (scale or recover).