[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]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.

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] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 fused experts under grouped_mm (merged)
    • [OMNIML-6073] [3/3] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 calibration on grouped_mm
      • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
  • [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 meta-DCP checkpoint load (merged)
  • [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 distribution plan, run group, collectives (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 Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
          • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
            • [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
              • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
                • [OMNIML-6076] 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, src=0) runs fn on src 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 in tests/unit/torch/utils/test_distributed.py, the module's test file. 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 rejection of a plan no load could carry out: an EP degree that does not divide the world, a degree below 1, 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;
  • 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, any dense layout accepted, the EP-degree check dropped, 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 src, 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.

• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). The line this PR's rewrite of the run's cleanup deletes is main's sys.exception() form now ([OMNIML-5802] Drop Python 3.10 support and adopt Python 3.11+ features #3007).

• Review fixes, 2026-10-27 (bot review): broadcast_from and all_reduce serialize each value once, inside the guard that sends a failure to every rank, and the collective carries the bytes, so a value that passes the check can no longer fail later in the collective on one rank; two-rank tests with a value only pickle protocol 2 takes fail without the change.

• CI flake fix, 2026-10-27: bounded_collectives runs a barrier on its new group before the job's body. A pooled job that runs no collective on the group (one that mocks barrier, or one refused before it calibrates) destroyed it at once, and a peer still connecting to it then failed in new_group with "Connection closed by peer", the "Worker(s) failed" flake of the pooled tests in CI. The bounded-barrier job on the 2-rank gloo pool, under load (six job loops pinned to two cores):

  Barrier on the new group	Failed jobs
  no	12 of 3000
  yes	0 of 3000

• Review suggestions, 2026-10-27 evening: broadcast_from returns, on the sending rank too, the value loaded from the bytes it sent, so no rank's result shares state with fn's; _encoded's docstring says what its load-back proves.

• Review suggestion, 2026-10-27 night: broadcast_from's and all_reduce's docstrings say what each rank gets back, on a distributed run and in a single process.

• Review suggestions, 2026-10-27, last round: distribution_plan's docstring says a launch's processes read as single processes before their group is set up, and _encoded's says its load-back runs the value's load once more on the sending rank.

• Review fixes, 2026-10-27, last round: a failure whose text cannot be formatted (an exception whose __str__ raises) still reaches the broadcast or gather, and both two-rank jobs cover it; the advice for a model with no plan points at set_distribution_plan() only, as load_pretrained records plans from [OMNIML-6054] Route load_pretrained by the distribution plan #2779 on; distributed_run's docstring limits abort to the run's own group; _decoded carries the inline justification SECURITY.md asks for.

• Review fix, 2026-10-27, last round: set_distribution_plan refuses a plan for another number of ranks than the run it is recorded in, and records nothing then, when the run has more than one rank; a run of one can still stand in for several, as the export tests' single-process runs do. The two-rank job's new check fails without the change.

• Review suggestion, 2026-10-28: the private collective helpers' docstrings state their contracts only. The rationale they carried:

  • A rank's SystemExit travels to the others as its exit code, so every rank exits. distributed_run ends an exit with a barrier but an error with abort, so a rank exiting while its peers abort would wait in that barrier for ranks that are gone; if any rank's failure is not an exit, every rank raises an error instead.
  • _encoded loads each value back on the sending rank, so a value that cannot round-trip fails inside the caller's guard, which tells every rank. That costs one more copy and a second load on the sender. A load that depends on the loading rank (a class only the sending rank can import) can still fail on a peer alone, which is acceptable for the small values these collectives carry: a plan, a flag, a size.
  • Every rank, the sender included, gets a copy loaded from the bytes sent, so a result never aliases what fn returned.

• Review fix, 2026-10-28, second round: distributed_run's report of the body's exception formats it through _text too, so an exception whose __str__ raises still reaches abort() instead of escaping the finally; test_the_run_ends_the_group_it_set_up gains that case, which fails without the change.

• Review suggestion, 2026-10-28, fifth round: distributed_run's docstring says a SystemExit has to be every rank's, as the collectives make it, since the exit path waits on every rank in the group's teardown.

• Rebased onto main 7092db6ac5b on 2026-10-28 (no conflicts). Bot review fixes, same day: two contract caveats are back in the docstrings (a broadcast's label mismatch fails only the receivers, and the load-back proves a value loads on the producing rank only), and the tests moved into test_distributed.py, as CONTRIBUTING.md has test files mirror their module.

• Review fixes, 2026-10-28, last round: _decoded loads with map_location="cpu", as _deserialize does, so a tensor in a value lands on the CPU on every rank, the sending one included, instead of on the sending rank's device, which on a peer is another rank's GPU or absent; broadcast_from's and all_reduce's docstrings say so. No test pins the mapping: a single process returns before any value is serialized, so observing it takes a two-rank run with a CUDA tensor, and no caller in the series sends a tensor (they send a DistributionPlan, a flag, a size). pre-commit, mypy included, passes at this head.

• Review fix, 2026-10-28, evening: distributed_run warns when a timeout is passed but the caller's group is already set up, since a group's timeout is fixed when it is set up, instead of dropping it silently. test_the_run_leaves_a_group_it_found_alone passes a timeout in all four of its cases and fails without the warning. pre-commit, mypy included, passes at this head.

• Review changes (kevalmorabia97), 2026-10-28, night: _decoded loads with weights_only=True, with DistributionPlan registered via add_safe_globals. A value of an unregistered class fails its load-back on the producing rank, inside the guard, so every rank raises; the two-rank tests now expect that refusal, and register their protocol-2 class. DistributedState.device is a torch.device on every path. Each test fails without its change; this PR's tests pass (52), and pre-commit, bandit included, passes against main.

• Rebased on 2026-10-28, night, onto main bc0b6e03237, without conflicts; this PR's own diff is unchanged. pre-commit passes against main, and this PR's tests pass (52).

• Review changes (kevalmorabia97), 2026-10-28, night: _encoded builds on _serialize and _decoded calls _deserialize, which takes the collectives' map_location, so one loader sets the load flags. broadcast_from's rank parameter is src, as it shadowed rank() (Claude). The tests are trimmed: the stand-in broadcast and gather tests and their helpers are gone, since the real two-rank jobs cover those paths. The encoding's failure modes are tested once in a single process, with one integration case per collective. bounded_collectives loses its own test but stays, because later tests rely on its 10 s bound, which keeps the worker pool alive where the pool's job timeout would kill it. The plan tests that restated the dataclass are gone, and the invalid-plan cases cover each branch once. The launch-state test is folded into the end-the-group matrix, whose launch rows still tell the local rank from the global one; its entered-in-except variant now runs on two rows. Test docstrings are one line. 36 tests pass, and pre-commit passes against main.

• Rebased on 2026-10-28, night, onto main e9f432affbb (Bump minimum torch to 2.10 and drop older-torch code paths #3091, the torch 2.10 floor). Bump minimum torch to 2.10 and drop older-torch code paths #3091 made _get_fsdp2_mesh call fully_shard.state, so [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088's two ownership tests, which counted those calls, fail on main. This PR carries the fix, counting _fsdp2_param_groups, the scan itself, instead. The same commit is in [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745: whichever merges first lands it, and the other's copy drops out on rebase. This PR's tests and test_utils.py pass (63), and pre-commit passes against main.

• Review suggestions (Claude), 2026-10-28, late night: a test round-trips a DistributionPlan through the collectives' encoding, the one class they carry; it fails with the add_safe_globals registration removed. The carried test_utils.py fix stays until [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 or [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merges. It counts the ownership scan itself, which is narrower than counting fully_shard.state (since Bump minimum torch to 2.10 and drop older-torch code paths #3091 also part of the mesh lookup), and that scan is what the cache saves. 37 tests pass, and pre-commit passes against main.

• Review suggestion (Claude), 2026-10-28, late night: distributed_run's docstring says a single process takes a timeout without a warning, since a run of one waits on no peer and hf_ptq.py ([OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754) passes one on every run, for its launches. The single-process test passes a timeout and fails if the run warns about it. 37 tests pass, and pre-commit passes against main.

• Review suggestion (Claude), 2026-10-28, late night: the unused-timeout warning is only for a caller's group of several ranks, so a single process in a one-rank group of the caller's takes a timeout quietly too; the single-process test gains that row, which fails with the warning. 38 tests pass, and pre-commit passes against main.

• Review suggestions (Claude), 2026-10-28, late night: the timeout warning's condition is is_distributed_run(), the same test before the run sets a group up; comments say a loaded plan skips __post_init__ (the sending rank built and checked it), and that the object collectives pickle the message while weights_only bounds its value, which the Security section now says too. Kept: the process-wide add_safe_globals registration, as the collectives ask of a caller's classes, and the all_reduce name (open for the code owners). 38 tests pass, and pre-commit passes against main.

• Rebased on 2026-10-28, late night, onto main 1ca8bd1fc4d after [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged: the test_utils.py fix both PRs carried landed with [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, so its copy here dropped out, and this PR's diff is distributed.py and its tests again. 65 tests pass (this PR's and test_utils.py), and pre-commit passes against main.

Security

The value loads with weights_only: _decoded loads each collective's payload with torch.load(..., weights_only=True), as _deserialize does. Built-in types and tensors load as they are; DistributionPlan, the one class the collectives carry, is registered with torch.serialization.add_safe_globals, and a caller that sends another class registers it. The transport is torch's object collectives (broadcast_object_list, all_gather_object), which pickle the message envelope (label, payload bytes, error text) between the job's ranks, one trust boundary; weights_only bounds the value inside it. The earlier weights_only=False exception is gone (review, 2026-10-28).

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?: ✅ At c3576fe83b5; the predicate and comment change after it is from that review.

Additional Information

Part of the expert-parallel PTQ series for Hugging Face MoE models. Based on main and independent of #2745; the series' other PRs build on it.

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

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

❌ Patch coverage is 97.97297% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.14%. Comparing base (bb13106) to head (c3576fe).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/utils/distributed.py	97.97%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2736      +/-   ##
==========================================
+ Coverage   73.82%   74.14%   +0.32%     
==========================================
  Files         662      650      -12     
  Lines       75949    72948    -3001     
==========================================
- Hits        56069    54087    -1982     
+ Misses      19880    18861    -1019     

Flag	Coverage Δ
examples	?
unit	61.93% <97.29%> (+0.08%)	⬆️

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-28 22:47 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 cf4597d