diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
new file mode 100644
index 00000000000..82f091be7a2
--- /dev/null
+++ b/modelopt/torch/export/distributed_writer.py
@@ -0,0 +1,491 @@
+# 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.
+
+"""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 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]
+        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))
+
+    # ONE collective for all the metadata: the key list (to catch divergence), this rank's expert
+    # 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)
+    # 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}"
+        )
+    # 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]
+
+    # 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])}"
+        )
+
+    # 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:
+        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 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, 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, 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)
+
+
+def _revert_whole_sd(model: nn.Module, local_sd: dict) -> bool:
+    """Reverse the transformers conversion mapping on this rank's whole tensors, in place.
+
+    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
+
+    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 ({failures[0]}); 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 _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 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
+
+    from safetensors.torch import save_file
+
+    from modelopt.torch.utils.distributed import all_reduce
+
+    export_dir = Path(export_dir)
+    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)
+
+    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(
+            "distributed export produced no shards on any rank: every tensor was dropped "
+            "before the write (state dict emptied upstream)"
+        )
+    base = sum(counts[:rank])
+
+    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 = 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
new file mode 100644
index 00000000000..236c2e5acb8
--- /dev/null
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -0,0 +1,416 @@
+# 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 contextlib
+import json
+import math
+from unittest import mock
+
+import pytest
+import torch
+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
+
+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,
+)
+
+
+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),
+        ("1GiB", 2**30),
+        (" 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):
+    """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
+
+
+@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)
+
+    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):
+    # 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
+
+
+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 _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),
+        mesh,
+        [Shard(0)],
+        run_check=False,
+        shape=torch.Size([3, 4]),
+        stride=(4, 1),
+    )
+
+
+@pytest.mark.parametrize(
+    ("build", "match"),
+    [
+        (
+            lambda mesh: distribute_tensor(torch.zeros(4, 4), mesh, [Shard(1)]),
+            "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()]
+            ),
+            "cannot materialize",
+        ),
+        (_padded_shard, "do not tile their global shape"),
+    ],
+    ids=["shard1", "2d-mesh", "padded"],
+)
+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, 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():
+    """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 _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)
+    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
+):
+    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)
+
+
+def _one_rank_fails_job(rank, world_size, tmp_dir, 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
+        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:
+        assert isinstance(raised, OSError), raised
+    else:
+        assert isinstance(raised, RuntimeError), raised
+        assert f"rank {failing_rank}" in str(raised) and "disk full" in str(raised), raised
+
+
+@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)
+
+
+def _cross_rank_collision_job(rank, world_size, tmp_dir):
+    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):
+    cpu_dist_workers.run(_cross_rank_collision_job, tmp_dir=tmp_path)