diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 5627754cb06..26892ef112d 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -22,9 +22,10 @@ import time import traceback from collections.abc import Callable, Iterator -from contextlib import suppress +from contextlib import contextmanager, suppress +from dataclasses import dataclass from datetime import timedelta -from typing import Any +from typing import Any, Literal, TypeVar from warnings import warn import torch @@ -34,17 +35,26 @@ __all__ = [ "DistributedProcessGroup", + "DistributedState", + "DistributionPlan", "ParallelState", + "all_reduce", "backend", "barrier", + "broadcast_from", + "distributed_run", + "distribution_plan", "fsdp2_wrap", "gather_bytes", "gather_state_dicts", "is_available", + "is_distributed_launch", + "is_distributed_run", "is_fsdp2_model", "is_initialized", "is_master", "rank", + "set_distribution_plan", "size", ] @@ -98,6 +108,25 @@ 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, 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 + + +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) @@ -106,13 +135,13 @@ def _serialize(obj: Any) -> torch.Tensor: return tensor -def _deserialize(tensor: torch.Tensor, size: int | None = None) -> Any: +def _deserialize(tensor: torch.Tensor, size: int | None = None, map_location=None) -> Any: buffer = tensor.numpy().tobytes() if size is not None: buffer = buffer[:size] # Explicit so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD cannot re-enable pickle for peer payloads. Custom # classes sent through broadcast/allgather must be registered with add_safe_globals. - return torch.load(io.BytesIO(buffer), weights_only=True) + return torch.load(io.BytesIO(buffer), map_location=map_location, weights_only=True) def _broadcast(tensor: torch.Tensor, src: int = 0, group=None) -> None: @@ -269,6 +298,374 @@ def setup(timeout: timedelta | None = None): torch.distributed.init_process_group("cpu:gloo,cuda:nccl", timeout=timeout) +@dataclass(frozen=True) +class DistributionPlan: + """How a model is distributed over the ranks of a run. + + * ``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 + shard is replicated ``edp = world_size / ep_size`` times. + """ + + 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.""" + 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"): + raise ValueError(f"dense={self.dense!r}: expected 'replicate' or 'fsdp'.") + + @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 FSDP2-sharded rather than replicated.""" + return self.dense == "fsdp" + + +# A plan is a value the collectives carry, and they load what the ranks send with weights_only. A +# loaded copy skips __post_init__: it is the plan the sending rank built, and checked there. +torch.serialization.add_safe_globals([DistributionPlan]) + +_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 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. 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 + 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: + where = f"{type(module).__name__} {name!r}" if name else type(module).__name__ + raise ValueError( + 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) + + +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. 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. 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: + 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. 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() + + +_V = TypeVar("_V") + + +def _sendable(error: BaseException | None) -> tuple[bool, Any] | None: + """``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 + if isinstance(error, SystemExit): + code = error.code + 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: + """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], src: int, label: str) -> BaseException: + """The exception to raise for what rank ``src`` sent (:func:`_sendable`) on failing ``label``.""" + exited, detail = sent + if exited: + return SystemExit(detail) + return RuntimeError(f"rank {src} failed to compute {label}: {detail}") + + +def _encoded(value: Any) -> torch.Tensor: + """``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. It loads on the producing rank + only: a load that depends on the loading rank (a class only the producer can import) can still + fail on a peer alone. + """ + payload = _serialize(value) + _decoded(payload) + return payload + + +def _decoded(payload: torch.Tensor) -> Any: + # On the CPU on every rank, the producing one too: a tensor's saved device is the producer's. + return _deserialize(payload, map_location="cpu") + + +def broadcast_from(label: str, fn: Callable[[], _V], src: int = 0) -> _V: + """Return ``fn()``, run on ``src`` 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: ``src`` its own error, the others a SystemExit with its + code if it exited, else a RuntimeError. Every rank, ``src`` too, gets a copy loaded from the + bytes sent, with any tensor in it on the CPU; the value must load with ``weights_only``, so a + class in it must be registered with ``torch.serialization.add_safe_globals``. Collective: every + rank calls it at the same point with the same ``label`` and ``src``; a receiving rank whose + ``label`` differs raises, but ``src`` itself cannot tell and returns its value. 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 <= src < size(): + raise ValueError(f"src={src} is not a rank of this {size()}-rank run.") + if not is_distributed_run(): + return fn() + sending = rank() == src + payload: Any = None + error = None + if sending: + try: + 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. + 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. The object collective pickles the + # message itself between the job's ranks; weights_only bounds the value inside it. + message = [(label, payload, _sendable(error))] + torch.distributed.broadcast_object_list(message, group_src=src) + sent_label, payload, sent_error = message[0] + if error is not None: + raise error + if sent_error is not None: + raise _received(sent_error, src, label) + if sent_label != label: + raise RuntimeError( + f"Ranks broadcast out of step: this rank broadcast {label!r} while rank {src} sent " + f"{sent_label!r}. Every rank must broadcast them in the same order." + ) + return _decoded(payload) + + +def all_reduce(label: str, fn: Callable[[], _V], combine: Callable[[list[_V]], _V]) -> _V: + """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. ``combine`` + gets copies loaded from the bytes the ranks sent, with any tensor in them on the CPU; in a + single process it gets ``fn()``'s value as it is. The values must load with ``weights_only``, so + a class in them must be registered with ``torch.serialization.add_safe_globals``. 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(): + return combine([fn()]) + payload, error = None, None + try: + 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. + 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] + if failures: + # 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]: + raise error + raise _received(sent, failed_rank, label) + 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." + ) + # 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) +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 + + @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, 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 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 + process, and this rank's place in a group the caller set up without the launch's environment. + A ``timeout`` cannot apply to a group of several ranks the caller set up, so passing one then + warns. A single process leaves it unused, without a warning, also in a one-rank group of the + caller's: no rank there waits on a peer, so an entry point may pass one for its launches only. + + 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 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. + + 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"): + raise ValueError( + f"distributed_run(device={device!r}): a distributed launch runs each rank on its own " + "CUDA device, chosen from LOCAL_RANK; pass device='cuda' or run a single process." + ) + group_was_there = is_initialized() + if timeout is not None and is_distributed_run(): + warn( + f"distributed_run(timeout={timeout}) does not apply: the process group was set up before " + "the run, and a group's timeout is fixed when it is set up. Pass it where the group is." + ) + 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, + # 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 torch.device(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. 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 {rank()}] body raised {type(exc).__name__}: {_text(exc)}", + flush=True, + ) + if owns_group and size() > 1: + abort() + if owns_group: + _destroy_group(exc) + + def cleanup(): """Cleans up the distributed environment. @@ -277,8 +674,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/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 45b59915351..c38d51ff1c6 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,46 @@ def reseeding_worker_teardown(rank, world_size): torch.manual_seed(1234) +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. + + 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}) + + +# 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. + + 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)) + 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) + 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: """Persistent worker pool that keeps distributed processes alive across multiple test dispatches. diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py index f88377e2a5f..25b3bc12268 100644 --- a/tests/unit/torch/utils/test_distributed.py +++ b/tests/unit/torch/utils/test_distributed.py @@ -13,17 +13,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pure-function tests for ``modelopt.torch.utils.distributed``.""" +"""Tests for ``modelopt.torch.utils.distributed``: its pure functions, the distribution plan and a +model's record of it, the collectives that fail together, and the process group's lifetime.""" import os import pickle +from contextlib import nullcontext +from datetime import timedelta import pytest import torch +import torch.distributed import torch.nn as nn +from _test_utils.torch.distributed.utils import bounded_collectives +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.nas.hparams.container import DepthHparam -from modelopt.torch.utils.distributed import _deserialize, _off_dtype_params, _serialize +from modelopt.torch.utils.distributed import ( + DistributedState, + DistributionPlan, + _deserialize, + _off_dtype_params, + _serialize, + all_reduce, + broadcast_from, + distributed_run, + distribution_plan, + set_distribution_plan, +) def _model(*sizes_and_dtypes) -> nn.Module: @@ -86,3 +105,439 @@ def test_deserialize_traced_hp_subclass(): hp = _deserialize(_serialize(DepthHparam([1, 2]))) assert type(hp) is DepthHparam assert hp.choices == [1, 2] + + +# --------------------------------------------------------------------------------- DistributionPlan + + +@pytest.mark.parametrize( + ("degrees", "match"), + [ + ({"world_size": 8, "ep_size": 3}, "must divide world_size"), + ({"world_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", "ep-negative", "unknown-dense"], +) +def test_an_invalid_plan_is_rejected(degrees, match): + """A plan no load could carry out fails at construction.""" + with pytest.raises(ValueError, match=match): + DistributionPlan(**degrees) + + +# ------------------------------------------------------------------- a model's recorded plan + + +def _plan_model(): + return nn.Sequential(nn.Linear(2, 2), nn.Sequential(nn.Linear(2, 2))) + + +@pytest.mark.parametrize( + ("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(recorded_on, again): + """Every submodule reads the recorded plan, and a second record is refused, recording nothing.""" + model, plan = _plan_model(), DistributionPlan(world_size=8, ep_size=8) + assert distribution_plan(model) == DistributionPlan() # no plan: a single process + holder = model if recorded_on == "model" else model[1] + set_distribution_plan(holder, plan) + holders = list(holder.modules()) + assert all(distribution_plan(module) is plan for module in holders) + + with pytest.raises(ValueError, match="its plan recorded once"): + set_distribution_plan(model, again) + for module in model.modules(): + assert distribution_plan(module) == (plan if module in holders else DistributionPlan()) + + +def _outcome(call): + """``(value, None)``, or ``(None, error)`` if ``call()`` raises, a ``SystemExit`` included.""" + try: + return call(), None + except (Exception, SystemExit) as error: + return None, error + + +def _no_plan_on_several_ranks_job(rank, size): + model = _plan_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 = _plan_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 + model = _plan_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): + """On two ranks a model needs a recorded plan, whole, FSDP2-sharded or added to later.""" + cpu_dist_workers.run(_no_plan_on_several_ranks_job) + + +# ------------------------------------------------------------------------------------ collectives + + +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") + ) + assert broadcast_from("x", lambda: 3) == 3 + assert all_reduce("smallest", lambda: 4, min) == 4 + with pytest.raises(ValueError, match="src=1"): + broadcast_from("x", lambda: 3, src=1) + + +def test_a_launch_without_its_group_refuses_the_collectives(monkeypatch): + """A rank of a launch whose group is not set up raises instead of computing its own value.""" + 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 == [] + + +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,) + + +# The collectives load what the ranks send with weights_only. +torch.serialization.add_safe_globals([_Protocol2Only]) + + +def _refuse_to_load(): + raise RuntimeError("this value cannot be loaded") + + +class _Unloadable: + """Pickles, and is refused when loaded back: weights_only runs no callable it was not given.""" + + 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 + + +@pytest.mark.parametrize( + ("value", "error"), + [((n for n in range(1)), TypeError), (_Unloadable(), pickle.UnpicklingError)], + ids=["does-not-pickle", "does-not-load"], +) +def test_the_collectives_encoding_refuses_a_value_that_does_not_round_trip(value, error): + """Both collectives encode through it, inside the guard that tells every rank.""" + with pytest.raises(error): + dist._encoded(value) + + +def test_the_collectives_carry_a_plan(): + """The one class they carry is registered, so it loads back under weights_only.""" + plan = DistributionPlan(world_size=8, ep_size=4, dense="fsdp") + assert dist._decoded(dist._encoded(plan)) == plan + + +def _broadcast_from_job(rank, size): + ran = [] + + def compute(): + ran.append(rank) + return f"rank {rank}'s" + + def fail(): + raise ValueError("bad degrees") + + def exit_2(): + raise SystemExit(2) + + with bounded_collectives(): + for sender in range(size): + ran.clear() + assert broadcast_from("x", compute, src=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, src=size)) + assert isinstance(error, ValueError), error + assert f"src={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 + _, error = _outcome(lambda: broadcast_from("x", exit_2)) + assert isinstance(error, SystemExit) and error.code == 2, 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 + # A value that does not round-trip fails on the sending rank, and reaches every rank. + _, error = _outcome(lambda: broadcast_from("x", _Unloadable)) + if rank == 0: + assert isinstance(error, pickle.UnpicklingError), error + else: + 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, on two ranks: failures, exits and ranks out of step reach every rank.""" + cpu_dist_workers.run(_broadcast_from_job) + + +def _all_reduce_job(rank, size): + def fail_on_rank_1(): + if rank == 1: + raise ValueError("no memory reading") + return 4 + + def exit_on_rank_1(): + if 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", _raise_unprintable if rank == 1 else lambda: 4, min) + ) + if rank == 1: + assert isinstance(error, _UnprintableError), type(error) + else: + assert "rank 1 failed to compute smallest: _UnprintableError" 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( + 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) + ) + assert isinstance(error, RuntimeError), error + assert "out of step" in str(error), error + + +def test_all_reduce_gives_every_rank_the_combined_value(cpu_dist_workers): + """Over a real gather, on two ranks: one rank's failure or exit, and ranks out of step, reach all.""" + cpu_dist_workers.run(_all_reduce_job) + + +# ---------------------------------------------------------------------------------- process group + + +@pytest.mark.parametrize("callers_group", [False, True], ids=["no-group", "callers-one-rank-group"]) +def test_a_single_process_is_rank_0_of_one_on_the_given_device(monkeypatch, recwarn, callers_group): + """No ranks to sync with: no group set up or torn down, and a timeout taken quietly.""" + monkeypatch.delenv("WORLD_SIZE", raising=False) + if callers_group: + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "rank", lambda group=None: 0) + monkeypatch.setattr(dist, "size", lambda group=None: 1) + monkeypatch.setattr(dist, "setup", lambda **kwargs: pytest.fail("set a group up")) + monkeypatch.setattr(dist, "_destroy_group", lambda exc: pytest.fail("tore a group down")) + with distributed_run(timeout=timedelta(hours=3), device="cpu") as state: + assert state == DistributedState(rank=0, world_size=1, device=torch.device("cpu")) + assert state.is_main + assert not [w for w in recwarn if "does not apply" in str(w.message)] + + +@pytest.mark.parametrize("device", ["cpu", "cuda:3"]) +def test_a_distributed_launch_refuses_another_device(monkeypatch, device): + """A launch runs each rank on the CUDA device LOCAL_RANK names, and refuses any other.""" + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setattr(dist, "setup", lambda **kwargs: pytest.fail("set a group up")) + with pytest.raises(ValueError, match=f"device='{device}'"), distributed_run(device=device): + pass + + +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", "in_except"), + [ + (True, None, ["barrier", "destroy"], False), + (True, ValueError, ["abort"], False), + (True, _UnprintableError, ["abort"], False), + (True, SystemExit, ["barrier", "destroy"], False), + (False, None, ["barrier", "destroy"], False), + (False, ValueError, ["destroy"], False), + (True, None, ["barrier", "destroy"], True), + (False, ValueError, ["destroy"], True), + ], + ids=[ + "launch", + "launch-raises", + "launch-raises-unprintable", + "launch-exits", + "single-process", + "single-process-raises", + "launch-entered-in-except", + "single-process-raises-entered-in-except", + ], +) +def test_the_run_ends_the_group_it_set_up( + monkeypatch, capsys, launch, body_raises, ends_by, in_except +): + """The group the run set up ends with a barrier and destroy, or with abort() if a rank raised.""" + # Under a launch, rank 2 of 4 on local rank 1: the device follows the local rank, not the global. + this_rank, world = (2, 4) if launch else (0, 1) + if launch: + monkeypatch.setenv("WORLD_SIZE", str(world)) + else: + monkeypatch.delenv("WORLD_SIZE", raising=False) + live, ended, timeouts = [], [], [] + + def setup(timeout=None): + timeouts.append(timeout) + live.append("the launch's group") + + monkeypatch.setattr(dist, "setup", setup) + monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) + monkeypatch.setattr(dist, "rank", lambda group=None: this_rank) + monkeypatch.setattr(dist, "size", lambda group=None: world) + monkeypatch.setattr(dist, "local_rank", lambda: 1) + 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 + expected = pytest.raises(escapes) if escapes else nullcontext() + + states = [] + + def run(): + with ( + expected, + distributed_run( + timeout=timedelta(hours=1), device="cuda" if launch else "cpu" + ) as state, + ): + states.append(state) + 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 + assert timeouts == ([timedelta(hours=1)] if launch else []) + device = torch.device("cuda:1" if launch else "cpu") + assert states == [DistributedState(rank=this_rank, world_size=world, device=device)] + assert states[0].is_main is not launch + # 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 (f"[distributed_run][rank {this_rank}] body raised" in out) == ( + body_raises not in (None, SystemExit) + ) + assert "Traceback" not in err + + +@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, recwarn, world_size, raises): + """A group the caller set up stays, the state is the rank's place in it, and a timeout warns.""" + bound: list[int] = [] + if world_size is None: + monkeypatch.delenv("WORLD_SIZE", raising=False) + else: + monkeypatch.setenv("WORLD_SIZE", world_size) + 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) + 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( + timeout=timedelta(hours=1), device="cpu" if world_size is None else "cuda" + ) as state, + ): + assert (state.rank, state.world_size) == (1, 2) + assert 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") + assert "timeout=1:00:00) does not apply" in str(recwarn.pop(UserWarning).message)