[OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank by shengliangxu · Pull Request #2751 · NVIDIA/Model-Optimizer

shengliangxu · 2026-10-12T00:57:17Z

Note

The expert-parallel PTQ series is now a tree instead of one stack (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [15/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.

[3/3] of this PR's split (2026-10-29), at its review's request: [1/3] #3100 (the reverse conversion's unmatched_merges hook), [2/3] #3101 (_materialize_owned, ownership and transport), and this PR (reverse conversion and shard writing). Merge in that order. [1/3] #3100 merged on 2026-10-29 (8ab63386152) and [2/3] #3101 the same day (4bef85279e4); this PR now targets main.

Base: main, on #3101's squash 4bef85279e4. The diff below is this PR's alone.

Re-cut: the gloo_world_of_one fixture its tests use is in #3101 now, in the same text #3099 adds, so either order merges cleanly. Review fixes since are listed at the end of Testing.

Series layout

main
- [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 fused experts under grouped_mm (merged)
  - [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 calibration on grouped_mm
    - [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
- [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 meta-DCP checkpoint load (merged)
- [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 distribution plan, run group, collectives (merged)
  - [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
  - [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold
  - [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100 conversion hook, [1/3] (merged)
    - [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101 ownership and transport, [2/3] (merged)
      - [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 shard writing, [3/3] (this PR)
    - [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
      - [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export
  - [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 plan decision
    - [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 load routing
      - [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 expert-parallel loader (on [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
        
        [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
        
        [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
        
        [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
        
        [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 expert parallelism in the model specs
        
        [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 on-the-fly decompression
        [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 Kimi-K2.5 and Kimi-VL adapters
        [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter
        
        [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)

What does this PR do?

Type of change: New feature

Main's FSDP2 export already writes each rank's share, but it first unshards each FSDP2 unit in turn, an all-gather of that unit to every rank. It also refuses conversion mappings that merge or split tensors, since one unit at a time cannot regroup them, and it does not handle expert-parallel shards. In this writer every tensor moves once, whole, to one owner rank and is written from there (modelopt/torch/export/distributed_writer.py). After the split this PR holds the last of its pieces: _materialize_owned is in #3101, and the conversion's unmatched_merges hook is in #3100.

_materialize_owned (now in [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101) gives each tensor exactly one owner rank and makes it whole there:
- replicated and plain values need no transport;
- dim-0 shards travel point-to-point, batched and bucketed per peer (one bucket per key deadlocked at 16 ranks);
- each group of keys the reverse conversion merges into one tensor (_merge_groups) goes to one owner, with its weights' quantization state, so the merge, or its refusal of a quantized group, happens there.
Placements it cannot assemble, and shards that do not tile the global shape, raise rather than risk a silently mis-assembled checkpoint; a zero-element tensor in such a layout is refused like a non-empty one. Experts arrive whole and plain, one disjoint set per rank: an expert that is still a DTensor is refused on every rank before anything moves.
_revert_whole_sd runs the reverse weight conversion on those whole tensors, so they go out under the checkpoint's names. transformers' conversion ops cannot express a placement. Ownership keeps each merge group whole on one rank, so no rank sees part of one. The reverse runs inside the failure-carrying all_reduce, so a refusal on any rank, such as a merge's refusal of a quantized or incomplete group, makes every rank keep its in-memory names together.
_write_even_shards:
- bin-packs each rank's set into evenly sized files numbered into one global range;
- refuses a name two ranks both hold, on every rank, in its first collective, before any file is written;
- writes the files with safetensors, both names of a tensor two keys share included (one is copied, as ShardedSafetensorsWriter does), which save_file would otherwise refuse;
- unions the index;
- never overwrites a checkpoint: a directory already holding model.safetensors, its index, or model-NNNNN-of-NNNNN.safetensors files is refused on every rank in the first collective, before anything is written. Other files, auxiliary weights such as mtp.safetensors included, are left alone;
- writes every file under a temporary name and renames them into place only once every rank has written all of its own, the index last, so no index names a shard that is not in place.
Failing together. The export directory's creation with the file counts and tensor names, the shard writes and rank 0's index write each run inside the failure-carrying all_reduce ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), so one rank's failure raises on every rank with that rank's error rather than leaving the others in an unmatched collective.

A tensor over max_shard_size gets a file of its own, as in HF, and the rest still pack.

The entry point that drives them arrives in #2752.

Why not a DCP writer, or main's streaming writer (_StreamingShardWriter, name_shards_and_write_index)?

DCP writes its own format, not safetensors shards under the hub's names with an index. vLLM and TRT-LLM would need a consolidation pass over it, and that pass is the gather this writer avoids.
The streaming writer serves main's FSDP2 export, which unshards each unit, an all-gather of it to every rank, and refuses conversion mappings that merge or split tensors. This writer moves each tensor once, whole, to one owner, so a merge can run there; that part has nothing to extend.
The file writing itself could share code with _StreamingShardWriter, which already stages to temporary files, given three things this writer needs:
- every step fails on every rank together, by running in the failure-carrying all_reduce, where the streaming path names its shards on rank 0 after a barrier and a gather_object;
- a name two ranks both hold is refused before anything is written;
- each rank's owned set is bin-packed into even files.
That sharing is deferred to a follow-up. It would change main's streaming export, a path this series does not otherwise touch, and no checkpoint format, so it can land on its own once this writer is in, without growing a PR already split for size.

The metadata collective carries each sharded tensor's dtype, every dense tensor's size and the merge units, which owners are assigned from; ranks that disagree on any of them are refused on every rank. The send buffers are freed once the transfers finish, before a rank assembles what it owns (a P2POp and its request each keep their tensor, so the transfer runs in a helper that drops both); on top of its own shards a rank holds at most about twice its owned set. On two gloo ranks with 128 MiB of shards each:

Send buffers	Peak RSS above the shards
kept while a rank assembles (previous revision)	257 MiB
freed before it assembles (this revision)	193 MiB

Usage

N/A: internal to the distributed export (#2752).

Testing

New tests/unit/torch/export/test_distributed_writer.py, on CPU under gloo, drives each piece directly:
- HF shard-size parsing, a unit in any case read as bytes, as the streaming export reads it (transformers reads a lowercase b as bits);
- binning: balanced, bins added until each fits, an oversized tensor isolated, no empty bins;
- DTensor sizes and the shard-0 placement check;
- ownership making every tensor whole and plain, and rejecting placements it cannot assemble;
- the reverse conversion renaming in place and keeping in-memory names when it fails;
- shard files and an index covering every tensor once;
- two ranks where one rank's directory or write fails, or two ranks hold one name: both raise, and the name collision leaves no shard file. Run on a 10 s gloo group, these fail in seconds on the code before the fix rather than waiting out the process group's timeout.
Mutation-checked: each test fails under at least one deliberate break of the code.
Each helper's tests are one table, and checks a remaining test covers are dropped; the two-rank jobs use the shared bounded_collectives. Each drop is checked by mutation.
On two ranks, _materialize_owned moves an uneven shard, an empty one and two dtypes to their owners, and a dtype disagreement, or an expert that only rank 1 still holds as a DTensor, is refused on every rank. An empty Shard(1) tensor is refused, and two names for one CPU tensor are both written. The expert, collision, empty-tensor and shared-tensor cases each fail without their fix (mutation-checked).
On two ranks, what a rank sent is freed before it stores an assembled tensor, a rank that received nothing included; a dense size that differs across ranks is refused on every rank; and _write_even_shards with uneven file counts numbers the files across ranks and indexes both ranks' tensors. Each fails without its fix (mutation-checked).
Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.
Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.
Review fixes, 2026-10-27: each merge group goes to one owner rank, so the reverse conversion merges it; the shard size is read with the streaming export's parser; the mesh memo's comment gives its real reason. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.
Review fixes, 2026-10-27, round 2: the merge groups are sorted before the units are built, so every rank derives the same owners whatever order its keys come in; a new case shows a quantized merge group, refused by the rank holding its scales, falling back on every rank. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.
Review fixes, 2026-10-27, round 3: a merged weight's quantization state stays in its merge group, so the rank that merges a quantized group refuses it even when the merge pattern does not match the scales; groups that share a key form one unit whatever their order; the ranks check in the metadata collective that they formed the same units. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.
Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8).
Review suggestions, 2026-10-27 evening: the refusal of ranks that group the keys differently lists each rank's first differing merge units.
Review fix, 2026-10-27 night: when the shard or index step fails, every rank removes the files it wrote, a partial one included, so a retry starts from the directory as it was. The one-rank-failure test checks the directory is left empty; its shard and index cases fail without the change.
Review suggestions, 2026-10-27, last round: a failed write's cleanup ignores an unlink error, so the error naming the failing rank is the one raised, and a comment says a retry rebuilds the state dict. The new test fails without the change.
Review suggestion, 2026-10-27, last round: files a failed write's cleanup cannot remove are named in one warning, and the write's own error is still the one raised. The extended test fails without the change.
The same gloo_world_of_one fix as [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's, in this PR's copy of the fixture, so the two stay identical.
The same test_meta_dcp_load.py fixture deletion as [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's, so the two stay identical.
Review suggestion, 2026-10-28, second round: each tensor is dropped right after its host copy, so a file's tensors are no longer held on the device while it is written; the new test fails without the change.
Review fix, 2026-10-28, seventh round: the comment by the drop now says it frees each tensor once its host copy exists. It does not lower the device peak, which is the allocation at entry with or without it, so a peak-memory measurement reads the same either way by construction and is no evidence against it; the memory still held when each file is written is what moves. Measured on 4 GPUs, in a distributed HF export of tiny FP8 and NVFP4 models that runs this writer on CUDA tensors (torch 2.15 pre-release, transformers 5.18; this PR's own writer test runs on CPU, where device memory does not move). The table gives the device bytes allocated (torch.cuda.memory_allocated) on each rank just before save_file. Each rank wrote one file, so a file was its whole owned set; a rank writing several files was not measured.

Model	Tensors per rank	Held before this change	Held with it	Released before the write
FP8, dense, tied embeddings	5 to 20	167,424 to 167,936	109,056	58,368 to 58,880 (35%)
FP8, MoE	39 to 53	526,848 to 528,384	275,456	251,392 to 252,928 (48%)
NVFP4, dense, tied embeddings	6 to 25	131,072 to 131,584	97,280	33,792 to 34,304 (26%)
NVFP4, MoE	49 to 71	371,712 to 376,320	220,672	151,040 to 155,648 (41%)

Rebased on 2026-10-28, late night, onto [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 at cf4597dcbeb, on main 1ca8bd1fc4d (after [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged), without conflicts; this PR's own diff is unchanged. Its tests pass (143), and pre-commit passes against main.
Rebased on 2026-10-28, late night, onto main bdc45368d10 after [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 merged; [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736's squash has the tree of its last head, so this PR's tree is unchanged.
Review (meenchen's review bot), 2026-10-29, all items:
- Split in three: [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101 and this PR. This branch merges [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101, so its own commits stay; source lines are about 40, 347 and 254.
- Rollback: a failed export no longer deletes a checkpoint already in the directory. Every file is written under a temporary name, and the ranks rename theirs into place only once every rank has written all of its own, the index last. A failure before then removes only the temporary files. The one-rank-failure test runs over an existing checkpoint too and checks it byte for byte; its shard and index cases fail without the change. A new test re-exports over a checkpoint.
- Imports: the writer's, the fixture's and a test's imports moved to module scope.
- Docstrings: the long ones keep their constraints; the narrative moved here.
- GPU tests: tests/gpu/torch/export/test_distributed_writer_gpu.py runs from two GPUs. It moves BF16, FP8 and uint8 shards of uneven and empty extent over NCCL and writes several files per rank. The GB200 runs and measurements follow below.
- Results: 148 tests pass, and pre-commit passes against main.
- Narrative moved from the docstrings: bucketing the transfers by peer replaced one operation per (key, peer), which at 235B was about 12.7k operations and 660 s for a few GB of payload.
GPU runs on GB200 of tests/gpu/torch/export/test_distributed_writer_gpu.py (its two tests), ran at f577fc05af0. The code is unchanged at this head; only [2/3]'s docstrings and tests changed since.

GPUs Result Time

1 2 skipped (need two GPUs) 0.5 s

2 2 passed 17.6 s

4 2 passed 23.5 s

Measured on GB200 at f577fc05af0 with a synthetic sharded state dict: 600 keys of shape (8741, 4096), 360 BF16, 150 FP8 e4m3 and 90 uint8, 32.0 GiB in total. 8741 rows divide by neither 4 nor 8, so the dim-0 shards are uneven. Timed _materialize_owned(sd, set(), rank, world), then _write_even_shards(sd, out_dir, "5GB", rank, world). Rank 0 is shown; all ranks were near-identical.

GPUs	Input per rank	Owned per rank	Transport	Peak device memory in the transport, above its input	Write	Files per rank	Bytes written per rank	Peak host RSS (process lifetime)
4 (one node)	8.01 GiB	8.00 GiB	6.22 s	12.01 GiB	4.36 s	2	8.00 GiB	8.68 GiB
8 (two nodes × 4)	4.00 GiB	4.00 GiB	14.15 s	7.01 GiB	2.57 s	1	4.00 GiB	8.49 GiB

The device peak is torch.cuda.max_memory_allocated() after reset_peak_memory_stats, minus the allocation before the transport. During the transfer a rank holds one send buffer and one receive buffer per (peer, dtype). The send buffers copy its shards of the tensors other ranks own, (w-1)/w of its input; the receive buffers hold the other ranks' shards of the tensors it owns, (w-1)/w of its owned set. Their sum is the measured peak at both sizes. In this run the state dict held the only reference to each input shard, so a rank freed its inputs as it assembled.
Host RSS is ru_maxrss, the maximum over the process's life. It was 2.1 GiB before the write, so the write's own peak is at most the number shown.
The 8 GPUs span two nodes, so that transport time includes cross-node NCCL traffic and does not scale like-for-like against the one-node run.

Re-measured on GB200 at 239260c7d8d with the input DTensors kept alive through the write, as the model keeps its parameters in a real export. The writer's code is the same as at f577fc05af0 (docstrings differ) and is unchanged at 2470bc4a4f7. Same 32.0 GiB state dict; rank 0 is shown, and all ranks were near-identical.

GPUs	Owned per rank	Transport	Peak device memory above the allocation before the transport (transport and write)	Write	Files per rank	Peak host RSS during the write
4 (one node)	8.00 GiB	6.83 s	14.01 GiB	4.12 s	2	8.65 GiB
8 (two nodes × 4)	4.00 GiB	14.76 s	7.51 GiB	2.52 s	1	8.21 GiB

With the inputs alive, the peak comes from the assembly: the receive buffers, (w-1)/w of the owned set, plus the assembled owned tensors. Their sum is the measured peak at both sizes, within _materialize_owned's bound of about twice the owned set.
The write never raised the peak above the transport's: it copies each tensor to the host and drops it, so it adds time, not device memory.
Host RSS here is windowed to the write (the high-water mark is reset just before it); the first table's is the process-lifetime ru_maxrss.

2026-10-29: merges [2/3]'s review round (46e42e10a74, rebased onto [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100's b9fbb242300). That brings [1/3]'s docstring and test, plus [2/3]'s merge-group and strided-shard tests and trimmed docstrings. No change of its own. Pre-commit passes against main; test_distributed_writer.py, test_quant_aware_conversion.py and test_meta_dcp_load.py pass (149).
2026-10-29, later (2470bc4a4f7, test only): test_the_index_goes_in_place_only_after_every_shard pins the commit order. Rank 1 fails to rename its shard into place after every file is written, and the index must stay as it was: absent in an empty directory, byte for byte over an existing checkpoint. So an index in the directory never names missing shards. With the shard and index commits swapped, both cases fail, and every other writer test stays green. The writer's code is unchanged since 239260c7d8d. test_distributed_writer.py passes (44), and pre-commit passes against main.
2026-10-29 (meenchen's review bot): why the writer is neither a DCP writer nor main's streaming writer, and why sharing its file writing with the latter is deferred, is now in the description above, moved from the review reply. Description only.
Review (meenchen's review bot), 2026-10-29, second round (285c30bf08e): once every rank has committed, rank 0 removes an earlier checkpoint's model-NNNNN-of-NNNNN.safetensors files the new index does not name, and a model.safetensors, which transformers loads ahead of the index. The replacement test runs over an earlier one-shard, three-shard and single-file checkpoint, with a config.json and an off-index weight file that must stay; the three-shard and single-file cases fail without the removal. The failure tests' existing checkpoint gained a model.safetensors, so a removal on a failure path fails them (checked by running it in the failure handler). The docstring states the commit's limit: a rename failing partway through the final commit can leave a checkpoint already there mixed; it is documented rather than fixed, for the reasons in the review thread. test_distributed_writer.py passes (46), and pre-commit passes against main. Source lines: 272.
2026-10-29, later: merges [2/3]'s second review round (055b57bd7cf): mixed-device shards and whole-tensor disagreements refused on every rank, each transfer round's buffers freed when it ends. No change of its own (1feacd040c9). The writer, conversion and meta-DCP tests pass (158), and pre-commit passes against main.
Review fix, 2026-10-29 (03ae6f5d5ff, meenchen's review bot): a reverse conversion that fails with an exception carrying no message now counts as a failure. Before, RuntimeError() read as success, so rank 0 committed hub names while rank 1 kept in-memory ones (reproduced on two gloo ranks with tiny Mixtral). The one-rank-failure test gains that case, which fails without the fix. test_distributed_writer.py passes (52), and pre-commit passes against main.
GPU run, 2026-10-29, ran at 1feacd040c9 (the same code as the current head). GB200, torch 2.11.0a0 (NGC 26.02), transformers 5.9.0:

tests/gpu/torch/export/test_distributed_writer_gpu.py (2 tests) 1 GPU 2 GPUs 4 GPUs

Result 2 skipped 2 passed, 20.3 s 2 passed, 22.5 s
Review (meenchen's review bot), 2026-10-29, third round (d9bf5aafbbb, the user's call): the writer refuses, on every rank and before writing anything, a directory that already holds a checkpoint (model.safetensors, its index, or model-NNNNN-of-NNNNN.safetensors), so a commit failure can no longer mix one. The refusal is decided in the first collective, so one rank's view of the directory stops them all. This replaces the documented commit limit and the stale-file cleanup of 285c30bf08e, which is removed with its tests; the staged .tmp writes stay. A two-rank test checks the refusal for a one-shard, three-shard and single-file earlier checkpoint, each left byte for byte, and fails without the refusal. Another exports beside a config.json and an auxiliary weights file, and fails if the pattern also matches the auxiliary file. test_distributed_writer.py passes (49), and pre-commit passes against main. Source lines: 270.
2026-10-29, later: merges [2/3] rebased onto main 2b900c8fa41 with its third review round (b73c052e45f): every rank initializes NCCL before a transfer only some ranks take part in, with a fresh-group GPU test, and the transport's memory bound names what a rank sends. No change of its own (8b6f3a275f7). The writer and conversion tests pass (128), and pre-commit passes against main.
GPU run, 2026-10-29, at 8b6f3a275f7. GB200, torch 2.11.0a0 (NGC 26.02), CUDA 13.1, NCCL 2.29.2, transformers 5.9.0:

tests/gpu/torch/export/test_distributed_writer_gpu.py 2 GPUs 4 GPUs

Result 2 passed, 1 skipped (the fresh-group test needs 3 or more), 20.4 s 3 passed, 46.1 s
2026-10-29, later: merges [2/3]'s docstring note that owners are weighed by bytes (764036cd64e). No change of its own (2a65d9fa364); pre-commit passes against main.
Rebased on 2026-10-29 onto [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101's squash 4bef85279e4 after it merged, with this PR's own diff unchanged. The history is now linear: one commit for the pieces as they stood when the split was made, then this PR's four commits since. [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100's and [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101's pre-squash commits are dropped. Pre-commit passes against main; the writer and conversion tests pass (128).
Claude review, 2026-10-29 (approved, two suggestions; 55b62f6fdf4, docstring and message only): _write_even_shards' docstring says that a failure during the final renames leaves the shards already renamed, which a retry is refused over until they are removed, and the refusal's message names that cause. It also says the caller drops tied weights by name. Pre-commit passes against main; test_distributed_writer.py passes (49).
Rebased on 2026-10-30 onto main 99c3834f62c (after [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099 and [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 merged, and [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737), own diff unchanged and no conflicts: main's new commits touch none of this PR's files. Pre-commit passes against main.

Before your PR is "Ready for review"

Is this change backward compatible?: ✅ Additive; nothing calls these yet.
If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
Did you write any new necessary tests?: ✅
Did you update Changelog?: N/A. The no-gather export's entry lands with [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752.
Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 15 of 21 of expert-parallel PTQ for Hugging Face MoE models; its [1/3] #3100 and [2/3] #3101 have merged.

Summary by CodeRabbit

New Features
- Added distributed checkpoint export in Hugging Face safetensors format, writing tensors across ranks without gathering the full model on every rank.
- Exported files are balanced by size, with oversized tensors placed in their own files and an index generated for the shards.
- Export supports eligible sharded tensors and reverses supported quantization-related weight conversions.
Bug Fixes
- Exports detect incompatible shards, tensor-name collisions, and write failures across ranks to help prevent incomplete or inconsistent checkpoints.

copy-pr-bot · 2026-10-12T00:57:21Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

coderabbitai · 2026-10-12T00:57:26Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Last reviewed commit: a8c56cf
Latest commit, not reviewed: f796006

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7ac65992-6997-4758-bcfc-5195bdc2d18d

📥 Commits

Reviewing files that changed from the base of the PR and between eca01ac and a8c56cf.

📒 Files selected for processing: 2

modelopt/torch/export/distributed_writer.py
tests/unit/torch/export/test_distributed_writer.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 9 remain after this review.

Walkthrough

Adds a distributed checkpoint writer that assigns tensors to ranks, transfers supported shards, and writes safetensors shards with a global index. It coordinates reverse weight conversion and propagates rank-local failures.

Changes: Distributed checkpoint export

Layer / File(s)	Summary
Tensor ownership and materialization `modelopt/torch/export/distributed_writer.py`, `tests/unit/torch/export/test_distributed_writer.py`	The writer parses shard-size limits, measures tensors, assigns dense tensors to ranks, validates supported dimension-0 shards, and transfers shards to their owners. Tests cover size parsing, tensor sizing, assignment, materialization, and unsupported layouts.
Coordinated reverse conversion `modelopt/torch/export/distributed_writer.py`, `tests/unit/torch/export/test_distributed_writer.py`	Ranks reverse-convert owned tensors. If any rank fails, all ranks retain the original local state dictionary. Tests cover key renaming and coordinated conversion failure.
Shard and index writing `modelopt/torch/export/distributed_writer.py`, `tests/unit/torch/export/test_distributed_writer.py`	Ranks bin and write owned tensors as safetensors shards, then gather shard metadata. Rank 0 writes the index. Tests cover checkpoint contents, empty output, distributed write failures, and name collisions.

Priority: Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Ranks as Export ranks
  participant Writer as Distributed checkpoint writer
  participant Shards as Safetensors shards
  participant Rank0 as Rank 0
  participant Index as Safetensors index
  Ranks->>Writer: Provide owned tensors
  Writer->>Shards: Write rank-local shard files
  Writer->>Rank0: Gather shard counts and tensor mappings
  Rank0->>Index: Write global index
  Index-->>Ranks: Complete index-writing collective

Merge Risk: Low · up to a8c56

This PR adds internal pieces of a distributed checkpoint writer that are not yet called by any export path. Failures on one rank now reach every rank, which prevents hangs. One limitation remains: for models whose weight conversion merges tensors, the export can fall back, with a warning, to writing in-memory tensor names. That checkpoint stays consistent but may not load as a standard Hugging Face checkpoint. The PR is mergeable if owners are aware of this.

Pre-merge checks: 5

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The PR changes only `modelopt/torch/export/distributed_writer.py` and its unit test. The production diff adds no `torch.load(..., weights_only=False)`, `numpy.load(..., allow_pickle=True)`, hard…
Title check		The title clearly identifies the distributed Hugging Face checkpoint shard-writing change and matches the main purpose of the pull request.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The PR changes only modelopt/torch/export/distributed_writer.py and its unit test. The production diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval(), exec(), or # nosec. It also adds no dependency files or PIP dependencies.

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

Commit to this branch
Create a new PR

🧪 Generate unit tests (beta)

Commit to this branch
Create a new PR

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

github-actions · 2026-10-12T02:50:47Z

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-30 17:51 UTC

codecov · 2026-10-12T03:13:43Z

Codecov Report

❌ Patch coverage is 98.62069% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.95%. Comparing base (99c3834) to head (f796006).

Files with missing lines	Patch %	Lines
modelopt/torch/export/distributed_writer.py	98.62%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2751      +/-   ##
==========================================
+ Coverage   71.55%   79.95%   +8.39%     
==========================================
  Files         664      664              
  Lines       76659    76804     +145     
==========================================
+ Hits        54856    61406    +6550     
+ Misses      21803    15398    -6405

Flag	Coverage Δ
examples-diffusers	`19.81% <0.00%> (-0.11%)`	⬇️
examples-gpt-oss	`13.41% <0.00%> (-0.06%)`	⬇️
examples-hf_ptq	`23.38% <0.00%> (-0.10%)`	⬇️
examples-llm_distill	`13.49% <0.00%> (-0.06%)`	⬇️
examples-llm_eval	`17.20% <0.00%> (-0.03%)`	⬇️
examples-llm_qat	`17.32% <0.00%> (-0.10%)`	⬇️
examples-llm_sparsity	`16.71% <0.00%> (-0.03%)`	⬇️
examples-megatron_bridge	`26.32% <0.00%> (-0.43%)`	⬇️
examples-specdec_bench	`13.20% <0.00%> (-0.05%)`	⬇️
examples-speculative_decoding	`17.44% <0.00%> (-0.10%)`	⬇️
examples-torch_onnx	`21.08% <0.00%> (-0.15%)`	⬇️
examples-torch_trt	`15.14% <0.00%> (-0.07%)`	⬇️
examples-vllm_serve	`17.42% <0.00%> (-0.07%)`	⬇️
gpu	`58.81% <12.41%> (+36.86%)`	⬆️
regression	`14.96% <0.00%> (-0.06%)`	⬇️
unit	`62.40% <98.62%> (+0.07%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

shengliangxu · 2026-10-12T03:14:18Z

/ok to test 6994370

shengliangxu · 2026-10-12T15:48:11Z

/ok to test 5b2ee4a

shengliangxu · 2026-10-13T03:41:55Z

/ok to test 812e4e1

shengliangxu · 2026-10-29T06:07:16Z

Thanks. Here is what each item became; it is pushed as 285c30b.

Preserving the old checkpoint through commit failures: no code change. _write_even_shards's docstring now states the exact guarantee: a failure while writing leaves a checkpoint already there untouched, but a rename failing partway through the final commit can leave it mixed. The final step is a series of same-directory renames, run only after every rank has written. Main's save_pretrained deletes the old shards before it writes, so it gives no such guarantee at all. Full safety would need per-export shard names, giving up HF's standard naming. The details are in the thread.
Retiring the old checkpoint's files: done.
- Once every rank has committed, rank 0 removes the model-NNNNN-of-NNNNN.safetensors files the new index does not name, plus a model.safetensors. Nothing else is touched, and nothing is removed on a failure.
- The replacement test covers a different shard count and an existing model.safetensors. It checks that a load sees only the new weights and that the other files are unchanged; both new cases fail without the removal.
DCP and streaming-writer rationale: it is in the description now, including why sharing the file writing with _StreamingShardWriter is deferred. That sharing would change main's streaming export, a path this series does not otherwise touch, and no checkpoint format. So it can land on its own once this writer is in, without growing a PR already split for size.

meenchen

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Changes requested: commit-failure corruption remains unresolved, and message-less reverse exceptions can publish mixed tensor names.

Needs action:

💬 Author documented the mixed-checkpoint limit — still preserve prior files or collectively refuse overwriting in _write_even_shards; extend shard/index-rename failure coverage to verify every prior file byte-for-byte.
Fix _revert_whole_sd.combine to check error is not None; extend the two-rank failure test with RuntimeError(). See inline.

No action needed:

✔️ Resolved since the last review: obsolete-weight cleanup and the design rationale. The split, staged write-time preservation, module-level imports, and GPU coverage remain intact.

meenchen · 2026-10-29T06:13:27Z

+            return str(exc), []
+        return None, unmatched
+
+    def combine(results: list) -> list:


Bot comment.

[P2] Detect reverse failures independently of their message. RuntimeError() has an empty string, so revert() returns ("", []), but if error discards that failure. With rank 1 raising this exception and rank 0 successfully reversing its Mixtral subset, combine returns no failures, both ranks report True, and rank 0 commits hub names while rank 1 keeps in-memory names. This bypasses the all-rank fallback without a warning. Check error is not None instead, and extend _one_rank_cannot_revert_job with a message-less exception that verifies both ranks retain their original names.

Fixed in 03ae6f5; you were right. combine now counts any rank's error as a failure (error is not None), however empty its message. The warning names the exception's type when there is no message.

_one_rank_cannot_revert_job now also runs with a message-less RuntimeError() on rank 1. Both ranks must return False and keep their original names. That case fails without the fix: rank 0 committed its hub names.

shengliangxu · 2026-10-29T06:33:37Z

/ok to test 1feacd0

shengliangxu · 2026-10-29T06:47:53Z

/ok to test 03ae6f5

shengliangxu · 2026-10-29T07:08:41Z

/ok to test c28c038

meenchen

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Changes requested: message-less reverse failures are fixed, but commit failures can still corrupt a previously valid checkpoint.

Needs action:

💬 Author documented the commit-failure limit — still preserve prior files or collectively refuse overwriting in _write_even_shards; extend shard/index-rename failure coverage to verify the entire prior checkpoint byte-for-byte. See the existing commit-failure inline.

No action needed:

✔️ 7 earlier concerns addressed, including message-less reverse failures and write-time preservation; obsolete-weight cleanup, design rationale, split, module-level imports and GPU coverage remain intact.

shengliangxu · 2026-10-29T07:48:07Z

/ok to test d9bf5aa

shengliangxu · 2026-10-29T07:48:28Z

Thanks. The remaining item is done in d9bf5aa, by collectively refusing to overwrite.

_write_even_shards now refuses, on every rank and before anything is written, a directory that already holds a checkpoint: model.safetensors, its index, or model-NNNNN-of-NNNNN.safetensors. It decides this in its first collective, so one rank's view of the directory is enough to stop all of them.
Other files, auxiliary weights included, are left alone.
With nothing to overwrite, a commit failure can no longer mix a prior checkpoint.
A two-rank test checks the refusal for a one-shard, three-shard and single-file earlier checkpoint, and that each survives byte for byte; all three fail without the refusal.
A second test checks that a directory with only a config.json and an auxiliary weights file still exports.
The earlier stale-file cleanup (thread on obsolete weight files) is superseded and removed, since overwriting is now refused.

shengliangxu · 2026-10-29T16:28:09Z

/ok to test 8b6f3a2

shengliangxu · 2026-10-29T17:28:58Z

/ok to test 2a65d9f

shengliangxu · 2026-10-29T19:09:48Z

/ok to test bb9dd1f

shengliangxu · 2026-10-29T22:33:11Z

/claude review

claude · 2026-10-29T22:40:58Z

+        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:


[SUGGESTION] The cleanup unlinks only <name>.tmp, so a failure during the final commit leaves the shards that were already renamed into place — and the new refusal then permanently blocks a retry into that directory.

Walk the path test_the_index_goes_in_place_only_after_every_shard exercises: every rank has written all its .tmp files, all_reduce("shards committed", ...) runs, rank 0 renames its shards successfully, rank 1's os.replace raises. Every rank raises, and each one's cleanup finds no .tmp left for a name it already renamed ("A file already renamed into place has no temporary file left"). The directory is left holding model-00001-of-00003.safetensors with no index. That test only asserts the index is absent, so the leftover shards are not pinned either way.

Why it matters: those leftovers match _CHECKPOINT_FILES, so the next run of this writer refuses with "<dir> already holds a checkpoint (model-00001-of-00003.safetensors), which the distributed export does not overwrite" — which reads as "someone else's checkpoint is here" when in fact it is debris from the export that just failed. The user has to work out which files to delete by hand before any retry, and the docstring's "never overwrites a checkpoint" now also means "a failed commit needs manual cleanup".

Unlike the stale-file removal dropped in 285c30bf08e, removing these is unambiguously safe: the refusal in plan_mine guarantees the directory held no model*.safetensors before this export started, so every name in this rank's staged that is now in place was written by this export and by this rank. Suggested change — unlink both the temporary and the committed name, per rank:

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. The commit renames a staged file, so remove both names: the # directory held no checkpoint file when the export started (plan_mine refuses one), so # a ``model*.safetensors`` under one of this rank's names is this export's own. for leftover in (export_dir / f"{name}.tmp", export_dir / name): try: leftover.unlink(missing_ok=True) except OSError: left.append(leftover.name) if left: warnings.warn( f"rank {rank}: could not remove {len(left)} file(s) the failed export " f"wrote to {export_dir}: {left[:5]}" ) raise

That makes "a failure leaves the directory as it was" hold on the commit path too, and test_the_index_goes_in_place_only_after_every_shard can then assert the stronger not list(tmp_path.iterdir()) the other failure test already asserts (it fails today on the shard-commit case). If you'd rather keep the commit non-reverting, the alternative is to say so in _write_even_shards' docstring and to have the refusal name the possibility, e.g. "... does not overwrite; remove those files (a previous export that failed during its final rename can leave some behind) or export to another directory."

I took the alternative you offered, in 55b62f6, and kept the commit non-reverting.

Docstring: _write_even_shards's docstring now says that a failure removes the temporary files, but one during the final renames leaves the shards already renamed, and a retry into the directory is refused over them until they are removed.

Error message: the refusal now names that cause: "remove those files (an export that failed during its final renames can leave some behind) or export to another directory".

Your reasoning is right: plan_mine's refusal means any model*.safetensors under one of this rank's staged names was written by this export, so deleting them on failure would be safe. I didn't make that change here, because the failure-path cleanup was decided explicitly for this PR; I've raised your version as an option instead.

claude · 2026-10-29T22:41:15Z

+                # 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())


[SUGGESTION] "every name has to reach the index" is the right call for this helper, but it makes the distributed export write a tied lm_head.weight as a full byte-for-byte duplicate, where both save_pretrained and main's streaming export drop it by name. Worth a pointer here so #2752 does not inherit it silently.

unified_export_hf_streaming.py:325 / :373 spell out the reason main's path uses names rather than pointers: "Tied weights are dropped by name from _tied_weights_keys (data_ptr is meaningless …)". This writer has no equivalent filter, and it cannot have one — it is handed a state dict, not a model — so for a tied-embedding model it emits model.embed_tokens.weight and lm_head.weight as two independent tensors. The checkpoint still loads (transformers re-ties after load), but it costs an extra vocab × hidden tensor on disk and in the index's total_size: ~0.5 GB in BF16 for a 1B-class tied model, and the PR description's own 4-GPU measurement table has two "tied embeddings" rows, so this is a layout the writer already meets.

Suggestion: no code change here, just make the constraint explicit so the entry point owns it — e.g. extend the comment to "... Copy rather than drop one: every name has to reach the index. Tied weights must therefore be dropped by name upstream, as unified_export_hf_streaming does from _tied_weights_keys; this writer cannot tell a tie from an alias." and add the same sentence to _write_even_shards' docstring, which is where #2752 will look.

Done in 55b62f6. _write_even_shards's docstring now states the contract: every name in local_sd is written, two names for one tensor included, so the caller drops tied weights by name beforehand, as the streaming export drops _tied_weights_keys.

The no-gather entry point in #2752 already does this. It builds a TiedWeightMap from the model and passes it to postprocess_state_dict, for both the expert and the dense state dicts, so a tied lm_head.weight is dropped by name before the writer sees it.

claude

Claude review — [3/3] shard writing from each tensors owner rank

Scope: full review (no scoping instructions in the trigger). 3 files changed (+741/-3), all read: distributed_writer.py (+270), tests/unit/torch/export/test_distributed_writer.py (+443/-2), tests/gpu/torch/export/test_distributed_writer_gpu.py (+28/-1). Read from the base for context: _materialize_owned / _is_full_world_shard0 / _nbytes and the unit_of construction in the same module; _apply_merge_rule / _apply_split_rule / _split_leaf_tensor / _build_reverse_rules / _QUANT_STATE_SUFFIXES / rename_source_key; unified_export_hf_streaming._parse_shard_size and its shared-tensor and tied-weight handling; utils/distributed.all_reduce. Grepped modelopt/ and examples/ for importers — none, as the description says (entry point is #2752). The checkout is shallow, so I re-derived the write-path arguments against the base rather than reviewing only the delta, with extra attention on what is new since the 2026-10-28 approval: the checkpoint refusal that replaced 285c30bf08es stale-file cleanup, the two-phase commit, and the empty-message fix in _revert_whole_sd.

Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2

Neither blocks.

The cleanup unlinks only <name>.tmp, so a rename failure partway through the final commit leaves shards already renamed into place, and _CHECKPOINT_FILES then blocks a retry into that directory with a message reading as "someone elses checkpoint is here". Since plan_mine guarantees the directory held no checkpoint file at entry, also removing this ranks committed names is safe and restores "a failure leaves the directory as it was".
Tied weights are written as full duplicates, where save_pretrained and mains streaming export drop them by _tied_weights_keys name. This helper cannot tell a tie from an alias, so the constraint is worth stating where the alias copy is already explained, for #2752.

Verified this round

The refusal is rank-identical and cannot be pre-empted. plan_mines mkdir/iterdir/FileExistsError runs inside the first all_reduce, which runs fn on every rank and gathers before any proceeds, so a directory only one rank can see still stops them all — what _refuses_a_checkpoint_job pins. The collision Counter, n_total == 0 and base = sum(counts[:rank]) all read the rank-identical planned, so those raises are identical too.
_CHECKPOINT_FILES matches what it should and nothing else. Under fullmatch: model.safetensors, its index, and model-NNNNN-of-NNNNN.safetensors. Not *.safetensors.tmp (so a previous runs staged files do not block a retry and are simply replaced by os.replace), not config.json / mtp.safetensors / adapter_model.safetensors. 4-digit numbering and other prefixes are out, matching transformers own %05d.
Commit ordering holds. all_reduce gathers after running fn, so the index can only land once every rank has committed every shard — no index ever names a missing shard, and both cases of test_the_index_goes_in_place_only_after_every_shard fail with the two commits swapped. The index commit is a pure barrier on non-zero ranks, which is what lets the GPU test read the checkpoint back without its own barrier.
_even_bins terminates and emits no empty or over-bound file. n is bumped at most to len(fits), where LPT places one item per bin; max(loads) never sees an empty sequence. Empty key set returns [] (a rank owning nothing writes no file and still joins every collective, leaving base/n_total contiguous); all-oversized floors n to 1 with a bin each; min(range(n), ...) takes the first minimum, so the plan is deterministic. sizes/fits/my_files are computed before write_mine drains the dict.
Every key reaches one file and the index. my_files partitions sorted(local_sd) by the > max_bytes predicate; base/n_total give disjoint, contiguous, correctly totalled numbers across ranks, pinned with uneven per-rank counts (1 + 2).
The alias guard survives the per-key drop. del local_sd[k1] drops only the dict entry while t1 holds the storage, so k2s copy still reports the same data_ptr() and is cloned — and the deleted key is the first, so the test exercises the order that matters. Address reuse cannot false-positive, since every pointer in ptrs belongs to a tensor still in tensors. Resetting ptrs per file is right (save_file only refuses sharing within one file). I checked the one case where data_ptr() and safetensors storage-identity predicate diverge — two contiguous views into one storage at different offsets — and the reverse conversion cannot produce it: _split_leaf_tensor ends in .clone() on both branches, and shard0/empty0 come out of torch.cat/torch.empty. Empty tensors are skipped by safetensors storage_size != 0 guard. The guard also matches _StreamingShardWriters existing data_ptr() approach, so it is not a new divergence.
_revert_whole_sd is atomic across ranks and the no-message fix is right. revert() catches Exception itself, so the decision is made in combine, where failures or ... short-circuits before an errored ranks empty unmatched could empty the intersection. str(exc) or type(exc).__name__ closes the case where a bare RuntimeError() read as success. apply_reverse_rules copies before mutating and reverted keeps its initial binding on failure, so the reverted is not local_sd guard means no checkpoint can mix hub and in-memory names.
_merge_groups matches what the merge demands. (rule_index, target) is the same granularity _apply_merge_rule groups by. Pulling each .weights _QUANT_STATE_SUFFIXES siblings in is what makes the guard at quant_aware_conversion.py:210-216 fire on the owning rank for a weight-anchored pattern that never matches a scale — the anchored case. Duplicate keys in a group are harmless (unit_of unions sets); expert keys are filtered by k in unit_of; regrouped catches a rank that formed different units before a byte moves. The unit_of write-back is order-independent (walked [a,b],[c,d],[b,c] and its reverse), so the dict-insertion order of _merge_groups cannot make ranks disagree.
Failure propagation. Every write step runs inside all_reduce, which re-raises on the failing rank and raises a rank-naming RuntimeError on its peers; staged.append precedes each write so a partial file is recorded before it exists; the except OSError around the unlink keeps the rank-naming error as the one raised. The three pre-commit failure cases assert not list(tmp_path.iterdir()); the commit case is the gap suggestion 1 covers.
Serialization and API surface. Every payload round-trips under weights_only; _in_rank_orders unwrap is correct in both the distributed and the single-process combine([fn()]) branch. No mode registration, config schema, modelopt_state key or public __init__.py export is touched, and nothing imports this module yet — though the new module-scope safetensors / quant_aware_conversion / unified_export_hf_streaming imports mean #2752 should keep its own import of it lazy if export/__init__.py is to stay importable without transformers 5.

Tests are proportionate: CPU-only unit file, every multi-process case on the shared cpu_dist_workers as CLAUDE.md requires, _even_bins as one parametrized table, and the GPU addition reuses the existing helpers and asserts byte identity plus no leftover .tmp.

Already raised, not re-filed. The raw all_gather_object in _materialize_owned outside the failure-carrying all_reduce; the receive-buffer device taken from one key per bucket; _is_full_world_shard0 still refusing the HSDP (Replicate(), Shard(0)) layout, which belongs with #2742/#2753. A max_shard_size of "0" divides by zero in math.ceil(fits / max_bytes) — pathological, not worth a guard.

Risk: low, and additive. Nothing dispatches to these helpers yet, so the blast radius before #2752 is zero. Every invariant whose violation would ship a checkpoint that loads and is quietly wrong — a name two ranks both write, a file number landing on another ranks file, a merge split across ranks, a half-reverted name space, an index naming a missing shard, an overwritten checkpoint — is either refused from the single gathered payload on every rank or covered by a world > 1 test the author reports as mutation-checked. 270 added source lines is within the sizing budget, and the split into #3100 / #3101 / this PR answers the earlier oversize note.

🤖 Generated with Claude Code

shengliangxu · 2026-10-29T22:48:01Z

/ok to test 55b62f6

meenchen

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Approved: collective refusal before writes resolves the remaining checkpoint-preservation concern, and reverse-conversion failures retain consistent names across ranks.

No action needed:

✔️ Resolved since the last review: prior-checkpoint corruption and message-less reverse failures.

Complex PR: 2 existing test files modified or removed. Looping in a human for approval.

[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>

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>

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>

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>

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 (285c30b) goes, with its tests. The staged .tmp writes stay. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…ghts 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>

shengliangxu · 2026-10-30T00:48:13Z

/ok to test f796006

meenchen

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

LGTM: the prior correctness concerns are resolved; a human should sign off on the expanded CPU/GPU test coverage.

Needs action:

Sign off on the test expansion in tests/unit/torch/export/test_distributed_writer.py and tests/gpu/torch/export/test_distributed_writer_gpu.py; existing coverage remains intact.

No action needed:

✔️ Resolved since the last review: checkpoint preservation and message-less reverse failures; the earlier split, import and merge-group fixes remain intact.

kevalmorabia97

Each tensor is written once from its owner, the shard plan and failures go through the #2736 collectives so all ranks agree, and the tests fail when key logic is broken (49 unit tests pass, plus the 2-GPU test).

shengliangxu added this pull request to stack #2755 October 12, 2026 00:57

shengliangxu force-pushed the shengliangx/hfep-12a-keep-fused-fold branch from 65d7878 to e832c67 Compare October 12, 2026 01:11

shengliangxu force-pushed the shengliangx/hfep-12b-writer-pieces branch 2 times, most recently from 8aef85d to 6994370 Compare October 12, 2026 02:44

shengliangxu force-pushed the shengliangx/hfep-12a-keep-fused-fold branch from e832c67 to 52b4592 Compare October 12, 2026 02:44

shengliangxu force-pushed the shengliangx/hfep-12a-keep-fused-fold branch from 52b4592 to d314e75 Compare October 12, 2026 08:39

shengliangxu force-pushed the shengliangx/hfep-12b-writer-pieces branch from 6994370 to 5b2ee4a Compare October 12, 2026 08:39

shengliangxu changed the title ~~[16/19] Add the pieces of a distributed HF checkpoint writer~~ [OMNIML-6082] [16/19] Add the pieces of a distributed HF checkpoint writer Oct 12, 2026

shengliangxu force-pushed the shengliangx/hfep-12b-writer-pieces branch from 5b2ee4a to f738cd0 Compare October 13, 2026 00:53

shengliangxu force-pushed the shengliangx/hfep-12a-keep-fused-fold branch from 636aea0 to dbb570a Compare October 13, 2026 01:32

shengliangxu force-pushed the shengliangx/hfep-12b-writer-pieces branch from f738cd0 to 812e4e1 Compare October 13, 2026 01:32

shengliangxu removed this pull request from stack #2755 October 13, 2026 17:42

shengliangxu added this pull request to stack #2775 October 13, 2026 17:43

shengliangxu changed the title ~~[OMNIML-6082] [16/19] Add the pieces of a distributed HF checkpoint writer~~ [OMNIML-6082] [16/20] Add the pieces of a distributed HF checkpoint writer Oct 13, 2026

shengliangxu force-pushed the shengliangx/hfep-12a-keep-fused-fold branch from dbb570a to e85bc68 Compare October 13, 2026 19:38

shengliangxu force-pushed the shengliangx/hfep-12b-writer-pieces branch from 812e4e1 to 6c1cb0e Compare October 13, 2026 19:38

shengliangxu removed this pull request from stack #2775 October 13, 2026 19:38

shengliangxu added this pull request to stack #2781 October 13, 2026 19:40

shengliangxu changed the title ~~[OMNIML-6082] [16/20] Add the pieces of a distributed HF checkpoint writer~~ [OMNIML-6082] [18/21] Add the pieces of a distributed HF checkpoint writer Oct 13, 2026

shengliangxu force-pushed the shengliangx/hfep-12a-keep-fused-fold branch from e85bc68 to be5d536 Compare October 13, 2026 21:12

shengliangxu force-pushed the shengliangx/hfep-12b-writer-pieces branch from 6c1cb0e to a8e0c27 Compare October 13, 2026 21:12

shengliangxu removed this pull request from stack #2781 October 14, 2026 00:13

shengliangxu force-pushed the shengliangx/hfep-12b-writer-pieces branch from a8e0c27 to a3be04a Compare October 14, 2026 00:14

meenchen reviewed Oct 29, 2026

View reviewed changes

claude Bot reviewed Oct 29, 2026

View reviewed changes

claude Bot approved these changes Oct 29, 2026

View reviewed changes

meenchen reviewed Oct 29, 2026

View reviewed changes

This was referenced Oct 30, 2026

[OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137

Open

[OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138

Open

shengliangxu added 6 commits October 30, 2026 00:47

meenchen reviewed Oct 30, 2026

View reviewed changes

kevalmorabia97 approved these changes Oct 30, 2026

View reviewed changes

This was referenced Oct 30, 2026

[OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136

Open

[OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133

Open

GPUs	Result	Time
1	2 skipped (need two GPUs)	0.5 s
2	2 passed	17.6 s
4	2 passed	23.5 s

Conversation

shengliangxu commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Summary by CodeRabbit

Uh oh!

copy-pr-bot Bot commented Oct 12, 2026

Uh oh!

coderabbitai Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Sequence Diagram(s)

Uh oh!

github-actions Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

shengliangxu commented Oct 13, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

meenchen left a comment

Choose a reason for hiding this comment

Uh oh!

meenchen Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

meenchen left a comment • edited by shengliangxu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

shengliangxu commented Oct 29, 2026

Uh oh!

shengliangxu commented Oct 12, 2026 •

edited

Loading

coderabbitai Bot commented Oct 12, 2026 •

edited

Loading

github-actions Bot commented Oct 12, 2026 •

edited

Loading

codecov Bot commented Oct 12, 2026 •

edited

Loading

meenchen left a comment •

edited by shengliangxu

Loading

meenchen left a comment •

edited by shengliangxu

Loading