[OMNIML-6082] [15/21] Add the pieces of a distributed HF checkpoint writer by shengliangxu · Pull Request #2751 · NVIDIA/Model-Optimizer

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

What does this PR do?

Type of change: New feature

Exporting an FSDP2 model gathers every shard to rank 0 first, so rank 0 needs host RAM for the whole model and the gather costs world x the bytes. An expert-parallel model cannot go through that path at all. These are the pieces of a writer where every rank writes its own share instead (modelopt/torch/export/distributed_writer.py):

_materialize_owned 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).
Placements it cannot assemble, and shards that do not tile the global shape, raise rather than risk a silently mis-assembled checkpoint.
_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.
_write_even_shards:
- bin-packs each rank's set into evenly sized files numbered into one global range;
- writes them with safetensors;
- unions the index, failing on a name two ranks both wrote, on every rank.
Failing together. The export directory's creation with the file counts, the shard writes and rank 0's index write each run inside the failure-carrying all_reduce ([4/21]), 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 [16/21].

Usage

N/A: internal to the distributed export ([16/21]).

Testing

New tests/unit/torch/export/test_distributed_writer.py, on CPU under gloo, drives each piece directly:
- HF shard-size parsing;
- 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 write one name: both raise. 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.
Locally on this revision, the reordered series rebased on main 264fd26: the full unit suite on transformers 5.18, and on 5.5 apart from the tests whose ONNX, TensorRT or diffusers packages that environment lacks; CI's transformers-free unit lane; pre-commit over the whole PR against main; and the Sphinx docs build. The same at every PR of the stack.

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 [16/21].
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; merges after [14/21].

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

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 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 47.37% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 38 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title accurately describes the main change: adding internal components for a distributed Hugging Face checkpoint writer. The progress marker and issue identifier add context without making the tit…
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…

Full details: Title check

Explanation

The title accurately describes the main change: adding internal components for a distributed Hugging Face checkpoint writer. The progress marker and issue identifier add context without making the title misleading.

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
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2751/
Built to branch `gh-pages` at 2026-10-24 05:22 UTC. Preview will be ready when the GitHub Pages deployment is complete.

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

Codecov Report

❌ Patch coverage is 84.71616% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.83%. Comparing base (730bbeb) to head (dd1344d).

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

Additional details and impacted files

@@                           Coverage Diff                            @@
##           shengliangx/hfep-12a-keep-fused-fold    #2751      +/-   ##
========================================================================
+ Coverage                                 61.75%   61.83%   +0.07%     
========================================================================
  Files                                       655      656       +1     
  Lines                                     73098    73327     +229     
========================================================================
+ Hits                                      45144    45344     +200     
- Misses                                    27954    27983      +29

