[shengliangxu]

Note

[2/3] of #2751's split (2026-10-29), at its review's request: [1/3] #3100 (the reverse conversion's hook), [2/3] this PR (ownership and transport), and [3/3] #2751 (reverse conversion and shard writing). Merge in that order.

Base: main. [1/3] #3100 merged on 2026-10-29 (8ab63386152), and this PR was rebased onto it with its own diff unchanged.

What does this PR do?

Type of change: New feature (part of the no-gather distributed export)

This is the first part of a checkpoint writer for sharded models (FSDP2 and expert parallelism) in which every tensor moves once, whole, to one owner rank and is written from there (modelopt/torch/export/distributed_writer.py).

Main's FSDP2 export already writes each rank's share, but each FSDP2 unit is first unsharded 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. Here no rank all-gathers anything, and the tensors a merge needs end up together, so the reverse conversion in [3/3] can run.

_materialize_owned:

• Owner assignment: gives each dense tensor one owner, largest first onto the lightest rank, seeded with each rank's expert bytes. Each group a merge needs (together) goes to one owner as a unit.

• Transport:

  • dim-0 shards travel once, to their owner, over batched point-to-point sends bucketed by (peer, dtype), so the operation count is O(world × dtypes), not O(keys × peers). Both sides build the buckets from the same sorted keys and owner map, so every send meets a receive of its length. Rounds of at most 4096 keys bound how many keys one transfer carries;
  • replicated and plain tensors are unwrapped by their owner and dropped elsewhere;
  • an empty tensor is created by its owner.

  The send buffers are freed before a rank assembles what it owns: the ops and their requests hold them, and they are dropped inside _run_to_completion. On top of the shards it is handed, which the model keeps, a rank therefore holds in each round copies of its shards of keys other ranks own plus the shards it receives, and then those received shards plus the tensors assembled from them. With balanced shards that is about twice its owned set; a layout that puts most rows on one rank (four one-row tensors over four ranks) makes that rank's outgoing copies the larger part. The mesh check is memoized per mesh: FSDP2 gives every parameter the same mesh, so a few hundred keys would otherwise each redo it.

• NCCL initialization: before the first transfer, every rank runs a one-element all-reduce on the transport's device. A batched P2P that is NCCL's first call on a group needs every rank of the group, and on the cpu:gloo,cuda:nccl group the setup creates, the metadata collective runs on Gloo, so a transfer only some ranks take part in could otherwise be that first call.

• Refusals: an unsupported placement, shards that do not tile their global shape, an expert still a DTensor, dim-0 shards on more than one device type (a CPU-offloaded shard beside one on the accelerator), and ranks that disagree on keys, on a shard's dtype or device, on a whole tensor's dtype or shape, or on groups are all refused. Every refusal is decided from one gathered metadata collective, so it is the same on every rank.

Also here, the one-rank gloo world the meta-DCP load tests used moves to the shared tests/unit/torch/conftest.py, in the same text #3099 adds.

Usage

