From 3297dcd725d95d7cec50ed84ceecb05d1034cf8c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 00:11:19 +0000
Subject: [PATCH 1/6] Move each tensor of a sharded model whole to one owner
 rank

The first pieces of the no-gather HF checkpoint writer. _materialize_owned
assigns every dense tensor to one owner, largest first onto the lightest
rank and seeded with each rank's expert bytes, keeping each group a
merge needs together, and moves it there: dim-0 shards travel once, to
their owner, over batched point-to-point sends bucketed by peer and
dtype, so no rank all-gathers the model; replicated and plain tensors
are unwrapped by their owner and dropped elsewhere. Every refusal, an
unsupported placement, shards that do not tile, ranks that disagree on
keys, dtypes, sizes or groups, is decided from one gathered metadata
collective, so it is the same on every rank.

The one-rank gloo world the meta-DCP load tests used moves to the shared
conftest. The GPU tests move BF16, FP8 and uint8 shards of uneven and
empty extent over NCCL.

[2/3] of #2751's split, at its review's request.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py   | 347 ++++++++++++++++++
 .../export/test_distributed_writer_gpu.py     |  89 +++++
 tests/unit/torch/conftest.py                  |  14 +
 .../torch/export/test_distributed_writer.py   | 245 +++++++++++++
 tests/unit/torch/utils/test_meta_dcp_load.py  |   9 -
 5 files changed, 695 insertions(+), 9 deletions(-)
 create mode 100644 modelopt/torch/export/distributed_writer.py
 create mode 100644 tests/gpu/torch/export/test_distributed_writer_gpu.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..195a4d75c44
