From 1ce245add4687d86c54dbd569465fd715cdbfc84 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 19:08:52 +0000
Subject: [PATCH 1/6] Add the shard writing of the no-gather distributed HF
 checkpoint writer

[3/3] of #2751's split, on [2/3]'s transport (#3101): the reverse
conversion on each rank's whole tensors, with each merge group and its
quantization state kept on one owner (_merge_groups, _revert_whole_sd),
and the shards' writing: even files numbered into one global range,
staged as .tmp files and renamed into place once every rank has written,
the index last (_even_bins, _write_even_shards). Every step a rank can
fail alone runs in the failure-carrying all_reduce.

This collects the branch's history from before the split, when it held
the whole writer; the commits after the split follow on their own.

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

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index fdb3674e257..098a9302a1b 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -20,10 +20,55 @@
 replicated across ranks.
 """
 
+import json
+import math
+import os
+import warnings
+from collections import Counter
+from pathlib import Path
+
 import torch
 import torch.distributed as dist
+import torch.nn as nn
+from safetensors.torch import save_file
 from torch.distributed.tensor import DTensor
 
+from modelopt.torch.utils.distributed import all_reduce
+
+from .quant_aware_conversion import (
+    _QUANT_STATE_SUFFIXES,
+    QuantConversionUnsupportedError,
+    _build_reverse_rules,
+    revert_weight_conversion_quant_aware,
+)
+from .unified_export_hf_streaming import _parse_shard_size
+
+
+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."""
@@ -359,3 +404,212 @@ def _materialize_owned(
         # 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)
+
+
+def _merge_groups(model: nn.Module, keys) -> list:
+    """The keys the model's reverse conversion merges into one tensor, one list per merged tensor.
+
+    A merge needs its whole group on one rank, so :func:`_materialize_owned` keeps each together.
+    A weight's quantization state goes with it: the merge refuses a quantized weight only where it
+    sees its scales, which a weight-specific pattern does not match. A mapping that cannot be
+    reversed has none: its reverse falls back to the in-memory names anyway. Groups are matched on
+    the names before the reverse splits run: a merge whose sources a split produces is whole on one
+    rank only because a split's outputs stay with the fused key they come from.
+    """
+    try:
+        merge_rules = _build_reverse_rules(model)[1]
+    except QuantConversionUnsupportedError:
+        return []
+    key_set = set(keys)
+    groups: dict = {}
+    for i, rule in enumerate(merge_rules):
+        for k in keys:
+            target, source_pattern = rule.converter.rename_source_key(k)
+            if source_pattern is None:
+                continue
+            group = groups.setdefault((i, target), [])
+            group.append(k)
+            if k.endswith(".weight"):
+                module = k.removesuffix(".weight")
+                group += [
+                    module + leaf for leaf in _QUANT_STATE_SUFFIXES if module + leaf in key_set
+                ]
+    return list(groups.values())
+
+
+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 needs its whole group, which ownership keeps on one rank
+    (:func:`_merge_groups`); a rank holding none of a rule's groups skips the rule, and a rule no
+    rank matches is refused, as the single-process path refuses one. 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.
+    """
+    reverted = local_sd
+
+    def revert() -> tuple[str | None, list]:
+        nonlocal reverted
+        unmatched: list = []
+        try:
+            reverted = revert_weight_conversion_quant_aware(
+                model, local_sd, unmatched_merges=unmatched
+            )
+        except Exception as exc:
+            return str(exc), []
+        return None, unmatched
+
+    def combine(results: list) -> list:
+        failures = [f"rank {r}: {error}" for r, (error, _) in enumerate(results) if error]
+        return failures or sorted(set.intersection(*(set(unmatched) for _, unmatched in results)))
+
+    failures = all_reduce("reverse weight conversion", revert, combine)
+    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 as HF safetensors shards; rank 0 writes the index.
+
+    Each rank splits its tensors into even shards, numbered into one global range. Runs after the
+    reverse conversion, whose renames and splits fix the names and sizes packed here.
+
+    Every file is written under a temporary name and renamed into place only once every rank has
+    written all of its own, the index last, so a failure before then leaves the directory as it was,
+    a checkpoint already there included. Each step a rank can fail alone runs in
+    :func:`~modelopt.torch.utils.distributed.all_reduce`, which raises on every rank when one fails.
+    """
+    export_dir = Path(export_dir)
+    max_bytes = _parse_shard_size(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 plan_mine() -> list:
+        # 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), write_keys)]
+
+    planned = all_reduce("shard plan", plan_mine, _in_rank_order)
+    # The reverse conversion checks for rename collisions, but only within one rank's slice. A
+    # collision ACROSS ranks would be silently swallowed by the index's dict update and ship a
+    # checkpoint quietly missing a tensor, so every rank refuses it here, before any file is written.
+    names = Counter(k for _, keys in planned for k in keys)
+    collisions = sorted(k for k, n in names.items() if n > 1)
+    if collisions:
+        raise RuntimeError(
+            f"tensor name collision across ranks: {len(collisions)} name(s) held by more than one "
+            f"rank; first few: {collisions[:5]}"
+        )
+    counts = [n for n, _ in planned]
+    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])
+    staged: list = []  # this rank's files, each written as <name>.tmp until the commit
+    index_name = "model.safetensors.index.json"
+
+    def stage(name: str) -> Path:
+        # Recorded before the write: one that fails midway leaves a partial file behind.
+        staged.append(name)
+        return export_dir / f"{name}.tmp"
+
+    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"
+            path = stage(fname)
+            tensors: dict = {}
+            ptrs: set = set()
+            for k in keys:
+                t = local_sd[k].detach().to("cpu").contiguous()
+                # Once copied this rank has no further use for its own tensor. The writer never
+                # allocates on the device, so dropping it here does not lower the peak (the owned
+                # set at entry); it frees the tensor now instead of after the whole file is written.
+                del local_sd[k]
+                # save_file refuses tensors sharing memory, and neither call copies a tensor that is
+                # already on the CPU and contiguous, so two names for one tensor arrive aliased. Copy
+                # rather than drop one: every name has to reach the index.
+                if t.data_ptr() in ptrs:
+                    t = t.clone()
+                ptrs.add(t.data_ptr())
+                tensors[k] = t
+            save_file(tensors, str(path), metadata={"format": "pt"})
+            for k in keys:
+                weight_map[k] = fname
+            my_bytes += sum(_nbytes(t) for t in tensors.values())
+            del tensors
+        return weight_map, my_bytes
+
+    def write_index() -> bool:
+        if rank == 0:
+            index = {"metadata": {"total_size": total_bytes}, "weight_map": full_map}
+            with open(stage(index_name), "w") as f:
+                json.dump(index, f, indent=2)
+        return True
+
+    def commit(names: list) -> bool:
+        for name in names:
+            os.replace(export_dir / f"{name}.tmp", export_dir / name)
+        return True
+
+    # A failed export removes the temporary files it wrote, and only those: a retry into the
+    # directory then finds it as it was. What was written is gone from ``local_sd``, so a retry needs
+    # the state dict built again.
+    try:
+        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
+        all_reduce("index written", write_index, all)
+        # Every file is written: only now does the directory change, the index last, so until
+        # then an index already there still names its own shards. Also the closing barrier.
+        all_reduce("shards committed", lambda: commit([n for n in staged if n != index_name]), all)
+        all_reduce("index committed", lambda: commit([n for n in staged if n == index_name]), all)
+    except BaseException:
+        left = []
+        for name in staged:
+            # Best effort: the filesystem that failed the write can fail this too, and its error
+            # would replace the one naming the rank that failed, so the files it keeps are named in
+            # a warning instead. A file already renamed into place has no temporary file left.
+            try:
+                (export_dir / f"{name}.tmp").unlink(missing_ok=True)
+            except OSError:
+                left.append(f"{name}.tmp")
+        if left:
+            warnings.warn(
+                f"rank {rank}: could not remove {len(left)} temporary file(s) the failed export "
+                f"wrote to {export_dir}: {left[:5]}"
+            )
+        raise
diff --git a/tests/gpu/torch/export/test_distributed_writer_gpu.py b/tests/gpu/torch/export/test_distributed_writer_gpu.py
index ae1b8c3621d..b918d2f29e0 100644
--- a/tests/gpu/torch/export/test_distributed_writer_gpu.py
+++ b/tests/gpu/torch/export/test_distributed_writer_gpu.py
@@ -15,16 +15,18 @@
 
 """The no-gather writer's transport and shard writing over NCCL, on every visible GPU."""
 
