[shengliangxu]

Note

[2/3] of #2754's split (2026-10-30), cut to the 500-line budget: [1/3] #3133 library pieces → [2/3] this PR, hf_ptq's previews and whole-model placement → [3/3] #2754, hf_ptq's expert-parallel run. They merge in that order.

Base: shengliangx/hfep-base-14-hf-ptq, a branch (not a PR) that merges #2746, #2742, #2753 and #3133. This PR needs only #2779 (load_pretrained(distribution=...)) and #2752 (the plan --use_fsdp2 records), both below it; the diff below is this PR's alone.

What does this PR do?

Type of change: Bug fix and new feature

hf_ptq's generation previews, and get_model's placement of a whole model, for a model distributed over several ranks. Both read the model's distribution plan. On this base --use_fsdp2 already records its plan (#2752), so its runs get the preview changes now. get_model's pinning has its first caller in #2754, which loads every model by the plan.

• Previews keep a sharded model's ranks decoding together.
  • Each rank previews the first sample of its own calibration shard, so the ranks finish decoding at different steps. When the plan shards the model (expert parallelism, or FSDP2 over several ranks), every forward runs a collective, so a rank that stopped decoding left the others blocked in it.
  • preview_generate_kwargs hands generate synced_gpus=True for such a plan, in the previews before and after quantization and in Nemotron-VL's.
  • The ranks vote on whether to run Nemotron-VL's image preview (an all_reduce of whether each found a sample image), so every rank previews or none does. Each node reads the images from its own checkout, and the repository ships none.
  • Under synced_gpus a failing preview raises instead of being skipped, since its peers would wait for it: the VL previews, and the preview after quantization, which then stops the run before the export with a message that names --skip_generate. A single process keeps the best-effort skip (NVBug 6752977).
  • On a run of several ranks, main aborts a rank that raised, so torchrun stops the others, rather than tear the process group down while they wait in a collective that rank never reaches.
• generate keeps a MoE's experts on their backend.
  • On a GPU, generate switches the experts of a grouped_mm MoE to batched_mm for decoding (_optimize_model_for_decode, transformers 5.0 and later). batched_mm bypasses the quantizers, so the preview after PTQ decoded unquantized experts. Transformers also refuses the switch for an FSDP2-wrapped model: FSDPQwen3MoeForCausalLM does not support setting experts implementation.
  • keep_the_experts_implementation keeps every module of the model that has the hook from switching. A model can generate through a submodule, as a Nemotron-VL model's chat does through its language model.
  • It puts back a hook a module holds itself, so a hook bound per instance survives and previews nest.
  • This holds for every run, a single process included: there too the preview after PTQ now decodes on the quantized experts.
