Skip to content

[OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shengliangxu wants to merge 6 commits into
base: shengliangx/hfep-base-14-hf-ptq
Choose a base branch
from
Open

[OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
20d7eee
Keep hf_ptq's previews of a sharded model in lockstep, on the quantiz…
shengliangxu Oct 30, 2026
3c5ffb9
Pin a whole-model load to the rank's own GPU by its distribution plan
shengliangxu Oct 30, 2026
6904778
Hand load_pretrained only a distribution plan get_model's caller gave
shengliangxu Oct 30, 2026
b9bd7c3
Refuse a plan that shards the model in get_model, which loads it whole
shengliangxu Oct 30, 2026
2d54de5
Stop the run when a sharded model's preview after quantization fails
shengliangxu Oct 30, 2026
739b185
Abort a failed rank of a multi-rank hf_ptq run instead of tearing the…
shengliangxu Oct 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
160 changes: 126 additions & 34 deletions examples/hf_ptq/example_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import os
import warnings
from collections.abc import Iterator
from contextlib import contextmanager
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
from datetime import timedelta
from pathlib import Path
Expand Down Expand Up @@ -56,6 +56,7 @@
)
from modelopt.torch.quantization.mode import get_modelike_from_algo_cfg
from modelopt.torch.utils import distributed as dist_utils
from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run
from modelopt.torch.utils.plugins.model_load_utils import is_fsdp2_text_only_vlm, load_pretrained

Expand Down Expand Up @@ -154,6 +155,48 @@ def validate_fsdp2_supported(args, config):
)


def preview_generate_kwargs(model) -> dict[str, Any]:
"""Extra ``generate`` kwargs that keep every rank of a sharded model decoding a preview.

Each rank previews the first sample of its own calibration shard, so the ranks finish at
different steps. A model its plan shards runs a collective in every forward (the routed experts'
all-to-all, FSDP2's all-gathers): a rank that stopped decoding would leave the others blocked in
it. ``synced_gpus`` keeps every rank stepping until all of them have finished.
"""
plan = distribution_plan(model)
sharded = plan.world_size > 1 and (plan.is_expert_parallel or plan.is_fsdp_dense)
return {"synced_gpus": True} if sharded else {}


@contextmanager
def keep_the_experts_implementation(model) -> Iterator[None]:
"""Keep ``generate`` from switching a MoE's experts to ``batched_mm`` while it decodes.

On a GPU, ``generate`` runs the experts of a ``grouped_mm`` MoE on ``batched_mm`` for the
decoding stage and switches them back (``_optimize_model_for_decode``, Transformers 5.0 and
later). Two things follow. ModelOpt's quantizers sit on the ``eager``, ``grouped_mm`` and
``torch_a2a`` forwards, which ``batched_mm`` bypasses, so the preview after PTQ would decode on
unquantized experts. And Transformers refuses the switch for an FSDP2-wrapped model, whose class
is a subclass built outside its modeling module: ``FSDPQwen3MoeForCausalLM does not support
setting experts implementation``. Every module of the model that has the hook is kept from it:
a model can generate through a submodule, as a Nemotron-VL model's ``chat`` generates through
its language model.
"""
modules = [m for m in model.modules() if hasattr(m, "_optimize_model_for_decode")]
# Put back what each module itself held, if anything: a nested use, or a hook bound per instance.
own = [(module, vars(module).get("_optimize_model_for_decode")) for module in modules]
for module in modules:
module._optimize_model_for_decode = nullcontext
try:
yield
finally:
for module, held in own:
if held is None:
del module._optimize_model_for_decode
else:
module._optimize_model_for_decode = held