Flag	Coverage Δ
unit	`61.83% <84.71%> (+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-20T21:10:28Z

/ok to test 992b8f3

shengliangxu · 2026-10-21T17:44:20Z

/ok to test ed900e5

shengliangxu · 2026-10-21T18:32:51Z

/ok to test b3cb0f2

shengliangxu · 2026-10-22T00:53:42Z

/ok to test fab7a67

shengliangxu · 2026-10-22T01:15:01Z

/ok to test 3eec2a9

shengliangxu · 2026-10-22T06:47:10Z

/ok to test 2d85f70

shengliangxu · 2026-10-22T16:55:59Z

/ok to test 52fcfa8

shengliangxu · 2026-10-22T18:45:19Z

/ok to test c6bf2af

shengliangxu · 2026-10-22T19:10:56Z

/ok to test 6b650ce

shengliangxu · 2026-10-22T20:24:39Z

/ok to test c274dc5

shengliangxu · 2026-10-22T20:59:36Z

/ok to test c69cff5

shengliangxu · 2026-10-23T01:16:39Z

/ok to test fa8c1e5

shengliangxu · 2026-10-23T05:30:46Z

/ok to test e1e746f

shengliangxu · 2026-10-23T07:22:17Z

/ok to test 7095ff2

shengliangxu · 2026-10-23T08:05:02Z

/ok to test 67caebe

shengliangxu · 2026-10-23T17:07:10Z

/ok to test b6d674b

shengliangxu · 2026-10-23T21:38:56Z

/ok to test 81aca54

shengliangxu · 2026-10-23T22:34:06Z

/ok to test 918e6e6

shengliangxu · 2026-10-23T23:54:34Z

/ok to test 2556c46

shengliangxu · 2026-10-24T00:42:56Z

/ok to test 675a8e0

shengliangxu · 2026-10-24T01:24:21Z

/ok to test b551cba

Exporting an FSDP2 model gathers every shard to rank 0 first, so rank 0 needs host RAM for the whole model and the gather costs world x the bytes, and an expert-parallel model cannot go through that path at all. These are the pieces of a writer where every rank writes its own share instead (modelopt/torch/export/distributed_writer.py): - _materialize_owned gives each tensor exactly one owner rank and makes it whole there: replicated and plain values need no transport, and dim-0 shards travel point-to-point, batched and bucketed per peer -- one bucket per key deadlocked at 16 ranks. Placements it cannot assemble, and shards that do not tile the global shape, raise rather than risk a silently mis-assembled checkpoint. - _revert_whole_sd runs the reverse weight conversion on those whole tensors, which transformers' conversion ops need: they cannot express a placement. - _write_even_shards bin-packs each rank's set into even files numbered into one global range, writes them with safetensors, and unions the index on rank 0, failing on a name two ranks both wrote. 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 follows. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

A rank whose shard write raised left the writer and met its peers in a different collective: they raised an unrelated error or, when the failure was rank 0's name-collision check or index write before the closing barrier, waited out the process group's timeout. The per-rank file counts, the shard writes and rank 0's index write now each run inside the failure-carrying all_reduce, so every rank raises with the failing rank's error, and every rank checks the union of the written names for a collision. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

_size_to_bytes is one table of a decimal and a binary unit, case and padding, bare numbers and the refusals, not a row per unit. _even_bins is one table of the file sizes it produces; the every-key-once test goes, since the shard writer's test fails when a tensor is dropped. The _materialize_owned refusals are one table: a dim-1 shard, a smaller or 2-D mesh, a padded chunk. The _is_full_world_shard0 accept/reject tests go with it, as the materialize test's sharded tensor and that table cover both answers. _nbytes keeps only its DTensor case, the revert fallback no longer pins its warning text, and the two pooled jobs use the shared bounded_collectives with plain asserts: the peer of a failing rank must get that rank's error. The local gloo_world_of_one goes for the conftest's. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

_write_even_shards created the export directory before its first collective, so a rank whose filesystem refused the mkdir raised alone while its peers waited in "shard counts" until the process-group timeout, against the function's own contract that every step a rank can fail alone runs inside all_reduce. Create it inside the "shard counts" all_reduce instead, which adds no collective. The one-rank failure test gains a "dir" case. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

_materialize_owned refused an unsupported placement before its metadata gather, and checked tiling against each rank's own view of the global shape. A rank whose view diverged (a packed-weight rewrap once re-inferred the global shape per rank) then raised alone while its peers waited in the gather or the point-to-point transport until the timeout. Gather every rank's placement verdict and view of the global shape in the one metadata collective, and decide every refusal from that data before anything moves, so all ranks raise. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Each rank of the no-gather writer reverses the conversion mapping on the tensors it owns. A merge needs its whole group, which the per-tensor ownership can spread over ranks, and a rank holding none of a group fails it. That rank kept its in-memory names while the others wrote hub names, so the checkpoint mixed the two. The ranks now agree on the outcome: if the reverse fails on any rank, every rank keeps the in-memory names, as the single-process export does on a failure. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu · 2026-10-24T05:14:58Z

/ok to test dd1344d

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

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

shengliangxu mentioned this pull request Oct 22, 2026

[OMNIML-6059] [1/21] Add the DP=EP all-to-all expert transport #2738

Merged

shengliangxu added 7 commits October 24, 2026 05:11

Date distributed_writer 2026, as LICENSE_HEADER does

dd1344d

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Conversation

shengliangxu commented Oct 12, 2026 • edited by coderabbitai Bot 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!

Built to branch gh-pages at 2026-10-24 05:22 UTC. Preview will be ready when the GitHub Pages deployment is complete.

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 20, 2026

Uh oh!

shengliangxu commented Oct 21, 2026

Uh oh!

shengliangxu commented Oct 21, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 22, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 24, 2026

Uh oh!

shengliangxu commented Oct 24, 2026

Uh oh!

shengliangxu commented Oct 24, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

shengliangxu commented Oct 12, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Oct 12, 2026 •

edited

Loading

github-actions Bot commented Oct 12, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-10-24 05:22 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

codecov Bot commented Oct 12, 2026 •

edited

Loading