--- /dev/null
+++ b/modelopt/torch/export/distributed_writer.py
@@ -0,0 +1,347 @@
+# 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 torch
+import torch.distributed as dist
+from torch.distributed.tensor import DTensor
+
+
+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: FSDP2 gives every parameter the
+    SAME mesh object, so without it a few hundred keys each redo the tensor-to-list conversion and
+    the range compare for one shared answer. (``mesh.mesh`` is a CPU tensor, so none of them 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 _run_to_completion(ops: list) -> None:
+    """Run the point-to-point ``ops`` and wait for all of them, then empty ``ops``.
+
+    This is what frees the send buffers, copies of everything this rank sends, before the caller
+    assembles what it received. A ``P2POp`` keeps its tensor, and so does its request after the
+    wait: the ops go with the emptied list, and the requests never leave this call.
+    """
+    if ops:
+        for req in dist.batch_isend_irecv(ops):
+            req.wait()
+    ops.clear()
+
+
+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.
+
+    Each shard travels once, to its owner, where ``full_tensor()`` would all-gather every tensor to
+    every rank. Everything a rank owes a peer goes in one flat buffer per (peer, dtype), one
+    operation each, so the operation count is O(world x dtypes), not O(keys x peers). Both sides
+    build those buckets from the same sorted keys and owner map, so every send meets a receive of
+    the same length. ``all_shapes[r][k]`` is rank ``r``'s dim-0 extent of key ``k``, whose shards
+    the caller has checked tile the global tensor.
+    """
+
+    def _numel(shape):
+        n = 1
+        for s in shape:
+            n *= s
+        return n
+
+    # A chunk bounds how many keys one round of transfers carries, not their bytes: within a chunk
+    # the op count is O(world x dtypes) however many keys it holds, and a model of fewer keys than a
+    # chunk moves in one round.
+    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 = []
+        pending: list = []
+        for peer, dt in sorted(send_groups, key=lambda kd: (kd[0], str(kd[1]))):
+            ks = sorted(send_groups[(peer, dt)])
+            # Packed straight into its op: no name here then holds the copy past the transfer.
+            ops.append(
+                dist.P2POp(
+                    dist.isend, torch.cat([local_sd[k].to_local().reshape(-1) for k in ks]), 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=local_sd[ks[0]].to_local().device)
+            pending.append((srcr, ks, buf))
+            ops.append(dist.P2POp(dist.irecv, buf, srcr))
+
+        _run_to_completion(ops)
+
+        # 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 pending, parts_by_key
+
+
+def _materialize_owned(
+    local_sd: dict, expert_key_set: set, rank: int, world: int, together: list = ()
+) -> None:
+    """Make every tensor in ``local_sd`` whole and owned by exactly one rank, in place.
+
+    Experts are already whole, plain and disjoint per rank; one still a DTensor is refused. Each
+    dense key goes to one owner, largest first onto the lightest rank (seeded with each rank's expert
+    bytes), and keys this rank does not own are dropped. Each group in ``together`` (the same on
+    every rank, which is checked) goes to one owner as a unit, with any group it shares a key with,
+    since the reverse conversion merges a group only where it is whole. Afterwards nothing is a
+    DTensor or partial. A rank holds at most about twice its owned set: the shards it receives and
+    the tensors assembled from them.
+    """
+    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 = {}
+
+    # Placements are local metadata: classify first, so the shard shapes ride in the one metadata
+    # collective below, which also catches a rank that classified differently.
+    replicated, shard0, unsupported, plain, empty0 = [], [], [], [], []
+    for k in dense_keys:
+        v = local_sd[k]
+        if not isinstance(v, DTensor):
+            plain.append(k)
+        elif all(p.is_replicate() for p in v.placements):
+            replicated.append(k)
+        elif _is_full_world_shard0(v, world, _mesh_memo):
+            # An empty global tensor cannot be reassembled: FSDP2 still hands one rank a padded
+            # 1-row chunk, which concatenating would turn into a row. Its owner creates it instead.
+            (empty0 if v.numel() == 0 else shard0).append(k)
+        else:
+            unsupported.append((k, str(v.placements), v.device_mesh.ndim))
+
+    # A key outside ``together`` is a unit of its own, which keeps the order of a model without merge
+    # groups. Groups sharing a key are one unit, whatever their order: split, one of them would be a
+    # merge its owner cannot complete.
+    unit_of = {k: (k,) for k in dense_keys}
+    for group in together:
+        unit = tuple(sorted({j for k in group if k in unit_of for j in unit_of[k]}))
+        for k in unit:
+            unit_of[k] = unit
+    merged_units = sorted({u for u in unit_of.values() if len(u) > 1})
+
+    # One collective for all the metadata the steps below need. Every refusal below is decided from
+    # it, so it is the same on every rank: one that raised alone would leave its peers waiting.
+    dense_sizes = {k: _nbytes(local_sd[k]) for k in dense_keys}
+    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]],
+        {k: str(local_sd[k].dtype) for k in shard0},
+        [k for k in expert_keys if isinstance(local_sd[k], DTensor)][:3],
+        dense_sizes,
+        merged_units,
+    )
+    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}"
+        )
+    # Experts skip the classification above and go to the write as they are, which cannot serialize
+    # a DTensor: refuse one here, before any rank moves or writes anything.
+    sharded_experts = {r: m[6] for r, m in enumerate(meta) if m[6]}
+    if sharded_experts:
+        raise RuntimeError(
+            f"rank {rank}: expert tensor(s) are still DTensors, but each rank writes its experts as "
+            f"it holds them; first few per rank: {sharded_experts}"
+        )
+    expert_bytes = [m[1] for m in meta]
+    all_shapes = [m[2] for m in meta]
+    global_shapes = [m[3] for m in meta]
+    dtypes = [m[5] for m in meta]
+
+    # all_shapes[r]'s keys are rank r's dim-0 classification: name a rank that disagrees, which would
+    # otherwise surface as a bare KeyError below.
+    _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 the shards rebuilds the tensor only if they tile it: dim-0 extents summing to
+    # the global dim 0, the other dims equal. A padded chunk from uneven sharding would concatenate
+    # into a wrong tensor rather than fail, so check, against the gathered global shapes.
+    untiled, mixed = [], []
+    for k in shard0:
+        # A shard is sent in the sender's dtype and received in the owner's: ranks that disagree
+        # would mismatch the transfer's byte counts and hang in it.
+        if len({dtypes[r][k] for r in range(world)}) > 1:
+            mixed.append((k, [dtypes[r][k] for r in range(world)][:3]))
+        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 mixed:
+        raise RuntimeError(
+            f"rank {rank}: {len(mixed)} dim-0 sharded tensor(s) have a different dtype on different "
+            f"ranks, so their shards cannot be moved to one rank; first few: {mixed[:2]}"
+        )
+    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]}"
+        )
+    # Every rank assigns the owners below from these sizes, so ranks that disagree on one (a plain
+    # buffer left in another dtype on one rank) would disagree on who keeps which tensor: one could
+    # be dropped everywhere or written twice, or a shard sent to a rank not waiting for it.
+    resized = {
+        r: sorted(k for k in dense_keys if m[7][k] != dense_sizes[k])[:3]
+        for r, m in enumerate(meta)
+        if m[7] != dense_sizes
+    }
+    if resized:
+        raise RuntimeError(
+            f"rank {rank}: dense tensor(s) have a different size on different ranks, so the ranks "
+            f"would not agree on which of them writes each; first few per rank: {resized}"
+        )
+    # So would ranks that group the merged keys differently.
+    regrouped = {
+        r: sorted(set(map(tuple, m[8])) ^ set(merged_units))[:3]
+        for r, m in enumerate(meta)
+        if m[8] != merged_units
+    }
+    if regrouped:
+        raise RuntimeError(
+            f"rank {rank}: ranks group the keys a merge needs on one rank differently, so the ranks "
+            "would not agree on which of them writes each; the units that differ from this rank's, "
+            f"first few per rank: {regrouped}"
+        )
+
+    # Deterministic on every rank: the sort and the argmin both break ties explicitly.
+    unit_bytes = {u: sum(dense_sizes[k] for k in u) for u in set(unit_of.values())}
+    loads = list(expert_bytes)
+    owner: dict = {}
+    for unit in sorted(unit_bytes, key=lambda u: (-unit_bytes[u], u)):
+        r = min(range(world), key=lambda i: (loads[i], i))
+        for k in unit:
+            owner[k] = r
+        loads[r] += unit_bytes[unit]
+
+    # 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=_v.to_local().device)
+        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/gpu/torch/export/test_distributed_writer_gpu.py b/tests/gpu/torch/export/test_distributed_writer_gpu.py
new file mode 100644
index 00000000000..2adc7ea1d69
--- /dev/null
+++ b/tests/gpu/torch/export/test_distributed_writer_gpu.py
@@ -0,0 +1,89 @@
+# 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 no-gather writer's transport and shard writing over NCCL, on every visible GPU."""
+
+import torch
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import DTensor, Replicate, Shard
+
+from modelopt.torch.export.distributed_writer import _materialize_owned
+
+
+def _raw_bytes(tensor):
+    return tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu()
+
+
+def _full_tensors(world: int) -> dict:
+    """What an export moves: BF16 and FP8 weights and packed NVFP4 bytes, sharded unevenly over the
+    world, one with a single row and one with none, a replicated scale and a plain buffer."""
+    torch.manual_seed(0)
+    rows = 2 * world + 1
+    return {
+        "bf16.weight": torch.randn(rows, 64, dtype=torch.bfloat16),
+        "fp8.weight": torch.randn(rows, 32).to(torch.float8_e4m3fn),
+        "nvfp4.weight": torch.randint(0, 256, (rows, 16), dtype=torch.uint8),
+        "one_row.weight": torch.randn(1, 8, dtype=torch.bfloat16),
+        "no_rows.weight": torch.empty(0, 8, dtype=torch.bfloat16),
+        **{f"layers.{i}.weight": torch.randn(4, 16, dtype=torch.bfloat16) for i in range(8)},
+        "scale": torch.tensor(0.5),
+        "buffer": torch.arange(6.0),
+    }
+
+
+def _sharded(full: dict, rank: int, world: int) -> dict:
+    """``full`` as this rank holds it: dim-0 shards in ``torch.chunk``'s layout, as FSDP2 cuts them,
+    built locally so the setup itself moves nothing; the scale replicated, the buffer plain."""
+    mesh = init_device_mesh("cuda", (world,))
+    sd = {}
+    for key, tensor in full.items():
+        tensor = tensor.cuda()
+        if key == "buffer":
+            sd[key] = tensor
+        elif key == "scale":
+            sd[key] = DTensor.from_local(tensor, mesh, [Replicate()], run_check=False)
+        else:
+            chunks = tensor.chunk(world) if tensor.shape[0] else ()
+            local = chunks[rank] if rank < len(chunks) else tensor[:0]
+            sd[key] = DTensor.from_local(
+                local.contiguous(),
+                mesh,
+                [Shard(0)],
+                run_check=False,
+                shape=tensor.shape,
+                stride=tensor.stride(),
+            )
+    return sd
+
+
+def _materialize_job(rank, world_size):
+    full = _full_tensors(world_size)
+    sd = _sharded(full, rank, world_size)
+
+    _materialize_owned(sd, set(), rank, world_size)
+
+    for key, tensor in sd.items():
+        assert not isinstance(tensor, DTensor) and tensor.is_cuda, key
+        assert (tensor.dtype, tuple(tensor.shape)) == (full[key].dtype, tuple(full[key].shape)), key
+        assert torch.equal(_raw_bytes(tensor), _raw_bytes(full[key])), key
+    owned: list = [None] * world_size
+    torch.distributed.all_gather_object(owned, sorted(sd))
+    assert sorted(key for keys in owned for key in keys) == sorted(full)
+
+
+def test_materialize_owned_moves_uneven_and_empty_shards_over_nccl(need_2_gpus, dist_workers):
+    """Every key ends whole, byte for byte, on exactly one rank: BF16, FP8 and uint8 shards of
+    uneven extent, empty ones included, a tensor with no rows, a replicated and a plain one."""
+    dist_workers.run(_materialize_job)
diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py
index 852412d2565..d751079bfd2 100644
--- a/tests/unit/torch/conftest.py
+++ b/tests/unit/torch/conftest.py
@@ -16,6 +16,8 @@
 from collections import defaultdict
 
 import pytest
+import torch
+from torch.distributed.device_mesh import init_device_mesh
 
 from modelopt.torch.opt.mode import _ModeRegistryCls
 
@@ -52,3 +54,15 @@ def use_channel_div_4():
         for name, kv_for_all in lookup.items():
             for regex, (k, v) in kv_for_all.items():
                 config_class.model_fields[name].default[regex][k] = v
+
+
+@pytest.fixture
+def gloo_world_of_one(skip_on_windows, tmp_path_factory):
+    """A one-rank gloo process group, so the code under test takes its distributed path, and its
+    CPU device mesh."""
+    store = tmp_path_factory.mktemp("pg") / "store"
+    torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1)
+    try:
+        yield init_device_mesh("cpu", (1,))
+    finally:
+        torch.distributed.destroy_process_group()
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..a081ae4121f
--- /dev/null
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -0,0 +1,245 @@
+# 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 or mis-assembled from its shards. These tests
+drive each piece directly, on CPU.
+"""
+
+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
+from torch.multiprocessing.reductions import StorageWeakRef
+
+pytest.importorskip("safetensors")
+pytest.importorskip("transformers")
+
+from modelopt.torch.export.distributed_writer import (
+    _is_full_world_shard0,
+    _materialize_owned,
+    _nbytes,
+)
+
+# ------------------------------------------------------------------ pure helpers
+
+
+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"),
+        # Empty, so there is nothing to assemble, but in a layout it would refuse a non-empty one in.
+        (
+            lambda mesh: distribute_tensor(torch.zeros(0, 4), mesh, [Shard(1)]),
+            "cannot materialize",
+        ),
+    ],
+    ids=["shard1", "2d-mesh", "padded", "empty-shard1"],
+)
+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)])
+    elif case == "size":
+        # A plain buffer rank 1 holds in another dtype: the ranks would assign owners by different sizes.
+        w = torch.zeros(2, dtype=torch.float32 if rank == 0 else torch.bfloat16)
+    elif case == "grouping":
+        # Only rank 1 keeps w with another key: the ranks would assign owners by different units.
+        w = torch.zeros(2)
+    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)
+        elif case == "dtype":
+            # Rank 1 cast its shard: the transfer's byte counts would not match.
+            dtype = torch.float32 if rank == 0 else torch.bfloat16
+            local, placement, gshape = torch.zeros(2, 4, dtype=dtype), Shard(0), (4, 4)
+        elif case == "global-shape":
+            # 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)
+        else:
+            # An expert, which each rank writes as it holds it: only rank 1's is still a DTensor.
+            local, placement, gshape = torch.zeros(2, 4), Shard(0), (4, 4)
+        w = DTensor.from_local(
+            local, mesh, [placement], run_check=False, shape=torch.Size(gshape), stride=(4, 1)
+        )
+        if case == "expert" and rank == 0:
+            w = w.to_local()
+    local_sd = {"w": w, "b": torch.zeros(1)} if case == "grouping" else {"w": w}
+    together = [["w", "b"]] if case == "grouping" and rank == 1 else ()
+    try:
+        with bounded_collectives():
+            _materialize_owned(
+                local_sd, {"w"} if case == "expert" else set(), rank, world_size, together
+            )
+    except RuntimeError as error:
+        raised = str(error)
+    else:
+        raised = "no error"
+    match = {
+        "global-shape": "do not tile their global shape",
+        "dtype": "different dtype",
+        "expert": "still DTensors",
+        "size": "different size",
+        "grouping": "group the keys",
+    }.get(case, "cannot materialize")
+    assert match in raised, raised
+
+
+@pytest.mark.parametrize(
+    "case", ["smaller-mesh", "placement", "global-shape", "dtype", "expert", "size", "grouping"]
+)
+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 _materialize_two_ranks_job(rank, world_size):
+    mesh = init_device_mesh("cpu", (world_size,))
+    torch.manual_seed(0)
+    full = {
+        "a.weight": torch.randn(7, 4),  # uneven: 4 + 3 rows
+        "b.weight": torch.randn(8, 3).to(torch.bfloat16),  # a second dtype between the same ranks
+        "c.weight": torch.randn(2, 5),
+        "d.weight": torch.randn(1, 6),  # rank 1's shard is empty
+    }
+    # Rank 0 owns the larger and holds all of its rows: it assembles having received nothing, so
+    # what it sent is all it could still be holding.
+    send_only = {"e.weight": torch.randn(1, 60), "f.weight": torch.randn(7, 4)}
+    # The copies this rank sends, and how many of them are still held at each assembled store.
+    sent, held = [], []
+    batch_isend_irecv = torch.distributed.batch_isend_irecv
+
+    def recording_sends(ops):
+        isends = [op for op in ops if op.op is torch.distributed.isend]
+        sent.extend(StorageWeakRef(op.tensor.untyped_storage()) for op in isends)
+        return batch_isend_irecv(ops)
+
+    class CountingHeldSends(dict):
+        def __setitem__(self, key, value):
+            held.append(sum(not ref.expired() for ref in sent))
+            super().__setitem__(key, value)
+
+    for layout in (full, send_only):
+        sd = CountingHeldSends(
+            {k: distribute_tensor(v, mesh, [Shard(0)]) for k, v in layout.items()}
+        )
+
+        with (
+            bounded_collectives(),
+            mock.patch.object(torch.distributed, "batch_isend_irecv", recording_sends),
+        ):
+            _materialize_owned(sd, set(), rank, world_size)
+
+        for k, v in sd.items():
+            assert not isinstance(v, DTensor), k
+            assert torch.equal(v, layout[k]), k
+        owned: list = [None] * world_size
+        torch.distributed.all_gather_object(owned, sorted(sd))
+        assert sorted(k for keys in owned for k in keys) == sorted(layout)
+    assert sent and held and not any(held), held
+
+
+def test_materialize_owned_moves_each_shard_to_its_owner_over_two_ranks(cpu_dist_workers):
+    """Each key ends whole on exactly one rank, its shards packed per peer and dtype, unpacked at
+    their offsets and joined in rank order, an uneven and an empty shard included. What a rank
+    sent is freed before it assembles what it owns."""
+    cpu_dist_workers.run(_materialize_two_ranks_job)
diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py
index 764bce545c6..4c642b1644c 100644
--- a/tests/unit/torch/utils/test_meta_dcp_load.py
+++ b/tests/unit/torch/utils/test_meta_dcp_load.py
@@ -474,15 +474,6 @@ def recording_load(state_dict, **kwargs):
 # ------------------------------------------------------------------- under a process group
 
 
-@pytest.fixture
-def gloo_world_of_one(skip_on_windows, tmp_path_factory):
-    """A one-rank gloo process group, so the loader takes its distributed path."""
-    store = tmp_path_factory.mktemp("pg") / "store"
-    torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1)
-    yield
-    torch.distributed.destroy_process_group()
-
-
 def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one):
     """The EP hybrid's layout: the decoder layers are ``fully_shard``-ed on meta before the load, the
     root is not. The layers' params must stay sharded ``DTensor``s holding the checkpoint values, the