+import json
 import os
 import threading
 
 import pytest
 import torch
 from _test_utils.torch.distributed.utils import spawn_multiprocess_job
+from safetensors.torch import load_file
 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
+from modelopt.torch.export.distributed_writer import _materialize_owned, _write_even_shards
 
 
 def _raw_bytes(tensor):
@@ -94,6 +96,31 @@ def test_materialize_owned_moves_uneven_and_empty_shards_over_nccl(need_2_gpus,
     dist_workers.run(_materialize_job)
 
 
+def _materialize_and_write_job(rank, world_size, tmp_dir):
+    full = _full_tensors(world_size)
+    sd = _sharded(full, rank, world_size)
+
+    _materialize_owned(sd, set(), rank, world_size)
+    # 256 bytes fits two of the 4x16 BF16 layers: an owner of several writes several files.
+    _write_even_shards(sd, tmp_dir, "256", rank, world_size)
+
+    assert not sd
+    index = json.loads((tmp_dir / "model.safetensors.index.json").read_text())
+    loaded = {}
+    for shard in sorted(set(index["weight_map"].values())):
+        loaded.update(load_file(str(tmp_dir / shard)))
+    assert sorted(loaded) == sorted(full)
+    for key, tensor in full.items():
+        assert torch.equal(_raw_bytes(loaded[key]), _raw_bytes(tensor)), key
+    assert not list(tmp_dir.glob("*.tmp"))
+
+
+def test_write_even_shards_writes_every_rank_s_files_over_nccl(need_2_gpus, dist_workers, tmp_path):
+    """Moved to their owners and written there, several files per rank, the checkpoint holds every
+    tensor byte for byte, and no temporary file is left."""
+    dist_workers.run(_materialize_and_write_job, tmp_dir=tmp_path)
+
+
 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,))
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index 45058290cd9..a0c2183ddf1 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -16,10 +16,14 @@
 """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.
+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
+import types
 from unittest import mock
 
 import pytest
@@ -32,18 +36,104 @@
 
 pytest.importorskip("safetensors")
 pytest.importorskip("transformers")
+from safetensors.torch import load_file, save_file
+from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig
+from transformers.core_model_loading import Chunk, WeightConverter
 
 from modelopt.torch.export.distributed_writer import (
+    _even_bins,
     _is_full_world_shard0,
     _materialize_owned,
+    _merge_groups,
     _nbytes,
+    _revert_whole_sd,
+    _write_even_shards,
 )
 
 _WRITER = "modelopt.torch.export.distributed_writer"
 
+
+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(
+    ("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(
@@ -311,3 +401,304 @@ def test_materialize_owned_moves_each_shard_to_its_owner_over_two_ranks(cpu_dist
     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)
+
+
+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(
+        f"{_WRITER}.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 _merge_group_job(rank, world_size, case):
+    # A checkpoint's fused qkv, which the model holds as three projections. Anchored to the weight,
+    # the reverse's patterns do not match the projections' scales.
+    anchor = "$" if case == "anchored" else ""
+    qkv = WeightConverter(
+        source_patterns="attn.qkv.weight",
+        target_patterns=[f"attn.{part}_proj.weight{anchor}" for part in "qkv"],
+        operations=[Chunk(dim=0)],
+    )
+    model = types.SimpleNamespace(_weight_conversions=[qkv])
+    module = "blk.mlp" if case == "unmatched" else "blk.attn"
+    # q is the largest key and k, v the smallest: placed one key at a time, the group would split
+    # over the ranks (q and v on one, k on the other).
+    full = {
+        f"{module}.q_proj.weight": torch.full((6, 4), 1.0),
+        f"{module}.k_proj.weight": torch.full((2, 4), 2.0),
+        f"{module}.v_proj.weight": torch.full((2, 4), 3.0),
+        "blk.o_proj.weight": torch.full((5, 4), 4.0),
+    }
+    if case in ("quantized", "anchored"):
+        # Quantized weights the reverse cannot merge: the owner must see their scales to refuse.
+        full.update({f"{module}.{part}_proj.weight_scale": torch.tensor(0.5) for part in "qkv"})
+    sd = dict(full)
+    with bounded_collectives():
+        _materialize_owned(sd, set(), rank, world_size, _merge_groups(model, sd))
+        if case == "merge":
+            assert _revert_whole_sd(model, sd) is True
+        elif case in ("quantized", "anchored"):
+            with pytest.warns(UserWarning, match="cannot merge quantization state"):
+                assert _revert_whole_sd(model, sd) is False
+        else:
+            with pytest.warns(UserWarning, match="matched no state-dict key"):
+                assert _revert_whole_sd(model, sd) is False
+    owned: list = [None] * world_size
+    torch.distributed.all_gather_object(owned, dict(sd))
+    written = {k: v for mine in owned for k, v in mine.items()}
+    if case in ("quantized", "anchored"):
+        # The rank holding the group, the one that refuses it, also holds the weights' scales.
+        group = {
+            f"{module}.{part}_proj.{leaf}" for part in "qkv" for leaf in ("weight", "weight_scale")
+        }
+        assert any(group <= set(mine) for mine in owned), [sorted(mine) for mine in owned]
+    if case == "merge":
+        assert sorted(written) == ["blk.attn.qkv.weight", "blk.o_proj.weight"]
+        want = torch.cat([full[f"blk.attn.{part}_proj.weight"] for part in "qkv"])
+        assert torch.equal(written["blk.attn.qkv.weight"], want)
+    else:
+        assert sorted(written) == sorted(full)
+
+
+@pytest.mark.parametrize("case", ["merge", "unmatched", "quantized", "anchored"])
+def test_revert_whole_sd_merges_a_group_ownership_keeps_on_one_rank(cpu_dist_workers, case):
+    """A merge group goes to one owner, which merges it while the other rank skips the rule, as the
+    gather export merges it. A rule no rank's keys match is refused, as the gather export refuses
+    it, and so is a quantized group, whose scales go to its owner with it, also when the merge's
+    patterns are anchored to the weights and do not match the scales."""
+    cpu_dist_workers.run(_merge_group_job, case=case)
+
+
+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_replaces_a_checkpoint_already_there(gloo_world_of_one, tmp_path):
+    """A re-export into the directory replaces the shards and index it finds there."""
+    _write_even_shards({"w": torch.zeros(2)}, tmp_path, "1MB", 0, 1)
+    tensors = {"w": torch.ones(2)}
+
+    _write_even_shards(dict(tensors), tmp_path, "1MB", 0, 1)
+
+    _assert_byte_identical(_load_checkpoint(tmp_path), tensors)
+    assert not list(tmp_path.glob("*.tmp"))
+
+
+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 test_a_failed_write_keeps_its_error_when_the_cleanup_fails_too(gloo_world_of_one, tmp_path):
+    """The filesystem that failed the write can fail the cleanup too: the write's error is the one
+    raised, not the cleanup's, and the file the cleanup could not remove is named in a warning."""
+    with (
+        mock.patch(f"{_WRITER}.save_file", side_effect=OSError("disk full")),
+        mock.patch("pathlib.Path.unlink", side_effect=OSError("read-only file system")),
+        pytest.warns(UserWarning, match="could not remove 1 temporary file"),
+        pytest.raises(OSError, match="disk full"),
+    ):
+        _write_even_shards({"w": torch.ones(2)}, tmp_path, "1MB", 0, 1)
+
+
+def test_write_even_shards_drops_each_tensor_before_writing_its_file(gloo_world_of_one, tmp_path):
+    """A rank drops its own tensor as soon as the host copy exists, so a file's tensors are not
+    held on the device while the file is written."""
+    local = {f"w{i}": torch.randn(16, 8) for i in range(4)}
+    still_held = []
+
+    def spy(tensors, *args, **kwargs):
+        still_held.extend(set(tensors) & set(local))
+        return save_file(tensors, *args, **kwargs)
+
+    with mock.patch(f"{_WRITER}.save_file", side_effect=spy):
+        _write_even_shards(local, tmp_path, "1100", 0, 1)
+
+    assert not still_held and not local
+
+
+def test_write_even_shards_writes_both_names_of_a_shared_tensor(gloo_world_of_one, tmp_path):
+    """Both names of one CPU tensor are written, though ``save_file`` refuses shared memory."""
+    shared = torch.arange(4.0)
+
+    _write_even_shards({"a": shared, "b": shared}, tmp_path, "1MB", 0, 1)
+
+    _assert_byte_identical(_load_checkpoint(tmp_path), {"a": shared, "b": shared})
+
+
+def _write_two_ranks_job(rank, world_size, tmp_dir):
+    torch.manual_seed(rank)
+    # Two of these 512-byte tensors fit a file, so rank 0 writes one file and rank 1 two: a file
+    # number off by a rank's count then lands on another rank's file, or past the total.
+    mine = {f"rank{rank}.w{i}": torch.randn(16, 8) for i in range(2 + rank)}
+    everyone: list = [None] * world_size
+    with bounded_collectives():
+        torch.distributed.all_gather_object(everyone, mine)
+        _write_even_shards(dict(mine), tmp_dir, "1100", rank, world_size)
+
+    _assert_byte_identical(
+        _load_checkpoint(tmp_dir), {k: v for m in everyone for k, v in m.items()}
+    )
+
+
+def test_write_even_shards_numbers_two_ranks_files_into_one_checkpoint(cpu_dist_workers, tmp_path):
+    """Each rank's files are numbered after those of the ranks before it, and the index rank 0
+    writes names every rank's tensors."""
+    cpu_dist_workers.run(_write_two_ranks_job, tmp_dir=tmp_path)
+
+    shards = sorted(p.name for p in tmp_path.glob("*.safetensors"))
+    assert shards == [f"model-{i:05d}-of-00003.safetensors" for i in (1, 2, 3)]
+
+
+def _one_rank_fails_job(rank, world_size, tmp_dir, fails):
+    failing_rank, target = {
+        "dir": (1, "pathlib.Path.mkdir"),
+        "shard": (1, f"{_WRITER}.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
+
+
+def _existing_checkpoint(export_dir):
+    """A checkpoint under the names the two-rank export below writes, holding other tensors."""
+    weight_map = {
+        "old.a": "model-00001-of-00002.safetensors",
+        "old.b": "model-00002-of-00002.safetensors",
+    }
+    for key, shard in weight_map.items():
+        save_file({key: torch.zeros(3)}, str(export_dir / shard))
+    index = {"metadata": {"total_size": 24}, "weight_map": weight_map}
+    (export_dir / "model.safetensors.index.json").write_text(json.dumps(index))
+    return {p.name: p.read_bytes() for p in export_dir.iterdir()}
+
+
+@pytest.mark.parametrize("existing", [False, True], ids=["empty", "over-a-checkpoint"])
+@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, existing
+):
+    """Not just on the rank that failed, while its peers wait in a collective it never joins; and
+    the directory is left as it was: no file of the failed export, and a checkpoint already there
+    byte for byte, though the export would have written its names."""
+    before = _existing_checkpoint(tmp_path) if existing else {}
+
+    cpu_dist_workers.run(_one_rank_fails_job, tmp_dir=tmp_path, fails=fails)
+
+    assert {p.name: p.read_bytes() for p in tmp_path.iterdir()} == before
+
+
+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):
+    """Before any rank writes a shard, which a retry into the directory would otherwise inherit."""
+    cpu_dist_workers.run(_cross_rank_collision_job, tmp_dir=tmp_path)
+
+    assert not list(tmp_path.glob("*.safetensors"))

