From 132eef1f28978d786f7820094794573554326823 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:22:09 +0000 Subject: [PATCH 01/48] Add the distributed PTQ session and its distribution plan Distributed PTQ needs a few process-level facts before a model loads, and one place to read them from instead of every caller re-deriving rank, world size and device from torch.distributed. - is_distributed_launch() reads the launch environment, so parse-time checks work before any process group exists; is_distributed_run() reads the live group. - DistributionPlan records how a HuggingFace MoE is distributed: the expert-parallel degree, and whether the dense part is replicated or FSDP2-sharded. An EP degree that cannot divide the world is rejected at construction rather than as a mesh error once every rank is up. - ModelOptSession holds this process's identity (distributed / rank / world_size / local_rank / device / is_main), populated once per run -- reading any of it earlier raises rather than returning a stale default -- plus the expert-parallel mesh the EP loader sets. - modelopt_session() owns the process group for a block: on a distributed launch it creates the group and destroys it on exit, even if the body raises. A body that raises prints its traceback before the cleanup barrier, so a rank that fails alone is visible rather than hidden behind peers blocked in a collective. get_active_mo_session(required=True) fails fast; required=False returns the singleton, for defensive reads deep in export. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 234 +++++++++++++++++- .../torch/utils/test_distributed_session.py | 105 ++++++++ 2 files changed, 338 insertions(+), 1 deletion(-) create mode 100644 tests/unit/torch/utils/test_distributed_session.py diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 5627754cb06..a00c661ab3a 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -22,7 +22,7 @@ import time import traceback from collections.abc import Callable, Iterator -from contextlib import suppress +from contextlib import contextmanager, suppress from datetime import timedelta from typing import Any from warnings import warn @@ -34,16 +34,22 @@ __all__ = [ "DistributedProcessGroup", + "DistributionPlan", + "ModelOptSession", "ParallelState", "backend", "barrier", "fsdp2_wrap", "gather_bytes", "gather_state_dicts", + "get_active_mo_session", "is_available", + "is_distributed_launch", + "is_distributed_run", "is_fsdp2_model", "is_initialized", "is_master", + "modelopt_session", "rank", "size", ] @@ -98,6 +104,26 @@ def is_last_process(group=None) -> bool: return rank(group=group) == size(group=group) - 1 +def is_distributed_launch() -> bool: + """True if launched distributed (WORLD_SIZE>1), readable BEFORE the process group is set up. + + Env-based (does NOT touch the guarded session -- see :func:`get_active_mo_session`), so parse-time code + can detect the launch mode before :func:`modelopt_session` exists. After the session is entered, prefer + ``get_active_mo_session().distributed``. + """ + return int(os.environ.get("WORLD_SIZE", "1")) > 1 + + +def is_distributed_run() -> bool: + """True if a real multi-rank job is running: the process group is initialized and world size > 1. + + Runtime state (checks the live process group), unlike :func:`is_distributed_launch` which reads + only the launch env. A single process or a 1-rank group returns False. Backend-agnostic -- true + for any backend that initialized ``torch.distributed`` (HF EP/DP loader, Megatron, plain DDP). + """ + return is_initialized() and size() > 1 + + def _serialize(obj: Any) -> torch.Tensor: buffer = io.BytesIO() torch.save(obj, buffer) @@ -269,6 +295,212 @@ def setup(timeout: timedelta | None = None): torch.distributed.init_process_group("cpu:gloo,cuda:nccl", timeout=timeout) +class DistributionPlan: + """How a HuggingFace MoE model is distributed for PTQ, as decided by ``decide_distribution_plan``. + + * ``ep_size`` -- expert-parallel degree: the routed experts are sharded over ``ep_size`` ranks. + * ``fsdp_size`` -- dense-FSDP degree: ``1`` replicates the dense (non-expert) part on every rank; ``>1`` + FSDP2-shards it alongside the expert-parallel experts. + + Every rank calibrates on its own data, so data-parallelism spans the whole world, and each expert + shard is replicated ``edp = world_size / ep_size`` times. + """ + + def __init__(self, world_size: int, ep_size: int = 1, fsdp_size: int = 1) -> None: + """Build the plan, rejecting an expert-parallel degree that cannot divide the world.""" + if world_size % ep_size != 0: + raise ValueError(f"ep_size={ep_size} must divide world_size={world_size}") + self.world_size = world_size + self.ep_size = ep_size + self.fsdp_size = fsdp_size + + @property + def edp(self) -> int: + """Expert-data-parallel degree: how many ranks hold a replica of each expert shard.""" + return self.world_size // self.ep_size + + @property + def is_expert_parallel(self) -> bool: + """Whether routed experts are sharded across ranks.""" + return self.ep_size > 1 + + @property + def is_fsdp_dense(self) -> bool: + """Whether the dense part is FSDP-sharded rather than replicated.""" + return self.fsdp_size > 1 + + def __repr__(self) -> str: + dense = f"fsdp_size={self.fsdp_size}" if self.is_fsdp_dense else "dense=replicate" + return ( + f"DistributionPlan(world={self.world_size}, ep_size={self.ep_size}, edp={self.edp}, " + f"{dense})" + ) + + +class ModelOptSession: + """The live distributed session of a PTQ run: this process's identity, plus the model's EP mesh. + + One class covers both cases: ``distributed`` is False for a single process (``WORLD_SIZE==1``) and + True for a real ``torch.distributed`` launch (``WORLD_SIZE>1``). Two parts, with different lifetimes: + + * process IDENTITY (``distributed``/``rank``/``world_size``/``local_rank``/``device``/``is_main``) -- + stable for the whole run, populated once by :func:`modelopt_session`. Reading ANY before + initialization raises, so code cannot silently consume a stale/zero default before the process + group is set up. + * ``ep_mesh`` -- the expert-parallel device mesh of the loaded model, set by ``load_and_prepare_ep`` + (``None`` until then, and for a model loaded without expert parallelism). + """ + + def __init__(self) -> None: + """Create an empty (uninitialized) session; populate identity with :meth:`initialize`.""" + self._values: dict[str, Any] | None = None + self.ep_mesh = None + + def initialize( + self, + *, + distributed: bool, + rank: int, + world_size: int, + local_rank: int, + device, + is_main: bool, + ) -> None: + """Populate the session fields (called by ``modelopt_session``, or directly for tests).""" + self._values = { + "distributed": distributed, + "rank": rank, + "world_size": world_size, + "local_rank": local_rank, + "device": device, + "is_main": is_main, + } + + def _reset(self) -> None: + self._values = None + self.ep_mesh = None + + @property + def initialized(self) -> bool: + """Whether ``modelopt_session`` has populated the state.""" + return self._values is not None + + def _get(self, key: str): + if self._values is None: + raise RuntimeError( + "session accessed before initialization; enter " + "modelopt.torch.utils.distributed.modelopt_session() first" + ) + return self._values[key] + + @property + def distributed(self) -> bool: + """True for a distributed launch (WORLD_SIZE>1), False for single-process.""" + return self._get("distributed") + + @property + def rank(self) -> int: + """Global rank of this process (0 single-process).""" + return self._get("rank") + + @property + def world_size(self) -> int: + """Total number of processes (1 single-process).""" + return self._get("world_size") + + @property + def local_rank(self) -> int: + """Node-local rank of this process (0 single-process).""" + return self._get("local_rank") + + @property + def device(self): + """This process's torch device (``cuda:local_rank`` distributed; the caller default otherwise).""" + return self._get("device") + + @property + def is_main(self) -> bool: + """Whether this is the main process (rank 0).""" + return self._get("is_main") + + +# The process-wide ACTIVE session: a single STABLE object created once at import and NEVER rebound. +# :func:`modelopt_session` mutates it in place via :meth:`ModelOptSession.initialize` / ``_reset`` rather +# than assigning a new object -- so any held reference stays valid and never goes stale. Module-PRIVATE +# (not exported): reach it only through :func:`get_active_mo_session`. It exists before initialization; +# reading its IDENTITY fields (``rank`` etc.) then raises via the per-field guard. +_active_mo_session: ModelOptSession = ModelOptSession() + + +def get_active_mo_session(required: bool = True) -> ModelOptSession: + """Return the process-wide session. + + * ``required=True`` (default) -- fail-fast: raise if no session has been entered yet by + :func:`modelopt_session`. Use when the caller genuinely needs a live session. + * ``required=False`` -- return the singleton regardless of initialization. Its ``ep_mesh`` is + readable (``None``) even before a session is entered, so this is for defensive reads deep in public + APIs (export / calibration) that may run with no active session. + """ + if required and not _active_mo_session.initialized: + raise RuntimeError( + "no active session; enter modelopt.torch.utils.distributed.modelopt_session() first" + ) + return _active_mo_session + + +@contextmanager +def modelopt_session(timeout: timedelta | None = None, device=None): + """Populate the active session for the block, then tear down. + + On a distributed launch (WORLD_SIZE>1) this initializes the process group (binding the local CUDA + device) with ``timeout`` and records rank/world_size/local_rank/``cuda:local_rank``/is_main; it + DESTROYS the group on exit even if the body raises. Single-process records the trivial single-rank + topology -- with ``device`` as the device (the caller's default, e.g. the ``--device`` CLI value, since + there is no local rank to derive one) -- and enter/exit are no-ops around the body. + + Yields the session (also reachable via :func:`get_active_mo_session`); it is reset on exit. + """ + owns_group = is_distributed_launch() + if owns_group: + setup(timeout=timeout) + _active_mo_session.initialize( + distributed=True, + rank=rank(), + world_size=size(), + local_rank=local_rank(), + device=torch.device(f"cuda:{local_rank()}"), + is_main=rank() == 0, + ) + else: + _active_mo_session.initialize( + distributed=False, + rank=0, + world_size=1, + local_rank=0, + device=torch.device(device) if device is not None else None, + is_main=True, + ) + try: + yield _active_mo_session + finally: + # If the body is unwinding due to an exception, PRINT it here BEFORE cleanup()'s collective + # barrier. A per-rank exception (only some ranks raise) otherwise makes those ranks enter the + # cleanup barrier while the rest are still inside a calibration a2a collective -> cross-rank + # DEADLOCK, with the original exception never surfaced (pending behind the hung barrier). + exc = sys.exc_info()[1] + if exc is not None: + print( + f"[modelopt_session][rank {os.environ.get('RANK', '?')}] body raised " + f"{type(exc).__name__}: {exc}", + flush=True, + ) + traceback.print_exc() + sys.stderr.flush() + if owns_group: + cleanup() + _active_mo_session._reset() + + def cleanup(): """Cleans up the distributed environment. diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_session.py new file mode 100644 index 00000000000..2ee9f70c707 --- /dev/null +++ b/tests/unit/torch/utils/test_distributed_session.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the distributed PTQ session: the distribution plan, the session handle, launch detection.""" + +import pytest + +from modelopt.torch.utils.distributed import ( + DistributionPlan, + ModelOptSession, + get_active_mo_session, + is_distributed_launch, +) + +# --------------------------------------------------------------------------------- DistributionPlan + + +def test_plain_world_replicates_everything(): + plan = DistributionPlan(world_size=8) + assert plan.edp == 8 + assert not plan.is_expert_parallel + assert not plan.is_fsdp_dense + + +def test_expert_parallel_splits_the_expert_replica_pool(): + plan = DistributionPlan(world_size=8, ep_size=2) + assert plan.edp == 4 # each expert shard is replicated 4x + assert plan.is_expert_parallel + + +def test_fsdp_dense_plan(): + plan = DistributionPlan(world_size=8, ep_size=8, fsdp_size=8) + assert plan.is_fsdp_dense + assert plan.edp == 1 + + +def test_an_ep_degree_that_does_not_divide_the_world_is_rejected(): + """An unsatisfiable plan must fail here, not as a confusing mesh error thousands of ranks later.""" + with pytest.raises(ValueError, match="must divide world_size"): + DistributionPlan(world_size=8, ep_size=3) + + +def test_repr_names_the_degrees(): + assert repr(DistributionPlan(world_size=8, ep_size=2)) == ( + "DistributionPlan(world=8, ep_size=2, edp=4, dense=replicate)" + ) + assert "fsdp_size=8" in repr(DistributionPlan(world_size=8, ep_size=8, fsdp_size=8)) + + +# ------------------------------------------------------------------------------------- session handle + + +def test_required_session_fails_fast_when_none_was_entered(): + with pytest.raises(RuntimeError, match="no active session"): + get_active_mo_session() + + +def test_optional_session_is_readable_without_a_session(): + """Export / calibration read the session defensively; they must not explode with no session.""" + session = get_active_mo_session(required=False) + assert session is not None + assert not session.initialized + assert session.ep_mesh is None + + +def test_identity_reads_raise_before_initialization(): + """A stale or zero default would silently mislabel every rank; reading early must fail instead.""" + with pytest.raises(RuntimeError, match="before initialization"): + _ = ModelOptSession().rank + + +def test_reset_clears_the_identity_and_the_ep_mesh(): + session = ModelOptSession() + session.initialize( + distributed=True, rank=1, world_size=2, local_rank=1, device=None, is_main=False + ) + session.ep_mesh = object() + assert (session.rank, session.world_size, session.is_main) == (1, 2, False) + session._reset() + assert not session.initialized + assert session.ep_mesh is None + + +# ------------------------------------------------------------------------------------ launch detect + + +@pytest.mark.parametrize(("world", "expected"), [(None, False), ("1", False), ("8", True)]) +def test_launch_detection_reads_the_env_before_any_process_group(monkeypatch, world, expected): + """Must work at parse time, before modelopt_session exists.""" + monkeypatch.delenv("WORLD_SIZE", raising=False) + if world is not None: + monkeypatch.setenv("WORLD_SIZE", world) + assert is_distributed_launch() is expected From df59de9c69175947a9e249c8609dbb1c9668d7c8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 17:56:42 +0000 Subject: [PATCH 02/48] Keep synced session attributes the same on every rank An attribute declared on ModelOptSession with synced() holds the same value on every rank of a group, all of them by default. Reading it is local; assigning it is collective, and every rank keeps the value the group's first rank assigned, so no rank drifts from the rest. session.decide(name, fn) sets one from fn() run on the first rank only, for a value that depends on what each rank reads for itself; if fn raises, every rank raises instead of waiting in a broadcast the first rank never reaches. Ranks that set different attributes at the same point fail instead of swapping values. Outside a distributed session a synced attribute is a plain one. The session's distribution plan is the first: session.plan. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 98 +++++++++++++ .../torch/utils/test_distributed_session.py | 132 +++++++++++++++++- 2 files changed, 229 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index a00c661ab3a..df1e17a6fa3 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -15,6 +15,7 @@ """Utility functions for using torch.distributed.""" +import copy import functools import io import os @@ -37,6 +38,7 @@ "DistributionPlan", "ModelOptSession", "ParallelState", + "SyncedAttribute", "backend", "barrier", "fsdp2_wrap", @@ -52,6 +54,7 @@ "modelopt_session", "rank", "size", + "synced", ] @@ -337,6 +340,44 @@ def __repr__(self) -> str: ) +class SyncedAttribute: + """A :class:`ModelOptSession` attribute every rank of a group holds the same value of. + + Declare one with :func:`synced`. Reading it is local. Assigning it is collective: every rank of + the group assigns it at the same point, and each then holds the value the group's first rank + assigned -- whatever the others passed -- so no rank drifts from the rest. + :meth:`ModelOptSession.decide` sets one from a function run on the first rank only. Assign a new + value rather than mutating the one held, which would change this rank's copy only. Outside a + distributed session it is a plain attribute. + """ + + def __init__(self, group: "Callable[[ModelOptSession], Any] | None" = None) -> None: + """Hold the attribute on ``group`` (see :func:`synced`).""" + self.group = group + self.name = "" + + def __set_name__(self, owner: type, name: str) -> None: + self.name = name + + def __get__(self, session, owner=None): + if session is None: + return self + return session._synced.get(self.name) + + def __set__(self, session, value) -> None: + session._sync(self, lambda: value, decide=False) + + +def synced(group: "Callable[[ModelOptSession], Any] | None" = None) -> SyncedAttribute: + """Declare a :class:`ModelOptSession` attribute that stays the same on every rank of a group. + + ``group`` picks the ranks: ``None`` for all of them, or a callable taking the session and returning + the ``ProcessGroup`` that holds the attribute, resolved at each assignment. See + :class:`SyncedAttribute`. + """ + return SyncedAttribute(group) + + class ModelOptSession: """The live distributed session of a PTQ run: this process's identity, plus the model's EP mesh. @@ -349,11 +390,17 @@ class ModelOptSession: group is set up. * ``ep_mesh`` -- the expert-parallel device mesh of the loaded model, set by ``load_and_prepare_ep`` (``None`` until then, and for a model loaded without expert parallelism). + + Attributes declared with :func:`synced`, such as ``plan``, hold the same value on every rank. """ + #: The :class:`DistributionPlan` of the model this session runs, the same on every rank. + plan = synced() + def __init__(self) -> None: """Create an empty (uninitialized) session; populate identity with :meth:`initialize`.""" self._values: dict[str, Any] | None = None + self._synced: dict[str, Any] = {} self.ep_mesh = None def initialize( @@ -378,8 +425,59 @@ def initialize( def _reset(self) -> None: self._values = None + self._synced = {} self.ep_mesh = None + def decide(self, name: str, decide_fn: Callable[[], Any]) -> Any: + """Set the synced attribute ``name`` to ``decide_fn()``, run on its group's first rank only. + + The other ranks take that value instead of computing their own, which could differ when it + depends on what each rank reads for itself: the environment, its GPU, files on its node. If + ``decide_fn`` raises, every rank of the group raises, instead of waiting in a collective the + first rank never joins. Collective, like assigning the attribute. Returns the value. + """ + attribute = getattr(type(self), name, None) + if not isinstance(attribute, SyncedAttribute): + raise AttributeError(f"ModelOptSession.{name} is not a synced attribute") + self._sync(attribute, decide_fn, decide=True) + return self._synced[name] + + def _sync( + self, attribute: SyncedAttribute, value_fn: Callable[[], Any], *, decide: bool + ) -> None: + """Store ``value_fn()`` from the group's first rank under ``attribute`` on every rank.""" + name = attribute.name + if not (self.initialized and self.distributed): + self._synced[name] = value_fn() + return + group = None if attribute.group is None else attribute.group(self) + first = torch.distributed.get_rank(group) == 0 + value, error = None, None + if first or not decide: + try: + value = value_fn() + except Exception as e: + if not first: + raise + error = e + # One message from the first rank: which attribute it set, so a rank setting another learns the + # two are out of step, and its value or its error. The error goes as text: the other ranks only + # need to stop, and an exception is not guaranteed to pickle. + message = [(name, value, None if error is None else f"{type(error).__name__}: {error}")] + torch.distributed.broadcast_object_list(message, group=group, group_src=0) + first_name, value, first_error = message[0] + if error is not None: + raise error + if first_error is not None: + raise RuntimeError(f"rank 0 failed to decide ModelOptSession.{name}: {first_error}") + if first_name != name: + raise RuntimeError( + f"Ranks set synced session attributes out of step: this rank set {name!r} while rank " + f"0 of its group set {first_name!r}. Every rank must set them in the same order." + ) + # The first rank's own object stays the caller's; a copy keeps a later mutation of it local. + self._synced[name] = copy.deepcopy(value) if first else value + @property def initialized(self) -> bool: """Whether ``modelopt_session`` has populated the state.""" diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_session.py index 2ee9f70c707..cf1d4f2a4d8 100644 --- a/tests/unit/torch/utils/test_distributed_session.py +++ b/tests/unit/torch/utils/test_distributed_session.py @@ -13,15 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the distributed PTQ session: the distribution plan, the session handle, launch detection.""" +"""Unit tests for the distributed PTQ session: the distribution plan, the session handle, its synced +attributes, launch detection.""" import pytest +import torch.distributed from modelopt.torch.utils.distributed import ( DistributionPlan, ModelOptSession, get_active_mo_session, is_distributed_launch, + synced, ) # --------------------------------------------------------------------------------- DistributionPlan @@ -93,6 +96,133 @@ def test_reset_clears_the_identity_and_the_ep_mesh(): assert session.ep_mesh is None +# ------------------------------------------------------------------------------- synced attributes + + +class _Session(ModelOptSession): + """A session with a second synced attribute, and one held on a group of its own.""" + + other = synced() + grouped = synced(group=lambda session: "ep-group") + + +def _session(rank=None): + """Rank ``rank`` of a two-rank session; ``None``: a single process.""" + session = _Session() + session.initialize( + distributed=rank is not None, + rank=rank or 0, + world_size=1 if rank is None else 2, + local_rank=rank or 0, + device=None, + is_main=not rank, + ) + return session + + +def _as_rank(monkeypatch, rank, rank0_message=None): + """Answer the session's collectives as ``rank``; returns what each broadcast was handed. + + A non-zero rank receives ``rank0_message``, as the group's first rank would have sent it. + """ + sent = [] + + def broadcast_object_list(message, group=None, group_src=None): + assert group_src == 0 + sent.append((message[0], group)) + if rank != 0: + message[0] = rank0_message + + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) + return sent + + +def test_a_synced_attribute_is_plain_in_a_single_process(monkeypatch): + """Nothing to sync with: no collective, and an uninitialized session behaves the same.""" + monkeypatch.setattr( + torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("a collective ran") + ) + for session in (_session(), _Session()): + plan = DistributionPlan(world_size=1) + session.plan = plan + assert session.plan is plan + assert session.decide("other", lambda: 3) == 3 + + +def test_rank0_sends_what_it_assigns_and_keeps_its_own_copy(monkeypatch): + sent = _as_rank(monkeypatch, 0) + session, plan = _session(rank=0), DistributionPlan(world_size=2, ep_size=2) + + session.plan = plan + + assert sent == [(("plan", plan, None), None)] + assert session.plan.ep_size == 2 + assert session.plan is not plan # mutating the caller's object must not reach the session + + +def test_other_ranks_hold_what_rank0_assigned(monkeypatch): + """Whatever a non-zero rank passes, it ends up with rank 0's value, as every rank does.""" + rank0_plan = DistributionPlan(world_size=2, ep_size=2) + _as_rank(monkeypatch, 1, rank0_message=("plan", rank0_plan, None)) + session = _session(rank=1) + + session.plan = DistributionPlan(world_size=2) + + assert session.plan is rank0_plan + + +def test_decide_runs_only_on_rank0(monkeypatch): + _as_rank(monkeypatch, 1, rank0_message=("other", 7, None)) + session = _session(rank=1) + + assert session.decide("other", lambda: pytest.fail("a non-zero rank decided")) == 7 + assert session.other == 7 + + +def test_rank0_failing_to_decide_still_reaches_the_broadcast(monkeypatch): + """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it.""" + sent = _as_rank(monkeypatch, 0) + + def fail(): + raise ValueError("bad degrees") + + with pytest.raises(ValueError, match="bad degrees"): + _session(rank=0).decide("other", fail) + assert sent == [(("other", None, "ValueError: bad degrees"), None)] + + +def test_other_ranks_raise_when_rank0_failed(monkeypatch): + _as_rank(monkeypatch, 1, rank0_message=("other", None, "ValueError: bad degrees")) + with pytest.raises(RuntimeError, match=r"rank 0 failed to decide ModelOptSession.other.*bad"): + _session(rank=1).decide("other", lambda: 0) + + +def test_ranks_setting_different_attributes_fail_instead_of_mixing_them(monkeypatch): + """Rank 0 setting one attribute while rank 1 sets another would cross the two values.""" + _as_rank(monkeypatch, 1, rank0_message=("other", 7, None)) + with pytest.raises(RuntimeError, match="out of step"): + _session(rank=1).plan = DistributionPlan(world_size=2) + + +def test_a_synced_attribute_is_held_on_its_group(monkeypatch): + sent = _as_rank(monkeypatch, 0) + _session(rank=0).grouped = 5 + assert sent == [(("grouped", 5, None), "ep-group")] + + +def test_only_synced_attributes_can_be_decided(): + with pytest.raises(AttributeError, match="not a synced attribute"): + _session().decide("ep_mesh", lambda: None) + + +def test_reset_clears_synced_attributes(): + session = _session() + session.plan = DistributionPlan(world_size=1) + session._reset() + assert session.plan is None + + # ------------------------------------------------------------------------------------ launch detect From e9f72dbc0119ec3cd77f4250e6965dce24c0b9f5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 22:17:27 +0000 Subject: [PATCH 03/48] Let a synced session attribute combine every rank's value A synced attribute so far took its group's first rank's value. Some facts need every rank's: a calibration batch size each rank's memory bounds must be the smallest of them, or rank 0's choice runs another rank out of memory. synced(reduce=fn) gathers the group's values and holds fn(values), the same on every rank; session.decide runs the function on every rank for such an attribute. A rank whose value fails still joins the gather and the others raise naming it, and ranks setting different attributes at the same point fail, as for the first-rank attributes. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 68 ++++++++++++--- .../torch/utils/test_distributed_session.py | 82 ++++++++++++++++++- 2 files changed, 134 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index df1e17a6fa3..4bc91cd0512 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -345,15 +345,21 @@ class SyncedAttribute: Declare one with :func:`synced`. Reading it is local. Assigning it is collective: every rank of the group assigns it at the same point, and each then holds the value the group's first rank - assigned -- whatever the others passed -- so no rank drifts from the rest. - :meth:`ModelOptSession.decide` sets one from a function run on the first rank only. Assign a new - value rather than mutating the one held, which would change this rank's copy only. Outside a - distributed session it is a plain attribute. + assigned -- whatever the others passed -- so no rank drifts from the rest. A reduced attribute + holds the group's values combined instead, the same on every rank. + :meth:`ModelOptSession.decide` sets one from a function run on the first rank only, or on every + rank for a reduced attribute. Assign a new value rather than mutating the one held, which would + change this rank's copy only. Outside a distributed session it is a plain attribute. """ - def __init__(self, group: "Callable[[ModelOptSession], Any] | None" = None) -> None: - """Hold the attribute on ``group`` (see :func:`synced`).""" + def __init__( + self, + group: "Callable[[ModelOptSession], Any] | None" = None, + reduce: "Callable[[list[Any]], Any] | None" = None, + ) -> None: + """Hold the attribute on ``group``, combining the ranks' values with ``reduce`` if given.""" self.group = group + self.reduce = reduce self.name = "" def __set_name__(self, owner: type, name: str) -> None: @@ -368,14 +374,19 @@ def __set__(self, session, value) -> None: session._sync(self, lambda: value, decide=False) -def synced(group: "Callable[[ModelOptSession], Any] | None" = None) -> SyncedAttribute: +def synced( + group: "Callable[[ModelOptSession], Any] | None" = None, + reduce: "Callable[[list[Any]], Any] | None" = None, +) -> SyncedAttribute: """Declare a :class:`ModelOptSession` attribute that stays the same on every rank of a group. ``group`` picks the ranks: ``None`` for all of them, or a callable taking the session and returning - the ``ProcessGroup`` that holds the attribute, resolved at each assignment. See - :class:`SyncedAttribute`. + the ``ProcessGroup`` that holds the attribute, resolved at each assignment. ``reduce`` makes every + rank's value count: it takes the group's values in rank order and returns the one every rank then + holds, e.g. ``min`` for a batch size each rank's memory bounds. Without it, the group's first rank + decides. See :class:`SyncedAttribute`. """ - return SyncedAttribute(group) + return SyncedAttribute(group, reduce) class ModelOptSession: @@ -432,9 +443,10 @@ def decide(self, name: str, decide_fn: Callable[[], Any]) -> Any: """Set the synced attribute ``name`` to ``decide_fn()``, run on its group's first rank only. The other ranks take that value instead of computing their own, which could differ when it - depends on what each rank reads for itself: the environment, its GPU, files on its node. If + depends on what each rank reads for itself: the environment, its GPU, files on its node. For a + reduced attribute ``decide_fn`` runs on every rank instead, and the results are combined. If ``decide_fn`` raises, every rank of the group raises, instead of waiting in a collective the - first rank never joins. Collective, like assigning the attribute. Returns the value. + failing rank never joins. Collective, like assigning the attribute. Returns the value. """ attribute = getattr(type(self), name, None) if not isinstance(attribute, SyncedAttribute): @@ -451,6 +463,9 @@ def _sync( self._synced[name] = value_fn() return group = None if attribute.group is None else attribute.group(self) + if attribute.reduce is not None: + self._synced[name] = self._reduce(name, attribute.reduce, value_fn, group) + return first = torch.distributed.get_rank(group) == 0 value, error = None, None if first or not decide: @@ -478,6 +493,35 @@ def _sync( # The first rank's own object stays the caller's; a copy keeps a later mutation of it local. self._synced[name] = copy.deepcopy(value) if first else value + def _reduce( + self, name: str, reduce: Callable[[list[Any]], Any], value_fn: Callable[[], Any], group + ) -> Any: + """Combine ``value_fn()`` from every rank of ``group`` with ``reduce``.""" + value, error = None, None + try: + value = value_fn() + except Exception as e: + # Still join the gather: the other ranks would otherwise wait in it for this one. + error = e + mine = (name, value, None if error is None else f"{type(error).__name__}: {error}") + gathered: list[Any] = [None] * torch.distributed.get_world_size(group) + torch.distributed.all_gather_object(gathered, mine, group=group) + if error is not None: + raise error + for rank, (_, _, rank_error) in enumerate(gathered): + if rank_error is not None: + raise RuntimeError( + f"rank {rank} failed to set ModelOptSession.{name}: {rank_error}" + ) + names = sorted({rank_name for rank_name, _, _ in gathered}) + if names != [name]: + raise RuntimeError( + f"Ranks set synced session attributes out of step: {names} at the same point. " + "Every rank must set them in the same order." + ) + # The gathered values are unpickled copies, so the result shares nothing with the callers'. + return reduce([rank_value for _, rank_value, _ in gathered]) + @property def initialized(self) -> bool: """Whether ``modelopt_session`` has populated the state.""" diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_session.py index cf1d4f2a4d8..307e2218126 100644 --- a/tests/unit/torch/utils/test_distributed_session.py +++ b/tests/unit/torch/utils/test_distributed_session.py @@ -100,10 +100,12 @@ def test_reset_clears_the_identity_and_the_ep_mesh(): class _Session(ModelOptSession): - """A session with a second synced attribute, and one held on a group of its own.""" + """A session with a second synced attribute, one held on a group of its own, and reduced ones.""" other = synced() grouped = synced(group=lambda session: "ep-group") + smallest = synced(reduce=min) + grouped_smallest = synced(group=lambda session: "ep-group", reduce=min) def _session(rank=None): @@ -138,16 +140,38 @@ def broadcast_object_list(message, group=None, group_src=None): return sent -def test_a_synced_attribute_is_plain_in_a_single_process(monkeypatch): - """Nothing to sync with: no collective, and an uninitialized session behaves the same.""" +def _gather_as_rank(monkeypatch, rank, others): + """Answer the session's gathers as ``rank``; returns what each gather was handed. + + ``others`` are the other ranks' messages in rank order, as they would have sent them. + """ + sent = [] + + def all_gather_object(gathered, mine, group=None): + sent.append((mine, group)) + gathered[:] = [*others[:rank], mine, *others[rank:]] + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: len(others) + 1) + monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) monkeypatch.setattr( - torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("a collective ran") + torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("rank 0 decided") ) + return sent + + +def test_a_synced_attribute_is_plain_in_a_single_process(monkeypatch): + """Nothing to sync with: no collective, and an uninitialized session behaves the same.""" + for collective in ("broadcast_object_list", "all_gather_object"): + monkeypatch.setattr( + torch.distributed, collective, lambda *a, **k: pytest.fail("a collective ran") + ) for session in (_session(), _Session()): plan = DistributionPlan(world_size=1) session.plan = plan assert session.plan is plan assert session.decide("other", lambda: 3) == 3 + session.smallest = 4 + assert session.smallest == 4 def test_rank0_sends_what_it_assigns_and_keeps_its_own_copy(monkeypatch): @@ -211,6 +235,56 @@ def test_a_synced_attribute_is_held_on_its_group(monkeypatch): assert sent == [(("grouped", 5, None), "ep-group")] +def test_a_reduced_attribute_holds_the_groups_combined_value(monkeypatch): + """Every rank's value counts, and every rank ends up holding the same result.""" + sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None), ("smallest", 2, None)]) + session = _session(rank=1) + + session.smallest = 4 + + assert sent == [(("smallest", 4, None), None)] + assert session.smallest == 2 + + +def test_decide_runs_on_every_rank_for_a_reduced_attribute(monkeypatch): + _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) + assert _session(rank=1).decide("smallest", lambda: 3) == 3 + + +def test_a_rank_failing_a_reduced_attribute_still_joins_the_gather(monkeypatch): + """A failing rank must still reach the gather, or every other rank waits in it.""" + sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) + + def fail(): + raise ValueError("no memory reading") + + with pytest.raises(ValueError, match="no memory reading"): + _session(rank=1).decide("smallest", fail) + assert sent == [(("smallest", None, "ValueError: no memory reading"), None)] + + +def test_the_other_ranks_raise_when_one_failed_a_reduced_attribute(monkeypatch): + _gather_as_rank(monkeypatch, 0, others=[("smallest", None, "ValueError: no memory reading")]) + with pytest.raises( + RuntimeError, match=r"rank 1 failed to set ModelOptSession.smallest.*memory" + ): + _session(rank=0).smallest = 4 + + +def test_reduced_attributes_set_out_of_step_fail(monkeypatch): + _gather_as_rank(monkeypatch, 0, others=[("other", 8, None)]) + with pytest.raises(RuntimeError, match="out of step"): + _session(rank=0).smallest = 4 + + +def test_a_reduced_attribute_is_held_on_its_group(monkeypatch): + sent = _gather_as_rank(monkeypatch, 0, others=[("grouped_smallest", 1, None)]) + session = _session(rank=0) + session.grouped_smallest = 5 + assert sent == [(("grouped_smallest", 5, None), "ep-group")] + assert session.grouped_smallest == 1 + + def test_only_synced_attributes_can_be_decided(): with pytest.raises(AttributeError, match="not a synced attribute"): _session().decide("ep_mesh", lambda: None) From 048dbc2de047ccd42c174786f8cb6454f5e39a10 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 00:34:36 +0000 Subject: [PATCH 04/48] Agree on a one-off value across ranks without keeping it on the session A synced attribute is for a fact a later stage reads back from the session. The reduce mode stored a value nothing read back: the only reduced attribute, the calibration batch size, is used where it is decided. ModelOptSession.agree(label, fn, combine) runs fn on every rank and returns the combined value, keeping nothing: with combine, e.g. min, every rank's value counts; without it, every rank must pass the same value, and a rank that differs makes every rank raise. As before, a rank whose fn fails still joins the gather and the others raise naming it, and ranks agreeing out of step fail. synced() keeps the first rank's decision only. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 90 +++++++++++-------- .../torch/utils/test_distributed_session.py | 65 ++++++++------ 2 files changed, 88 insertions(+), 67 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 4bc91cd0512..5e5b9e3a9f9 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -343,23 +343,18 @@ def __repr__(self) -> str: class SyncedAttribute: """A :class:`ModelOptSession` attribute every rank of a group holds the same value of. - Declare one with :func:`synced`. Reading it is local. Assigning it is collective: every rank of - the group assigns it at the same point, and each then holds the value the group's first rank - assigned -- whatever the others passed -- so no rank drifts from the rest. A reduced attribute - holds the group's values combined instead, the same on every rank. - :meth:`ModelOptSession.decide` sets one from a function run on the first rank only, or on every - rank for a reduced attribute. Assign a new value rather than mutating the one held, which would - change this rank's copy only. Outside a distributed session it is a plain attribute. + Declare one with :func:`synced`, for a fact a later stage reads back from the session; a one-off + decision that nothing reads back is :meth:`ModelOptSession.agree`. Reading it is local. Assigning + it is collective: every rank of the group assigns it at the same point, and each then holds the + value the group's first rank assigned -- whatever the others passed -- so no rank drifts from the + rest. :meth:`ModelOptSession.decide` sets one from a function run on the first rank only. Assign a + new value rather than mutating the one held, which would change this rank's copy only. Outside a + distributed session it is a plain attribute. """ - def __init__( - self, - group: "Callable[[ModelOptSession], Any] | None" = None, - reduce: "Callable[[list[Any]], Any] | None" = None, - ) -> None: - """Hold the attribute on ``group``, combining the ranks' values with ``reduce`` if given.""" + def __init__(self, group: "Callable[[ModelOptSession], Any] | None" = None) -> None: + """Hold the attribute on ``group``.""" self.group = group - self.reduce = reduce self.name = "" def __set_name__(self, owner: type, name: str) -> None: @@ -374,19 +369,24 @@ def __set__(self, session, value) -> None: session._sync(self, lambda: value, decide=False) -def synced( - group: "Callable[[ModelOptSession], Any] | None" = None, - reduce: "Callable[[list[Any]], Any] | None" = None, -) -> SyncedAttribute: +def synced(group: "Callable[[ModelOptSession], Any] | None" = None) -> SyncedAttribute: """Declare a :class:`ModelOptSession` attribute that stays the same on every rank of a group. ``group`` picks the ranks: ``None`` for all of them, or a callable taking the session and returning - the ``ProcessGroup`` that holds the attribute, resolved at each assignment. ``reduce`` makes every - rank's value count: it takes the group's values in rank order and returns the one every rank then - holds, e.g. ``min`` for a batch size each rank's memory bounds. Without it, the group's first rank + the ``ProcessGroup`` that holds the attribute, resolved at each assignment. The group's first rank decides. See :class:`SyncedAttribute`. """ - return SyncedAttribute(group, reduce) + return SyncedAttribute(group) + + +def _agreed(label: str, values: list[Any]) -> Any: + """The value every rank passed to :meth:`ModelOptSession.agree`; raise if any rank passed another.""" + for rank, value in enumerate(values): + if value != values[0]: + raise RuntimeError( + f"Ranks disagree on {label}: rank 0 has {values[0]!r}, rank {rank} has {value!r}." + ) + return values[0] class ModelOptSession: @@ -443,8 +443,7 @@ def decide(self, name: str, decide_fn: Callable[[], Any]) -> Any: """Set the synced attribute ``name`` to ``decide_fn()``, run on its group's first rank only. The other ranks take that value instead of computing their own, which could differ when it - depends on what each rank reads for itself: the environment, its GPU, files on its node. For a - reduced attribute ``decide_fn`` runs on every rank instead, and the results are combined. If + depends on what each rank reads for itself: the environment, its GPU, files on its node. If ``decide_fn`` raises, every rank of the group raises, instead of waiting in a collective the failing rank never joins. Collective, like assigning the attribute. Returns the value. """ @@ -463,9 +462,6 @@ def _sync( self._synced[name] = value_fn() return group = None if attribute.group is None else attribute.group(self) - if attribute.reduce is not None: - self._synced[name] = self._reduce(name, attribute.reduce, value_fn, group) - return first = torch.distributed.get_rank(group) == 0 value, error = None, None if first or not decide: @@ -493,34 +489,50 @@ def _sync( # The first rank's own object stays the caller's; a copy keeps a later mutation of it local. self._synced[name] = copy.deepcopy(value) if first else value - def _reduce( - self, name: str, reduce: Callable[[list[Any]], Any], value_fn: Callable[[], Any], group + def agree( + self, + label: str, + value_fn: Callable[[], Any], + combine: "Callable[[list[Any]], Any] | None" = None, + group=None, ) -> Any: - """Combine ``value_fn()`` from every rank of ``group`` with ``reduce``.""" + """Return what every rank of ``group`` agrees on from ``value_fn()``, run on each; keep nothing. + + For a decision no later stage reads back from the session (a fact that one does is a + :func:`synced` attribute). ``combine`` takes the group's values in rank order and returns the + one every rank gets, e.g. ``min`` for a batch size each rank's memory bounds, where one rank's + choice could run another out of memory. Without it every rank must pass the same value, e.g. + the library versions each runs: a rank that differs makes every rank raise, where adopting + one rank's value would hide the difference until a collective hangs on it. If ``value_fn`` + raises on any rank, every rank raises, instead of waiting in the gather that rank never joins. + Ranks agreeing on different labels at the same point fail. Collective over ``group``, every + rank by default; outside a distributed session ``value_fn`` runs alone. + """ + combine = combine or functools.partial(_agreed, label) + if not (self.initialized and self.distributed): + return combine([value_fn()]) value, error = None, None try: value = value_fn() except Exception as e: # Still join the gather: the other ranks would otherwise wait in it for this one. error = e - mine = (name, value, None if error is None else f"{type(error).__name__}: {error}") + mine = (label, value, None if error is None else f"{type(error).__name__}: {error}") gathered: list[Any] = [None] * torch.distributed.get_world_size(group) torch.distributed.all_gather_object(gathered, mine, group=group) if error is not None: raise error for rank, (_, _, rank_error) in enumerate(gathered): if rank_error is not None: - raise RuntimeError( - f"rank {rank} failed to set ModelOptSession.{name}: {rank_error}" - ) - names = sorted({rank_name for rank_name, _, _ in gathered}) - if names != [name]: + raise RuntimeError(f"rank {rank} failed to agree on {label}: {rank_error}") + labels = sorted({rank_label for rank_label, _, _ in gathered}) + if labels != [label]: raise RuntimeError( - f"Ranks set synced session attributes out of step: {names} at the same point. " - "Every rank must set them in the same order." + f"Ranks agreed out of step: {labels} at the same point. Every rank must agree on " + "them in the same order." ) # The gathered values are unpickled copies, so the result shares nothing with the callers'. - return reduce([rank_value for _, rank_value, _ in gathered]) + return combine([rank_value for _, rank_value, _ in gathered]) @property def initialized(self) -> bool: diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_session.py index 307e2218126..6a30b2efd2b 100644 --- a/tests/unit/torch/utils/test_distributed_session.py +++ b/tests/unit/torch/utils/test_distributed_session.py @@ -100,12 +100,10 @@ def test_reset_clears_the_identity_and_the_ep_mesh(): class _Session(ModelOptSession): - """A session with a second synced attribute, one held on a group of its own, and reduced ones.""" + """A session with a second synced attribute, and one held on a group of its own.""" other = synced() grouped = synced(group=lambda session: "ep-group") - smallest = synced(reduce=min) - grouped_smallest = synced(group=lambda session: "ep-group", reduce=min) def _session(rank=None): @@ -170,8 +168,8 @@ def test_a_synced_attribute_is_plain_in_a_single_process(monkeypatch): session.plan = plan assert session.plan is plan assert session.decide("other", lambda: 3) == 3 - session.smallest = 4 - assert session.smallest == 4 + assert session.agree("smallest", lambda: 4, min) == 4 + assert session.agree("versions", lambda: "2.6") == "2.6" def test_rank0_sends_what_it_assigns_and_keeps_its_own_copy(monkeypatch): @@ -235,23 +233,38 @@ def test_a_synced_attribute_is_held_on_its_group(monkeypatch): assert sent == [(("grouped", 5, None), "ep-group")] -def test_a_reduced_attribute_holds_the_groups_combined_value(monkeypatch): - """Every rank's value counts, and every rank ends up holding the same result.""" +def test_agree_gives_every_rank_the_combined_value(monkeypatch): + """Every rank's value counts, and every rank gets the same result.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None), ("smallest", 2, None)]) - session = _session(rank=1) - - session.smallest = 4 + assert _session(rank=1).agree("smallest", lambda: 4, min) == 2 assert sent == [(("smallest", 4, None), None)] - assert session.smallest == 2 -def test_decide_runs_on_every_rank_for_a_reduced_attribute(monkeypatch): +def test_agree_keeps_nothing_on_the_session(monkeypatch): _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) - assert _session(rank=1).decide("smallest", lambda: 3) == 3 + session = _session(rank=1) + + session.agree("smallest", lambda: 4, min) + + assert session._synced == {} + + +def test_without_combine_every_rank_must_pass_the_same_value(monkeypatch): + _gather_as_rank(monkeypatch, 1, others=[("versions", "2.6", None)]) + assert _session(rank=1).agree("versions", lambda: "2.6") == "2.6" -def test_a_rank_failing_a_reduced_attribute_still_joins_the_gather(monkeypatch): +def test_ranks_that_disagree_all_raise(monkeypatch): + """Adopting rank 0's value instead would hide the difference until a collective hangs on it.""" + _gather_as_rank(monkeypatch, 0, others=[("versions", "2.5", None)]) + with pytest.raises( + RuntimeError, match=r"disagree on versions: rank 0 has '2.6', rank 1 has '2.5'" + ): + _session(rank=0).agree("versions", lambda: "2.6") + + +def test_a_rank_failing_to_agree_still_joins_the_gather(monkeypatch): """A failing rank must still reach the gather, or every other rank waits in it.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) @@ -259,30 +272,26 @@ def fail(): raise ValueError("no memory reading") with pytest.raises(ValueError, match="no memory reading"): - _session(rank=1).decide("smallest", fail) + _session(rank=1).agree("smallest", fail, min) assert sent == [(("smallest", None, "ValueError: no memory reading"), None)] -def test_the_other_ranks_raise_when_one_failed_a_reduced_attribute(monkeypatch): +def test_the_other_ranks_raise_when_one_failed_to_agree(monkeypatch): _gather_as_rank(monkeypatch, 0, others=[("smallest", None, "ValueError: no memory reading")]) - with pytest.raises( - RuntimeError, match=r"rank 1 failed to set ModelOptSession.smallest.*memory" - ): - _session(rank=0).smallest = 4 + with pytest.raises(RuntimeError, match=r"rank 1 failed to agree on smallest.*memory"): + _session(rank=0).agree("smallest", lambda: 4, min) -def test_reduced_attributes_set_out_of_step_fail(monkeypatch): +def test_ranks_agreeing_out_of_step_fail(monkeypatch): _gather_as_rank(monkeypatch, 0, others=[("other", 8, None)]) with pytest.raises(RuntimeError, match="out of step"): - _session(rank=0).smallest = 4 + _session(rank=0).agree("smallest", lambda: 4, min) -def test_a_reduced_attribute_is_held_on_its_group(monkeypatch): - sent = _gather_as_rank(monkeypatch, 0, others=[("grouped_smallest", 1, None)]) - session = _session(rank=0) - session.grouped_smallest = 5 - assert sent == [(("grouped_smallest", 5, None), "ep-group")] - assert session.grouped_smallest == 1 +def test_agree_gathers_on_its_group(monkeypatch): + sent = _gather_as_rank(monkeypatch, 0, others=[("smallest", 1, None)]) + assert _session(rank=0).agree("smallest", lambda: 5, min, group="ep-group") == 1 + assert sent == [(("smallest", 5, None), "ep-group")] def test_only_synced_attributes_can_be_decided(): From 70c007b69e15da99306c0ada4b55dab41bc060fc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 00:54:59 +0000 Subject: [PATCH 05/48] Declare a synced session attribute by its type: Synced[T] A synced attribute was declared with synced(), which hid its type, and could be held on a group no attribute used. Annotating it now carries both: plan: Synced[DistributionPlan | None] is typed as the plan for readers and type checkers, and the session turns each Synced[...] annotation into the synced attribute, so assigning it is collective. A plain annotation stays a plain attribute. Synced is Annotated[T, mark]; the session's base installs the attribute for each marked annotation when a subclass is defined. Every synced attribute and agree() now span every rank. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 96 +++++++++++-------- .../torch/utils/test_distributed_session.py | 38 ++++---- 2 files changed, 76 insertions(+), 58 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 5e5b9e3a9f9..790a8c45f6f 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -17,6 +17,7 @@ import copy import functools +import inspect import io import os import sys @@ -25,7 +26,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager, suppress from datetime import timedelta -from typing import Any +from typing import Annotated, Any, TypeVar from warnings import warn import torch @@ -38,6 +39,7 @@ "DistributionPlan", "ModelOptSession", "ParallelState", + "Synced", "SyncedAttribute", "backend", "barrier", @@ -54,7 +56,6 @@ "modelopt_session", "rank", "size", - "synced", ] @@ -340,24 +341,35 @@ def __repr__(self) -> str: ) +_T = TypeVar("_T") + + +class _SyncedMark: + """What :data:`Synced` attaches to an annotation.""" + + def __repr__(self) -> str: + return "synced" + + +#: Annotate a session attribute with its type as ``Synced[]`` to make every rank hold the same +#: value of it: ``plan: Synced[DistributionPlan | None]``. See :class:`SyncedAttribute`. +Synced = Annotated[_T, _SyncedMark()] + + class SyncedAttribute: - """A :class:`ModelOptSession` attribute every rank of a group holds the same value of. - - Declare one with :func:`synced`, for a fact a later stage reads back from the session; a one-off - decision that nothing reads back is :meth:`ModelOptSession.agree`. Reading it is local. Assigning - it is collective: every rank of the group assigns it at the same point, and each then holds the - value the group's first rank assigned -- whatever the others passed -- so no rank drifts from the - rest. :meth:`ModelOptSession.decide` sets one from a function run on the first rank only. Assign a - new value rather than mutating the one held, which would change this rank's copy only. Outside a - distributed session it is a plain attribute. + """A session attribute annotated :data:`Synced`, which every rank holds the same value of. + + Declare one by its type, ``plan: Synced[DistributionPlan | None]``, for a fact a later stage reads + back from the session; a one-off decision that nothing reads back is + :meth:`ModelOptSession.agree`. Reading it is local. Assigning it is collective: every rank assigns + it at the same point, and each then holds the value rank 0 assigned -- whatever the others passed + -- so no rank drifts from the rest. :meth:`ModelOptSession.decide` sets one from a function run on + rank 0 only. Assign a new value rather than mutating the one held, which would change this rank's + copy only. Outside a distributed session it is a plain attribute. """ - def __init__(self, group: "Callable[[ModelOptSession], Any] | None" = None) -> None: - """Hold the attribute on ``group``.""" - self.group = group - self.name = "" - - def __set_name__(self, owner: type, name: str) -> None: + def __init__(self, name: str) -> None: + """The synced attribute ``name``.""" self.name = name def __get__(self, session, owner=None): @@ -369,14 +381,16 @@ def __set__(self, session, value) -> None: session._sync(self, lambda: value, decide=False) -def synced(group: "Callable[[ModelOptSession], Any] | None" = None) -> SyncedAttribute: - """Declare a :class:`ModelOptSession` attribute that stays the same on every rank of a group. +class _SyncedHolder: + """Base of a class with synced attributes: each ``Synced[...]`` annotation of a subclass becomes one.""" - ``group`` picks the ranks: ``None`` for all of them, or a callable taking the session and returning - the ``ProcessGroup`` that holds the attribute, resolved at each assignment. The group's first rank - decides. See :class:`SyncedAttribute`. - """ - return SyncedAttribute(group) + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + for name, annotation in inspect.get_annotations(cls, eval_str=True).items(): + if any( + isinstance(mark, _SyncedMark) for mark in getattr(annotation, "__metadata__", ()) + ): + setattr(cls, name, SyncedAttribute(name)) def _agreed(label: str, values: list[Any]) -> Any: @@ -389,7 +403,7 @@ def _agreed(label: str, values: list[Any]) -> Any: return values[0] -class ModelOptSession: +class ModelOptSession(_SyncedHolder): """The live distributed session of a PTQ run: this process's identity, plus the model's EP mesh. One class covers both cases: ``distributed`` is False for a single process (``WORLD_SIZE==1``) and @@ -402,11 +416,11 @@ class ModelOptSession: * ``ep_mesh`` -- the expert-parallel device mesh of the loaded model, set by ``load_and_prepare_ep`` (``None`` until then, and for a model loaded without expert parallelism). - Attributes declared with :func:`synced`, such as ``plan``, hold the same value on every rank. + Attributes annotated :data:`Synced`, such as ``plan``, hold the same value on every rank. """ #: The :class:`DistributionPlan` of the model this session runs, the same on every rank. - plan = synced() + plan: Synced[DistributionPlan | None] def __init__(self) -> None: """Create an empty (uninitialized) session; populate identity with :meth:`initialize`.""" @@ -440,12 +454,12 @@ def _reset(self) -> None: self.ep_mesh = None def decide(self, name: str, decide_fn: Callable[[], Any]) -> Any: - """Set the synced attribute ``name`` to ``decide_fn()``, run on its group's first rank only. + """Set the synced attribute ``name`` to ``decide_fn()``, run on rank 0 only. The other ranks take that value instead of computing their own, which could differ when it depends on what each rank reads for itself: the environment, its GPU, files on its node. If - ``decide_fn`` raises, every rank of the group raises, instead of waiting in a collective the - failing rank never joins. Collective, like assigning the attribute. Returns the value. + ``decide_fn`` raises, every rank raises, instead of waiting in a collective rank 0 never + joins. Collective, like assigning the attribute. Returns the value. """ attribute = getattr(type(self), name, None) if not isinstance(attribute, SyncedAttribute): @@ -456,13 +470,12 @@ def decide(self, name: str, decide_fn: Callable[[], Any]) -> Any: def _sync( self, attribute: SyncedAttribute, value_fn: Callable[[], Any], *, decide: bool ) -> None: - """Store ``value_fn()`` from the group's first rank under ``attribute`` on every rank.""" + """Store ``value_fn()`` from rank 0 under ``attribute`` on every rank.""" name = attribute.name if not (self.initialized and self.distributed): self._synced[name] = value_fn() return - group = None if attribute.group is None else attribute.group(self) - first = torch.distributed.get_rank(group) == 0 + first = torch.distributed.get_rank() == 0 value, error = None, None if first or not decide: try: @@ -475,7 +488,7 @@ def _sync( # two are out of step, and its value or its error. The error goes as text: the other ranks only # need to stop, and an exception is not guaranteed to pickle. message = [(name, value, None if error is None else f"{type(error).__name__}: {error}")] - torch.distributed.broadcast_object_list(message, group=group, group_src=0) + torch.distributed.broadcast_object_list(message, group_src=0) first_name, value, first_error = message[0] if error is not None: raise error @@ -484,7 +497,7 @@ def _sync( if first_name != name: raise RuntimeError( f"Ranks set synced session attributes out of step: this rank set {name!r} while rank " - f"0 of its group set {first_name!r}. Every rank must set them in the same order." + f"0 set {first_name!r}. Every rank must set them in the same order." ) # The first rank's own object stays the caller's; a copy keeps a later mutation of it local. self._synced[name] = copy.deepcopy(value) if first else value @@ -494,19 +507,18 @@ def agree( label: str, value_fn: Callable[[], Any], combine: "Callable[[list[Any]], Any] | None" = None, - group=None, ) -> Any: - """Return what every rank of ``group`` agrees on from ``value_fn()``, run on each; keep nothing. + """Return what every rank agrees on from ``value_fn()``, run on each; keep nothing. For a decision no later stage reads back from the session (a fact that one does is a - :func:`synced` attribute). ``combine`` takes the group's values in rank order and returns the + :data:`Synced` attribute). ``combine`` takes the ranks' values in rank order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's memory bounds, where one rank's choice could run another out of memory. Without it every rank must pass the same value, e.g. the library versions each runs: a rank that differs makes every rank raise, where adopting one rank's value would hide the difference until a collective hangs on it. If ``value_fn`` raises on any rank, every rank raises, instead of waiting in the gather that rank never joins. - Ranks agreeing on different labels at the same point fail. Collective over ``group``, every - rank by default; outside a distributed session ``value_fn`` runs alone. + Ranks agreeing on different labels at the same point fail. Collective; outside a distributed + session ``value_fn`` runs alone. """ combine = combine or functools.partial(_agreed, label) if not (self.initialized and self.distributed): @@ -518,8 +530,8 @@ def agree( # Still join the gather: the other ranks would otherwise wait in it for this one. error = e mine = (label, value, None if error is None else f"{type(error).__name__}: {error}") - gathered: list[Any] = [None] * torch.distributed.get_world_size(group) - torch.distributed.all_gather_object(gathered, mine, group=group) + gathered: list[Any] = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(gathered, mine) if error is not None: raise error for rank, (_, _, rank_error) in enumerate(gathered): diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_session.py index 6a30b2efd2b..03af7b0cc83 100644 --- a/tests/unit/torch/utils/test_distributed_session.py +++ b/tests/unit/torch/utils/test_distributed_session.py @@ -22,9 +22,10 @@ from modelopt.torch.utils.distributed import ( DistributionPlan, ModelOptSession, + Synced, + SyncedAttribute, get_active_mo_session, is_distributed_launch, - synced, ) # --------------------------------------------------------------------------------- DistributionPlan @@ -100,10 +101,10 @@ def test_reset_clears_the_identity_and_the_ep_mesh(): class _Session(ModelOptSession): - """A session with a second synced attribute, and one held on a group of its own.""" + """A session with a second synced attribute, and a plain one.""" - other = synced() - grouped = synced(group=lambda session: "ep-group") + other: Synced[int] + note: str def _session(rank=None): @@ -157,6 +158,23 @@ def all_gather_object(gathered, mine, group=None): return sent +def test_only_attributes_annotated_synced_are_synced(): + """The annotation carries both the value's type and the synced semantic.""" + assert isinstance(vars(_Session)["other"], SyncedAttribute) + assert isinstance(vars(ModelOptSession)["plan"], SyncedAttribute) + assert "note" not in vars(_Session) + + +def test_a_plain_attribute_assigns_without_a_collective(monkeypatch): + for collective in ("broadcast_object_list", "all_gather_object"): + monkeypatch.setattr( + torch.distributed, collective, lambda *a, **k: pytest.fail("a collective ran") + ) + session = _session(rank=1) + session.note = "local" + assert session.note == "local" + + def test_a_synced_attribute_is_plain_in_a_single_process(monkeypatch): """Nothing to sync with: no collective, and an uninitialized session behaves the same.""" for collective in ("broadcast_object_list", "all_gather_object"): @@ -227,12 +245,6 @@ def test_ranks_setting_different_attributes_fail_instead_of_mixing_them(monkeypa _session(rank=1).plan = DistributionPlan(world_size=2) -def test_a_synced_attribute_is_held_on_its_group(monkeypatch): - sent = _as_rank(monkeypatch, 0) - _session(rank=0).grouped = 5 - assert sent == [(("grouped", 5, None), "ep-group")] - - def test_agree_gives_every_rank_the_combined_value(monkeypatch): """Every rank's value counts, and every rank gets the same result.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None), ("smallest", 2, None)]) @@ -288,12 +300,6 @@ def test_ranks_agreeing_out_of_step_fail(monkeypatch): _session(rank=0).agree("smallest", lambda: 4, min) -def test_agree_gathers_on_its_group(monkeypatch): - sent = _gather_as_rank(monkeypatch, 0, others=[("smallest", 1, None)]) - assert _session(rank=0).agree("smallest", lambda: 5, min, group="ep-group") == 1 - assert sent == [(("smallest", 5, None), "ep-group")] - - def test_only_synced_attributes_can_be_decided(): with pytest.raises(AttributeError, match="not a synced attribute"): _session().decide("ep_mesh", lambda: None) From 4ddd7fdbc1882ac27dd268f9439ce90fefd28d59 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 02:33:51 +0000 Subject: [PATCH 06/48] Set a synced fact through its handle: session.plan.decide(fn, rank) A synced attribute read as its value and set by assignment, a collective hidden behind an =, with decide on the session naming it by string. Annotating a fact Synced[T] now makes it a typed handle: session.plan.value reads it locally, session.plan.decide(fn, rank=0) sets it to fn() run on that rank alone, and session.plan.agree(fn, combine) to what every rank agrees on; assigning it raises. ModelOptSession.decide(label, fn, rank) and agree(label, fn, combine) return the value without keeping it, for a one-off decision nothing reads back. decide can now run fn on any rank, not only rank 0. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 186 ++++++++++-------- .../torch/utils/test_distributed_session.py | 133 +++++++------ 2 files changed, 180 insertions(+), 139 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 790a8c45f6f..60b20f7947a 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -26,7 +26,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager, suppress from datetime import timedelta -from typing import Annotated, Any, TypeVar +from typing import Any, Generic, TypeVar from warnings import warn import torch @@ -40,7 +40,6 @@ "ModelOptSession", "ParallelState", "Synced", - "SyncedAttribute", "backend", "barrier", "fsdp2_wrap", @@ -342,55 +341,86 @@ def __repr__(self) -> str: _T = TypeVar("_T") +_V = TypeVar("_V") -class _SyncedMark: - """What :data:`Synced` attaches to an annotation.""" +class Synced(Generic[_T]): + """A fact of the session every rank holds the same value of, declared by its type. - def __repr__(self) -> str: - return "synced" + ``plan: Synced[DistributionPlan]`` on the session makes ``session.plan`` this handle. Read the + value with ``session.plan.value``: local, ``None`` until it is set. Set it with :meth:`decide`, one + rank computing it, or :meth:`agree`, every rank does; both are collective, so every rank calls them + at the same point. Assigning ``session.plan`` raises, so nothing communicates behind an ``=``. + Declare one only for a fact a later stage reads back: a one-off decision nothing reads back is + :meth:`ModelOptSession.decide` or :meth:`ModelOptSession.agree`, which keep nothing. + """ + def __init__(self, holder: "_SyncedHolder", name: str) -> None: + """The fact ``name`` of ``holder``.""" + self._holder = holder + self._name = name -#: Annotate a session attribute with its type as ``Synced[]`` to make every rank hold the same -#: value of it: ``plan: Synced[DistributionPlan | None]``. See :class:`SyncedAttribute`. -Synced = Annotated[_T, _SyncedMark()] + @property + def value(self) -> _T | None: + """The value every rank holds, read locally; ``None`` until one is set.""" + return self._holder._synced.get(self._name) + def decide(self, fn: Callable[[], _T], rank: int = 0) -> _T: + """Set the fact to ``fn()`` run on ``rank`` only; every rank holds and returns that value. -class SyncedAttribute: - """A session attribute annotated :data:`Synced`, which every rank holds the same value of. + See :meth:`ModelOptSession.decide`. + """ + label = self._holder._label(self._name) + return self._hold(self._holder._sync_session().decide(label, fn, rank)) + + def agree(self, fn: Callable[[], _T], combine: "Callable[[list[_T]], _T] | None" = None) -> _T: + """Set the fact to what every rank agrees on from ``fn()``, run on each. + + See :meth:`ModelOptSession.agree`. + """ + label = self._holder._label(self._name) + return self._hold(self._holder._sync_session().agree(label, fn, combine)) + + def _hold(self, value: _T) -> _T: + # The holder keeps its own copy: a later mutation of the caller's object stays the caller's. + self._holder._synced[self._name] = copy.deepcopy(value) + return self._holder._synced[self._name] - Declare one by its type, ``plan: Synced[DistributionPlan | None]``, for a fact a later stage reads - back from the session; a one-off decision that nothing reads back is - :meth:`ModelOptSession.agree`. Reading it is local. Assigning it is collective: every rank assigns - it at the same point, and each then holds the value rank 0 assigned -- whatever the others passed - -- so no rank drifts from the rest. :meth:`ModelOptSession.decide` sets one from a function run on - rank 0 only. Assign a new value rather than mutating the one held, which would change this rank's - copy only. Outside a distributed session it is a plain attribute. - """ + +class _SyncedAttribute: + """What a ``Synced[...]`` annotation installs: reading it gives the fact's :class:`Synced` handle.""" def __init__(self, name: str) -> None: - """The synced attribute ``name``.""" self.name = name - def __get__(self, session, owner=None): - if session is None: - return self - return session._synced.get(self.name) + def __get__(self, holder, owner=None): + return self if holder is None else Synced(holder, self.name) - def __set__(self, session, value) -> None: - session._sync(self, lambda: value, decide=False) + def __set__(self, holder, value) -> None: + raise AttributeError( + f"{type(holder).__name__}.{self.name} is synced: set it with .decide(fn) or " + ".agree(fn), which every rank calls at the same point." + ) class _SyncedHolder: - """Base of a class with synced attributes: each ``Synced[...]`` annotation of a subclass becomes one.""" + """Base of a class with synced facts: each ``Synced[...]`` annotation of a subclass becomes one.""" + + _synced: dict[str, Any] def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) for name, annotation in inspect.get_annotations(cls, eval_str=True).items(): - if any( - isinstance(mark, _SyncedMark) for mark in getattr(annotation, "__metadata__", ()) - ): - setattr(cls, name, SyncedAttribute(name)) + if getattr(annotation, "__origin__", None) is Synced: + setattr(cls, name, _SyncedAttribute(name)) + + def _sync_session(self) -> "ModelOptSession": + """The session whose ranks hold this object's synced facts.""" + raise NotImplementedError + + def _label(self, name: str) -> str: + """How the fact ``name`` is named to the other ranks, which must name the same one.""" + return name def _agreed(label: str, values: list[Any]) -> Any: @@ -416,11 +446,11 @@ class ModelOptSession(_SyncedHolder): * ``ep_mesh`` -- the expert-parallel device mesh of the loaded model, set by ``load_and_prepare_ep`` (``None`` until then, and for a model loaded without expert parallelism). - Attributes annotated :data:`Synced`, such as ``plan``, hold the same value on every rank. + Facts annotated :class:`Synced`, such as ``plan``, hold the same value on every rank. """ #: The :class:`DistributionPlan` of the model this session runs, the same on every rank. - plan: Synced[DistributionPlan | None] + plan: Synced[DistributionPlan] def __init__(self) -> None: """Create an empty (uninitialized) session; populate identity with :meth:`initialize`.""" @@ -453,72 +483,66 @@ def _reset(self) -> None: self._synced = {} self.ep_mesh = None - def decide(self, name: str, decide_fn: Callable[[], Any]) -> Any: - """Set the synced attribute ``name`` to ``decide_fn()``, run on rank 0 only. + def _sync_session(self) -> "ModelOptSession": + return self + + def decide(self, label: str, fn: Callable[[], _V], rank: int = 0) -> _V: + """Return ``fn()`` run on ``rank`` only, on every rank; keep nothing. - The other ranks take that value instead of computing their own, which could differ when it - depends on what each rank reads for itself: the environment, its GPU, files on its node. If - ``decide_fn`` raises, every rank raises, instead of waiting in a collective rank 0 never - joins. Collective, like assigning the attribute. Returns the value. + For a value that depends on what each rank reads for itself -- the environment, its GPU, files + on its node -- so that ranks deciding for themselves could disagree. If ``fn`` raises, every + rank raises, instead of waiting in a broadcast ``rank`` never sends. Ranks deciding different + labels at the same point fail. Collective: every rank calls it at the same point, with the same + ``rank``. Outside a distributed session ``fn`` runs alone. A fact a later stage reads back is a + :class:`Synced` attribute, set with its own ``decide``. """ - attribute = getattr(type(self), name, None) - if not isinstance(attribute, SyncedAttribute): - raise AttributeError(f"ModelOptSession.{name} is not a synced attribute") - self._sync(attribute, decide_fn, decide=True) - return self._synced[name] - - def _sync( - self, attribute: SyncedAttribute, value_fn: Callable[[], Any], *, decide: bool - ) -> None: - """Store ``value_fn()`` from rank 0 under ``attribute`` on every rank.""" - name = attribute.name if not (self.initialized and self.distributed): - self._synced[name] = value_fn() - return - first = torch.distributed.get_rank() == 0 - value, error = None, None - if first or not decide: + return fn() + if not 0 <= rank < self.world_size: + raise ValueError(f"rank={rank} is not a rank of this {self.world_size}-rank run.") + deciding = torch.distributed.get_rank() == rank + value: Any = None + error = None + if deciding: try: - value = value_fn() + value = fn() except Exception as e: - if not first: - raise + # Still send: the other ranks would otherwise wait in the broadcast for this one. error = e - # One message from the first rank: which attribute it set, so a rank setting another learns the + # One message from the deciding rank: what it decided, so a rank deciding another learns the # two are out of step, and its value or its error. The error goes as text: the other ranks only # need to stop, and an exception is not guaranteed to pickle. - message = [(name, value, None if error is None else f"{type(error).__name__}: {error}")] - torch.distributed.broadcast_object_list(message, group_src=0) - first_name, value, first_error = message[0] + message = [(label, value, None if error is None else f"{type(error).__name__}: {error}")] + torch.distributed.broadcast_object_list(message, group_src=rank) + decided_label, value, decided_error = message[0] if error is not None: raise error - if first_error is not None: - raise RuntimeError(f"rank 0 failed to decide ModelOptSession.{name}: {first_error}") - if first_name != name: + if decided_error is not None: + raise RuntimeError(f"rank {rank} failed to decide {label}: {decided_error}") + if decided_label != label: raise RuntimeError( - f"Ranks set synced session attributes out of step: this rank set {name!r} while rank " - f"0 set {first_name!r}. Every rank must set them in the same order." + f"Ranks decided out of step: this rank decided {label!r} while rank {rank} decided " + f"{decided_label!r}. Every rank must decide them in the same order." ) - # The first rank's own object stays the caller's; a copy keeps a later mutation of it local. - self._synced[name] = copy.deepcopy(value) if first else value + return value def agree( self, label: str, - value_fn: Callable[[], Any], - combine: "Callable[[list[Any]], Any] | None" = None, - ) -> Any: + value_fn: Callable[[], _V], + combine: "Callable[[list[_V]], _V] | None" = None, + ) -> _V: """Return what every rank agrees on from ``value_fn()``, run on each; keep nothing. - For a decision no later stage reads back from the session (a fact that one does is a - :data:`Synced` attribute). ``combine`` takes the ranks' values in rank order and returns the - one every rank gets, e.g. ``min`` for a batch size each rank's memory bounds, where one rank's - choice could run another out of memory. Without it every rank must pass the same value, e.g. - the library versions each runs: a rank that differs makes every rank raise, where adopting - one rank's value would hide the difference until a collective hangs on it. If ``value_fn`` - raises on any rank, every rank raises, instead of waiting in the gather that rank never joins. - Ranks agreeing on different labels at the same point fail. Collective; outside a distributed - session ``value_fn`` runs alone. + For a decision no later stage reads back from the session; a fact that one does is a + :class:`Synced` attribute, set with its own ``agree``. ``combine`` takes the ranks' values in rank + order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's memory + bounds, where one rank's choice could run another out of memory. Without it every rank must pass + the same value, e.g. the library versions each runs: a rank that differs makes every rank raise, + where adopting one rank's value would hide the difference until a collective hangs on it. If + ``value_fn`` raises on any rank, every rank raises, instead of waiting in the gather that rank + never joins. Ranks agreeing on different labels at the same point fail. Collective; outside a + distributed session ``value_fn`` runs alone. """ combine = combine or functools.partial(_agreed, label) if not (self.initialized and self.distributed): diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_session.py index 03af7b0cc83..8be3dbba3e9 100644 --- a/tests/unit/torch/utils/test_distributed_session.py +++ b/tests/unit/torch/utils/test_distributed_session.py @@ -14,7 +14,7 @@ # limitations under the License. """Unit tests for the distributed PTQ session: the distribution plan, the session handle, its synced -attributes, launch detection.""" +facts and one-off decisions, launch detection.""" import pytest import torch.distributed @@ -23,7 +23,6 @@ DistributionPlan, ModelOptSession, Synced, - SyncedAttribute, get_active_mo_session, is_distributed_launch, ) @@ -97,11 +96,11 @@ def test_reset_clears_the_identity_and_the_ep_mesh(): assert session.ep_mesh is None -# ------------------------------------------------------------------------------- synced attributes +# ----------------------------------------------------------------------------------- synced facts class _Session(ModelOptSession): - """A session with a second synced attribute, and a plain one.""" + """A session with a second synced fact, and a plain attribute.""" other: Synced[int] note: str @@ -121,18 +120,17 @@ def _session(rank=None): return session -def _as_rank(monkeypatch, rank, rank0_message=None): - """Answer the session's collectives as ``rank``; returns what each broadcast was handed. +def _as_rank(monkeypatch, rank, decided=None): + """Answer the session's broadcasts as ``rank``; returns ``(message, sending rank)`` for each. - A non-zero rank receives ``rank0_message``, as the group's first rank would have sent it. + A rank that is not the sending one receives ``decided``, as the sending rank would have sent it. """ sent = [] def broadcast_object_list(message, group=None, group_src=None): - assert group_src == 0 - sent.append((message[0], group)) - if rank != 0: - message[0] = rank0_message + sent.append((message[0], group_src)) + if rank != group_src: + message[0] = decided monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank) monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) @@ -147,24 +145,31 @@ def _gather_as_rank(monkeypatch, rank, others): sent = [] def all_gather_object(gathered, mine, group=None): - sent.append((mine, group)) + sent.append(mine) gathered[:] = [*others[:rank], mine, *others[rank:]] monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: len(others) + 1) monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) monkeypatch.setattr( - torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("rank 0 decided") + torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("a rank decided") ) return sent -def test_only_attributes_annotated_synced_are_synced(): +def test_only_facts_annotated_synced_are_synced(): """The annotation carries both the value's type and the synced semantic.""" - assert isinstance(vars(_Session)["other"], SyncedAttribute) - assert isinstance(vars(ModelOptSession)["plan"], SyncedAttribute) + session = _session() + assert isinstance(session.other, Synced) + assert isinstance(session.plan, Synced) assert "note" not in vars(_Session) +def test_a_synced_fact_cannot_be_assigned(): + """Setting one is collective, so it goes through decide or agree, never a bare assignment.""" + with pytest.raises(AttributeError, match=r"plan is synced: set it with \.decide"): + _session().plan = DistributionPlan(world_size=1) + + def test_a_plain_attribute_assigns_without_a_collective(monkeypatch): for collective in ("broadcast_object_list", "all_gather_object"): monkeypatch.setattr( @@ -175,7 +180,7 @@ def test_a_plain_attribute_assigns_without_a_collective(monkeypatch): assert session.note == "local" -def test_a_synced_attribute_is_plain_in_a_single_process(monkeypatch): +def test_a_single_process_decides_and_agrees_alone(monkeypatch): """Nothing to sync with: no collective, and an uninitialized session behaves the same.""" for collective in ("broadcast_object_list", "all_gather_object"): monkeypatch.setattr( @@ -183,66 +188,75 @@ def test_a_synced_attribute_is_plain_in_a_single_process(monkeypatch): ) for session in (_session(), _Session()): plan = DistributionPlan(world_size=1) - session.plan = plan - assert session.plan is plan - assert session.decide("other", lambda: 3) == 3 + assert session.plan.value is None + assert session.plan.decide(lambda: plan).world_size == 1 + assert session.plan.value.world_size == 1 + assert session.other.agree(lambda: 4, min) == 4 + assert session.decide("x", lambda: 3) == 3 assert session.agree("smallest", lambda: 4, min) == 4 assert session.agree("versions", lambda: "2.6") == "2.6" -def test_rank0_sends_what_it_assigns_and_keeps_its_own_copy(monkeypatch): +def test_the_deciding_rank_sends_its_value_and_keeps_its_own_copy(monkeypatch): sent = _as_rank(monkeypatch, 0) session, plan = _session(rank=0), DistributionPlan(world_size=2, ep_size=2) - session.plan = plan + decided = session.plan.decide(lambda: plan) - assert sent == [(("plan", plan, None), None)] - assert session.plan.ep_size == 2 - assert session.plan is not plan # mutating the caller's object must not reach the session + assert sent == [(("plan", plan, None), 0)] + assert decided is session.plan.value + assert decided.ep_size == 2 + assert decided is not plan, "mutating the caller's object must not reach the session" -def test_other_ranks_hold_what_rank0_assigned(monkeypatch): - """Whatever a non-zero rank passes, it ends up with rank 0's value, as every rank does.""" +def test_other_ranks_hold_the_deciding_ranks_value_without_running_fn(monkeypatch): rank0_plan = DistributionPlan(world_size=2, ep_size=2) - _as_rank(monkeypatch, 1, rank0_message=("plan", rank0_plan, None)) + _as_rank(monkeypatch, 1, decided=("plan", rank0_plan, None)) session = _session(rank=1) - session.plan = DistributionPlan(world_size=2) + decided = session.plan.decide(lambda: pytest.fail("a rank that does not decide ran fn")) - assert session.plan is rank0_plan + assert decided.ep_size == session.plan.value.ep_size == 2 -def test_decide_runs_only_on_rank0(monkeypatch): - _as_rank(monkeypatch, 1, rank0_message=("other", 7, None)) - session = _session(rank=1) +def test_any_rank_can_decide(monkeypatch): + """With rank=1, rank 1 runs fn and sends it; rank 0 takes rank 1's value.""" + sent = _as_rank(monkeypatch, 1) + assert _session(rank=1).decide("x", lambda: "rank 1's", rank=1) == "rank 1's" + assert sent == [(("x", "rank 1's", None), 1)] + + _as_rank(monkeypatch, 0, decided=("x", "rank 1's", None)) + assert _session(rank=0).decide("x", lambda: pytest.fail("rank 0 decided"), rank=1) == "rank 1's" + - assert session.decide("other", lambda: pytest.fail("a non-zero rank decided")) == 7 - assert session.other == 7 +def test_a_rank_outside_the_run_cannot_decide(): + with pytest.raises(ValueError, match="rank=2 is not a rank of this 2-rank run"): + _session(rank=0).decide("x", lambda: 0, rank=2) -def test_rank0_failing_to_decide_still_reaches_the_broadcast(monkeypatch): - """Rank 0 must still reach the broadcast on failure, or every other rank hangs in it.""" +def test_the_deciding_rank_failing_still_reaches_the_broadcast(monkeypatch): + """The deciding rank must still send on failure, or every other rank hangs in the broadcast.""" sent = _as_rank(monkeypatch, 0) def fail(): raise ValueError("bad degrees") with pytest.raises(ValueError, match="bad degrees"): - _session(rank=0).decide("other", fail) - assert sent == [(("other", None, "ValueError: bad degrees"), None)] + _session(rank=0).decide("x", fail) + assert sent == [(("x", None, "ValueError: bad degrees"), 0)] -def test_other_ranks_raise_when_rank0_failed(monkeypatch): - _as_rank(monkeypatch, 1, rank0_message=("other", None, "ValueError: bad degrees")) - with pytest.raises(RuntimeError, match=r"rank 0 failed to decide ModelOptSession.other.*bad"): - _session(rank=1).decide("other", lambda: 0) +def test_other_ranks_raise_when_the_deciding_rank_failed(monkeypatch): + _as_rank(monkeypatch, 1, decided=("x", None, "ValueError: bad degrees")) + with pytest.raises(RuntimeError, match=r"rank 0 failed to decide x.*bad degrees"): + _session(rank=1).decide("x", lambda: 0) -def test_ranks_setting_different_attributes_fail_instead_of_mixing_them(monkeypatch): - """Rank 0 setting one attribute while rank 1 sets another would cross the two values.""" - _as_rank(monkeypatch, 1, rank0_message=("other", 7, None)) +def test_ranks_deciding_out_of_step_fail(monkeypatch): + """Rank 0 deciding one fact while rank 1 decides another would cross the two values.""" + _as_rank(monkeypatch, 1, decided=("other", 7, None)) with pytest.raises(RuntimeError, match="out of step"): - _session(rank=1).plan = DistributionPlan(world_size=2) + _session(rank=1).plan.decide(lambda: DistributionPlan(world_size=2)) def test_agree_gives_every_rank_the_combined_value(monkeypatch): @@ -250,10 +264,18 @@ def test_agree_gives_every_rank_the_combined_value(monkeypatch): sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None), ("smallest", 2, None)]) assert _session(rank=1).agree("smallest", lambda: 4, min) == 2 - assert sent == [(("smallest", 4, None), None)] + assert sent == [("smallest", 4, None)] -def test_agree_keeps_nothing_on_the_session(monkeypatch): +def test_a_facts_agree_holds_the_combined_value(monkeypatch): + _gather_as_rank(monkeypatch, 1, others=[("other", 8, None)]) + session = _session(rank=1) + + assert session.other.agree(lambda: 4, min) == 4 + assert session.other.value == 4 + + +def test_session_decide_and_agree_keep_nothing(monkeypatch): _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) session = _session(rank=1) @@ -285,7 +307,7 @@ def fail(): with pytest.raises(ValueError, match="no memory reading"): _session(rank=1).agree("smallest", fail, min) - assert sent == [(("smallest", None, "ValueError: no memory reading"), None)] + assert sent == [("smallest", None, "ValueError: no memory reading")] def test_the_other_ranks_raise_when_one_failed_to_agree(monkeypatch): @@ -300,16 +322,11 @@ def test_ranks_agreeing_out_of_step_fail(monkeypatch): _session(rank=0).agree("smallest", lambda: 4, min) -def test_only_synced_attributes_can_be_decided(): - with pytest.raises(AttributeError, match="not a synced attribute"): - _session().decide("ep_mesh", lambda: None) - - -def test_reset_clears_synced_attributes(): +def test_reset_clears_synced_facts(): session = _session() - session.plan = DistributionPlan(world_size=1) + session.plan.decide(lambda: DistributionPlan(world_size=1)) session._reset() - assert session.plan is None + assert session.plan.value is None # ------------------------------------------------------------------------------------ launch detect From b9080c4f5dcba1fcbb2bdd7b66ac81e3c1f478e6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 03:04:29 +0000 Subject: [PATCH 07/48] Name the session's collectives after what they do: broadcast_from and all_reduce decide and agree did not say which collective runs. broadcast_from(fn, rank) runs fn on one rank and broadcasts its result; all_reduce(fn, combine) runs fn on every rank and gives each the combined result. all_reduce now takes its combine: nothing needs the old default, which required every rank to pass the same value. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 107 +++++++--------- .../torch/utils/test_distributed_session.py | 120 ++++++++---------- 2 files changed, 98 insertions(+), 129 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 60b20f7947a..12057d5da08 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -348,11 +348,12 @@ class Synced(Generic[_T]): """A fact of the session every rank holds the same value of, declared by its type. ``plan: Synced[DistributionPlan]`` on the session makes ``session.plan`` this handle. Read the - value with ``session.plan.value``: local, ``None`` until it is set. Set it with :meth:`decide`, one - rank computing it, or :meth:`agree`, every rank does; both are collective, so every rank calls them - at the same point. Assigning ``session.plan`` raises, so nothing communicates behind an ``=``. - Declare one only for a fact a later stage reads back: a one-off decision nothing reads back is - :meth:`ModelOptSession.decide` or :meth:`ModelOptSession.agree`, which keep nothing. + value with ``session.plan.value``: local, ``None`` until it is set. Set it with + :meth:`broadcast_from`, one rank computing it, or :meth:`all_reduce`, every rank computing its own + and the results combined; both are collective, so every rank calls them at the same point. + Assigning ``session.plan`` raises, so nothing communicates behind an ``=``. Declare one only for a + fact a later stage reads back: a one-off value nothing reads back is + :meth:`ModelOptSession.broadcast_from` or :meth:`ModelOptSession.all_reduce`, which keep nothing. """ def __init__(self, holder: "_SyncedHolder", name: str) -> None: @@ -365,21 +366,21 @@ def value(self) -> _T | None: """The value every rank holds, read locally; ``None`` until one is set.""" return self._holder._synced.get(self._name) - def decide(self, fn: Callable[[], _T], rank: int = 0) -> _T: + def broadcast_from(self, fn: Callable[[], _T], rank: int = 0) -> _T: """Set the fact to ``fn()`` run on ``rank`` only; every rank holds and returns that value. - See :meth:`ModelOptSession.decide`. + See :meth:`ModelOptSession.broadcast_from`. """ label = self._holder._label(self._name) - return self._hold(self._holder._sync_session().decide(label, fn, rank)) + return self._hold(self._holder._sync_session().broadcast_from(label, fn, rank)) - def agree(self, fn: Callable[[], _T], combine: "Callable[[list[_T]], _T] | None" = None) -> _T: - """Set the fact to what every rank agrees on from ``fn()``, run on each. + def all_reduce(self, fn: Callable[[], _T], combine: Callable[[list[_T]], _T]) -> _T: + """Set the fact to ``combine`` of every rank's ``fn()``; every rank holds and returns it. - See :meth:`ModelOptSession.agree`. + See :meth:`ModelOptSession.all_reduce`. """ label = self._holder._label(self._name) - return self._hold(self._holder._sync_session().agree(label, fn, combine)) + return self._hold(self._holder._sync_session().all_reduce(label, fn, combine)) def _hold(self, value: _T) -> _T: # The holder keeps its own copy: a later mutation of the caller's object stays the caller's. @@ -398,8 +399,8 @@ def __get__(self, holder, owner=None): def __set__(self, holder, value) -> None: raise AttributeError( - f"{type(holder).__name__}.{self.name} is synced: set it with .decide(fn) or " - ".agree(fn), which every rank calls at the same point." + f"{type(holder).__name__}.{self.name} is synced: set it with .broadcast_from(fn) or " + ".all_reduce(fn, combine), which every rank calls at the same point." ) @@ -423,16 +424,6 @@ def _label(self, name: str) -> str: return name -def _agreed(label: str, values: list[Any]) -> Any: - """The value every rank passed to :meth:`ModelOptSession.agree`; raise if any rank passed another.""" - for rank, value in enumerate(values): - if value != values[0]: - raise RuntimeError( - f"Ranks disagree on {label}: rank 0 has {values[0]!r}, rank {rank} has {value!r}." - ) - return values[0] - - class ModelOptSession(_SyncedHolder): """The live distributed session of a PTQ run: this process's identity, plus the model's EP mesh. @@ -486,70 +477,62 @@ def _reset(self) -> None: def _sync_session(self) -> "ModelOptSession": return self - def decide(self, label: str, fn: Callable[[], _V], rank: int = 0) -> _V: + def broadcast_from(self, label: str, fn: Callable[[], _V], rank: int = 0) -> _V: """Return ``fn()`` run on ``rank`` only, on every rank; keep nothing. For a value that depends on what each rank reads for itself -- the environment, its GPU, files - on its node -- so that ranks deciding for themselves could disagree. If ``fn`` raises, every - rank raises, instead of waiting in a broadcast ``rank`` never sends. Ranks deciding different - labels at the same point fail. Collective: every rank calls it at the same point, with the same - ``rank``. Outside a distributed session ``fn`` runs alone. A fact a later stage reads back is a - :class:`Synced` attribute, set with its own ``decide``. + on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, every + rank raises, instead of waiting in a broadcast ``rank`` never sends. Ranks broadcasting + different labels at the same point fail. Collective: every rank calls it at the same point, + with the same ``rank``. Outside a distributed session ``fn`` runs alone. A fact a later stage + reads back is a :class:`Synced` attribute, set with its own ``broadcast_from``. """ if not (self.initialized and self.distributed): return fn() if not 0 <= rank < self.world_size: raise ValueError(f"rank={rank} is not a rank of this {self.world_size}-rank run.") - deciding = torch.distributed.get_rank() == rank + sending = torch.distributed.get_rank() == rank value: Any = None error = None - if deciding: + if sending: try: value = fn() except Exception as e: # Still send: the other ranks would otherwise wait in the broadcast for this one. error = e - # One message from the deciding rank: what it decided, so a rank deciding another learns the - # two are out of step, and its value or its error. The error goes as text: the other ranks only + # One message from the sending rank: its label, so a rank broadcasting another learns the two + # are out of step, and its value or its error. The error goes as text: the other ranks only # need to stop, and an exception is not guaranteed to pickle. message = [(label, value, None if error is None else f"{type(error).__name__}: {error}")] torch.distributed.broadcast_object_list(message, group_src=rank) - decided_label, value, decided_error = message[0] + sent_label, value, sent_error = message[0] if error is not None: raise error - if decided_error is not None: - raise RuntimeError(f"rank {rank} failed to decide {label}: {decided_error}") - if decided_label != label: + if sent_error is not None: + raise RuntimeError(f"rank {rank} failed to compute {label}: {sent_error}") + if sent_label != label: raise RuntimeError( - f"Ranks decided out of step: this rank decided {label!r} while rank {rank} decided " - f"{decided_label!r}. Every rank must decide them in the same order." + f"Ranks broadcast out of step: this rank broadcast {label!r} while rank {rank} sent " + f"{sent_label!r}. Every rank must broadcast them in the same order." ) return value - def agree( - self, - label: str, - value_fn: Callable[[], _V], - combine: "Callable[[list[_V]], _V] | None" = None, - ) -> _V: - """Return what every rank agrees on from ``value_fn()``, run on each; keep nothing. - - For a decision no later stage reads back from the session; a fact that one does is a - :class:`Synced` attribute, set with its own ``agree``. ``combine`` takes the ranks' values in rank - order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's memory - bounds, where one rank's choice could run another out of memory. Without it every rank must pass - the same value, e.g. the library versions each runs: a rank that differs makes every rank raise, - where adopting one rank's value would hide the difference until a collective hangs on it. If - ``value_fn`` raises on any rank, every rank raises, instead of waiting in the gather that rank - never joins. Ranks agreeing on different labels at the same point fail. Collective; outside a - distributed session ``value_fn`` runs alone. + def all_reduce(self, label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _V]) -> _V: + """Return ``combine`` of every rank's ``fn()``, on every rank; keep nothing. + + For a value no later stage reads back from the session; a fact that one does is a + :class:`Synced` attribute, set with its own ``all_reduce``. ``combine`` takes the ranks' values + in rank order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's + memory bounds, where one rank's choice could run another out of memory. If ``fn`` raises on any + rank, every rank raises, instead of waiting in the gather that rank never joins. Ranks reducing + different labels at the same point fail. Collective; outside a distributed session ``fn`` runs + alone. """ - combine = combine or functools.partial(_agreed, label) if not (self.initialized and self.distributed): - return combine([value_fn()]) + return combine([fn()]) value, error = None, None try: - value = value_fn() + value = fn() except Exception as e: # Still join the gather: the other ranks would otherwise wait in it for this one. error = e @@ -560,11 +543,11 @@ def agree( raise error for rank, (_, _, rank_error) in enumerate(gathered): if rank_error is not None: - raise RuntimeError(f"rank {rank} failed to agree on {label}: {rank_error}") + raise RuntimeError(f"rank {rank} failed to compute {label}: {rank_error}") labels = sorted({rank_label for rank_label, _, _ in gathered}) if labels != [label]: raise RuntimeError( - f"Ranks agreed out of step: {labels} at the same point. Every rank must agree on " + f"Ranks reduced out of step: {labels} at the same point. Every rank must reduce " "them in the same order." ) # The gathered values are unpickled copies, so the result shares nothing with the callers'. diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_session.py index 8be3dbba3e9..de3dcee327b 100644 --- a/tests/unit/torch/utils/test_distributed_session.py +++ b/tests/unit/torch/utils/test_distributed_session.py @@ -14,7 +14,7 @@ # limitations under the License. """Unit tests for the distributed PTQ session: the distribution plan, the session handle, its synced -facts and one-off decisions, launch detection.""" +facts and one-off collectives, launch detection.""" import pytest import torch.distributed @@ -120,17 +120,17 @@ def _session(rank=None): return session -def _as_rank(monkeypatch, rank, decided=None): +def _as_rank(monkeypatch, rank, received=None): """Answer the session's broadcasts as ``rank``; returns ``(message, sending rank)`` for each. - A rank that is not the sending one receives ``decided``, as the sending rank would have sent it. + A rank that is not the sending one receives ``received``, as the sending rank would have sent it. """ sent = [] def broadcast_object_list(message, group=None, group_src=None): sent.append((message[0], group_src)) if rank != group_src: - message[0] = decided + message[0] = received monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank) monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) @@ -151,7 +151,7 @@ def all_gather_object(gathered, mine, group=None): monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: len(others) + 1) monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) monkeypatch.setattr( - torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("a rank decided") + torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("a rank broadcast") ) return sent @@ -165,8 +165,8 @@ def test_only_facts_annotated_synced_are_synced(): def test_a_synced_fact_cannot_be_assigned(): - """Setting one is collective, so it goes through decide or agree, never a bare assignment.""" - with pytest.raises(AttributeError, match=r"plan is synced: set it with \.decide"): + """Setting one is collective: broadcast_from or all_reduce, never a bare assignment.""" + with pytest.raises(AttributeError, match=r"plan is synced: set it with \.broadcast_from"): _session().plan = DistributionPlan(world_size=1) @@ -180,7 +180,7 @@ def test_a_plain_attribute_assigns_without_a_collective(monkeypatch): assert session.note == "local" -def test_a_single_process_decides_and_agrees_alone(monkeypatch): +def test_a_single_process_computes_alone(monkeypatch): """Nothing to sync with: no collective, and an uninitialized session behaves the same.""" for collective in ("broadcast_object_list", "all_gather_object"): monkeypatch.setattr( @@ -189,116 +189,102 @@ def test_a_single_process_decides_and_agrees_alone(monkeypatch): for session in (_session(), _Session()): plan = DistributionPlan(world_size=1) assert session.plan.value is None - assert session.plan.decide(lambda: plan).world_size == 1 + assert session.plan.broadcast_from(lambda: plan).world_size == 1 assert session.plan.value.world_size == 1 - assert session.other.agree(lambda: 4, min) == 4 - assert session.decide("x", lambda: 3) == 3 - assert session.agree("smallest", lambda: 4, min) == 4 - assert session.agree("versions", lambda: "2.6") == "2.6" + assert session.other.all_reduce(lambda: 4, min) == 4 + assert session.broadcast_from("x", lambda: 3) == 3 + assert session.all_reduce("smallest", lambda: 4, min) == 4 -def test_the_deciding_rank_sends_its_value_and_keeps_its_own_copy(monkeypatch): +def test_the_sending_rank_sends_its_value_and_keeps_its_own_copy(monkeypatch): sent = _as_rank(monkeypatch, 0) session, plan = _session(rank=0), DistributionPlan(world_size=2, ep_size=2) - decided = session.plan.decide(lambda: plan) + held = session.plan.broadcast_from(lambda: plan) assert sent == [(("plan", plan, None), 0)] - assert decided is session.plan.value - assert decided.ep_size == 2 - assert decided is not plan, "mutating the caller's object must not reach the session" + assert held is session.plan.value + assert held.ep_size == 2 + assert held is not plan, "mutating the caller's object must not reach the session" -def test_other_ranks_hold_the_deciding_ranks_value_without_running_fn(monkeypatch): +def test_other_ranks_hold_the_sending_ranks_value_without_running_fn(monkeypatch): rank0_plan = DistributionPlan(world_size=2, ep_size=2) - _as_rank(monkeypatch, 1, decided=("plan", rank0_plan, None)) + _as_rank(monkeypatch, 1, received=("plan", rank0_plan, None)) session = _session(rank=1) - decided = session.plan.decide(lambda: pytest.fail("a rank that does not decide ran fn")) + held = session.plan.broadcast_from(lambda: pytest.fail("a rank that does not send ran fn")) - assert decided.ep_size == session.plan.value.ep_size == 2 + assert held.ep_size == session.plan.value.ep_size == 2 -def test_any_rank_can_decide(monkeypatch): +def test_any_rank_can_broadcast(monkeypatch): """With rank=1, rank 1 runs fn and sends it; rank 0 takes rank 1's value.""" sent = _as_rank(monkeypatch, 1) - assert _session(rank=1).decide("x", lambda: "rank 1's", rank=1) == "rank 1's" + assert _session(rank=1).broadcast_from("x", lambda: "rank 1's", rank=1) == "rank 1's" assert sent == [(("x", "rank 1's", None), 1)] - _as_rank(monkeypatch, 0, decided=("x", "rank 1's", None)) - assert _session(rank=0).decide("x", lambda: pytest.fail("rank 0 decided"), rank=1) == "rank 1's" + _as_rank(monkeypatch, 0, received=("x", "rank 1's", None)) + held = _session(rank=0).broadcast_from("x", lambda: pytest.fail("rank 0 ran fn"), rank=1) + assert held == "rank 1's" -def test_a_rank_outside_the_run_cannot_decide(): +def test_a_rank_outside_the_run_cannot_broadcast(): with pytest.raises(ValueError, match="rank=2 is not a rank of this 2-rank run"): - _session(rank=0).decide("x", lambda: 0, rank=2) + _session(rank=0).broadcast_from("x", lambda: 0, rank=2) -def test_the_deciding_rank_failing_still_reaches_the_broadcast(monkeypatch): - """The deciding rank must still send on failure, or every other rank hangs in the broadcast.""" +def test_the_sending_rank_failing_still_reaches_the_broadcast(monkeypatch): + """The sending rank must still send on failure, or every other rank hangs in the broadcast.""" sent = _as_rank(monkeypatch, 0) def fail(): raise ValueError("bad degrees") with pytest.raises(ValueError, match="bad degrees"): - _session(rank=0).decide("x", fail) + _session(rank=0).broadcast_from("x", fail) assert sent == [(("x", None, "ValueError: bad degrees"), 0)] -def test_other_ranks_raise_when_the_deciding_rank_failed(monkeypatch): - _as_rank(monkeypatch, 1, decided=("x", None, "ValueError: bad degrees")) - with pytest.raises(RuntimeError, match=r"rank 0 failed to decide x.*bad degrees"): - _session(rank=1).decide("x", lambda: 0) +def test_other_ranks_raise_when_the_sending_rank_failed(monkeypatch): + _as_rank(monkeypatch, 1, received=("x", None, "ValueError: bad degrees")) + with pytest.raises(RuntimeError, match=r"rank 0 failed to compute x.*bad degrees"): + _session(rank=1).broadcast_from("x", lambda: 0) -def test_ranks_deciding_out_of_step_fail(monkeypatch): - """Rank 0 deciding one fact while rank 1 decides another would cross the two values.""" - _as_rank(monkeypatch, 1, decided=("other", 7, None)) +def test_ranks_broadcasting_out_of_step_fail(monkeypatch): + """Rank 0 sending one fact while rank 1 expects another would cross the two values.""" + _as_rank(monkeypatch, 1, received=("other", 7, None)) with pytest.raises(RuntimeError, match="out of step"): - _session(rank=1).plan.decide(lambda: DistributionPlan(world_size=2)) + _session(rank=1).plan.broadcast_from(lambda: DistributionPlan(world_size=2)) -def test_agree_gives_every_rank_the_combined_value(monkeypatch): +def test_all_reduce_gives_every_rank_the_combined_value(monkeypatch): """Every rank's value counts, and every rank gets the same result.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None), ("smallest", 2, None)]) - assert _session(rank=1).agree("smallest", lambda: 4, min) == 2 + assert _session(rank=1).all_reduce("smallest", lambda: 4, min) == 2 assert sent == [("smallest", 4, None)] -def test_a_facts_agree_holds_the_combined_value(monkeypatch): +def test_a_facts_all_reduce_holds_the_combined_value(monkeypatch): _gather_as_rank(monkeypatch, 1, others=[("other", 8, None)]) session = _session(rank=1) - assert session.other.agree(lambda: 4, min) == 4 + assert session.other.all_reduce(lambda: 4, min) == 4 assert session.other.value == 4 -def test_session_decide_and_agree_keep_nothing(monkeypatch): +def test_the_sessions_own_all_reduce_keeps_nothing(monkeypatch): _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) session = _session(rank=1) - session.agree("smallest", lambda: 4, min) + session.all_reduce("smallest", lambda: 4, min) assert session._synced == {} -def test_without_combine_every_rank_must_pass_the_same_value(monkeypatch): - _gather_as_rank(monkeypatch, 1, others=[("versions", "2.6", None)]) - assert _session(rank=1).agree("versions", lambda: "2.6") == "2.6" - - -def test_ranks_that_disagree_all_raise(monkeypatch): - """Adopting rank 0's value instead would hide the difference until a collective hangs on it.""" - _gather_as_rank(monkeypatch, 0, others=[("versions", "2.5", None)]) - with pytest.raises( - RuntimeError, match=r"disagree on versions: rank 0 has '2.6', rank 1 has '2.5'" - ): - _session(rank=0).agree("versions", lambda: "2.6") - - -def test_a_rank_failing_to_agree_still_joins_the_gather(monkeypatch): +def test_a_rank_failing_to_compute_still_joins_the_gather(monkeypatch): """A failing rank must still reach the gather, or every other rank waits in it.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) @@ -306,25 +292,25 @@ def fail(): raise ValueError("no memory reading") with pytest.raises(ValueError, match="no memory reading"): - _session(rank=1).agree("smallest", fail, min) + _session(rank=1).all_reduce("smallest", fail, min) assert sent == [("smallest", None, "ValueError: no memory reading")] -def test_the_other_ranks_raise_when_one_failed_to_agree(monkeypatch): +def test_the_other_ranks_raise_when_one_failed_to_compute(monkeypatch): _gather_as_rank(monkeypatch, 0, others=[("smallest", None, "ValueError: no memory reading")]) - with pytest.raises(RuntimeError, match=r"rank 1 failed to agree on smallest.*memory"): - _session(rank=0).agree("smallest", lambda: 4, min) + with pytest.raises(RuntimeError, match=r"rank 1 failed to compute smallest.*memory"): + _session(rank=0).all_reduce("smallest", lambda: 4, min) -def test_ranks_agreeing_out_of_step_fail(monkeypatch): +def test_ranks_reducing_out_of_step_fail(monkeypatch): _gather_as_rank(monkeypatch, 0, others=[("other", 8, None)]) with pytest.raises(RuntimeError, match="out of step"): - _session(rank=0).agree("smallest", lambda: 4, min) + _session(rank=0).all_reduce("smallest", lambda: 4, min) def test_reset_clears_synced_facts(): session = _session() - session.plan.decide(lambda: DistributionPlan(world_size=1)) + session.plan.broadcast_from(lambda: DistributionPlan(world_size=1)) session._reset() assert session.plan.value is None From 6797d3a9908dee05228e23f1a92e64570e5a59d1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:12:46 +0000 Subject: [PATCH 08/48] Drop the session object; keep the process group's lifetime and the collectives as functions Library APIs take what they need as arguments, so the session's identity, synced facts and records leave this series, to be designed on their own later. What the series uses stays, as functions: distributed_run() owns the process group for a block and prints a failing rank's error before the cleanup barrier; broadcast_from and all_reduce run fn on one rank or on every rank, and when any rank fails, every rank raises instead of waiting in the collective. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 376 ++++-------------- ...ted_session.py => test_distributed_run.py} | 211 ++++------ 2 files changed, 142 insertions(+), 445 deletions(-) rename tests/unit/torch/utils/{test_distributed_session.py => test_distributed_run.py} (52%) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 12057d5da08..8c189c23f9e 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -15,9 +15,7 @@ """Utility functions for using torch.distributed.""" -import copy import functools -import inspect import io import os import sys @@ -26,7 +24,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager, suppress from datetime import timedelta -from typing import Any, Generic, TypeVar +from typing import Any, TypeVar from warnings import warn import torch @@ -37,22 +35,21 @@ __all__ = [ "DistributedProcessGroup", "DistributionPlan", - "ModelOptSession", "ParallelState", - "Synced", + "all_reduce", "backend", "barrier", + "broadcast_from", + "distributed_run", "fsdp2_wrap", "gather_bytes", "gather_state_dicts", - "get_active_mo_session", "is_available", "is_distributed_launch", "is_distributed_run", "is_fsdp2_model", "is_initialized", "is_master", - "modelopt_session", "rank", "size", ] @@ -110,9 +107,8 @@ def is_last_process(group=None) -> bool: def is_distributed_launch() -> bool: """True if launched distributed (WORLD_SIZE>1), readable BEFORE the process group is set up. - Env-based (does NOT touch the guarded session -- see :func:`get_active_mo_session`), so parse-time code - can detect the launch mode before :func:`modelopt_session` exists. After the session is entered, prefer - ``get_active_mo_session().distributed``. + Env-based, so parse-time code can detect the launch mode before :func:`distributed_run` sets the + process group up. Once it has, :func:`is_distributed_run` reads the live group. """ return int(os.environ.get("WORLD_SIZE", "1")) > 1 @@ -340,321 +336,96 @@ def __repr__(self) -> str: ) -_T = TypeVar("_T") _V = TypeVar("_V") -class Synced(Generic[_T]): - """A fact of the session every rank holds the same value of, declared by its type. +def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: + """Return ``fn()`` run on ``rank`` only, on every rank. - ``plan: Synced[DistributionPlan]`` on the session makes ``session.plan`` this handle. Read the - value with ``session.plan.value``: local, ``None`` until it is set. Set it with - :meth:`broadcast_from`, one rank computing it, or :meth:`all_reduce`, every rank computing its own - and the results combined; both are collective, so every rank calls them at the same point. - Assigning ``session.plan`` raises, so nothing communicates behind an ``=``. Declare one only for a - fact a later stage reads back: a one-off value nothing reads back is - :meth:`ModelOptSession.broadcast_from` or :meth:`ModelOptSession.all_reduce`, which keep nothing. + For a value that depends on what each rank reads for itself -- the environment, its GPU, files + on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, every + rank raises, instead of waiting in a broadcast ``rank`` never sends. Ranks broadcasting + different labels at the same point fail. Collective: every rank calls it at the same point, + with the same ``rank``. Outside a distributed run ``fn`` runs alone. """ - - def __init__(self, holder: "_SyncedHolder", name: str) -> None: - """The fact ``name`` of ``holder``.""" - self._holder = holder - self._name = name - - @property - def value(self) -> _T | None: - """The value every rank holds, read locally; ``None`` until one is set.""" - return self._holder._synced.get(self._name) - - def broadcast_from(self, fn: Callable[[], _T], rank: int = 0) -> _T: - """Set the fact to ``fn()`` run on ``rank`` only; every rank holds and returns that value. - - See :meth:`ModelOptSession.broadcast_from`. - """ - label = self._holder._label(self._name) - return self._hold(self._holder._sync_session().broadcast_from(label, fn, rank)) - - def all_reduce(self, fn: Callable[[], _T], combine: Callable[[list[_T]], _T]) -> _T: - """Set the fact to ``combine`` of every rank's ``fn()``; every rank holds and returns it. - - See :meth:`ModelOptSession.all_reduce`. - """ - label = self._holder._label(self._name) - return self._hold(self._holder._sync_session().all_reduce(label, fn, combine)) - - def _hold(self, value: _T) -> _T: - # The holder keeps its own copy: a later mutation of the caller's object stays the caller's. - self._holder._synced[self._name] = copy.deepcopy(value) - return self._holder._synced[self._name] - - -class _SyncedAttribute: - """What a ``Synced[...]`` annotation installs: reading it gives the fact's :class:`Synced` handle.""" - - def __init__(self, name: str) -> None: - self.name = name - - def __get__(self, holder, owner=None): - return self if holder is None else Synced(holder, self.name) - - def __set__(self, holder, value) -> None: - raise AttributeError( - f"{type(holder).__name__}.{self.name} is synced: set it with .broadcast_from(fn) or " - ".all_reduce(fn, combine), which every rank calls at the same point." - ) - - -class _SyncedHolder: - """Base of a class with synced facts: each ``Synced[...]`` annotation of a subclass becomes one.""" - - _synced: dict[str, Any] - - def __init_subclass__(cls, **kwargs: Any) -> None: - super().__init_subclass__(**kwargs) - for name, annotation in inspect.get_annotations(cls, eval_str=True).items(): - if getattr(annotation, "__origin__", None) is Synced: - setattr(cls, name, _SyncedAttribute(name)) - - def _sync_session(self) -> "ModelOptSession": - """The session whose ranks hold this object's synced facts.""" - raise NotImplementedError - - def _label(self, name: str) -> str: - """How the fact ``name`` is named to the other ranks, which must name the same one.""" - return name - - -class ModelOptSession(_SyncedHolder): - """The live distributed session of a PTQ run: this process's identity, plus the model's EP mesh. - - One class covers both cases: ``distributed`` is False for a single process (``WORLD_SIZE==1``) and - True for a real ``torch.distributed`` launch (``WORLD_SIZE>1``). Two parts, with different lifetimes: - - * process IDENTITY (``distributed``/``rank``/``world_size``/``local_rank``/``device``/``is_main``) -- - stable for the whole run, populated once by :func:`modelopt_session`. Reading ANY before - initialization raises, so code cannot silently consume a stale/zero default before the process - group is set up. - * ``ep_mesh`` -- the expert-parallel device mesh of the loaded model, set by ``load_and_prepare_ep`` - (``None`` until then, and for a model loaded without expert parallelism). - - Facts annotated :class:`Synced`, such as ``plan``, hold the same value on every rank. - """ - - #: The :class:`DistributionPlan` of the model this session runs, the same on every rank. - plan: Synced[DistributionPlan] - - def __init__(self) -> None: - """Create an empty (uninitialized) session; populate identity with :meth:`initialize`.""" - self._values: dict[str, Any] | None = None - self._synced: dict[str, Any] = {} - self.ep_mesh = None - - def initialize( - self, - *, - distributed: bool, - rank: int, - world_size: int, - local_rank: int, - device, - is_main: bool, - ) -> None: - """Populate the session fields (called by ``modelopt_session``, or directly for tests).""" - self._values = { - "distributed": distributed, - "rank": rank, - "world_size": world_size, - "local_rank": local_rank, - "device": device, - "is_main": is_main, - } - - def _reset(self) -> None: - self._values = None - self._synced = {} - self.ep_mesh = None - - def _sync_session(self) -> "ModelOptSession": - return self - - def broadcast_from(self, label: str, fn: Callable[[], _V], rank: int = 0) -> _V: - """Return ``fn()`` run on ``rank`` only, on every rank; keep nothing. - - For a value that depends on what each rank reads for itself -- the environment, its GPU, files - on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, every - rank raises, instead of waiting in a broadcast ``rank`` never sends. Ranks broadcasting - different labels at the same point fail. Collective: every rank calls it at the same point, - with the same ``rank``. Outside a distributed session ``fn`` runs alone. A fact a later stage - reads back is a :class:`Synced` attribute, set with its own ``broadcast_from``. - """ - if not (self.initialized and self.distributed): - return fn() - if not 0 <= rank < self.world_size: - raise ValueError(f"rank={rank} is not a rank of this {self.world_size}-rank run.") - sending = torch.distributed.get_rank() == rank - value: Any = None - error = None - if sending: - try: - value = fn() - except Exception as e: - # Still send: the other ranks would otherwise wait in the broadcast for this one. - error = e - # One message from the sending rank: its label, so a rank broadcasting another learns the two - # are out of step, and its value or its error. The error goes as text: the other ranks only - # need to stop, and an exception is not guaranteed to pickle. - message = [(label, value, None if error is None else f"{type(error).__name__}: {error}")] - torch.distributed.broadcast_object_list(message, group_src=rank) - sent_label, value, sent_error = message[0] - if error is not None: - raise error - if sent_error is not None: - raise RuntimeError(f"rank {rank} failed to compute {label}: {sent_error}") - if sent_label != label: - raise RuntimeError( - f"Ranks broadcast out of step: this rank broadcast {label!r} while rank {rank} sent " - f"{sent_label!r}. Every rank must broadcast them in the same order." - ) - return value - - def all_reduce(self, label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _V]) -> _V: - """Return ``combine`` of every rank's ``fn()``, on every rank; keep nothing. - - For a value no later stage reads back from the session; a fact that one does is a - :class:`Synced` attribute, set with its own ``all_reduce``. ``combine`` takes the ranks' values - in rank order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's - memory bounds, where one rank's choice could run another out of memory. If ``fn`` raises on any - rank, every rank raises, instead of waiting in the gather that rank never joins. Ranks reducing - different labels at the same point fail. Collective; outside a distributed session ``fn`` runs - alone. - """ - if not (self.initialized and self.distributed): - return combine([fn()]) - value, error = None, None + if not is_distributed_run(): + return fn() + if not 0 <= rank < size(): + raise ValueError(f"rank={rank} is not a rank of this {size()}-rank run.") + sending = torch.distributed.get_rank() == rank + value: Any = None + error = None + if sending: try: value = fn() except Exception as e: - # Still join the gather: the other ranks would otherwise wait in it for this one. + # Still send: the other ranks would otherwise wait in the broadcast for this one. error = e - mine = (label, value, None if error is None else f"{type(error).__name__}: {error}") - gathered: list[Any] = [None] * torch.distributed.get_world_size() - torch.distributed.all_gather_object(gathered, mine) - if error is not None: - raise error - for rank, (_, _, rank_error) in enumerate(gathered): - if rank_error is not None: - raise RuntimeError(f"rank {rank} failed to compute {label}: {rank_error}") - labels = sorted({rank_label for rank_label, _, _ in gathered}) - if labels != [label]: - raise RuntimeError( - f"Ranks reduced out of step: {labels} at the same point. Every rank must reduce " - "them in the same order." - ) - # The gathered values are unpickled copies, so the result shares nothing with the callers'. - return combine([rank_value for _, rank_value, _ in gathered]) - - @property - def initialized(self) -> bool: - """Whether ``modelopt_session`` has populated the state.""" - return self._values is not None - - def _get(self, key: str): - if self._values is None: - raise RuntimeError( - "session accessed before initialization; enter " - "modelopt.torch.utils.distributed.modelopt_session() first" - ) - return self._values[key] - - @property - def distributed(self) -> bool: - """True for a distributed launch (WORLD_SIZE>1), False for single-process.""" - return self._get("distributed") - - @property - def rank(self) -> int: - """Global rank of this process (0 single-process).""" - return self._get("rank") - - @property - def world_size(self) -> int: - """Total number of processes (1 single-process).""" - return self._get("world_size") - - @property - def local_rank(self) -> int: - """Node-local rank of this process (0 single-process).""" - return self._get("local_rank") - - @property - def device(self): - """This process's torch device (``cuda:local_rank`` distributed; the caller default otherwise).""" - return self._get("device") - - @property - def is_main(self) -> bool: - """Whether this is the main process (rank 0).""" - return self._get("is_main") - - -# The process-wide ACTIVE session: a single STABLE object created once at import and NEVER rebound. -# :func:`modelopt_session` mutates it in place via :meth:`ModelOptSession.initialize` / ``_reset`` rather -# than assigning a new object -- so any held reference stays valid and never goes stale. Module-PRIVATE -# (not exported): reach it only through :func:`get_active_mo_session`. It exists before initialization; -# reading its IDENTITY fields (``rank`` etc.) then raises via the per-field guard. -_active_mo_session: ModelOptSession = ModelOptSession() + # One message from the sending rank: its label, so a rank broadcasting another learns the two + # are out of step, and its value or its error. The error goes as text: the other ranks only + # need to stop, and an exception is not guaranteed to pickle. + message = [(label, value, None if error is None else f"{type(error).__name__}: {error}")] + torch.distributed.broadcast_object_list(message, group_src=rank) + sent_label, value, sent_error = message[0] + if error is not None: + raise error + if sent_error is not None: + raise RuntimeError(f"rank {rank} failed to compute {label}: {sent_error}") + if sent_label != label: + raise RuntimeError( + f"Ranks broadcast out of step: this rank broadcast {label!r} while rank {rank} sent " + f"{sent_label!r}. Every rank must broadcast them in the same order." + ) + return value -def get_active_mo_session(required: bool = True) -> ModelOptSession: - """Return the process-wide session. +def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _V]) -> _V: + """Return ``combine`` of every rank's ``fn()``, on every rank. - * ``required=True`` (default) -- fail-fast: raise if no session has been entered yet by - :func:`modelopt_session`. Use when the caller genuinely needs a live session. - * ``required=False`` -- return the singleton regardless of initialization. Its ``ep_mesh`` is - readable (``None``) even before a session is entered, so this is for defensive reads deep in public - APIs (export / calibration) that may run with no active session. + ``combine`` takes the ranks' values in rank order and returns the one every rank gets, e.g. + ``min`` for a batch size each rank's memory bounds, where one rank's choice could run another + out of memory. If ``fn`` raises on any rank, every rank raises, instead of waiting in the gather + that rank never joins. Ranks reducing different labels at the same point fail. Collective; + outside a distributed run ``fn`` runs alone. """ - if required and not _active_mo_session.initialized: + if not is_distributed_run(): + return combine([fn()]) + value, error = None, None + try: + value = fn() + except Exception as e: + # Still join the gather: the other ranks would otherwise wait in it for this one. + error = e + mine = (label, value, None if error is None else f"{type(error).__name__}: {error}") + gathered: list[Any] = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(gathered, mine) + if error is not None: + raise error + for rank_, (_, _, rank_error) in enumerate(gathered): + if rank_error is not None: + raise RuntimeError(f"rank {rank_} failed to compute {label}: {rank_error}") + labels = sorted({rank_label for rank_label, _, _ in gathered}) + if labels != [label]: raise RuntimeError( - "no active session; enter modelopt.torch.utils.distributed.modelopt_session() first" + f"Ranks reduced out of step: {labels} at the same point. Every rank must reduce " + "them in the same order." ) - return _active_mo_session + # The gathered values are unpickled copies, so the result shares nothing with the callers'. + return combine([rank_value for _, rank_value, _ in gathered]) @contextmanager -def modelopt_session(timeout: timedelta | None = None, device=None): - """Populate the active session for the block, then tear down. +def distributed_run(timeout: timedelta | None = None): + """Own the process group for the block on a distributed launch; do nothing otherwise. - On a distributed launch (WORLD_SIZE>1) this initializes the process group (binding the local CUDA - device) with ``timeout`` and records rank/world_size/local_rank/``cuda:local_rank``/is_main; it - DESTROYS the group on exit even if the body raises. Single-process records the trivial single-rank - topology -- with ``device`` as the device (the caller's default, e.g. the ``--device`` CLI value, since - there is no local rank to derive one) -- and enter/exit are no-ops around the body. - - Yields the session (also reachable via :func:`get_active_mo_session`); it is reset on exit. + On a distributed launch (``WORLD_SIZE>1``) this sets the process group up, binding this rank's + CUDA device, with ``timeout``, and destroys it on exit even if the body raises. """ owns_group = is_distributed_launch() if owns_group: setup(timeout=timeout) - _active_mo_session.initialize( - distributed=True, - rank=rank(), - world_size=size(), - local_rank=local_rank(), - device=torch.device(f"cuda:{local_rank()}"), - is_main=rank() == 0, - ) - else: - _active_mo_session.initialize( - distributed=False, - rank=0, - world_size=1, - local_rank=0, - device=torch.device(device) if device is not None else None, - is_main=True, - ) try: - yield _active_mo_session + yield finally: # If the body is unwinding due to an exception, PRINT it here BEFORE cleanup()'s collective # barrier. A per-rank exception (only some ranks raise) otherwise makes those ranks enter the @@ -663,7 +434,7 @@ def modelopt_session(timeout: timedelta | None = None, device=None): exc = sys.exc_info()[1] if exc is not None: print( - f"[modelopt_session][rank {os.environ.get('RANK', '?')}] body raised " + f"[distributed_run][rank {os.environ.get('RANK', '?')}] body raised " f"{type(exc).__name__}: {exc}", flush=True, ) @@ -671,7 +442,6 @@ def modelopt_session(timeout: timedelta | None = None, device=None): sys.stderr.flush() if owns_group: cleanup() - _active_mo_session._reset() def cleanup(): diff --git a/tests/unit/torch/utils/test_distributed_session.py b/tests/unit/torch/utils/test_distributed_run.py similarity index 52% rename from tests/unit/torch/utils/test_distributed_session.py rename to tests/unit/torch/utils/test_distributed_run.py index de3dcee327b..74600954bab 100644 --- a/tests/unit/torch/utils/test_distributed_session.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -13,17 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the distributed PTQ session: the distribution plan, the session handle, its synced -facts and one-off collectives, launch detection.""" + +"""Unit tests for distributed PTQ: the distribution plan, the collectives that fail together, the +process group's lifetime, launch detection.""" + +from datetime import timedelta import pytest import torch.distributed +import modelopt.torch.utils.distributed as dist from modelopt.torch.utils.distributed import ( DistributionPlan, - ModelOptSession, - Synced, - get_active_mo_session, + all_reduce, + broadcast_from, + distributed_run, is_distributed_launch, ) @@ -62,69 +66,21 @@ def test_repr_names_the_degrees(): assert "fsdp_size=8" in repr(DistributionPlan(world_size=8, ep_size=8, fsdp_size=8)) -# ------------------------------------------------------------------------------------- session handle - - -def test_required_session_fails_fast_when_none_was_entered(): - with pytest.raises(RuntimeError, match="no active session"): - get_active_mo_session() - - -def test_optional_session_is_readable_without_a_session(): - """Export / calibration read the session defensively; they must not explode with no session.""" - session = get_active_mo_session(required=False) - assert session is not None - assert not session.initialized - assert session.ep_mesh is None - - -def test_identity_reads_raise_before_initialization(): - """A stale or zero default would silently mislabel every rank; reading early must fail instead.""" - with pytest.raises(RuntimeError, match="before initialization"): - _ = ModelOptSession().rank - - -def test_reset_clears_the_identity_and_the_ep_mesh(): - session = ModelOptSession() - session.initialize( - distributed=True, rank=1, world_size=2, local_rank=1, device=None, is_main=False - ) - session.ep_mesh = object() - assert (session.rank, session.world_size, session.is_main) == (1, 2, False) - session._reset() - assert not session.initialized - assert session.ep_mesh is None +# ------------------------------------------------------------------------------------ collectives -# ----------------------------------------------------------------------------------- synced facts - - -class _Session(ModelOptSession): - """A session with a second synced fact, and a plain attribute.""" - - other: Synced[int] - note: str - - -def _session(rank=None): - """Rank ``rank`` of a two-rank session; ``None``: a single process.""" - session = _Session() - session.initialize( - distributed=rank is not None, - rank=rank or 0, - world_size=1 if rank is None else 2, - local_rank=rank or 0, - device=None, - is_main=not rank, - ) - return session +def _run_with_ranks(monkeypatch, world_size): + """Make this process look like one rank of a live ``world_size``-rank run.""" + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size) def _as_rank(monkeypatch, rank, received=None): - """Answer the session's broadcasts as ``rank``; returns ``(message, sending rank)`` for each. + """Answer the broadcasts of a two-rank run as ``rank``; returns ``(message, sending rank)`` for each. A rank that is not the sending one receives ``received``, as the sending rank would have sent it. """ + _run_with_ranks(monkeypatch, 2) sent = [] def broadcast_object_list(message, group=None, group_src=None): @@ -138,17 +94,17 @@ def broadcast_object_list(message, group=None, group_src=None): def _gather_as_rank(monkeypatch, rank, others): - """Answer the session's gathers as ``rank``; returns what each gather was handed. + """Answer the gathers as ``rank``; returns what each gather was handed. ``others`` are the other ranks' messages in rank order, as they would have sent them. """ + _run_with_ranks(monkeypatch, len(others) + 1) sent = [] def all_gather_object(gathered, mine, group=None): sent.append(mine) gathered[:] = [*others[:rank], mine, *others[rank:]] - monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: len(others) + 1) monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) monkeypatch.setattr( torch.distributed, "broadcast_object_list", lambda *a, **k: pytest.fail("a rank broadcast") @@ -156,82 +112,47 @@ def all_gather_object(gathered, mine, group=None): return sent -def test_only_facts_annotated_synced_are_synced(): - """The annotation carries both the value's type and the synced semantic.""" - session = _session() - assert isinstance(session.other, Synced) - assert isinstance(session.plan, Synced) - assert "note" not in vars(_Session) - - -def test_a_synced_fact_cannot_be_assigned(): - """Setting one is collective: broadcast_from or all_reduce, never a bare assignment.""" - with pytest.raises(AttributeError, match=r"plan is synced: set it with \.broadcast_from"): - _session().plan = DistributionPlan(world_size=1) - - -def test_a_plain_attribute_assigns_without_a_collective(monkeypatch): - for collective in ("broadcast_object_list", "all_gather_object"): - monkeypatch.setattr( - torch.distributed, collective, lambda *a, **k: pytest.fail("a collective ran") - ) - session = _session(rank=1) - session.note = "local" - assert session.note == "local" - - def test_a_single_process_computes_alone(monkeypatch): - """Nothing to sync with: no collective, and an uninitialized session behaves the same.""" + """Nothing to sync with: no collective.""" for collective in ("broadcast_object_list", "all_gather_object"): monkeypatch.setattr( torch.distributed, collective, lambda *a, **k: pytest.fail("a collective ran") ) - for session in (_session(), _Session()): - plan = DistributionPlan(world_size=1) - assert session.plan.value is None - assert session.plan.broadcast_from(lambda: plan).world_size == 1 - assert session.plan.value.world_size == 1 - assert session.other.all_reduce(lambda: 4, min) == 4 - assert session.broadcast_from("x", lambda: 3) == 3 - assert session.all_reduce("smallest", lambda: 4, min) == 4 + assert broadcast_from("x", lambda: 3) == 3 + assert all_reduce("smallest", lambda: 4, min) == 4 -def test_the_sending_rank_sends_its_value_and_keeps_its_own_copy(monkeypatch): +def test_the_sending_rank_sends_its_value(monkeypatch): sent = _as_rank(monkeypatch, 0) - session, plan = _session(rank=0), DistributionPlan(world_size=2, ep_size=2) - - held = session.plan.broadcast_from(lambda: plan) + plan = DistributionPlan(world_size=2, ep_size=2) + assert broadcast_from("plan", lambda: plan) is plan assert sent == [(("plan", plan, None), 0)] - assert held is session.plan.value - assert held.ep_size == 2 - assert held is not plan, "mutating the caller's object must not reach the session" -def test_other_ranks_hold_the_sending_ranks_value_without_running_fn(monkeypatch): +def test_other_ranks_take_the_sending_ranks_value_without_running_fn(monkeypatch): rank0_plan = DistributionPlan(world_size=2, ep_size=2) _as_rank(monkeypatch, 1, received=("plan", rank0_plan, None)) - session = _session(rank=1) - held = session.plan.broadcast_from(lambda: pytest.fail("a rank that does not send ran fn")) + plan = broadcast_from("plan", lambda: pytest.fail("a rank that does not send ran fn")) - assert held.ep_size == session.plan.value.ep_size == 2 + assert plan.ep_size == 2 def test_any_rank_can_broadcast(monkeypatch): """With rank=1, rank 1 runs fn and sends it; rank 0 takes rank 1's value.""" sent = _as_rank(monkeypatch, 1) - assert _session(rank=1).broadcast_from("x", lambda: "rank 1's", rank=1) == "rank 1's" + assert broadcast_from("x", lambda: "rank 1's", rank=1) == "rank 1's" assert sent == [(("x", "rank 1's", None), 1)] _as_rank(monkeypatch, 0, received=("x", "rank 1's", None)) - held = _session(rank=0).broadcast_from("x", lambda: pytest.fail("rank 0 ran fn"), rank=1) - assert held == "rank 1's" + assert broadcast_from("x", lambda: pytest.fail("rank 0 ran fn"), rank=1) == "rank 1's" -def test_a_rank_outside_the_run_cannot_broadcast(): +def test_a_rank_outside_the_run_cannot_broadcast(monkeypatch): + _run_with_ranks(monkeypatch, 2) with pytest.raises(ValueError, match="rank=2 is not a rank of this 2-rank run"): - _session(rank=0).broadcast_from("x", lambda: 0, rank=2) + broadcast_from("x", lambda: 0, rank=2) def test_the_sending_rank_failing_still_reaches_the_broadcast(monkeypatch): @@ -242,48 +163,31 @@ def fail(): raise ValueError("bad degrees") with pytest.raises(ValueError, match="bad degrees"): - _session(rank=0).broadcast_from("x", fail) + broadcast_from("x", fail) assert sent == [(("x", None, "ValueError: bad degrees"), 0)] def test_other_ranks_raise_when_the_sending_rank_failed(monkeypatch): _as_rank(monkeypatch, 1, received=("x", None, "ValueError: bad degrees")) with pytest.raises(RuntimeError, match=r"rank 0 failed to compute x.*bad degrees"): - _session(rank=1).broadcast_from("x", lambda: 0) + broadcast_from("x", lambda: 0) def test_ranks_broadcasting_out_of_step_fail(monkeypatch): - """Rank 0 sending one fact while rank 1 expects another would cross the two values.""" + """Rank 0 sending one value while rank 1 expects another would cross the two.""" _as_rank(monkeypatch, 1, received=("other", 7, None)) with pytest.raises(RuntimeError, match="out of step"): - _session(rank=1).plan.broadcast_from(lambda: DistributionPlan(world_size=2)) + broadcast_from("plan", lambda: DistributionPlan(world_size=2)) def test_all_reduce_gives_every_rank_the_combined_value(monkeypatch): """Every rank's value counts, and every rank gets the same result.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None), ("smallest", 2, None)]) - assert _session(rank=1).all_reduce("smallest", lambda: 4, min) == 2 + assert all_reduce("smallest", lambda: 4, min) == 2 assert sent == [("smallest", 4, None)] -def test_a_facts_all_reduce_holds_the_combined_value(monkeypatch): - _gather_as_rank(monkeypatch, 1, others=[("other", 8, None)]) - session = _session(rank=1) - - assert session.other.all_reduce(lambda: 4, min) == 4 - assert session.other.value == 4 - - -def test_the_sessions_own_all_reduce_keeps_nothing(monkeypatch): - _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) - session = _session(rank=1) - - session.all_reduce("smallest", lambda: 4, min) - - assert session._synced == {} - - def test_a_rank_failing_to_compute_still_joins_the_gather(monkeypatch): """A failing rank must still reach the gather, or every other rank waits in it.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) @@ -292,27 +196,50 @@ def fail(): raise ValueError("no memory reading") with pytest.raises(ValueError, match="no memory reading"): - _session(rank=1).all_reduce("smallest", fail, min) + all_reduce("smallest", fail, min) assert sent == [("smallest", None, "ValueError: no memory reading")] def test_the_other_ranks_raise_when_one_failed_to_compute(monkeypatch): _gather_as_rank(monkeypatch, 0, others=[("smallest", None, "ValueError: no memory reading")]) with pytest.raises(RuntimeError, match=r"rank 1 failed to compute smallest.*memory"): - _session(rank=0).all_reduce("smallest", lambda: 4, min) + all_reduce("smallest", lambda: 4, min) def test_ranks_reducing_out_of_step_fail(monkeypatch): _gather_as_rank(monkeypatch, 0, others=[("other", 8, None)]) with pytest.raises(RuntimeError, match="out of step"): - _session(rank=0).all_reduce("smallest", lambda: 4, min) + all_reduce("smallest", lambda: 4, min) + + +# ---------------------------------------------------------------------------------- process group + + +def test_a_single_process_sets_no_group_up(monkeypatch): + monkeypatch.delenv("WORLD_SIZE", raising=False) + monkeypatch.setattr(dist, "setup", lambda **kwargs: pytest.fail("set a group up")) + monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("tore a group down")) + with distributed_run(): + pass + + +def test_a_distributed_launch_owns_the_group_even_when_the_body_raises(monkeypatch, capsys): + """The group is torn down on the way out, and the body's error is printed first: a rank failing + alone would otherwise wait in the cleanup barrier while the rest sit in a collective.""" + monkeypatch.setenv("WORLD_SIZE", "2") + calls = [] + monkeypatch.setattr(dist, "setup", lambda timeout=None: calls.append(("setup", timeout))) + monkeypatch.setattr(dist, "cleanup", lambda: calls.append(("cleanup", capsys.readouterr().out))) + with ( + pytest.raises(ValueError, match="calibration failed"), + distributed_run(timeout=timedelta(hours=1)), + ): + raise ValueError("calibration failed") -def test_reset_clears_synced_facts(): - session = _session() - session.plan.broadcast_from(lambda: DistributionPlan(world_size=1)) - session._reset() - assert session.plan.value is None + assert calls[0] == ("setup", timedelta(hours=1)) + assert calls[1][0] == "cleanup" + assert "calibration failed" in calls[1][1], "the error must be printed before the cleanup" # ------------------------------------------------------------------------------------ launch detect @@ -320,7 +247,7 @@ def test_reset_clears_synced_facts(): @pytest.mark.parametrize(("world", "expected"), [(None, False), ("1", False), ("8", True)]) def test_launch_detection_reads_the_env_before_any_process_group(monkeypatch, world, expected): - """Must work at parse time, before modelopt_session exists.""" + """Must work at parse time, before distributed_run sets the process group up.""" monkeypatch.delenv("WORLD_SIZE", raising=False) if world is not None: monkeypatch.setenv("WORLD_SIZE", world) From 82a4e2bedc0347671a30e913110c8c5a25efb260 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 19:23:20 +0000 Subject: [PATCH 09/48] Record a model's distribution plan on it, and make the plan a value with a dense layout DistributionPlan is now how any model is distributed, not only an MoE: a frozen value of world_size, ep_size and dense ('replicate' or 'fsdp'), which replaces fsdp_size. A single process is the default plan, and FSDP2 on one rank is expressible. set_distribution_plan records a plan on a model and every submodule; distribution_plan reads it back. A model with none recorded is a single process's, unless it is FSDP2-sharded or held in a run of several ranks, which raises instead of being read as one. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 75 ++++++++++++++----- .../unit/torch/utils/test_distributed_run.py | 53 +++++++++++-- 2 files changed, 100 insertions(+), 28 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 8c189c23f9e..cded777ced0 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -23,8 +23,9 @@ import traceback from collections.abc import Callable, Iterator from contextlib import contextmanager, suppress +from dataclasses import dataclass from datetime import timedelta -from typing import Any, TypeVar +from typing import Any, Literal, TypeVar from warnings import warn import torch @@ -41,6 +42,7 @@ "barrier", "broadcast_from", "distributed_run", + "distribution_plan", "fsdp2_wrap", "gather_bytes", "gather_state_dicts", @@ -51,6 +53,7 @@ "is_initialized", "is_master", "rank", + "set_distribution_plan", "size", ] @@ -294,24 +297,30 @@ def setup(timeout: timedelta | None = None): torch.distributed.init_process_group("cpu:gloo,cuda:nccl", timeout=timeout) +@dataclass(frozen=True) class DistributionPlan: - """How a HuggingFace MoE model is distributed for PTQ, as decided by ``decide_distribution_plan``. + """How a model is distributed over the ranks of a run. - * ``ep_size`` -- expert-parallel degree: the routed experts are sharded over ``ep_size`` ranks. - * ``fsdp_size`` -- dense-FSDP degree: ``1`` replicates the dense (non-expert) part on every rank; ``>1`` - FSDP2-shards it alongside the expert-parallel experts. + * ``world_size`` -- the ranks that hold the model; ``1`` for a single process. + * ``ep_size`` -- the expert-parallel degree: an MoE's routed experts are sharded over groups of + ``ep_size`` ranks. ``1`` shards none. + * ``dense`` -- how the rest of the model, all of it without expert parallelism, is placed: + ``"replicate"`` gives every rank a whole copy, ``"fsdp"`` FSDP2-shards it. - Every rank calibrates on its own data, so data-parallelism spans the whole world, and each expert + Every rank calibrates on its own data, so data parallelism spans the whole world, and each expert shard is replicated ``edp = world_size / ep_size`` times. """ - def __init__(self, world_size: int, ep_size: int = 1, fsdp_size: int = 1) -> None: - """Build the plan, rejecting an expert-parallel degree that cannot divide the world.""" - if world_size % ep_size != 0: - raise ValueError(f"ep_size={ep_size} must divide world_size={world_size}") - self.world_size = world_size - self.ep_size = ep_size - self.fsdp_size = fsdp_size + world_size: int = 1 + ep_size: int = 1 + dense: Literal["replicate", "fsdp"] = "replicate" + + def __post_init__(self) -> None: + """Reject a plan no load could carry out, before any rank starts on it.""" + if self.world_size % self.ep_size != 0: + raise ValueError(f"ep_size={self.ep_size} must divide world_size={self.world_size}") + if self.dense not in ("replicate", "fsdp"): + raise ValueError(f"dense={self.dense!r}: expected 'replicate' or 'fsdp'.") @property def edp(self) -> int: @@ -325,15 +334,41 @@ def is_expert_parallel(self) -> bool: @property def is_fsdp_dense(self) -> bool: - """Whether the dense part is FSDP-sharded rather than replicated.""" - return self.fsdp_size > 1 + """Whether the dense part is FSDP2-sharded rather than replicated.""" + return self.dense == "fsdp" - def __repr__(self) -> str: - dense = f"fsdp_size={self.fsdp_size}" if self.is_fsdp_dense else "dense=replicate" - return ( - f"DistributionPlan(world={self.world_size}, ep_size={self.ep_size}, edp={self.edp}, " - f"{dense})" + +_PLAN_ATTR = "_modelopt_distribution_plan" + + +def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> None: + """Record that ``model`` is distributed as ``plan``, on it and on every submodule. + + Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed + the whole model. ``load_pretrained`` records the plan it loads by; code that distributes a model + itself records it here before calibrating or exporting the model. + """ + for module in model.modules(): + setattr(module, _PLAN_ATTR, plan) + + +def distribution_plan(model: torch.nn.Module) -> DistributionPlan: + """How ``model`` is distributed: the plan recorded on it (:func:`set_distribution_plan`). + + A model with no plan recorded is a single process's, unless it shows otherwise: one that is + FSDP2-sharded, or held in a run of several ranks, raises rather than be taken for one. + """ + plan = getattr(model, _PLAN_ATTR, None) + if plan is not None: + return plan + sharded = is_fsdp2_model(model) + if sharded or size() > 1: + how = "is FSDP2-sharded" if sharded else f"is held in a run of {size()} ranks" + raise RuntimeError( + f"{type(model).__name__} {how}, but has no distribution plan recorded. Load it with " + "load_pretrained, or record how it is distributed with set_distribution_plan()." ) + return DistributionPlan() _V = TypeVar("_V") diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 74600954bab..005967a4dc6 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -14,13 +14,14 @@ # limitations under the License. -"""Unit tests for distributed PTQ: the distribution plan, the collectives that fail together, the -process group's lifetime, launch detection.""" +"""Unit tests for distributed PTQ: the distribution plan and a model's record of it, the collectives +that fail together, the process group's lifetime, launch detection.""" from datetime import timedelta import pytest import torch.distributed +from torch import nn import modelopt.torch.utils.distributed as dist from modelopt.torch.utils.distributed import ( @@ -28,7 +29,9 @@ all_reduce, broadcast_from, distributed_run, + distribution_plan, is_distributed_launch, + set_distribution_plan, ) # --------------------------------------------------------------------------------- DistributionPlan @@ -48,22 +51,56 @@ def test_expert_parallel_splits_the_expert_replica_pool(): def test_fsdp_dense_plan(): - plan = DistributionPlan(world_size=8, ep_size=8, fsdp_size=8) + plan = DistributionPlan(world_size=8, ep_size=8, dense="fsdp") assert plan.is_fsdp_dense assert plan.edp == 1 +def test_the_default_plan_is_a_single_process(): + assert DistributionPlan() == DistributionPlan(world_size=1, ep_size=1, dense="replicate") + + def test_an_ep_degree_that_does_not_divide_the_world_is_rejected(): """An unsatisfiable plan must fail here, not as a confusing mesh error thousands of ranks later.""" with pytest.raises(ValueError, match="must divide world_size"): DistributionPlan(world_size=8, ep_size=3) -def test_repr_names_the_degrees(): - assert repr(DistributionPlan(world_size=8, ep_size=2)) == ( - "DistributionPlan(world=8, ep_size=2, edp=4, dense=replicate)" - ) - assert "fsdp_size=8" in repr(DistributionPlan(world_size=8, ep_size=8, fsdp_size=8)) +def test_an_unknown_dense_layout_is_rejected(): + with pytest.raises(ValueError, match="expected 'replicate' or 'fsdp'"): + DistributionPlan(world_size=8, dense="sharded") + + +# ------------------------------------------------------------------- a model's recorded plan + + +def _model(): + return nn.Sequential(nn.Linear(2, 2), nn.Sequential(nn.Linear(2, 2))) + + +def test_every_submodule_reads_the_recorded_plan(): + model, plan = _model(), DistributionPlan(world_size=8, ep_size=8) + set_distribution_plan(model, plan) + assert all(distribution_plan(module) is plan for module in model.modules()) + + +def test_a_model_with_no_plan_is_a_single_process(): + assert distribution_plan(_model()) == DistributionPlan() + + +def test_an_fsdp2_model_with_no_plan_raises(monkeypatch): + monkeypatch.setattr(dist, "is_fsdp2_model", lambda model: True) + with pytest.raises(RuntimeError, match="is FSDP2-sharded, but has no distribution plan"): + distribution_plan(_model()) + + +def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(monkeypatch): + monkeypatch.setattr(dist, "size", lambda group=None: 2) + with pytest.raises(RuntimeError, match="held in a run of 2 ranks"): + distribution_plan(_model()) + model, plan = _model(), DistributionPlan(world_size=2) + set_distribution_plan(model, plan) + assert distribution_plan(model) is plan # ------------------------------------------------------------------------------------ collectives From b11d6d87dbb362bd8931d388ac414c40fcc995a5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 22:16:35 +0000 Subject: [PATCH 10/48] Yield this process's DistributedState from distributed_run A run's code needs this process's rank, the run's size and the device it computes on; hf_ptq kept that in its own DistributedState, filled from the group distributed_run had set up. distributed_run now builds it: this rank's on a distributed launch, on its CUDA device, and rank 0 of a run of one, on the caller's device, for a single process. is_main is read from the rank, so a state cannot claim the wrong main rank. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 32 ++++++++++++++++--- .../unit/torch/utils/test_distributed_run.py | 20 ++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index cded777ced0..bef71e2f876 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -35,6 +35,7 @@ __all__ = [ "DistributedProcessGroup", + "DistributedState", "DistributionPlan", "ParallelState", "all_reduce", @@ -449,18 +450,41 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ return combine([rank_value for _, rank_value, _ in gathered]) +@dataclass(frozen=True) +class DistributedState: + """This process's place in a run: its rank, the run's size, and the device it computes on. + + :func:`distributed_run` yields it, for a single process too (rank 0 of a run of one). + """ + + rank: int + world_size: int + device: torch.device | str + + @property + def is_main(self) -> bool: + """Whether this is rank 0, the rank that writes the run's outputs.""" + return self.rank == 0 + + @contextmanager -def distributed_run(timeout: timedelta | None = None): - """Own the process group for the block on a distributed launch; do nothing otherwise. +def distributed_run(timeout: timedelta | None = None, device: torch.device | str = "cuda"): + """Own the process group for the block on a distributed launch, and yield this process's state. On a distributed launch (``WORLD_SIZE>1``) this sets the process group up, binding this rank's - CUDA device, with ``timeout``, and destroys it on exit even if the body raises. + CUDA device, with ``timeout``, and destroys it on exit even if the body raises. It yields this + rank's :class:`DistributedState`, on that CUDA device. A single process sets nothing up and is + rank 0 of a run of one, on ``device``. """ owns_group = is_distributed_launch() if owns_group: setup(timeout=timeout) try: - yield + yield ( + DistributedState(rank(), size(), torch.device(f"cuda:{local_rank()}")) + if owns_group + else DistributedState(0, 1, device) + ) finally: # If the body is unwinding due to an exception, PRINT it here BEFORE cleanup()'s collective # barrier. A per-rank exception (only some ranks raise) otherwise makes those ranks enter the diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 005967a4dc6..bfe2e10edf8 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -25,6 +25,7 @@ import modelopt.torch.utils.distributed as dist from modelopt.torch.utils.distributed import ( + DistributedState, DistributionPlan, all_reduce, broadcast_from, @@ -260,6 +261,25 @@ def test_a_single_process_sets_no_group_up(monkeypatch): pass +def test_a_single_process_is_rank_0_of_one_on_the_given_device(monkeypatch): + monkeypatch.delenv("WORLD_SIZE", raising=False) + with distributed_run(device="cpu") as state: + assert state == DistributedState(rank=0, world_size=1, device="cpu") + assert state.is_main + + +def test_a_distributed_launch_yields_this_ranks_state_on_its_cuda_device(monkeypatch): + monkeypatch.setenv("WORLD_SIZE", "4") + monkeypatch.setattr(dist, "setup", lambda timeout=None: None) + monkeypatch.setattr(dist, "cleanup", lambda: None) + monkeypatch.setattr(dist, "rank", lambda group=None: 2) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + monkeypatch.setattr(dist, "local_rank", lambda: 1) + with distributed_run(device="cpu") as state: + assert state == DistributedState(rank=2, world_size=4, device=torch.device("cuda:1")) + assert not state.is_main + + def test_a_distributed_launch_owns_the_group_even_when_the_body_raises(monkeypatch, capsys): """The group is torn down on the way out, and the body's error is printed first: a rank failing alone would otherwise wait in the cleanup barrier while the rest sit in a collective.""" From 7f46ef6101d78cdc5368deea319a3ab291753c75 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 22:32:09 +0000 Subject: [PATCH 11/48] Own every group a run sets up, and set up FSDP2's one-rank group here distributed_run destroyed only the group of a distributed launch; the one-rank group FSDP2 needs in a single process was hf_ptq's to set up and tear down. distributed_run now destroys whatever group the run set up, its own or one the body set up, and leaves a group that was there before the run alone. setup_fsdp2_process_group(plan) sets that one-rank group up when a single process's plan FSDP2-shards the model, so the run owns it like a launch's. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 26 ++++++++--- .../unit/torch/utils/test_distributed_run.py | 45 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index bef71e2f876..f01b3f85c89 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -55,6 +55,7 @@ "is_master", "rank", "set_distribution_plan", + "setup_fsdp2_process_group", "size", ] @@ -469,14 +470,18 @@ def is_main(self) -> bool: @contextmanager def distributed_run(timeout: timedelta | None = None, device: torch.device | str = "cuda"): - """Own the process group for the block on a distributed launch, and yield this process's state. + """Own the run's process group for the block, and yield this process's state. On a distributed launch (``WORLD_SIZE>1``) this sets the process group up, binding this rank's - CUDA device, with ``timeout``, and destroys it on exit even if the body raises. It yields this - rank's :class:`DistributedState`, on that CUDA device. A single process sets nothing up and is - rank 0 of a run of one, on ``device``. + CUDA device, with ``timeout``. It yields this rank's :class:`DistributedState`, on that CUDA + device. A single process sets nothing up and is rank 0 of a run of one, on ``device``. + + On exit, even if the body raises, it destroys the group the run set up: the launch's, or one the + body set up, such as the one-rank group of :func:`setup_fsdp2_process_group`. A group that was + already there when the run began is the caller's, and is left as it is. """ owns_group = is_distributed_launch() + group_was_there = is_initialized() if owns_group: setup(timeout=timeout) try: @@ -499,10 +504,21 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str ) traceback.print_exc() sys.stderr.flush() - if owns_group: + if owns_group or (not group_was_there and is_initialized()): cleanup() +def setup_fsdp2_process_group(plan: DistributionPlan, timeout: timedelta | None = None) -> None: + """Set up the process group a single process needs to FSDP2-shard by ``plan``. + + FSDP2 shards through a process group even on one rank. A distributed launch already has its + group (:func:`distributed_run`), and a plan that does not FSDP2-shard needs none. Inside a + :func:`distributed_run`, the run destroys the group on exit. + """ + if plan.is_fsdp_dense and not is_initialized(): + setup(timeout=timeout) + + def cleanup(): """Cleans up the distributed environment. diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index bfe2e10edf8..ffbbddc3923 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -33,6 +33,7 @@ distribution_plan, is_distributed_launch, set_distribution_plan, + setup_fsdp2_process_group, ) # --------------------------------------------------------------------------------- DistributionPlan @@ -280,6 +281,50 @@ def test_a_distributed_launch_yields_this_ranks_state_on_its_cuda_device(monkeyp assert not state.is_main +def test_the_run_destroys_a_group_its_body_set_up(monkeypatch): + """FSDP2 on one rank sets a group up mid-run; the run tears it down like a launch's.""" + monkeypatch.delenv("WORLD_SIZE", raising=False) + live, destroyed = [], [] + monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) + monkeypatch.setattr(dist, "setup", lambda timeout=None: live.append(timeout)) + monkeypatch.setattr(dist, "cleanup", lambda: destroyed.append(0)) + + with distributed_run(device="cpu"): + setup_fsdp2_process_group(DistributionPlan(dense="fsdp"), timeout=timedelta(hours=2)) + + assert live == [timedelta(hours=2)] + assert len(destroyed) == 1 + + +def test_the_run_leaves_a_group_it_found_alone(monkeypatch): + monkeypatch.delenv("WORLD_SIZE", raising=False) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("destroyed the caller's group")) + with distributed_run(device="cpu"): + pass + + +@pytest.mark.parametrize( + ("plan", "live_group", "setups"), + [ + (DistributionPlan(dense="fsdp"), False, 1), + (DistributionPlan(2, dense="fsdp"), True, 0), + (DistributionPlan(), False, 0), + ], + ids=["fsdp2_on_one_rank", "launched_group", "replicated"], +) +def test_fsdp2_on_one_rank_sets_up_a_group_of_its_own(monkeypatch, plan, live_group, setups): + """FSDP2 shards through a process group even on one rank; a distributed launch has its group, + and a plan that does not FSDP2-shard needs none.""" + calls = [] + monkeypatch.setattr(dist, "is_initialized", lambda: live_group) + monkeypatch.setattr(dist, "setup", lambda timeout=None: calls.append(timeout)) + + setup_fsdp2_process_group(plan, timeout=timedelta(hours=2)) + + assert calls == [timedelta(hours=2)] * setups + + def test_a_distributed_launch_owns_the_group_even_when_the_body_raises(monkeypatch, capsys): """The group is torn down on the way out, and the body's error is printed first: a rank failing alone would otherwise wait in the cleanup barrier while the rest sit in a collective.""" From 8bf70e23937cc36de62824ff55b2254013496b8f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 22:55:21 +0000 Subject: [PATCH 12/48] Drop setup_fsdp2_process_group for the FSDP2 load to set its group up A single process's FSDP2 plan needs a one-rank process group, and the load that carries the plan out is where it is needed: load_pretrained sets it up there ([3/21]), and distributed_run destroys it with the run. A separate call for every caller to remember is one more way to get it wrong. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 16 ++--------- .../unit/torch/utils/test_distributed_run.py | 28 ++----------------- 2 files changed, 4 insertions(+), 40 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index f01b3f85c89..f883b4bcbd7 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -55,7 +55,6 @@ "is_master", "rank", "set_distribution_plan", - "setup_fsdp2_process_group", "size", ] @@ -477,8 +476,8 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str device. A single process sets nothing up and is rank 0 of a run of one, on ``device``. On exit, even if the body raises, it destroys the group the run set up: the launch's, or one the - body set up, such as the one-rank group of :func:`setup_fsdp2_process_group`. A group that was - already there when the run began is the caller's, and is left as it is. + body set up, such as the one-rank group a single process's FSDP2 load shards through. A group + that was already there when the run began is the caller's, and is left as it is. """ owns_group = is_distributed_launch() group_was_there = is_initialized() @@ -508,17 +507,6 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str cleanup() -def setup_fsdp2_process_group(plan: DistributionPlan, timeout: timedelta | None = None) -> None: - """Set up the process group a single process needs to FSDP2-shard by ``plan``. - - FSDP2 shards through a process group even on one rank. A distributed launch already has its - group (:func:`distributed_run`), and a plan that does not FSDP2-shard needs none. Inside a - :func:`distributed_run`, the run destroys the group on exit. - """ - if plan.is_fsdp_dense and not is_initialized(): - setup(timeout=timeout) - - def cleanup(): """Cleans up the distributed environment. diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index ffbbddc3923..fa545d3bff5 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -33,7 +33,6 @@ distribution_plan, is_distributed_launch, set_distribution_plan, - setup_fsdp2_process_group, ) # --------------------------------------------------------------------------------- DistributionPlan @@ -282,17 +281,15 @@ def test_a_distributed_launch_yields_this_ranks_state_on_its_cuda_device(monkeyp def test_the_run_destroys_a_group_its_body_set_up(monkeypatch): - """FSDP2 on one rank sets a group up mid-run; the run tears it down like a launch's.""" + """An FSDP2 load on one rank sets a group up mid-run; the run tears it down like a launch's.""" monkeypatch.delenv("WORLD_SIZE", raising=False) live, destroyed = [], [] monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) - monkeypatch.setattr(dist, "setup", lambda timeout=None: live.append(timeout)) monkeypatch.setattr(dist, "cleanup", lambda: destroyed.append(0)) with distributed_run(device="cpu"): - setup_fsdp2_process_group(DistributionPlan(dense="fsdp"), timeout=timedelta(hours=2)) + live.append("the body's group") - assert live == [timedelta(hours=2)] assert len(destroyed) == 1 @@ -304,27 +301,6 @@ def test_the_run_leaves_a_group_it_found_alone(monkeypatch): pass -@pytest.mark.parametrize( - ("plan", "live_group", "setups"), - [ - (DistributionPlan(dense="fsdp"), False, 1), - (DistributionPlan(2, dense="fsdp"), True, 0), - (DistributionPlan(), False, 0), - ], - ids=["fsdp2_on_one_rank", "launched_group", "replicated"], -) -def test_fsdp2_on_one_rank_sets_up_a_group_of_its_own(monkeypatch, plan, live_group, setups): - """FSDP2 shards through a process group even on one rank; a distributed launch has its group, - and a plan that does not FSDP2-shard needs none.""" - calls = [] - monkeypatch.setattr(dist, "is_initialized", lambda: live_group) - monkeypatch.setattr(dist, "setup", lambda timeout=None: calls.append(timeout)) - - setup_fsdp2_process_group(plan, timeout=timedelta(hours=2)) - - assert calls == [timedelta(hours=2)] * setups - - def test_a_distributed_launch_owns_the_group_even_when_the_body_raises(monkeypatch, capsys): """The group is torn down on the way out, and the body's error is printed first: a rank failing alone would otherwise wait in the cleanup barrier while the rest sit in a collective.""" From bd31620fb9b085c4ce02abab9b7f31f55e79c18c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:40:36 +0000 Subject: [PATCH 13/48] Refuse to record a second, different plan on a model A model is distributed once, by the load that decides how. Recording a plan over a different one already on the model, or on any submodule, would leave the readers of the first one disagreeing with those of the second; it raises instead, and leaves the model as it was. Recording the same plan again is harmless. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 10 +++++++++ .../unit/torch/utils/test_distributed_run.py | 21 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index f883b4bcbd7..6b56ab811ec 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -348,7 +348,17 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed the whole model. ``load_pretrained`` records the plan it loads by; code that distributes a model itself records it here before calibrating or exporting the model. + + A model is distributed once: a plan other than the one already recorded on the model, or on any + of its submodules, raises instead of replacing it, and leaves the model as it was. """ + for module in model.modules(): + recorded = getattr(module, _PLAN_ATTR, None) + if recorded is not None and recorded != plan: + raise ValueError( + f"{type(module).__name__} is already recorded as distributed by {recorded!r}; a " + f"model is distributed once, so {plan!r} cannot replace it." + ) for module in model.modules(): setattr(module, _PLAN_ATTR, plan) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index fa545d3bff5..1fb280087a0 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -85,6 +85,27 @@ def test_every_submodule_reads_the_recorded_plan(): assert all(distribution_plan(module) is plan for module in model.modules()) +def test_a_recorded_plan_is_not_replaced_by_another(): + """A model is distributed once; a second, different plan would leave readers disagreeing.""" + model, plan = _model(), DistributionPlan(world_size=8, ep_size=8) + set_distribution_plan(model, plan) + set_distribution_plan(model, DistributionPlan(world_size=8, ep_size=8)) # the same plan again + + with pytest.raises(ValueError, match="a model is distributed once"): + set_distribution_plan(model, DistributionPlan(world_size=8, dense="fsdp")) + assert all(distribution_plan(module) == plan for module in model.modules()) + + +def test_a_plan_recorded_on_a_submodule_is_not_replaced_either(): + """Nothing is recorded when any submodule already carries another plan.""" + model = _model() + set_distribution_plan(model[1], DistributionPlan(world_size=8, ep_size=8)) + + with pytest.raises(ValueError, match="a model is distributed once"): + set_distribution_plan(model, DistributionPlan(world_size=8)) + assert distribution_plan(model) == DistributionPlan() + + def test_a_model_with_no_plan_is_a_single_process(): assert distribution_plan(_model()) == DistributionPlan() From 67cd9bf77e83291750daaab292c86aa44b0d259c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:52:55 +0000 Subject: [PATCH 14/48] Reject a plan degree below 1 ep_size=-2 divided a world of 4 and read as no expert parallelism, so an EP knob of -2 silently replicated an MoE it asked to shard; ep_size=0 divided by zero, and world_size=0 was accepted. Every degree must be at least 1. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 5 +++++ tests/unit/torch/utils/test_distributed_run.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 6b56ab811ec..d327917ef10 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -318,6 +318,11 @@ class DistributionPlan: def __post_init__(self) -> None: """Reject a plan no load could carry out, before any rank starts on it.""" + for name in ("world_size", "ep_size"): + # Below 1 a degree is meaningless: ep_size=-2 would divide the world and read as no + # expert parallelism, ep_size=0 would divide by zero. + if getattr(self, name) < 1: + raise ValueError(f"{name}={getattr(self, name)} must be at least 1.") if self.world_size % self.ep_size != 0: raise ValueError(f"ep_size={self.ep_size} must divide world_size={self.world_size}") if self.dense not in ("replicate", "fsdp"): diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 1fb280087a0..e9b9942c379 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -67,6 +67,22 @@ def test_an_ep_degree_that_does_not_divide_the_world_is_rejected(): DistributionPlan(world_size=8, ep_size=3) +@pytest.mark.parametrize( + "degrees", + [ + {"world_size": 0}, + {"world_size": -4}, + {"world_size": 4, "ep_size": 0}, + {"world_size": 4, "ep_size": -2}, + ], +) +def test_a_degree_below_one_is_rejected(degrees): + """ep_size=-2 divides a world of 4 and would read as no expert parallelism, replicating an MoE + an EP knob of -2 asked to shard; ep_size=0 would divide by zero.""" + with pytest.raises(ValueError, match="must be at least 1"): + DistributionPlan(**degrees) + + def test_an_unknown_dense_layout_is_rejected(): with pytest.raises(ValueError, match="expected 'replicate' or 'fsdp'"): DistributionPlan(world_size=8, dense="sharded") From 967342d5ded92581c16ed7ab67a9db4c2112f66e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 01:07:58 +0000 Subject: [PATCH 15/48] Record a model's plan once, refusing a reload as well as a redistribution Nothing reloads or re-distributes a model it loaded: load_pretrained records the plan once, on the model it just built. A second record, of the same plan or another, can only mean a second writer, so it raises like a different plan already did, and leaves the model as it was. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 10 ++++++---- .../unit/torch/utils/test_distributed_run.py | 19 ++++++++++++------- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index d327917ef10..7abafcd287f 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -354,15 +354,17 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non the whole model. ``load_pretrained`` records the plan it loads by; code that distributes a model itself records it here before calibrating or exporting the model. - A model is distributed once: a plan other than the one already recorded on the model, or on any - of its submodules, raises instead of replacing it, and leaves the model as it was. + A model is distributed once, and its plan recorded once: recording a plan on a model that, or + any of whose submodules, already records one raises, the same plan included, and leaves the + model as it was. No code reloads or re-distributes a model it loaded. """ for module in model.modules(): recorded = getattr(module, _PLAN_ATTR, None) - if recorded is not None and recorded != plan: + if recorded is not None: raise ValueError( f"{type(module).__name__} is already recorded as distributed by {recorded!r}; a " - f"model is distributed once, so {plan!r} cannot replace it." + f"model is distributed and its plan recorded once, so {plan!r} cannot be recorded " + "over it." ) for module in model.modules(): setattr(module, _PLAN_ATTR, plan) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index e9b9942c379..5b15802a335 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -101,15 +101,20 @@ def test_every_submodule_reads_the_recorded_plan(): assert all(distribution_plan(module) is plan for module in model.modules()) -def test_a_recorded_plan_is_not_replaced_by_another(): - """A model is distributed once; a second, different plan would leave readers disagreeing.""" +@pytest.mark.parametrize( + "again", + [DistributionPlan(world_size=8, ep_size=8), DistributionPlan(world_size=8, dense="fsdp")], + ids=["reload", "redistribute"], +) +def test_a_model_records_its_plan_once(again): + """A model is distributed once, by its load: recording any plan over a recorded one, the same + or another, means a second writer, and is refused.""" model, plan = _model(), DistributionPlan(world_size=8, ep_size=8) set_distribution_plan(model, plan) - set_distribution_plan(model, DistributionPlan(world_size=8, ep_size=8)) # the same plan again - with pytest.raises(ValueError, match="a model is distributed once"): - set_distribution_plan(model, DistributionPlan(world_size=8, dense="fsdp")) - assert all(distribution_plan(module) == plan for module in model.modules()) + with pytest.raises(ValueError, match="its plan recorded once"): + set_distribution_plan(model, again) + assert all(distribution_plan(module) is plan for module in model.modules()) def test_a_plan_recorded_on_a_submodule_is_not_replaced_either(): @@ -117,7 +122,7 @@ def test_a_plan_recorded_on_a_submodule_is_not_replaced_either(): model = _model() set_distribution_plan(model[1], DistributionPlan(world_size=8, ep_size=8)) - with pytest.raises(ValueError, match="a model is distributed once"): + with pytest.raises(ValueError, match="its plan recorded once"): set_distribution_plan(model, DistributionPlan(world_size=8)) assert distribution_plan(model) == DistributionPlan() From cccb71dd65f4ff7706eafde1f35c1ac01bd02f63 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:43:55 +0000 Subject: [PATCH 16/48] Leave a process group the caller set up before the run, under a launch too distributed_run took ownership of the group whenever the launch was distributed, so a caller that had set its group up before entering lost it on exit under torchrun, against the docstring. Clean up only a group that was not there when the run began. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 8 +++--- .../unit/torch/utils/test_distributed_run.py | 26 ++++++++++++++++--- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7abafcd287f..0bcc14f5d19 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -496,14 +496,14 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str body set up, such as the one-rank group a single process's FSDP2 load shards through. A group that was already there when the run began is the caller's, and is left as it is. """ - owns_group = is_distributed_launch() + launched = is_distributed_launch() group_was_there = is_initialized() - if owns_group: + if launched: setup(timeout=timeout) try: yield ( DistributedState(rank(), size(), torch.device(f"cuda:{local_rank()}")) - if owns_group + if launched else DistributedState(0, 1, device) ) finally: @@ -520,7 +520,7 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str ) traceback.print_exc() sys.stderr.flush() - if owns_group or (not group_was_there and is_initialized()): + if not group_was_there and is_initialized(): cleanup() diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 5b15802a335..51295eb59d8 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -335,8 +335,17 @@ def test_the_run_destroys_a_group_its_body_set_up(monkeypatch): assert len(destroyed) == 1 -def test_the_run_leaves_a_group_it_found_alone(monkeypatch): - monkeypatch.delenv("WORLD_SIZE", raising=False) +@pytest.mark.parametrize("world_size", [None, "2"], ids=["single-process", "launch"]) +def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size): + """Under a launch too: a caller that set its group up before the run keeps it.""" + if world_size is None: + monkeypatch.delenv("WORLD_SIZE", raising=False) + else: + monkeypatch.setenv("WORLD_SIZE", world_size) + monkeypatch.setattr(dist, "setup", lambda timeout=None: None) + monkeypatch.setattr(dist, "rank", lambda group=None: 0) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(dist, "local_rank", lambda: 0) monkeypatch.setattr(dist, "is_initialized", lambda: True) monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("destroyed the caller's group")) with distributed_run(device="cpu"): @@ -347,8 +356,17 @@ def test_a_distributed_launch_owns_the_group_even_when_the_body_raises(monkeypat """The group is torn down on the way out, and the body's error is printed first: a rank failing alone would otherwise wait in the cleanup barrier while the rest sit in a collective.""" monkeypatch.setenv("WORLD_SIZE", "2") - calls = [] - monkeypatch.setattr(dist, "setup", lambda timeout=None: calls.append(("setup", timeout))) + calls, live = [], [] + + def setup(timeout=None): + calls.append(("setup", timeout)) + live.append("the launch's group") + + monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) + monkeypatch.setattr(dist, "setup", setup) + monkeypatch.setattr(dist, "rank", lambda group=None: 0) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(dist, "local_rank", lambda: 0) monkeypatch.setattr(dist, "cleanup", lambda: calls.append(("cleanup", capsys.readouterr().out))) with ( From 05695ae4645fe851c92e6153e5b7ee4248b52cda Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 02:15:30 +0000 Subject: [PATCH 17/48] Exit a rank whose body raised in a distributed run instead of destroying its group distributed_run cleaned up after a body that raised: cleanup() skips its barrier then, but destroy_process_group still blocks while the peers sit in a collective this rank never reaches, which cleanup()'s own docstring says error paths must avoid by calling abort(). A rank of a run of several ranks that owns its group now aborts: abort() prints the traceback, which the run no longer prints first, and exits, so torchrun tears the job down. A single process and an intentional exit still clean up, and a group the caller set up is still left alone. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 19 ++-- .../unit/torch/utils/test_distributed_run.py | 100 ++++++++++-------- 2 files changed, 71 insertions(+), 48 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 0bcc14f5d19..3222becb1a3 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -494,7 +494,10 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str On exit, even if the body raises, it destroys the group the run set up: the launch's, or one the body set up, such as the one-rank group a single process's FSDP2 load shards through. A group - that was already there when the run began is the caller's, and is left as it is. + that was already there when the run began is the caller's, and is left as it is. A rank whose + body raised in a run of several ranks exits through :func:`abort` instead, and the launcher + tears the job down: destroying the group would wait on peers that may sit in a collective this + rank never reaches. """ launched = is_distributed_launch() group_was_there = is_initialized() @@ -507,20 +510,24 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str else DistributedState(0, 1, device) ) finally: - # If the body is unwinding due to an exception, PRINT it here BEFORE cleanup()'s collective - # barrier. A per-rank exception (only some ranks raise) otherwise makes those ranks enter the - # cleanup barrier while the rest are still inside a calibration a2a collective -> cross-rank - # DEADLOCK, with the original exception never surfaced (pending behind the hung barrier). + # If the body is unwinding due to an exception, PRINT it here BEFORE anything that waits on + # the peers. A per-rank exception (only some ranks raise) leaves the rest inside a + # calibration a2a collective, so this rank's cleanup barrier or destroy_process_group would + # DEADLOCK with the original exception never surfaced. A run of several ranks that owns its + # group therefore aborts the rank: abort() prints the traceback and exits. exc = sys.exc_info()[1] + owns_group = not group_was_there and is_initialized() if exc is not None: print( f"[distributed_run][rank {os.environ.get('RANK', '?')}] body raised " f"{type(exc).__name__}: {exc}", flush=True, ) + if owns_group and not isinstance(exc, SystemExit) and size() > 1: + abort() traceback.print_exc() sys.stderr.flush() - if not group_was_there and is_initialized(): + if owns_group: cleanup() diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 51295eb59d8..686e5286c86 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -17,6 +17,7 @@ """Unit tests for distributed PTQ: the distribution plan and a model's record of it, the collectives that fail together, the process group's lifetime, launch detection.""" +from contextlib import nullcontext from datetime import timedelta import pytest @@ -312,32 +313,72 @@ def test_a_single_process_is_rank_0_of_one_on_the_given_device(monkeypatch): def test_a_distributed_launch_yields_this_ranks_state_on_its_cuda_device(monkeypatch): monkeypatch.setenv("WORLD_SIZE", "4") - monkeypatch.setattr(dist, "setup", lambda timeout=None: None) + timeouts = [] + monkeypatch.setattr(dist, "setup", lambda timeout=None: timeouts.append(timeout)) monkeypatch.setattr(dist, "cleanup", lambda: None) monkeypatch.setattr(dist, "rank", lambda group=None: 2) monkeypatch.setattr(dist, "size", lambda group=None: 4) monkeypatch.setattr(dist, "local_rank", lambda: 1) - with distributed_run(device="cpu") as state: + with distributed_run(timeout=timedelta(hours=1), device="cpu") as state: assert state == DistributedState(rank=2, world_size=4, device=torch.device("cuda:1")) assert not state.is_main + assert timeouts == [timedelta(hours=1)] -def test_the_run_destroys_a_group_its_body_set_up(monkeypatch): - """An FSDP2 load on one rank sets a group up mid-run; the run tears it down like a launch's.""" - monkeypatch.delenv("WORLD_SIZE", raising=False) - live, destroyed = [], [] +class _ExitedError(Exception): + """Raised by a stand-in for abort(), which exits the process instead of returning.""" + + +@pytest.mark.parametrize( + ("launch", "body_raises", "ends_by"), + [ + (True, None, "cleanup"), + (True, ValueError, "abort"), + (True, SystemExit, "cleanup"), + (False, None, "cleanup"), + (False, ValueError, "cleanup"), + ], + ids=["launch", "launch-raises", "launch-exits", "single-process", "single-process-raises"], +) +def test_the_run_ends_the_group_it_set_up(monkeypatch, launch, body_raises, ends_by): + """The launch's group, or the one a single process's FSDP2 load sets up mid-run, is destroyed on + the way out. A rank of several whose body raised exits through abort() instead, and torchrun tears + the job down: the cleanup's barrier and destroy_process_group would wait on peers that may sit in + a collective this rank never reaches. Every rank reaches an intentional exit.""" + if launch: + monkeypatch.setenv("WORLD_SIZE", "2") + else: + monkeypatch.delenv("WORLD_SIZE", raising=False) + live, ended = [], [] + monkeypatch.setattr(dist, "setup", lambda timeout=None: live.append("the launch's group")) monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) - monkeypatch.setattr(dist, "cleanup", lambda: destroyed.append(0)) + monkeypatch.setattr(dist, "rank", lambda group=None: 0) + monkeypatch.setattr(dist, "size", lambda group=None: 2 if launch else 1) + monkeypatch.setattr(dist, "local_rank", lambda: 0) + monkeypatch.setattr(dist, "cleanup", lambda: ended.append("cleanup")) + + def abort(): + ended.append("abort") + raise _ExitedError - with distributed_run(device="cpu"): - live.append("the body's group") + monkeypatch.setattr(dist, "abort", abort) + escapes = _ExitedError if ends_by == "abort" else body_raises + expected = pytest.raises(escapes) if escapes else nullcontext() - assert len(destroyed) == 1 + with expected, distributed_run(device="cpu"): + if not launch: + live.append("the body's group") + if body_raises: + raise body_raises + assert ended == [ends_by] + +@pytest.mark.parametrize("raises", [False, True], ids=["returns", "raises"]) @pytest.mark.parametrize("world_size", [None, "2"], ids=["single-process", "launch"]) -def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size): - """Under a launch too: a caller that set its group up before the run keeps it.""" +def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size, raises): + """Under a launch too, and when the body raises: a caller that set its group up before the run + keeps it, and gets the error.""" if world_size is None: monkeypatch.delenv("WORLD_SIZE", raising=False) else: @@ -348,36 +389,11 @@ def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size): monkeypatch.setattr(dist, "local_rank", lambda: 0) monkeypatch.setattr(dist, "is_initialized", lambda: True) monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("destroyed the caller's group")) - with distributed_run(device="cpu"): - pass - - -def test_a_distributed_launch_owns_the_group_even_when_the_body_raises(monkeypatch, capsys): - """The group is torn down on the way out, and the body's error is printed first: a rank failing - alone would otherwise wait in the cleanup barrier while the rest sit in a collective.""" - monkeypatch.setenv("WORLD_SIZE", "2") - calls, live = [], [] - - def setup(timeout=None): - calls.append(("setup", timeout)) - live.append("the launch's group") - - monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) - monkeypatch.setattr(dist, "setup", setup) - monkeypatch.setattr(dist, "rank", lambda group=None: 0) - monkeypatch.setattr(dist, "size", lambda group=None: 2) - monkeypatch.setattr(dist, "local_rank", lambda: 0) - monkeypatch.setattr(dist, "cleanup", lambda: calls.append(("cleanup", capsys.readouterr().out))) - - with ( - pytest.raises(ValueError, match="calibration failed"), - distributed_run(timeout=timedelta(hours=1)), - ): - raise ValueError("calibration failed") - - assert calls[0] == ("setup", timedelta(hours=1)) - assert calls[1][0] == "cleanup" - assert "calibration failed" in calls[1][1], "the error must be printed before the cleanup" + monkeypatch.setattr(dist, "abort", lambda: pytest.fail("exited the caller's process")) + expected = pytest.raises(ValueError, match="calibration failed") if raises else nullcontext() + with expected, distributed_run(device="cpu"): + if raises: + raise ValueError("calibration failed") # ------------------------------------------------------------------------------------ launch detect From 693ef98456ef8a42d4c483db99996ce32bf19bc8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:07:34 +0000 Subject: [PATCH 18/48] Test broadcast_from, all_reduce and the plan guard on two real ranks The tests of the collectives answered torch.distributed with stand-ins and asserted the messages handed to them, so a broadcast or gather that behaved differently on a real process group went unnoticed. They now run on the session's two gloo workers: either rank sending, only the sending rank computing, a rank outside the run refused, a failure on one rank stopping every rank, and ranks out of step. So do the guards that keep an unrecorded model from being read as a single process's, with the model held whole and FSDP2-sharded for real. Only the two hang guards keep stand-ins: a rank that fails to send or to join the gather would leave a real peer waiting out its timeout, which a stand-in catches at once. bounded_collectives gives the pooled jobs a group that times out in seconds, so a rank left waiting fails the job instead of hanging the session. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/distributed/utils.py | 21 +++ .../unit/torch/utils/test_distributed_run.py | 166 ++++++++++-------- 2 files changed, 112 insertions(+), 75 deletions(-) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 45b59915351..28b72d5e9c7 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -22,6 +22,10 @@ import socket import time import traceback +from contextlib import ExitStack, contextmanager +from datetime import timedelta +from functools import partial +from unittest import mock import torch import torch.distributed as dist @@ -150,6 +154,23 @@ def reseeding_worker_teardown(rank, world_size): torch.manual_seed(1234) +@contextmanager +def bounded_collectives(seconds: float = 10): + """Run the object collectives on a new group that gives up after ``seconds``, within the block. + + For a pooled job driving code that runs ``broadcast_object_list``, ``all_gather_object`` or + ``barrier`` on the default group: a rank left waiting in one its peer never joins then raises in + seconds instead of hanging for the default group's timeout, and the default group, which the + workers keep for the next job, stays in step. Every rank of the job must enter it. + """ + group = dist.new_group(timeout=timedelta(seconds=seconds)) + with ExitStack() as stack: + for name in ("broadcast_object_list", "all_gather_object", "barrier"): + collective = partial(getattr(dist, name), group=group) + stack.enter_context(mock.patch.object(dist, name, collective)) + yield + + class DistributedWorkerPool: """Persistent worker pool that keeps distributed processes alive across multiple test dispatches. diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 686e5286c86..94ecbfe1277 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -22,7 +22,10 @@ import pytest import torch.distributed +from _test_utils.torch.distributed.utils import bounded_collectives from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard import modelopt.torch.utils.distributed as dist from modelopt.torch.utils.distributed import ( @@ -132,21 +135,38 @@ def test_a_model_with_no_plan_is_a_single_process(): assert distribution_plan(_model()) == DistributionPlan() -def test_an_fsdp2_model_with_no_plan_raises(monkeypatch): - monkeypatch.setattr(dist, "is_fsdp2_model", lambda model: True) - with pytest.raises(RuntimeError, match="is FSDP2-sharded, but has no distribution plan"): - distribution_plan(_model()) +def _outcome(call): + """``(value, None)``, or ``(None, error)`` if ``call()`` raises. + + For pooled jobs, which cannot use ``pytest.raises``: its failure is a ``BaseException``, which + ends the worker instead of failing the job. + """ + try: + return call(), None + except Exception as error: + return None, error -def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(monkeypatch): - monkeypatch.setattr(dist, "size", lambda group=None: 2) - with pytest.raises(RuntimeError, match="held in a run of 2 ranks"): - distribution_plan(_model()) - model, plan = _model(), DistributionPlan(world_size=2) +def _no_plan_on_several_ranks_job(rank, size): + _, error = _outcome(lambda: distribution_plan(_model())) + assert isinstance(error, RuntimeError), error + assert f"is held in a run of {size} ranks" in str(error), error + model = _model() + fully_shard(model, mesh=init_device_mesh("cpu", (size,))) + _, error = _outcome(lambda: distribution_plan(model)) + assert isinstance(error, RuntimeError), error + assert "is FSDP2-sharded" in str(error), error + plan = DistributionPlan(world_size=size, dense="fsdp") set_distribution_plan(model, plan) assert distribution_plan(model) is plan +def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers): + """Held whole or FSDP2-sharded, rather than be taken for a single process's; once a plan is + recorded, that plan is what it reads.""" + cpu_dist_workers.run(_no_plan_on_several_ranks_job) + + # ------------------------------------------------------------------------------------ collectives @@ -156,18 +176,13 @@ def _run_with_ranks(monkeypatch, world_size): monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size) -def _as_rank(monkeypatch, rank, received=None): - """Answer the broadcasts of a two-rank run as ``rank``; returns ``(message, sending rank)`` for each. - - A rank that is not the sending one receives ``received``, as the sending rank would have sent it. - """ +def _as_rank(monkeypatch, rank): + """Record the broadcasts of a two-rank run as ``rank``; returns ``(message, sending rank)`` for each.""" _run_with_ranks(monkeypatch, 2) sent = [] def broadcast_object_list(message, group=None, group_src=None): sent.append((message[0], group_src)) - if rank != group_src: - message[0] = received monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank) monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) @@ -203,41 +218,47 @@ def test_a_single_process_computes_alone(monkeypatch): assert all_reduce("smallest", lambda: 4, min) == 4 -def test_the_sending_rank_sends_its_value(monkeypatch): - sent = _as_rank(monkeypatch, 0) - plan = DistributionPlan(world_size=2, ep_size=2) - - assert broadcast_from("plan", lambda: plan) is plan - assert sent == [(("plan", plan, None), 0)] - - -def test_other_ranks_take_the_sending_ranks_value_without_running_fn(monkeypatch): - rank0_plan = DistributionPlan(world_size=2, ep_size=2) - _as_rank(monkeypatch, 1, received=("plan", rank0_plan, None)) - - plan = broadcast_from("plan", lambda: pytest.fail("a rank that does not send ran fn")) - - assert plan.ep_size == 2 +def _broadcast_from_job(rank, size): + ran = [] + def compute(): + ran.append(rank) + return f"rank {rank}'s" -def test_any_rank_can_broadcast(monkeypatch): - """With rank=1, rank 1 runs fn and sends it; rank 0 takes rank 1's value.""" - sent = _as_rank(monkeypatch, 1) - assert broadcast_from("x", lambda: "rank 1's", rank=1) == "rank 1's" - assert sent == [(("x", "rank 1's", None), 1)] - - _as_rank(monkeypatch, 0, received=("x", "rank 1's", None)) - assert broadcast_from("x", lambda: pytest.fail("rank 0 ran fn"), rank=1) == "rank 1's" + def fail(): + raise ValueError("bad degrees") - -def test_a_rank_outside_the_run_cannot_broadcast(monkeypatch): - _run_with_ranks(monkeypatch, 2) - with pytest.raises(ValueError, match="rank=2 is not a rank of this 2-rank run"): - broadcast_from("x", lambda: 0, rank=2) + with bounded_collectives(): + for sender in range(size): + ran.clear() + assert broadcast_from("x", compute, rank=sender) == f"rank {sender}'s" + assert ran == ([rank] if rank == sender else []), "a rank that does not send ran fn" + _, error = _outcome(lambda: broadcast_from("x", compute, rank=size)) + assert isinstance(error, ValueError), error + assert f"rank={size} is not a rank of this {size}-rank run" in str(error), error + _, error = _outcome(lambda: broadcast_from("x", fail)) + if rank == 0: + assert isinstance(error, ValueError), error + else: + assert isinstance(error, RuntimeError), error + assert "rank 0 failed to compute x" in str(error), error + sent, error = _outcome(lambda: broadcast_from("plan" if rank == 0 else "other", lambda: 7)) + if rank == 0: + assert (sent, error) == (7, None) + else: + assert isinstance(error, RuntimeError), error + assert "out of step" in str(error), error + + +def test_broadcast_from_gives_every_rank_the_sending_ranks_value(cpu_dist_workers): + """Over a real broadcast: any rank can send, and only it runs fn. A rank outside the run is + refused, the sending rank's failure stops every rank, and ranks out of step fail.""" + cpu_dist_workers.run(_broadcast_from_job) def test_the_sending_rank_failing_still_reaches_the_broadcast(monkeypatch): - """The sending rank must still send on failure, or every other rank hangs in the broadcast.""" + """The sending rank must still send on failure, or every other rank hangs in the broadcast. A + stand-in broadcast catches that at once, where a real one would wait out its timeout.""" sent = _as_rank(monkeypatch, 0) def fail(): @@ -248,29 +269,36 @@ def fail(): assert sent == [(("x", None, "ValueError: bad degrees"), 0)] -def test_other_ranks_raise_when_the_sending_rank_failed(monkeypatch): - _as_rank(monkeypatch, 1, received=("x", None, "ValueError: bad degrees")) - with pytest.raises(RuntimeError, match=r"rank 0 failed to compute x.*bad degrees"): - broadcast_from("x", lambda: 0) - - -def test_ranks_broadcasting_out_of_step_fail(monkeypatch): - """Rank 0 sending one value while rank 1 expects another would cross the two.""" - _as_rank(monkeypatch, 1, received=("other", 7, None)) - with pytest.raises(RuntimeError, match="out of step"): - broadcast_from("plan", lambda: DistributionPlan(world_size=2)) - +def _all_reduce_job(rank, size): + def fail_on_rank_1(): + if rank == 1: + raise ValueError("no memory reading") + return 4 + + with bounded_collectives(): + assert all_reduce("smallest", lambda: [8, 2][rank], min) == 2 + _, error = _outcome(lambda: all_reduce("smallest", fail_on_rank_1, min)) + if rank == 1: + assert isinstance(error, ValueError), error + else: + assert isinstance(error, RuntimeError), error + assert "rank 1 failed to compute smallest" in str(error), error + _, error = _outcome( + lambda: all_reduce("smallest" if rank == 0 else "other", lambda: 4, min) + ) + assert isinstance(error, RuntimeError), error + assert "out of step" in str(error), error -def test_all_reduce_gives_every_rank_the_combined_value(monkeypatch): - """Every rank's value counts, and every rank gets the same result.""" - sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None), ("smallest", 2, None)]) - assert all_reduce("smallest", lambda: 4, min) == 2 - assert sent == [("smallest", 4, None)] +def test_all_reduce_gives_every_rank_the_combined_value(cpu_dist_workers): + """Over a real gather: every rank's value counts and every rank gets the same result. One rank's + failure stops every rank, and ranks out of step fail.""" + cpu_dist_workers.run(_all_reduce_job) def test_a_rank_failing_to_compute_still_joins_the_gather(monkeypatch): - """A failing rank must still reach the gather, or every other rank waits in it.""" + """A failing rank must still reach the gather, or every other rank waits in it. A stand-in + gather catches that at once, where a real one would wait out its timeout.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) def fail(): @@ -281,18 +309,6 @@ def fail(): assert sent == [("smallest", None, "ValueError: no memory reading")] -def test_the_other_ranks_raise_when_one_failed_to_compute(monkeypatch): - _gather_as_rank(monkeypatch, 0, others=[("smallest", None, "ValueError: no memory reading")]) - with pytest.raises(RuntimeError, match=r"rank 1 failed to compute smallest.*memory"): - all_reduce("smallest", lambda: 4, min) - - -def test_ranks_reducing_out_of_step_fail(monkeypatch): - _gather_as_rank(monkeypatch, 0, others=[("other", 8, None)]) - with pytest.raises(RuntimeError, match="out of step"): - all_reduce("smallest", lambda: 4, min) - - # ---------------------------------------------------------------------------------- process group From ad2999ed3f263dda6ea251d946772bf25076fefe Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:49:49 +0000 Subject: [PATCH 19/48] Trim the distribution plan's and distributed_run's tests Fold near-duplicates into one test per behavior, and drop one test other tests already cover: - test_a_plan_splits_the_world takes the plain-world, expert-parallel and FSDP2-dense property tests. - test_an_invalid_plan_is_rejected takes the EP-degree-not-dividing, degree-below-one and unknown-dense refusals. - test_a_model_records_its_plan_once takes the submodule case. - test_a_single_process_is_rank_0_of_one_on_the_given_device also checks that no group is set up or torn down. - test_launch_detection_reads_the_env_before_any_process_group goes: the distributed_run tests fail if WORLD_SIZE=1 counts as a launch or if the launch is read from the live process group. Signed-off-by: Shengliang Xu --- .../unit/torch/utils/test_distributed_run.py | 122 +++++++----------- 1 file changed, 47 insertions(+), 75 deletions(-) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 94ecbfe1277..7f3cb8596c6 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -15,7 +15,7 @@ """Unit tests for distributed PTQ: the distribution plan and a model's record of it, the collectives -that fail together, the process group's lifetime, launch detection.""" +that fail together, the process group's lifetime.""" from contextlib import nullcontext from datetime import timedelta @@ -35,63 +35,52 @@ broadcast_from, distributed_run, distribution_plan, - is_distributed_launch, set_distribution_plan, ) # --------------------------------------------------------------------------------- DistributionPlan -def test_plain_world_replicates_everything(): - plan = DistributionPlan(world_size=8) - assert plan.edp == 8 - assert not plan.is_expert_parallel - assert not plan.is_fsdp_dense - - -def test_expert_parallel_splits_the_expert_replica_pool(): - plan = DistributionPlan(world_size=8, ep_size=2) - assert plan.edp == 4 # each expert shard is replicated 4x - assert plan.is_expert_parallel - - -def test_fsdp_dense_plan(): - plan = DistributionPlan(world_size=8, ep_size=8, dense="fsdp") - assert plan.is_fsdp_dense - assert plan.edp == 1 +@pytest.mark.parametrize( + ("plan", "edp", "expert_parallel", "fsdp_dense"), + [ + (DistributionPlan(world_size=8), 8, False, False), + # each expert shard is replicated 4x + (DistributionPlan(world_size=8, ep_size=2), 4, True, False), + (DistributionPlan(world_size=8, ep_size=8, dense="fsdp"), 1, True, True), + ], + ids=["plain-world", "expert-parallel", "fsdp-dense"], +) +def test_a_plan_splits_the_world(plan, edp, expert_parallel, fsdp_dense): + assert plan.edp == edp + assert plan.is_expert_parallel is expert_parallel + assert plan.is_fsdp_dense is fsdp_dense def test_the_default_plan_is_a_single_process(): assert DistributionPlan() == DistributionPlan(world_size=1, ep_size=1, dense="replicate") -def test_an_ep_degree_that_does_not_divide_the_world_is_rejected(): - """An unsatisfiable plan must fail here, not as a confusing mesh error thousands of ranks later.""" - with pytest.raises(ValueError, match="must divide world_size"): - DistributionPlan(world_size=8, ep_size=3) - - @pytest.mark.parametrize( - "degrees", + ("degrees", "match"), [ - {"world_size": 0}, - {"world_size": -4}, - {"world_size": 4, "ep_size": 0}, - {"world_size": 4, "ep_size": -2}, + ({"world_size": 8, "ep_size": 3}, "must divide world_size"), + ({"world_size": 0}, "must be at least 1"), + ({"world_size": -4}, "must be at least 1"), + ({"world_size": 4, "ep_size": 0}, "must be at least 1"), + ({"world_size": 4, "ep_size": -2}, "must be at least 1"), + ({"world_size": 8, "dense": "sharded"}, "expected 'replicate' or 'fsdp'"), ], + ids=["ep-not-dividing", "world-0", "world-negative", "ep-0", "ep-negative", "unknown-dense"], ) -def test_a_degree_below_one_is_rejected(degrees): - """ep_size=-2 divides a world of 4 and would read as no expert parallelism, replicating an MoE - an EP knob of -2 asked to shard; ep_size=0 would divide by zero.""" - with pytest.raises(ValueError, match="must be at least 1"): +def test_an_invalid_plan_is_rejected(degrees, match): + """An unsatisfiable plan must fail here, not as a confusing mesh error thousands of ranks later. + ep_size=-2 divides a world of 4 and would read as no expert parallelism, replicating an MoE an + EP knob of -2 asked to shard; ep_size=0 would divide by zero.""" + with pytest.raises(ValueError, match=match): DistributionPlan(**degrees) -def test_an_unknown_dense_layout_is_rejected(): - with pytest.raises(ValueError, match="expected 'replicate' or 'fsdp'"): - DistributionPlan(world_size=8, dense="sharded") - - # ------------------------------------------------------------------- a model's recorded plan @@ -106,29 +95,29 @@ def test_every_submodule_reads_the_recorded_plan(): @pytest.mark.parametrize( - "again", - [DistributionPlan(world_size=8, ep_size=8), DistributionPlan(world_size=8, dense="fsdp")], - ids=["reload", "redistribute"], + ("recorded_on", "again"), + [ + ("model", DistributionPlan(world_size=8, ep_size=8)), + ("model", DistributionPlan(world_size=8, dense="fsdp")), + ("submodule", DistributionPlan(world_size=8)), + ], + ids=["reload", "redistribute", "submodule"], ) -def test_a_model_records_its_plan_once(again): +def test_a_model_records_its_plan_once(recorded_on, again): """A model is distributed once, by its load: recording any plan over a recorded one, the same - or another, means a second writer, and is refused.""" + or another, means a second writer, and is refused, also when only a submodule carries one. A + refusal records nothing.""" model, plan = _model(), DistributionPlan(world_size=8, ep_size=8) - set_distribution_plan(model, plan) + set_distribution_plan(model if recorded_on == "model" else model[1], plan) + holders = [module for module in model.modules() if distribution_plan(module) is plan] with pytest.raises(ValueError, match="its plan recorded once"): set_distribution_plan(model, again) - assert all(distribution_plan(module) is plan for module in model.modules()) - - -def test_a_plan_recorded_on_a_submodule_is_not_replaced_either(): - """Nothing is recorded when any submodule already carries another plan.""" - model = _model() - set_distribution_plan(model[1], DistributionPlan(world_size=8, ep_size=8)) - - with pytest.raises(ValueError, match="its plan recorded once"): - set_distribution_plan(model, DistributionPlan(world_size=8)) - assert distribution_plan(model) == DistributionPlan() + for module in model.modules(): + if module in holders: + assert distribution_plan(module) is plan + else: + assert distribution_plan(module) == DistributionPlan() def test_a_model_with_no_plan_is_a_single_process(): @@ -312,16 +301,11 @@ def fail(): # ---------------------------------------------------------------------------------- process group -def test_a_single_process_sets_no_group_up(monkeypatch): +def test_a_single_process_is_rank_0_of_one_on_the_given_device(monkeypatch): + """With no ranks to sync with, the run sets no group up and tears none down.""" monkeypatch.delenv("WORLD_SIZE", raising=False) monkeypatch.setattr(dist, "setup", lambda **kwargs: pytest.fail("set a group up")) monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("tore a group down")) - with distributed_run(): - pass - - -def test_a_single_process_is_rank_0_of_one_on_the_given_device(monkeypatch): - monkeypatch.delenv("WORLD_SIZE", raising=False) with distributed_run(device="cpu") as state: assert state == DistributedState(rank=0, world_size=1, device="cpu") assert state.is_main @@ -410,15 +394,3 @@ def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size, raises): with expected, distributed_run(device="cpu"): if raises: raise ValueError("calibration failed") - - -# ------------------------------------------------------------------------------------ launch detect - - -@pytest.mark.parametrize(("world", "expected"), [(None, False), ("1", False), ("8", True)]) -def test_launch_detection_reads_the_env_before_any_process_group(monkeypatch, world, expected): - """Must work at parse time, before distributed_run sets the process group up.""" - monkeypatch.delenv("WORLD_SIZE", raising=False) - if world is not None: - monkeypatch.setenv("WORLD_SIZE", world) - assert is_distributed_launch() is expected From 7a92514005a3e5906374175c88e5f43a13b87745 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:52:29 +0000 Subject: [PATCH 20/48] Destroy the bounded-timeout group when its block exits bounded_collectives creates a gloo group per use and never destroyed it. The cpu_dist_workers pool keeps its workers for the whole session, so each job left a group and its sockets behind: 61 groups and 330 open files per rank after 60 jobs. Destroy the group in a finally. A peer still waiting in it then fails at once instead of at the timeout; the pool's default group is untouched. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/distributed/utils.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 28b72d5e9c7..831b1661ed9 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -164,11 +164,16 @@ def bounded_collectives(seconds: float = 10): workers keep for the next job, stays in step. Every rank of the job must enter it. """ group = dist.new_group(timeout=timedelta(seconds=seconds)) - with ExitStack() as stack: - for name in ("broadcast_object_list", "all_gather_object", "barrier"): - collective = partial(getattr(dist, name), group=group) - stack.enter_context(mock.patch.object(dist, name, collective)) - yield + try: + with ExitStack() as stack: + for name in ("broadcast_object_list", "all_gather_object", "barrier"): + collective = partial(getattr(dist, name), group=group) + stack.enter_context(mock.patch.object(dist, name, collective)) + yield + finally: + # The workers outlive the job, and would keep the group and its connections. Destroying it + # also fails a peer still waiting in it at once, instead of at the timeout. + dist.destroy_process_group(group) class DistributedWorkerPool: From 629afebaa76e250fe8d0cb03d4d1f3751d447db6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 19:05:25 +0000 Subject: [PATCH 21/48] Date test_distributed_run 2026, as LICENSE_HEADER does Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_distributed_run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 7f3cb8596c6..485bb4fa47a 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); From 8fce8c9e75cc8aab8c77982b6ad72cd88ff67b8b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 19:57:44 +0000 Subject: [PATCH 22/48] Refuse a non-CUDA device under a launch, and a rank outside the run distributed_run documented that a launch ignores its device argument, so `torchrun hf_ptq.py --device cpu` ran on the GPUs without a word; a launch now refuses a non-CUDA device before setting anything up. broadcast_from checked its sending rank only in a run of several ranks, so a single process accepted rank=1; it now checks first. distributed_run gains its return annotation, and set_distribution_plan's docstring says it records on the submodules the model has when it is called. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 23 +++++++++++++------ .../unit/torch/utils/test_distributed_run.py | 17 +++++++++++--- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 3222becb1a3..434bfe9812a 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -348,11 +348,12 @@ def is_fsdp_dense(self) -> bool: def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> None: - """Record that ``model`` is distributed as ``plan``, on it and on every submodule. + """Record that ``model`` is distributed as ``plan``, on it and on every submodule it has now. Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed - the whole model. ``load_pretrained`` records the plan it loads by; code that distributes a model - itself records it here before calibrating or exporting the model. + the whole model; a module added afterwards records none. ``load_pretrained`` records the plan it + loads by once the model is built; code that distributes a model itself records it here before + calibrating or exporting the model. A model is distributed once, and its plan recorded once: recording a plan on a model that, or any of whose submodules, already records one raises, the same plan included, and leaves the @@ -401,10 +402,10 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: different labels at the same point fail. Collective: every rank calls it at the same point, with the same ``rank``. Outside a distributed run ``fn`` runs alone. """ - if not is_distributed_run(): - return fn() if not 0 <= rank < size(): raise ValueError(f"rank={rank} is not a rank of this {size()}-rank run.") + if not is_distributed_run(): + return fn() sending = torch.distributed.get_rank() == rank value: Any = None error = None @@ -485,12 +486,15 @@ def is_main(self) -> bool: @contextmanager -def distributed_run(timeout: timedelta | None = None, device: torch.device | str = "cuda"): +def distributed_run( + timeout: timedelta | None = None, device: torch.device | str = "cuda" +) -> Iterator[DistributedState]: """Own the run's process group for the block, and yield this process's state. On a distributed launch (``WORLD_SIZE>1``) this sets the process group up, binding this rank's CUDA device, with ``timeout``. It yields this rank's :class:`DistributedState`, on that CUDA - device. A single process sets nothing up and is rank 0 of a run of one, on ``device``. + device, and refuses any other ``device``. A single process sets nothing up and is rank 0 of a run + of one, on ``device``. On exit, even if the body raises, it destroys the group the run set up: the launch's, or one the body set up, such as the one-rank group a single process's FSDP2 load shards through. A group @@ -500,6 +504,11 @@ def distributed_run(timeout: timedelta | None = None, device: torch.device | str rank never reaches. """ launched = is_distributed_launch() + if launched and torch.device(device).type != "cuda": + raise ValueError( + f"distributed_run(device={device!r}): a distributed launch runs each rank on its own " + "CUDA device; run a single process for another device." + ) group_was_there = is_initialized() if launched: setup(timeout=timeout) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 485bb4fa47a..cec0ca897e7 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -205,6 +205,8 @@ def test_a_single_process_computes_alone(monkeypatch): ) assert broadcast_from("x", lambda: 3) == 3 assert all_reduce("smallest", lambda: 4, min) == 4 + with pytest.raises(ValueError, match="rank=1"): + broadcast_from("x", lambda: 3, rank=1) def _broadcast_from_job(rank, size): @@ -319,12 +321,21 @@ def test_a_distributed_launch_yields_this_ranks_state_on_its_cuda_device(monkeyp monkeypatch.setattr(dist, "rank", lambda group=None: 2) monkeypatch.setattr(dist, "size", lambda group=None: 4) monkeypatch.setattr(dist, "local_rank", lambda: 1) - with distributed_run(timeout=timedelta(hours=1), device="cpu") as state: + with distributed_run(timeout=timedelta(hours=1)) as state: assert state == DistributedState(rank=2, world_size=4, device=torch.device("cuda:1")) assert not state.is_main assert timeouts == [timedelta(hours=1)] +def test_a_distributed_launch_refuses_another_device(monkeypatch): + """Every rank runs on its own CUDA device, so another device is refused, before any setup, + rather than ignored.""" + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setattr(dist, "setup", lambda **kwargs: pytest.fail("set a group up")) + with pytest.raises(ValueError, match="device='cpu'"), distributed_run(device="cpu"): + pass + + class _ExitedError(Exception): """Raised by a stand-in for abort(), which exits the process instead of returning.""" @@ -365,7 +376,7 @@ def abort(): escapes = _ExitedError if ends_by == "abort" else body_raises expected = pytest.raises(escapes) if escapes else nullcontext() - with expected, distributed_run(device="cpu"): + with expected, distributed_run(device="cuda" if launch else "cpu"): if not launch: live.append("the body's group") if body_raises: @@ -391,6 +402,6 @@ def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size, raises): monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("destroyed the caller's group")) monkeypatch.setattr(dist, "abort", lambda: pytest.fail("exited the caller's process")) expected = pytest.raises(ValueError, match="calibration failed") if raises else nullcontext() - with expected, distributed_run(device="cpu"): + with expected, distributed_run(device="cpu" if world_size is None else "cuda"): if raises: raise ValueError("calibration failed") From d90a8f8f62f2c6a340c0f9eaf87029c943dc49a0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:04:38 +0000 Subject: [PATCH 23/48] Join the collective on a SystemExit too, and bound a barrier named group=None broadcast_from and all_reduce caught Exception only, so a SystemExit or KeyboardInterrupt from fn, which argparse-style checks raise, skipped the collective and left the other ranks waiting in it. They now catch BaseException, join the collective, and re-raise the original. bounded_collectives bound its collectives with a partial's group keyword, which a caller's own group=None overrode: modelopt's barrier() passes it, so that barrier escaped the bound. The group is now forced over the caller's. distributed_run refuses a CUDA device other than the one LOCAL_RANK names under a launch, sets up no group when the caller brought one (setup would rebind the caller's device), and no longer reports a SystemExit, which every rank reaches, as an error. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 27 ++++---- tests/_test_utils/torch/distributed/utils.py | 8 ++- .../unit/torch/utils/test_distributed_run.py | 61 +++++++++++++------ 3 files changed, 65 insertions(+), 31 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 434bfe9812a..a189d458cec 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -412,8 +412,9 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: if sending: try: value = fn() - except Exception as e: - # Still send: the other ranks would otherwise wait in the broadcast for this one. + except BaseException as e: + # Still send: the other ranks would otherwise wait in the broadcast for this one. A + # SystemExit or KeyboardInterrupt too, which argparse-style checks raise. error = e # One message from the sending rank: its label, so a rank broadcasting another learns the two # are out of step, and its value or its error. The error goes as text: the other ranks only @@ -447,8 +448,9 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ value, error = None, None try: value = fn() - except Exception as e: - # Still join the gather: the other ranks would otherwise wait in it for this one. + except BaseException as e: + # Still join the gather: the other ranks would otherwise wait in it for this one, a + # SystemExit or KeyboardInterrupt included. error = e mine = (label, value, None if error is None else f"{type(error).__name__}: {error}") gathered: list[Any] = [None] * torch.distributed.get_world_size() @@ -492,9 +494,9 @@ def distributed_run( """Own the run's process group for the block, and yield this process's state. On a distributed launch (``WORLD_SIZE>1``) this sets the process group up, binding this rank's - CUDA device, with ``timeout``. It yields this rank's :class:`DistributedState`, on that CUDA - device, and refuses any other ``device``. A single process sets nothing up and is rank 0 of a run - of one, on ``device``. + CUDA device, with ``timeout``, unless the caller set one up already. It yields this rank's + :class:`DistributedState`, on that CUDA device, and refuses any other ``device``. A single + process sets nothing up and is rank 0 of a run of one, on ``device``. On exit, even if the body raises, it destroys the group the run set up: the launch's, or one the body set up, such as the one-rank group a single process's FSDP2 load shards through. A group @@ -504,13 +506,13 @@ def distributed_run( rank never reaches. """ launched = is_distributed_launch() - if launched and torch.device(device).type != "cuda": + if launched and torch.device(device) != torch.device("cuda"): raise ValueError( f"distributed_run(device={device!r}): a distributed launch runs each rank on its own " - "CUDA device; run a single process for another device." + "CUDA device, chosen from LOCAL_RANK; pass device='cuda' or run a single process." ) group_was_there = is_initialized() - if launched: + if launched and not group_was_there: setup(timeout=timeout) try: yield ( @@ -526,13 +528,14 @@ def distributed_run( # group therefore aborts the rank: abort() prints the traceback and exits. exc = sys.exc_info()[1] owns_group = not group_was_there and is_initialized() - if exc is not None: + # Every rank reaches a SystemExit, an intentional exit: it is not an error to report. + if exc is not None and not isinstance(exc, SystemExit): print( f"[distributed_run][rank {os.environ.get('RANK', '?')}] body raised " f"{type(exc).__name__}: {exc}", flush=True, ) - if owns_group and not isinstance(exc, SystemExit) and size() > 1: + if owns_group and size() > 1: abort() traceback.print_exc() sys.stderr.flush() diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 831b1661ed9..4333f1d0c8c 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -154,6 +154,12 @@ def reseeding_worker_teardown(rank, world_size): torch.manual_seed(1234) +def _on_group(collective, bound, /, *args, **kwargs): + """``collective`` on ``bound``, whatever group the caller named: modelopt's ``barrier()`` passes + ``group=None`` explicitly, which would override a ``partial``'s keyword.""" + return collective(*args, **{**kwargs, "group": bound}) + + @contextmanager def bounded_collectives(seconds: float = 10): """Run the object collectives on a new group that gives up after ``seconds``, within the block. @@ -167,7 +173,7 @@ def bounded_collectives(seconds: float = 10): try: with ExitStack() as stack: for name in ("broadcast_object_list", "all_gather_object", "barrier"): - collective = partial(getattr(dist, name), group=group) + collective = partial(_on_group, getattr(dist, name), group) stack.enter_context(mock.patch.object(dist, name, collective)) yield finally: diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index cec0ca897e7..9c99e127164 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -19,6 +19,7 @@ from contextlib import nullcontext from datetime import timedelta +from unittest import mock import pytest import torch.distributed @@ -247,17 +248,19 @@ def test_broadcast_from_gives_every_rank_the_sending_ranks_value(cpu_dist_worker cpu_dist_workers.run(_broadcast_from_job) -def test_the_sending_rank_failing_still_reaches_the_broadcast(monkeypatch): - """The sending rank must still send on failure, or every other rank hangs in the broadcast. A - stand-in broadcast catches that at once, where a real one would wait out its timeout.""" +@pytest.mark.parametrize("error", [ValueError, SystemExit]) +def test_the_sending_rank_failing_still_reaches_the_broadcast(monkeypatch, error): + """The sending rank must still send on failure, or every other rank hangs in the broadcast; a + SystemExit from an argparse-style check too. A stand-in broadcast catches that at once, where a + real one would wait out its timeout.""" sent = _as_rank(monkeypatch, 0) def fail(): - raise ValueError("bad degrees") + raise error("bad degrees") - with pytest.raises(ValueError, match="bad degrees"): + with pytest.raises(error, match="bad degrees"): broadcast_from("x", fail) - assert sent == [(("x", None, "ValueError: bad degrees"), 0)] + assert sent == [(("x", None, f"{error.__name__}: bad degrees"), 0)] def _all_reduce_job(rank, size): @@ -287,17 +290,36 @@ def test_all_reduce_gives_every_rank_the_combined_value(cpu_dist_workers): cpu_dist_workers.run(_all_reduce_job) -def test_a_rank_failing_to_compute_still_joins_the_gather(monkeypatch): - """A failing rank must still reach the gather, or every other rank waits in it. A stand-in - gather catches that at once, where a real one would wait out its timeout.""" +@pytest.mark.parametrize("error", [ValueError, SystemExit]) +def test_a_rank_failing_to_compute_still_joins_the_gather(monkeypatch, error): + """A failing rank must still reach the gather, or every other rank waits in it; a SystemExit + too. A stand-in gather catches that at once, where a real one would wait out its timeout.""" sent = _gather_as_rank(monkeypatch, 1, others=[("smallest", 8, None)]) def fail(): - raise ValueError("no memory reading") + raise error("no memory reading") - with pytest.raises(ValueError, match="no memory reading"): + with pytest.raises(error, match="no memory reading"): all_reduce("smallest", fail, min) - assert sent == [("smallest", None, "ValueError: no memory reading")] + assert sent == [("smallest", None, f"{error.__name__}: no memory reading")] + + +def _bounded_barrier_job(rank, size): + groups = [] + with ( + mock.patch.object( + torch.distributed, "barrier", lambda *a, group=None, **k: groups.append(group) + ), + bounded_collectives(), + ): + torch.distributed.barrier(group=None) # as modelopt's barrier() calls it + assert groups and groups[0] is not None, groups + + +def test_bounded_collectives_bound_a_barrier_called_with_group_none(cpu_dist_workers): + """modelopt's ``barrier()`` names ``group=None`` explicitly; the bound holds for it too, or a + rank left waiting in it would wait out the default group's timeout.""" + cpu_dist_workers.run(_bounded_barrier_job) # ---------------------------------------------------------------------------------- process group @@ -327,12 +349,13 @@ def test_a_distributed_launch_yields_this_ranks_state_on_its_cuda_device(monkeyp assert timeouts == [timedelta(hours=1)] -def test_a_distributed_launch_refuses_another_device(monkeypatch): - """Every rank runs on its own CUDA device, so another device is refused, before any setup, - rather than ignored.""" +@pytest.mark.parametrize("device", ["cpu", "cuda:3"]) +def test_a_distributed_launch_refuses_another_device(monkeypatch, device): + """Every rank runs on the CUDA device its LOCAL_RANK names, so another device, or another index, + is refused, before any setup, rather than ignored.""" monkeypatch.setenv("WORLD_SIZE", "2") monkeypatch.setattr(dist, "setup", lambda **kwargs: pytest.fail("set a group up")) - with pytest.raises(ValueError, match="device='cpu'"), distributed_run(device="cpu"): + with pytest.raises(ValueError, match=f"device='{device}'"), distributed_run(device=device): pass @@ -351,7 +374,7 @@ class _ExitedError(Exception): ], ids=["launch", "launch-raises", "launch-exits", "single-process", "single-process-raises"], ) -def test_the_run_ends_the_group_it_set_up(monkeypatch, launch, body_raises, ends_by): +def test_the_run_ends_the_group_it_set_up(monkeypatch, capsys, launch, body_raises, ends_by): """The launch's group, or the one a single process's FSDP2 load sets up mid-run, is destroyed on the way out. A rank of several whose body raised exits through abort() instead, and torchrun tears the job down: the cleanup's barrier and destroy_process_group would wait on peers that may sit in @@ -383,6 +406,8 @@ def abort(): raise body_raises assert ended == [ends_by] + # An intentional exit is not reported as an error. + assert ("body raised" in capsys.readouterr().out) == (body_raises is ValueError) @pytest.mark.parametrize("raises", [False, True], ids=["returns", "raises"]) @@ -394,7 +419,7 @@ def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size, raises): monkeypatch.delenv("WORLD_SIZE", raising=False) else: monkeypatch.setenv("WORLD_SIZE", world_size) - monkeypatch.setattr(dist, "setup", lambda timeout=None: None) + monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set it up again")) monkeypatch.setattr(dist, "rank", lambda group=None: 0) monkeypatch.setattr(dist, "size", lambda group=None: 2) monkeypatch.setattr(dist, "local_rank", lambda: 0) From 0acb6bcf3fb229f241598fb707ea94d2bee98c7f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:33:54 +0000 Subject: [PATCH 24/48] Say that a module added after a plan is recorded must read it off its model set_distribution_plan records the plan on the model and the submodules it has then; a module added afterwards has none, and a module holds no reference to the model it was added to, so a lookup cannot climb to it. distribution_plan now says so, and its refusal tells the reader to take the plan from the model. A pooled test pins that such a module raises in a run of several ranks rather than read as a single process's. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 8 ++++++-- tests/unit/torch/utils/test_distributed_run.py | 8 +++++++- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index a189d458cec..f6970334b40 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -375,7 +375,9 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: """How ``model`` is distributed: the plan recorded on it (:func:`set_distribution_plan`). A model with no plan recorded is a single process's, unless it shows otherwise: one that is - FSDP2-sharded, or held in a run of several ranks, raises rather than be taken for one. + FSDP2-sharded, or held in a run of several ranks, raises rather than be taken for one. That + includes a module added to a model after its plan was recorded: a module does not know the + model it was added to, so read the plan off that model. """ plan = getattr(model, _PLAN_ATTR, None) if plan is not None: @@ -385,7 +387,9 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: how = "is FSDP2-sharded" if sharded else f"is held in a run of {size()} ranks" raise RuntimeError( f"{type(model).__name__} {how}, but has no distribution plan recorded. Load it with " - "load_pretrained, or record how it is distributed with set_distribution_plan()." + "load_pretrained, or record how it is distributed with set_distribution_plan(); a " + "module added to a model after its plan was recorded has none, so read the plan off " + "that model." ) return DistributionPlan() diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 9c99e127164..064b8c44acc 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -149,11 +149,17 @@ def _no_plan_on_several_ranks_job(rank, size): plan = DistributionPlan(world_size=size, dense="fsdp") set_distribution_plan(model, plan) assert distribution_plan(model) is plan + model = _model() + set_distribution_plan(model, DistributionPlan(world_size=size)) + model.add_module("added", nn.Linear(2, 2)) + _, error = _outcome(lambda: distribution_plan(model.added)) + assert isinstance(error, RuntimeError), error def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers): """Held whole or FSDP2-sharded, rather than be taken for a single process's; once a plan is - recorded, that plan is what it reads.""" + recorded, that plan is what it reads. A module added after the plan was recorded has none, and + raises too, rather than read as a single process's.""" cpu_dist_workers.run(_no_plan_on_several_ranks_job) From 07c0fdb26261275c9d76934538a3dbcb9a75754c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:13:29 +0000 Subject: [PATCH 25/48] Abort a run only on its body's own error, not on one its caller handles distributed_run read sys.exc_info() in its finally, which also reports an exception a caller up the stack is handling. A run entered from an except block took its clean exit for a failure: it printed "body raised" with the caller's exception and, on a launch, aborted a healthy rank. The run now records the exception its body raised and acts on that alone. bounded_collectives bounds the collectives on the default group only, leaving a group a collective names to that group. broadcast_from and all_reduce say that fn's value must pickle, as the collective pickles it before it communicates, and broadcast_from that a label mismatch fails the receiving ranks only. set_distribution_plan says what the plan describes, next to the ParallelState a quantized module holds. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 23 ++++++++---- tests/_test_utils/torch/distributed/utils.py | 7 ++-- .../unit/torch/utils/test_distributed_run.py | 35 +++++++++++++------ 3 files changed, 45 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index f6970334b40..165d6d67c4f 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -353,7 +353,8 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed the whole model; a module added afterwards records none. ``load_pretrained`` records the plan it loads by once the model is built; code that distributes a model itself records it here before - calibrating or exporting the model. + calibrating or exporting the model. The plan says how the model is placed over the run; the + :class:`ParallelState` a quantized module holds names the process groups its collectives run on. A model is distributed once, and its plan recorded once: recording a plan on a model that, or any of whose submodules, already records one raises, the same plan included, and leaves the @@ -402,9 +403,11 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: For a value that depends on what each rank reads for itself -- the environment, its GPU, files on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, every - rank raises, instead of waiting in a broadcast ``rank`` never sends. Ranks broadcasting - different labels at the same point fail. Collective: every rank calls it at the same point, - with the same ``rank``. Outside a distributed run ``fn`` runs alone. + rank raises, instead of waiting in a broadcast ``rank`` never sends; its value must pickle, + which the broadcast does before it sends anything. A rank broadcasting another label than + ``rank`` at the same point fails; ``rank`` sees no other rank's label, and returns. Collective: + every rank calls it at the same point, with the same ``rank``. Outside a distributed run ``fn`` + runs alone. """ if not 0 <= rank < size(): raise ValueError(f"rank={rank} is not a rank of this {size()}-rank run.") @@ -444,8 +447,9 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ ``combine`` takes the ranks' values in rank order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's memory bounds, where one rank's choice could run another out of memory. If ``fn`` raises on any rank, every rank raises, instead of waiting in the gather - that rank never joins. Ranks reducing different labels at the same point fail. Collective; - outside a distributed run ``fn`` runs alone. + that rank never joins; its value must pickle, which the gather does before it gathers anything. + Ranks reducing different labels at the same point fail. Collective; outside a distributed run + ``fn`` runs alone. """ if not is_distributed_run(): return combine([fn()]) @@ -518,19 +522,24 @@ def distributed_run( group_was_there = is_initialized() if launched and not group_was_there: setup(timeout=timeout) + # The body's own exception: sys.exc_info() would also report one a caller up the stack is + # handling, which would make a clean exit from a run entered in an except block abort the rank. + exc: BaseException | None = None try: yield ( DistributedState(rank(), size(), torch.device(f"cuda:{local_rank()}")) if launched else DistributedState(0, 1, device) ) + except BaseException as e: + exc = e + raise finally: # If the body is unwinding due to an exception, PRINT it here BEFORE anything that waits on # the peers. A per-rank exception (only some ranks raise) leaves the rest inside a # calibration a2a collective, so this rank's cleanup barrier or destroy_process_group would # DEADLOCK with the original exception never surfaced. A run of several ranks that owns its # group therefore aborts the rank: abort() prints the traceback and exits. - exc = sys.exc_info()[1] owns_group = not group_was_there and is_initialized() # Every rank reaches a SystemExit, an intentional exit: it is not an error to report. if exc is not None and not isinstance(exc, SystemExit): diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 4333f1d0c8c..d29dd8c3408 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -155,9 +155,10 @@ def reseeding_worker_teardown(rank, world_size): def _on_group(collective, bound, /, *args, **kwargs): - """``collective`` on ``bound``, whatever group the caller named: modelopt's ``barrier()`` passes - ``group=None`` explicitly, which would override a ``partial``'s keyword.""" - return collective(*args, **{**kwargs, "group": bound}) + """``collective`` on ``bound`` unless the caller named a group of its own: modelopt's + ``barrier()`` passes ``group=None`` explicitly, which would override a ``partial``'s keyword.""" + group = kwargs.get("group") + return collective(*args, **{**kwargs, "group": bound if group is None else group}) @contextmanager diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 064b8c44acc..3dc80cadb7f 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -311,7 +311,7 @@ def fail(): def _bounded_barrier_job(rank, size): - groups = [] + groups, own = [], object() with ( mock.patch.object( torch.distributed, "barrier", lambda *a, group=None, **k: groups.append(group) @@ -319,12 +319,14 @@ def _bounded_barrier_job(rank, size): bounded_collectives(), ): torch.distributed.barrier(group=None) # as modelopt's barrier() calls it - assert groups and groups[0] is not None, groups + torch.distributed.barrier(group=own) + assert len(groups) == 2 and groups[0] not in (None, own) and groups[1] is own, groups def test_bounded_collectives_bound_a_barrier_called_with_group_none(cpu_dist_workers): """modelopt's ``barrier()`` names ``group=None`` explicitly; the bound holds for it too, or a - rank left waiting in it would wait out the default group's timeout.""" + rank left waiting in it would wait out the default group's timeout. A collective on a group of + its own keeps that group.""" cpu_dist_workers.run(_bounded_barrier_job) @@ -369,6 +371,7 @@ class _ExitedError(Exception): """Raised by a stand-in for abort(), which exits the process instead of returning.""" +@pytest.mark.parametrize("in_except", [False, True], ids=["entered-plainly", "entered-in-except"]) @pytest.mark.parametrize( ("launch", "body_raises", "ends_by"), [ @@ -380,11 +383,14 @@ class _ExitedError(Exception): ], ids=["launch", "launch-raises", "launch-exits", "single-process", "single-process-raises"], ) -def test_the_run_ends_the_group_it_set_up(monkeypatch, capsys, launch, body_raises, ends_by): +def test_the_run_ends_the_group_it_set_up( + monkeypatch, capsys, launch, body_raises, ends_by, in_except +): """The launch's group, or the one a single process's FSDP2 load sets up mid-run, is destroyed on the way out. A rank of several whose body raised exits through abort() instead, and torchrun tears the job down: the cleanup's barrier and destroy_process_group would wait on peers that may sit in - a collective this rank never reaches. Every rank reaches an intentional exit.""" + a collective this rank never reaches. Every rank reaches an intentional exit. Only the body's own + exception counts, not one the caller is handling where it entered the run.""" if launch: monkeypatch.setenv("WORLD_SIZE", "2") else: @@ -405,11 +411,20 @@ def abort(): escapes = _ExitedError if ends_by == "abort" else body_raises expected = pytest.raises(escapes) if escapes else nullcontext() - with expected, distributed_run(device="cuda" if launch else "cpu"): - if not launch: - live.append("the body's group") - if body_raises: - raise body_raises + def run(): + with expected, distributed_run(device="cuda" if launch else "cpu"): + if not launch: + live.append("the body's group") + if body_raises: + raise body_raises + + if in_except: + try: + raise KeyError("the caller's") + except KeyError: + run() + else: + run() assert ended == [ends_by] # An intentional exit is not reported as an error. From dcde6a2028b67837aac7578f53454e55b523b48f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:48:23 +0000 Subject: [PATCH 26/48] Read a run's rank from its live group; barrier unless its body raised distributed_run destroyed the group it set up through cleanup(), which decides its barrier by sys.exc_info(). That also reports an exception a caller up the stack is handling, so a clean exit from a run entered in an except block skipped the barrier and destroyed the group while peers could still be in their body. The teardown is now _destroy_group(exc): cleanup() calls it with sys.exc_info()[1], as before, and distributed_run with its body's own exception. The yielded DistributedState took its rank and size from the launch environment (WORLD_SIZE), while the collectives in the same block run on the live group. In a group a caller set up without that environment, every rank was rank 0 of a run of one, and is_main on all of them. The state now reads rank() and size() either way; without a group they are 0 and 1. The run-lifetime test runs the real teardown and records the barrier and destroy_process_group, so its entered-in-except cases check the barrier; the found-group test checks that the state is the group's. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 27 ++++++++---- .../unit/torch/utils/test_distributed_run.py | 44 +++++++++++-------- 2 files changed, 43 insertions(+), 28 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 165d6d67c4f..229395a6330 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -503,8 +503,10 @@ def distributed_run( On a distributed launch (``WORLD_SIZE>1``) this sets the process group up, binding this rank's CUDA device, with ``timeout``, unless the caller set one up already. It yields this rank's - :class:`DistributedState`, on that CUDA device, and refuses any other ``device``. A single - process sets nothing up and is rank 0 of a run of one, on ``device``. + :class:`DistributedState`, on that CUDA device, and refuses any other ``device``. Without a + launch it sets nothing up, and the state is on ``device``. The state's rank and size are always + the live group's, which the block's collectives run on: rank 0 of a run of one for a single + process, and this rank's place in a group the caller set up without the launch's environment. On exit, even if the body raises, it destroys the group the run set up: the launch's, or one the body set up, such as the one-rank group a single process's FSDP2 load shards through. A group @@ -523,13 +525,12 @@ def distributed_run( if launched and not group_was_there: setup(timeout=timeout) # The body's own exception: sys.exc_info() would also report one a caller up the stack is - # handling, which would make a clean exit from a run entered in an except block abort the rank. + # handling, which would make a clean exit from a run entered in an except block abort the rank, + # or destroy its group without the barrier. exc: BaseException | None = None try: - yield ( - DistributedState(rank(), size(), torch.device(f"cuda:{local_rank()}")) - if launched - else DistributedState(0, 1, device) + yield DistributedState( + rank(), size(), torch.device(f"cuda:{local_rank()}") if launched else device ) except BaseException as e: exc = e @@ -553,7 +554,7 @@ def distributed_run( traceback.print_exc() sys.stderr.flush() if owns_group: - cleanup() + _destroy_group(exc) def cleanup(): @@ -564,8 +565,16 @@ def cleanup(): That is not sufficient on its own -- ``destroy_process_group`` below blocks for the same reason -- so error paths must call :func:`abort` before reaching this ``finally``. """ + _destroy_group(sys.exc_info()[1]) + + +def _destroy_group(exc: BaseException | None) -> None: + """:func:`cleanup` for a rank unwinding from ``exc``, or from nothing if it is ``None``. + + For a caller that knows what it unwinds from: :func:`cleanup` reads ``sys.exc_info()``, which + also reports an exception a caller up the stack is handling. + """ if is_initialized(): - exc = sys.exception() if exc is None or isinstance(exc, SystemExit): with suppress(Exception): barrier() diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 3dc80cadb7f..7693314bce7 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -337,7 +337,7 @@ def test_a_single_process_is_rank_0_of_one_on_the_given_device(monkeypatch): """With no ranks to sync with, the run sets no group up and tears none down.""" monkeypatch.delenv("WORLD_SIZE", raising=False) monkeypatch.setattr(dist, "setup", lambda **kwargs: pytest.fail("set a group up")) - monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("tore a group down")) + monkeypatch.setattr(dist, "_destroy_group", lambda exc: pytest.fail("tore a group down")) with distributed_run(device="cpu") as state: assert state == DistributedState(rank=0, world_size=1, device="cpu") assert state.is_main @@ -347,7 +347,7 @@ def test_a_distributed_launch_yields_this_ranks_state_on_its_cuda_device(monkeyp monkeypatch.setenv("WORLD_SIZE", "4") timeouts = [] monkeypatch.setattr(dist, "setup", lambda timeout=None: timeouts.append(timeout)) - monkeypatch.setattr(dist, "cleanup", lambda: None) + monkeypatch.setattr(dist, "_destroy_group", lambda exc: None) monkeypatch.setattr(dist, "rank", lambda group=None: 2) monkeypatch.setattr(dist, "size", lambda group=None: 4) monkeypatch.setattr(dist, "local_rank", lambda: 1) @@ -375,11 +375,11 @@ class _ExitedError(Exception): @pytest.mark.parametrize( ("launch", "body_raises", "ends_by"), [ - (True, None, "cleanup"), - (True, ValueError, "abort"), - (True, SystemExit, "cleanup"), - (False, None, "cleanup"), - (False, ValueError, "cleanup"), + (True, None, ["barrier", "destroy"]), + (True, ValueError, ["abort"]), + (True, SystemExit, ["barrier", "destroy"]), + (False, None, ["barrier", "destroy"]), + (False, ValueError, ["destroy"]), ], ids=["launch", "launch-raises", "launch-exits", "single-process", "single-process-raises"], ) @@ -387,10 +387,11 @@ def test_the_run_ends_the_group_it_set_up( monkeypatch, capsys, launch, body_raises, ends_by, in_except ): """The launch's group, or the one a single process's FSDP2 load sets up mid-run, is destroyed on - the way out. A rank of several whose body raised exits through abort() instead, and torchrun tears - the job down: the cleanup's barrier and destroy_process_group would wait on peers that may sit in - a collective this rank never reaches. Every rank reaches an intentional exit. Only the body's own - exception counts, not one the caller is handling where it entered the run.""" + the way out, after a barrier unless the body raised. A rank of several whose body raised exits + through abort() instead, and torchrun tears the job down: the barrier and destroy_process_group + would wait on peers that may sit in a collective this rank never reaches. Every rank reaches an + intentional exit. Only the body's own exception counts, not one the caller is handling where it + entered the run.""" if launch: monkeypatch.setenv("WORLD_SIZE", "2") else: @@ -401,14 +402,15 @@ def test_the_run_ends_the_group_it_set_up( monkeypatch.setattr(dist, "rank", lambda group=None: 0) monkeypatch.setattr(dist, "size", lambda group=None: 2 if launch else 1) monkeypatch.setattr(dist, "local_rank", lambda: 0) - monkeypatch.setattr(dist, "cleanup", lambda: ended.append("cleanup")) + monkeypatch.setattr(dist, "barrier", lambda group=None: ended.append("barrier")) + monkeypatch.setattr(torch.distributed, "destroy_process_group", lambda: ended.append("destroy")) def abort(): ended.append("abort") raise _ExitedError monkeypatch.setattr(dist, "abort", abort) - escapes = _ExitedError if ends_by == "abort" else body_raises + escapes = _ExitedError if ends_by == ["abort"] else body_raises expected = pytest.raises(escapes) if escapes else nullcontext() def run(): @@ -426,7 +428,7 @@ def run(): else: run() - assert ended == [ends_by] + assert ended == ends_by # An intentional exit is not reported as an error. assert ("body raised" in capsys.readouterr().out) == (body_raises is ValueError) @@ -435,19 +437,23 @@ def run(): @pytest.mark.parametrize("world_size", [None, "2"], ids=["single-process", "launch"]) def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size, raises): """Under a launch too, and when the body raises: a caller that set its group up before the run - keeps it, and gets the error.""" + keeps it, and gets the error. The state is this rank's place in that group, also when the + caller set it up without the launch's environment.""" if world_size is None: monkeypatch.delenv("WORLD_SIZE", raising=False) else: monkeypatch.setenv("WORLD_SIZE", world_size) monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set it up again")) - monkeypatch.setattr(dist, "rank", lambda group=None: 0) - monkeypatch.setattr(dist, "size", lambda group=None: 2) monkeypatch.setattr(dist, "local_rank", lambda: 0) monkeypatch.setattr(dist, "is_initialized", lambda: True) - monkeypatch.setattr(dist, "cleanup", lambda: pytest.fail("destroyed the caller's group")) + monkeypatch.setattr(dist, "rank", lambda group=None: 1) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr( + dist, "_destroy_group", lambda exc: pytest.fail("destroyed the caller's group") + ) monkeypatch.setattr(dist, "abort", lambda: pytest.fail("exited the caller's process")) expected = pytest.raises(ValueError, match="calibration failed") if raises else nullcontext() - with expected, distributed_run(device="cpu" if world_size is None else "cuda"): + with expected, distributed_run(device="cpu" if world_size is None else "cuda") as state: + assert (state.rank, state.world_size) == (1, 2) if raises: raise ValueError("calibration failed") From b1dd500c1c7d8460fd691c6d7d2e9cd1c724153a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:33:05 +0000 Subject: [PATCH 27/48] Exit every rank on a collective's SystemExit; fail unpicklable values broadcast_from and all_reduce sent a failing rank's error to the others as text, which they raised as a RuntimeError. A SystemExit on one rank then took distributed_run's exit path there (a barrier, then destroy_process_group) and the abort path on its peers, leaving the exiting rank in a barrier its aborted peers never reach. The error now goes with its kind: the peers of a rank that exited raise SystemExit with its code (a code other than an int or None as its text, which the interpreter prints for it), so every rank exits the same way. fn's value was pickled only inside broadcast_object_list or all_gather_object, so a value that did not pickle raised on its own rank before the collective sent anything, and left the other ranks waiting in it. The value is now also pickled inside fn's try, with torch.save into an io.BytesIO (the bandit hook refuses `import pickle` under modelopt/), so it fails as that rank's error, which the collective carries to every rank. set_distribution_plan's refusal also names the submodule that already records a plan by its qualified name, which says where in a real model the plan is. The two-rank broadcast_from and all_reduce tests cover an exit and a value that does not pickle; both fail without the fixes, the latter fast rather than hanging. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 79 +++++++++++++------ .../unit/torch/utils/test_distributed_run.py | 53 ++++++++++--- 2 files changed, 101 insertions(+), 31 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 229395a6330..37ccb89525c 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -360,13 +360,13 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non any of whose submodules, already records one raises, the same plan included, and leaves the model as it was. No code reloads or re-distributes a model it loaded. """ - for module in model.modules(): + for name, module in model.named_modules(): recorded = getattr(module, _PLAN_ATTR, None) if recorded is not None: + where = f"{type(module).__name__} {name!r}" if name else type(module).__name__ raise ValueError( - f"{type(module).__name__} is already recorded as distributed by {recorded!r}; a " - f"model is distributed and its plan recorded once, so {plan!r} cannot be recorded " - "over it." + f"{where} is already recorded as distributed by {recorded!r}; a model is " + f"distributed and its plan recorded once, so {plan!r} cannot be recorded over it." ) for module in model.modules(): setattr(module, _PLAN_ATTR, plan) @@ -398,16 +398,43 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: _V = TypeVar("_V") +def _sendable(error: BaseException | None) -> tuple[bool, Any] | None: + """``error`` as this rank sends it to the others, ``None`` for none; see :func:`_received`. + + An error goes as its text: the other ranks only need to stop, and an exception is not + guaranteed to pickle. A SystemExit goes as its code, for every rank to exit with: + :func:`distributed_run` ends an exit with a barrier but an error with :func:`abort`, so a rank + exiting while its peers abort would wait in that barrier for ranks that are gone. A code other + than an int or ``None`` goes as its text, which the interpreter prints for it before exiting + with 1. + """ + if error is None: + return None + if isinstance(error, SystemExit): + code = error.code + return True, code if code is None or isinstance(code, int) else str(code) + return False, f"{type(error).__name__}: {error}" + + +def _received(sent: tuple[bool, Any], rank: int, label: str) -> BaseException: + """The exception to raise for what ``rank`` sent (:func:`_sendable`) on failing ``label``.""" + exited, detail = sent + if exited: + return SystemExit(detail) + return RuntimeError(f"rank {rank} failed to compute {label}: {detail}") + + def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: """Return ``fn()`` run on ``rank`` only, on every rank. For a value that depends on what each rank reads for itself -- the environment, its GPU, files - on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, every - rank raises, instead of waiting in a broadcast ``rank`` never sends; its value must pickle, - which the broadcast does before it sends anything. A rank broadcasting another label than - ``rank`` at the same point fails; ``rank`` sees no other rank's label, and returns. Collective: - every rank calls it at the same point, with the same ``rank``. Outside a distributed run ``fn`` - runs alone. + on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, or + returns a value that does not pickle, as the broadcast needs, every rank raises instead of + waiting in a broadcast ``rank`` never sends: ``rank`` its own error, the others a SystemExit + with its code if that is what it was, so that every rank exits, else a RuntimeError. A rank + broadcasting another label than ``rank`` at the same point fails; ``rank`` sees no other rank's + label, and returns. Collective: every rank calls it at the same point, with the same ``rank``. + Outside a distributed run ``fn`` runs alone. """ if not 0 <= rank < size(): raise ValueError(f"rank={rank} is not a rank of this {size()}-rank run.") @@ -419,20 +446,23 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: if sending: try: value = fn() + # Pickle it here, as the broadcast will: a value that does not pickle then fails as this + # rank's error, which the broadcast sends, rather than in the broadcast, on this rank + # alone, leaving the others waiting. + torch.save(value, io.BytesIO()) except BaseException as e: # Still send: the other ranks would otherwise wait in the broadcast for this one. A # SystemExit or KeyboardInterrupt too, which argparse-style checks raise. - error = e + value, error = None, e # One message from the sending rank: its label, so a rank broadcasting another learns the two - # are out of step, and its value or its error. The error goes as text: the other ranks only - # need to stop, and an exception is not guaranteed to pickle. - message = [(label, value, None if error is None else f"{type(error).__name__}: {error}")] + # are out of step, and its value or its error. + message = [(label, value, _sendable(error))] torch.distributed.broadcast_object_list(message, group_src=rank) sent_label, value, sent_error = message[0] if error is not None: raise error if sent_error is not None: - raise RuntimeError(f"rank {rank} failed to compute {label}: {sent_error}") + raise _received(sent_error, rank, label) if sent_label != label: raise RuntimeError( f"Ranks broadcast out of step: this rank broadcast {label!r} while rank {rank} sent " @@ -446,28 +476,33 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ ``combine`` takes the ranks' values in rank order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's memory bounds, where one rank's choice could run another - out of memory. If ``fn`` raises on any rank, every rank raises, instead of waiting in the gather - that rank never joins; its value must pickle, which the gather does before it gathers anything. - Ranks reducing different labels at the same point fail. Collective; outside a distributed run - ``fn`` runs alone. + out of memory. If ``fn`` raises on any rank, or returns a value that does not pickle, as the + gather needs, every rank raises instead of waiting in a gather that rank never joins: a failing + rank its own error, the others the first failing rank's, a SystemExit as a SystemExit with its + code, anything else as a RuntimeError. Ranks reducing different labels at the same point fail. + Collective; outside a distributed run ``fn`` runs alone. """ if not is_distributed_run(): return combine([fn()]) value, error = None, None try: value = fn() + # Pickle it here, as the gather will: a value that does not pickle then fails as this + # rank's error, which the gather carries, rather than in the gather, on this rank alone, + # leaving the others waiting. + torch.save(value, io.BytesIO()) except BaseException as e: # Still join the gather: the other ranks would otherwise wait in it for this one, a # SystemExit or KeyboardInterrupt included. - error = e - mine = (label, value, None if error is None else f"{type(error).__name__}: {error}") + value, error = None, e + mine = (label, value, _sendable(error)) gathered: list[Any] = [None] * torch.distributed.get_world_size() torch.distributed.all_gather_object(gathered, mine) if error is not None: raise error for rank_, (_, _, rank_error) in enumerate(gathered): if rank_error is not None: - raise RuntimeError(f"rank {rank_} failed to compute {label}: {rank_error}") + raise _received(rank_error, rank_, label) labels = sorted({rank_label for rank_label, _, _ in gathered}) if labels != [label]: raise RuntimeError( diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 7693314bce7..fbfff2a4baa 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -126,14 +126,14 @@ def test_a_model_with_no_plan_is_a_single_process(): def _outcome(call): - """``(value, None)``, or ``(None, error)`` if ``call()`` raises. + """``(value, None)``, or ``(None, error)`` if ``call()`` raises, a ``SystemExit`` included. For pooled jobs, which cannot use ``pytest.raises``: its failure is a ``BaseException``, which - ends the worker instead of failing the job. + ends the worker instead of failing the job, as an uncaught ``SystemExit`` would. """ try: return call(), None - except Exception as error: + except (Exception, SystemExit) as error: return None, error @@ -226,6 +226,9 @@ def compute(): def fail(): raise ValueError("bad degrees") + def exit_2(): + raise SystemExit(2) + with bounded_collectives(): for sender in range(size): ran.clear() @@ -240,6 +243,14 @@ def fail(): else: assert isinstance(error, RuntimeError), error assert "rank 0 failed to compute x" in str(error), error + _, error = _outcome(lambda: broadcast_from("x", exit_2)) + assert isinstance(error, SystemExit) and error.code == 2, error + _, error = _outcome(lambda: broadcast_from("x", lambda: (n for n in range(1)))) + if rank == 0: + assert isinstance(error, TypeError), error + else: + assert isinstance(error, RuntimeError), error + assert "rank 0 failed to compute x: TypeError" in str(error), error sent, error = _outcome(lambda: broadcast_from("plan" if rank == 0 else "other", lambda: 7)) if rank == 0: assert (sent, error) == (7, None) @@ -250,15 +261,17 @@ def fail(): def test_broadcast_from_gives_every_rank_the_sending_ranks_value(cpu_dist_workers): """Over a real broadcast: any rank can send, and only it runs fn. A rank outside the run is - refused, the sending rank's failure stops every rank, and ranks out of step fail.""" + refused, the sending rank's failure stops every rank and its exit exits every rank with its + code, so that distributed_run tears every rank down alike. A value that does not pickle fails as + fn raising does, rather than in the broadcast, and ranks out of step fail.""" cpu_dist_workers.run(_broadcast_from_job) @pytest.mark.parametrize("error", [ValueError, SystemExit]) def test_the_sending_rank_failing_still_reaches_the_broadcast(monkeypatch, error): """The sending rank must still send on failure, or every other rank hangs in the broadcast; a - SystemExit from an argparse-style check too. A stand-in broadcast catches that at once, where a - real one would wait out its timeout.""" + SystemExit from an argparse-style check too, which goes as its code. A stand-in broadcast + catches that at once, where a real one would wait out its timeout.""" sent = _as_rank(monkeypatch, 0) def fail(): @@ -266,7 +279,8 @@ def fail(): with pytest.raises(error, match="bad degrees"): broadcast_from("x", fail) - assert sent == [(("x", None, f"{error.__name__}: bad degrees"), 0)] + told = (True, "bad degrees") if error is SystemExit else (False, "ValueError: bad degrees") + assert sent == [(("x", None, told), 0)] def _all_reduce_job(rank, size): @@ -275,6 +289,11 @@ def fail_on_rank_1(): raise ValueError("no memory reading") return 4 + def exit_on_rank_1(): + if rank == 1: + raise SystemExit(3) + return 4 + with bounded_collectives(): assert all_reduce("smallest", lambda: [8, 2][rank], min) == 2 _, error = _outcome(lambda: all_reduce("smallest", fail_on_rank_1, min)) @@ -283,6 +302,16 @@ def fail_on_rank_1(): else: assert isinstance(error, RuntimeError), error assert "rank 1 failed to compute smallest" in str(error), error + _, error = _outcome(lambda: all_reduce("smallest", exit_on_rank_1, min)) + assert isinstance(error, SystemExit) and error.code == 3, error + _, error = _outcome( + lambda: all_reduce("smallest", lambda: (n for n in range(1)) if rank == 1 else 4, min) + ) + if rank == 1: + assert isinstance(error, TypeError), error + else: + assert isinstance(error, RuntimeError), error + assert "rank 1 failed to compute smallest: TypeError" in str(error), error _, error = _outcome( lambda: all_reduce("smallest" if rank == 0 else "other", lambda: 4, min) ) @@ -292,7 +321,8 @@ def fail_on_rank_1(): def test_all_reduce_gives_every_rank_the_combined_value(cpu_dist_workers): """Over a real gather: every rank's value counts and every rank gets the same result. One rank's - failure stops every rank, and ranks out of step fail.""" + failure stops every rank and its exit exits every rank with its code. A value that does not + pickle fails as fn raising does, rather than in the gather, and ranks out of step fail.""" cpu_dist_workers.run(_all_reduce_job) @@ -307,7 +337,12 @@ def fail(): with pytest.raises(error, match="no memory reading"): all_reduce("smallest", fail, min) - assert sent == [("smallest", None, f"{error.__name__}: no memory reading")] + told = ( + (True, "no memory reading") + if error is SystemExit + else (False, "ValueError: no memory reading") + ) + assert sent == [("smallest", None, told)] def _bounded_barrier_job(rank, size): From ad2689c96f6baf4a1fb00f0ba485aad57f835e3a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 05:23:50 +0000 Subject: [PATCH 28/48] Fail every rank of an all_reduce alike when ranks fail differently all_reduce raised a failing rank's own exception before it read the gathered outcomes. When one rank raised a SystemExit and another an ordinary error in the same call, the ranks split: the erring rank aborted through distributed_run, while the exiting rank, and every healthy rank that found it first, took the exit path into a barrier that waits for the aborted rank. Every rank now reads the gathered outcomes and raises by one rule: an error other than a SystemExit on any rank makes every rank raise an error (a rank with one its own, the others a RuntimeError naming the first), else every rank exits. broadcast_from has one sender, so one kind of failure, and is unchanged. distributed_run reports a body's error under rank(), not the RANK variable, which a group the caller set up need not have, and no longer prints the traceback itself where the error goes on to the caller; abort() still prints it on its own path. all_reduce's docstring says why it runs fn itself, unlike allreduce. The two-rank all_reduce test adds an exit on rank 0 with an error on rank 1; it fails without the fix, rank 0 exiting. The run-lifetime test checks the reported rank and that no traceback prints. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 38 ++++++++++++------- .../unit/torch/utils/test_distributed_run.py | 29 +++++++++----- 2 files changed, 43 insertions(+), 24 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 37ccb89525c..958f286e5eb 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -477,10 +477,15 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ ``combine`` takes the ranks' values in rank order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's memory bounds, where one rank's choice could run another out of memory. If ``fn`` raises on any rank, or returns a value that does not pickle, as the - gather needs, every rank raises instead of waiting in a gather that rank never joins: a failing - rank its own error, the others the first failing rank's, a SystemExit as a SystemExit with its - code, anything else as a RuntimeError. Ranks reducing different labels at the same point fail. - Collective; outside a distributed run ``fn`` runs alone. + gather needs, every rank raises instead of waiting in a gather that rank never joins, and every + rank alike: if any rank's error is not a SystemExit, a rank with such an error raises its own + and the others a RuntimeError naming the first one; else every rank exits, a rank that exited + with its own SystemExit and the others with the first one's code. Ranks reducing different + labels at the same point fail. Collective; outside a distributed run ``fn`` runs alone. + + It runs ``fn`` itself so that a rank whose ``fn`` fails still joins the gather, which + :func:`allreduce`, handed values the ranks already hold, cannot do; :func:`broadcast_from` + differs from :func:`broadcast` and :func:`master_only` in the same way. """ if not is_distributed_run(): return combine([fn()]) @@ -498,11 +503,18 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ mine = (label, value, _sendable(error)) gathered: list[Any] = [None] * torch.distributed.get_world_size() torch.distributed.all_gather_object(gathered, mine) - if error is not None: - raise error - for rank_, (_, _, rank_error) in enumerate(gathered): - if rank_error is not None: - raise _received(rank_error, rank_, label) + failures = [(rank_, sent) for rank_, (_, _, sent) in enumerate(gathered) if sent is not None] + if failures: + # Every rank raises by one rule, for distributed_run to end every rank alike. A rank whose + # error is not a SystemExit aborts, so if any rank's is not, every rank raises an error, an + # exiting rank too, rather than wait in the exit's barrier for the aborted one; else every + # rank exits. A failing rank raises its own exception when it is of that kind. Each sent is + # _sendable's (exited, detail). + errors = [(rank_, sent) for rank_, sent in failures if not sent[0]] + failed_rank, sent = (errors or failures)[0] + if error is not None and isinstance(error, SystemExit) == sent[0]: + raise error + raise _received(sent, failed_rank, label) labels = sorted({rank_label for rank_label, _, _ in gathered}) if labels != [label]: raise RuntimeError( @@ -575,19 +587,17 @@ def distributed_run( # the peers. A per-rank exception (only some ranks raise) leaves the rest inside a # calibration a2a collective, so this rank's cleanup barrier or destroy_process_group would # DEADLOCK with the original exception never surfaced. A run of several ranks that owns its - # group therefore aborts the rank: abort() prints the traceback and exits. + # group therefore aborts the rank: abort() prints the traceback and exits. Anywhere else the + # error goes on to the caller, traceback and all. owns_group = not group_was_there and is_initialized() # Every rank reaches a SystemExit, an intentional exit: it is not an error to report. if exc is not None and not isinstance(exc, SystemExit): print( - f"[distributed_run][rank {os.environ.get('RANK', '?')}] body raised " - f"{type(exc).__name__}: {exc}", + f"[distributed_run][rank {rank()}] body raised {type(exc).__name__}: {exc}", flush=True, ) if owns_group and size() > 1: abort() - traceback.print_exc() - sys.stderr.flush() if owns_group: _destroy_group(exc) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index fbfff2a4baa..a9dd19ffdce 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -294,14 +294,18 @@ def exit_on_rank_1(): raise SystemExit(3) return 4 + def exit_on_rank_0_fail_on_rank_1(): + raise SystemExit(3) if rank == 0 else ValueError("no memory reading") + with bounded_collectives(): assert all_reduce("smallest", lambda: [8, 2][rank], min) == 2 - _, error = _outcome(lambda: all_reduce("smallest", fail_on_rank_1, min)) - if rank == 1: - assert isinstance(error, ValueError), error - else: - assert isinstance(error, RuntimeError), error - assert "rank 1 failed to compute smallest" in str(error), error + for fail in (fail_on_rank_1, exit_on_rank_0_fail_on_rank_1): + _, error = _outcome(lambda: all_reduce("smallest", fail, min)) + if rank == 1: + assert isinstance(error, ValueError), error + else: + assert isinstance(error, RuntimeError), error + assert "rank 1 failed to compute smallest: ValueError" in str(error), error _, error = _outcome(lambda: all_reduce("smallest", exit_on_rank_1, min)) assert isinstance(error, SystemExit) and error.code == 3, error _, error = _outcome( @@ -321,8 +325,10 @@ def exit_on_rank_1(): def test_all_reduce_gives_every_rank_the_combined_value(cpu_dist_workers): """Over a real gather: every rank's value counts and every rank gets the same result. One rank's - failure stops every rank and its exit exits every rank with its code. A value that does not - pickle fails as fn raising does, rather than in the gather, and ranks out of step fail.""" + failure stops every rank and its exit exits every rank with its code. A rank's error outranks + another's exit, so that no rank waits in the exit's barrier for one that aborted. A value that + does not pickle fails as fn raising does, rather than in the gather, and ranks out of step + fail.""" cpu_dist_workers.run(_all_reduce_job) @@ -464,8 +470,11 @@ def run(): run() assert ended == ends_by - # An intentional exit is not reported as an error. - assert ("body raised" in capsys.readouterr().out) == (body_raises is ValueError) + # An intentional exit is not reported as an error. An error is reported once, under this rank's + # place in the group; its traceback is abort()'s to print, or goes on with it to the caller. + out, err = capsys.readouterr() + assert ("[distributed_run][rank 0] body raised" in out) == (body_raises is ValueError) + assert "Traceback" not in err @pytest.mark.parametrize("raises", [False, True], ids=["returns", "raises"]) From a99099def72806441100604706a6bcd47df0a322 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:51 +0000 Subject: [PATCH 29/48] Bind the rank's CUDA device under a launch, also to a group the caller set up distributed_run left a group the caller had set up alone, device included, yet reported the rank's own CUDA device. A caller that set its group up without binding one left every rank on cuda:0, which NCCL then runs on. setup() binds the device whoever set the group up, so the run calls it on every launch; it still sets a group up only when there is none. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 8 +++++--- tests/unit/torch/utils/test_distributed_run.py | 14 +++++++++++--- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 958f286e5eb..e82a13cf42e 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -548,8 +548,8 @@ def distributed_run( ) -> Iterator[DistributedState]: """Own the run's process group for the block, and yield this process's state. - On a distributed launch (``WORLD_SIZE>1``) this sets the process group up, binding this rank's - CUDA device, with ``timeout``, unless the caller set one up already. It yields this rank's + On a distributed launch (``WORLD_SIZE>1``) this binds this rank's CUDA device, and sets the + process group up with ``timeout`` unless the caller set one up already. It yields this rank's :class:`DistributedState`, on that CUDA device, and refuses any other ``device``. Without a launch it sets nothing up, and the state is on ``device``. The state's rank and size are always the live group's, which the block's collectives run on: rank 0 of a run of one for a single @@ -569,7 +569,9 @@ def distributed_run( "CUDA device, chosen from LOCAL_RANK; pass device='cuda' or run a single process." ) group_was_there = is_initialized() - if launched and not group_was_there: + if launched: + # setup() binds the device even to a group the caller set up, which may not have bound one: + # the state's device is then the one this rank's NCCL collectives run on. setup(timeout=timeout) # The body's own exception: sys.exc_info() would also report one a caller up the stack is # handling, which would make a clean exit from a run entered in an except block abort the rank, diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index a9dd19ffdce..c8d33604137 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -482,13 +482,19 @@ def run(): def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size, raises): """Under a launch too, and when the body raises: a caller that set its group up before the run keeps it, and gets the error. The state is this rank's place in that group, also when the - caller set it up without the launch's environment.""" + caller set it up without the launch's environment. Under a launch the run binds this rank's + CUDA device all the same, since the caller may not have: the state's device is then the one + NCCL runs on.""" + bound: list[int] = [] if world_size is None: monkeypatch.delenv("WORLD_SIZE", raising=False) else: monkeypatch.setenv("WORLD_SIZE", world_size) - monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set it up again")) - monkeypatch.setattr(dist, "local_rank", lambda: 0) + monkeypatch.setattr( + torch.distributed, "init_process_group", lambda *a, **k: pytest.fail("set it up again") + ) + monkeypatch.setattr(torch.cuda, "set_device", bound.append) + monkeypatch.setattr(dist, "local_rank", lambda: 1) monkeypatch.setattr(dist, "is_initialized", lambda: True) monkeypatch.setattr(dist, "rank", lambda group=None: 1) monkeypatch.setattr(dist, "size", lambda group=None: 2) @@ -499,5 +505,7 @@ def test_the_run_leaves_a_group_it_found_alone(monkeypatch, world_size, raises): expected = pytest.raises(ValueError, match="calibration failed") if raises else nullcontext() with expected, distributed_run(device="cpu" if world_size is None else "cuda") as state: assert (state.rank, state.world_size) == (1, 2) + assert torch.device(state.device) == torch.device("cpu" if world_size is None else "cuda:1") + assert bound == ([] if world_size is None else [1]) if raises: raise ValueError("calibration failed") From 1a53333f3189543381d362431de5bad31e630714 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:51 +0000 Subject: [PATCH 30/48] Refuse the collectives on a launch whose process group is not set up broadcast_from and all_reduce ran fn on each rank alone there, as in a single process, so the ranks could compute different values with nothing telling them. They raise instead, naming the launch as the loaders' own refusal does, and say to compute the value inside distributed_run(). Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 21 +++++++++++++++++-- .../unit/torch/utils/test_distributed_run.py | 15 +++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index e82a13cf42e..16c481f10b2 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -416,6 +416,19 @@ def _sendable(error: BaseException | None) -> tuple[bool, Any] | None: return False, f"{type(error).__name__}: {error}" +def _refuse_a_launch_without_its_group(label: str) -> None: + """Raise if this process is a rank of a multi-rank launch whose process group is not set up. + + Each rank would otherwise compute ``label`` alone, and the ranks could disagree with nothing + telling them so. + """ + if is_distributed_launch() and not is_initialized(): + raise RuntimeError( + f"{label!r} needs this {os.environ['WORLD_SIZE']}-rank launch's process group, which is " + "not set up yet; compute it inside distributed_run(), which sets the group up." + ) + + def _received(sent: tuple[bool, Any], rank: int, label: str) -> BaseException: """The exception to raise for what ``rank`` sent (:func:`_sendable`) on failing ``label``.""" exited, detail = sent @@ -434,8 +447,10 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: with its code if that is what it was, so that every rank exits, else a RuntimeError. A rank broadcasting another label than ``rank`` at the same point fails; ``rank`` sees no other rank's label, and returns. Collective: every rank calls it at the same point, with the same ``rank``. - Outside a distributed run ``fn`` runs alone. + In a single process ``fn`` runs alone; a rank of a launch whose process group is not set up yet + raises instead. """ + _refuse_a_launch_without_its_group(label) if not 0 <= rank < size(): raise ValueError(f"rank={rank} is not a rank of this {size()}-rank run.") if not is_distributed_run(): @@ -481,12 +496,14 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ rank alike: if any rank's error is not a SystemExit, a rank with such an error raises its own and the others a RuntimeError naming the first one; else every rank exits, a rank that exited with its own SystemExit and the others with the first one's code. Ranks reducing different - labels at the same point fail. Collective; outside a distributed run ``fn`` runs alone. + labels at the same point fail. Collective; in a single process ``fn`` runs alone, and a rank of a + launch whose process group is not set up yet raises instead. It runs ``fn`` itself so that a rank whose ``fn`` fails still joins the gather, which :func:`allreduce`, handed values the ranks already hold, cannot do; :func:`broadcast_from` differs from :func:`broadcast` and :func:`master_only` in the same way. """ + _refuse_a_launch_without_its_group(label) if not is_distributed_run(): return combine([fn()]) value, error = None, None diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index c8d33604137..9652cfd3aae 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -216,6 +216,21 @@ def test_a_single_process_computes_alone(monkeypatch): broadcast_from("x", lambda: 3, rank=1) +def test_a_launch_without_its_group_refuses_the_collectives(monkeypatch): + """A rank of a multi-rank launch whose group is not set up yet would compute its own value, + which another rank's could differ from with nothing telling them; it raises instead.""" + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setattr(dist, "is_initialized", lambda: False) + ran = [] + for collective in ( + lambda: broadcast_from("x", lambda: ran.append("broadcast_from")), + lambda: all_reduce("x", lambda: ran.append("all_reduce"), min), + ): + with pytest.raises(RuntimeError, match="'x' needs this 2-rank launch's process group"): + collective() + assert ran == [] + + def _broadcast_from_job(rank, size): ran = [] From def9c44480f7a2207c5cf47f1b50e8dfce19ee29 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:51 +0000 Subject: [PATCH 31/48] Say why the collectives check and unpickle values as they do torch.save stands in for the object collectives' pickle, which bandit refuses under modelopt/, and those collectives unpickle with plain pickle rather than _deserialize's weights_only load. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 16c481f10b2..718efbfeaa1 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -463,14 +463,17 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: value = fn() # Pickle it here, as the broadcast will: a value that does not pickle then fails as this # rank's error, which the broadcast sends, rather than in the broadcast, on this rank - # alone, leaving the others waiting. + # alone, leaving the others waiting. torch.save stands in for the broadcast's pickle, + # which bandit refuses under modelopt/, so the check is close rather than exact. torch.save(value, io.BytesIO()) except BaseException as e: # Still send: the other ranks would otherwise wait in the broadcast for this one. A # SystemExit or KeyboardInterrupt too, which argparse-style checks raise. value, error = None, e # One message from the sending rank: its label, so a rank broadcasting another learns the two - # are out of step, and its value or its error. + # are out of step, and its value or its error. The object collectives unpickle with plain + # pickle, not _deserialize's weights_only load: the peers are this job's own ranks, and a + # DistributionPlan crosses without add_safe_globals. message = [(label, value, _sendable(error))] torch.distributed.broadcast_object_list(message, group_src=rank) sent_label, value, sent_error = message[0] From 23549a6e071c0ba981612782dc911df426c7673e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 16:48:44 +0000 Subject: [PATCH 32/48] Serialize a collective's value once, where the failure is still sent broadcast_from and all_reduce checked that a value pickles with torch.save, protocol 2, and then handed the value itself to the object collective, which pickled it again with plain pickle's default protocol, outside the guard. A value that passed the check could then fail inside the collective on one rank alone, and the other ranks waited for it. Each value is now serialized once, inside the guard, and loaded back there too, so a value the other ranks could not load fails alike; the collective carries the bytes, which always pickle, and every receiving rank loads them. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 90 +++++++++++-------- .../unit/torch/utils/test_distributed_run.py | 28 +++++- 2 files changed, 78 insertions(+), 40 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 718efbfeaa1..ed713a3cf0d 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -437,18 +437,44 @@ def _received(sent: tuple[bool, Any], rank: int, label: str) -> BaseException: return RuntimeError(f"rank {rank} failed to compute {label}: {detail}") +def _encoded(value: Any) -> bytes: + """``value`` as :func:`broadcast_from` and :func:`all_reduce` send it: serialized once, here. + + It is also loaded back here: a value that does not serialize, or that the other ranks could not + load (a class a forced weights-only load refuses, say), then fails inside the caller's guard, + which sends the failure to every rank, rather than inside the collective, on one rank alone. The + collective carries the bytes, which always pickle. torch.save stands in for plain pickle, which + bandit refuses under modelopt/. + """ + buffer = io.BytesIO() + torch.save(value, buffer) + payload = buffer.getvalue() + _decoded(payload) + return payload + + +def _decoded(payload: bytes) -> Any: + """The value :func:`_encoded` serialized into ``payload``. + + A full load, not :func:`_deserialize`'s weights_only one: the peers are this job's own ranks, + as the object collectives' own unpickling would trust them, and a DistributionPlan crosses + without add_safe_globals. + """ + return torch.load(io.BytesIO(payload), weights_only=False) + + def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: """Return ``fn()`` run on ``rank`` only, on every rank. For a value that depends on what each rank reads for itself -- the environment, its GPU, files on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, or - returns a value that does not pickle, as the broadcast needs, every rank raises instead of - waiting in a broadcast ``rank`` never sends: ``rank`` its own error, the others a SystemExit - with its code if that is what it was, so that every rank exits, else a RuntimeError. A rank - broadcasting another label than ``rank`` at the same point fails; ``rank`` sees no other rank's - label, and returns. Collective: every rank calls it at the same point, with the same ``rank``. - In a single process ``fn`` runs alone; a rank of a launch whose process group is not set up yet - raises instead. + returns a value that does not serialize and load back (:func:`_encoded`), every rank raises + instead of waiting in a broadcast ``rank`` never sends: ``rank`` its own error, the others a + SystemExit with its code if that is what it was, so that every rank exits, else a RuntimeError. + A rank broadcasting another label than ``rank`` at the same point fails; ``rank`` sees no other + rank's label, and returns. Collective: every rank calls it at the same point, with the same + ``rank``. In a single process ``fn`` runs alone; a rank of a launch whose process group is not + set up yet raises instead. """ _refuse_a_launch_without_its_group(label) if not 0 <= rank < size(): @@ -457,26 +483,21 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: return fn() sending = torch.distributed.get_rank() == rank value: Any = None + payload: Any = None error = None if sending: try: value = fn() - # Pickle it here, as the broadcast will: a value that does not pickle then fails as this - # rank's error, which the broadcast sends, rather than in the broadcast, on this rank - # alone, leaving the others waiting. torch.save stands in for the broadcast's pickle, - # which bandit refuses under modelopt/, so the check is close rather than exact. - torch.save(value, io.BytesIO()) + payload = _encoded(value) except BaseException as e: # Still send: the other ranks would otherwise wait in the broadcast for this one. A # SystemExit or KeyboardInterrupt too, which argparse-style checks raise. - value, error = None, e + value, payload, error = None, None, e # One message from the sending rank: its label, so a rank broadcasting another learns the two - # are out of step, and its value or its error. The object collectives unpickle with plain - # pickle, not _deserialize's weights_only load: the peers are this job's own ranks, and a - # DistributionPlan crosses without add_safe_globals. - message = [(label, value, _sendable(error))] + # are out of step, and its serialized value or its error. + message = [(label, payload, _sendable(error))] torch.distributed.broadcast_object_list(message, group_src=rank) - sent_label, value, sent_error = message[0] + sent_label, payload, sent_error = message[0] if error is not None: raise error if sent_error is not None: @@ -486,7 +507,7 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: f"Ranks broadcast out of step: this rank broadcast {label!r} while rank {rank} sent " f"{sent_label!r}. Every rank must broadcast them in the same order." ) - return value + return value if sending else _decoded(payload) def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _V]) -> _V: @@ -494,13 +515,13 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ ``combine`` takes the ranks' values in rank order and returns the one every rank gets, e.g. ``min`` for a batch size each rank's memory bounds, where one rank's choice could run another - out of memory. If ``fn`` raises on any rank, or returns a value that does not pickle, as the - gather needs, every rank raises instead of waiting in a gather that rank never joins, and every - rank alike: if any rank's error is not a SystemExit, a rank with such an error raises its own - and the others a RuntimeError naming the first one; else every rank exits, a rank that exited - with its own SystemExit and the others with the first one's code. Ranks reducing different - labels at the same point fail. Collective; in a single process ``fn`` runs alone, and a rank of a - launch whose process group is not set up yet raises instead. + out of memory. If ``fn`` raises on any rank, or returns a value that does not serialize and load + back (:func:`_encoded`), every rank raises instead of waiting in a gather that rank never joins, + and every rank alike: if any rank's error is not a SystemExit, a rank with such an error raises + its own and the others a RuntimeError naming the first one; else every rank exits, a rank that + exited with its own SystemExit and the others with the first one's code. Ranks reducing + different labels at the same point fail. Collective; in a single process ``fn`` runs alone, and + a rank of a launch whose process group is not set up yet raises instead. It runs ``fn`` itself so that a rank whose ``fn`` fails still joins the gather, which :func:`allreduce`, handed values the ranks already hold, cannot do; :func:`broadcast_from` @@ -509,18 +530,14 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ _refuse_a_launch_without_its_group(label) if not is_distributed_run(): return combine([fn()]) - value, error = None, None + payload, error = None, None try: - value = fn() - # Pickle it here, as the gather will: a value that does not pickle then fails as this - # rank's error, which the gather carries, rather than in the gather, on this rank alone, - # leaving the others waiting. - torch.save(value, io.BytesIO()) + payload = _encoded(fn()) except BaseException as e: # Still join the gather: the other ranks would otherwise wait in it for this one, a # SystemExit or KeyboardInterrupt included. - value, error = None, e - mine = (label, value, _sendable(error)) + payload, error = None, e + mine = (label, payload, _sendable(error)) gathered: list[Any] = [None] * torch.distributed.get_world_size() torch.distributed.all_gather_object(gathered, mine) failures = [(rank_, sent) for rank_, (_, _, sent) in enumerate(gathered) if sent is not None] @@ -541,8 +558,9 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ f"Ranks reduced out of step: {labels} at the same point. Every rank must reduce " "them in the same order." ) - # The gathered values are unpickled copies, so the result shares nothing with the callers'. - return combine([rank_value for _, rank_value, _ in gathered]) + # Every rank's value is loaded from what it sent, this rank's too, so the result shares nothing + # with the callers'. + return combine([_decoded(rank_payload) for _, rank_payload, _ in gathered]) @dataclass(frozen=True) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 9652cfd3aae..d6a269bc443 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -231,6 +231,18 @@ def test_a_launch_without_its_group_refuses_the_collectives(monkeypatch): assert ran == [] +class _Protocol2Only: + """Pickles at protocol 2, as torch.save does, and refuses plain pickle's higher default.""" + + def __init__(self, n): + self.n = n + + def __reduce_ex__(self, protocol): + if protocol > 2: + raise TypeError("protocol 2 only") + return _Protocol2Only, (self.n,) + + def _broadcast_from_job(rank, size): ran = [] @@ -266,6 +278,8 @@ def exit_2(): else: assert isinstance(error, RuntimeError), error assert "rank 0 failed to compute x: TypeError" in str(error), error + # Serialized once, where the sending rank catches a failure, not again by the broadcast. + assert broadcast_from("x", lambda: _Protocol2Only(7)).n == 7 sent, error = _outcome(lambda: broadcast_from("plan" if rank == 0 else "other", lambda: 7)) if rank == 0: assert (sent, error) == (7, None) @@ -277,8 +291,9 @@ def exit_2(): def test_broadcast_from_gives_every_rank_the_sending_ranks_value(cpu_dist_workers): """Over a real broadcast: any rank can send, and only it runs fn. A rank outside the run is refused, the sending rank's failure stops every rank and its exit exits every rank with its - code, so that distributed_run tears every rank down alike. A value that does not pickle fails as - fn raising does, rather than in the broadcast, and ranks out of step fail.""" + code, so that distributed_run tears every rank down alike. A value that does not serialize + fails as fn raising does, rather than in the broadcast; one only torch.save's pickle protocol + takes crosses, serialized once; and ranks out of step fail.""" cpu_dist_workers.run(_broadcast_from_job) @@ -331,6 +346,11 @@ def exit_on_rank_0_fail_on_rank_1(): else: assert isinstance(error, RuntimeError), error assert "rank 1 failed to compute smallest: TypeError" in str(error), error + # Serialized once, where each rank catches a failure, not again by the gather. + reduced = all_reduce( + "x", lambda: _Protocol2Only(rank), lambda values: [v.n for v in values] + ) + assert reduced == list(range(size)), reduced _, error = _outcome( lambda: all_reduce("smallest" if rank == 0 else "other", lambda: 4, min) ) @@ -342,8 +362,8 @@ def test_all_reduce_gives_every_rank_the_combined_value(cpu_dist_workers): """Over a real gather: every rank's value counts and every rank gets the same result. One rank's failure stops every rank and its exit exits every rank with its code. A rank's error outranks another's exit, so that no rank waits in the exit's barrier for one that aborted. A value that - does not pickle fails as fn raising does, rather than in the gather, and ranks out of step - fail.""" + does not serialize fails as fn raising does, rather than in the gather; one only torch.save's + pickle protocol takes crosses, serialized once; and ranks out of step fail.""" cpu_dist_workers.run(_all_reduce_job) From f405e480b8ec6b93300ce7255f5b28f40d283475 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 17:35:54 +0000 Subject: [PATCH 33/48] Test that a collective's value the other ranks could not load fails on every rank The serialized value is loaded back before it is sent, so a value that pickles but cannot be loaded fails inside the guard, which sends the failure to every rank. Nothing tested that half: without the load-back the sending rank returned and only the receiving ranks raised. A value whose reduction raises on load now crosses both broadcast_from and all_reduce on two ranks. Signed-off-by: Shengliang Xu --- .../unit/torch/utils/test_distributed_run.py | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index d6a269bc443..dc5d7d5214b 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -243,6 +243,17 @@ def __reduce_ex__(self, protocol): return _Protocol2Only, (self.n,) +def _refuse_to_load(): + raise RuntimeError("this value cannot be loaded") + + +class _Unloadable: + """Pickles, and raises when loaded back.""" + + def __reduce__(self): + return _refuse_to_load, () + + def _broadcast_from_job(rank, size): ran = [] @@ -280,6 +291,12 @@ def exit_2(): assert "rank 0 failed to compute x: TypeError" in str(error), error # Serialized once, where the sending rank catches a failure, not again by the broadcast. assert broadcast_from("x", lambda: _Protocol2Only(7)).n == 7 + # Loaded back there too: a value the other ranks could not load fails on the sending rank. + _, error = _outcome(lambda: broadcast_from("x", _Unloadable)) + assert isinstance(error, RuntimeError), error + assert ("cannot be loaded" if rank == 0 else "rank 0 failed to compute x") in str(error), ( + error + ) sent, error = _outcome(lambda: broadcast_from("plan" if rank == 0 else "other", lambda: 7)) if rank == 0: assert (sent, error) == (7, None) @@ -292,8 +309,9 @@ def test_broadcast_from_gives_every_rank_the_sending_ranks_value(cpu_dist_worker """Over a real broadcast: any rank can send, and only it runs fn. A rank outside the run is refused, the sending rank's failure stops every rank and its exit exits every rank with its code, so that distributed_run tears every rank down alike. A value that does not serialize - fails as fn raising does, rather than in the broadcast; one only torch.save's pickle protocol - takes crosses, serialized once; and ranks out of step fail.""" + fails as fn raising does, rather than in the broadcast, and so does one the other ranks could + not load back; one only torch.save's pickle protocol takes crosses, serialized once; and ranks + out of step fail.""" cpu_dist_workers.run(_broadcast_from_job) @@ -351,6 +369,11 @@ def exit_on_rank_0_fail_on_rank_1(): "x", lambda: _Protocol2Only(rank), lambda values: [v.n for v in values] ) assert reduced == list(range(size)), reduced + _, error = _outcome(lambda: all_reduce("x", _Unloadable if rank == 1 else lambda: 4, min)) + assert isinstance(error, RuntimeError), error + assert ("cannot be loaded" if rank == 1 else "rank 1 failed to compute x") in str(error), ( + error + ) _, error = _outcome( lambda: all_reduce("smallest" if rank == 0 else "other", lambda: 4, min) ) @@ -362,8 +385,9 @@ def test_all_reduce_gives_every_rank_the_combined_value(cpu_dist_workers): """Over a real gather: every rank's value counts and every rank gets the same result. One rank's failure stops every rank and its exit exits every rank with its code. A rank's error outranks another's exit, so that no rank waits in the exit's barrier for one that aborted. A value that - does not serialize fails as fn raising does, rather than in the gather; one only torch.save's - pickle protocol takes crosses, serialized once; and ranks out of step fail.""" + does not serialize, or that the other ranks could not load back, fails as fn raising does, rather + than in the gather; one only torch.save's pickle protocol takes crosses, serialized once; and + ranks out of step fail.""" cpu_dist_workers.run(_all_reduce_job) From dcf56285c4f185e616171d36beda81fc3bc270f1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:10:59 +0000 Subject: [PATCH 34/48] Connect a bounded group on every rank before a job can destroy it gloo's new_group can return on one rank while a peer is still connecting to the group. A pooled job that runs no collective on its bounded group (one that mocks barrier, or one refused before it calibrates) destroys the group at once, and the peer's new_group then fails with "Gloo connectFullMesh failed ... Connection closed by peer". That is the "Worker(s) failed" flake of the bounded-barrier and expert-parallel calibration tests in CI. bounded_collectives now runs a barrier on the new group before the job's body, with the real barrier, which the job may have mocked by then. On a two-rank gloo pool under load (six job loops pinned to two cores), the bounded-barrier job failed 12 of 3000 times without the barrier and 0 of 3000 with it. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/distributed/utils.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index d29dd8c3408..42598f3d6e2 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -161,6 +161,10 @@ def _on_group(collective, bound, /, *args, **kwargs): return collective(*args, **{**kwargs, "group": bound if group is None else group}) +# The real barrier: a job may have replaced dist.barrier by the time it enters bounded_collectives. +_barrier = dist.barrier + + @contextmanager def bounded_collectives(seconds: float = 10): """Run the object collectives on a new group that gives up after ``seconds``, within the block. @@ -172,6 +176,10 @@ def bounded_collectives(seconds: float = 10): """ group = dist.new_group(timeout=timedelta(seconds=seconds)) try: + # gloo's new_group can return on one rank while a peer is still connecting to the group, and + # destroying it then, as a job that runs no collective on it does at once, fails the peer's + # new_group ("Connection closed by peer"). Past this barrier, every rank has connected. + _barrier(group=group) with ExitStack() as stack: for name in ("broadcast_object_list", "all_gather_object", "barrier"): collective = partial(_on_group, getattr(dist, name), group) From 1d3d06316b33637de98f05e77a276b74138c458a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:00:29 +0000 Subject: [PATCH 35/48] Say what the collectives' load-back check proves, and why it loads fully Loading the bytes back on the computing rank proves they load there: a value whose load depends on the loading rank, a tensor's, which lands on the device it was saved from, can still fail on a peer alone. The collectives carry small values, and the check holds one more copy of the value while it runs. The full load's reason is that the values are of any class, as through torch's object collectives. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index ed713a3cf0d..a71b00d81f6 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -440,11 +440,14 @@ def _received(sent: tuple[bool, Any], rank: int, label: str) -> BaseException: def _encoded(value: Any) -> bytes: """``value`` as :func:`broadcast_from` and :func:`all_reduce` send it: serialized once, here. - It is also loaded back here: a value that does not serialize, or that the other ranks could not - load (a class a forced weights-only load refuses, say), then fails inside the caller's guard, - which sends the failure to every rank, rather than inside the collective, on one rank alone. The - collective carries the bytes, which always pickle. torch.save stands in for plain pickle, which - bandit refuses under modelopt/. + It is also loaded back here: a value that does not serialize, or that every rank's load refuses + alike (a class a forced weights-only load rejects, say), then fails inside the caller's guard, + which sends the failure to every rank, rather than inside the collective, on one rank alone. A + load that depends on the loading rank can still fail on a peer alone, a tensor's, which lands on + the device it was saved from, say: these collectives carry small values, such as a plan, a flag + or a size. The check holds one more copy of the value while it runs. The collective carries the + bytes, which always pickle. torch.save stands in for plain pickle, which bandit refuses under + modelopt/. """ buffer = io.BytesIO() torch.save(value, buffer) @@ -456,9 +459,9 @@ def _encoded(value: Any) -> bytes: def _decoded(payload: bytes) -> Any: """The value :func:`_encoded` serialized into ``payload``. - A full load, not :func:`_deserialize`'s weights_only one: the peers are this job's own ranks, - as the object collectives' own unpickling would trust them, and a DistributionPlan crosses - without add_safe_globals. + A full load, not :func:`_deserialize`'s weights_only one: the values are whatever the callers + compute, of any class, and the peers are this job's own ranks, whose payloads torch's object + collectives unpickle just as fully. """ return torch.load(io.BytesIO(payload), weights_only=False) From ae4fd1bd31a235d8668c4ec55e29112ce668732c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:03:25 +0000 Subject: [PATCH 36/48] Give broadcast_from's sending rank the loaded copy too The sending rank returned the object fn built, still shared with whatever fn closed over, while every other rank got a copy loaded from the bytes sent: a caller mutating the result would change state on rank 0 alone. Every rank now returns the loaded copy, as all_reduce does. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 10 +++++----- tests/unit/torch/utils/test_distributed_run.py | 8 ++++++-- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index a71b00d81f6..0c184b2116c 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -474,6 +474,8 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: returns a value that does not serialize and load back (:func:`_encoded`), every rank raises instead of waiting in a broadcast ``rank`` never sends: ``rank`` its own error, the others a SystemExit with its code if that is what it was, so that every rank exits, else a RuntimeError. + Every rank, ``rank`` too, gets the value loaded from the bytes sent, so the result shares nothing + with what ``fn`` returned, on any rank, as with :func:`all_reduce`. A rank broadcasting another label than ``rank`` at the same point fails; ``rank`` sees no other rank's label, and returns. Collective: every rank calls it at the same point, with the same ``rank``. In a single process ``fn`` runs alone; a rank of a launch whose process group is not @@ -485,17 +487,15 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: if not is_distributed_run(): return fn() sending = torch.distributed.get_rank() == rank - value: Any = None payload: Any = None error = None if sending: try: - value = fn() - payload = _encoded(value) + payload = _encoded(fn()) except BaseException as e: # Still send: the other ranks would otherwise wait in the broadcast for this one. A # SystemExit or KeyboardInterrupt too, which argparse-style checks raise. - value, payload, error = None, None, e + payload, error = None, e # One message from the sending rank: its label, so a rank broadcasting another learns the two # are out of step, and its serialized value or its error. message = [(label, payload, _sendable(error))] @@ -510,7 +510,7 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: f"Ranks broadcast out of step: this rank broadcast {label!r} while rank {rank} sent " f"{sent_label!r}. Every rank must broadcast them in the same order." ) - return value if sending else _decoded(payload) + return _decoded(payload) def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _V]) -> _V: diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index dc5d7d5214b..0d39403af24 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -291,6 +291,10 @@ def exit_2(): assert "rank 0 failed to compute x: TypeError" in str(error), error # Serialized once, where the sending rank catches a failure, not again by the broadcast. assert broadcast_from("x", lambda: _Protocol2Only(7)).n == 7 + # Every rank gets the value loaded from the bytes sent, the sending rank too. + built: list = [] + got = broadcast_from("x", lambda: built) + assert got == [] and got is not built # Loaded back there too: a value the other ranks could not load fails on the sending rank. _, error = _outcome(lambda: broadcast_from("x", _Unloadable)) assert isinstance(error, RuntimeError), error @@ -310,8 +314,8 @@ def test_broadcast_from_gives_every_rank_the_sending_ranks_value(cpu_dist_worker refused, the sending rank's failure stops every rank and its exit exits every rank with its code, so that distributed_run tears every rank down alike. A value that does not serialize fails as fn raising does, rather than in the broadcast, and so does one the other ranks could - not load back; one only torch.save's pickle protocol takes crosses, serialized once; and ranks - out of step fail.""" + not load back; one only torch.save's pickle protocol takes crosses, serialized once; every rank, + the sending one too, gets a copy loaded from the bytes; and ranks out of step fail.""" cpu_dist_workers.run(_broadcast_from_job) From 18f8f8eb2b945b5dca987fccdfb797c4d3f04726 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:50:24 +0000 Subject: [PATCH 37/48] Say a single process gets fn's value as it is, unserialized broadcast_from's docstring promised every rank the value loaded from the bytes sent, and all_reduce's that it shares nothing with the callers', but with no distributed run both hand back fn's own value, neither serialized nor copied. The docstrings now say which path gives which. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 0c184b2116c..7add05e1be7 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -474,12 +474,12 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: returns a value that does not serialize and load back (:func:`_encoded`), every rank raises instead of waiting in a broadcast ``rank`` never sends: ``rank`` its own error, the others a SystemExit with its code if that is what it was, so that every rank exits, else a RuntimeError. - Every rank, ``rank`` too, gets the value loaded from the bytes sent, so the result shares nothing - with what ``fn`` returned, on any rank, as with :func:`all_reduce`. + On a distributed run every rank, ``rank`` too, gets the value loaded from the bytes sent, so the + result shares nothing with what ``fn`` returned, as with :func:`all_reduce`. A rank broadcasting another label than ``rank`` at the same point fails; ``rank`` sees no other rank's label, and returns. Collective: every rank calls it at the same point, with the same - ``rank``. In a single process ``fn`` runs alone; a rank of a launch whose process group is not - set up yet raises instead. + ``rank``. In a single process ``fn`` runs alone, and its value comes back as it is, neither + serialized nor copied; a rank of a launch whose process group is not set up yet raises instead. """ _refuse_a_launch_without_its_group(label) if not 0 <= rank < size(): @@ -523,8 +523,9 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ and every rank alike: if any rank's error is not a SystemExit, a rank with such an error raises its own and the others a RuntimeError naming the first one; else every rank exits, a rank that exited with its own SystemExit and the others with the first one's code. Ranks reducing - different labels at the same point fail. Collective; in a single process ``fn`` runs alone, and - a rank of a launch whose process group is not set up yet raises instead. + different labels at the same point fail. Collective. In a single process ``fn`` runs alone and + ``combine`` gets its value as it is, neither serialized nor copied; a rank of a launch whose + process group is not set up yet raises instead. It runs ``fn`` itself so that a rank whose ``fn`` fails still joins the gather, which :func:`allreduce`, handed values the ranks already hold, cannot do; :func:`broadcast_from` From f0153accc28e7418994a39013e59b75089bfd393 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 20:57:25 +0000 Subject: [PATCH 38/48] Say how distribution_plan reads a launch before its group, and what the load-back costs Before a launch sets its group up, each process runs alone with a model of its own, so a model with no plan reads as a single process's there; the load-back check in _encoded loads the value once more on the sending rank, so a load with side effects has them twice there. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7add05e1be7..7fa407c95f0 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -378,7 +378,9 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: A model with no plan recorded is a single process's, unless it shows otherwise: one that is FSDP2-sharded, or held in a run of several ranks, raises rather than be taken for one. That includes a module added to a model after its plan was recorded: a module does not know the - model it was added to, so read the plan off that model. + model it was added to, so read the plan off that model. A run is its live process group: before + a launch sets its group up, each of its processes runs alone with a model of its own, so a model + with no plan reads as a single process's there. """ plan = getattr(model, _PLAN_ATTR, None) if plan is not None: @@ -445,9 +447,10 @@ def _encoded(value: Any) -> bytes: which sends the failure to every rank, rather than inside the collective, on one rank alone. A load that depends on the loading rank can still fail on a peer alone, a tensor's, which lands on the device it was saved from, say: these collectives carry small values, such as a plan, a flag - or a size. The check holds one more copy of the value while it runs. The collective carries the - bytes, which always pickle. torch.save stands in for plain pickle, which bandit refuses under - modelopt/. + or a size. The check holds one more copy of the value while it runs, and loads it once more on + the sending rank: a value whose load has side effects has them twice there. The collective + carries the bytes, which always pickle. torch.save stands in for plain pickle, which bandit + refuses under modelopt/. """ buffer = io.BytesIO() torch.save(value, buffer) From 356bbcde4131a355d35b5b260b1be207f1218b93 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:43:04 +0000 Subject: [PATCH 39/48] Say which runs abort a rank whose body raised distributed_run aborts only in a run of several ranks on a group that was not there when it began; a run that found the caller's group leaves the error to the caller, which owns the peers. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7fa407c95f0..334692cfc19 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -603,9 +603,10 @@ def distributed_run( On exit, even if the body raises, it destroys the group the run set up: the launch's, or one the body set up, such as the one-rank group a single process's FSDP2 load shards through. A group that was already there when the run began is the caller's, and is left as it is. A rank whose - body raised in a run of several ranks exits through :func:`abort` instead, and the launcher - tears the job down: destroying the group would wait on peers that may sit in a collective this - rank never reaches. + body raised, in a run of several ranks on a group that was not there when it began, exits + through :func:`abort` instead, and the launcher tears the job down: destroying the group would + wait on peers that may sit in a collective this rank never reaches. In a run that found the + caller's group, the error goes on to the caller, which owns the peers. """ launched = is_distributed_launch() if launched and torch.device(device) != torch.device("cuda"): From 4a68b4a038d1b85ed17460bbd2a39e4814d7a6e0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:43:19 +0000 Subject: [PATCH 40/48] Say a job names its group by keyword under bounded_collectives _on_group reads the caller's group from the keyword only; a group passed by position would also get the bound one, and fail as a duplicate argument. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/distributed/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 42598f3d6e2..c38d51ff1c6 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -156,7 +156,10 @@ def reseeding_worker_teardown(rank, world_size): def _on_group(collective, bound, /, *args, **kwargs): """``collective`` on ``bound`` unless the caller named a group of its own: modelopt's - ``barrier()`` passes ``group=None`` explicitly, which would override a ``partial``'s keyword.""" + ``barrier()`` passes ``group=None`` explicitly, which would override a ``partial``'s keyword. + + A job names its group as ``group=``: one passed by position, as ``dist.barrier(group)`` can, is + not seen here and fails as a duplicate argument.""" group = kwargs.get("group") return collective(*args, **{**kwargs, "group": bound if group is None else group}) From df5f78a2316878ee5396ac4d375d6591ce2a41ad Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:43:26 +0000 Subject: [PATCH 41/48] Clear WORLD_SIZE in the single-process collectives test Both collectives read the launch's environment now, so an inherited WORLD_SIZE would make them refuse; the test says it runs alone, as its siblings do. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_distributed_run.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 0d39403af24..47d2058bc1a 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -206,6 +206,7 @@ def all_gather_object(gathered, mine, group=None): def test_a_single_process_computes_alone(monkeypatch): """Nothing to sync with: no collective.""" + monkeypatch.delenv("WORLD_SIZE", raising=False) for collective in ("broadcast_object_list", "all_gather_object"): monkeypatch.setattr( torch.distributed, collective, lambda *a, **k: pytest.fail("a collective ran") From f433158f2d6d5c797fe3fd20cb6e0a3578da7fa3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:01:22 +0000 Subject: [PATCH 42/48] Send a failure whose text cannot be formatted, rather than fail before the collective _sendable formatted the error outside the guarded block, so an exception whose __str__ raises, or a SystemExit whose code does, failed its rank before the broadcast or gather and left the other ranks waiting. Its text now falls back to a stand-in naming the type. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 15 +++++++++-- .../unit/torch/utils/test_distributed_run.py | 26 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 334692cfc19..d9c0d9e115b 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -414,8 +414,19 @@ def _sendable(error: BaseException | None) -> tuple[bool, Any] | None: return None if isinstance(error, SystemExit): code = error.code - return True, code if code is None or isinstance(code, int) else str(code) - return False, f"{type(error).__name__}: {error}" + return True, code if code is None or isinstance(code, int) else _text(code) + return False, f"{type(error).__name__}: {_text(error)}" + + +def _text(value: Any) -> str: + """``str(value)``, or a stand-in if that raises. + + Formatting a failure must not keep its rank from the collective that tells the others. + """ + try: + return str(value) + except Exception: + return f"" def _refuse_a_launch_without_its_group(label: str) -> None: diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 47d2058bc1a..fcf3301f916 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -255,6 +255,17 @@ def __reduce__(self): return _refuse_to_load, () +class _UnprintableError(Exception): + """An error whose text cannot be formatted.""" + + def __str__(self): + raise ValueError("no text") + + +def _raise_unprintable(): + raise _UnprintableError + + def _broadcast_from_job(rank, size): ran = [] @@ -296,6 +307,14 @@ def exit_2(): built: list = [] got = broadcast_from("x", lambda: built) assert got == [] and got is not built + # An error whose text cannot be formatted still reaches the broadcast. + _, error = _outcome(lambda: broadcast_from("x", _raise_unprintable)) + if rank == 0: + assert isinstance(error, _UnprintableError), type(error) + else: + assert "rank 0 failed to compute x: _UnprintableError: Date: Tue, 27 Oct 2026 22:01:43 +0000 Subject: [PATCH 43/48] Point a model without a plan at set_distribution_plan only load_pretrained records the plan only from the load routing on; until then the advice to load through it would leave a multi-rank model as planless as before. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index d9c0d9e115b..400bf59b68a 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -351,10 +351,10 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non """Record that ``model`` is distributed as ``plan``, on it and on every submodule it has now. Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed - the whole model; a module added afterwards records none. ``load_pretrained`` records the plan it - loads by once the model is built; code that distributes a model itself records it here before - calibrating or exporting the model. The plan says how the model is placed over the run; the - :class:`ParallelState` a quantized module holds names the process groups its collectives run on. + the whole model; a module added afterwards records none. Code that distributes a model records + its plan here before calibrating or exporting the model. The plan says how the model is placed + over the run; the :class:`ParallelState` a quantized module holds names the process groups its + collectives run on. A model is distributed once, and its plan recorded once: recording a plan on a model that, or any of whose submodules, already records one raises, the same plan included, and leaves the @@ -389,10 +389,9 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: if sharded or size() > 1: how = "is FSDP2-sharded" if sharded else f"is held in a run of {size()} ranks" raise RuntimeError( - f"{type(model).__name__} {how}, but has no distribution plan recorded. Load it with " - "load_pretrained, or record how it is distributed with set_distribution_plan(); a " - "module added to a model after its plan was recorded has none, so read the plan off " - "that model." + f"{type(model).__name__} {how}, but has no distribution plan recorded. Record how it " + "is distributed with set_distribution_plan(); a module added to a model after its plan " + "was recorded has none, so read the plan off that model." ) return DistributionPlan() From 8f722deee63109a2d7e3fc181bf739ed930cddaa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:01:57 +0000 Subject: [PATCH 44/48] Say inline why _decoded loads with weights_only=False, as SECURITY.md asks Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 400bf59b68a..528a23992bc 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -476,7 +476,10 @@ def _decoded(payload: bytes) -> Any: compute, of any class, and the peers are this job's own ranks, whose payloads torch's object collectives unpickle just as fully. """ - return torch.load(io.BytesIO(payload), weights_only=False) + return torch.load( + io.BytesIO(payload), + weights_only=False, # this job's own ranks sent the bytes over its process group: no file + ) def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: From 4520ab654fe5fc38b9513fe568cfe171381691ef Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:52:10 +0000 Subject: [PATCH 45/48] Refuse a plan recorded in a run of another number of ranks Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 10 ++++++++++ tests/unit/torch/utils/test_distributed_run.py | 10 +++++++--- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 528a23992bc..09d497f3c79 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -359,7 +359,17 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non A model is distributed once, and its plan recorded once: recording a plan on a model that, or any of whose submodules, already records one raises, the same plan included, and leaves the model as it was. No code reloads or re-distributes a model it loaded. + + In a run of several ranks, a plan for another number of ranks raises too, and records nothing: + a single process's plan there would be taken at its word, and the model replicated. A plan + recorded before the run's group is set up is not checked, nor one in a run of one, where a rank + can stand in for several. """ + if size() > 1 and plan.world_size != size(): + raise ValueError( + f"{plan!r} is for {plan.world_size} ranks, but the run it is recorded in has " + f"{size()}; a plan's world_size is the ranks that hold the model." + ) for name, module in model.named_modules(): recorded = getattr(module, _PLAN_ATTR, None) if recorded is not None: diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index fcf3301f916..126391ef5a8 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -138,7 +138,10 @@ def _outcome(call): def _no_plan_on_several_ranks_job(rank, size): - _, error = _outcome(lambda: distribution_plan(_model())) + model = _model() + _, error = _outcome(lambda: set_distribution_plan(model, DistributionPlan())) + assert isinstance(error, ValueError), error + _, error = _outcome(lambda: distribution_plan(model)) assert isinstance(error, RuntimeError), error assert f"is held in a run of {size} ranks" in str(error), error model = _model() @@ -158,8 +161,9 @@ def _no_plan_on_several_ranks_job(rank, size): def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers): """Held whole or FSDP2-sharded, rather than be taken for a single process's; once a plan is - recorded, that plan is what it reads. A module added after the plan was recorded has none, and - raises too, rather than read as a single process's.""" + recorded, that plan is what it reads. A single process's plan cannot be recorded there either, + and records nothing. A module added after the plan was recorded has none, and raises too, rather + than read as a single process's.""" cpu_dist_workers.run(_no_plan_on_several_ranks_job) From d2c8924721532db8be4d73499964813bf162b415 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:06:56 +0000 Subject: [PATCH 46/48] Keep the collective helpers' docstrings to their contracts The derivations behind them (why a SystemExit travels as its code, what the load-back costs and why a rank-dependent load can still fail on a peer alone) move to the PR description, as CONTRIBUTING.md asks of rationale. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 84 +++++++++++------------------ 1 file changed, 31 insertions(+), 53 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 09d497f3c79..b6b643857ba 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -410,14 +410,10 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: def _sendable(error: BaseException | None) -> tuple[bool, Any] | None: - """``error`` as this rank sends it to the others, ``None`` for none; see :func:`_received`. - - An error goes as its text: the other ranks only need to stop, and an exception is not - guaranteed to pickle. A SystemExit goes as its code, for every rank to exit with: - :func:`distributed_run` ends an exit with a barrier but an error with :func:`abort`, so a rank - exiting while its peers abort would wait in that barrier for ranks that are gone. A code other - than an int or ``None`` goes as its text, which the interpreter prints for it before exiting - with 1. + """``error`` as this rank sends it to the others (see :func:`_received`), ``None`` for none. + + An error goes as its text, since an exception need not pickle; a SystemExit as its code, for + every rank to exit with. """ if error is None: return None @@ -460,17 +456,11 @@ def _received(sent: tuple[bool, Any], rank: int, label: str) -> BaseException: def _encoded(value: Any) -> bytes: - """``value`` as :func:`broadcast_from` and :func:`all_reduce` send it: serialized once, here. - - It is also loaded back here: a value that does not serialize, or that every rank's load refuses - alike (a class a forced weights-only load rejects, say), then fails inside the caller's guard, - which sends the failure to every rank, rather than inside the collective, on one rank alone. A - load that depends on the loading rank can still fail on a peer alone, a tensor's, which lands on - the device it was saved from, say: these collectives carry small values, such as a plan, a flag - or a size. The check holds one more copy of the value while it runs, and loads it once more on - the sending rank: a value whose load has side effects has them twice there. The collective - carries the bytes, which always pickle. torch.save stands in for plain pickle, which bandit - refuses under modelopt/. + """``value`` serialized for :func:`broadcast_from` and :func:`all_reduce`, and loaded back once. + + The load-back makes a value that cannot round-trip fail inside the caller's guard, which tells + every rank, rather than inside the collective on one rank alone. torch.save stands in for + pickle, which bandit refuses under modelopt/. """ buffer = io.BytesIO() torch.save(value, buffer) @@ -493,19 +483,15 @@ def _decoded(payload: bytes) -> Any: def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: - """Return ``fn()`` run on ``rank`` only, on every rank. - - For a value that depends on what each rank reads for itself -- the environment, its GPU, files - on its node -- so that ranks computing it for themselves could disagree. If ``fn`` raises, or - returns a value that does not serialize and load back (:func:`_encoded`), every rank raises - instead of waiting in a broadcast ``rank`` never sends: ``rank`` its own error, the others a - SystemExit with its code if that is what it was, so that every rank exits, else a RuntimeError. - On a distributed run every rank, ``rank`` too, gets the value loaded from the bytes sent, so the - result shares nothing with what ``fn`` returned, as with :func:`all_reduce`. - A rank broadcasting another label than ``rank`` at the same point fails; ``rank`` sees no other - rank's label, and returns. Collective: every rank calls it at the same point, with the same - ``rank``. In a single process ``fn`` runs alone, and its value comes back as it is, neither - serialized nor copied; a rank of a launch whose process group is not set up yet raises instead. + """Return ``fn()``, run on ``rank`` only, on every rank. + + For a value ranks computing it for themselves could disagree on (their environment, GPU or + files). If ``fn`` raises or its value does not round-trip (:func:`_encoded`), every rank raises + rather than wait for the broadcast: ``rank`` its own error, the others a SystemExit with its + code if it exited, else a RuntimeError. Every rank, ``rank`` too, gets a copy loaded from the + bytes sent. Collective: every rank calls it at the same point with the same ``label`` and + ``rank``; a rank out of step raises. In a single process ``fn``'s value comes back as it is; a + rank of a launch whose process group is not set up raises. """ _refuse_a_launch_without_its_group(label) if not 0 <= rank < size(): @@ -540,22 +526,17 @@ def broadcast_from(label: str, fn: Callable[[], _V], rank: int = 0) -> _V: def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _V]) -> _V: - """Return ``combine`` of every rank's ``fn()``, on every rank. - - ``combine`` takes the ranks' values in rank order and returns the one every rank gets, e.g. - ``min`` for a batch size each rank's memory bounds, where one rank's choice could run another - out of memory. If ``fn`` raises on any rank, or returns a value that does not serialize and load - back (:func:`_encoded`), every rank raises instead of waiting in a gather that rank never joins, - and every rank alike: if any rank's error is not a SystemExit, a rank with such an error raises - its own and the others a RuntimeError naming the first one; else every rank exits, a rank that - exited with its own SystemExit and the others with the first one's code. Ranks reducing - different labels at the same point fail. Collective. In a single process ``fn`` runs alone and - ``combine`` gets its value as it is, neither serialized nor copied; a rank of a launch whose - process group is not set up yet raises instead. - - It runs ``fn`` itself so that a rank whose ``fn`` fails still joins the gather, which - :func:`allreduce`, handed values the ranks already hold, cannot do; :func:`broadcast_from` - differs from :func:`broadcast` and :func:`master_only` in the same way. + """Return ``combine`` of every rank's ``fn()``, in rank order, on every rank. + + E.g. ``combine=min`` for a batch size each rank's memory bounds. If ``fn`` fails on any rank, or + its value does not round-trip (:func:`_encoded`), every rank raises instead of waiting in the + gather: if any rank's failure is not a SystemExit, a rank with such an error raises it and the + others a RuntimeError naming the first; else each rank exits, with its own code or the first + one's. Collective, with the same ``label`` on every rank; ranks out of step raise. In a single + process ``combine`` gets ``fn()``'s value as it is; a rank of a launch whose process group is + not set up raises. + + Unlike :func:`allreduce`, it runs ``fn`` itself, so a rank whose ``fn`` fails still joins. """ _refuse_a_launch_without_its_group(label) if not is_distributed_run(): @@ -572,11 +553,8 @@ def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _ torch.distributed.all_gather_object(gathered, mine) failures = [(rank_, sent) for rank_, (_, _, sent) in enumerate(gathered) if sent is not None] if failures: - # Every rank raises by one rule, for distributed_run to end every rank alike. A rank whose - # error is not a SystemExit aborts, so if any rank's is not, every rank raises an error, an - # exiting rank too, rather than wait in the exit's barrier for the aborted one; else every - # rank exits. A failing rank raises its own exception when it is of that kind. Each sent is - # _sendable's (exited, detail). + # One rule on every rank: any error that is not an exit makes every rank raise one, for an + # exiting rank would wait in distributed_run's exit barrier for an aborted peer. errors = [(rank_, sent) for rank_, sent in failures if not sent[0]] failed_rank, sent = (errors or failures)[0] if error is not None and isinstance(error, SystemExit) == sent[0]: From 4cb61ecd59bb708920296a48498ae19c6a5ede4e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 01:06:55 +0000 Subject: [PATCH 47/48] Format distributed_run's own error report through _text too The report of the body's exception in distributed_run's finally formatted it with str(); an exception whose __str__ raises escaped there, skipping abort() and leaving the other ranks to wait out the group's timeout. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 2 +- tests/unit/torch/utils/test_distributed_run.py | 14 ++++++++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index b6b643857ba..aa0b2e31a76 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -642,7 +642,7 @@ def distributed_run( # Every rank reaches a SystemExit, an intentional exit: it is not an error to report. if exc is not None and not isinstance(exc, SystemExit): print( - f"[distributed_run][rank {rank()}] body raised {type(exc).__name__}: {exc}", + f"[distributed_run][rank {rank()}] body raised {type(exc).__name__}: {_text(exc)}", flush=True, ) if owns_group and size() > 1: diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 126391ef5a8..ff9ba4729b0 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -512,11 +512,19 @@ class _ExitedError(Exception): [ (True, None, ["barrier", "destroy"]), (True, ValueError, ["abort"]), + (True, _UnprintableError, ["abort"]), (True, SystemExit, ["barrier", "destroy"]), (False, None, ["barrier", "destroy"]), (False, ValueError, ["destroy"]), ], - ids=["launch", "launch-raises", "launch-exits", "single-process", "single-process-raises"], + ids=[ + "launch", + "launch-raises", + "launch-raises-unprintable", + "launch-exits", + "single-process", + "single-process-raises", + ], ) def test_the_run_ends_the_group_it_set_up( monkeypatch, capsys, launch, body_raises, ends_by, in_except @@ -567,7 +575,9 @@ def run(): # An intentional exit is not reported as an error. An error is reported once, under this rank's # place in the group; its traceback is abort()'s to print, or goes on with it to the caller. out, err = capsys.readouterr() - assert ("[distributed_run][rank 0] body raised" in out) == (body_raises is ValueError) + assert ("[distributed_run][rank 0] body raised" in out) == ( + body_raises not in (None, SystemExit) + ) assert "Traceback" not in err From 73172fbec6a5c12bca253823f6b52c869da34832 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 08:00:02 +0000 Subject: [PATCH 48/48] Say distributed_run's exit path assumes every rank exits A SystemExit out of the body is not reported: the group is destroyed as on a normal exit, whose barrier waits on every rank. The collectives carry one rank's exit to all of them; a rank that exits alone waits in that barrier until the group times out. The docstring now states that precondition. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index aa0b2e31a76..136b8d5df99 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -608,6 +608,10 @@ def distributed_run( through :func:`abort` instead, and the launcher tears the job down: destroying the group would wait on peers that may sit in a collective this rank never reaches. In a run that found the caller's group, the error goes on to the caller, which owns the peers. + + A ``SystemExit`` is not an error: the group is destroyed as on a normal exit, which waits on + every rank, so the exit has to be every rank's, as one raised inside :func:`broadcast_from` or + :func:`all_reduce` is. A rank that exits alone waits in that teardown until the group times out. """ launched = is_distributed_launch() if launched and torch.device(device) != torch.device("cuda"):