From f409c2e1cbb14aa885f5ab3ce7e1e185d6f64421 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 00:47:25 +0000
Subject: [PATCH 2/6] Test merge groups that share a key onto one owner, and a
 strided shard's refusal

The two-rank transport test gains two merge groups that share a key, sized so
that byte balancing alone would split them, and checks one rank owns all
three. A _StridedShard on dim 0 is not Shard(0) in torch>=2.10, and the
refusal test now pins that it is refused rather than joined in rank order.
The docstrings keep the constraint and drop the cost derivations, which are
in the PR description; they say the caller hands experts over disjoint, and
that the memory bound is on top of the shards, which the model keeps.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py   | 39 +++++++------------
 .../torch/export/test_distributed_writer.py   | 29 +++++++++++---
 2 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 195a4d75c44..c11852a3986 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -37,10 +37,7 @@ def _is_full_world_shard0(v, world: int, _mesh_memo: dict | None = None) -> bool
     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: FSDP2 gives every parameter the
-    SAME mesh object, so without it a few hundred keys each redo the tensor-to-list conversion and
-    the range compare for one shared answer. (``mesh.mesh`` is a CPU tensor, so none of them syncs.)
+    ``_mesh_memo`` keeps the answer per mesh, since FSDP2 gives every parameter the same mesh.
     """
     if len(v.placements) != 1 or not v.placements[0].is_shard(0):
         return False
@@ -58,9 +55,7 @@ def _is_full_world_shard0(v, world: int, _mesh_memo: dict | None = None) -> bool
 def _run_to_completion(ops: list) -> None:
     """Run the point-to-point ``ops`` and wait for all of them, then empty ``ops``.
 
