diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index fdb3674e257..0f7e63436ea 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -20,10 +20,59 @@
 replicated across ranks.
 """
 
+import json
+import math
+import os
+import re
+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
+
+# 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:
+    """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 +408,228 @@ 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:
+            # 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 is not None
+        ]
+        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. It never overwrites a checkpoint: a directory that already
+    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)
+
+    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:
+        # 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 (an export that failed "
+                "during its final renames can leave some behind) or export to another directory."
+            )
+        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. 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 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:
+        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..8dd8ec6413c 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -16,10 +16,15 @@
 """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 os
+import types
 from unittest import mock
 
 import pytest
@@ -32,18 +37,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("model*.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 +402,353 @@ 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, 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(message)
+    )
+    with (
+        bounded_collectives(),
+        failure if rank == 1 else contextlib.nullcontext(),
+        pytest.warns(UserWarning, match=f"rank 1: {message or 'RuntimeError'}"),
+    ):
+        assert _revert_whole_sd(model, sd) is False
+
+    assert sd == before
+
+
+@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. 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):
+    # 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 _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 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)
+    tensors = {"w": torch.ones(2)}
+
+    _write_even_shards(dict(tensors), tmp_path, "1MB", 0, 1)
+
+    _assert_byte_identical(_load_checkpoint(tmp_path), tensors)
+    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(
+    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
+
+
+@pytest.mark.parametrize("fails", ["dir", "shard", "index"])
+def test_a_write_that_fails_on_one_rank_fails_on_every_rank(cpu_dist_workers, tmp_path, fails):
+    """Not just on the rank that failed, while its peers wait in a collective it never joins; and
+    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 not list(tmp_path.iterdir())
+
+
+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)
+
+
+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."""
+    cpu_dist_workers.run(_shard_commit_fails_job, tmp_dir=tmp_path)
+
+    assert not (tmp_path / "model.safetensors.index.json").exists()
+
+
+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"))