config: spell the parallel config tier at the call site (#36250)
This commit is contained in:
@@ -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).
|
||||
|
||||
|
||||
Reference in New Issue
Block a user