From d078d94ae41e7215851852d732a5afb59fd35297 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 11:10:28 +0000 Subject: [PATCH 1/7] Add the pieces of a distributed HF checkpoint writer Exporting an FSDP2 model gathers every shard to rank 0 first, so rank 0 needs host RAM for the whole model and the gather costs world x the bytes, and an expert-parallel model cannot go through that path at all. These are the pieces of a writer where every rank writes its own share instead (modelopt/torch/export/distributed_writer.py): - _materialize_owned gives each tensor exactly one owner rank and makes it whole there: replicated and plain values need no transport, and dim-0 shards travel point-to-point, batched and bucketed per peer -- one bucket per key deadlocked at 16 ranks. Placements it cannot assemble, and shards that do not tile the global shape, raise rather than risk a silently mis-assembled checkpoint. - _revert_whole_sd runs the reverse weight conversion on those whole tensors, which transformers' conversion ops need: they cannot express a placement. - _write_even_shards bin-packs each rank's set into even files numbered into one global range, writes them with safetensors, and unions the index on rank 0, failing on a name two ranks both wrote. A tensor over max_shard_size gets a file of its own, as in HF, and the rest still pack. The entry point that drives them follows. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 443 ++++++++++++++++++ .../torch/export/test_distributed_writer.py | 361 ++++++++++++++ 2 files changed, 804 insertions(+) create mode 100644 modelopt/torch/export/distributed_writer.py create mode 100644 tests/unit/torch/export/test_distributed_writer.py diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py new file mode 100644 index 00000000000..9c029ecec62 --- /dev/null +++ b/modelopt/torch/export/distributed_writer.py @@ -0,0 +1,443 @@ +# 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. + +"""Pieces of a no-gather distributed HF checkpoint writer. + +Every tensor is sent to exactly one owner rank over batched point-to-point sends and +written from there, so no rank ever materializes the whole model and no unit is +replicated across ranks. +""" + +import json +import warnings +from pathlib import Path + +import torch +import torch.nn as nn + + +def _size_to_bytes(size: "str | int") -> int: + """Parse an HF-style shard-size string (``"5GB"``, ``"500MB"``, ``"1GiB"``) to bytes. + + Matches transformers' decimal convention (GB == 10**9). Bare ints pass through. + """ + if isinstance(size, int): + return size + s = str(size).strip().upper() + units = { + "KIB": 2**10, "MIB": 2**20, "GIB": 2**30, "TIB": 2**40, + "KB": 10**3, "MB": 10**6, "GB": 10**9, "TB": 10**12, + } # fmt: skip + for unit in ("KIB", "MIB", "GIB", "TIB", "KB", "MB", "GB", "TB"): + if s.endswith(unit): + return int(float(s[: -len(unit)]) * units[unit]) + return int(float(s)) + + +def _even_bins(keys: list, sizes: dict, n: int, max_bytes: int) -> list: + """Partition ``keys`` into ~``n`` EVEN-sized bins (by bytes), each <= ``max_bytes`` where possible. + + Longest-processing-time greedy (largest item -> lightest bin) balances the bins, so N files come out + ~``total/N`` each rather than (N-1) full files + a small remainder -- ``max_shard_size`` is an UPPER + BOUND, not the target. ``n`` counts the bins for the tensors within the bound, and is bumped (up to + one file per tensor) if a lumpy item pushes a bin over it. A tensor larger than ``max_bytes`` gets + a file of its own on top, as in HF, and the rest still pack around it. + Returns the non-empty bins (``list[list[key]]``). + """ + oversized = [[k] for k in keys if sizes[k] > max_bytes] + fits = [k for k in keys if sizes[k] <= max_bytes] + n = max(1, min(n, len(fits))) if fits else 1 + ordered = sorted(fits, key=lambda k: -sizes[k]) + while True: + bins: list = [[] for _ in range(n)] + loads = [0] * n + for k in ordered: + j = min(range(n), key=lambda i: loads[i]) + bins[j].append(k) + loads[j] += sizes[k] + if n >= len(fits) or max(loads) <= max_bytes: + return oversized + [b for b in bins if b] + n += 1 + + +def _nbytes(v) -> int: + """Byte size. DTensor.numel() is GLOBAL, so this is consistent across ranks.""" + return int(v.numel() * v.element_size()) + + +def _is_full_world_shard0(v, world: int, _mesh_memo: dict | None = None) -> bool: + """True when ``v`` is sharded on dim 0 over a 1-D mesh that is exactly global ranks ``0..world-1``. + + Restricting the point-to-point fast path to that case keeps the peer in ``P2POp`` unambiguous -- + a plain global rank on the default process group, with no mesh-local/global translation to get + wrong. Any other layout is rejected by :func:`_materialize_owned` rather than routed through a + slower path, so no untested transport can silently produce a mis-assembled tensor. + + ``_mesh_memo`` caches the mesh-layout answer by mesh identity. Reading ``mesh.mesh`` back to the + host is a device sync, and FSDP2 gives every parameter the SAME mesh object, so without the memo + a few hundred keys cost a few hundred redundant syncs. + """ + if len(v.placements) != 1 or not v.placements[0].is_shard(0): + return False + mesh = v.device_mesh + if mesh.ndim != 1: + return False + if _mesh_memo is None: + return mesh.mesh.flatten().tolist() == list(range(world)) + key = id(mesh) + if key not in _mesh_memo: + _mesh_memo[key] = mesh.mesh.flatten().tolist() == list(range(world)) + return _mesh_memo[key] + + +def _gather_shard0_to_owners( + local_sd: dict, keys: list, owner: dict, all_shapes: list, rank: int, world: int +) -> None: + """Gather each dim-0-sharded DTensor to its owner and drop it elsewhere, in place. + + ``full_tensor()`` is an ALL-gather: it lands a full copy on every rank, so moving N bytes costs + ``world x N`` and every rank allocates the whole tensor even though one rank keeps it. Here each + shard travels once, to its owner only. + + Transfers are bucketed BY PEER rather than per key: everything this rank owes a given peer is + concatenated into one flat buffer per (peer, dtype) and sent as a single operation. One op per + (key, peer) pair does not survive leaving NVLink -- at 235B that was ~12.7k operations and took + 660s, with the payload only a few GB. Bucketing makes it O(world x dtypes) operations. + + ``all_shapes[r][k]`` is rank ``r``'s local dim-0 extent for key ``k``; the caller has already + checked that the shards tile the global tensor exactly. + + Deadlock safety: both sides derive their buckets from the same globally-sorted key list and the + same owner map, so every send has a matching receive of exactly the same length. + """ + import torch.distributed as dist + + def _numel(shape): + n = 1 + for s in shape: + n *= s + return n + + device = local_sd[keys[0]].to_local().device + + # Chunk the key list so the in-flight receive buffers stay bounded. Within a chunk the op count + # is O(world x dtypes) regardless of how many keys it holds, so chunks can be large. + chunk = 4096 + for base in range(0, len(keys), chunk): + part = keys[base : base + chunk] + + # (peer, dtype) -> keys. A rank sends its shard of every key the peer owns; the owner + # receives, from every other rank, its shard of every key this rank owns. Both sides build + # the identical grouping from shared state. Empty shards are skipped on both sides. + send_groups: dict = {} + recv_groups: dict = {} + for k in part: + dt = local_sd[k].dtype + o = owner[k] + if o != rank: + if _numel(all_shapes[rank][k]): + send_groups.setdefault((o, dt), []).append(k) + else: + for r in range(world): + if r != rank and _numel(all_shapes[r][k]): + recv_groups.setdefault((r, dt), []).append(k) + + ops: list = [] + keepalive: list = [] + pending: list = [] + for peer, dt in sorted(send_groups, key=lambda kd: (kd[0], str(kd[1]))): + ks = sorted(send_groups[(peer, dt)]) + buf = torch.cat([local_sd[k].to_local().reshape(-1) for k in ks]) + keepalive.append(buf) # must outlive the wait + ops.append(dist.P2POp(dist.isend, buf, peer)) + for srcr, dt in sorted(recv_groups, key=lambda kd: (kd[0], str(kd[1]))): + ks = sorted(recv_groups[(srcr, dt)]) + total = sum(_numel(all_shapes[srcr][k]) for k in ks) + buf = torch.empty(total, dtype=dt, device=device) + pending.append((srcr, ks, buf)) + ops.append(dist.P2POp(dist.irecv, buf, srcr)) + + if ops: + for req in dist.batch_isend_irecv(ops): + req.wait() + + # Unpack each peer's flat buffer back into per-key views, in the same order it was packed. + parts_by_key: dict = {} + for srcr, ks, buf in pending: + off = 0 + for k in ks: + shape = all_shapes[srcr][k] + n = _numel(shape) + parts_by_key.setdefault(k, {})[srcr] = buf[off : off + n].view(shape) + off += n + + # Rank order is mesh order (checked by _is_full_world_shard0), so concatenating the parts in + # rank order reproduces exactly what full_tensor() would have built. + for k in part: + if owner[k] != rank: + del local_sd[k] + continue + got = parts_by_key.get(k, {}) + pieces = [] + for r in range(world): + if r == rank: + pieces.append(local_sd[k].to_local()) + elif r in got: + pieces.append(got[r]) + local_sd[k] = torch.cat(pieces, dim=0) + del keepalive, pending, parts_by_key + + +def _materialize_owned(local_sd: dict, expert_key_set: set, rank: int, world: int) -> None: + """Make every tensor in ``local_sd`` WHOLE and owned by exactly one rank, in place. + + Experts are already disjoint per rank (each rank split its own shard). Every dense fqn is + assigned to ONE rank -- largest-first onto the currently-lightest rank, seeded with each rank's + expert bytes so the per-rank totals level out -- and then moved there by the cheapest transport + its layout allows (see :func:`_gather_shard0_to_owners`). Keys this rank does not own are dropped. + + This is the step that lets the reverse conversion run: afterwards nothing is a DTensor and no + tensor is partial, so transformers' conversion ops see the whole tensors they were written for. + It is NOT a gather to a single rank -- peak memory is this rank's owned set, not the model. + """ + import torch.distributed as dist + from torch.distributed.tensor import DTensor + + expert_keys = sorted(k for k in local_sd if k in expert_key_set) + dense_keys = sorted(k for k in local_sd if k not in expert_key_set) + _mesh_memo: dict = {} + + # Classify FIRST: placements are local DTensor metadata, so this needs no communication, and + # doing it up front means the shard shapes can ride along in the single metadata collective + # below. Every rank sees the same placements and so agrees on the buckets. + replicated, shard0, unsupported, plain, empty0 = [], [], [], [], [] + for k in dense_keys: + v = local_sd[k] + if not isinstance(v, DTensor): + plain.append(k) + elif v.numel() == 0: + # A zero-element global tensor (e.g. a projection that is absent in some layers, so its + # dim 0 is 0) cannot be reassembled from shards: FSDP2 still hands one rank a PADDED + # 1-row chunk while the rest get genuinely empty ones, so concatenating would fabricate + # a row the global tensor does not have. There is no data to move -- the owner emits the + # empty tensor directly. + empty0.append(k) + elif all(p.is_replicate() for p in v.placements): + replicated.append(k) + elif _is_full_world_shard0(v, world, _mesh_memo): + shard0.append(k) + else: + unsupported.append((k, str(v.placements), v.device_mesh.ndim)) + + # Reject rather than silently taking a slower, untested route. The point-to-point transport + # covers replicated, plain, and dim-0 shards on a 1-D full-world mesh; anything else (2-D mesh, + # non-dim-0 shard, partial placement) has no correct handling here. Raising names the offending + # keys, which is recoverable; emitting a mis-assembled tensor would not be. + if unsupported: + raise RuntimeError( + f"rank {rank}: {len(unsupported)} tensor(s) have a placement the no-gather export " + f"cannot materialize; first few: " + f"{[(k, pl, f'mesh.ndim={nd}') for k, pl, nd in unsupported[:3]]}" + ) + + # ONE collective for all the metadata: the key list (to catch divergence), this rank's expert + # byte count (to seed load balancing), and its local dim-0 extents (to size receive buffers). + # These were three separate all_gather_object round trips, each pickling a ~900-entry payload. + my_meta = ( + dense_keys, + sum(_nbytes(local_sd[k]) for k in expert_keys), + {k: tuple(local_sd[k].to_local().shape) for k in shard0}, + ) + meta: list = [None] * world + dist.all_gather_object(meta, my_meta) + # Everything below is keyed on the dense key list, so a rank that disagreed about it would HANG + # in a collective rather than fail. Check once, up front, and make that a diagnosable error. + if any(m[0] != dense_keys for m in meta): + mine = set(dense_keys) + diff = {r: sorted(set(m[0]) ^ mine)[:5] for r, m in enumerate(meta) if m[0] != dense_keys} + raise RuntimeError( + f"rank {rank}: dense key sets differ across ranks, so the export would deadlock; " + f"sample divergences per rank: {diff}" + ) + expert_bytes = [m[1] for m in meta] + all_shapes = [m[2] for m in meta] + + dense_sizes = {k: _nbytes(local_sd[k]) for k in dense_keys} + # Deterministic on every rank: the sort and the argmin both break ties explicitly. + loads = list(expert_bytes) + owner: dict = {} + for k in sorted(dense_keys, key=lambda k: (-dense_sizes[k], k)): + r = min(range(world), key=lambda i: (loads[i], i)) + owner[k] = r + loads[r] += dense_sizes[k] + + # Replicated and plain values are already whole on every rank -- the owner just unwraps, the + # rest just forget. No communication at all. + for k in empty0: + if owner[k] == rank: + _v = local_sd[k] + local_sd[k] = torch.empty(tuple(_v.shape), dtype=_v.dtype, device="cpu") + else: + del local_sd[k] + for k in replicated: + if owner[k] == rank: + local_sd[k] = local_sd[k].to_local() + else: + del local_sd[k] + for k in plain: + if owner[k] != rank: + del local_sd[k] + + # Concatenating local shards reproduces the global tensor only if they tile it exactly: dim-0 + # extents summing to the global dim 0, every other dim equal. Uneven sharding can leave + # to_local() returning a PADDED chunk, which would concatenate into a wrong (too large, + # misaligned) tensor rather than failing -- so verify instead of assuming. Derived from the + # same gathered shapes on every rank, so the split is agreed. + # all_shapes[r] is keyed by rank r's shard0 bucket, so its key SET is that rank's classification. + # Those buckets come from per-rank DTensor metadata, and a rank that disagrees makes the loop + # below raise a bare KeyError from deep inside the writer -- which says nothing about the real + # mistake. (Seen for real: a packed-weight rewrap re-inferred the global shape as local x world, + # so under uneven sharding each rank computed a different global shape and classified + # differently.) Name the divergence instead. + _mine = set(all_shapes[rank]) + _bad = { + r: sorted(set(all_shapes[r]) ^ _mine)[:3] + for r in range(world) + if set(all_shapes[r]) != _mine + } + if _bad: + raise RuntimeError( + f"rank {rank}: ranks disagree on which tensors are dim-0 sharded, so the export cannot " + f"assemble them; per-rank symmetric differences (first few): {dict(list(_bad.items())[:3])}" + ) + + untiled = [] + for k in shard0: + gshape = tuple(local_sd[k].shape) + shard_shapes = [all_shapes[r][k] for r in range(world)] + if not ( + sum(s[0] for s in shard_shapes) == gshape[0] + and all(tuple(s[1:]) == gshape[1:] for s in shard_shapes) + ): + untiled.append((k, gshape, shard_shapes[:3])) + if untiled: + raise RuntimeError( + f"rank {rank}: {len(untiled)} dim-0 sharded tensor(s) do not tile their global shape, " + f"so concatenating the local shards would not reconstruct them (to_local() may be " + f"returning padded chunks); first few: {untiled[:2]}" + ) + if shard0: + _gather_shard0_to_owners(local_sd, shard0, owner, all_shapes, rank, world) + + +def _revert_whole_sd(model: nn.Module, local_sd: dict) -> bool: + """Reverse the transformers conversion mapping on this rank's whole tensors, in place. + + Safe to run per-rank on a disjoint subset: renames are per-key, and a split's outputs stay on + the rank that owned the input. Every value is a whole plain tensor by now, so no placement + handling is needed. Best-effort and atomic, mirroring the single-process path -- on any + failure the in-memory names are kept for the whole dict rather than half-applied. + """ + from .quant_aware_conversion import revert_weight_conversion_quant_aware + + try: + reverted = revert_weight_conversion_quant_aware(model, local_sd) + except Exception as exc: + warnings.warn( + f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor " + "names may not match the original HF hub checkpoint." + ) + return False + # revert_weight_conversion_quant_aware returns the INPUT dict itself when the model has no + # reverse rules (a dense model with no conversion mapping). Clearing then updating would then + # read from the dict just emptied and wipe the state dict, so only rebuild on a fresh object. + if reverted is not local_sd: + local_sd.clear() + local_sd.update(reverted) + return True + + +def _write_even_shards( + local_sd: dict, export_dir: "str | Path", max_shard_size: "str | int", rank: int, world: int +) -> None: + """Write this rank's whole tensors directly as final HF safetensors; rank 0 unions the index. + + Every tensor is whole on exactly one rank, so there is no consolidation step. This rank's set is + split into even shards, numbered into a global range (prefix-sum of per-rank file counts), and + written with ``safetensors.save_file``. Must run AFTER the reverse conversion -- the reverse + renames keys and splits tensors, so the sizes and names bin-packed here have to be the final ones. + """ + import math + + import torch.distributed as dist + from safetensors.torch import save_file + + export_dir = Path(export_dir) + export_dir.mkdir(parents=True, exist_ok=True) + max_bytes = _size_to_bytes(max_shard_size) + + write_keys = sorted(local_sd) + sizes = {k: _nbytes(local_sd[k]) for k in write_keys} + # Tensors over the bound get files of their own, so only the rest decide how many to aim for. + fits = sum(size for size in sizes.values() if size <= max_bytes) + my_files = _even_bins(write_keys, sizes, max(1, math.ceil(fits / max_bytes)), max_bytes) + n_local = len(my_files) + counts: list = [None] * world + dist.all_gather_object(counts, n_local) + n_total = sum(counts) + if n_total == 0: + raise RuntimeError( + "distributed export produced no shards on any rank: every tensor was dropped " + "before the write (state dict emptied upstream)" + ) + base = sum(counts[:rank]) + + weight_map: dict = {} + my_bytes = 0 + for i, keys in enumerate(my_files): + fname = f"model-{base + i + 1:05d}-of-{n_total:05d}.safetensors" + tensors = {k: local_sd[k].detach().to("cpu").contiguous() for k in keys} + save_file(tensors, str(export_dir / fname), metadata={"format": "pt"}) + for k in keys: + weight_map[k] = fname + # Once a tensor is on disk this rank has no further use for it. Dropping it here keeps + # peak memory at roughly one shard above the owned set instead of holding all of them. + del local_sd[k] + my_bytes += sum(_nbytes(t) for t in tensors.values()) + del tensors + + # One collective, not two: at this point each rank has already finished writing, so the latency + # of these round trips is pure overhead on the critical path. + gathered: list = [None] * world + dist.all_gather_object(gathered, (weight_map, my_bytes)) + if rank == 0: + full_map: dict = {} + total_bytes = 0 + for m, b in gathered: + full_map.update(m) + total_bytes += b + # The reverse conversion checks for rename collisions, but only within one rank's slice. + # A collision ACROSS ranks would otherwise be silently swallowed by this dict update and + # ship a checkpoint that is quietly missing a tensor. + n_written = sum(len(m) for m, _ in gathered) + if n_written != len(full_map): + raise RuntimeError( + f"tensor name collision across ranks: {n_written} tensors written but only " + f"{len(full_map)} distinct names in the index" + ) + index = {"metadata": {"total_size": total_bytes}, "weight_map": full_map} + with open(export_dir / "model.safetensors.index.json", "w") as f: + json.dump(index, f, indent=2) + dist.barrier() diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py new file mode 100644 index 00000000000..5d1d421070e --- /dev/null +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -0,0 +1,361 @@ +# 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"); +# 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. + +"""The pieces of the no-gather HF checkpoint writer (``modelopt.torch.export.distributed_writer``). + +The writer moves every tensor to exactly one owner rank and writes it from there, so what it can +get wrong is structural: a tensor missing, duplicated, mis-assembled from its shards, renamed, or +left out of the weight index. These tests drive each piece directly, on CPU. +""" + +import json +import math + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +pytest.importorskip("safetensors") +pytest.importorskip("transformers") +from safetensors.torch import load_file +from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig + +from modelopt.torch.export import quant_aware_conversion +from modelopt.torch.export.distributed_writer import ( + _even_bins, + _is_full_world_shard0, + _materialize_owned, + _nbytes, + _revert_whole_sd, + _size_to_bytes, + _write_even_shards, +) + + +@pytest.fixture +def gloo_world_of_one(skip_on_windows, tmp_path_factory): + """A one-rank gloo process group and its CPU device mesh.""" + store = tmp_path_factory.mktemp("pg") / "store" + dist.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + yield init_device_mesh("cpu", (1,)) + dist.destroy_process_group() + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8) + + +def _assert_byte_identical(got, want): + assert set(got) == set(want), (sorted(set(want) - set(got)), sorted(set(got) - set(want))) + for key, ref in want.items(): + assert got[key].dtype == ref.dtype, key + assert got[key].shape == ref.shape, key + assert torch.equal(_raw_bytes(got[key]), _raw_bytes(ref.cpu())), key + + +def _load_checkpoint(export_dir): + """Every tensor the index names, read from the shard the index names for it.""" + index = json.loads((export_dir / "model.safetensors.index.json").read_text()) + weight_map = index["weight_map"] + shards = sorted(p.name for p in export_dir.glob("*.safetensors")) + assert set(weight_map.values()) == set(shards), "index and shard files disagree" + tensors = {} + for shard in shards: + loaded = load_file(str(export_dir / shard)) + assert {k for k, f in weight_map.items() if f == shard} == set(loaded), shard + tensors.update(loaded) + assert index["metadata"]["total_size"] == sum(_nbytes(t) for t in tensors.values()) + return tensors + + +def _mixtral_config(**overrides): + """A tiny Mixtral: fused experts in memory, per-expert ``block_sparse_moe`` names on the hub.""" + kwargs = { + "vocab_size": 64, + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "max_position_embeddings": 32, + **overrides, + } + return MixtralConfig(**kwargs) + + +def _llama_config(): + # 63 vocab rows: the embedding and lm_head shard unevenly over two ranks. + return LlamaConfig( + vocab_size=63, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + ) + + +# ------------------------------------------------------------------ pure helpers + + +@pytest.mark.parametrize( + ("size", "nbytes"), + [ + ("5GB", 5 * 10**9), + ("500MB", 500 * 10**6), + ("10KB", 10 * 10**3), + ("2TB", 2 * 10**12), + ("1GiB", 2**30), + ("3MiB", 3 * 2**20), + ("4KiB", 4 * 2**10), + ("1TiB", 2**40), + ("1.5GB", 1_500_000_000), + (" 2 mb ", 2 * 10**6), + ("123", 123), + (4096, 4096), + ], +) +def test_size_to_bytes_reads_hf_shard_sizes(size, nbytes): + assert _size_to_bytes(size) == nbytes + + +@pytest.mark.parametrize("size", ["5B", "5 gigs", ""]) +def test_size_to_bytes_rejects_an_unknown_unit(size): + with pytest.raises(ValueError): + _size_to_bytes(size) + + +def test_even_bins_places_every_key_once_within_the_bound(): + sizes = {f"t{i}": i + 1 for i in range(20)} # 210 bytes in all + + bins = _even_bins(list(sizes), sizes, math.ceil(210 / 50), 50) + + assert sorted(k for b in bins for k in b) == sorted(sizes) + assert all(0 < sum(sizes[k] for k in b) <= 50 for b in bins) + + +def test_even_bins_balances_the_files_instead_of_filling_them_in_turn(): + """``max_shard_size`` is a bound, not a target: two files of 8 rather than 10 + 6.""" + sizes = {"a": 4, "b": 3, "c": 3, "d": 2, "e": 2, "f": 2} + + bins = _even_bins(list(sizes), sizes, 2, 10) + + assert sorted(sum(sizes[k] for k in b) for b in bins) == [8, 8] + + +def test_even_bins_adds_bins_until_every_bin_fits(): + sizes = {"a": 6, "b": 6, "c": 6} + + bins = _even_bins(list(sizes), sizes, 1, 10) + + assert sorted(map(sorted, bins)) == [["a"], ["b"], ["c"]] + + +def test_even_bins_gives_a_tensor_larger_than_the_bound_a_bin_of_its_own(): + sizes = {"big": 100, "a": 1, "b": 1, "c": 1} + + bins = _even_bins(list(sizes), sizes, 1, 10) + + assert ["big"] in bins + assert sorted(k for b in bins for k in b) == sorted(sizes) + + +def test_even_bins_still_packs_the_rest_when_one_tensor_exceeds_the_bound(): + sizes = {"big": 100, "a": 1, "b": 1, "c": 1} + + bins = _even_bins(list(sizes), sizes, 1, 10) + + assert len(bins) == 2 + + +def test_even_bins_never_returns_an_empty_bin(): + sizes = {"a": 1, "b": 1} + assert sorted(map(sorted, _even_bins(list(sizes), sizes, 5, 10))) == [["a"], ["b"]] + assert _even_bins([], {}, 3, 10) == [] + + +def test_nbytes_counts_the_global_size_of_a_dtensor(gloo_world_of_one): + assert _nbytes(torch.zeros(3, 4, dtype=torch.bfloat16)) == 24 + assert _nbytes(torch.zeros(5, dtype=torch.uint8)) == 5 + # This rank holds 2 of 6 rows; every rank must agree on the size, so all 6 count. + shard = DTensor.from_local( + torch.zeros(2, 4), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=torch.Size([6, 4]), + stride=(4, 1), + ) + assert _nbytes(shard) == 6 * 4 * 4 + + +@pytest.mark.parametrize( + ("placement", "expected"), + [(Shard(0), True), (Shard(1), False), (Replicate(), False)], + ids=["shard0", "shard1", "replicate"], +) +def test_is_full_world_shard0_accepts_only_a_dim0_shard_over_the_whole_world( + gloo_world_of_one, placement, expected +): + tensor = distribute_tensor(torch.arange(12.0).reshape(4, 3), gloo_world_of_one, [placement]) + assert _is_full_world_shard0(tensor, 1) is expected + + +def test_is_full_world_shard0_rejects_a_mesh_that_is_not_the_whole_world(gloo_world_of_one): + tensor = distribute_tensor(torch.zeros(4, 3), gloo_world_of_one, [Shard(0)]) + assert _is_full_world_shard0(tensor, 2) is False + mesh_2d = init_device_mesh("cpu", (1, 1)) + tensor_2d = distribute_tensor(torch.zeros(4, 3), mesh_2d, [Shard(0), Replicate()]) + assert _is_full_world_shard0(tensor_2d, 1) is False + + +def test_is_full_world_shard0_answers_from_the_mesh_memo(gloo_world_of_one): + tensor = distribute_tensor(torch.zeros(4, 3), gloo_world_of_one, [Shard(0)]) + + memo = {} + assert _is_full_world_shard0(tensor, 1, memo) is True + assert memo == {id(gloo_world_of_one): True} + # A cached answer is used as is: the mesh is not read back again. + assert _is_full_world_shard0(tensor, 1, {id(gloo_world_of_one): False}) is False + + +def test_materialize_owned_leaves_every_tensor_whole_and_plain(gloo_world_of_one): + mesh = gloo_world_of_one + full = torch.arange(12.0).reshape(4, 3) + expert_key, expert = "model.layers.0.mlp.experts.0.up_proj.weight", torch.full((2,), 7.0) + sd = { + "plain": torch.ones(3), + "replicated": distribute_tensor(full, mesh, [Replicate()]), + "sharded": distribute_tensor(full, mesh, [Shard(0)]), + "empty": distribute_tensor(torch.zeros(0, 4, dtype=torch.bfloat16), mesh, [Shard(0)]), + expert_key: expert, + } + + _materialize_owned(sd, {expert_key}, 0, 1) + + assert set(sd) == {"plain", "replicated", "sharded", "empty", expert_key} + assert not any(isinstance(v, DTensor) for v in sd.values()) + assert torch.equal(sd["plain"], torch.ones(3)) + assert torch.equal(sd["replicated"], full) + assert torch.equal(sd["sharded"], full) + assert sd["empty"].shape == (0, 4) and sd["empty"].dtype == torch.bfloat16 + assert sd[expert_key] is expert # experts are already disjoint per rank: left alone + + +def test_materialize_owned_rejects_a_placement_it_cannot_assemble(gloo_world_of_one): + sd = {"w": distribute_tensor(torch.zeros(4, 4), gloo_world_of_one, [Shard(1)])} + with pytest.raises(RuntimeError, match=r"cannot materialize.*'w'"): + _materialize_owned(sd, set(), 0, 1) + + +def test_materialize_owned_rejects_shards_that_do_not_tile_the_global_shape(gloo_world_of_one): + """A padded local chunk would concatenate into a wrong tensor; it must raise instead.""" + padded = DTensor.from_local( + torch.zeros(2, 4), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=torch.Size([3, 4]), + stride=(4, 1), + ) + with pytest.raises(RuntimeError, match="do not tile their global shape"): + _materialize_owned({"w": padded}, set(), 0, 1) + + +def test_revert_whole_sd_keeps_a_state_dict_with_nothing_to_revert(): + """A model without a conversion mapping gets its input dict back from the reverse pass.""" + model = AutoModelForCausalLM.from_config(_llama_config()) + sd = dict(model.state_dict()) + before = dict(sd) + + assert _revert_whole_sd(model, sd) is True + + assert sd == before + + +def test_revert_whole_sd_renames_to_hub_names_in_place(): + model = AutoModelForCausalLM.from_config(_mixtral_config()) + gate, w1, w3, w2 = (torch.full((2, 2), float(i)) for i in range(4)) + sd = { + "model.layers.0.mlp.gate.weight": gate, + "model.layers.0.mlp.experts.0.gate_proj.weight": w1, + "model.layers.0.mlp.experts.0.up_proj.weight": w3, + "model.layers.0.mlp.experts.0.down_proj.weight": w2, + } + same_dict = sd + + assert _revert_whole_sd(model, sd) is True + + assert sd is same_dict + hub = "model.layers.0.block_sparse_moe" + assert set(sd) == { + f"{hub}.gate.weight", + f"{hub}.experts.0.w1.weight", + f"{hub}.experts.0.w3.weight", + f"{hub}.experts.0.w2.weight", + } + assert torch.equal(sd[f"{hub}.experts.0.w1.weight"], w1) + assert torch.equal(sd[f"{hub}.experts.0.w2.weight"], w2) + + +def test_revert_whole_sd_keeps_in_memory_names_when_the_reverse_fails(monkeypatch): + def _fail(model, state_dict): + raise RuntimeError("unsupported op") + + monkeypatch.setattr(quant_aware_conversion, "revert_weight_conversion_quant_aware", _fail) + sd = {"model.layers.0.mlp.experts.0.up_proj.weight": torch.ones(2)} + before = dict(sd) + + with pytest.warns(UserWarning, match="reverse weight conversion skipped"): + assert _revert_whole_sd(torch.nn.Module(), sd) is False + + assert sd == before + + +def test_write_even_shards_writes_even_shards_and_an_index_of_every_tensor( + gloo_world_of_one, tmp_path +): + torch.manual_seed(0) + tensors = {f"model.layers.{i}.weight": torch.randn(16, 8 + 4 * i) for i in range(10)} + tensors["model.layers.0.weight_scale"] = torch.tensor(0.5) + tensors["model.layers.1.weight"] = torch.randn(16, 16).to(torch.float8_e4m3fn) + tensors["model.norm.weight"] = torch.randn(16, dtype=torch.bfloat16) + tensors["lm_head.weight"] = torch.randn(16, 8).t() # not contiguous + max_bytes = 4096 + assert max(_nbytes(t) for t in tensors.values()) <= max_bytes + local = dict(tensors) + + _write_even_shards(local, tmp_path, "4096", 0, 1) + + assert not local, "a written tensor must be dropped from the rank's set" + shards = sorted(p.name for p in tmp_path.glob("*.safetensors")) + n = len(shards) + assert n == math.ceil(sum(_nbytes(t) for t in tensors.values()) / max_bytes) + assert shards == [f"model-{i:05d}-of-{n:05d}.safetensors" for i in range(1, n + 1)] + _assert_byte_identical(_load_checkpoint(tmp_path), tensors) + for shard in shards: + assert sum(_nbytes(t) for t in load_file(str(tmp_path / shard)).values()) <= max_bytes + + +def test_write_even_shards_fails_loudly_when_no_rank_has_anything_to_write( + gloo_world_of_one, tmp_path +): + with pytest.raises(RuntimeError, match="produced no shards"): + _write_even_shards({}, tmp_path, "1MB", 0, 1) From 2eddb1aa3cd71de5d6d1b2cff80e07ced59e6708 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:53:51 +0000 Subject: [PATCH 2/7] Fail the no-gather write on every rank when one rank's write fails A rank whose shard write raised left the writer and met its peers in a different collective: they raised an unrelated error or, when the failure was rank 0's name-collision check or index write before the closing barrier, waited out the process group's timeout. The per-rank file counts, the shard writes and rank 0's index write now each run inside the failure-carrying all_reduce, so every rank raises with the failing rank's error, and every rank checks the union of the written names for a collision. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 92 +++++++++++-------- .../torch/export/test_distributed_writer.py | 55 +++++++++++ 2 files changed, 108 insertions(+), 39 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 9c029ecec62..e2c58719257 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -369,21 +369,31 @@ def _revert_whole_sd(model: nn.Module, local_sd: dict) -> bool: return True +def _in_rank_order(values: list[list]) -> list: + """``all_reduce``'s combine that keeps every rank's value: each rank's own is wrapped in a list.""" + return [value for (value,) in values] + + def _write_even_shards( local_sd: dict, export_dir: "str | Path", max_shard_size: "str | int", rank: int, world: int ) -> None: - """Write this rank's whole tensors directly as final HF safetensors; rank 0 unions the index. + """Write this rank's whole tensors directly as final HF safetensors; rank 0 writes the index. Every tensor is whole on exactly one rank, so there is no consolidation step. This rank's set is split into even shards, numbered into a global range (prefix-sum of per-rank file counts), and written with ``safetensors.save_file``. Must run AFTER the reverse conversion -- the reverse renames keys and splits tensors, so the sizes and names bin-packed here have to be the final ones. + + Each step a rank can fail alone runs inside :func:`~modelopt.torch.utils.distributed.all_reduce`, + which raises on every rank when one fails: a rank leaving early would otherwise meet its peers + in a different collective, and they would hang or raise something unrelated. """ import math - import torch.distributed as dist from safetensors.torch import save_file + from modelopt.torch.utils.distributed import all_reduce + export_dir = Path(export_dir) export_dir.mkdir(parents=True, exist_ok=True) max_bytes = _size_to_bytes(max_shard_size) @@ -393,9 +403,7 @@ def _write_even_shards( # Tensors over the bound get files of their own, so only the rest decide how many to aim for. fits = sum(size for size in sizes.values() if size <= max_bytes) my_files = _even_bins(write_keys, sizes, max(1, math.ceil(fits / max_bytes)), max_bytes) - n_local = len(my_files) - counts: list = [None] * world - dist.all_gather_object(counts, n_local) + counts = all_reduce("shard counts", lambda: [len(my_files)], _in_rank_order) n_total = sum(counts) if n_total == 0: raise RuntimeError( @@ -404,40 +412,46 @@ def _write_even_shards( ) base = sum(counts[:rank]) - weight_map: dict = {} - my_bytes = 0 - for i, keys in enumerate(my_files): - fname = f"model-{base + i + 1:05d}-of-{n_total:05d}.safetensors" - tensors = {k: local_sd[k].detach().to("cpu").contiguous() for k in keys} - save_file(tensors, str(export_dir / fname), metadata={"format": "pt"}) - for k in keys: - weight_map[k] = fname - # Once a tensor is on disk this rank has no further use for it. Dropping it here keeps - # peak memory at roughly one shard above the owned set instead of holding all of them. - del local_sd[k] - my_bytes += sum(_nbytes(t) for t in tensors.values()) - del tensors + def write_mine() -> tuple[dict, int]: + weight_map: dict = {} + my_bytes = 0 + for i, keys in enumerate(my_files): + fname = f"model-{base + i + 1:05d}-of-{n_total:05d}.safetensors" + tensors = {k: local_sd[k].detach().to("cpu").contiguous() for k in keys} + save_file(tensors, str(export_dir / fname), metadata={"format": "pt"}) + for k in keys: + weight_map[k] = fname + # Once a tensor is on disk this rank has no further use for it. Dropping it here + # keeps peak memory at roughly one shard above the owned set instead of holding all. + del local_sd[k] + my_bytes += sum(_nbytes(t) for t in tensors.values()) + del tensors + return weight_map, my_bytes # One collective, not two: at this point each rank has already finished writing, so the latency # of these round trips is pure overhead on the critical path. - gathered: list = [None] * world - dist.all_gather_object(gathered, (weight_map, my_bytes)) - if rank == 0: - full_map: dict = {} - total_bytes = 0 - for m, b in gathered: - full_map.update(m) - total_bytes += b - # The reverse conversion checks for rename collisions, but only within one rank's slice. - # A collision ACROSS ranks would otherwise be silently swallowed by this dict update and - # ship a checkpoint that is quietly missing a tensor. - n_written = sum(len(m) for m, _ in gathered) - if n_written != len(full_map): - raise RuntimeError( - f"tensor name collision across ranks: {n_written} tensors written but only " - f"{len(full_map)} distinct names in the index" - ) - index = {"metadata": {"total_size": total_bytes}, "weight_map": full_map} - with open(export_dir / "model.safetensors.index.json", "w") as f: - json.dump(index, f, indent=2) - dist.barrier() + gathered = all_reduce("shards written", lambda: [write_mine()], _in_rank_order) + full_map: dict = {} + total_bytes = 0 + for m, b in gathered: + full_map.update(m) + total_bytes += b + # The reverse conversion checks for rename collisions, but only within one rank's slice. + # A collision ACROSS ranks would otherwise be silently swallowed by this dict update and + # ship a checkpoint that is quietly missing a tensor. Every rank checks, so all of them raise. + n_written = sum(len(m) for m, _ in gathered) + if n_written != len(full_map): + raise RuntimeError( + f"tensor name collision across ranks: {n_written} tensors written but only " + f"{len(full_map)} distinct names in the index" + ) + + def write_index() -> bool: + if rank == 0: + index = {"metadata": {"total_size": total_bytes}, "weight_map": full_map} + with open(export_dir / "model.safetensors.index.json", "w") as f: + json.dump(index, f, indent=2) + return True + + # Also the closing barrier: no rank returns before the index is on disk. + all_reduce("index written", write_index, all) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 5d1d421070e..96a5a4fefc5 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -20,8 +20,12 @@ left out of the weight index. These tests drive each piece directly, on CPU. """ +import contextlib import json import math +from datetime import timedelta +from functools import partial +from unittest import mock import pytest import torch @@ -359,3 +363,54 @@ def test_write_even_shards_fails_loudly_when_no_rank_has_anything_to_write( ): with pytest.raises(RuntimeError, match="produced no shards"): _write_even_shards({}, tmp_path, "1MB", 0, 1) + + +@contextlib.contextmanager +def _bounded_collectives(seconds=10): + """Run the writer's default-group collectives on a new group that gives up after ``seconds``. + + A rank that fails alone can leave a peer waiting in a collective it never joins. On this group + that wait raises in seconds instead of hanging for the default group's timeout, and the default + group, which the worker pool keeps for the next test, never falls out of step. + """ + group = dist.new_group(timeout=timedelta(seconds=seconds)) + with ( + mock.patch.object(dist, "all_gather_object", partial(dist.all_gather_object, group=group)), + mock.patch.object(dist, "barrier", partial(dist.barrier, group=group)), + ): + yield + + +def _one_rank_fails_job(rank, world_size, tmp_dir, fails): + failing_rank, step, target = { + "shard": (1, "shards written", "safetensors.torch.save_file"), + "index": (0, "index written", "json.dump"), + }[fails] + if rank == failing_rank: + failure = mock.patch(target, side_effect=OSError("disk full")) + expected = pytest.raises(OSError, match="disk full") + else: + failure = contextlib.nullcontext() + expected = pytest.raises( + RuntimeError, match=f"rank {failing_rank} failed to compute {step}: OSError: disk full" + ) + with _bounded_collectives(), failure, expected: + _write_even_shards({f"rank{rank}.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size) + + +@pytest.mark.parametrize("fails", ["shard", "index"]) +def test_a_write_that_fails_on_one_rank_fails_on_every_rank(cpu_dist_workers, tmp_path, fails): + """Not just on the rank that failed, while its peers wait in a collective it never joins.""" + cpu_dist_workers.run(_one_rank_fails_job, tmp_dir=tmp_path, fails=fails) + + +def _cross_rank_collision_job(rank, world_size, tmp_dir): + with ( + _bounded_collectives(), + pytest.raises(RuntimeError, match="name collision across ranks"), + ): + _write_even_shards({"model.norm.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size) + + +def test_a_name_two_ranks_both_write_fails_the_write_on_every_rank(cpu_dist_workers, tmp_path): + cpu_dist_workers.run(_cross_rank_collision_job, tmp_dir=tmp_path) From f379f0ab64b95f9e4bb00dfbee70f71255afd4bc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:55:04 +0000 Subject: [PATCH 3/7] Trim the writer pieces' tests into tables _size_to_bytes is one table of a decimal and a binary unit, case and padding, bare numbers and the refusals, not a row per unit. _even_bins is one table of the file sizes it produces; the every-key-once test goes, since the shard writer's test fails when a tensor is dropped. The _materialize_owned refusals are one table: a dim-1 shard, a smaller or 2-D mesh, a padded chunk. The _is_full_world_shard0 accept/reject tests go with it, as the materialize test's sharded tensor and that table cover both answers. _nbytes keeps only its DTensor case, the revert fallback no longer pins its warning text, and the two pooled jobs use the shared bounded_collectives with plain asserts: the peer of a failing rank must get that rank's error. The local gloo_world_of_one goes for the conftest's. Signed-off-by: Shengliang Xu --- .../torch/export/test_distributed_writer.py | 240 ++++++++---------- 1 file changed, 99 insertions(+), 141 deletions(-) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 96a5a4fefc5..a5e4b843b84 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -23,13 +23,11 @@ import contextlib import json import math -from datetime import timedelta -from functools import partial from unittest import mock import pytest import torch -import torch.distributed as dist +from _test_utils.torch.distributed.utils import bounded_collectives from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor @@ -50,15 +48,6 @@ ) -@pytest.fixture -def gloo_world_of_one(skip_on_windows, tmp_path_factory): - """A one-rank gloo process group and its CPU device mesh.""" - store = tmp_path_factory.mktemp("pg") / "store" - dist.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) - yield init_device_mesh("cpu", (1,)) - dist.destroy_process_group() - - def _raw_bytes(tensor): return tensor.detach().contiguous().reshape(-1).view(torch.uint8) @@ -123,81 +112,46 @@ def _llama_config(): ("size", "nbytes"), [ ("5GB", 5 * 10**9), - ("500MB", 500 * 10**6), - ("10KB", 10 * 10**3), - ("2TB", 2 * 10**12), ("1GiB", 2**30), - ("3MiB", 3 * 2**20), - ("4KiB", 4 * 2**10), - ("1TiB", 2**40), - ("1.5GB", 1_500_000_000), - (" 2 mb ", 2 * 10**6), + (" 1.5 mb ", 1_500_000), ("123", 123), (4096, 4096), + ("5B", ValueError), + ("5 gigs", ValueError), + ("", ValueError), ], ) def test_size_to_bytes_reads_hf_shard_sizes(size, nbytes): - assert _size_to_bytes(size) == nbytes - - -@pytest.mark.parametrize("size", ["5B", "5 gigs", ""]) -def test_size_to_bytes_rejects_an_unknown_unit(size): - with pytest.raises(ValueError): - _size_to_bytes(size) - - -def test_even_bins_places_every_key_once_within_the_bound(): - sizes = {f"t{i}": i + 1 for i in range(20)} # 210 bytes in all - - bins = _even_bins(list(sizes), sizes, math.ceil(210 / 50), 50) - - assert sorted(k for b in bins for k in b) == sorted(sizes) - assert all(0 < sum(sizes[k] for k in b) <= 50 for b in bins) - - -def test_even_bins_balances_the_files_instead_of_filling_them_in_turn(): - """``max_shard_size`` is a bound, not a target: two files of 8 rather than 10 + 6.""" - sizes = {"a": 4, "b": 3, "c": 3, "d": 2, "e": 2, "f": 2} - - bins = _even_bins(list(sizes), sizes, 2, 10) - - assert sorted(sum(sizes[k] for k in b) for b in bins) == [8, 8] - - -def test_even_bins_adds_bins_until_every_bin_fits(): - sizes = {"a": 6, "b": 6, "c": 6} - - bins = _even_bins(list(sizes), sizes, 1, 10) - - assert sorted(map(sorted, bins)) == [["a"], ["b"], ["c"]] - - -def test_even_bins_gives_a_tensor_larger_than_the_bound_a_bin_of_its_own(): - sizes = {"big": 100, "a": 1, "b": 1, "c": 1} - - bins = _even_bins(list(sizes), sizes, 1, 10) - - assert ["big"] in bins - assert sorted(k for b in bins for k in b) == sorted(sizes) - - -def test_even_bins_still_packs_the_rest_when_one_tensor_exceeds_the_bound(): - sizes = {"big": 100, "a": 1, "b": 1, "c": 1} - - bins = _even_bins(list(sizes), sizes, 1, 10) + """Decimal and binary units, any case and padding, a bare number; anything else is refused.""" + if nbytes is ValueError: + with pytest.raises(ValueError): + _size_to_bytes(size) + else: + assert _size_to_bytes(size) == nbytes - assert len(bins) == 2 +@pytest.mark.parametrize( + ("sizes", "n", "file_sizes"), + [ + # max_shard_size is a bound, not a target: two files of 8 rather than 10 + 6. + ({"a": 4, "b": 3, "c": 3, "d": 2, "e": 2, "f": 2}, 2, [8, 8]), + # One file per tensor once no two fit together. + ({"a": 6, "b": 6, "c": 6}, 1, [6, 6, 6]), + # A tensor over the bound gets a file of its own, and the rest still share one. + ({"big": 100, "a": 1, "b": 1, "c": 1}, 1, [3, 100]), + # Never an empty file, asked for more files than tensors or for none. + ({"a": 1, "b": 1}, 5, [1, 1]), + ({}, 3, []), + ], + ids=["balanced", "grows", "oversized", "more-files-than-tensors", "nothing"], +) +def test_even_bins_balances_the_files_within_the_bound(sizes, n, file_sizes): + bins = _even_bins(list(sizes), sizes, n, 10) -def test_even_bins_never_returns_an_empty_bin(): - sizes = {"a": 1, "b": 1} - assert sorted(map(sorted, _even_bins(list(sizes), sizes, 5, 10))) == [["a"], ["b"]] - assert _even_bins([], {}, 3, 10) == [] + assert sorted(sum(sizes[k] for k in b) for b in bins) == file_sizes def test_nbytes_counts_the_global_size_of_a_dtensor(gloo_world_of_one): - assert _nbytes(torch.zeros(3, 4, dtype=torch.bfloat16)) == 24 - assert _nbytes(torch.zeros(5, dtype=torch.uint8)) == 5 # This rank holds 2 of 6 rows; every rank must agree on the size, so all 6 count. shard = DTensor.from_local( torch.zeros(2, 4), @@ -210,26 +164,6 @@ def test_nbytes_counts_the_global_size_of_a_dtensor(gloo_world_of_one): assert _nbytes(shard) == 6 * 4 * 4 -@pytest.mark.parametrize( - ("placement", "expected"), - [(Shard(0), True), (Shard(1), False), (Replicate(), False)], - ids=["shard0", "shard1", "replicate"], -) -def test_is_full_world_shard0_accepts_only_a_dim0_shard_over_the_whole_world( - gloo_world_of_one, placement, expected -): - tensor = distribute_tensor(torch.arange(12.0).reshape(4, 3), gloo_world_of_one, [placement]) - assert _is_full_world_shard0(tensor, 1) is expected - - -def test_is_full_world_shard0_rejects_a_mesh_that_is_not_the_whole_world(gloo_world_of_one): - tensor = distribute_tensor(torch.zeros(4, 3), gloo_world_of_one, [Shard(0)]) - assert _is_full_world_shard0(tensor, 2) is False - mesh_2d = init_device_mesh("cpu", (1, 1)) - tensor_2d = distribute_tensor(torch.zeros(4, 3), mesh_2d, [Shard(0), Replicate()]) - assert _is_full_world_shard0(tensor_2d, 1) is False - - def test_is_full_world_shard0_answers_from_the_mesh_memo(gloo_world_of_one): tensor = distribute_tensor(torch.zeros(4, 3), gloo_world_of_one, [Shard(0)]) @@ -263,24 +197,49 @@ def test_materialize_owned_leaves_every_tensor_whole_and_plain(gloo_world_of_one assert sd[expert_key] is expert # experts are already disjoint per rank: left alone -def test_materialize_owned_rejects_a_placement_it_cannot_assemble(gloo_world_of_one): - sd = {"w": distribute_tensor(torch.zeros(4, 4), gloo_world_of_one, [Shard(1)])} - with pytest.raises(RuntimeError, match=r"cannot materialize.*'w'"): - _materialize_owned(sd, set(), 0, 1) - - -def test_materialize_owned_rejects_shards_that_do_not_tile_the_global_shape(gloo_world_of_one): - """A padded local chunk would concatenate into a wrong tensor; it must raise instead.""" - padded = DTensor.from_local( +def _padded_shard(mesh): + """This rank's 2 rows of a 3-row tensor: a padded local chunk, which would concatenate wrong.""" + return DTensor.from_local( torch.zeros(2, 4), - gloo_world_of_one, + mesh, [Shard(0)], run_check=False, shape=torch.Size([3, 4]), stride=(4, 1), ) - with pytest.raises(RuntimeError, match="do not tile their global shape"): - _materialize_owned({"w": padded}, set(), 0, 1) + + +@pytest.mark.parametrize( + ("build", "world", "match"), + [ + ( + lambda mesh: distribute_tensor(torch.zeros(4, 4), mesh, [Shard(1)]), + 1, + "cannot materialize", + ), + # A dim-0 shard over a mesh that is not the whole world: smaller, or 2-D. + ( + lambda mesh: distribute_tensor(torch.zeros(4, 3), mesh, [Shard(0)]), + 2, + "cannot materialize", + ), + ( + lambda mesh: distribute_tensor( + torch.zeros(4, 3), init_device_mesh("cpu", (1, 1)), [Shard(0), Replicate()] + ), + 1, + "cannot materialize", + ), + (_padded_shard, 1, "do not tile their global shape"), + ], + ids=["shard1", "smaller-mesh", "2d-mesh", "padded"], +) +def test_materialize_owned_rejects_a_tensor_it_cannot_assemble( + gloo_world_of_one, build, world, match +): + """Refused before anything moves, rather than assembled wrong.""" + with pytest.raises(RuntimeError, match=match): + _materialize_owned({"w": build(gloo_world_of_one)}, set(), 0, world) def test_revert_whole_sd_keeps_a_state_dict_with_nothing_to_revert(): @@ -327,7 +286,7 @@ def _fail(model, state_dict): sd = {"model.layers.0.mlp.experts.0.up_proj.weight": torch.ones(2)} before = dict(sd) - with pytest.warns(UserWarning, match="reverse weight conversion skipped"): + with pytest.warns(UserWarning): assert _revert_whole_sd(torch.nn.Module(), sd) is False assert sd == before @@ -365,37 +324,31 @@ def test_write_even_shards_fails_loudly_when_no_rank_has_anything_to_write( _write_even_shards({}, tmp_path, "1MB", 0, 1) -@contextlib.contextmanager -def _bounded_collectives(seconds=10): - """Run the writer's default-group collectives on a new group that gives up after ``seconds``. - - A rank that fails alone can leave a peer waiting in a collective it never joins. On this group - that wait raises in seconds instead of hanging for the default group's timeout, and the default - group, which the worker pool keeps for the next test, never falls out of step. - """ - group = dist.new_group(timeout=timedelta(seconds=seconds)) - with ( - mock.patch.object(dist, "all_gather_object", partial(dist.all_gather_object, group=group)), - mock.patch.object(dist, "barrier", partial(dist.barrier, group=group)), - ): - yield - - def _one_rank_fails_job(rank, world_size, tmp_dir, fails): - failing_rank, step, target = { - "shard": (1, "shards written", "safetensors.torch.save_file"), - "index": (0, "index written", "json.dump"), - }[fails] + failing_rank, target = {"shard": (1, "safetensors.torch.save_file"), "index": (0, "json.dump")}[ + fails + ] + failure = ( + mock.patch(target, side_effect=OSError("disk full")) + if rank == failing_rank + else contextlib.nullcontext() + ) + try: + with bounded_collectives(), failure: + _write_even_shards( + {f"rank{rank}.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size + ) + except (OSError, RuntimeError) as error: + raised = error + else: + raised = None + # The failing rank raises its own error; its peer gets that error, naming the rank, instead of + # waiting for it. if rank == failing_rank: - failure = mock.patch(target, side_effect=OSError("disk full")) - expected = pytest.raises(OSError, match="disk full") + assert isinstance(raised, OSError), raised else: - failure = contextlib.nullcontext() - expected = pytest.raises( - RuntimeError, match=f"rank {failing_rank} failed to compute {step}: OSError: disk full" - ) - with _bounded_collectives(), failure, expected: - _write_even_shards({f"rank{rank}.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size) + assert isinstance(raised, RuntimeError), raised + assert f"rank {failing_rank}" in str(raised) and "disk full" in str(raised), raised @pytest.mark.parametrize("fails", ["shard", "index"]) @@ -405,11 +358,16 @@ def test_a_write_that_fails_on_one_rank_fails_on_every_rank(cpu_dist_workers, tm def _cross_rank_collision_job(rank, world_size, tmp_dir): - with ( - _bounded_collectives(), - pytest.raises(RuntimeError, match="name collision across ranks"), - ): - _write_even_shards({"model.norm.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size) + try: + with bounded_collectives(): + _write_even_shards( + {"model.norm.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size + ) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + assert "name collision across ranks" in raised, raised def test_a_name_two_ranks_both_write_fails_the_write_on_every_rank(cpu_dist_workers, tmp_path): From e6b1278f2d5a5514da9cd3013f79fb9429583395 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:51:38 +0000 Subject: [PATCH 4/7] Create the export directory inside the writer's first collective _write_even_shards created the export directory before its first collective, so a rank whose filesystem refused the mkdir raised alone while its peers waited in "shard counts" until the process-group timeout, against the function's own contract that every step a rank can fail alone runs inside all_reduce. Create it inside the "shard counts" all_reduce instead, which adds no collective. The one-rank failure test gains a "dir" case. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 10 ++++++++-- tests/unit/torch/export/test_distributed_writer.py | 10 ++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index e2c58719257..75d69e8850e 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -395,7 +395,6 @@ def _write_even_shards( from modelopt.torch.utils.distributed import all_reduce export_dir = Path(export_dir) - export_dir.mkdir(parents=True, exist_ok=True) max_bytes = _size_to_bytes(max_shard_size) write_keys = sorted(local_sd) @@ -403,7 +402,14 @@ def _write_even_shards( # Tensors over the bound get files of their own, so only the rest decide how many to aim for. fits = sum(size for size in sizes.values() if size <= max_bytes) my_files = _even_bins(write_keys, sizes, max(1, math.ceil(fits / max_bytes)), max_bytes) - counts = all_reduce("shard counts", lambda: [len(my_files)], _in_rank_order) + + def count_mine() -> list[int]: + # The directory can fail to appear on one rank alone (its node's filesystem), so it is + # created inside the first collective, which carries that failure to every rank. + export_dir.mkdir(parents=True, exist_ok=True) + return [len(my_files)] + + counts = all_reduce("shard counts", count_mine, _in_rank_order) n_total = sum(counts) if n_total == 0: raise RuntimeError( diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index a5e4b843b84..7eae9bb9b34 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -325,9 +325,11 @@ def test_write_even_shards_fails_loudly_when_no_rank_has_anything_to_write( def _one_rank_fails_job(rank, world_size, tmp_dir, fails): - failing_rank, target = {"shard": (1, "safetensors.torch.save_file"), "index": (0, "json.dump")}[ - fails - ] + failing_rank, target = { + "dir": (1, "pathlib.Path.mkdir"), + "shard": (1, "safetensors.torch.save_file"), + "index": (0, "json.dump"), + }[fails] failure = ( mock.patch(target, side_effect=OSError("disk full")) if rank == failing_rank @@ -351,7 +353,7 @@ def _one_rank_fails_job(rank, world_size, tmp_dir, fails): assert f"rank {failing_rank}" in str(raised) and "disk full" in str(raised), raised -@pytest.mark.parametrize("fails", ["shard", "index"]) +@pytest.mark.parametrize("fails", ["dir", "shard", "index"]) def test_a_write_that_fails_on_one_rank_fails_on_every_rank(cpu_dist_workers, tmp_path, fails): """Not just on the rank that failed, while its peers wait in a collective it never joins.""" cpu_dist_workers.run(_one_rank_fails_job, tmp_dir=tmp_path, fails=fails) From 7187792fc67fa185bd07c126171a1edf4b644dbd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:24:06 +0000 Subject: [PATCH 5/7] Refuse an unassemblable tensor on every rank of the no-gather writer _materialize_owned refused an unsupported placement before its metadata gather, and checked tiling against each rank's own view of the global shape. A rank whose view diverged (a packed-weight rewrap once re-inferred the global shape per rank) then raised alone while its peers waited in the gather or the point-to-point transport until the timeout. Gather every rank's placement verdict and view of the global shape in the one metadata collective, and decide every refusal from that data before anything moves, so all ranks raise. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 108 ++++++++++-------- .../torch/export/test_distributed_writer.py | 58 +++++++--- 2 files changed, 102 insertions(+), 64 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 75d69e8850e..766f38da229 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -221,7 +221,8 @@ def _materialize_owned(local_sd: dict, expert_key_set: set, rank: int, world: in # Classify FIRST: placements are local DTensor metadata, so this needs no communication, and # doing it up front means the shard shapes can ride along in the single metadata collective - # below. Every rank sees the same placements and so agrees on the buckets. + # below. Every rank should see the same placements; one that does not is caught below, from the + # gathered metadata, so the refusal happens on every rank. replicated, shard0, unsupported, plain, empty0 = [], [], [], [], [] for k in dense_keys: v = local_sd[k] @@ -241,24 +242,18 @@ def _materialize_owned(local_sd: dict, expert_key_set: set, rank: int, world: in else: unsupported.append((k, str(v.placements), v.device_mesh.ndim)) - # Reject rather than silently taking a slower, untested route. The point-to-point transport - # covers replicated, plain, and dim-0 shards on a 1-D full-world mesh; anything else (2-D mesh, - # non-dim-0 shard, partial placement) has no correct handling here. Raising names the offending - # keys, which is recoverable; emitting a mis-assembled tensor would not be. - if unsupported: - raise RuntimeError( - f"rank {rank}: {len(unsupported)} tensor(s) have a placement the no-gather export " - f"cannot materialize; first few: " - f"{[(k, pl, f'mesh.ndim={nd}') for k, pl, nd in unsupported[:3]]}" - ) - # ONE collective for all the metadata: the key list (to catch divergence), this rank's expert - # byte count (to seed load balancing), and its local dim-0 extents (to size receive buffers). - # These were three separate all_gather_object round trips, each pickling a ~900-entry payload. + # byte count (to seed load balancing), its local dim-0 extents (to size receive buffers), its view + # of each dim-0 shard's global shape, and the placements it cannot materialize. The first three + # were separate all_gather_object round trips, each pickling a ~900-entry payload. Every refusal below + # is decided from this gathered data, so it is the same on every rank: a rank that raised alone + # would leave its peers waiting in a collective it never joins. my_meta = ( dense_keys, sum(_nbytes(local_sd[k]) for k in expert_keys), {k: tuple(local_sd[k].to_local().shape) for k in shard0}, + {k: tuple(local_sd[k].shape) for k in shard0}, + [(k, pl, f"mesh.ndim={nd}") for k, pl, nd in unsupported[:3]], ) meta: list = [None] * world dist.all_gather_object(meta, my_meta) @@ -271,40 +266,20 @@ def _materialize_owned(local_sd: dict, expert_key_set: set, rank: int, world: in f"rank {rank}: dense key sets differ across ranks, so the export would deadlock; " f"sample divergences per rank: {diff}" ) + # Reject rather than silently taking a slower, untested route. The point-to-point transport + # covers replicated, plain, and dim-0 shards on a 1-D full-world mesh; anything else (2-D mesh, + # non-dim-0 shard, partial placement) has no correct handling here. Raising names the offending + # keys, which is recoverable; emitting a mis-assembled tensor would not be. + unsupported_by_rank = {r: m[4] for r, m in enumerate(meta) if m[4]} + if unsupported_by_rank: + raise RuntimeError( + f"rank {rank}: tensor(s) have a placement the no-gather export cannot materialize; " + f"first few per rank: {unsupported_by_rank}" + ) expert_bytes = [m[1] for m in meta] all_shapes = [m[2] for m in meta] + global_shapes = [m[3] for m in meta] - dense_sizes = {k: _nbytes(local_sd[k]) for k in dense_keys} - # Deterministic on every rank: the sort and the argmin both break ties explicitly. - loads = list(expert_bytes) - owner: dict = {} - for k in sorted(dense_keys, key=lambda k: (-dense_sizes[k], k)): - r = min(range(world), key=lambda i: (loads[i], i)) - owner[k] = r - loads[r] += dense_sizes[k] - - # Replicated and plain values are already whole on every rank -- the owner just unwraps, the - # rest just forget. No communication at all. - for k in empty0: - if owner[k] == rank: - _v = local_sd[k] - local_sd[k] = torch.empty(tuple(_v.shape), dtype=_v.dtype, device="cpu") - else: - del local_sd[k] - for k in replicated: - if owner[k] == rank: - local_sd[k] = local_sd[k].to_local() - else: - del local_sd[k] - for k in plain: - if owner[k] != rank: - del local_sd[k] - - # Concatenating local shards reproduces the global tensor only if they tile it exactly: dim-0 - # extents summing to the global dim 0, every other dim equal. Uneven sharding can leave - # to_local() returning a PADDED chunk, which would concatenate into a wrong (too large, - # misaligned) tensor rather than failing -- so verify instead of assuming. Derived from the - # same gathered shapes on every rank, so the split is agreed. # all_shapes[r] is keyed by rank r's shard0 bucket, so its key SET is that rank's classification. # Those buckets come from per-rank DTensor metadata, and a rank that disagrees makes the loop # below raise a bare KeyError from deep inside the writer -- which says nothing about the real @@ -323,21 +298,56 @@ def _materialize_owned(local_sd: dict, expert_key_set: set, rank: int, world: in f"assemble them; per-rank symmetric differences (first few): {dict(list(_bad.items())[:3])}" ) + # Concatenating local shards reproduces the global tensor only if they tile it exactly: dim-0 + # extents summing to the global dim 0, every other dim equal. Uneven sharding can leave + # to_local() returning a PADDED chunk, which would concatenate into a wrong (too large, + # misaligned) tensor rather than failing -- so verify instead of assuming. The global shape is + # the gathered one too: ranks that inferred different global shapes (the rewrap above) make the + # tensor untiled on every rank, not only on the rank whose own view does not add up. untiled = [] for k in shard0: - gshape = tuple(local_sd[k].shape) + gshapes = {global_shapes[r][k] for r in range(world)} + gshape = next(iter(gshapes)) shard_shapes = [all_shapes[r][k] for r in range(world)] - if not ( + if len(gshapes) != 1 or not ( sum(s[0] for s in shard_shapes) == gshape[0] and all(tuple(s[1:]) == gshape[1:] for s in shard_shapes) ): - untiled.append((k, gshape, shard_shapes[:3])) + untiled.append((k, sorted(gshapes), shard_shapes[:3])) if untiled: raise RuntimeError( f"rank {rank}: {len(untiled)} dim-0 sharded tensor(s) do not tile their global shape, " f"so concatenating the local shards would not reconstruct them (to_local() may be " - f"returning padded chunks); first few: {untiled[:2]}" + f"returning padded chunks, or the ranks disagree on the global shape); first few: " + f"{untiled[:2]}" ) + + dense_sizes = {k: _nbytes(local_sd[k]) for k in dense_keys} + # Deterministic on every rank: the sort and the argmin both break ties explicitly. + loads = list(expert_bytes) + owner: dict = {} + for k in sorted(dense_keys, key=lambda k: (-dense_sizes[k], k)): + r = min(range(world), key=lambda i: (loads[i], i)) + owner[k] = r + loads[r] += dense_sizes[k] + + # Replicated and plain values are already whole on every rank -- the owner just unwraps, the + # rest just forget. No communication at all. + for k in empty0: + if owner[k] == rank: + _v = local_sd[k] + local_sd[k] = torch.empty(tuple(_v.shape), dtype=_v.dtype, device="cpu") + else: + del local_sd[k] + for k in replicated: + if owner[k] == rank: + local_sd[k] = local_sd[k].to_local() + else: + del local_sd[k] + for k in plain: + if owner[k] != rank: + del local_sd[k] + if shard0: _gather_shard0_to_owners(local_sd, shard0, owner, all_shapes, rank, world) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 7eae9bb9b34..6ab3048ad97 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -210,36 +210,64 @@ def _padded_shard(mesh): @pytest.mark.parametrize( - ("build", "world", "match"), + ("build", "match"), [ ( lambda mesh: distribute_tensor(torch.zeros(4, 4), mesh, [Shard(1)]), - 1, - "cannot materialize", - ), - # A dim-0 shard over a mesh that is not the whole world: smaller, or 2-D. - ( - lambda mesh: distribute_tensor(torch.zeros(4, 3), mesh, [Shard(0)]), - 2, "cannot materialize", ), + # A dim-0 shard over a 2-D mesh; a smaller mesh is a two-rank case below. ( lambda mesh: distribute_tensor( torch.zeros(4, 3), init_device_mesh("cpu", (1, 1)), [Shard(0), Replicate()] ), - 1, "cannot materialize", ), - (_padded_shard, 1, "do not tile their global shape"), + (_padded_shard, "do not tile their global shape"), ], - ids=["shard1", "smaller-mesh", "2d-mesh", "padded"], + ids=["shard1", "2d-mesh", "padded"], ) -def test_materialize_owned_rejects_a_tensor_it_cannot_assemble( - gloo_world_of_one, build, world, match -): +def test_materialize_owned_rejects_a_tensor_it_cannot_assemble(gloo_world_of_one, build, match): """Refused before anything moves, rather than assembled wrong.""" with pytest.raises(RuntimeError, match=match): - _materialize_owned({"w": build(gloo_world_of_one)}, set(), 0, world) + _materialize_owned({"w": build(gloo_world_of_one)}, set(), 0, 1) + + +def _refusing_materialize_job(rank, world_size, case): + if case == "smaller-mesh": + # A dim-0 shard over a mesh smaller than the world: each rank's own 1-rank sub-mesh. + sub = init_device_mesh("cpu", (world_size, 1), mesh_dim_names=("dp", "sub"))["sub"] + w = distribute_tensor(torch.zeros(4, 3), sub, [Shard(0)]) + else: + mesh = init_device_mesh("cpu", (world_size,)) + if case == "placement": + # Only rank 1 holds a placement the export cannot materialize. + local, placement = ( + (torch.zeros(2, 4), Shard(0)) if rank == 0 else (torch.zeros(4, 2), Shard(1)) + ) + gshape = (4, 4) + else: + # A rewrap that re-inferred the global shape per rank: rank 0's own view tiles, rank 1's not. + local, placement, gshape = torch.zeros(2, 4), Shard(0), (4 if rank == 0 else 3, 4) + w = DTensor.from_local( + local, mesh, [placement], run_check=False, shape=torch.Size(gshape), stride=(4, 1) + ) + try: + with bounded_collectives(): + _materialize_owned({"w": w}, set(), rank, world_size) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + match = "do not tile their global shape" if case == "global-shape" else "cannot materialize" + assert match in raised, raised + + +@pytest.mark.parametrize("case", ["smaller-mesh", "placement", "global-shape"]) +def test_materialize_owned_refuses_on_every_rank(cpu_dist_workers, case): + """Every rank refuses, the ranks whose own view looks fine included, instead of waiting in a + collective the refusing rank never joins.""" + cpu_dist_workers.run(_refusing_materialize_job, case=case) def test_revert_whole_sd_keeps_a_state_dict_with_nothing_to_revert(): From b44b4d8e66d0c226dbb3b388fd6556f9fd25b6f2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 05:19:14 +0000 Subject: [PATCH 6/7] Keep the in-memory names on every rank when one rank cannot revert them Each rank of the no-gather writer reverses the conversion mapping on the tensors it owns. A merge needs its whole group, which the per-tensor ownership can spread over ranks, and a rank holding none of a group fails it. That rank kept its in-memory names while the others wrote hub names, so the checkpoint mixed the two. The ranks now agree on the outcome: if the reverse fails on any rank, every rank keeps the in-memory names, as the single-process export does on a failure. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 34 ++++++++++++++----- .../torch/export/test_distributed_writer.py | 28 ++++++++++----- 2 files changed, 46 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 766f38da229..53ec16f5fdc 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -355,18 +355,36 @@ def _materialize_owned(local_sd: dict, expert_key_set: set, rank: int, world: in def _revert_whole_sd(model: nn.Module, local_sd: dict) -> bool: """Reverse the transformers conversion mapping on this rank's whole tensors, in place. - Safe to run per-rank on a disjoint subset: renames are per-key, and a split's outputs stay on - the rank that owned the input. Every value is a whole plain tensor by now, so no placement - handling is needed. Best-effort and atomic, mirroring the single-process path -- on any - failure the in-memory names are kept for the whole dict rather than half-applied. + Each rank reverts the disjoint subset it owns. Every value is a whole plain tensor by now, so no + placement handling is needed. Renames and splits act on one tensor at a time, so a subset is + all they need. A merge is different: it needs its whole group, ownership may spread a group + over ranks, and a rank holding none of it fails. Best-effort and atomic across the ranks, + mirroring the single-process path, which keeps the in-memory names for the whole dict on any + failure. If the reverse fails on any rank, every rank keeps them, rather than writing a + checkpoint that mixes the two. Collective. """ + from modelopt.torch.utils.distributed import all_reduce + from .quant_aware_conversion import revert_weight_conversion_quant_aware - try: - reverted = revert_weight_conversion_quant_aware(model, local_sd) - except Exception as exc: + reverted = local_sd + + def revert() -> str | None: + nonlocal reverted + try: + reverted = revert_weight_conversion_quant_aware(model, local_sd) + except Exception as exc: + return str(exc) + return None + + failures = all_reduce( + "reverse weight conversion", + revert, + lambda errors: [f"rank {r}: {error}" for r, error in enumerate(errors) if error], + ) + if failures: warnings.warn( - f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor " + f"Quant-aware reverse weight conversion skipped ({failures[0]}); exported tensor " "names may not match the original HF hub checkpoint." ) return False diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 6ab3048ad97..236c2e5acb8 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -306,20 +306,32 @@ def test_revert_whole_sd_renames_to_hub_names_in_place(): assert torch.equal(sd[f"{hub}.experts.0.w2.weight"], w2) -def test_revert_whole_sd_keeps_in_memory_names_when_the_reverse_fails(monkeypatch): - def _fail(model, state_dict): - raise RuntimeError("unsupported op") - - monkeypatch.setattr(quant_aware_conversion, "revert_weight_conversion_quant_aware", _fail) +def _one_rank_cannot_revert_job(rank, world_size): + # Rank 0's subset reverts; rank 1's fails, as a merge whose group ownership split does. + model = AutoModelForCausalLM.from_config(_mixtral_config()) sd = {"model.layers.0.mlp.experts.0.up_proj.weight": torch.ones(2)} before = dict(sd) - - with pytest.warns(UserWarning): - assert _revert_whole_sd(torch.nn.Module(), sd) is False + failure = mock.patch.object( + quant_aware_conversion, + "revert_weight_conversion_quant_aware", + side_effect=RuntimeError("incomplete merge"), + ) + with ( + bounded_collectives(), + failure if rank == 1 else contextlib.nullcontext(), + pytest.warns(UserWarning, match="rank 1: incomplete merge"), + ): + assert _revert_whole_sd(model, sd) is False assert sd == before +def test_revert_whole_sd_keeps_in_memory_names_on_every_rank_when_one_fails(cpu_dist_workers): + """Rank 0 could revert its own subset, but a checkpoint mixing hub and in-memory names would + load neither way.""" + cpu_dist_workers.run(_one_rank_cannot_revert_job) + + def test_write_even_shards_writes_even_shards_and_an_index_of_every_tensor( gloo_world_of_one, tmp_path ): From dd1344d0a75022a2433b0b2fc8bf73996872d444 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 19:05:26 +0000 Subject: [PATCH 7/7] Date distributed_writer 2026, as LICENSE_HEADER does Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 53ec16f5fdc..82f091be7a2 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.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");