From 1cd9e4036f254117481c11e54ea05229312bbe5c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 19:08:52 +0000
Subject: [PATCH 2/6] Test that the index goes in place only after every rank's
 shards

Swapping the shard and index commits kept every writer test green. A rank
that fails to rename its shard into place must leave the index as it was,
so an index in the directory never names shards that are missing. The new
test fails with the two commits swapped.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/test_distributed_writer.py   | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index a0c2183ddf1..f2063100b08 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -23,6 +23,7 @@
 import contextlib
 import json
 import math
+import os
 import types
 from unittest import mock
 
@@ -684,6 +685,36 @@ def test_a_write_that_fails_on_one_rank_fails_on_every_rank(
     assert {p.name: p.read_bytes() for p in tmp_path.iterdir()} == before
 
 
+def _shard_commit_fails_job(rank, world_size, tmp_dir):
+    # Every file is written; rank 1 then fails to rename its shard into place.
+    replace = os.replace
+
+    def failing_shard_rename(src, dst):
+        if str(dst).endswith(".safetensors"):
+            raise OSError("disk full")
+        replace(src, dst)
+
+    failure = (
+        mock.patch(f"{_WRITER}.os.replace", side_effect=failing_shard_rename)
+        if rank == 1
+        else contextlib.nullcontext()
+    )
+    with bounded_collectives(), failure, pytest.raises((OSError, RuntimeError), match="disk full"):
+        _write_even_shards({f"rank{rank}.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size)
+
+
+@pytest.mark.parametrize("existing", [False, True], ids=["empty", "over-a-checkpoint"])
+def test_the_index_goes_in_place_only_after_every_shard(cpu_dist_workers, tmp_path, existing):
+    """A rank that fails to put its shard in place leaves the index as it was, so an index in the
+    directory never names shards that are missing."""
+    index = tmp_path / "model.safetensors.index.json"
+    before = _existing_checkpoint(tmp_path)[index.name] if existing else None
+
+    cpu_dist_workers.run(_shard_commit_fails_job, tmp_dir=tmp_path)
+
+    assert (index.read_bytes() if index.exists() else None) == before
+
+
 def _cross_rank_collision_job(rank, world_size, tmp_dir):
     try:
         with bounded_collectives():

From ecd92b609214aa9267605d25c1ba35e4277b88ad Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 19:08:52 +0000
Subject: [PATCH 3/6] Remove an earlier checkpoint's weight files once a
 re-export is in place

A re-export into a directory replaced the shards it wrote under the same
names and left the rest: an earlier export's shards beyond the new count,
and a single-file model.safetensors, which transformers loads ahead of the
new index. Once every rank has committed, rank 0 removes the
model-NNNNN-of-NNNNN.safetensors files the new index does not name, as
save_pretrained does, and a model.safetensors. Nothing else is touched, and
nothing is removed on a failure.

The docstring states the commit's guarantee: a failure while writing leaves
a checkpoint already there untouched, but a rename failing partway through
the final commit can leave it mixed.

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

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 098a9302a1b..fb3944ef5e6 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -23,6 +23,7 @@
 import json
 import math
 import os
+import re
 import warnings
 from collections import Counter
 from pathlib import Path
@@ -43,6 +44,9 @@
 )
 from .unified_export_hf_streaming import _parse_shard_size
 
+# A checkpoint's weight files, sharded as save_pretrained names them or in one file.
+_CHECKPOINT_WEIGHTS = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors")
+
 
 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.
@@ -496,8 +500,10 @@ def _write_even_shards(
     reverse conversion, whose renames and splits fix the names and sizes packed here.
 
     Every file is written under a temporary name and renamed into place only once every rank has
-    written all of its own, the index last, so a failure before then leaves the directory as it was,
-    a checkpoint already there included. Each step a rank can fail alone runs in
+    written all of its own, the index last. A failure while writing leaves the directory as it was, a
+    checkpoint already there included; a rename failing partway through that final commit can leave
+    such a checkpoint mixed. Once every rank's files are in place, rank 0 removes the earlier
+    checkpoint's weight files the new index does not name. Each step a rank can fail alone runs in
     :func:`~modelopt.torch.utils.distributed.all_reduce`, which raises on every rank when one fails.
     """
     export_dir = Path(export_dir)
@@ -582,6 +588,16 @@ def commit(names: list) -> bool:
             os.replace(export_dir / f"{name}.tmp", export_dir / name)
         return True
 
+    def remove_stale() -> bool:
+        # Left in place, an earlier checkpoint's shards would outlive it, and its model.safetensors
+        # would be loaded ahead of the new index.
+        if rank == 0:
+            current = set(full_map.values())
+            for path in export_dir.iterdir():
+                if path.name not in current and _CHECKPOINT_WEIGHTS.fullmatch(path.name):
+                    path.unlink()
+        return True
+
     # A failed export removes the temporary files it wrote, and only those: a retry into the
     # directory then finds it as it was. What was written is gone from ``local_sd``, so a retry needs
     # the state dict built again.
@@ -594,7 +610,7 @@ def commit(names: list) -> bool:
             total_bytes += b
         all_reduce("index written", write_index, all)
         # Every file is written: only now does the directory change, the index last, so until
-        # then an index already there still names its own shards. Also the closing barrier.
+        # then an index already there still names its own shards.
         all_reduce("shards committed", lambda: commit([n for n in staged if n != index_name]), all)
         all_reduce("index committed", lambda: commit([n for n in staged if n == index_name]), all)
     except BaseException:
@@ -613,3 +629,5 @@ def commit(names: list) -> bool:
                 f"wrote to {export_dir}: {left[:5]}"
             )
         raise
+    # Only once every rank has committed, and never on a failure. Also the closing barrier.
+    all_reduce("stale weights removed", remove_stale, all)
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index f2063100b08..78768548e0a 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -67,10 +67,11 @@ def _assert_byte_identical(got, want):
 
 
 def _load_checkpoint(export_dir):
-    """Every tensor the index names, read from the shard the index names for it."""
+    """Every tensor the index names, read from the shard the index names for it; no other weight file
+    of the checkpoint is left, a ``model.safetensors`` included, which transformers loads first."""
     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"))
+    shards = sorted(p.name for p in export_dir.glob("model*.safetensors"))
     assert set(weight_map.values()) == set(shards), "index and shard files disagree"
     tensors = {}
     for shard in shards:
@@ -549,15 +550,29 @@ def test_write_even_shards_writes_even_shards_and_an_index_of_every_tensor(
         assert sum(_nbytes(t) for t in load_file(str(tmp_path / shard)).values()) <= max_bytes
 
 
-def test_write_even_shards_replaces_a_checkpoint_already_there(gloo_world_of_one, tmp_path):
-    """A re-export into the directory replaces the shards and index it finds there."""
-    _write_even_shards({"w": torch.zeros(2)}, tmp_path, "1MB", 0, 1)
+@pytest.mark.parametrize("earlier", ["one-shard", "three-shards", "single-file"])
+def test_write_even_shards_replaces_a_checkpoint_already_there(
+    gloo_world_of_one, tmp_path, earlier
+):
+    """A re-export into the directory replaces the checkpoint it finds there: its shards and index,
+    any shards the new index does not name, and a single ``model.safetensors``. Other files stay."""
+    if earlier == "single-file":
+        save_file({"old": torch.zeros(2)}, str(tmp_path / "model.safetensors"))
+    else:
+        n = 1 if earlier == "one-shard" else 3
+        _write_even_shards({f"old{i}": torch.zeros(256) for i in range(n)}, tmp_path, "1100", 0, 1)
+    others = {"config.json": b"{}", "mtp.safetensors": b"weights the index does not cover"}
+    for name, data in others.items():
+        (tmp_path / name).write_bytes(data)
     tensors = {"w": torch.ones(2)}
 
     _write_even_shards(dict(tensors), tmp_path, "1MB", 0, 1)
 
     _assert_byte_identical(_load_checkpoint(tmp_path), tensors)
-    assert not list(tmp_path.glob("*.tmp"))
+    assert sorted(p.name for p in tmp_path.iterdir()) == sorted(
+        ["model-00001-of-00001.safetensors", "model.safetensors.index.json", *others]
+    )
+    assert all((tmp_path / name).read_bytes() == data for name, data in others.items())
 
 
 def test_write_even_shards_fails_loudly_when_no_rank_has_anything_to_write(
@@ -658,13 +673,15 @@ def _one_rank_fails_job(rank, world_size, tmp_dir, fails):
 
 
 def _existing_checkpoint(export_dir):
-    """A checkpoint under the names the two-rank export below writes, holding other tensors."""
+    """A checkpoint under the names the two-rank export below writes, holding other tensors, and a
+    ``model.safetensors`` a successful export would remove."""
     weight_map = {
         "old.a": "model-00001-of-00002.safetensors",
         "old.b": "model-00002-of-00002.safetensors",
     }
     for key, shard in weight_map.items():
         save_file({key: torch.zeros(3)}, str(export_dir / shard))
+    save_file({"old.c": torch.zeros(3)}, str(export_dir / "model.safetensors"))
     index = {"metadata": {"total_size": 24}, "weight_map": weight_map}
     (export_dir / "model.safetensors.index.json").write_text(json.dumps(index))
     return {p.name: p.read_bytes() for p in export_dir.iterdir()}

From edbe49c486acd813b5400416d6d3e8d455b7af23 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 19:08:52 +0000
Subject: [PATCH 4/6] Count a reverse conversion that fails with no message as
 a failure

The ranks' reverse verdicts told a failure by its message, and an exception
with none, as RuntimeError() is, read as success: the other ranks then
committed hub names while that rank kept its in-memory ones. A failure is
now any rank's error, whatever it says, and the warning names the type of
one with no message. The one-rank-failure test gains that case, which fails
without the change.

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

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index fb3944ef5e6..1e161196be8 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -463,11 +463,14 @@ def revert() -> tuple[str | None, list]:
                 model, local_sd, unmatched_merges=unmatched
             )
         except Exception as exc:
-            return str(exc), []
+            # An exception may carry no message: name its type, so the warning says something.
+            return str(exc) or type(exc).__name__, []
         return None, unmatched
 
     def combine(results: list) -> list:
-        failures = [f"rank {r}: {error}" for r, (error, _) in enumerate(results) if error]
+        failures = [
+            f"rank {r}: {error}" for r, (error, _) in enumerate(results) if error is not None
+        ]
         return failures or sorted(set.intersection(*(set(unmatched) for _, unmatched in results)))
 
     failures = all_reduce("reverse weight conversion", revert, combine)
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index 78768548e0a..153a1e11f1e 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -441,29 +441,31 @@ def test_revert_whole_sd_renames_to_hub_names_in_place():
     assert torch.equal(sd[f"{hub}.experts.0.w2.weight"], w2)
 
 
-def _one_rank_cannot_revert_job(rank, world_size):
+def _one_rank_cannot_revert_job(rank, world_size, message):
     # 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(
-        f"{_WRITER}.revert_weight_conversion_quant_aware",
-        side_effect=RuntimeError("incomplete merge"),
+        f"{_WRITER}.revert_weight_conversion_quant_aware", side_effect=RuntimeError(message)
     )
     with (
         bounded_collectives(),
         failure if rank == 1 else contextlib.nullcontext(),
-        pytest.warns(UserWarning, match="rank 1: incomplete merge"),
+        pytest.warns(UserWarning, match=f"rank 1: {message or 'RuntimeError'}"),
     ):
         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):
+@pytest.mark.parametrize("message", ["incomplete merge", ""], ids=["message", "no-message"])
+def test_revert_whole_sd_keeps_in_memory_names_on_every_rank_when_one_fails(
+    cpu_dist_workers, message
+):
     """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)
+    load neither way. An exception with no message is a failure all the same."""
+    cpu_dist_workers.run(_one_rank_cannot_revert_job, message=message)
 
 
 def _merge_group_job(rank, world_size, case):

From c666ea57fff3198fa608e9c3a3c413ffc4e51414 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 19:08:53 +0000
Subject: [PATCH 5/6] Refuse to write over a checkpoint already in the export
 directory

A rename failing partway through the final commit could leave a checkpoint
already in the directory mixed: some of its shards replaced, its index
naming tensors they no longer hold. The writer now refuses, on every rank
and before anything is written, a directory holding model.safetensors, its
index, or model-NNNNN-of-NNNNN.safetensors shards. The check runs in the
first collective, so a rank whose view of the directory alone shows them
fails the others too. Other files, auxiliary weights included, are left
alone and do not stop the export.

With no checkpoint ever written over, the removal of an earlier
checkpoint's stale weight files (285c30bf08e) goes, with its tests. The
staged .tmp writes stay.

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

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 1e161196be8..18967662d9c 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -44,8 +44,8 @@
 )
 from .unified_export_hf_streaming import _parse_shard_size
 
-# A checkpoint's weight files, sharded as save_pretrained names them or in one file.
-_CHECKPOINT_WEIGHTS = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors")
+# A checkpoint's files: its weights, sharded as save_pretrained names them or in one file, and index.
+_CHECKPOINT_FILES = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors(\.index\.json)?")
 
 
 def _even_bins(keys: list, sizes: dict, n: int, max_bytes: int) -> list:
@@ -503,11 +503,10 @@ def _write_even_shards(
     reverse conversion, whose renames and splits fix the names and sizes packed here.
 
     Every file is written under a temporary name and renamed into place only once every rank has
-    written all of its own, the index last. A failure while writing leaves the directory as it was, a
-    checkpoint already there included; a rename failing partway through that final commit can leave
-    such a checkpoint mixed. Once every rank's files are in place, rank 0 removes the earlier
-    checkpoint's weight files the new index does not name. Each step a rank can fail alone runs in
-    :func:`~modelopt.torch.utils.distributed.all_reduce`, which raises on every rank when one fails.
+    written all of its own, the index last. It never overwrites a checkpoint: a directory that already
+    holds one is refused on every rank before anything is written. Each step a rank can fail alone
+    runs in :func:`~modelopt.torch.utils.distributed.all_reduce`, which raises on every rank when one
+    fails.
     """
     export_dir = Path(export_dir)
     max_bytes = _parse_shard_size(max_shard_size)
@@ -519,9 +518,18 @@ def _write_even_shards(
     my_files = _even_bins(write_keys, sizes, max(1, math.ceil(fits / max_bytes)), max_bytes)
 
     def plan_mine() -> list:
-        # 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.
+        # In the first collective, which carries one rank's failure to every rank: its filesystem can
+        # fail to create the directory, or show a checkpoint there, which the export never overwrites.
         export_dir.mkdir(parents=True, exist_ok=True)
+        existing = sorted(
+            p.name for p in export_dir.iterdir() if _CHECKPOINT_FILES.fullmatch(p.name)
+        )
+        if existing:
+            raise FileExistsError(
+                f"{export_dir} already holds a checkpoint ({', '.join(existing[:3])}), which the "
+                "distributed export does not overwrite; remove those files or export to another "
+                "directory."
+            )
         return [(len(my_files), write_keys)]
 
     planned = all_reduce("shard plan", plan_mine, _in_rank_order)
@@ -591,19 +599,8 @@ def commit(names: list) -> bool:
             os.replace(export_dir / f"{name}.tmp", export_dir / name)
         return True
 
-    def remove_stale() -> bool:
-        # Left in place, an earlier checkpoint's shards would outlive it, and its model.safetensors
-        # would be loaded ahead of the new index.
-        if rank == 0:
-            current = set(full_map.values())
-            for path in export_dir.iterdir():
-                if path.name not in current and _CHECKPOINT_WEIGHTS.fullmatch(path.name):
-                    path.unlink()
-        return True
-
-    # A failed export removes the temporary files it wrote, and only those: a retry into the
-    # directory then finds it as it was. What was written is gone from ``local_sd``, so a retry needs
-    # the state dict built again.
+    # A failed export removes the temporary files it wrote, and only those. What was written is gone
+    # from ``local_sd``, so a retry needs the state dict built again.
     try:
         gathered = all_reduce("shards written", lambda: [write_mine()], _in_rank_order)
         full_map: dict = {}
@@ -612,8 +609,8 @@ def remove_stale() -> bool:
             full_map.update(m)
             total_bytes += b
         all_reduce("index written", write_index, all)
-        # Every file is written: only now does the directory change, the index last, so until
-        # then an index already there still names its own shards.
+        # Every file is written: only now does the directory change, the index last, so no index
+        # names a shard that is not in place. Also the closing barrier.
         all_reduce("shards committed", lambda: commit([n for n in staged if n != index_name]), all)
         all_reduce("index committed", lambda: commit([n for n in staged if n == index_name]), all)
     except BaseException:
@@ -632,5 +629,3 @@ def remove_stale() -> bool:
                 f"wrote to {export_dir}: {left[:5]}"
             )
         raise
-    # Only once every rank has committed, and never on a failure. Also the closing barrier.
-    all_reduce("stale weights removed", remove_stale, all)
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index 153a1e11f1e..8dd8ec6413c 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -67,8 +67,7 @@ def _assert_byte_identical(got, want):
 
 
 def _load_checkpoint(export_dir):
-    """Every tensor the index names, read from the shard the index names for it; no other weight file
-    of the checkpoint is left, a ``model.safetensors`` included, which transformers loads first."""
+    """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("model*.safetensors"))
@@ -552,17 +551,46 @@ def test_write_even_shards_writes_even_shards_and_an_index_of_every_tensor(
         assert sum(_nbytes(t) for t in load_file(str(tmp_path / shard)).values()) <= max_bytes
 
 
-@pytest.mark.parametrize("earlier", ["one-shard", "three-shards", "single-file"])
-def test_write_even_shards_replaces_a_checkpoint_already_there(
-    gloo_world_of_one, tmp_path, earlier
-):
-    """A re-export into the directory replaces the checkpoint it finds there: its shards and index,
-    any shards the new index does not name, and a single ``model.safetensors``. Other files stay."""
-    if earlier == "single-file":
-        save_file({"old": torch.zeros(2)}, str(tmp_path / "model.safetensors"))
+def _earlier_checkpoint(export_dir, layout):
+    """A checkpoint an earlier export left in ``export_dir``: one shard, three, or one file."""
+    export_dir.mkdir(parents=True, exist_ok=True)
+    if layout == "single-file":
+        save_file({"old": torch.zeros(2)}, str(export_dir / "model.safetensors"))
     else:
-        n = 1 if earlier == "one-shard" else 3
-        _write_even_shards({f"old{i}": torch.zeros(256) for i in range(n)}, tmp_path, "1100", 0, 1)
+        n = 1 if layout == "one-shard" else 3
+        names = [f"model-{i:05d}-of-{n:05d}.safetensors" for i in range(1, n + 1)]
+        for i, name in enumerate(names):
+            save_file({f"old{i}": torch.zeros(2)}, str(export_dir / name))
+        weight_map = {f"old{i}": name for i, name in enumerate(names)}
+        index = {"metadata": {"total_size": 8 * n}, "weight_map": weight_map}
+        (export_dir / "model.safetensors.index.json").write_text(json.dumps(index))
+    return {p.name: p.read_bytes() for p in export_dir.iterdir()}
+
+
+def _refuses_a_checkpoint_job(rank, world_size, tmp_dir, layout):
+    # Only rank 1's view of the directory holds the checkpoint: rank 0 must not write either.
+    export_dir = tmp_dir / f"rank{rank}"
+    before = _earlier_checkpoint(export_dir, layout) if rank == 1 else {}
+    with bounded_collectives(), pytest.raises((FileExistsError, RuntimeError)) as raised:
+        _write_even_shards(
+            {f"rank{rank}.weight": torch.ones(2)}, export_dir, "1MB", rank, world_size
+        )
+
+    assert "already holds a checkpoint" in str(raised.value), raised.value
+    assert {p.name: p.read_bytes() for p in export_dir.iterdir()} == before
+
+
+@pytest.mark.parametrize("layout", ["one-shard", "three-shards", "single-file"])
+def test_write_even_shards_refuses_a_checkpoint_already_there(cpu_dist_workers, tmp_path, layout):
+    """Every rank refuses before anything is written, though only one rank sees the checkpoint, and
+    the checkpoint is left byte for byte: a rename failing partway through the commit would leave
+    it mixed."""
+    cpu_dist_workers.run(_refuses_a_checkpoint_job, tmp_dir=tmp_path, layout=layout)
+
+
+def test_write_even_shards_writes_beside_other_files(gloo_world_of_one, tmp_path):
+    """A config and an auxiliary weights file the index does not cover are no checkpoint: the export
+    writes beside them and leaves them as they were."""
     others = {"config.json": b"{}", "mtp.safetensors": b"weights the index does not cover"}
     for name, data in others.items():
         (tmp_path / name).write_bytes(data)
@@ -571,9 +599,6 @@ def test_write_even_shards_replaces_a_checkpoint_already_there(
     _write_even_shards(dict(tensors), tmp_path, "1MB", 0, 1)
 
     _assert_byte_identical(_load_checkpoint(tmp_path), tensors)
-    assert sorted(p.name for p in tmp_path.iterdir()) == sorted(
-        ["model-00001-of-00001.safetensors", "model.safetensors.index.json", *others]
-    )
     assert all((tmp_path / name).read_bytes() == data for name, data in others.items())
 
 
@@ -674,34 +699,13 @@ def _one_rank_fails_job(rank, world_size, tmp_dir, fails):
         assert f"rank {failing_rank}" in str(raised) and "disk full" in str(raised), raised
 
 
-def _existing_checkpoint(export_dir):
-    """A checkpoint under the names the two-rank export below writes, holding other tensors, and a
-    ``model.safetensors`` a successful export would remove."""
-    weight_map = {
-        "old.a": "model-00001-of-00002.safetensors",
-        "old.b": "model-00002-of-00002.safetensors",
-    }
-    for key, shard in weight_map.items():
-        save_file({key: torch.zeros(3)}, str(export_dir / shard))
-    save_file({"old.c": torch.zeros(3)}, str(export_dir / "model.safetensors"))
-    index = {"metadata": {"total_size": 24}, "weight_map": weight_map}
-    (export_dir / "model.safetensors.index.json").write_text(json.dumps(index))
-    return {p.name: p.read_bytes() for p in export_dir.iterdir()}
-
-
-@pytest.mark.parametrize("existing", [False, True], ids=["empty", "over-a-checkpoint"])
 @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, existing
-):
+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; and
-    the directory is left as it was: no file of the failed export, and a checkpoint already there
-    byte for byte, though the export would have written its names."""
-    before = _existing_checkpoint(tmp_path) if existing else {}
-
+    the directory is left as it was, with no file of the failed export."""
     cpu_dist_workers.run(_one_rank_fails_job, tmp_dir=tmp_path, fails=fails)
 
-    assert {p.name: p.read_bytes() for p in tmp_path.iterdir()} == before
+    assert not list(tmp_path.iterdir())
 
 
 def _shard_commit_fails_job(rank, world_size, tmp_dir):
@@ -722,16 +726,12 @@ def failing_shard_rename(src, dst):
         _write_even_shards({f"rank{rank}.weight": torch.ones(2)}, tmp_dir, "1MB", rank, world_size)
 
 
-@pytest.mark.parametrize("existing", [False, True], ids=["empty", "over-a-checkpoint"])
-def test_the_index_goes_in_place_only_after_every_shard(cpu_dist_workers, tmp_path, existing):
-    """A rank that fails to put its shard in place leaves the index as it was, so an index in the
+def test_the_index_goes_in_place_only_after_every_shard(cpu_dist_workers, tmp_path):
+    """A rank that fails to put its shard in place leaves no index behind, so an index in the
     directory never names shards that are missing."""
-    index = tmp_path / "model.safetensors.index.json"
-    before = _existing_checkpoint(tmp_path)[index.name] if existing else None
-
     cpu_dist_workers.run(_shard_commit_fails_job, tmp_dir=tmp_path)
 
-    assert (index.read_bytes() if index.exists() else None) == before
+    assert not (tmp_path / "model.safetensors.index.json").exists()
 
 
 def _cross_rank_collision_job(rank, world_size, tmp_dir):

From f796006d1ffcebe07278b4d963393f55faa4bc07 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 22:47:48 +0000
Subject: [PATCH 6/6] Say what a failure during the final renames leaves, and
 that tied weights are the caller's to drop

A failure removes the writer's temporary files, but one during the final
renames leaves the shards already renamed, and the refusal of a directory
holding a checkpoint then refuses a retry over them. The docstring says so,
and the refusal's message names that cause. The docstring also says that
every name is written, two names for one tensor included, so tied weights
are dropped by name before the writer, as the streaming export does.

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

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 18967662d9c..0f7e63436ea 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -504,9 +504,13 @@ def _write_even_shards(
 
     Every file is written under a temporary name and renamed into place only once every rank has
     written all of its own, the index last. It never overwrites a checkpoint: a directory that already
-    holds one is refused on every rank before anything is written. Each step a rank can fail alone
-    runs in :func:`~modelopt.torch.utils.distributed.all_reduce`, which raises on every rank when one
-    fails.
+    holds one is refused on every rank before anything is written. A failure removes the temporary
+    files, but one during the final renames leaves the shards already renamed, which a retry into the
+    directory is refused over until they are removed. Each step a rank can fail alone runs in
+    :func:`~modelopt.torch.utils.distributed.all_reduce`, which raises on every rank when one fails.
+
+    Every name in ``local_sd`` is written, two names for one tensor included: the caller drops tied
+    weights by name beforehand, as the streaming export drops ``_tied_weights_keys``.
     """
     export_dir = Path(export_dir)
     max_bytes = _parse_shard_size(max_shard_size)
@@ -527,8 +531,8 @@ def plan_mine() -> list:
         if existing:
             raise FileExistsError(
                 f"{export_dir} already holds a checkpoint ({', '.join(existing[:3])}), which the "
-                "distributed export does not overwrite; remove those files or export to another "
-                "directory."
+                "distributed export does not overwrite; remove those files (an export that failed "
+                "during its final renames can leave some behind) or export to another directory."
             )
         return [(len(my_files), write_keys)]