def run_nemotron_vl_preview(
full_model,
tokenizer,
Expand All @@ -178,32 +221,37 @@ def run_nemotron_vl_preview(
"""
from vlm_utils import run_text_only_generation, run_vl_preview_generation

print(f"Running text-only preview generation for Nemotron VL model ({stage_name})...")
question = tokenizer.decode(input_ids[0], skip_special_tokens=True)
generation_config = {
"max_new_tokens": 100,
"do_sample": False,
"eos_token_id": tokenizer.eos_token_id,
}

# Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse)
text_response = run_text_only_generation(
full_model, tokenizer, question, generation_config, pyt_ckpt_path, trust_remote_code
)

generated_ids = None
if text_response is not None:
print(f"✅ Text-only generation successful: {text_response[:100]}...")
generated_ids = text_response
elif allow_fallback:
print("Text-only generation failed, falling back to standard generate...")
generated_ids = full_model.generate(input_ids, max_new_tokens=100)
with keep_the_experts_implementation(full_model):
print(f"Running text-only preview generation for Nemotron VL model ({stage_name})...")
question = tokenizer.decode(input_ids[0], skip_special_tokens=True)
lockstep = preview_generate_kwargs(full_model)
generation_config = {
"max_new_tokens": 100,
"do_sample": False,
"eos_token_id": tokenizer.eos_token_id,
**lockstep,
}

# Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse)
text_response = run_text_only_generation(
full_model, tokenizer, question, generation_config, pyt_ckpt_path, trust_remote_code
)

# Run additional VL test with images
print(f"Running additional VL test with images ({stage_name})...")
run_vl_preview_generation(full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code)
generated_ids = None
if text_response is not None:
print(f"✅ Text-only generation successful: {text_response[:100]}...")
generated_ids = text_response
elif allow_fallback:
print("Text-only generation failed, falling back to standard generate...")
generated_ids = full_model.generate(input_ids, max_new_tokens=100, **lockstep)

# Run additional VL test with images
print(f"Running additional VL test with images ({stage_name})...")
run_vl_preview_generation(
full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code, lockstep
)

return generated_ids
return generated_ids


def _is_multimodal_config(config):
Expand Down Expand Up @@ -604,8 +652,19 @@ def get_model(
offload_folder=None,
max_cpu_memory_gb=None,
max_gpu_memory_gb=None,
distribution: DistributionPlan | None = None,
):
print(f"Initializing model from {ckpt_path}")
# The plan to load the whole model by on every rank; a single process's by default. Only a plan
# the caller gave reaches load_pretrained: without one it decides itself as before, and refuses
# MODELOPT_* knobs asking to shard over a launch whose process group is not set up.
plan = distribution or DistributionPlan()

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

if plan.is_expert_parallel or plan.is_fsdp_dense:
raise ValueError(
f"{plan!r} shards the model over the ranks, but get_model loads the whole model on "
"every rank: load it with load_pretrained."
)
routing = {} if distribution is None else {"distribution": distribution}

_disk_offload = offload_folder is not None
if _disk_offload and max_cpu_memory_gb is None:
Expand All @@ -616,9 +675,25 @@ def get_model(
UserWarning,
)

device_map = "auto"
if device == "cpu":
device_map = "cpu"
# Under torchrun each rank asks for its own "cuda:<local_rank>". Do not fall through to "auto":
# accelerate runs its placement heuristic INDEPENDENTLY in every rank, and on a node whose GPUs
# are all idle every rank scores them identically and picks the SAME card. All ranks' weights,
# and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL collective
# binds the whole data-parallel group to that single device and aborts with "Duplicate GPU
# detected : rank 1 and rank 0 both on CUDA device ...". Every load path below honours the pin
# except those that load with no device map (Nemotron VL): each rank builds a whole copy in host
# memory and only then moves it to its own GPU, so the node needs one host copy per rank.
# A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under
# data-parallel calibration each rank must own its own GPU. A single process keeps "auto".
# hf_ptq hands in its run's torch.device, which never equals a string such as "cpu".
on_cpu = torch.device(device).type == "cpu"
if plan.world_size > 1 and not on_cpu and ":" not in str(device):
raise ValueError(
f"{plan!r} has each rank hold its own copy of the model on its own GPU, but "
f"device={device!r} names no GPU index; pass the rank's own GPU, cuda:<local_rank>."
)
pinned = {"": str(device)} if not on_cpu and plan.world_size > 1 else None
Comment on lines +689 to +695

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@cjluo-nv cjluo-nv Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

device_map = "cpu" if on_cpu else pinned or "auto"
Comment on lines +689 to +696

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


# Prepare config kwargs for loading
config_kwargs = {"trust_remote_code": trust_remote_code} if trust_remote_code else {}
Expand Down Expand Up @@ -654,13 +729,24 @@ def get_model(
# modules together for checkpoints that fit; larger ones can still spill and split a
# tied pair, and need an explicit single-device map. Multi-GPU only: a single-GPU split
# cannot separate a tied pair, and sequential would needlessly cap max_memory there.
if device != "cpu" and torch.cuda.device_count() > 1 and is_diffusion_gemma(hf_config):
if (
not pinned
and not on_cpu
and torch.cuda.device_count() > 1
and is_diffusion_gemma(hf_config)
):
print(
"Detected DiffusionGemma model. Using device_map='sequential'; the balanced "
"'auto' mapping can split its tied encoder/decoder weights across GPUs."
)
use_seq_device_map = True

if use_seq_device_map and pinned:
raise ValueError(
f"--use_seq_device_map spreads one model over every GPU, but each rank of this run "
f"holds its own copy on its own GPU ({device}). Drop it, or shard the model with "
"--use_fsdp2."
)
if use_seq_device_map:
device_map = "sequential"
# If we use sequential, set max_memory limit to ensure that the model does not occupy the full GPU
Expand Down Expand Up @@ -712,6 +798,7 @@ def has_pack_quantized_config(config):
ckpt_path,
AutoModelForCausalLM,
model_type=hf_config.model_type,
**routing,
device_map=device_map,
**model_kwargs,
)
Expand All @@ -723,7 +810,8 @@ def has_pack_quantized_config(config):
ckpt_path,
AutoModelForCausalLM,
model_type=hf_config.model_type,
device_map="auto",
**routing,
device_map=pinned or "auto",
trust_remote_code=trust_remote_code,
dtype="auto",
**experts_kwargs,
Expand All @@ -750,7 +838,8 @@ def has_pack_quantized_config(config):
ckpt_path,
AutoModelForCausalLM,
model_type=hf_config.model_type,
device_map="cpu" if device == "cpu" else "sequential",
**routing,
device_map="cpu" if on_cpu else pinned or "sequential",
**model_kwargs,
)
else:
Expand Down Expand Up @@ -836,7 +925,9 @@ def has_pack_quantized_config(config):
f"Offload folder: {offload_folder}\n"
"Weights exceeding GPU+CPU budgets will be streamed from disk."
)
else:
elif not pinned:
# A pinned copy goes onto its rank's GPU whole, by a map that ignores max_memory: sizing
# it over the node's GPUs would print a limit the load does not apply.
inferred_device_map = infer_auto_device_map(model, max_memory=max_memory)
if "cpu" in inferred_device_map.values():
for _device in max_memory:
Expand All @@ -858,6 +949,7 @@ def has_pack_quantized_config(config):
ckpt_path,
auto_model_module,
model_type=hf_config.model_type,
**routing,
device_map=device_map,
**model_kwargs2,
)
Expand All @@ -866,11 +958,11 @@ def has_pack_quantized_config(config):
_unpack_compressed_linear_weights(model, ckpt_path)

# If device_map was disabled (None), manually move model to target device
if device_map is None and device != "cpu":
if device_map is None and not on_cpu:
print(f"Moving model to {device} device...")
model = model.to(device)

if device == "cuda" and not is_model_on_gpu(model):
if torch.device(device).type == "cuda" and not is_model_on_gpu(model):
print("Warning: Some parameters are not on a GPU. Calibration can be slow or hit OOM")

return model
Expand Down
45 changes: 34 additions & 11 deletions examples/hf_ptq/hf_ptq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,9 @@
get_tokenizer,
is_nemotron_vl,
is_trtllm_enc_dec_export,
keep_the_experts_implementation,
mlflow_run,
preview_generate_kwargs,
recipe_layerwise_blocks,
run_nemotron_vl_preview,
save_source_config,
Expand Down Expand Up @@ -99,7 +101,7 @@
get_max_batch_size,
get_supported_datasets,
)
from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
from modelopt.torch.utils.distributed import DistributionPlan, abort, set_distribution_plan
from modelopt.torch.utils.memory_monitor import launch_memory_monitor
from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args
from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
Expand All @@ -116,6 +118,10 @@
RAND_SEED = 1234
# Most tokens per calibration batch when layerwise calibration picks the batch size (64 x 512).
LAYERWISE_CALIB_TOKENS = 32768
_PREVIEW_STOPS_THE_RUN = (
"The preview after quantization failed, which stops the run before the export, so no "
"checkpoint is written; --skip_generate exports without the preview."
)


def _layerwise_calib_batch_size(calib_seq: int, num_samples: int) -> int:
Expand Down Expand Up @@ -809,11 +815,13 @@ def pre_quantize(
trust_remote_code=args.trust_remote_code,
)
else:
generated_ids_before_ptq = full_model.generate(
preview_input_ids,
attention_mask=preview_attention_mask,
max_new_tokens=100,
)
with keep_the_experts_implementation(full_model):
generated_ids_before_ptq = full_model.generate(
preview_input_ids,
attention_mask=preview_attention_mask,
max_new_tokens=100,
**preview_generate_kwargs(full_model),
)

return preview_input_ids, preview_attention_mask, generated_ids_before_ptq

Expand Down Expand Up @@ -866,13 +874,20 @@ def post_quantize(
# part of the model to CPU (seen on unified-memory single-GPU hosts) can make a
# quantized layer run on CPU, which some kernels (e.g. NVFP4 dynamic block
# quantization) don't support. Don't let that discard the completed calibration.
lockstep = preview_generate_kwargs(full_model)
try:
generated_ids_after_ptq = full_model.generate(
preview_input_ids,
attention_mask=preview_attention_mask,
max_new_tokens=100,
)
with keep_the_experts_implementation(full_model):
generated_ids_after_ptq = full_model.generate(
preview_input_ids,
attention_mask=preview_attention_mask,
max_new_tokens=100,
**lockstep,
)
except Exception as e:

@cjluo-nv cjluo-nv Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

if lockstep:
# The ranks decode in lockstep: a rank that skipped the preview would go on to the
# export's collectives while the others still wait in generate's, so the run stops.
raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e
Comment on lines 886 to +890

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
Comment on lines 886 to 891

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines 886 to 891

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

elif is_nemotron_vl_model and tokenizer is not None:
generated_ids_after_ptq = run_nemotron_vl_preview(
Expand Down Expand Up @@ -1577,6 +1592,14 @@ def main(args: argparse.Namespace):
tokenizer,
device,
)
except BaseException:
# A rank that fails alone, such as in a sharded model's lockstep preview, leaves the others
# in a collective it never reaches, and cleanup's barrier and destroy_process_group would
# wait on them with the traceback unprinted. abort() prints it and exits, so the launcher
# stops the other ranks.
if args.dist_state.world_size > 1:
abort()
raise
finally:
cleanup_distributed(args)

Expand Down
Loading
Loading