Internal to the no-gather export (#2752).

Testing

• tests/unit/torch/export/test_distributed_writer.py, on CPU under gloo:

  • sizes and the placement check;
  • every tensor whole and plain after ownership;
  • each unsupported layout refused, on every rank;
  • two ranks moving uneven and empty shards of two dtypes, with what a rank sent freed before it assembles.

• tests/gpu/torch/export/test_distributed_writer_gpu.py runs on every visible GPU, from two. It moves BF16, FP8 and packed-NVFP4 uint8 shards of uneven and empty extent, a tensor with no rows, a replicated and a plain tensor over NCCL, and checks each ends whole, byte for byte, on exactly one rank. The same logic passes on two gloo ranks locally; GPU runs are requested.

• These tests come from [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751's reviewed suite, moved here with their code. 43 pass locally, and pre-commit passes against main.

• GPU runs on GB200 of test_distributed_writer_gpu.py, ran at b4baa594bdf (the transport code is unchanged since; only docstrings and tests changed):

  GPUs	Result	Time
  1	skipped (needs two GPUs)	0.5 s
  2	passed	21.1 s
  4	passed	23.2 s

• Transport measured on GB200 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. _materialize_owned(sd, set(), rank, world) ran at [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751's f577fc05af0, whose transport is this PR's. Rank 0 is shown; all ranks were near-identical.

  GPUs	Input per rank	Owned per rank	Transport	Peak device memory above the input
  4 (one node)	8.01 GiB	8.00 GiB	6.22 s	12.01 GiB
  8 (two nodes × 4)	4.00 GiB	4.00 GiB	14.15 s	7.01 GiB

  The peak is torch.cuda.max_memory_allocated() after reset_peak_memory_stats, minus the allocation before the transport. During the transfer a rank holds a send buffer and a 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. The 8 GPUs span two nodes, so that time includes cross-node traffic and does not scale like-for-like against the one-node run.

• Re-measured at [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751's 239260c7d8d (same transport code) with the input DTensors kept alive, as the model keeps its parameters in a real export. Same state dict and hardware:

  GPUs	Owned per rank	Transport	Peak device memory above the input (transport, and the write that follows)
  4 (one node)	8.00 GiB	6.83 s	14.01 GiB
  8 (two nodes × 4)	4.00 GiB	14.76 s	7.51 GiB

  The peak now 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 the docstring's bound of about twice the owned set.

• Review round, 2026-10-29 (46e42e10a74, rebased onto [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100's b9fbb242300): the two-rank test checks that two merge groups sharing a key end on one owner, where byte balancing alone would split them; a _StridedShard on dim 0 is pinned as refused (in torch>=2.10 it is not a Shard, so it was already refused); the docstrings keep the constraints, and the cost derivations moved here. Pre-commit passes against main; test_distributed_writer.py passes (16).

• Review round 2, 2026-10-29 (055b57bd7cf, cjluo-nv's review bot):

  • Dim-0 shards on more than one device type are refused on every rank before anything moves; so is a key whose ranks disagree on its dtype or device. A round's transfers are one batch, which cannot mix devices.
  • Every dense tensor's dtype and shape are compared across ranks, not only its byte count.
  • Each round of transfers runs in a helper, so its receive buffers are freed before the next round starts.
  • The refusal test gains a shard on another device on one rank, shards on two device types, and three whole-tensor disagreements at equal byte counts (dtype, shape, empty shapes). The transfer test runs rounds of two keys and checks that no earlier round's receive buffer is alive when the next starts. Each fails without its change.
  • test_distributed_writer.py passes (21), and pre-commit passes against main. The transport's GPU test was re-run at this round's head; see below.

• GPU run, 2026-10-29, ran at 055b57bd7cf (the same code as the rebased head ba9f1c5773b). GB200, torch 2.11.0a0 (NGC 26.02), transformers 5.9.0:

  tests/gpu/torch/export/test_distributed_writer_gpu.py (1 test)	1 GPU	2 GPUs	4 GPUs
  Result	skipped (needs 2 or more)	passed, 22.3 s	passed, 24.0 s

• Rebased on 2026-10-29 onto main 2b900c8fa41 (the [hf] floor bump and the QAT Diffusers-attention fix), own diff unchanged.

• Review round 3, 2026-10-29 (bb80474fe8f, b73c052e45f; cjluo-nv's review bot):

  • Every rank initializes NCCL with a one-element all-reduce on the transport's device before the first transfer. A new GPU test, test_a_first_transfer_most_ranks_sit_out_finishes_on_a_fresh_group, spawns a fresh cpu:gloo,cuda:nccl group on 3 or more GPUs, runs no collective on it first, and transfers a two-row tensor that only ranks 0 and 1 hold. A rank stuck in NCCL's initialization exits through a watchdog, so the test fails instead of hanging. Its logic passes on three gloo CPU ranks; its GPU run is below.
  • The docstring's memory bound names a rank's outgoing shard copies as well as what it receives and assembles; with skewed dim-0 shards, the outgoing copies can exceed the owned set.
  • Pre-commit passes against main; test_distributed_writer.py passes (21).

• GPU runs, 2026-10-29, at b73c052e45f (cjluo-nv's review bot asked for the fresh-group test with and without the NCCL initialization). GB200, torch 2.11.0a0 (NGC 26.02), CUDA 13.1, NCCL 2.29.2, transformers 5.9.0.

  test_a_first_transfer_most_ranks_sit_out_finishes_on_a_fresh_group, on 4 GPUs:

  Arm	Result
  With the NCCL initialization (b73c052e45f)	passed, 23.2 s
  Same head with only the initialization removed (84548cd6d4c)	failed, 75.8 s: a rank exited non-zero through the 60 s watchdog

  The file's pool-based tests:

  test_distributed_writer_gpu.py	2 GPUs	4 GPUs
  Result	1 passed, 1 skipped (the fresh-group test needs 3 or more), 21.5 s	2 passed, 44.0 s

• Review (juhi10071998), 2026-10-29 (764036cd64e): the docstring says ranks are weighed by bytes, not free device memory, so a large embedding or lm_head must fit on its owner. Docstring only; pre-commit passes against main.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive; nothing calls it 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; its code was reviewed inside [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751.

Additional Information

Split out of #2751, part of the expert-parallel PTQ series for Hugging Face MoE models.

Summary by CodeRabbit

• New Features
  • Distributed checkpoint materialization can assemble supported sharded tensors into full tensors without gathering all data onto one rank.
  • Dense tensors are assigned to a single rank, while replicated and ordinary tensors are handled locally. Supported dimension-0 shards, including uneven and empty shards, are combined in rank order.
  • Expert tensors remain on their existing ranks, and assembled tensors are retained by their assigned owner.
• Bug Fixes
  • Unsupported layouts and inconsistent tensor metadata now produce errors instead of being assembled incorrectly.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: shengliangx/hfep-12b2-writer-transport
• File: .coderabbit.yaml
• Commit: 764036c
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 5d686467-ddf4-418b-b437-eef97d902a82

Commits

Reviewing files that changed from the base of the PR and between b73c052 and 764036c.

Files selected for processing: 1

• modelopt/torch/export/distributed_writer.py

Files skipped from review as they are similar to previous changes: 1

• modelopt/torch/export/distributed_writer.py

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

Walkthrough

The pull request adds a distributed checkpoint materializer. It validates tensor metadata across ranks, assigns dense tensors and merge groups to owners, and transfers supported dimension-0 shards in bounded rounds. Unit and GPU tests cover materialization, ownership, transfer behavior, and invalid inputs.

Changes: Distributed tensor materialization

Layer / File(s)	Summary
Classification and collective validation modelopt/torch/export/distributed_writer.py, tests/unit/torch/export/test_distributed_writer.py, tests/unit/torch/conftest.py, tests/unit/torch/utils/test_meta_dcp_load.py	The writer adds tensor and mesh helpers and checks rank agreement on keys, layouts, devices, shapes, dtypes, and merge groups. Unit tests cover helper behavior and invalid inputs. The one-rank Gloo fixture is defined in shared torch test configuration and removed from the meta DCP test module.
Dense ownership and local materialization modelopt/torch/export/distributed_writer.py, tests/unit/torch/export/test_distributed_writer.py	The writer assigns dense units to ranks by size and current load, including expert bytes in the initial load. Owners retain plain tensors, unwrap replicated tensors, or create empty tensors; other ranks drop their copies.
Shard transfer and assembly modelopt/torch/export/distributed_writer.py, tests/unit/torch/export/test_distributed_writer.py, tests/gpu/torch/export/test_distributed_writer_gpu.py	The writer transfers supported dimension-0 shards in batched flat buffers and concatenates them on owners in rank order. Unit and GPU tests check assembled values, unique ownership, sparse transfers, and buffer release.

Priority: Normal

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

Change: Feature

Possibly related issues

• [RFC] NVIDIA Model Optimizer — Product Roadmap #1699: This PR adds distributed checkpoint export ownership and shard transport without all-gathering tensors.

Merge Risk: Minimal · up to 76403

No actionable risk introduced by this PR remains. The existing export paths do not expose the previously reported companion-scale scenario.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	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 names the main change: moving each tensor of a sharded model whole to a single owner rank. It matches the distributed writer ownership and transport work described in the summary and objecti…
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 adds one modelopt Python file and test-only Python changes. The added modelopt code contains no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_re…

Full details: Title check

Explanation

The title names the main change: moving each tensor of a sharded model whole to a single owner rank. It matches the distributed writer ownership and transport work described in the summary and objectives. It also carries the [OMNIML-6082] [2/3] tracking prefix, which is acceptable. The title is a single, specific sentence and does not use vague terms.

Full details: Security Anti-Patterns

Explanation

PASS. The PR adds one modelopt Python file and test-only Python changes. The added modelopt code contains no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval(), exec(), or # nosec. The diff adds no pyproject.toml or requirements.txt dependencies.

Full details: Docstring Coverage

Explanation

Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 16 functions across 4 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

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

[shengliangxu]

/ok to test b4baa59

[shengliangxu]

/claude review

[claude]

[SUGGESTION Performance] Non-owned local shards stay alive through the transfer, so the peak is higher than the docstring claims

torch.cat(...) at line 121 already copies every shard this rank owes a peer into the send buffer, so from that point the source local_sd[k] entries for owner[k] != rank are dead weight — but they are only dropped in the assembly loop at line 143, after the recv buffers are allocated and the transfer has run.

Peak therefore holds, concurrently: all local shards of the chunk (L), the send copies (≈ L·(world-1)/world), and the recv buffers (R). The docstring at line 169 says "A rank holds at most about twice its owned set: the shards it receives and the tensors assembled from them", which accounts for R + O but not the L + sends that coexist with R. Since L and the owned set are both ≈ model/world, the real peak is closer to 3× the per-rank share than 2×. For a feature whose whole value is the memory bound, I'd either correct the docstring or tighten the actual peak.

Tightening is cheap — after the send ops are built and before the recv buffers are allocated, nothing in the chunk needs the non-owned sources any more:

        # Packed into the send ops above, so the sources are dead: drop them before the recv
        # buffers are allocated, rather than after the transfer.
        for k in part:
            if owner[k] != rank:
                del local_sd[k]

        for srcr, dt in sorted(recv_groups, key=lambda kd: (kd[0], str(kd[1]))):
            ...

with the if owner[k] != rank: del local_sd[k]; continue at lines 142-145 then reducing to if owner[k] != rank: continue. Note the recv-buffer device at line 127 is taken from local_sd[ks[0]], which is a key this rank owns, so it is unaffected by the earlier deletion.

The existing two-rank test already measures live send storages at each assembled store; extending it to assert the non-owned sources are gone before assembly would pin the tighter bound.

[shengliangxu]

I'm not changing the transport here.

The values _materialize_owned receives are views of the model's own parameters: #2752 takes them from get_model_state_dict, and the model keeps them. Deleting them from the dict earlier therefore frees nothing in the export. The docstring's bound counts what the transport holds on top of those shards, and 46e42e1 makes the docstring say so.

Measured on GB200 with a 32 GiB synthetic sharded state dict:

• 600 keys of shape (8741, 4096): BF16, FP8 and uint8.
• Uneven dim-0 shards.
• Timed _materialize_owned at the head before this round. The transport code is unchanged since.

GPUs	Owned per rank	Transport	Peak device memory above the input
4 (one node)	8.0 GiB	6.2 s	12.0 GiB
8 (two nodes)	4.0 GiB	14.2 s	7.0 GiB

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, which is (w-1)/w of its input.
• The receive buffers hold the other ranks' shards of the tensors it owns, which is (w-1)/w of its owned set.

That is 6 + 6 GiB at 4 GPUs and 3.5 + 3.5 GiB at 8, as measured, and within the docstring's bound.

[claude]

[SUGGESTION] Expert disjointness is assumed but, unlike every other cross-rank precondition here, not refused

The docstring states experts are "whole, plain and disjoint per rank", and the function does check two of those three from the metadata collective: still-a-DTensor is refused (m[6], line 232), and wholeness follows from plainness. Disjointness is the one that is only asserted in prose.

It is violable in a configuration this series targets. In an EP × FSDP hybrid with EP group size < world size — world 8, EP 4, FSDP 2 — each expert tensor lives on two ranks. _materialize_owned leaves every expert exactly where it is on every rank, so both ranks would carry that key into the write, and [3/3] would emit the same tensor into two shard files with one index.json entry. That is a silently duplicated/ambiguous checkpoint, in a function whose stated design is to refuse anything it cannot handle from the single gathered metadata rather than risk a bad write.

I am filing this as a suggestion rather than a blocker because nothing calls this yet and the caller arrives in [3/3] — a global key→file index built there has to gather every key anyway and would surface the duplicate. Two reasonable resolutions:

• Cheapest: reword the docstring from a description into an explicit precondition on the caller ("the caller must guarantee expert_key_set members are held by exactly one rank; not checked here"), and note where it is enforced.
• Or check it: the gather already carries sum(_nbytes(...)) over experts (m[1]). Adding a per-rank expert key count plus a stable digest of the sorted expert key list lets you refuse when the summed counts exceed the union — without paying the O(world × experts) string payload that gathering the full expert key sets would cost on a 128-expert MoE.

Either way the reader learns whether this is guarded or delegated; right now the sentence reads as a guarantee the code does not make.

[shengliangxu]

The docstring now states this as the caller's precondition (46e42e1): experts must come whole, plain and disjoint per rank, as the caller splits them.

The refusal does exist, in [3/3] (#2751). _write_even_shards gathers every rank's key list in its first collective and refuses a name held by more than one rank before any file is written. A duplicated expert fails there, on every rank, at no extra cost. That is the global key check you describe, and it covers dense keys as well.

[claude]

Claude review — [2/3] Move each tensor of a sharded model whole to one owner rank

Findings: CRITICAL 0 · IMPORTANT 1 · SUGGESTION 4

Scope. Full review (trigger comment carried no scoping instructions). 5 files / +695 −9; I read all of them — modelopt/torch/export/distributed_writer.py in full, plus both new test files and the two test-infra moves. No prior Claude review on this PR to dedupe against; CodeRabbit auto-review is disabled on this non-default base branch, so nothing was skipped as already-covered.

Most impactful finding

is_shard(0) accepts _StridedShard (line 45) — IMPORTANT. _StridedShard subclasses Shard, so is_shard(0) is true for it, and a strided placement on a 1-D full-world mesh passes the gate into shard0. Its chunks interleave by split_factor rather than sitting contiguously at rank offsets, but the tiling check only verifies the dim-0 extents sum to the global dim 0 — a strided layout tiles just as well. So assembling with torch.cat(pieces, dim=0) in rank order yields a silently row-permuted tensor, which is the one outcome the module's own docstring says it refuses to risk. Reachability is narrow (the 2-D-mesh and full-world-mesh checks already catch FSDP2+TP and FSDP sub-meshes), but the failure mode is unrecoverable corruption rather than an error, and the fix is to make the check exact instead of subclass-tolerant.

The four suggestions: the real memory peak is ~3× the per-rank share rather than the ~2× the docstring claims, and dropping non-owned sources right after torch.cat packs them would actually buy the tighter bound; together's happy path is untested (only its disagreement refusal is); expert disjointness is asserted in prose but, unlike the other preconditions, not refused; and a third of the file is prose, some of it cost derivations that CONTRIBUTING.md places in the PR description.

What I checked and found correct

The cross-rank reasoning is the risky part of this change and it holds up:

• Every refusal is symmetric. For the key-set, dim-0-classification, size and grouping checks, each rank compares the gathered values against its own. If the values are not all equal then every rank necessarily finds a differing peer, so all of them raise — no rank is left in a collective its peer never joins. The cpu_dist_workers suite covers this on both the agreeing and the disagreeing rank for all seven cases.
• Send/recv buckets pair exactly. Sender (owner[k], dtype) groups and the owner's (src, dtype) groups are built from the same sorted key list, the same owner map and the same all_shapes emptiness test, both sorted by (peer, str(dtype)). Element counts match by construction (sum(_numel(all_shapes[sender][k])) on both sides), the dtype equality the buffer sizing depends on is refused up front, and the (peer, dtype) bucket is one op per side so in-order p2p matching per pair is preserved even with two dtypes between the same ranks.
• Chunk boundaries are aligned. shard0 is derived from sorted dense_keys and its membership is verified across ranks, so the chunk = 4096 rounds are identical everywhere — no rank runs a transfer round its peer skipped.
• Assembly order is shard order. _is_full_world_shard0 pins mesh.mesh == range(world), so global-rank order is mesh order and the torch.cat reproduces what full_tensor() would have built. The padded-chunk hazard is caught by the tiling check (short-circuits correctly when ranks disagree on the global shape, so a multi-element gshapes cannot produce a false pass), and empty shards are skipped symmetrically on both sides rather than concatenated.
• Owner assignment is deterministic. The unit sort (-bytes, u) and the min key (loads[i], i) both break ties explicitly, over inputs (dense_sizes, expert_bytes, merged_units) that are all verified identical — so no tensor can be dropped everywhere or claimed twice.
• The together union pass is transitively correct. The new unit is the union of every member's old unit and all of its members are rewritten, so [A,B], [C,D], [B,C] collapses to one unit regardless of group order.
• No host syncs on the metadata path. _nbytes, numel() == 0 and placements are all metadata; the one tensor read (mesh.mesh.tolist()) is CPU and memoized. Freeing the send buffers right after req.wait() is safe under NCCL because ProcessGroupNCCL::pointToPoint calls recordStream on them, so the caching allocator defers reuse.
• torch.distributed.tensor is fine against the declared torch>=2.10; the module is intentionally absent from export/__init__.py's star imports (all-private, no caller yet), and the gloo_world_of_one move to tests/unit/torch/conftest.py is a clean lift that keeps the meta-DCP tests working while adding the CPU mesh the new tests need.

Risk

Low-to-moderate. Additive and unreferenced, so nothing regresses today; the GPU test covers BF16/FP8/packed-uint8, uneven and empty shards, a no-row tensor, replicated and plain over NCCL, and the CPU suite covers the refusals on every rank. The residual risk is concentrated in the one placement gate above — this code's contract is "refuse anything not exactly Shard(0) on a full-world 1-D mesh", and that gate is currently one subclass wider than the contract. Worth closing before [3/3] starts depending on the transport.

🤖 Generated with Claude Code

[shengliangxu]

/ok to test 46e42e1

[shengliangxu]

/claude review

[shengliangxu]

Thanks. Here is what each item became; it is pushed as 46e42e1, rebased onto #3100's b9fbb24.

• IMPORTANT, _StridedShard: already refused. In torch>=2.10, _StridedShard is bound on Placement, not on Shard, and is_shard(0) is False for it. The case you suggested is now in the refusal test. Details are in the thread.
• Memory peak: the transport is unchanged. Its inputs are views of the model's parameters, so dropping them earlier frees nothing in the export. The docstring now says its bound is on top of those shards. A GB200 measurement is in the thread and the description.
• together's happy path: now tested in the two-rank job, and mutation-checked.
• Expert disjointness: the docstring now states it as the caller's precondition. [3/3] refuses a key held by two ranks before writing.
• Docstring length: trimmed. The cost derivations moved to the description.

[claude]

Claude review — [2/3] Move each tensor of a sharded model whole to one owner rank

Claude review passed — no blocking issues found. LGTM.

Findings: CRITICAL 0 · IMPORTANT 0 · SUGGESTION 2

Scope. Full review (the trigger comment carried no scoping instructions). 5 files / +705 −9; I read all of them — modelopt/torch/export/distributed_writer.py in full, both new test files, and the two test-infra moves. CodeRabbit's auto-review is disabled on this non-default base branch, so nothing was skipped as already-covered.

Since the previous round (b4baa59 → 46e42e1)

The prior review's five findings, re-checked against the current code:

Prior finding	Status now
IMPORTANT — is_shard(0) accepts _StridedShard	Addressed by test, downgraded. A strided-shard0 refusal case was added. The gate itself is unchanged, so the refusal still rests on torch's internal class hierarchy rather than on an explicit check — re-raised as a SUGGESTION, not a blocker, since CI decides whether the refusal fires.
SUGGESTION — together's happy path untested	Addressed. The two-rank test's third layout drives [["x","y"],["y","z"]], with sizes chosen so per-key balancing alone would split the group, and asserts all three land on one rank.
SUGGESTION — memory peak is ~3× not ~2×	Withdrawn — the docstring was right. _run_to_completion clears ops before the caller assembles, so the send copies are freed before the torch.cats allocate. Peak extra is max(send + recv, recv + assembled), and both terms are ~1× the owned set. No change needed here.
SUGGESTION — expert disjointness asserted in prose, not refused	Still open; non-blocking.
SUGGESTION — prose volume / cost derivations belong in the PR body	Still open; non-blocking.

This round's findings

1. [SUGGESTION] distributed_writer.py:42 — make the placement gate exact (type(p) is Shard) instead of relying on is_shard(0) excluding _StridedShard. The new test passes only under that torch-internal behaviour; stating the contract in code makes the refusal version-independent and turns the test into documentation of an intent.
2. [SUGGESTION] test_distributed_writer.py:33-34 — drop the safetensors/transformers importorskips. The module under test imports neither, so on a torch-only install they silently skip the seven cross-rank refusal cases, which reads as green.

What I checked and found correct

The cross-rank reasoning is the risky part of this change, and it holds up:

• Every refusal is symmetric. For the key-set, dim-0-classification, size, dtype, global-shape and grouping checks, each rank compares the gathered values against its own. If the values are not all equal, every rank necessarily finds a differing peer, so all of them raise — no rank is left in a collective its peer never joins. gshape = next(iter(gshapes)) being computed before the len(gshapes) != 1 test is harmless: the or short-circuits, so the arbitrarily-chosen element is never compared. The cpu_dist_workers suite covers all seven cases on both the agreeing and the disagreeing rank, under bounded_collectives() so a hang fails fast rather than stalling CI.
• Send/recv buckets pair exactly. Sender (owner[k], dtype) groups and the owner's (src, dtype) groups are built from the same sorted key list, the same owner map and the same all_shapes emptiness test, both sorted by (peer, str(dtype)). Element counts match by construction, the dtype equality the buffer sizing depends on is refused up front, and one op per (peer, dtype) per side preserves in-order p2p matching per pair even with two dtypes between the same ranks. chunk = 4096 rounds are derived from the verified-identical sorted shard0 list, so no rank runs a round its peer skipped.
• Assembly reproduces full_tensor(). _is_full_world_shard0 pins mesh.mesh == range(world), so global-rank order is mesh order; torch.chunk is what Shard._split_tensor uses, so the GPU test's hand-built shards match FSDP2's cut. Padded chunks are caught by the tiling check, and empty shards are skipped symmetrically on both sides rather than concatenated.
• Owner assignment is deterministic and total. The unit sort (-bytes, u) and the min key (loads[i], i) both break ties explicitly, over inputs (dense_sizes, expert_bytes, merged_units) that are all verified identical, and unit_of covers every dense key — so no tensor can be dropped everywhere or claimed twice. Asymmetric expert sets are fine by design: expert_bytes is read from the gathered metadata, not computed locally.
• The together union pass is transitively correct. Each new unit is the union of every member's current unit and all of its members are rewritten, so unit_of stays a valid partition regardless of group order — [A,B], [C,D], [B,C] collapses to one unit either way round.
• empty0 cannot desynchronize. DTensor.numel() is global, so ranks agree; a global shape disagreement surfaces as a shard0/empty0 classification split, which the all_shapes key-set check names.
• No host syncs on the metadata path. _nbytes, numel() == 0 and placements are all metadata; the one tensor read (mesh.mesh.tolist()) is CPU and memoized, and id(mesh) cannot collide because local_sd pins every mesh alive through the classification loop. Freeing the send buffers right after req.wait() is safe under NCCL because ProcessGroupNCCL::pointToPoint calls recordStream on them, so the caching allocator defers reuse — and the two-rank test proves the free actually happens, via StorageWeakRefs sampled at each assembled store.
• The module is intentionally absent from export/__init__.py (all-private, no caller yet — grep confirms nothing in modelopt/ references it), torch.distributed.tensor is fine against the declared torch>=2.10, and the gloo_world_of_one lift to tests/unit/torch/conftest.py keeps the meta-DCP tests working while adding the CPU mesh the new tests need, now with try/finally teardown.

Risk

Low. Additive and unreferenced, so nothing regresses today. The GPU test covers BF16/FP8/packed-uint8, uneven and empty shards, a no-row tensor, replicated and plain over NCCL; the CPU suite covers the refusals on every rank plus the merge-group and send-buffer-lifetime properties. Both suggestions are robustness and test hygiene, not behaviour. Fine to merge ahead of [3/3] once the strided gate is either made explicit or confirmed green in CI.

🤖 Generated with Claude Code

[shengliangxu]

Thanks for the approval. On the items the review lists as still open, and the two new ones:

• Expert disjointness: _materialize_owned's docstring now states it as the caller's precondition ("experts must come whole, plain and disjoint per rank, as the caller splits them"). [3/3]'s _write_even_shards refuses, on every rank, a name held by more than one rank before it writes anything; this is in 46e42e1.
• Prose volume: the three places you listed are trimmed in 46e42e1, and the cost derivations are in the description now.
• Exact placement gate: kept as is_shard(0). In torch>=2.10, is_shard is itself the exact check (a C++ typeid comparison), and type(p) is Shard would refuse a Shard built on the C++ side. Details are in the thread.
• The importorskips: kept, because [3/3]'s tests in this file import transformers at module level. Details are in the thread.

[cjluo-nv]

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

Changes requested: mixed-device shards can fail before their peers complete transport.

Needs action:

• Fix device-aware packing in distributed_writer.py:90–121; add a mixed CPU-offloaded/CUDA shard test, or refuse that layout collectively before mutation.
• Validate every dense tensor’s dtype and global shape in distributed_writer.py:197–206; test disagreements that preserve byte counts, including empty tensors.
• Release round-local receive-buffer references in distributed_writer.py:126–147; extend the lifetime regression across the 4096-key boundary.
• Confirm NCCL first-use works when some ranks have no P2P operations; add a cold-group sparse-participation test.

No action needed:

• Standard headers match LICENSE_HEADER; the fixture relocation preserves existing assertions.
• sdk-gpt61-20261029T014604Z/PR-3101: static review only; pinned SHAs and reported test runs could not be independently verified.

[shengliangxu]

/ok to test 055b57b

[shengliangxu]

Thanks. The first three items are done in 055b57b, with details in their threads; the fourth needs no change.

• Mixed devices: dim-0 shards on more than one device type, or a key whose ranks disagree on its dtype or device, are now refused on every rank before anything moves.
• Whole-tensor metadata: every dense tensor's dtype and shape are compared across ranks, not just its byte count.
• Round-local buffers: each round runs in a helper, and the transfer test checks across a round boundary.
• NCCL first use with sparse participation: the batched transfer is never the group's first collective. _materialize_owned runs dist.all_gather_object(meta, my_meta) on the default group (distributed_writer.py:217) before it calls _gather_shard0_to_owners (line 356). Every rank reaches that collective, and the P2P ops use the same default group, so the communicator exists before any rank skips a send. The GPU test ran over NCCL on 2 and 4 GB200 GPUs, with uneven and empty shards, and passed.

[juhi10071998]

Reviewed distributed_writer.py: I found no blocking bug.

What I checked: I ran _materialize_owned on 2, 3 and 4 gloo ranks (torch 2.10, CPU) with:

• uneven dim-0 shards (7 rows), and fewer rows than ranks (empty shards)
• uint8, bf16, and fp8 viewed as uint8, side by side
• replicated and plain tensors, plus per-rank expert tensors
• overlapping merge groups chained transitively ((a, d), (d, g))

In every case each tensor ended up whole on exactly one rank, byte-equal to the original, with no duplicates, and each merge group landed on a single owner. Send and receive grouping are built from the same gathered metadata, and every refusal is decided from the single all_gather_object, so all ranks agree and none is left hanging.

Non-blocking:

• Owners are balanced by bytes, not by free device memory, so a large embedding or lm_head could OOM its owner. Consider documenting this, or using a memory-aware bound.
• This PR and [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099 add the same gloo_world_of_one fixture to tests/unit/torch/conftest.py, so expect a trivial rebase conflict for whichever merges second.

[shengliangxu]

/ok to test 764036c

[shengliangxu]

@juhi10071998, thank you for the multi-rank checks. Running _materialize_owned on 2, 3 and 4 ranks, with empty shards, mixed dtypes and chained merge groups, covers layouts the unit tests take one at a time.

• Owners balanced by bytes, not free memory: documented in 764036c. _materialize_owned's docstring now says ranks are weighed by bytes, not by free device memory, so a large embedding or lm_head must fit on its owner. I did not build a memory-aware bound. It would need each rank's free memory gathered in the metadata collective, and that number changes as the export runs.
• The duplicate gloo_world_of_one fixture: agreed. [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099 and this PR add it in byte-identical text, so whichever merges second will be rebased onto main and drop its copy. That happens once one of them merges.

[kevalmorabia97]

The owner assignment and the batched send/receive pairing are consistent across ranks, the earlier cold-group NCCL finding is fixed, and the unit tests pass here (49 passed). Two description lines are stale: the Testing note that the transport is unchanged since the GB200 run at b4baa594, and the "Claude approval: not yet" checklist line.