-    This is what frees the send buffers, copies of everything this rank sends, before the caller
-    assembles what it received. A ``P2POp`` keeps its tensor, and so does its request after the
-    wait: the ops go with the emptied list, and the requests never leave this call.
+    Emptying the list frees the send buffers, which the ops hold, before the caller assembles.
     """
     if ops:
         for req in dist.batch_isend_irecv(ops):
@@ -73,12 +68,9 @@ def _gather_shard0_to_owners(
 ) -> None:
     """Gather each dim-0-sharded DTensor to its owner and drop it elsewhere, in place.
 
-    Each shard travels once, to its owner, where ``full_tensor()`` would all-gather every tensor to
-    every rank. Everything a rank owes a peer goes in one flat buffer per (peer, dtype), one
-    operation each, so the operation count is O(world x dtypes), not O(keys x peers). Both sides
-    build those buckets from the same sorted keys and owner map, so every send meets a receive of
-    the same length. ``all_shapes[r][k]`` is rank ``r``'s dim-0 extent of key ``k``, whose shards
-    the caller has checked tile the global tensor.
+    What a rank owes a peer goes in one flat buffer per (peer, dtype), which both sides build from
+    the same sorted keys and owner map, so every send meets a receive of its length.
+    ``all_shapes[r][k]`` is the shape of rank ``r``'s shard of ``k``; the caller checked they tile.
     """
 
     def _numel(shape):
@@ -87,9 +79,7 @@ def _numel(shape):
             n *= s
         return n
 
-    # A chunk bounds how many keys one round of transfers carries, not their bytes: within a chunk
-    # the op count is O(world x dtypes) however many keys it holds, and a model of fewer keys than a
-    # chunk moves in one round.
+    # Keys per round of transfers: a round's op count does not grow with its keys.
     chunk = 4096
     for base in range(0, len(keys), chunk):
         part = keys[base : base + chunk]
@@ -161,13 +151,14 @@ def _materialize_owned(
 ) -> None:
     """Make every tensor in ``local_sd`` whole and owned by exactly one rank, in place.
 