• get_model(..., distribution=plan) pins a whole model to the rank's GPU.
  • When the plan gives each of several ranks the whole model, every load path places it on the rank's own cuda:<local_rank>: the plain load, the pack-quantized and MXFP4 loads, and DiffusionGemma's sequential map. Accelerate's "auto" would run its placement in every rank alike and could put every rank's copy on the same idle GPU.
  • There, --use_seq_device_map and a device that names no GPU index are refused. A pinned copy is not sized against the node's GPUs, a limit its one-device map would ignore.
  • A single process keeps "auto", as before. The device is compared by type, since a run's torch.device never equals a string such as "cpu".
  • A plan that shards the model (expert parallelism, or FSDP2) is refused: get_model loads the whole model on every rank, and hf_ptq hands such plans to load_pretrained ([OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754).

Usage

# Each rank of a replicated run loads the whole model onto its own GPU.
model = get_model(ckpt_path, device=f"cuda:{local_rank}", distribution=DistributionPlan(world_size))

Testing

• These tests moved with their code from [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754's test_hf_ptq_expert_parallel.py, where each was mutation-checked:
  • tests/examples/hf_ptq/test_example_utils.py: the pin on every load path, a pinned copy not held to a node-wide memory limit, a single process left on "auto", a CPU device given as a torch.device, and the refusals of the sequential map and of a device without a GPU index.
  • New tests/examples/hf_ptq/test_hf_ptq_previews.py:
    • the previews before and after quantization, and Nemotron-VL's text, fallback and image generations, use synced_gpus exactly when the plan shards the model;
    • the VL image preview passes the kwargs it is handed to chat;
    • the ranks skip the image preview together when one has no image;
    • a failing text-only VL preview raises under synced_gpus.
• tests/examples/hf_ptq/test_example_utils.py also covers the experts guard:
  • the previews keep generate from switching the experts implementation, also on a language model that a model around it generates through (fails without the fix);
  • a model without the decode hook is previewed as it is;
  • a decode hook the model holds itself is put back, also when previews nest.
• tests/examples/hf_ptq/test_example_utils.py also checks that get_model without a plan leaves the decision to load_pretrained, which refuses a sharding knob over a launch without a process group.
• tests/examples/hf_ptq/test_hf_ptq_model_type.py: the preview tests' stand-in models are modules, since the previews read their plan. test_hf_ptq_args.py: the stand-in whose generate fails carries a plan.
• Pre-commit, mypy included, passes against main. The hf_ptq example tests above pass on CPU: 183 passed with transformers 5.18, and 183 with 5.5.4.
• Rebased on 2026-10-30 onto the rebuilt hfep-base-14-hf-ptq (every parent rebased onto main 99c3834): no conflicts in this PR's commits. Review fix (Claude): get_model hands load_pretrained only a plan its caller gave, so without one the load decides as before and refuses MODELOPT_* knobs that ask to shard over a launch whose process group is not set up; the new test fails without it. Pre-commit passes against main; the hf_ptq example tests above pass on CPU (transformers 5.18: 184 passed; 5.5.4: 184 passed).
• Review fixes, 2026-10-30 (cjluo-nv's bot): get_model refuses a plan that shards the model, which hf_ptq loads through load_pretrained instead; the new test fails without the check. This description no longer says that a single process previews as before (see the compatibility note). Pre-commit passes against main; test_example_utils.py, test_hf_ptq_previews.py, test_hf_ptq_model_type.py and test_hf_ptq_args.py pass on CPU (transformers 5.18 and 5.5.4: 178 passed each).
• Review fix, 2026-10-30 (cjluo-nv's bot and Claude, on the maintainers' decision): a sharded model's failed preview after quantization stops the run before the export, with a message that names --skip_generate; a single process keeps the skip, and NVBug 6752977's test is unchanged. The new test next to it, with an FSDP2 plan of two ranks, fails without the change. Pre-commit passes against main; the same four test files pass on CPU (transformers 5.18 and 5.5.4: 179 passed each).
• Review fix, 2026-10-30 (Claude round 4): on a run of several ranks, main exits a rank that raised through abort(), which prints the error and exits, instead of its teardown, whose barrier and destroy_process_group would wait on ranks blocked in a collective it never reaches; so every lockstep raise stops the run. The new test fails without the change. Pre-commit passes against main; the four test files pass on CPU (transformers 5.18 and 5.5.4: 181 passed each).
• Rebased on 2026-10-30 onto main 059588b, on hfep-base-14-hf-ptq rebuilt from the rebased [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742, [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 and [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133: no conflicts, own diff unchanged. Pre-commit passes against main; the four test files pass on CPU (transformers 5.18 and 5.5.4: 185 passed each).
• Rebased on 2026-10-30 onto main d464d30, on hfep-base-14-hf-ptq rebuilt from the rebased [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742, [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 and [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133: no conflicts, own diff unchanged. Pre-commit passes against main; the four test files pass on CPU (transformers 5.18 and 5.5.4: 185 passed each).
• Rebased on 2026-10-31 onto main e352cac, on hfep-base-14-hf-ptq rebuilt from the rebased [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742, [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 and [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133: no conflicts, own diff unchanged. Pre-commit passes against main; the four test files pass on CPU (transformers 5.18 and 5.5.4: 185 passed each).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ A single process loads as before. On transformers 5.0 and later, a preview of a grouped_mm MoE on a GPU now keeps the experts' backend, in a single process too: the preview after PTQ decodes on the quantized experts, where generate used to switch them to batched_mm, which bypasses the quantizers. --use_fsdp2 previews on several GPUs also decode in lockstep now. One failure path changes: a --use_fsdp2 run on several GPUs whose preview after quantization fails now stops before the export, where main skipped the preview and exported; --skip_generate runs it without the previews.
• 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 user-facing feature is [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754's entry)
• Did you get Claude approval on this PR?: ❌ Not yet (requested).

Additional Information

Split out of #2754 on review (meenchen's bot, "split this PR into stacked [x/3] PRs"). Merge order: #3133 [1/3], then this PR, then #2754 [3/3].

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review: 3

• main
• release/.*
• feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: shengliangx/hfep-14b-hf-ptq-previews
• File: .coderabbit.yaml
• Commit: 739b185
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 1b1b963f-e9aa-4338-8736-e2b8d996319f

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• Trigger review

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

[shengliangxu]

/ok to test f81be89

[shengliangxu]

/claude review

[claude]

[CRITICAL Distributed] The post-PTQ preview now runs generate(synced_gpus=True) on a sharded model, but its failure is still swallowed — so a rank that fails leaves every peer hanging.

Under synced_gpus, generate issues a dist.all_reduce of this_peer_finished on every decode step. If this generate raises on one rank only (which is exactly what the comment above anticipates: an offloaded layer landing on CPU, an OOM, an unsupported kernel), that rank warns, returns from post_quantize, and walks into the export's own collectives, while its peers sit in the lockstep all-reduce waiting for it. With --use_fsdp2 the group is created with a 2-hour timeout (setup_distributed_args), so the run wedges for two hours and then aborts with a collective mismatch instead of reporting the preview failure.

This is the same hazard the PR deliberately fixes for the VL previews — vlm_utils.py:204 and vlm_utils.py:281 both re-raise under synced_gpus with "the other ranks of a sharded model would wait for this one in lockstep". The non-VL post-PTQ preview is the one path left swallowing it, and lockstep is already in scope one line above:

Suggested change

	except Exception as e:
	warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
	except Exception as e:
	if lockstep:
	# The other ranks are in generate's lockstep all-reduce waiting for this one;
	# skipping the check here would leave them there until the group times out.
	raise
	warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

[shengliangxu]

Kept as is, on purpose; a maintainer should decide whether to change it.

• This is main's deliberate --use_fsdp2 best-effort skip, from the NVBug 6752977 fix (test_post_quantize_export_survives_a_failed_sanity_generate): a failed post-PTQ preview warns, and the export still runs.
• The hang for a failure on one rank alone is not new with this PR. Without synced_gpus, the other ranks wait in that rank's next FSDP2 all-gather instead of the lockstep all-reduce.
• Raising under lockstep would stop a failure every rank meets at the same step, such as an unsupported kernel, with no checkpoint, where main exports today. That is a behaviour change for --use_fsdp2, so it should be a maintainer's call rather than a review fix.

The expert-parallel route (#2754) does raise on a failed preview, since it runs inside distributed_run, which aborts the rank. If the maintainers want --use_fsdp2 to stop the run the same way, it is a small change there. Maintainers, please weigh in.

[claude]

[IMPORTANT Compatibility] Defaulting to DistributionPlan() and always forwarding it as distribution= silently disables load_pretrained's env-driven routing, including its explicit refusals.

Until now get_model passed no distribution, so load_pretrained reached resolve_distribution_plan (model_load_utils.py:833-844) and consulted MODELOPT_EP_DEGREE / MODELOPT_FSDP_SIZE. Two behaviours change:

1. Under torchrun without --use_fsdp2 the process group is not set up, so alone is true. Previously distribution is None hit the guard at model_load_utils.py:748-762, which raises "MODELOPT_FSDP_SIZE=… ask to shard the model over this N-rank launch, whose process group is not set up; load inside distributed_run()". With a one-rank plan handed in, handed() short-circuits that branch entirely: the env var is ignored without a word, plan.world_size == 1 so pinned is None, and every rank falls through to device_map="auto" — producing precisely the "Duplicate GPU detected" failure the new comment block below describes. That is the opposite of the intent, and CONTRIBUTING/CLAUDE guidance is to fail loudly rather than silently ignore an option or route.
2. In a single process, the memory-based auto-FSDP decision is also skipped. That part is arguably the point (a device_map would be refused by the sharded route), so forcing the whole-model plan is defensible — but it should be explicit, not a side effect.

Suggestion: keep the forced whole-model plan, and refuse rather than ignore a sharding request the caller made through the environment — e.g. when distribution is None and MODELOPT_EP_DEGREE/MODELOPT_FSDP_SIZE ask for a sharded layout, raise pointing at --use_fsdp2. A short comment on this line saying why a plan is always handed in (to keep the whole-model route rather than let rank 0's GPU memory choose) would also stop the next reader from "fixing" it back to distribution.

[shengliangxu]

Fixed in 06d7121: get_model forwards distribution only when its caller gave one, so without a plan load_pretrained decides as before and refuses a MODELOPT_EP_DEGREE / MODELOPT_FSDP_SIZE that asks to shard over a launch whose process group is not set up. Pinning still reads the plan (a single process's by default), and a comment on that line says why. New test test_without_a_plan_the_load_decides_and_refuses_sharding_knobs_it_cannot_honor (WORLD_SIZE=2, MODELOPT_FSDP_SIZE=world, no group) fails with the old forwarding.

[claude]

[SUGGESTION] pinned keys off plan.world_size > 1 alone, so a sharded plan also gets a device_map the sharded load cannot take.

preview_generate_kwargs (line 128) distinguishes the two cases properly — plan.world_size > 1 and (plan.is_expert_parallel or plan.is_fsdp_dense) — but here any multi-rank plan is treated as "each rank holds its own copy". Hand get_model(..., distribution=DistributionPlan(8, ep_size=8)) or DistributionPlan(8, dense="fsdp") and pinned becomes {"": "cuda:1"}, which load_pretrained then rejects in _refuse_sharded_load_options (model_load_utils.py:548) because device_map is not in _SHARDED_LOAD_OPTIONS. The error is survivable (it names device_map and suggests dense='replicate'), so this is not a correctness hole today, and every new test uses DistributionPlan(4) — but all the function's prose, the refusals, and the max_memory skip below assume replication, and nothing says so.

Using the same predicate makes the assumption explicit and lets get_model refuse early with a message about --use_fsdp2 instead of surfacing a load-option TypeError deep in the loader:

replicated = plan.world_size > 1 and not (plan.is_expert_parallel or plan.is_fsdp_dense)
pinned = {"": str(device)} if not on_cpu and replicated else None

Relatedly, the use_seq_device_map refusal at line 706 tells the user to "shard the model with --use_fsdp2", which reads oddly if the plan it is refusing is already an FSDP2 one.

[shengliangxu]

Left as is: get_model loads a whole model on every rank, and its one caller with a plan (#2754's load_model) hands it only replicated plans; expert-parallel and FSDP2 plans go to load_pretrained directly. A sharded plan passed here by mistake still fails loudly, with load_pretrained's refusal naming device_map and dense='replicate'.

[claude]

Claude review — 1 CRITICAL, 1 IMPORTANT, 1 SUGGESTION

Scope: full review (the trigger was a bare /claude review). 7 files changed (+596/−88); reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files. The base is the branch shengliangx/hfep-base-14-hf-ptq, so CodeRabbit auto-review is off for this PR — I covered a bit more ground than usual, but still skipped style/typo territory.

Findings

#	Severity	Where	Issue
1	CRITICAL	hf_ptq.py:882	Post-PTQ preview failure is still swallowed while generate now runs with synced_gpus=True → peers hang in the lockstep all-reduce until the 2h group timeout
2	IMPORTANT	example_utils.py:629	Always forwarding distribution=DistributionPlan() silently disables load_pretrained's env routing and its "process group is not set up" refusal
3	SUGGESTION	example_utils.py:651-658	pinned keys off world_size > 1 alone, so a sharded plan also gets a device_map the sharded load rejects

Most impactful

Finding 1 is the one worth fixing before merge, because it is the same bug the PR exists to fix, left in the one place it was not applied. vlm_utils.py:204 and :281 both now re-raise under synced_gpus — "the other ranks of a sharded model would wait for this one in lockstep" — but post_quantize's non-VL preview keeps its except Exception: warnings.warn(...). Under synced_gpus, generate all-reduces this_peer_finished on every decode step, so a rank whose preview raises (the offload-to-CPU case the comment right above it anticipates) warns, proceeds into the export's collectives, and leaves its peers blocked. lockstep is already bound one line above, so the fix is three lines.

Finding 2 is narrower but goes against the repo's "fail loudly rather than silently ignore an option or route" rule: torchrun + MODELOPT_FSDP_SIZE without --use_fsdp2 used to raise a pointed error from model_load_utils.py:748; with a one-rank plan handed in it now falls through to device_map="auto" and reproduces the "Duplicate GPU detected" failure the new comment block itself warns about.

What held up

• preview_generate_kwargs's sharded predicate is right: world_size > 1 and (is_expert_parallel or is_fsdp_dense). A replicated multi-rank plan runs no forward collectives, and DistributionPlan(1, dense="fsdp") correctly needs no lockstep — both parametrized in the new tests.
• keep_the_experts_implementation is sound: it shadows the class method with an instance attribute, so nullcontext is called without self (matching the real _optimize_model_for_decode(self) signature), vars(module).get(...) distinguishes a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__ for a plain attribute, and the finally restores on a failed generate. Walking model.modules() rather than just the root is what makes Nemotron-VL's chat-through-the-language-model case work, and the wrapped parametrization pins it.
• The image-preview vote is correctly ordered: _load_sample_image runs first, the all_reduce is unconditional across ranks once synced_gpus is set (which is itself derived from the plan, identical on every rank), and combine=all means every rank previews or none does.
• on_cpu = torch.device(device).type == "cpu" is a real fix for the torch.device("cpu") != "cpu" comparison, and a no-op for today's callers since --device is a plain string and the --use_fsdp2 path does not reach get_model.
• Test placement is right (tests/examples/, not tests/unit), and the new cases each pin a distinct behaviour rather than restating the implementation.

One thing I looked at and am not filing: under synced_gpus, a rank whose run_text_only_generation or AutoImageProcessor.from_pretrained fails re-raises while its peers are still heading for the all_reduce/chat. That window is bounded — torchrun tears down the group when a process exits nonzero — so the re-raise is the right call; it is only finding 1's swallow that converts the same situation into a silent multi-hour wedge.

Risk

Moderate. The blast radius is confined to examples/hf_ptq, nothing under modelopt/ changed, and the single-process path is unchanged by construction. The pinning code is inert until #2754 lands a caller, so finding 3 is forward-looking. Finding 1, though, is live today for any --use_fsdp2 run on more than one GPU.

🤖 Generated with Claude Code

[shengliangxu]

/ok to test 06d7121

[shengliangxu]

/claude review

[claude]

[SUGGESTION] The guard catches a device that names no GPU, but not the failure the comment block above it is written to prevent: a device that names the wrong GPU.

cuda:0 passed on every rank satisfies ":" not in str(device) and yields pinned = {"": "cuda:0"} on all four ranks — every rank's weights and _amax buffers land on one card and the first collective aborts with exactly the "Duplicate GPU detected : rank 1 and rank 0 both on CUDA device" the comment describes. And because get_model now always passes device_map= explicitly, it suppresses load_pretrained's own fallback pin (kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}), modelopt/torch/utils/plugins/model_load_utils.py:906), which would have placed the copy correctly regardless of what the caller named. So this is the only thing standing between the plan and that abort, and today it only checks for a colon.

Two small things that would close it:

1. Ask the question directly — torch.device(device).index is None is exactly "names no GPU index", and unlike the string test it is also right for device=0 (where str(device) == "0" has no colon, yet torch.device(0) is cuda:0, and pinned would become the invalid {"": "0"}).
2. Check the index against the rank: under torchrun the only correct value is this rank's own GPU, and dist_utils.local_rank() is already imported here. Refusing index != dist_utils.local_rank() makes the pin verify what the error message already tells the caller to pass ("pass the rank's own GPU, cuda:<local_rank>"), instead of trusting it.

Worth doing in this PR rather than in [3/3], since this is where the only caller's device will come from.

[shengliangxu]

Left as is: no caller can reach the wrong-GPU case, and the proposed rank check would refuse valid launches.

• The only caller that passes a plan, [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754's load_model, hands get_model the device distributed_run yields for this rank, cuda:<local_rank>. It never takes it from the command line, so it cannot name another rank's GPU.
• Checking index == local_rank() would refuse launchers that show each rank only its own GPU: there every rank computes on cuda:0 while LOCAL_RANK differs. get_max_batch_size reads torch.cuda.current_device() rather than LOCAL_RANK for the same reason.
• An integer device such as 0 is refused loudly with "names no GPU index", not loaded wrongly, so the string test fails safe. torch.device(device).index is None would accept it. That is a fair tidy-up for a later change, but no run needs it now.

[claude]

Claude review (round 2) — 1 CRITICAL (carried over), 0 IMPORTANT, 2 SUGGESTION

Scope: full review (bare /claude review). 7 files, +613/−88; re-reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files at 06d7121. CodeRabbit auto-review is off on this base branch, so I covered a little more ground, still skipping style/typo territory.

One caveat up front: gh api is not available to me in this environment, so I could read the previous round's review bodies but not the inline threads or your replies to them (you submitted three review comments at 01:04). I judged the two remaining items on the code at HEAD; if you already pushed back in-thread on either, treat the paragraphs below as my reaffirmation, not as a fresh claim you have not seen.

Resolved since round 1

• IMPORTANT (example_utils.py:629) — fixed and verified. routing = {} if distribution is None else {"distribution": distribution} with plan = distribution or DistributionPlan() is exactly right: without a caller plan nothing reaches load_pretrained, so it decides as before and still refuses MODELOPT_FSDP_SIZE=world over a launch whose group is not set up, while plan.world_size == 1 locally keeps device_map="auto" unchanged. test_without_a_plan_the_load_decides_and_refuses_sharding_knobs_it_cannot_honor fails without the change. Nothing else regressed: all four load_pretrained call sites in get_model (speculative, pack-quantized, MXFP4, general) carry **routing, so every load path still records a plan for the previews to read.

Still open

1. CRITICAL — hf_ptq.py:882: the post-PTQ preview still swallows a failure while generate runs with synced_gpus=True.

Unchanged at HEAD. I want to address the strongest argument for leaving it, because it is a reasonable one: a rank-local exception anywhere in hf_ptq strands its peers, this is not new, and the swallow exists so a unified-memory host does not lose a completed calibration. The first half is true, but it is what makes this case different: everywhere else the rank exits non-zero and torchrun tears the group down, so the peers die in seconds. Here the rank warns, keeps going into the export's collectives, and never exits — so nothing tears anything down, and the peers sit in _has_unfinished_sequences' all_reduce until the group timeout. The swallow converts a fast, loud failure into a silent multi-hour wedge, which is strictly worse than the pre-existing behaviour, not equal to it.

It also needs an asymmetric failure to bite, and the symmetric case the comment anticipates (a quantized layer on an offloaded device) stays harmless — every rank warns and continues. But the asymmetric ones are real on this path: each rank previews the first sample of its own shard, so prompt lengths differ, and an OOM or a data-dependent kernel error during decode hits one rank and not its peers. That is the same input asymmetry this PR exists to accommodate.

The asymmetry with your own vlm_utils.py changes is the other reason to close it: :204 and :281 both now re-raise under synced_gpus, with the comment "the other ranks of a sharded model would wait for this one in lockstep". post_quantize's non-VL preview is the one place that reasoning was not applied, and lockstep is already bound one line above:

        except Exception as e:
            if lockstep:
                # The peers are waiting for this rank in generate's lockstep all-reduce; a warn
                # here would leave them there until the group timeout.
                raise
            warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

If losing the calibration on a sharded run is the real worry, the alternative that keeps both properties is to vote: all_reduce("post-PTQ preview", lambda: ok, all) and have every rank skip the preview together — the same shape as the image-preview vote you added in vlm_utils.py. Either way the test is cheap (a stand-in whose generate raises, with and without a sharding plan).

2. SUGGESTION — example_utils.py:654-660: the pin guard refuses a device that names no GPU but accepts one that names the wrong GPU, which is the "Duplicate GPU detected" abort the comment block above it is written to prevent — and the explicit device_map= suppresses load_pretrained's own fallback pin (model_load_utils.py:906) that would have placed the copy correctly anyway. Details and two one-line options inline.

3. SUGGESTION (carried over) — example_utils.py:660: pinned keys off plan.world_size > 1 alone, so an expert-parallel or FSDP2 plan also gets a one-device device_map. I checked where that lands: the sharded route builds and places the model itself and "refuses any other keyword rather than ignore it" (model_load_utils.py:696-700, _sharded_load_kwargs), so a sharded plan handed to get_model fails loudly rather than mis-placing anything — which is why this stays a suggestion. But [3/3]'s description says it "loads every model by the plan", so the refusal that caller meets would come from deep inside load_pretrained, naming a keyword it never passed. Either gate pinned on not (plan.is_expert_parallel or plan.is_fsdp_dense) and refuse a sharding plan in get_model with a message pointing at the sharded loader, or say in the distribution parameter's docs that get_model takes a replicating plan only.

What held up this round

• keep_the_experts_implementation: shadowing the class method with an instance attribute means nullcontext is called without self, vars(module).get(...) tells a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__ for a plain attribute, and the finally restores on a failed generate. Walking model.modules() rather than the root is what makes Nemotron-VL's chat-through-the-language-model case work; nesting restores correctly (inner sees nullcontext as "own", outer deletes). Both transformers~=5.5.0 and ~=5.18.0 — the noxfile's tf_min/tf_latest — have the hook, so the monkeypatch.setattr(GenerationMixin, ...) in the test will not AttributeError on the supported floor.
• preview_generate_kwargs' predicate (world_size > 1 and (is_expert_parallel or is_fsdp_dense)) is right, and every preview branch in pre_quantize/post_quantize is selected by args flags and model_type only — nothing rank-dependent decides whether generate is entered, so the ranks enter it the same number of times. print_quant_summary / save_expert_token_count_table under is_main are local-only (no collectives), so the main-rank-only block above the preview does not start one.
• The image-preview vote is correctly ordered and unconditional across ranks once synced_gpus is set (itself derived from the plan, identical everywhere), combine=all gives every-or-none, and a failure inside the vote re-raises through the same synced_gpus branch.
• on_cpu = torch.device(device).type == "cpu" is a real fix for the torch.device("cpu") != "cpu" comparison; the device_map=None paths (Nemotron VL, and also bart/t5) still land on the rank's GPU through the existing model.to(device), so the pin holds there too.
• Tests sit in tests/examples/, not tests/unit, each new case pins a distinct behaviour, and the parametrizations (single_process / replicated / expert_parallel / fsdp2 / one_rank_fsdp2, and wrapped) are the ones that would catch a regression.

Risk

Moderate, and unchanged from round 1. Blast radius is confined to examples/hf_ptq; nothing under modelopt/ changed; the single-process path is unchanged by construction; the pinning is inert until [3/3] lands a caller. Finding 1 is the only one live today, for any --use_fsdp2 run on more than one GPU, and it is three lines.

🤖 Generated with Claude Code

[shengliangxu]

Replies to the round-2 review (the inline threads have the details):

1. CRITICAL, the post-PTQ preview's best-effort skip: kept, by the maintainer's decision. This is main's deliberate --use_fsdp2 behaviour from the NVBug 6752977 fix: a failed preview warns and the export still runs. The hang when only one rank fails is not new with this PR. Without synced_gpus, the peers wait in that rank's next FSDP2 all-gather; in this PR, a rank that raised would also wait, in cleanup_distributed's destroy_process_group. Raising under lockstep, or voting to skip, would stop a failure every rank meets at the same step with no checkpoint. That is a behaviour change for --use_fsdp2, which maintainers can choose to make. Details: [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 (comment)
2. SUGGESTION, a device that names the wrong GPU: left as is. The only caller that passes a plan ([OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754) hands get_model this rank's own cuda:<local_rank> from distributed_run. Checking the index against local_rank() would refuse launchers that show each rank only its own GPU. [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 (comment)
3. SUGGESTION, pinned for a sharded plan: left as is. get_model loads a whole model on every rank, and [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 sends expert-parallel and FSDP2 plans to load_pretrained directly. A sharded plan passed here by mistake still fails loudly. [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 (comment)

[cjluo-nv]

Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.

Changes requested: the lockstep previews are sound, but the main post-PTQ preview still lets one rank's failure send it on to export alone, which is the hang this PR exists to prevent.

Needs action:

• Fix post_quantize in examples/hf_ptq/hf_ptq.py: re-raise the sanity-generate failure when lockstep is set, as the VL paths now do, and add a test (inline).
• Explain the change to the run_text_only_generation fallback under synced_gpus: a failure that every rank hits (the Nemotron-Parse case the comment names) now aborts the run instead of falling back. Consider agreeing on the failure with all_reduce (inline).
• Decide whether get_model should refuse EP/FSDP plans, since it pins the whole model to one GPU for any plan with world_size > 1 (inline).
• Correct the PR body's "a single process previews as before". keep_the_experts_implementation also changes single-GPU previews on transformers 5.x: the post-PTQ preview now decodes on quantized grouped_mm experts, not batched_mm. Consider a CHANGELOG fix entry.

No action needed:

• About 184 added logic lines, under budget. The edits to the stand-ins in test_hf_ptq_args.py and test_hf_ptq_model_type.py are justified: distribution_plan and the experts guard call model.modules(), and the assertions are unchanged.

[cjluo-nv]

Bot comment.

Under lockstep (the --use_fsdp2 path this PR turns on), this except lets the rank that failed warn and go on to export_quantized, while its peers are still inside generate's synced step / FSDP all-gathers. That ends in mismatched collectives or a hang until the 2 h timeout. The VL paths in this PR re-raise under synced_gpus for exactly this reason. Suggest if lockstep: raise here, plus a test next to test_post_quantize_export_survives_a_failed_sanity_generate that uses a sharded plan.

[shengliangxu]

Changed in 7772c0d: under lockstep, a failed post-PTQ preview now stops the run with a message that names --skip_generate, as the VL previews and the expert-parallel route do. A single process keeps the skip: NVBug 6752977's case is a single-GPU device_map="auto" load. New test next to test_post_quantize_export_survives_a_failed_sanity_generate, with an FSDP2 plan of two ranks: the run raises and the export does not run; it fails without the change.

[cjluo-nv]

Bot comment.

The caller's comment says text-only generation "may fail for encoder-decoder models like Nemotron-Parse". That failure is deterministic, so every rank hits it. Under synced_gpus it now aborts the run instead of falling back. If the failure happens before any collective, an all_reduce vote on success (as the image preview does) would keep the fallback. Is aborting intended here?

[shengliangxu]

Aborting is intended. Under synced_gpus a rank cannot skip a failure on its own: if the failure is on that rank alone (an out-of-memory error mid-decode, say), the others are still inside generate's collectives, and whatever collective this rank issues next, a vote included, pairs with theirs. A vote is safe only when no rank can still be inside generate, which a rank that caught an exception cannot tell.

The runs this reaches: --use_fsdp2 refuses Nemotron VL checkpoints (validate_fsdp2_supported), and Nemotron-Parse, the comment's example, is not an MoE, so no expert-parallel plan shards it. Its previews never set synced_gpus and go on without the text-only preview, as before. Only an expert-parallel Nemotron VL MoE (#2754) stops, and --skip_generate runs it without the previews.

[cjluo-nv]

Bot comment.

pinned is set for any plan with world_size > 1, including EP or dense="fsdp" plans, but the comments describe a plan that gives each rank the whole model. If #2754 passes an EP/FSDP plan through get_model, every rank gets device_map={"": cuda:i} and the --use_seq_device_map refusal, and that refusal's advice ("shard the model with --use_fsdp2") wouldn't fit. Should this check not plan.is_expert_parallel and not plan.is_fsdp_dense, or refuse such plans?

[shengliangxu]

Refused now, in 28b6aba: get_model raises a ValueError for an expert-parallel or FSDP2 plan before it loads anything, since it loads the whole model on every rank. hf_ptq hands those plans to load_pretrained directly (#2754's load_model), so the pin and the --use_seq_device_map advice only ever see a plan that gives every rank the whole model. New test test_a_plan_that_shards_the_model_is_refused, for both plans; it fails without the check.

[shengliangxu]

/ok to test 28b6aba

[shengliangxu]

/claude review

[claude]

[CRITICAL Algorithm] Still open at 28b6aba, and now reaffirmed by a CODEOWNERS reviewer — the post-PTQ preview swallows a failure while generate runs with synced_gpus=True.

I'm not re-litigating the argument here; it was made in round 2 and answered in your reply ("kept, by the maintainer's decision"). The maintainer has now decided the other way: cjluo-nv's review lists "Fix post_quantize in examples/hf_ptq/hf_ptq.py: re-raise the sanity-generate failure when lockstep is set, as the VL paths now do, and add a test" as the first required action. Per the repo's own rule for review threads — if a reviewer reaffirms after a pushback, that settles it — this one is settled, so it stays CRITICAL until the code changes or the maintainer withdraws the request.

The one-line recap of why: under synced_gpus, generate all-reduces this_peer_finished every decode step. A rank that warns here does not exit, so torchrun tears nothing down, and the peers sit in that all-reduce until the 2h group timeout. lockstep is already bound one line above:

        except Exception as e:
            if lockstep:
                # The peers are waiting in generate's lockstep all-reduce; warning here would
                # leave them there until the group timeout.
                raise
            warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

If keeping the completed calibration on a sharded run is the blocker, the all_reduce vote you already use for the image preview preserves both properties. Either way the test is a stand-in whose generate raises, parametrized with and without a sharding plan.

[shengliangxu]

Changed in 7772c0d: under lockstep, a failed post-PTQ preview now stops the run with a message that names --skip_generate, as the VL previews and the expert-parallel route do. A single process keeps the skip: NVBug 6752977's case is a single-GPU device_map="auto" load. New test next to test_post_quantize_export_survives_a_failed_sanity_generate, with an FSDP2 plan of two ranks: the run raises and the export does not run; it fails without the change.

[claude]

[IMPORTANT Compatibility] Under synced_gpus this re-raise turns the documented, expected text-only failure into an aborted PTQ run, for a failure every rank hits identically — so it buys no hang protection and costs the whole job.

The caller at example_utils.py:203-210 says what this None is for: "Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse)". Both run_nemotron_vl_preview call sites in hf_ptq.py (:804 pre-quantize, :885 post-quantize) pass allow_fallback=False, so on main that failure makes generated_ids_before_ptq None, post_quantize takes its if generated_ids_before_ptq is None: pass branch, both previews are skipped, and the quantization and export still run. With a sharded plan (--use_fsdp2 over >1 GPU) the raise now propagates out of pre_quantize and kills the run before any calibration — for exactly the model family the comment names.

Crucially this failure is symmetric: model.chat(...) not existing, or an encoder-decoder refusing text-only input, happens on every rank at the same point, before any decode collective. Nobody is left waiting, so there is no peer to protect.

The two asymmetric cases where the re-raise is right — AutoProcessor.from_pretrained failing on one node, or a data-dependent error mid-decode — are the ones worth keeping. The all_reduce vote you added for the image preview is the shape that separates them, as cjluo-nv also suggests:

    except Exception as e:
        if generation_config.get("synced_gpus"):
            # Every rank of a sharded model agrees: a failure they all meet (an encoder-decoder
            # that has no text-only path) skips the preview everywhere, one only this rank meets
            # raises, since its peers would wait for it in generate's lockstep all-reduce.
            if all_reduce("VL text-only preview", lambda: False, any):
                raise
            return None
        print(f"Text-only generation failed: {e}")
        return None

One caveat to weigh: a rank that fails mid-decode cannot rendezvous in a fresh all_gather_object while its peers are inside generate's all-reduce, so the vote only closes cleanly for pre-decode failures. If that makes the vote unattractive, the narrower alternative is to re-raise only once decoding has started — e.g. catch the "no text-only path" case (not hasattr(model, "chat") / processor-construction failure) before entering generate and return None from there, leaving the re-raise for everything inside generate. Either way, please make sure the Nemotron-Parse path still reaches export under --use_fsdp2, and pin it with a test.

[shengliangxu]

Kept, for the reasons in cjluo-nv's thread on the same lines (#3139 (comment)). The code that bears on this one:

• No --use_fsdp2 run reaches this path. load_model runs validate_fsdp2_supported on the config before the FSDP2 load, and it refuses Nemotron VL and other multimodal checkpoints (the is_nemotron_vl(config) or _is_multimodal_config(config) check in example_utils.py). A Nemotron-Parse checkpoint is refused under --use_fsdp2 before anything loads, as on main, so it never gets here with synced_gpus.
• Nemotron-Parse is not an MoE, so the expert-parallel route ([OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754) does not shard it either: its previews run without synced_gpus, and this except returns None as before.
• The suggested vote runs only on the ranks that failed. It closes when every rank fails before decoding, but when one rank's processor load fails and the others enter generate, that rank's all_reduce pairs with their first decode collective. Returning None before generate without a vote leaves the same gap: that rank goes on to the image preview's vote while the others are in generate.

[claude]

Claude review (round 3) — 1 CRITICAL (carried, now maintainer-reaffirmed), 1 IMPORTANT, 1 SUGGESTION

Scope: full review (bare /claude review). 7 files, +632/−88 at 28b6aba; re-reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files. CodeRabbit auto-review is off on this base branch (shengliangx/hfep-base-14-hf-ptq), so I covered a little more ground, still skipping style/typo territory. I was able to read the previous rounds' review bodies and your round-2 reply comment, but not the inline threads themselves.

Resolved since round 2

• SUGGESTION 3 (get_model and sharded plans) — fixed. get_model now refuses an EP or FSDP2 plan up front, before any mutation, with a message that points at load_pretrained; test_a_plan_that_shards_the_model_is_refused pins both plan shapes. That also makes pinned's plan.world_size > 1 guard exact rather than merely safe-by-accident, since a plan reaching the pin can now only be a replicating one. This is the right place for the refusal — the previous behaviour surfaced it from inside _sharded_load_kwargs, naming a keyword the caller never passed.
• SUGGESTION 2 (a device naming the wrong GPU) — declined, and I accept the reason. Comparing the index against local_rank() would break launchers that expose only one GPU per rank. Not re-filing.

Still open

1. CRITICAL — hf_ptq.py:882-883: the post-PTQ preview still swallows a failure while generate runs with synced_gpus=True. (inline)

The substance was argued in round 2 and you answered it ("kept, by the maintainer's decision"). The maintainer has since decided the other way — cjluo-nv's review lists re-raising when lockstep is set as the first required action. Per this repo's own rule for review threads, a reviewer reaffirming after a pushback settles it, so I'm carrying this as CRITICAL rather than re-arguing it. lockstep is already bound one line above; three lines, plus a stand-in whose generate raises, parametrized with and without a sharding plan.

2. IMPORTANT — vlm_utils.py:280-282: the run_text_only_generation re-raise aborts the run for a failure every rank hits. (inline)

This is the item cjluo-nv raised that I explicitly declined to file in round 2 — I had looked only at the hang window, not at what the None was for, and on re-reading the caller I think the maintainer is right. example_utils.py:203 says it: "may fail for encoder-decoder models like Nemotron-Parse". Both run_nemotron_vl_preview call sites pass allow_fallback=False, so on main that failure yields generated_ids_before_ptq = None, post_quantize skips both previews, and quantization and export still run. Under a sharded plan the raise now propagates out of pre_quantize and kills the job before calibration — for the exact model family the comment names. And the failure is symmetric (chat missing, or an encoder-decoder refusing text-only input, happens on every rank before any decode collective), so no peer is left waiting and the re-raise protects nothing. The inline comment has the all_reduce vote the maintainer suggested, plus the mid-decode caveat that limits it and a narrower alternative.

3. SUGGESTION — the PR body's backward-compatibility line is wrong for single-process runs.

Is this change backward compatible?: ✅ A single process previews and loads as before.

keep_the_experts_implementation is applied unconditionally in pre_quantize, post_quantize and run_nemotron_vl_preview — it is not gated on the plan — so on transformers 5.x a single-GPU post-PTQ preview now decodes on quantized grouped_mm experts instead of batched_mm. That is the fix working as intended (the old output came from unquantized experts, which made the preview misleading), but it is a visible single-process change: preview text for a grouped_mm MoE will differ from main's, and users who compared pre/post previews will see a different, and more honest, diff. Worth correcting in the body, and worth a one-sentence CHANGELOG.rst fix entry under Quantization — "the post-PTQ generation preview of a grouped_mm MoE now runs on the quantized experts" reads as a user-visible bug fix rather than an intra-cycle regression.

What held up this round

• The new refusal is correctly ordered: plan = distribution or DistributionPlan(), refuse, then compute pinned — a refusal leaves no state touched. routing = {} if distribution is None else {...} still keeps load_pretrained's own routing and its "process group is not set up" refusal intact without a caller plan, and all four load_pretrained sites in get_model (:767, :779, :807, :918) carry **routing, so every path records a plan for the previews to read.
• The pin composes correctly with the other device-map routes: use_seq_device_map and pinned raises; the DiffusionGemma branch is gated on not pinned so it cannot re-enter "sequential" behind the pin; the device_map = None routes (:680, :729, :735) still land on the rank's GPU through the existing model.to(device); and elif not pinned around infer_auto_device_map avoids printing a max_memory limit a one-device map would ignore.
• preview_generate_kwargs cannot raise on today's paths: setup_distributed_args (example_utils.py:75-89) initialises the group only under --use_fsdp2, so a non-FSDP2 run has size() == 1 and distribution_plan returns DistributionPlan() rather than refusing a planless model; under --use_fsdp2 the load records the plan (model_load_utils.py:916). The predicate world_size > 1 and (is_expert_parallel or is_fsdp_dense) is right — a replicated multi-rank plan runs no forward collective, and DistributionPlan(1, dense="fsdp") needs no lockstep. Both parametrized.
• keep_the_experts_implementation re-checked: shadowing the class method with an instance attribute means nullcontext is called without self, vars(module).get(...) distinguishes a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__ for a plain attribute, the finally restores on a failed generate, and nesting restores correctly. Walking model.modules() rather than the root is what makes Nemotron-VL's chat-through-the-language-model case work, and the wrapped parametrization pins it.
• The image-preview vote is ordered correctly (_load_sample_image first, then an unconditional all_reduce once synced_gpus is set, which is itself derived from the identical plan on every rank), and combine=all gives every-or-none.
• Tests sit in tests/examples/, not tests/unit; each new case pins a distinct behaviour rather than restating the implementation.

Risk

Moderate, and the shape has shifted slightly: the pinning half of the PR is now well-fenced and inert until [3/3] lands a caller, while both open findings are in the preview half and live today for --use_fsdp2 on more than one GPU. Finding 1 turns one rank's failure into a multi-hour silent wedge; finding 2 turns a symmetric, previously-tolerated failure into a lost run. Nothing under modelopt/ changed.

🤖 Generated with Claude Code

[shengliangxu]

Replies to the review (the inline threads have the details):

1. post_quantize under lockstep: changed in 7772c0d, so a failed preview of a sharded model now stops the run with a message that names --skip_generate; a single process keeps the skip. [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 (comment)
2. The text-only preview's failure under synced_gpus: aborting is intended, since a rank cannot tell whether the others are still inside generate's collectives. [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 (comment)
3. get_model with a plan that shards the model: refused now, in 28b6aba. [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 (comment)
4. The PR body no longer says a single process previews as before: on transformers 5.0 and later, a single process's preview after PTQ of a grouped_mm MoE now decodes on the quantized experts. No CHANGELOG entry: the preview is a printed sanity check, the exported checkpoint is the same either way, and the changelog takes fixes for critical or known bugs.

[shengliangxu]

/ok to test 7772c0d

[shengliangxu]

/claude review

[claude]

[CRITICAL Algorithm] The raise does not stop the run — the rank wedges in destroy_process_group instead, with nothing printed.

This is the fix for round 3's finding, and the direction is right, but on a sharded launch the raise never reaches the launcher. main (hf_ptq.py:1553-1596) is try: ... finally: cleanup_distributed(args) with no except that aborts, and cleanup_distributed → dist_utils.cleanup() → _destroy_group(exc). Per this repo's own docstrings:

• cleanup() (modelopt/torch/utils/distributed.py:969-976): "The barrier is skipped when unwinding from an error... That is not sufficient on its own — destroy_process_group below blocks for the same reason — so error paths must call abort before reaching this finally."
• abort() (:993-1000): "Both a barrier and destroy_process_group stall when peers are blocked in a collective this rank will never reach, and the traceback only prints once the enclosing finally returns — so the run looks hung rather than failed."

So in the asymmetric case this branch targets (one rank's decode OOMs on its own shard's sample): the peers sit in generate's _has_unfinished_sequences all-reduce, this rank sits in destroy_process_group, nothing exits non-zero, torchrun tears nothing down, and _PREVIEW_STOPS_THE_RUN is never printed because the finally never returns. setup_distributed_args sets timeout=timedelta(hours=2), so that is a two-hour silent wedge — observably worse than the warnings.warn it replaces, which at least printed before hanging in the export's collectives.

The same gap applies to the other two lockstep raises this PR adds (vlm_utils.py:202-203 and :281-282) and to pre_quantize's un-wrapped generate: every one of them unwinds into the same finally.

Fix: follow the convention the repo already uses for every other distributed entrypoint (examples/megatron_bridge/quantize.py:545-549, distill.py:1005-1009, and distributed_run itself at distributed.py:680-690) — abort at main, which covers all four sites at once rather than one branch:

# examples/hf_ptq/hf_ptq.py, in main()
    try:
        ...
    except BaseException:
        # A per-rank failure leaves the peers inside a collective this rank will never reach
        # (generate's lockstep all-reduce, calibration's a2a), so cleanup's barrier and
        # destroy_process_group would stall with the traceback still unprinted. Abort instead and
        # let the launcher terminate the peers. SystemExit is re-raised by abort().
        if args.dist_state.world_size > 1:
            dist_utils.abort()
        raise
    finally:
        cleanup_distributed(args)

abort() prints the active exception itself, so the _PREVIEW_STOPS_THE_RUN message and its __cause__ do reach the user. test_post_quantize_stops_the_run_when_a_sharded_models_preview_fails keeps passing unchanged (it exercises post_quantize directly); the new behaviour wants a main-level test that a rank-local failure under world_size > 1 aborts instead of returning through cleanup_distributed — patch dist_utils.abort and assert it was called, with cleanup asserting it was not reached.

[shengliangxu]

Fixed in a8b2bae, at main, as suggested: on a run of several ranks, an exception leaving the run goes through abort(), which prints it, the --skip_generate message and its cause included, and exits, so torchrun stops the other ranks instead of this one waiting in cleanup_distributed's teardown. A single process raises as before. This covers the other lockstep raises and pre_quantize's preview too, since they all unwind through main. New test test_a_rank_that_fails_aborts_a_run_of_several_ranks: with two ranks a failed load reaches abort() before the teardown, and a single process raises; the two-rank case fails without the change. In [3/3] #2754, main runs inside distributed_run, which aborts a failed rank the same way, so that PR drops this except and its test.

[claude]

Claude review (round 4) — 1 CRITICAL, 0 IMPORTANT, 0 SUGGESTION

Scope: full review (bare /claude review). 7 files, +666/−88 at 7772c0d; re-reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files. CodeRabbit auto-review is off on this base branch (shengliangx/hfep-base-14-hf-ptq), so I covered a little more ground, still skipping style/typo territory. I could read the previous rounds' review bodies and your reply comments, but not the inline threads themselves.

Resolved since round 3

• IMPORTANT (vlm_utils.py:281, the run_text_only_generation re-raise) — closed by your explanation, and I accept it. A rank cannot tell whether its peers are still inside generate's collectives, so the all_reduce vote cjluo-nv suggested cannot cover the mid-decode case; and the one model family the comment names (Nemotron-VL / Nemotron-Parse) is refused under --use_fsdp2 by validate_fsdp2_supported (example_utils.py:106-107), so the aborting path is not reachable until [3/3]'s EP plans. Not re-filing.
• SUGGESTION (the PR body's backward-compatibility line) — fixed. The body now states the single-process change: on transformers 5.0+ the post-PTQ preview of a grouped_mm MoE decodes on the quantized experts. Your reason for no CHANGELOG.rst entry (the preview is a printed sanity check; the exported checkpoint is unchanged either way) is consistent with the repo's "critical or known bugs" rule. Settled.

Still open — one new finding

CRITICAL — hf_ptq.py:886-890: the new raise does not stop the run; the rank wedges in destroy_process_group with the message unprinted. (inline)

The direction of 7772c0d is right — swallowing under lockstep was the round-1/2/3 CRITICAL, and cjluo-nv asked for the raise. But raising is only half of it on this entrypoint. main is try: ... finally: cleanup_distributed(args) with no except that aborts (:1553-1596), and the library you added in [1/3] says in three places that this is not enough:

cleanup(): "…error paths must call abort before reaching this finally."
abort(): "Both a barrier and destroy_process_group stall when peers are blocked in a collective this rank will never reach, and the traceback only prints once the enclosing finally returns — so the run looks hung rather than failed."

So in the asymmetric case this branch exists for — one rank's decode fails on its own shard's sample — the peers stay in generate's all-reduce, this rank stays in destroy_process_group, no process exits non-zero, torchrun tears nothing down, and _PREVIEW_STOPS_THE_RUN is never printed. With setup_distributed_args' timeout=timedelta(hours=2) that is a two-hour wedge with no output, where main's warnings.warn at least printed before hanging in the export. This is the same conclusion you reached yourself in the round-2 reply ("a rank that raised would also wait, in cleanup_distributed's destroy_process_group") — the raise landed, the abort it depends on did not.

It is not specific to this branch: the two vlm_utils.py lockstep raises (:202, :281) and pre_quantize's un-wrapped generate all unwind into the same finally. So the fix belongs at main, matching what every other distributed entrypoint in this repo already does (examples/megatron_bridge/quantize.py:545-549, distill.py:1005-1009, distributed_run at distributed.py:680-690): except BaseException: if args.dist_state.world_size > 1: dist_utils.abort(); raise. Four lines, one place, and it makes all four raise sites behave as their comments claim. The inline comment has the snippet and a test shape.

What held up this round

• The lockstep gate itself is right: preview_generate_kwargs' predicate (world_size > 1 and (is_expert_parallel or is_fsdp_dense)) excludes a replicated multi-rank plan (no forward collective) and DistributionPlan(1, dense="fsdp"), both parametrized. lockstep is computed once, before the try, so the branch cannot disagree with what generate was handed.
• post_quantize's single-process path keeps the NVBug 6752977 best-effort skip verbatim, and test_post_quantize_export_survives_a_failed_sanity_generate still pins it; the new test_post_quantize_stops_the_run_when_a_sharded_models_preview_fails asserts the __cause__ and export_calls == [], so it fails without the change.
• get_model's refusal of a sharding plan is correctly ordered — plan = distribution or DistributionPlan(), refuse, then compute pinned — so a refusal leaves no state touched, and it makes pinned's world_size > 1 guard exact rather than safe-by-accident. routing = {} if distribution is None else {...} keeps load_pretrained's own env routing and its "process group is not set up" refusal intact without a caller plan, and all four load_pretrained sites carry **routing.
• The pin composes with the other device-map routes: use_seq_device_map and pinned raises; the DiffusionGemma branch is gated on not pinned; device_map=None routes still land on the rank's GPU through model.to(device); elif not pinned around infer_auto_device_map avoids printing a max_memory limit a one-device map would ignore (and model_kwargs["max_memory"] is only set inside that branch, so the pinned load passes none).
• keep_the_experts_implementation re-checked once more: shadowing the class method with an instance attribute means nullcontext is called without self, vars(module).get(...) distinguishes a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__, the finally restores on a failed generate, and nesting restores correctly. Walking model.modules() is what makes Nemotron-VL's chat-through-the-language-model case work; the wrapped parametrization pins it and fails without the fix.
• The image-preview vote is ordered correctly (_load_sample_image first, then an unconditional all_reduce once synced_gpus is set — itself derived from the identical plan on every rank), combine=all gives every-or-none, and the all_reduce cannot hit its planless-launch refusal since synced_gpus implies a set-up group.
• Tests sit in tests/examples/, not tests/unit; the stand-in changes in test_hf_ptq_args.py / test_hf_ptq_model_type.py are the minimum the model.modules() walk requires, with assertions unchanged.

Risk

Moderate. The pinning half is well-fenced and inert until [3/3] lands a caller. The preview half is live today for --use_fsdp2 on more than one GPU, and the one open finding is that its failure handling still ends in a silent multi-hour hang rather than a fast, loud one — the same class of bug the PR exists to close, now one layer further out. Nothing under modelopt/ changed.

🤖 Generated with Claude Code

[shengliangxu]

/ok to test a8b2bae

[shengliangxu]

/claude review

[claude]

Claude review (round 5) — 0 CRITICAL, 0 IMPORTANT, 0 SUGGESTION

Scope: full review (bare /claude review). 7 files, +712/−89 at a8b2bae; re-reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files. CodeRabbit auto-review is off on this base branch (shengliangx/hfep-base-14-hf-ptq), so I covered a little more ground, still skipping style/typo territory. I read the previous rounds' review bodies and your reply comments, but not the inline threads themselves.

Resolved since round 4

CRITICAL (hf_ptq.py, the raise that did not stop the run) — fixed, and I traced the whole path. main now carries the handler the library's own docstrings ask for:

    except BaseException:
        if args.dist_state.world_size > 1:
            abort()
        raise

What I checked, rather than taking it on the comment:

• It is reached. Nothing between post_quantize's raise RuntimeError(_PREVIEW_STOPS_THE_RUN) and this handler swallows it: hf_ptq.py has only two other except blocks (:862 for print_quant_summary, main-rank-only and non-collective, and :886, which is the raise site itself), and quantize_main wraps neither preview. The same holds for pre_quantize's un-wrapped generate and the two vlm_utils.py lockstep raises (:202, :281) — all four now unwind into this one place, which is what made fixing it at main the right call.
• Nothing blocks on the way. The unwind crosses with mlflow_run(args) first. tracked_run gates everything on logger.enabled = bool(args.mlflow) and is_main, so a non-main failing rank's exit does only the is_main-gated drop_experiment_json (nothing), and _closing_run's finish("FAILED") on rank 0 is MLflow network I/O with no collective. The failing rank therefore reaches abort() instead of stalling in front of it.
• It aborts before the teardown, not after. abort() ends in os._exit(exit_code), which bypasses finally: cleanup_distributed(args) — so the barrier and destroy_process_group that round 4's finding was about are skipped, traceback.print_exc() plus the stdout/stderr flushes put _PREVIEW_STOPS_THE_RUN on the log before the exit, and torchrun sees a non-zero rank and tears the peers down. The two-hour silent wedge is gone.
• The guard is exact. args.dist_state is attached by setup_distributed_args, which is called before the try, so the handler cannot AttributeError over the original exception; and it sets world_size > 1 only in the --use_fsdp2 branch, after dist_utils.setup(), so abort() is never called on a launch whose group was not created. A single process keeps raise plus the clean teardown.
• An intentional exit still unwinds cleanly. abort() re-raises SystemExit rather than exiting, so sys.exit(0) from inside the try propagates and cleanup_distributed still runs.
• It matches what every other distributed entrypoint here does (examples/megatron_bridge/quantize.py:543-551, distill.py, prune_minitron.py, and distributed_run at distributed.py:672-691) — and is slightly better than the megatron_bridge copies, which swallow instead of re-raising on the single-process path.
• The test pins it. test_a_rank_that_fails_aborts_a_run_of_several_ranks parametrizes world_size 1/2, stands abort in as a raising BaseException subclass (the real one exits), and asserts the event order ["abort", "cleanup"] against ["cleanup"] — so it fails if the guard inverts or the handler is dropped. monkeypatch.setattr(hf_ptq, "abort", ...) bites because the call is a module-global lookup.

What held up this round

• post_quantize's split is right: lockstep is computed once before the try, so the branch cannot disagree with what generate was handed; the single-process path keeps the NVBug 6752977 best-effort skip verbatim (test_post_quantize_export_survives_a_failed_sanity_generate), and test_post_quantize_stops_the_run_when_a_sharded_models_preview_fails asserts both __cause__ == "boom" and export_calls == [].
• Nothing rank-dependent decides whether a preview is entered: pre_quantize's branches key off args.skip_generate, args.specdec_offline_dataset and model_type only, and post_quantize's generated_ids_before_ptq is None short-circuit is reachable only through those same rank-uniform gates (the VL path that could return None asymmetrically now raises under lockstep, and is refused under --use_fsdp2 by validate_fsdp2_supported). So the ranks enter generate the same number of times.
• preview_generate_kwargs' predicate (world_size > 1 and (is_expert_parallel or is_fsdp_dense)) still excludes a replicated multi-rank plan and DistributionPlan(1, dense="fsdp"); all five plan shapes are parametrized.
• The image-preview vote calls all_reduce(label, fn, combine) with the real signature (distributed.py:553), is ordered after _load_sample_image, is unconditional across ranks once synced_gpus is set (itself derived from the identical plan everywhere), and combine=all gives every-or-none. A failure inside it re-raises through the same synced_gpus branch and now aborts.
• keep_the_experts_implementation re-checked: shadowing the class method with an instance attribute means nullcontext is called without self, nn.Module.__setattr__ stores a plain class object in __dict__, vars(module).get(...) distinguishes a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__, the finally restores on a failed generate, and nesting restores correctly. Walking model.modules() is what makes Nemotron-VL's chat-through-the-language-model case work.
• get_model: the sharding-plan refusal is still correctly ordered (plan = distribution or DistributionPlan(), refuse, then compute pinned), routing = {} if distribution is None else {...} keeps load_pretrained's own env routing and its "process group is not set up" refusal, all four load_pretrained sites carry **routing, use_seq_device_map and pinned raises, the DiffusionGemma branch is gated on not pinned, the device_map=None routes still land on the rank's GPU through model.to(device), and elif not pinned keeps max_memory out of a pinned load.
• Tests sit in tests/examples/, not tests/unit; each new case pins a distinct behaviour. Added source is ~206 lines, well under the 500-line budget.

One note, not a finding

except BaseException aborts on every non-SystemExit error of a multi-rank run, including symmetric ones (a bad --export_path, say) where all ranks would have met the barrier and torn the group down cleanly. Telling the two apart would need a collective, which is exactly what hangs when the failure is asymmetric — so aborting both is the sound choice, it is what megatron_bridge's entrypoints already do unconditionally, and the cost is a skipped destroy_process_group on a process that is exiting anyway. Recorded only so the trade-off is visible; no change requested.

Risk

Low, and down from the previous rounds. The blast radius is confined to examples/hf_ptq, nothing under modelopt/ changed, the pinning half is well-fenced and inert until [3/3] lands a caller, and the preview half's failure paths now all terminate loudly instead of wedging. The documented behaviour changes — a single-process grouped_mm MoE post-PTQ preview decoding on the quantized experts, and a sharded run's failed post-PTQ preview stopping before the export — are both in the PR body, and the second is maintainer-requested.

🤖 Generated with Claude Code

[claude]

Claude review passed — no blocking issues found. LGTM

[cjluo-nv]

Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.

Every earlier concern, the round-4 abort included, is now fixed in the diff or settled by a reply I checked. A human still needs to sign off on the edits to existing tests.

Needs action:

• Sign off on the existing-test edits in tests/examples/hf_ptq/test_hf_ptq_args.py and test_hf_ptq_model_type.py. The SimpleNamespace stand-ins became nn.Modules so the previews can read the plan; the assertions are unchanged.

No action needed:

• ✔️ Resolved since the last review:
  • the post-PTQ preview no longer swallows a failure under lockstep;
  • main now calls abort() before teardown on a run of several ranks;
  • get_model refuses EP/FSDP plans;
  • get_model forwards a plan only when the caller gave one.
    The run_text_only_generation abort and the wrong-GPU guard were settled by the author's replies.
• About 206 added source lines, under the budget. The new test file carries the standard NVIDIA header.

[shengliangxu]

/ok to test ab91b2a

[shengliangxu]

/ok to test 7f987fb

[shengliangxu]

/ok to test 739b185