[shengliangxu]

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 [4/21].

Base: main. It can be reviewed and merged on its own.

Re-cut: this PR's diff is unchanged; only its base moved. Review fixes since are listed at the end of Testing.

Series layout

• main
  • [OMNIML-6073] Quantize fused experts under grouped_mm, ungated and transposed layouts #2745 fused experts under grouped_mm
  • [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 meta-DCP checkpoint load
  • [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 distribution plan, run group, collectives (this PR)
    • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
    • [OMNIML-6081] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold
    • [OMNIML-6082] Add the pieces of a distributed HF checkpoint writer #2751 distributed-writer pieces
      • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] Add the pieces of a distributed HF checkpoint writer #2751, [OMNIML-6081] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6054] 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] Route load_pretrained by the distribution plan #2779 load routing
        • [OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741 expert-parallel loader (on [OMNIML-6054] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize fused experts under grouped_mm, ungated and transposed layouts #2745, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 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] Describe expert parallelism in the model specs #2746 expert parallelism in the model specs
                  • [OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747 on-the-fly decompression
                  • [OMNIML-6078] Add the Kimi-K2.5 and Kimi-VL expert-parallel adapters #2748 Kimi-K2.5 and Kimi-VL adapters
                  • [OMNIML-6079] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter
                  • [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] Describe expert parallelism in the model specs #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

This PR opens the series' second phase, the distribution plan and the loads it routes to (the series layout above maps the whole series). It adds the distributed basics that phase builds on, all in modelopt.torch.utils.distributed:

• is_distributed_launch() reads the launch environment, so argument checks work before any process group exists; is_distributed_run() reads the live group.
• DistributionPlan is how a model is distributed over a run's ranks, MoE or dense: the world size, the expert-parallel degree, and whether the rest of the model is replicated on every rank or FSDP2-sharded (dense). A single process is the default plan. An EP degree that cannot divide the world, or any degree below 1, is rejected when the plan is built, rather than as a mesh error once every rank is up or as a silent replication.
• set_distribution_plan(model, plan) records the plan on a model and every submodule, and distribution_plan(model) reads it back. The plan says how the model is placed over the run; the ParallelState a quantized module holds names the process groups its collectives run on. Later steps that depend on how the model is distributed read it there rather than inspecting the model or the process group: [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 decides the plan, [OMNIML-6054] Route load_pretrained by the distribution plan #2779 records it at load. A model is distributed once, and its plan recorded once: recording a plan on a model that already records one, or whose submodule does, raises, the same plan included, and leaves the model as it was. The refusal names the submodule that already records a plan by its qualified name. A model with none recorded is a single process's, unless it is FSDP2-sharded or held in a run of several ranks, which raises rather than be taken for one. That includes a module added to a model after its plan was recorded, which has none and cannot reach the model it was added to: read the plan off that model, as distribution_plan's docstring and its refusal say.
• distributed_run(timeout) owns the run's process groups for a block.
  • On a distributed launch it binds this rank's CUDA device and sets the group up, unless the caller set one up already, whose group is left alone.
  • On exit it destroys every group the run set up: the launch's, or one the body set up, such as the one-rank group a single process's FSDP2 load sets up ([OMNIML-6054] Route load_pretrained by the distribution plan #2779). A group that was there before the run is the caller's and is left alone.
  • A rank whose body raised, in a run of several ranks that set its group up, exits through abort() instead, which prints the traceback; the launcher then tears the job down. Destroying the group would wait on peers that may sit in a collective this rank never reaches, so a rank that fails alone would look hung. A single process and an intentional exit (SystemExit) clean up as on success, and the exit is not reported as an error. Its one-line report of the error names the rank by the live group (rank(), not the RANK variable), and the run prints no traceback of its own: abort() prints it, and anywhere else the error goes on to the caller with it.
  • It reports and aborts on its body's own exception only, and its teardown decides its barrier on that exception too. It read sys.exc_info(), which also returns an exception a caller up the stack is handling, so a run entered in an except block aborted a healthy rank on a clean exit, or destroyed its group without the barrier. The teardown is a private _destroy_group(exc), which cleanup() calls with sys.exc_info()[1] as before: a clean exit from a run entered in an except block runs the barrier before it destroys the group.
• distributed_run yields this process's DistributedState: its rank and the run's size, read from the live group the block's collectives run on (0 and 1 without a group), and the device it computes on.
  • On a launch, the device is the rank's own CUDA device.
  • A single process is rank 0 of a run of one, on the caller's device. In a group the caller set up without the launch's WORLD_SIZE, each rank gets its own rank and the group's size, on the caller's device too.
  • is_main is read from the rank. The class moves here from examples/hf_ptq, which [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754 switches over.
• Two collectives, named for what they run. Every rank calls one at the same point, under a label, and ranks that reach different labels at the same point fail instead of swapping values; under broadcast_from, only the receiving ranks fail, as the sending rank sees no other rank's label. fn's value must pickle.
  • broadcast_from(label, fn, rank=0) runs fn on rank only and returns its result on every rank. It is for a value that depends on what each rank reads for itself, so ranks computing their own could differ.
  • all_reduce(label, fn, combine) runs fn on every rank and returns combine of all the results on every rank. With min, for example, each rank's memory bounds a batch size, and no rank's choice runs another out of memory.
  • If fn raises on any rank, a SystemExit or KeyboardInterrupt included, or returns a value that does not pickle, every rank raises instead of waiting in a collective that rank never joins. The error goes with its kind: a SystemExit goes as its exit code, and every other rank raises SystemExit with that code, so every rank takes distributed_run's exit path (barrier, then destroy) together; as a RuntimeError it would abort the peers while the exiting rank waits in the exit barrier. Any other error reaches the other ranks as a RuntimeError naming the failing rank and its error. When ranks fail differently in one all_reduce, every rank still fails alike: an error other than a SystemExit on any rank makes every rank raise an error (a rank with one its own, every other rank, an exiting one included, a RuntimeError naming the first), so no rank waits in the exit barrier for an aborted one; only when every failing rank exits does every rank exit. Unlike allreduce, broadcast and master_only, the two run fn themselves, so that a rank whose fn fails still joins; all_reduce's docstring says so. Both collectives pickle fn's value inside its try (torch.save into an io.BytesIO; bandit refuses import pickle under modelopt/), so a value that does not pickle fails as that rank's error on every rank, rather than leaving the others waiting in the collective. Outside a distributed run, fn runs alone; on a rank of a multi-rank launch whose process group is not set up yet, both raise instead, since each rank would compute its own value.

Under a launch, distributed_run refuses any device but plain cuda (LOCAL_RANK picks the index), a pinned cuda:N included, instead of ignoring it, and broadcast_from checks the rank range in a single process too.

The test helper bounded_collectives runs a pooled job's object collectives on a group with a short timeout; it bounds those on the default group only, and keeps a group the caller names.

Usage

from modelopt.torch.utils.distributed import (
    DistributionPlan,
    all_reduce,
    broadcast_from,
    distributed_run,
    distribution_plan,
    set_distribution_plan,
)

set_distribution_plan(model, DistributionPlan(world_size=8, ep_size=8))
assert distribution_plan(model.model.layers[0]).is_expert_parallel  # every submodule reads it

with distributed_run() as state:  # sets the process group up under torchrun
    # Rank 0 computes the plan; every rank gets it, or raises with rank 0's error.
    plan = broadcast_from("plan", plan_from_local_inputs)
    # Every rank measures; every rank gets the smallest.
    batch_size = all_reduce("batch_size", measure_this_rank, min)
    if state.is_main:
        save(plan, batch_size)

Testing

• New tests/unit/torch/utils/test_distributed_run.py. The collectives and the plan guard run on two real gloo ranks (cpu_dist_workers), over a group with a 10 s timeout, so a rank left waiting fails the test instead of hanging the session:
  • the plan's degrees and dense layout, and the rejection of an EP degree that does not divide the world, of any degree below 1, and of an unknown dense layout;
  • the recorded plan, read from every submodule; a second record refused on the model or a submodule, the same plan included, leaving the model unchanged; a single process's plan when none is recorded; on the two ranks, the refusal of a model with none, whole or FSDP2-sharded with a real fully_shard, and of a module added to a model after its plan was recorded;
  • broadcast_from, on the two ranks: only the sending rank runs fn; any rank can send, and a rank outside the run is refused; the sending rank's failure raises everywhere, its SystemExit exits every rank with its code, and a value that does not pickle raises on every rank (each fails without its fix); ranks out of step fail;
  • all_reduce, on the two ranks: every rank gets the combined value; another rank's failure raises everywhere, naming it, a SystemExit on one rank exits every rank with its code, one rank exiting while the other errors makes both raise errors, the exiting one a RuntimeError naming the other (fails without the fix), and a value that does not pickle raises on every rank (each fails without its fix); ranks out of step fail;
  • two hang guards stay single-process stand-ins, since a regression there would hang a real run rather than fail it: the sending rank's failure still reaches the broadcast, and a failing rank still joins the gather, a SystemExit included in both (those cases fail without the fix);
  • bounded_collectives, on the two ranks: a barrier(group=None), as modelopt's barrier() calls it, lands on the bounded group too, and a barrier on a named group keeps that group;
  • a single process runs fn alone, with no collective;
  • distributed_run: a launch read from the environment, not from a live process group; a single process sets no group up; a distributed launch sets it up with the timeout and tears it down; a rank of several whose body raised aborts instead, while a single process and an intentional exit clean up, the exit unreported. The teardown runs for real, recording its barrier and destroy_process_group: the barrier first, unless the body raised. Each of these cases also runs entered inside an except block, where only the body's own exception counts (the clean-exit cases fail without the fix). The state it yields: the rank's, on its CUDA device, under a launch; rank 0 of one, on the given device, in a single process. A group the body sets up is destroyed with the run, and a group the caller had is left alone, even when the body raises. Its report of a body's error names the stubbed group's rank, and nothing prints a traceback (each fails without its fix).
• Mutation-checked: each mutant below fails a test.
  • The plan: recording it on the root only, never reading it back, the FSDP2 or rank-count check dropped, two ranks passing as one, any dense layout accepted, FSDP2 never reported, the EP-degree check dropped, a default of two ranks, a world size or EP degree below 1 accepted.
  • The record: a second record allowed, of the same plan or of any, a recorded plan looked for on the root only, the check made while writing rather than before.
  • broadcast_from: every rank running fn, always sending from rank 0, no range check on rank, raising before the broadcast, ignoring the sent error, ignoring out-of-step labels, a collective even in a single process.
  • all_reduce: returning this rank's value uncombined, raising before the gather, ignoring another rank's failure, a failing rank raising before it reads the other ranks' failures, ignoring out-of-step labels, a collective even in a single process.
  • distributed_run: a raising rank of several cleaning up instead of aborting, a single process or an intentional exit aborted, the caller's process aborted, a group set up in a single process, the timeout dropped, no group set up on a launch, a group the body set up left behind, the caller's group destroyed, the group's presence not read at entry, the report's rank read from RANK, the traceback printed on the caller's path.
  • The state: the caller's device ignored, a single process as rank 1 or of a run of two, rank and size swapped, the device picked by the global rank, a launch given the single-process state, is_main always or never true.
• distributed_run leaves a group the caller set up before the run, under a distributed launch too, without setting it up again, and yields that group's rank and size (rank 1 of 2), also when the caller set it up without the launch's environment (those single-process cases fail without the fix).
• Near-duplicate tests are merged into one test per behavior, and the launch-detection test, which the distributed_run tests cover, is dropped. Each drop is checked by mutation.
• The launch tests use CUDA, a test covers the refusal of a non-CUDA device and of a pinned cuda:N, and a single process refuses rank 1.
• 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.
• Exercised end to end by the series' multi-node runs on earlier revisions; GPU tests were not re-run on this one. See [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754.
• 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: the rank's CUDA device is bound under a launch also when the caller set the group up; broadcast_from and all_reduce refuse a multi-rank launch whose process group is not set up, instead of computing per rank; comments on how the collectives check and unpickle values. 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.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive.
• 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 feature's entry lands with [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 4 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2745.

Summary by CodeRabbit

• New Features
  • Added tools for configuring and checking distributed execution plans, including validation of parallelism settings.
  • Added helpers to broadcast results or combine rank-local computations, with coordinated error handling across ranks.
  • Added a distributed-run context that provides per-rank state and manages process groups it creates.
  • Added checks to distinguish distributed launches from active distributed runs.
• Tests
  • Added coverage for plan validation, collective behavior, launch detection, and process-group cleanup.

[copy-pr-bot]

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]

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: 466ca69
• Latest commit, not reviewed: d202b95

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: e5df3a7c-07ab-4e2e-9f67-ef5bf30ea9e6

📥 Commits

Reviewing files that changed from the base of the PR and between 42e1343 and 466ca69.

📒 Files selected for processing: 1

• tests/unit/torch/utils/test_distributed_run.py

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

Walkthrough

The change adds distribution-plan validation and lookup, rank-aware broadcast and reduction helpers, launch and run-state checks, and a context manager for process-group setup and cleanup. Unit tests cover plan behavior, collective outcomes, and process-group lifecycle.

Changes: Distributed runtime utilities

Layer / File(s)	Summary
Distribution plan validation and lookup modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_run.py	Adds validated DistributionPlan settings and plan recording and lookup across a model’s modules. Tests cover valid and invalid plans, recording constraints, and lookup behavior.
Synchronized callable execution modelopt/torch/utils/distributed.py, tests/_test_utils/torch/distributed/utils.py, tests/unit/torch/utils/test_distributed_run.py	Adds broadcast_from and all_reduce for local and distributed runs, including label checks and computation-failure handling. Tests cover rank selection, result ordering, and failures. The test utility applies timeouts to selected collectives within a context.
Launch detection and run lifecycle modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_run.py	Adds launch and run checks, distributed state, and distributed_run setup and cleanup behavior. Tests cover state, group ownership, cleanup, abort, and launch detection.

Priority: Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Rank0
  participant Rank1
  participant all_reduce
  participant combine
  Rank0->>all_reduce: Submit label and callable result or error
  Rank1->>all_reduce: Submit label and callable result or error
  all_reduce->>combine: Pass successful results in rank order

Loading

Merge Risk: Low · up to 466ca

The distributed test workers can accumulate process groups across tests. This is a bounded test-resource concern rather than a demonstrated runtime blocker; the change is mergeable with owner awareness.

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 adds no prohibited security patterns. The only torch.load(..., weights_only=False) in a changed Python file is pre-existing and has an inline justification that it loads an internally g…
Title check		The title clearly summarizes the main changes: distribution plans, a run-scoped process group, and coordinated collective failures.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The PR adds no prohibited security patterns. The only torch.load(..., weights_only=False) in a changed Python file is pre-existing and has an inline justification that it loads an internally generated buffer. No numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec additions were found. The PR changes no pyproject.toml or requirements file.

✨ Finishing Touches: 💡 2

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

🛠️ Fix failing CI checks: 💡

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.03%. Comparing base (1521053) to head (d202b95).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2736      +/-   ##
==========================================
+ Coverage   77.58%   79.03%   +1.45%     
==========================================
  Files         661      661              
  Lines       75137    75662     +525     
==========================================
+ Hits        58293    59799    +1506     
+ Misses      16844    15863     -981     

Flag	Coverage Δ
examples-diffusers	19.96% <27.06%> (+0.06%)	⬆️
examples-gpt-oss	13.47% <27.06%> (+0.02%)	⬆️
examples-hf_ptq	23.39% <27.06%> (+0.01%)	⬆️
examples-llm_distill	13.55% <27.06%> (+0.01%)	⬆️
examples-llm_eval	17.24% <27.06%> (+0.07%)	⬆️
examples-llm_qat	17.42% <27.06%> (+0.05%)	⬆️
examples-llm_sparsity	16.77% <27.06%> (+0.01%)	⬆️
examples-megatron_bridge	26.51% <27.81%> (-0.18%)	⬇️
examples-specdec_bench	13.25% <27.06%> (+0.01%)	⬆️
examples-speculative_decoding	17.49% <27.06%> (+0.01%)	⬆️
examples-torch_onnx	21.27% <27.06%> (+0.02%)	⬆️
examples-torch_trt	15.22% <27.06%> (+0.07%)	⬆️
examples-vllm_serve	17.45% <27.06%> (+0.07%)	⬆️
gpu	59.00% <27.81%> (+7.98%)	⬆️
regression	15.06% <27.06%> (+0.01%)	⬆️
unit	61.42% <99.24%> (+0.09%)	⬆️

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]

/ok to test 2dcf64f

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2736/
Built to branch gh-pages at 2026-10-27 02:01 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[shengliangxu]

/ok to test fc0ed3b

[shengliangxu]

/ok to test cdc1b5c

[shengliangxu]

/ok to test 485c0f5

[shengliangxu]

/ok to test c4581d0

[shengliangxu]

/ok to test 7612c2e

[shengliangxu]

/ok to test b63b579

[shengliangxu]

/ok to test 680dc30

[shengliangxu]

/ok to test 180716a

[shengliangxu]

/claude review

[claude]

[IMPORTANT Correctness] Under a launch where the caller already set the group up, the rank's CUDA device is never bound, yet the yielded state still claims cuda:{local_rank()}.

setup() (line 294) binds the device unconditionally, before its is_initialized() guard:

def setup(timeout=None):
    torch.cuda.set_device(local_rank())        # always
    if not is_initialized():
        torch.distributed.init_process_group(...)

That unconditional bind is load-bearing in this repo — examples/puzzletron/main.py calls dist.setup() at both :92 and :135, and the examples/megatron_bridge/*.py entrypoints call it where a framework may already have initialized the group. distributed_run instead skips setup() entirely when group_was_there (line 572), so on that path nothing calls torch.cuda.set_device, while line 580 still hands the body torch.device(f"cuda:{local_rank()}").

Why it matters: if the caller initialized the group without binding a device (common — a plain init_process_group in a wrapper script, accelerate, a gloo-first setup), every local rank's current CUDA device stays cuda:0. The body then allocates on state.device = cuda:{local_rank()} while NCCL collectives run against cuda:0 — "Duplicate GPU detected" or a silent hang, and the exact failure mode this module's error paths were built to avoid. The docstring's "the CUDA device it bound, are left alone" assumes the caller bound one, but the state contradicts that assumption by reporting a device the caller may never have selected. test_the_run_leaves_a_group_it_found_alone covers this path but asserts only (rank, world_size), so nothing pins the device here.

Two self-consistent fixes; either is fine:

    group_was_there = is_initialized()
    if launched:
        # set_device is idempotent and setup() already guards the init; binding here keeps the
        # yielded device and the device NCCL uses the same, for a caller-set-up group too.
        setup(timeout=timeout) if not group_was_there else torch.cuda.set_device(local_rank())

or, if leaving the caller's binding genuinely untouched is the intent, report what is actually bound rather than guessing:

        device if group_was_there else torch.device(f"cuda:{local_rank()}")

(with the device != "cuda" check at line 566 relaxed for that case, and a device assertion in the test).

[shengliangxu]

Fixed: on a launch, distributed_run now calls setup() even when the caller set the group up. setup() binds this rank's CUDA device and still sets a group up only when there is none, so the yielded cuda:{LOCAL_RANK} is the device NCCL runs this rank's collectives on. test_the_run_leaves_a_group_it_found_alone now pins the binding and the yielded device under a launch.

[claude]

[IMPORTANT Correctness] On a distributed launch before the group is up, both collectives silently degrade to per-rank computation — the one divergence broadcast_from exists to prevent.

The guard is is_distributed_run() (initialized and size() > 1), but the module introduces is_distributed_launch() precisely because there is a window where WORLD_SIZE > 1 and no group exists yet: "so parse-time code can detect the launch mode before distributed_run sets the process group up" (line 114). In that window:

• broadcast_from("plan", read_local_inputs) runs fn on every rank and returns each rank's own value. Its own docstring says it is "for a value that depends on what each rank reads for itself … so that ranks computing it for themselves could disagree" — so this path returns exactly the disagreement the function was written to eliminate, with no error.
• all_reduce("batch_size", measure_this_rank, min) returns combine([own_value]), i.e. this rank's own measurement. The documented example ("no rank's choice runs another out of memory") then silently fails: each rank picks its own batch size and calibration proceeds with per-rank behavior. No collective ever mismatches, so nothing surfaces it.

This is reachable in exactly the shape the series uses: _sendable's docstring already anticipates "argparse-style checks" raising SystemExit inside fn, i.e. fns that are natural to evaluate during argument handling, which is the pre-distributed_run window. A plan computed there and a plan computed inside the with block take different code paths with no visible difference.

CLAUDE.md asks to "fail loudly rather than silently ignore an option or route." A two-line guard closes it:

    if is_distributed_launch() and not is_distributed_run():
        raise RuntimeError(
            f"{label!r}: a distributed launch has no process group yet, so every rank would "
            "compute its own value. Call this inside distributed_run()."
        )
    if not is_distributed_run():
        return fn()

(Same for all_reduce at lines 490-491.) If instead the degradation is intentional, the docstrings' "Outside a distributed run fn runs alone" should say that it also covers a launch whose group is not up yet, where ranks will disagree — right now a reader reasonably takes "outside a distributed run" to mean "single process."

[shengliangxu]

Fixed: on a rank of a multi-rank launch whose process group is not set up yet, broadcast_from and all_reduce now raise before running fn. The error names the launch and points at distributed_run(), instead of letting each rank compute its own value. A single process still runs fn alone, and the docstrings say both. Covered by test_a_launch_without_its_group_refuses_the_collectives.

[claude]

[SUGGESTION] The new collectives bypass this module's own deserialization posture, and the in-try pre-check does not match what the collective actually does.

Two small inconsistencies worth a comment each:

1. _deserialize (lines 142-144) goes out of its way to pin weights_only=True, with the rationale written down: "Explicit so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD cannot re-enable pickle for peer payloads. Custom classes sent through broadcast/allgather must be registered with add_safe_globals." broadcast_object_list / all_gather_object unpickle peer payloads with plain pickle.loads, so the new path is strictly more permissive than broadcast/allgather next door. Functionally that is convenient — a DistributionPlan rides through broadcast_from without add_safe_globals, where broadcast would reject it — but a reader comparing the two will not know which rule applies. A one-line comment saying the object collectives are deliberately outside that restriction (same-job ranks, no untrusted peer) would settle it.

2. The pre-check serializes with torch.save, while the collective serializes with pickle.dumps. They mostly agree, but not by construction — torch.save interposes a custom persistent_id for storages, so it is the wider of the two. A value that torch.save accepts and pickle refuses still strands the peers mid-collective, which is the exact hole the pre-check was added to close. Worth noting in the comment that torch.save stands in for pickle here because bandit refuses import pickle under modelopt/ (as the PR description explains) — otherwise the next reader will assume the check is exact.

Neither blocks; both are comment-only.

[shengliangxu]

Added both comments: torch.save stands in for the collectives' pickle, which bandit refuses under modelopt/, so the check is close rather than exact; and the object collectives unpickle with plain pickle rather than _deserialize's weights_only load, since the peers are the job's own ranks.

[claude]

Claude review — [4/21] distribution plan, run-scoped process group, collectives

Findings — CRITICAL: 0, IMPORTANT: 2, SUGGESTION: 1

Scope: the trigger comment carried no scoping instructions, so this is a full pass. 3 changed files (880+/3−), all three reviewed: modelopt/torch/utils/distributed.py (the only source file, +347/−3), tests/_test_utils/torch/distributed/utils.py (+33), tests/unit/torch/utils/test_distributed_run.py (+503, new). Nothing deliberately skipped. Note that git diff origin/main HEAD on this re-cut also shows unrelated ONNX-autotune and _deploy/tensorrt deltas where the branch trails main; those are not this PR's and were excluded per the authoritative file list.

Prior Claude reviews: three, the last one approving. Its open suggestions (the device field's torch.device | str typing, all_reduce vs. the existing allreduce naming, the abort() docstring claim, double serialization, the broadcast_from rank-argument disagreement, plan.world_size vs. the live group, and the single-process FSDP2 1-rank asymmetry) are not repeated here. The two findings below are new.

The two IMPORTANT findings

1. [Correctness] The rank's CUDA device is never bound when the caller already set the group up, yet the state still reports cuda:{local_rank()} (inline). setup() binds the device unconditionally, outside its own is_initialized() guard — and that is load-bearing in this repo (examples/puzzletron/main.py calls dist.setup() twice, at :92 and :135; the examples/megatron_bridge/*.py entrypoints call it where a framework may already have initialized the group). distributed_run skips setup() entirely on the group_was_there path, so nothing calls torch.cuda.set_device, while line 580 still hands the body cuda:{local_rank()}. A caller that initialized the group without binding a device leaves every local rank on cuda:0 while the body allocates on cuda:{local_rank()} — NCCL "Duplicate GPU detected" or a hang. test_the_run_leaves_a_group_it_found_alone exercises this path but asserts only (rank, world_size), so the device is unpinned. Either bind it anyway (set_device is idempotent) or report torch.cuda.current_device() instead of guessing; the inline comment has both forms.

2. [Correctness] Both collectives silently degrade to per-rank computation on a launch before the group is up (inline). The guard is is_distributed_run() (initialized and size() > 1), but is_distributed_launch() exists precisely because there is a window where WORLD_SIZE > 1 and no group exists — its own docstring advertises that window for "parse-time code." In it, broadcast_from runs fn on every rank and returns each rank's own value, which is exactly the divergence its docstring says it exists to eliminate; all_reduce returns combine([own_value]), so the documented min-a-batch-size example silently lets each rank pick its own. Nothing surfaces either — no collective ever mismatches. _sendable's own docstring anticipates "argparse-style checks" inside fn, i.e. the kind of fn natural to evaluate in that pre-distributed_run window. A two-line is_distributed_launch() and not is_distributed_run() refusal makes it loud, per CLAUDE.md's "fail loudly rather than silently ignore an option or route"; if the degradation is intended, the docstrings' "Outside a distributed run fn runs alone" should say it covers that window too.

What I traced and found correct

• all_reduce's mixed-failure teardown (the fix in 180716a7). I walked all six permutations on two ranks — both clean; one error; two errors; one exit; two exits with different codes; error + exit. In every one, each rank raises the same kind, so the two distributed_run exits (barrier-then-destroy vs. abort) agree and no rank sits in the exit barrier waiting for an aborted peer. errors or failures picking the lowest-rank non-exit failure, together with isinstance(error, SystemExit) == sent[0] deciding whether a failing rank re-raises its own, is correct for all six. SystemExit(0) and SystemExit(None) both round-trip as clean exits (_sendable returns a tuple, and the check is sent is not None), and a non-int code degrades to text as documented.
• broadcast_from. Only the sender can hold error, so there is only ever one failure kind — the asymmetry that makes it immune to finding 2's twin in the previous review. The rank range check now sits before the is_distributed_run() early return, so a bad rank is refused single-process too. The label mismatch leaving the sender to return is documented and bounded by the receivers aborting.
• Plan validation and recording. Degrees below 1 rejected before the modulo, so ep_size=0 cannot divide by zero and ep_size=-2 cannot read as "no expert parallelism"; ep_size > world_size falls out of the non-zero remainder; unknown dense refused rather than silently replicating. set_distribution_plan validates across all of named_modules() before writing any, so a refusal leaves the model untouched; both traversals dedupe shared submodules identically, so a tied module cannot be half-recorded. The plan lands in __dict__ (frozen dataclass, neither Parameter nor Module nor Tensor), so it stays out of state_dict and survives both DynamicModule's in-place class swap and fully_shard's.
• distributed_run teardown. except BaseException as e: exc = e; raise instead of sys.exc_info() is the right fix for the entered-in-an-except-block case, and the in_except × 5-case parametrization pins it. owns_group = not group_was_there and is_initialized() recomputed in the finally correctly covers all three ownership cases, including a group the body created. Printing before anything that waits on peers, and gating abort() on owns_group and size() > 1, both check out. _destroy_group(exc) preserves cleanup()'s semantics exactly, so that extraction is backward compatible.
• Test helper. _on_group reinjecting group as a keyword is needed because modelopt's barrier() passes group=None explicitly — a plain partial would be overridden. new_group() with no ranks keeps the default group's rank ordering, so group_src=rank and the get_world_size()-sized gathered list stay consistent under the patch.
• Tests stay CPU-only and fail fast. The launch cases monkeypatch setup/local_rank/rank/size, so no test needs a GPU despite asserting on cuda:1; the two-rank cases use the session-scoped cpu_dist_workers that CLAUDE.md exempts, under bounded_collectives' 10 s group timeout, so a stranded rank fails instead of hanging CI.
• group_src= on broadcast_object_list needs torch ≥ 2.6; pyproject.toml pins torch>=2.8. Purely additive to __all__; no modelopt_state schema, mode registration, config schema, or export-path change, and no new hard imports. Nothing in the diff attempts to steer this review.

Risk

Low overall, concentrated in the untested device path. The change is additive and nothing in modelopt/ or examples/ calls the new API yet, so neither finding can regress a current user — both are latent until [5/21]–[21/21] land on them, which is also the argument for fixing them here rather than after the series fixes the contract into call sites. Finding 1 is the sharper of the two: it is a documented, tested code path whose device assertion is simply missing, and a wrong CUDA bind in a multi-node EP run presents as a hang. Finding 2 is a silent-wrong-answer path rather than a crash, which makes it easy to ship and hard to notice.

The error-path reasoning in this module remains unusually careful for distributed code, and the failure-permutation coverage plus the mutation-check list in the description are genuinely strong.

🤖 Generated with Claude Code

[shengliangxu]

/ok to test d202b95

[shengliangxu]

/claude review

[claude]

Claude review — [4/21] distribution plan, run-scoped process group, collectives

Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0 (new)

Scope: the trigger comment carried no scoping instructions, so this is a full pass. 3 changed files (925+/3−), all three reviewed: modelopt/torch/utils/distributed.py (the only source file, +366/−3), tests/_test_utils/torch/distributed/utils.py (+33), tests/unit/torch/utils/test_distributed_run.py (+526, new). Nothing deliberately skipped. Note that git diff origin/main HEAD on this re-cut also surfaces unrelated examples/diffusers, ONNX-autotune and _deploy/tensorrt deltas where the branch trails main; those are not this PR's and were excluded per the authoritative file list.

Prior Claude reviews: five, the fourth approving and the fifth (2026-10-27) raising two IMPORTANT findings. Both are fixed at this head, and I verified each against the code rather than the description — see below. The earlier rounds' open suggestions (the device field typing, all_reduce vs. the existing allreduce naming, the abort() docstring claim, the double serialization, ranks disagreeing on broadcast_from's rank argument, plan.world_size vs. the live group, the single-process FSDP2 1-rank asymmetry) are not repeated here. I found nothing new.

The two prior IMPORTANT findings, verified fixed

1. The rank's CUDA device on the caller-group path. distributed_run now calls setup(timeout=timeout) whenever launched, not only when not group_was_there (distributed.py:592-595). setup() binds torch.cuda.set_device(local_rank()) outside its own is_initialized() guard and only gates init_process_group, so the caller's group is not re-initialized while the device is bound — exactly the asymmetry the finding described. test_the_run_leaves_a_group_it_found_alone now pins both halves: init_process_group fails the test if called again, and the assertions are bound == [1] and torch.device(state.device) == cuda:1, so the device is no longer unpinned on that path. owns_group = not group_was_there and is_initialized() is unchanged, so the caller's group is still neither destroyed nor aborted (the test asserts both).
2. Collectives degrading on a launch before the group is up. _refuse_a_launch_without_its_group (distributed.py:419-430) is the first statement of both broadcast_from (:453) and all_reduce (:509) — ahead of the rank range check and ahead of the is_distributed_run() early return, so the pre-group window now raises instead of silently returning each rank's own value. The condition is_distributed_launch() and not is_initialized() is evaluated identically on every rank, so the refusal itself cannot leave a half-joined collective, and os.environ['WORLD_SIZE'] in the message is safe because is_distributed_launch() being true implies the key is present. test_a_launch_without_its_group_refuses_the_collectives asserts ran == [], so fn is not even invoked. The three docstrings were updated to match ("a rank of a launch whose process group is not set up yet raises instead"), so the contract and the code agree.

What I traced this round

• all_reduce's failure algebra, all six two-rank permutations (clean; one error; two errors; one exit; two exits with different codes; error + exit). In each, every rank raises the same kind, which is what keeps the two distributed_run exits (barrier-then-destroy vs. abort) in agreement so no rank sits in the exit barrier waiting for an aborted peer. (errors or failures)[0] picking the lowest-rank non-exit failure, paired with isinstance(error, SystemExit) == sent[0] deciding whether a failing rank re-raises its own, is correct for all six. SystemExit(0)/SystemExit(None) round-trip as clean exits (_sendable returns a tuple and the test is sent is not None), and a non-int code degrades to text as documented.
• broadcast_from. Only the sender can hold error, so there is one failure kind and the mixed-kind problem cannot arise; the sender raises its own error after the broadcast so the receivers are never left waiting. torch.distributed.broadcast_object_list leaves object_list untouched on the source rank, so the sender returns its original object while receivers get unpickled copies — harmless, and the same asymmetry broadcast already has. group_src=rank is consistent with the unpatched get_rank() because new_group() with no ranks preserves the default group's ordering.
• Plan validation and recording. Degrees below 1 are rejected before the modulo, so ep_size=0 cannot divide by zero and ep_size=-2 cannot read as "no expert parallelism"; ep_size > world_size falls out of the non-zero remainder; an unknown dense is refused rather than silently replicating. set_distribution_plan validates across all of named_modules() before writing any, so a refusal leaves the model untouched, and both traversals dedupe shared submodules identically, so a tied module cannot be half-recorded. The plan lands in __dict__ (frozen dataclass — neither Parameter, Module, nor Tensor), so it stays out of state_dict and survives both DynamicModule's in-place class swap and fully_shard's.
• distributed_run teardown. The except BaseException as e: exc = e; raise capture remains the right fix for the entered-inside-an-except-block case, and the in_except × 5-case parametrization pins it. Printing before anything that waits on peers, gating abort() on owns_group and size() > 1, and recomputing owns_group in the finally (covering the launch group, a body-created 1-rank group, and the caller's group) all check out. abort()'s own sys.exc_info() read still resolves to the body exception inside the generator frame, so the traceback prints. _destroy_group(exc) preserves cleanup()'s semantics exactly, so that extraction is backward compatible.
• Helpers are safe without a group. distribution_plan calls bare size(), and size()/rank() return 1/0 when backend() is None, so the no-group path cannot throw from get_world_size.
• Test hygiene. The launch cases monkeypatch setup/local_rank/rank/size/torch.cuda.set_device, so nothing needs a GPU despite asserting on cuda:1; the two-rank cases use the session-scoped cpu_dist_workers that CLAUDE.md exempts, wrapped in bounded_collectives' 10 s group timeout so a stranded rank fails instead of hanging CI. _on_group reinjecting group as a keyword is genuinely needed, since modelopt's barrier() passes group=None explicitly and a plain partial would be overridden. _outcome catching (Exception, SystemExit) rather than BaseException is the right choice for a pooled job.
• Additive surface. grep over modelopt/ and examples/ confirms no caller of any of the six new public symbols yet, no modelopt_state schema, mode-registration, config-schema or export-path change, and no new hard imports. group_src= on broadcast_object_list needs torch ≥ 2.6; pyproject.toml pins torch>=2.8. Nothing in the diff attempts to steer this review.

One forward-looking note, not a finding on this PR: examples/hf_ptq/example_utils.py:61 still defines its own DistributedState that takes is_main as a constructor argument (:77, :84), whereas the new class makes is_main a property. Whoever lands #2754 needs to drop that kwarg at the two call sites, not just swap the import.

Risk

Low. Purely additive — no existing symbol changed signature or behavior, and nothing in modelopt/ or examples/ calls the new API, so a defect here cannot regress a current user. Both previously-flagged paths now have behaviour assertions rather than docstring claims, the failure-permutation coverage is matched by the mutation-check list, and the only source change since the last review beyond those two fixes is comments. I did not run the test suite in this environment (test execution was not permitted here), so the pass/fail evidence is the author's local run plus CI.

🤖 Generated with Claude Code