-    Experts are already whole, plain and disjoint per rank; one still a DTensor is refused. Each
-    dense key goes to one owner, largest first onto the lightest rank (seeded with each rank's expert
-    bytes), and keys this rank does not own are dropped. Each group in ``together`` (the same on
-    every rank, which is checked) goes to one owner as a unit, with any group it shares a key with,
-    since the reverse conversion merges a group only where it is whole. Afterwards nothing is a
-    DTensor or partial. A rank holds at most about twice its owned set: the shards it receives and
-    the tensors assembled from them.
+    Experts must come whole, plain and disjoint per rank, as the caller splits them; one still a
+    DTensor is refused. Each dense key goes to one owner, largest first onto the lightest rank
+    (seeded with each rank's expert bytes), and keys this rank does not own are dropped. Each group
+    in ``together`` (the same on every rank, which is checked) goes to one owner as a unit, with any
+    group it shares a key with, since the reverse conversion merges a group only where it is whole.
+    Afterwards nothing is a DTensor or partial. On top of the shards it is handed, which the model
+    keeps, a rank holds at most about twice its owned set: the shards it receives and the tensors
+    assembled from them.
     """
     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)
@@ -191,7 +182,7 @@ def _materialize_owned(
 
     # A key outside ``together`` is a unit of its own, which keeps the order of a model without merge
     # groups. Groups sharing a key are one unit, whatever their order: split, one of them would be a
-    # merge its owner cannot complete.
+    # merge its owner cannot complete. Expert keys stay out: a rank holds its experts whole already.
     unit_of = {k: (k,) for k in dense_keys}
     for group in together:
         unit = tuple(sorted({j for k in group if k in unit_of for j in unit_of[k]}))
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index a081ae4121f..58feb8a9571 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -27,6 +27,7 @@
 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
+from torch.distributed.tensor.placement_types import _StridedShard
 from torch.multiprocessing.reductions import StorageWeakRef
 
 pytest.importorskip("safetensors")
@@ -114,13 +115,25 @@ def _padded_shard(mesh):
             "cannot materialize",
         ),
         (_padded_shard, "do not tile their global shape"),
+        # Its shards interleave, so joining them in rank order would permute the rows.
+        (
+            lambda mesh: DTensor.from_local(
+                torch.zeros(4, 3),
+                mesh,
+                [_StridedShard(0, split_factor=2)],
+                run_check=False,
+                shape=torch.Size([4, 3]),
+                stride=(3, 1),
+            ),
+            "cannot materialize",
+        ),
         # Empty, so there is nothing to assemble, but in a layout it would refuse a non-empty one in.
         (
             lambda mesh: distribute_tensor(torch.zeros(0, 4), mesh, [Shard(1)]),
             "cannot materialize",
         ),
     ],
-    ids=["shard1", "2d-mesh", "padded", "empty-shard1"],
+    ids=["shard1", "2d-mesh", "padded", "strided-shard0", "empty-shard1"],
 )
 def test_materialize_owned_rejects_a_tensor_it_cannot_assemble(gloo_world_of_one, build, match):
     """Refused before anything moves, rather than assembled wrong."""
@@ -204,6 +217,10 @@ def _materialize_two_ranks_job(rank, world_size):
     # Rank 0 owns the larger and holds all of its rows: it assembles having received nothing, so
     # what it sent is all it could still be holding.
     send_only = {"e.weight": torch.randn(1, 60), "f.weight": torch.randn(7, 4)}
+    # Two merge groups sharing a key: one owner for all three, where their sizes alone would put x
+    # on one rank and y and z on the other.
+    grouped = {k: torch.randn(n, 4) for k, n in (("x", 6), ("y", 5), ("z", 4))}
+    together = [["x", "y"], ["y", "z"]]
     # The copies this rank sends, and how many of them are still held at each assembled store.
     sent, held = [], []
     batch_isend_irecv = torch.distributed.batch_isend_irecv
@@ -218,7 +235,7 @@ def __setitem__(self, key, value):
             held.append(sum(not ref.expired() for ref in sent))
             super().__setitem__(key, value)
 
-    for layout in (full, send_only):
+    for layout, groups in ((full, ()), (send_only, ()), (grouped, together)):
         sd = CountingHeldSends(
             {k: distribute_tensor(v, mesh, [Shard(0)]) for k, v in layout.items()}
         )
@@ -227,7 +244,7 @@ def __setitem__(self, key, value):
             bounded_collectives(),
             mock.patch.object(torch.distributed, "batch_isend_irecv", recording_sends),
         ):
-            _materialize_owned(sd, set(), rank, world_size)
+            _materialize_owned(sd, set(), rank, world_size, groups)
 
         for k, v in sd.items():
             assert not isinstance(v, DTensor), k
@@ -235,11 +252,13 @@ def __setitem__(self, key, value):
         owned: list = [None] * world_size
         torch.distributed.all_gather_object(owned, sorted(sd))
         assert sorted(k for keys in owned for k in keys) == sorted(layout)
+        if groups:
+            assert sorted(layout) in owned, owned
     assert sent and held and not any(held), held
 
 
 def test_materialize_owned_moves_each_shard_to_its_owner_over_two_ranks(cpu_dist_workers):
     """Each key ends whole on exactly one rank, its shards packed per peer and dtype, unpacked at
-    their offsets and joined in rank order, an uneven and an empty shard included. What a rank
-    sent is freed before it assembles what it owns."""
+    their offsets and joined in rank order, an uneven and an empty shard included, and merge groups
+    that share a key on one rank. What a rank sent is freed before it assembles what it owns."""
     cpu_dist_workers.run(_materialize_two_ranks_job)

From 138a74c6c94b7f699381b129ac80302e6fe37074 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 06:30:51 +0000
Subject: [PATCH 3/6] Refuse dim-0 shards on two device types and whole tensors
 the ranks see unlike, and free each round's receive buffers

A peer's shards of one dtype travel in one buffer and a round's transfers in
one batch, so shards on two device types, as a CPU-offloaded one beside one
the export moved to the accelerator, made the sender's torch.cat raise while
its owner waited. The metadata collective now carries each shard's device
with its dtype, and every rank refuses a key whose ranks disagree or a set
of shards on more than one device type, before anything moves.

The ranks also compared only the byte counts of whole (plain, replicated
and empty) tensors, so two ranks holding one at the same size but another
dtype or shape went on to write the owner's version. They now compare every
dense tensor's dtype and shape.

A round's receive buffers stayed referenced, by the loop's last buffer and
the last owned key's parts, into the next round. Each round runs in a
helper now, so they are freed when it returns; the transfer test runs rounds
of two keys and checks it.

The batched transfer is never the default group's first collective, which
NCCL needs when not every rank takes part: the metadata collective runs on
it first.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py   | 188 ++++++++++--------
 .../torch/export/test_distributed_writer.py   |  71 ++++++-
 2 files changed, 163 insertions(+), 96 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index c11852a3986..3c77a042ebe 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -63,6 +63,17 @@ def _run_to_completion(ops: list) -> None:
     ops.clear()
 
 
+# Keys per round of transfers: a round's op count does not grow with its keys.
+_KEYS_PER_ROUND = 4096
+
+
+def _numel(shape) -> int:
+    n = 1
+    for s in shape:
+        n *= s
+    return n
+
+
 def _gather_shard0_to_owners(
     local_sd: dict, keys: list, owner: dict, all_shapes: list, rank: int, world: int
 ) -> None:
@@ -72,78 +83,75 @@ def _gather_shard0_to_owners(
     the same sorted keys and owner map, so every send meets a receive of its length.
     ``all_shapes[r][k]`` is the shape of rank ``r``'s shard of ``k``; the caller checked they tile.
     """
+    for base in range(0, len(keys), _KEYS_PER_ROUND):
+        _transfer_round(
+            local_sd, keys[base : base + _KEYS_PER_ROUND], owner, all_shapes, rank, world
+        )
 
-    def _numel(shape):
-        n = 1
-        for s in shape:
-            n *= s
-        return n
-
-    # Keys per round of transfers: a round's op count does not grow with its keys.
-    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 = []
-        pending: list = []
-        for peer, dt in sorted(send_groups, key=lambda kd: (kd[0], str(kd[1]))):
-            ks = sorted(send_groups[(peer, dt)])
-            # Packed straight into its op: no name here then holds the copy past the transfer.
-            ops.append(
-                dist.P2POp(
-                    dist.isend, torch.cat([local_sd[k].to_local().reshape(-1) for k in ks]), 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=local_sd[ks[0]].to_local().device)
-            pending.append((srcr, ks, buf))
-            ops.append(dist.P2POp(dist.irecv, buf, srcr))
-
-        _run_to_completion(ops)
-
-        # 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 = []
+
+def _transfer_round(
+    local_sd: dict, part: list, owner: dict, all_shapes: list, rank: int, world: int
+) -> None:
+    """One round of :func:`_gather_shard0_to_owners`, whose buffers are freed when it returns."""
+    # (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:
-                    pieces.append(local_sd[k].to_local())
-                elif r in got:
-                    pieces.append(got[r])
-            local_sd[k] = torch.cat(pieces, dim=0)
-        del pending, parts_by_key
+                if r != rank and _numel(all_shapes[r][k]):
+                    recv_groups.setdefault((r, dt), []).append(k)
+
+    ops: list = []
+    pending: list = []
+    for peer, dt in sorted(send_groups, key=lambda kd: (kd[0], str(kd[1]))):
+        ks = sorted(send_groups[(peer, dt)])
+        # Packed straight into its op: no name here then holds the copy past the transfer.
+        ops.append(
+            dist.P2POp(
+                dist.isend, torch.cat([local_sd[k].to_local().reshape(-1) for k in ks]), 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=local_sd[ks[0]].to_local().device)
+        pending.append((srcr, ks, buf))
+        ops.append(dist.P2POp(dist.irecv, buf, srcr))
+
+    _run_to_completion(ops)
+
+    # 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)
 
 
 def _materialize_owned(
@@ -193,15 +201,16 @@ def _materialize_owned(
     # One collective for all the metadata the steps below need. Every refusal below is decided from
     # it, so it is the same on every rank: one that raised alone would leave its peers waiting.
     dense_sizes = {k: _nbytes(local_sd[k]) for k in dense_keys}
+    dense_kinds = {k: (str(local_sd[k].dtype), tuple(local_sd[k].shape)) for k in dense_keys}
     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]],
-        {k: str(local_sd[k].dtype) for k in shard0},
+        {k: (str(local_sd[k].dtype), local_sd[k].to_local().device.type) for k in shard0},
         [k for k in expert_keys if isinstance(local_sd[k], DTensor)][:3],
-        dense_sizes,
+        dense_kinds,
         merged_units,
     )
     meta: list = [None] * world
@@ -236,7 +245,7 @@ def _materialize_owned(
     expert_bytes = [m[1] for m in meta]
     all_shapes = [m[2] for m in meta]
     global_shapes = [m[3] for m in meta]
-    dtypes = [m[5] for m in meta]
+    shard_kinds = [m[5] for m in meta]
 
     # all_shapes[r]'s keys are rank r's dim-0 classification: name a rank that disagrees, which would
     # otherwise surface as a bare KeyError below.
@@ -258,9 +267,9 @@ def _materialize_owned(
     untiled, mixed = [], []
     for k in shard0:
         # A shard is sent in the sender's dtype and received in the owner's: ranks that disagree
-        # would mismatch the transfer's byte counts and hang in it.
-        if len({dtypes[r][k] for r in range(world)}) > 1:
-            mixed.append((k, [dtypes[r][k] for r in range(world)][:3]))
+        # would mismatch the transfer's byte counts and hang in it. So would its device.
+        if len({shard_kinds[r][k] for r in range(world)}) > 1:
+            mixed.append((k, [shard_kinds[r][k] for r in range(world)][:3]))
         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)]
@@ -271,8 +280,16 @@ def _materialize_owned(
             untiled.append((k, sorted(gshapes), shard_shapes[:3]))
     if mixed:
         raise RuntimeError(
-            f"rank {rank}: {len(mixed)} dim-0 sharded tensor(s) have a different dtype on different "
-            f"ranks, so their shards cannot be moved to one rank; first few: {mixed[:2]}"
+            f"rank {rank}: {len(mixed)} dim-0 sharded tensor(s) have a different dtype or device on "
+            f"different ranks, so their shards cannot be moved to one rank; first few: {mixed[:2]}"
+        )
+    # A peer's shards of one dtype travel in one buffer, and a round's transfers in one batch: neither
+    # spans devices, as a CPU-offloaded shard beside one moved to the accelerator would need.
+    devices = sorted({shard_kinds[rank][k][1] for k in shard0})
+    if len(devices) > 1:
+        raise RuntimeError(
+            f"rank {rank}: dim-0 sharded tensors are on more than one device type ({devices}), "
+            "which the transport cannot move in one batch; put them on one device first."
         )
     if untiled:
         raise RuntimeError(
@@ -283,16 +300,17 @@ def _materialize_owned(
         )
     # Every rank assigns the owners below from these sizes, so ranks that disagree on one (a plain
     # buffer left in another dtype on one rank) would disagree on who keeps which tensor: one could
-    # be dropped everywhere or written twice, or a shard sent to a rank not waiting for it.
-    resized = {
-        r: sorted(k for k in dense_keys if m[7][k] != dense_sizes[k])[:3]
+    # be dropped everywhere or written twice, or a shard sent to a rank not waiting for it. A dtype or
+    # shape that differs at the same size would leave the owner's version in the checkpoint.
+    reshaped = {
+        r: sorted(k for k in dense_keys if m[7][k] != dense_kinds[k])[:3]
         for r, m in enumerate(meta)
-        if m[7] != dense_sizes
+        if m[7] != dense_kinds
     }
-    if resized:
+    if reshaped:
         raise RuntimeError(
-            f"rank {rank}: dense tensor(s) have a different size on different ranks, so the ranks "
-            f"would not agree on which of them writes each; first few per rank: {resized}"
+            f"rank {rank}: dense tensor(s) have a different dtype or shape on different ranks, so "
+            f"the ranks would not agree on which of them writes each; first few per rank: {reshaped}"
         )
     # So would ranks that group the merged keys differently.
     regrouped = {
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index 58feb8a9571..45058290cd9 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -39,6 +39,8 @@
     _nbytes,
 )
 
+_WRITER = "modelopt.torch.export.distributed_writer"
+
 # ------------------------------------------------------------------ pure helpers
 
 
@@ -141,14 +143,25 @@ def test_materialize_owned_rejects_a_tensor_it_cannot_assemble(gloo_world_of_one
         _materialize_owned({"w": build(gloo_world_of_one)}, set(), 0, 1)
 
 
+# A whole tensor rank 1 holds unlike rank 0: the ranks would assign owners by different sizes, or at
+# an equal size write the owner's version of it.
+_PLAIN_DISAGREEMENTS = {
+    "size": lambda rank: torch.zeros(2, dtype=torch.float32 if rank == 0 else torch.bfloat16),
+    "same-size-dtype": lambda rank: torch.zeros(
+        2, dtype=torch.float32 if rank == 0 else torch.int32
+    ),
+    "same-size-shape": lambda rank: torch.zeros((2, 4) if rank == 0 else (4, 2)),
+    "empty-shape": lambda rank: torch.zeros((0, 4) if rank == 0 else (0, 2)),
+}
+
+
 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)])
-    elif case == "size":
-        # A plain buffer rank 1 holds in another dtype: the ranks would assign owners by different sizes.
-        w = torch.zeros(2, dtype=torch.float32 if rank == 0 else torch.bfloat16)
+    elif case in _PLAIN_DISAGREEMENTS:
+        w = _PLAIN_DISAGREEMENTS[case](rank)
     elif case == "grouping":
         # Only rank 1 keeps w with another key: the ranks would assign owners by different units.
         w = torch.zeros(2)
@@ -167,6 +180,14 @@ def _refusing_materialize_job(rank, world_size, case):
         elif case == "global-shape":
             # 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)
+        elif case in ("device", "devices"):
+            # Only rank 1's shard on another device; or both ranks' shards of a second key.
+            on_meta = case == "device" and rank == 1
+            local, placement, gshape = (
+                torch.zeros(2, 4, device="meta" if on_meta else "cpu"),
+                Shard(0),
+                (4, 4),
+            )
         else:
             # An expert, which each rank writes as it holds it: only rank 1's is still a DTensor.
             local, placement, gshape = torch.zeros(2, 4), Shard(0), (4, 4)
@@ -176,6 +197,15 @@ def _refusing_materialize_job(rank, world_size, case):
         if case == "expert" and rank == 0:
             w = w.to_local()
     local_sd = {"w": w, "b": torch.zeros(1)} if case == "grouping" else {"w": w}
+    if case == "devices":
+        local_sd["m"] = DTensor.from_local(
+            torch.zeros(2, 4, device="meta"),
+            mesh,
+            [Shard(0)],
+            run_check=False,
+            shape=(4, 4),
+            stride=(4, 1),
+        )
     together = [["w", "b"]] if case == "grouping" and rank == 1 else ()
     try:
         with bounded_collectives():
@@ -188,16 +218,29 @@ def _refusing_materialize_job(rank, world_size, case):
         raised = "no error"
     match = {
         "global-shape": "do not tile their global shape",
-        "dtype": "different dtype",
+        "dtype": "different dtype or device",
+        "device": "different dtype or device",
+        "devices": "more than one device type",
         "expert": "still DTensors",
-        "size": "different size",
         "grouping": "group the keys",
+        **dict.fromkeys(_PLAIN_DISAGREEMENTS, "different dtype or shape"),
     }.get(case, "cannot materialize")
     assert match in raised, raised
 
 
 @pytest.mark.parametrize(
-    "case", ["smaller-mesh", "placement", "global-shape", "dtype", "expert", "size", "grouping"]
+    "case",
+    [
+        "smaller-mesh",
+        "placement",
+        "global-shape",
+        "dtype",
+        "device",
+        "devices",
+        "expert",
+        *_PLAIN_DISAGREEMENTS,
+        "grouping",
+    ],
 )
 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
@@ -221,13 +264,16 @@ def _materialize_two_ranks_job(rank, world_size):
     # on one rank and y and z on the other.
     grouped = {k: torch.randn(n, 4) for k, n in (("x", 6), ("y", 5), ("z", 4))}
     together = [["x", "y"], ["y", "z"]]
-    # The copies this rank sends, and how many of them are still held at each assembled store.
-    sent, held = [], []
+    # The copies this rank sends, and how many of them are still held at each assembled store; the
+    # buffers it receives into, and how many of them are still held when the next round starts.
+    sent, held, received, stale = [], [], [], []
     batch_isend_irecv = torch.distributed.batch_isend_irecv
 
     def recording_sends(ops):
-        isends = [op for op in ops if op.op is torch.distributed.isend]
-        sent.extend(StorageWeakRef(op.tensor.untyped_storage()) for op in isends)
+        stale.append(sum(not ref.expired() for ref in received))
+        for op in ops:
+            mine = sent if op.op is torch.distributed.isend else received
+            mine.append(StorageWeakRef(op.tensor.untyped_storage()))
         return batch_isend_irecv(ops)
 
     class CountingHeldSends(dict):
@@ -243,6 +289,7 @@ def __setitem__(self, key, value):
         with (
             bounded_collectives(),
             mock.patch.object(torch.distributed, "batch_isend_irecv", recording_sends),
+            mock.patch(f"{_WRITER}._KEYS_PER_ROUND", 2),  # rounds of two keys
         ):
             _materialize_owned(sd, set(), rank, world_size, groups)
 
@@ -255,10 +302,12 @@ def __setitem__(self, key, value):
         if groups:
             assert sorted(layout) in owned, owned
     assert sent and held and not any(held), held
+    assert received and len(stale) > 1 and not any(stale), stale
 
 
 def test_materialize_owned_moves_each_shard_to_its_owner_over_two_ranks(cpu_dist_workers):
     """Each key ends whole on exactly one rank, its shards packed per peer and dtype, unpacked at
     their offsets and joined in rank order, an uneven and an empty shard included, and merge groups
-    that share a key on one rank. What a rank sent is freed before it assembles what it owns."""
+    that share a key on one rank. What a rank sent is freed before it assembles what it owns, and what
+    it received before the next round of transfers starts."""
     cpu_dist_workers.run(_materialize_two_ranks_job)

From bb80474fe8f9ab525431a859ffb1760b23a495fa Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 16:24:06 +0000
Subject: [PATCH 4/6] Bring every rank into NCCL before a transfer only some
 ranks take part in

Batched P2P that is NCCL's first call on a group needs every rank of it,
and on the cpu:gloo,cuda:nccl group the setup creates, the metadata
collective runs on Gloo and leaves NCCL uninitialized. A first transfer
whose rows sit on some ranks only (two rows over four ranks) left the
others out of NCCL's initialization. Every rank now runs a one-element
all-reduce on the transport's device before the first transfer.

The new GPU test spawns a fresh group on three or more GPUs, runs no
collective on it first, and transfers a two-row tensor; a rank stuck in
NCCL's initialization exits through a watchdog, so it fails instead of
hanging.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py   |  3 ++
 .../export/test_distributed_writer_gpu.py     | 44 +++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 3c77a042ebe..4579fccf2ce 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -353,4 +353,7 @@ def _materialize_owned(
             del local_sd[k]
 
     if shard0:
+        # Batched P2P that is NCCL's first call on the group needs every rank in it, and the object
+        # collective above runs on Gloo where the group has it: every rank initializes NCCL here.
+        dist.all_reduce(torch.zeros(1, device=local_sd[shard0[0]].to_local().device))
         _gather_shard0_to_owners(local_sd, shard0, owner, all_shapes, rank, world)
diff --git a/tests/gpu/torch/export/test_distributed_writer_gpu.py b/tests/gpu/torch/export/test_distributed_writer_gpu.py
index 2adc7ea1d69..ae1b8c3621d 100644
--- a/tests/gpu/torch/export/test_distributed_writer_gpu.py
+++ b/tests/gpu/torch/export/test_distributed_writer_gpu.py
@@ -15,7 +15,12 @@
 
 """The no-gather writer's transport and shard writing over NCCL, on every visible GPU."""
 
+import os
+import threading
+
+import pytest
 import torch
+from _test_utils.torch.distributed.utils import spawn_multiprocess_job
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.tensor import DTensor, Replicate, Shard
 
@@ -87,3 +92,42 @@ def test_materialize_owned_moves_uneven_and_empty_shards_over_nccl(need_2_gpus,
     """Every key ends whole, byte for byte, on exactly one rank: BF16, FP8 and uint8 shards of
     uneven extent, empty ones included, a tensor with no rows, a replicated and a plain one."""
     dist_workers.run(_materialize_job)
+
+
+def _sparse_first_transfer_job(rank, world_size):
+    # A rank stuck in NCCL's initialization exits, so the test fails instead of hanging.
+    watchdog = threading.Timer(60, os._exit, args=(3,))
+    watchdog.start()
+    try:
+        mesh = init_device_mesh("cuda", (world_size,))
+        full = torch.arange(16, dtype=torch.bfloat16, device="cuda").reshape(2, 8)
+        # Rows on ranks 0 and 1 only: the other ranks have nothing to send or receive.
+        local = full[rank : rank + 1] if rank < 2 else full[:0]
+        sd = {
+            "w": DTensor.from_local(
+                local.contiguous(),
+                mesh,
+                [Shard(0)],
+                run_check=False,
+                shape=full.shape,
+                stride=full.stride(),
+            )
+        }
+
+        _materialize_owned(sd, set(), rank, world_size)
+        torch.cuda.synchronize()
+
+        assert torch.equal(sd["w"], full) if rank == 0 else not sd, sorted(sd)
+    finally:
+        watchdog.cancel()
+
+
+@pytest.mark.timeout(300)
+def test_a_first_transfer_most_ranks_sit_out_finishes_on_a_fresh_group():
+    """On a fresh ``cpu:gloo,cuda:nccl`` group, with no collective run on it before, a transfer that
+    only two of three or more ranks take part in finishes: NCCL's first batched P2P on a group needs
+    every rank of it, and the metadata collective runs on Gloo."""
+    world_size = torch.cuda.device_count()
+    if world_size < 3:
+        pytest.skip("needs 3 or more GPUs, so that a rank holds no shard")
+    spawn_multiprocess_job(world_size, _sparse_first_transfer_job, backend="nccl")

From b73c052e45f12b98a8e5945632ac069722cc6d63 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 16:24:10 +0000
Subject: [PATCH 5/6] Bound the transport's memory by what a rank sends as well
 as what it owns

The docstring promised about twice a rank's owned set, which holds only
when the outgoing shards are balanced like the owned ones. Under dim-0
sharding they need not be: four one-row tensors over four ranks put every
row on rank 0, which then packs three send buffers while it owns one
tensor. The bound now names the copies of a rank's shards of keys other
ranks own, the shards it receives, and the tensors it assembles.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 4579fccf2ce..34d5021d67c 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -165,8 +165,8 @@ def _materialize_owned(
     in ``together`` (the same on every rank, which is checked) goes to one owner as a unit, with any
     group it shares a key with, since the reverse conversion merges a group only where it is whole.
     Afterwards nothing is a DTensor or partial. On top of the shards it is handed, which the model
-    keeps, a rank holds at most about twice its owned set: the shards it receives and the tensors
-    assembled from them.
+    keeps, a rank holds in each round of transfers copies of its shards of keys other ranks own plus
+    the shards it receives, and then those received shards plus the tensors assembled from them.
     """
     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)

From 764036cd64eb2e0aaa0e28791cf7986ea9d3e983 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 17:28:21 +0000
Subject: [PATCH 6/6] Say that the transport weighs owners by bytes, not by
 free device memory

A rank's share is balanced by the bytes it owns, so a large embedding or
lm_head lands whole on one owner whatever its free memory, and must fit
there. The docstring says so (a review's note).

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 34d5021d67c..fdb3674e257 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -161,9 +161,11 @@ def _materialize_owned(
 
     Experts must come whole, plain and disjoint per rank, as the caller splits them; one still a
     DTensor is refused. Each dense key goes to one owner, largest first onto the lightest rank
-    (seeded with each rank's expert bytes), and keys this rank does not own are dropped. Each group
-    in ``together`` (the same on every rank, which is checked) goes to one owner as a unit, with any
-    group it shares a key with, since the reverse conversion merges a group only where it is whole.
+    (seeded with each rank's expert bytes), and keys this rank does not own are dropped. Ranks are
+    weighed by bytes, not free device memory: a large embedding or ``lm_head`` must fit on its owner.
+    Each group in ``together`` (the same on every rank, which is checked) goes to one owner as a
+    unit, with any group it shares a key with, since the reverse conversion merges a group only where
+    it is whole.
     Afterwards nothing is a DTensor or partial. On top of the shards it is handed, which the model
     keeps, a rank holds in each round of transfers copies of its shards of keys other ranks own plus
     the shards it receives, and then those received shards plus the tensors assembled from them.