FSDP2 calibration with hf_ptq.py [1/2] (#1563)

### What does this PR do?

Type of change: New feature (+ refactor/removal of the legacy multi-node
path) <!-- Use one of the following: Bug fix, new feature, new example,
new tests, documentation. -->

<!-- Details about the change. -->
Consolidates multi-node FSDP2 post-training quantization into the
standard hf_ptq.py entry point behind a --use_fsdp2 flag, and removes
the separate
Accelerate-based multinode_ptq.py script and fsdp2.yaml config. FSDP2
PTQ is now launched with torchrun and supports single-node multi-GPU and
multi-node out
  of the box.

  Highlights:
- --use_fsdp2 / --cpu_offload on hf_ptq.py — calibration runs under
PyTorch FSDP2; decoder layers are sharded (root stays replicated), with
optional CPU
  offload of decoder shards between forwards.
- New modelopt/torch/utils/model_load_utils.py — parallel, round-robin
safetensors loading: each rank reads only the decoder layers it owns
from disk, then
broadcasts them so every rank can shard its slice. Non-decoder weights
(embed/lm_head/norm) are read on rank 0 and broadcast.
- New FSDP2 helpers in modelopt/torch/utils/distributed.py — fsdp2_wrap,
shard_dataloader, fsdp_aware_forward_loop, broadcast_state_dict.
- Distributed export (unified_export_hf.py) — replaces the Accelerate
get_state_dict gather with get_model_state_dict(full_state_dict,
cpu_offload,
broadcast_from_rank0); rank 0 writes files, other ranks sync on a
barrier.
- core_utils.py — relaxes the stale "root must be an FSDPModule" assert
(only decoder layers are wrapped), and adds a CPU↔GPU mirror so weight
access/writeback
  works for CPU-offloaded shards.
- Docs — rewritten examples/llm_ptq/README.md FSDP2 section and the
SLURM PTQ reference, both using torchrun --use_fsdp2.

v1 scope: standard causal-LM checkpoints only. --use_fsdp2 raises
NotImplementedError for VILA, multimodal/VL,
pack-quantized/compressed-tensors, speculative decoding, MTP,
auto-quantize, and sparsity.

Design Notes
  
1. Why a custom parallel-safetensors loader instead of HF device_map /
accelerate

AutoModelForCausalLM.from_pretrained(device_map="auto" | "cpu") and
accelerate.load_checkpoint_in_model both load the full checkpoint on
every rank from disk (per-rank CPU peak ≈ full model size). For 70B
that's ~140 GiB/rank; for 200B+ it OOMs the node before sharding can
run. parallel_load_and_prepare_fsdp2 round-robins decoder layers across
ranks so each rank reads only ~model_size / world_size from disk, then
per-layer broadcasts to peers. Per-rank CPU peak is bounded by the
largest single layer + transient broadcast, not the full model. This is
what makes 200B+ FSDP2 PTQ feasible on commodity nodes; HF/accelerate's
existing entry points don't expose this composition.

2. Why a custom FSDP2 stack instead of keeping multinode_ptq.py +
accelerate launch

The deleted multinode_ptq.py ran on accelerate launch --config_file
fsdp2.yaml and duplicated hf_ptq.py's load → calibrate → export path.
Two consequences:
- Two divergent scripts for the same operation: CLI surface, calibration
loop, and export rewrites had to land twice. They had already drifted.
  - Users had to know which script applied at which scale.

The new code path unifies under hf_ptq.py --use_fsdp2. Going direct to
FSDP2 primitives (fully_shard, CPUOffloadPolicy, MixedPrecisionPolicy)
instead of
  routing through accelerate's wrappers buys:
- Direct control over mp_policy and offload_policy (accelerate's plugin
layer hides them).
- The parallel-read loader above (incompatible with accelerate's
per-rank full load).
  - No YAML config file in the example dir.
  - 
The "custom stack" is intentionally thin: fsdp2_wrap is a 5-line wrapper
over fully_shard; the rest is pure torch.distributed. We're not
reimplementing FSDP2 —just composing the public PyTorch surface directly
rather than via accelerate's adapter.

3. fsdp_aware_forward_loop ↔ transformers_trainer.py:_quantize_model
duplication
Both implement the same trick: mtq.quantize unwraps the FSDP module
before calling the user's forward_loop, and forwarding through the
unwrapped module bypasses FSDP2's pre/post-forward hooks (no all-gather
→ broken calibration). Both capture the outer wrapped model and forward
through it instead. This PR extracts the pattern into
fsdp_aware_forward_loop (in distributed.py) as the canonical helper. The
QLoRA path (transformers_trainer.py:_quantize_model) keeps its inlined
version this PR because the QLoRA forward loop has training-specific
quirks (batch shape, loss accumulation, gradient flow) that need a
careful pass to share the helper cleanly. The TODO in the helper's
docstring marks the consolidation target.

### Usage

```python
 # Single node, multiple GPUs
  torchrun --standalone --nproc_per_node=<num_gpus> hf_ptq.py \
      --pyt_ckpt_path <model> \
      --qformat nvfp4 \
      --export_path <out> \
      --use_fsdp2

  # Multi-node (run on each node)
  torchrun \
      --nnodes=<N> --node_rank=<rank> \
      --master_addr=<node0_ip> --master_port=<port> \
      --nproc_per_node=<gpus_per_node> \
      hf_ptq.py \
      --pyt_ckpt_path <model> --qformat nvfp4 \
      --export_path <out> --use_fsdp2 --cpu_offload   # --cpu_offload for very large models
```

### Testing
<!-- Mention how have you tested your change if applicable. -->
- tests/gpu/torch/quantization/test_fsdp2.py:
test_writeback_root_unwrapped (assert relaxation, root unwrapped) and
test_writeback_cpu_offload (CPU↔GPU mirror
  round-trip under CPUOffloadPolicy).
- tests/unit/torch/utils/test_model_load_utils.py: pure-function tests
for weight_map_for (sharded / single-file / missing) and
read_safetensors_subset.
- Manual end-to-end PTQ runs on single- and multi-node torchrun (FP8 /
NVFP4 / NVFP4 layerwise), with and without --cpu_offload.

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ❌ <!--- If ❌, explain why. -->
hf_ptq.py is fully backward compatible (FSDP2 is opt-in via a new flag),
but the legacy examples/llm_ptq/multinode_ptq.py script and fsdp2.yaml
are removed. Users of the old multi-node entry point must migrate to
hf_ptq.py --use_fsdp2
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A <!---
Mandatory -->
- Did you write any new necessary tests?: ✅ <!--- Mandatory for new
features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌<!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->
- Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run
`/claude review`. NVIDIA org members can self-trigger for complex
changes; orthogonal to CodeRabbit. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* FSDP2-enabled PTQ: torchrun single-/multi-node execution, CPU offload,
and NVFP4 layerwise calibration; distributed model loading and export
support.

* **Documentation**
* Rewritten PTQ guide and SLURM notes with explicit torchrun/FSDP2
instructions.

* **Removed**
  * Legacy Accelerate-based multinode PTQ workflow and YAML config.

* **Tests**
* Added FSDP2-focused tests for quantization writeback and safetensors
distributed loading.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: sugunav14 <178320438+sugunav14@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com>

Fix prequant layernorm export without scales (#1838)

- Skip pre-quant LayerNorm fusion when the representative input
quantizer has no `_pre_quant_scale` buffer.
- Add CPU unit coverage for the no-scale weight-only path and the
existing fusion/removal behavior when pre-quant scales are present.
- Compose the public `general/ptq/int4_blockwise_weight_only` recipe
from the shared recipe units so it stays aligned with
`INT4_BLOCKWISE_WEIGHT_ONLY_CFG`.

NVBug: https://nvbugspro.nvidia.com/bug/6311597

NVBug 6311597 reports a deterministic HF export crash for
`general/ptq/int4_blockwise_weight_only` on Llama-3.1-8B. That
weight-only recipe disables input quantization and skips calibration, so
`_pre_quant_scale` is not registered, but export still reaches
`fuse_prequant_layernorm` through the AWQ-like format path.

This PR also carries the recipe-sync diff that was previously in draft
PR #1836; PR #1836 has been closed after moving that diff here.

- `pre-commit run --files modelopt/torch/export/quant_utils.py
tests/unit/torch/export/test_unified_export_hf.py
modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml
tests/unit/recipe/test_loader.py`
- `pytest_pwd
tests/unit/torch/export/test_unified_export_hf.py::test_fuse_prequant_layernorm_skips_modules_without_pre_quant_scale
tests/unit/recipe/test_loader.py::test_load_recipe_all_builtins
tests/unit/recipe/test_loader.py::test_general_ptq_yaml_matches_config_dicts
tests/unit/recipe/test_presets.py -q` -> 29 passed

Full GB10/Llama repro was not run locally.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

* **Bug Fixes**
* Improved handling of layer norm fusion so it safely skips cases where
pre-quantization scale data is unavailable, avoiding unexpected errors.
* Updated the layer norm fusion flow to correctly apply scaling when
pre-quantization data is present.

* **Tests**
* Expanded coverage for export and recipe loading scenarios, including
cases with and without pre-quantization scale data.
  * Added validation for the new int4 blockwise weight-only recipe.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: realAsma <akuriparambi@nvidia.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Fix prequant layernorm export without scales (#1838)

- Skip pre-quant LayerNorm fusion when the representative input
quantizer has no `_pre_quant_scale` buffer.
- Add CPU unit coverage for the no-scale weight-only path and the
existing fusion/removal behavior when pre-quant scales are present.
- Compose the public `general/ptq/int4_blockwise_weight_only` recipe
from the shared recipe units so it stays aligned with
`INT4_BLOCKWISE_WEIGHT_ONLY_CFG`.

NVBug: https://nvbugspro.nvidia.com/bug/6311597

NVBug 6311597 reports a deterministic HF export crash for
`general/ptq/int4_blockwise_weight_only` on Llama-3.1-8B. That
weight-only recipe disables input quantization and skips calibration, so
`_pre_quant_scale` is not registered, but export still reaches
`fuse_prequant_layernorm` through the AWQ-like format path.

This PR also carries the recipe-sync diff that was previously in draft
PR #1836; PR #1836 has been closed after moving that diff here.

- `pre-commit run --files modelopt/torch/export/quant_utils.py
tests/unit/torch/export/test_unified_export_hf.py
modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml
tests/unit/recipe/test_loader.py`
- `pytest_pwd
tests/unit/torch/export/test_unified_export_hf.py::test_fuse_prequant_layernorm_skips_modules_without_pre_quant_scale
tests/unit/recipe/test_loader.py::test_load_recipe_all_builtins
tests/unit/recipe/test_loader.py::test_general_ptq_yaml_matches_config_dicts
tests/unit/recipe/test_presets.py -q` -> 29 passed

Full GB10/Llama repro was not run locally.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

* **Bug Fixes**
* Improved handling of layer norm fusion so it safely skips cases where
pre-quantization scale data is unavailable, avoiding unexpected errors.
* Updated the layer norm fusion flow to correctly apply scaling when
pre-quantization data is present.

* **Tests**
* Expanded coverage for export and recipe loading scenarios, including
cases with and without pre-quantization scale data.
  * Added validation for the new int4 blockwise weight-only recipe.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: realAsma <akuriparambi@nvidia.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Back-port non-gated fused-MoE wrapping for NemotronH (#1756 subset)

The unit test added by #1877
(test_nemotron_h_experts_only_recipes_target_routed_experts) fails on
release/0.45.0 because the non-gated fused-MoE-experts support from #1756
was never cherry-picked: release's fused-experts detector keys on
gate_up_proj (gated only), so NemotronH's NemotronHExperts (single up_proj
+ down_proj) are never wrapped and no routed-expert quantizers exist, so
`assert routed_expert_quantizers` gets {}.

Back-port only #1756's huggingface.py detection/wrapping, which is
decoupled from the feature stack (#1578/#1605/#1659/#1707/#1381) that
blocked the full #1756:

- Generalize _QuantFusedExperts with _first_proj_attr / _is_gated hooks
  (default gate_up_proj/gated, so the gated path is behavior-identical).
- Add _QuantNonGatedFusedExperts (up_proj, non-gated) for NemotronH.
- Detect both layouts via _fused_experts_wrapper_class and register the
  matching wrapper in register_fused_experts_on_the_fly.

Recipe matching already works: _normalize_fused_experts_quantizer_name is
layout-agnostic, so `*.experts.*weight_quantizer` matches up_proj_*
quantizers too. NemotronH export is out of scope (not exercised by the
test); this covers detection/wrapping/calibration only.

Rename the expert-index test reference to _get_expert_idx_from_first_proj.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

chore: stop tracking .claude/scheduled_tasks.lock (#1758)

## What
Remove `.claude/scheduled_tasks.lock` from version control and add a
`.gitignore` rule so it is never committed again.

## Why
This file is an **ephemeral Claude Code scheduler lock** — its contents
are runtime process state (`sessionId`, `pid`, `procStart`,
`acquiredAt`), not source. It was accidentally committed in #1623 and is
currently tracked on `main`.

Reported by @sychen52 in [review of
#1623](#1623 (review)).

## Changes
- `git rm --cached .claude/scheduled_tasks.lock`
- Add `.claude/scheduled_tasks.lock` to `.gitignore`

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Chores**
* Updated repository configuration to exclude internal runtime lock
files from version control.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Ye Yu <yeyu@nvidia.com>

Adds AutoQuant support for VLM / Qwen3.5-Qwen3.6 style models (#1381)

### What does this PR do?

Type of change: new feature, bug fix, new tests

### Details

- Enables AutoQuant search over fused MoE expert containers by
snapshotting/restoring their per-expert quantizers.
- Adds Qwen3.5/3.6 linear-attention grouping rules so fused deployment
layers keep compatible quant formats.
- Supports `w4a16_nvfp4` as an AutoQuant search format.
- Preserves disabled AutoQuant layer patterns in generated configs while
allowing selected modules like `lm_head` to override default disables.
- Keeps recipe-mode and AutoQuantize VLM paths on the outer CausalLM so
Qwen3.5/3.6-MoE `lm_head` remains visible.
- Skips `parent_class`-scoped quant config entries during AutoQuant bare
quantizer matching, preventing class-scoped global entries from
last-match overriding every selected module.
- Adds temporary hardcoded Qwen/VLM AutoQuant disabled-layer patterns in
`hf_ptq.py` with a TODO to refactor into the config system.

### Usage

```bash
python examples/llm_ptq/hf_ptq.py \
  --pyt_ckpt_path <model_path> \
  --qformat fp8,w4a16_nvfp4 \
  --auto_quantize_bits 5.0 \
  --auto_quantize_cost_model active_moe \
  --auto_quantize_checkpoint <autoquant_state.pt> \
  --export_path <output_dir>
```

### Testing

- `/Users/weimingc/miniconda3/envs/modelopt/bin/python -m pytest
tests/unit/torch/quantization/test_autoquant.py::test_get_auto_quantize_config_keeps_selected_lm_head_enabled
tests/unit/torch/quantization/test_config_validation.py::TestMatchQuantizerCfg::test_parent_class_scoped_entries_are_ignored_for_bare_autoquant_lookup`
- `/Users/weimingc/miniconda3/envs/modelopt/bin/python -m pytest
tests/unit/torch/quantization/test_autoquant.py
tests/unit/torch/quantization/test_config_validation.py -k "not
data_parallel"` (`120 passed, 1 deselected`)
- `/Users/weimingc/miniconda3/envs/modelopt/bin/python -m py_compile
examples/llm_ptq/hf_ptq.py modelopt/torch/quantization/algorithms.py
modelopt/torch/quantization/_auto_quantize_cost.py
tests/unit/torch/quantization/test_autoquant.py
tests/unit/torch/quantization/test_config_validation.py`
- Full local affected-file pytest without `-k "not data_parallel"` only
failed `test_data_parallel_auto_quantize` because this local sandbox
cannot bind a free socket (`PermissionError: Operation not permitted`).
- Ran Qwen3.6 35B AutoQuant e2e with `fp8,w4a16_nvfp4` and exported a
checkpoint.
- Verified exported checkpoint loads in vLLM nightly without local
patches.

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅
- 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A

### Additional Information

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added w4a16_nvfp4 quantization format and optional cost-exclusion
patterns for AutoQuantize.

* **Improvements**
* Safer multimodal/VLM handling and AutoQuantize now runs on the full
outer model when applicable.
* Better fused-MoE support, more accurate weight accounting, and refined
attention-grouping for improved quantization choices.
  * Dynamic layer-disabling support for targeted disables.

* **Tests**
* New unit tests covering cost-model exclusions, fused-MoE accounting,
and config selection.

* **Documentation**
  * Updated cost-constraint example to show exclusion-pattern usage.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

[OMNIML-4788] specdec_bench/Qwen3.5-4B: throughput_32k benchmark + S3…

… upload step (#1564)

### What does this PR do?

Type of change: enhancement (follow-up to
[#1531](#1531)).

Extends the merged Qwen3.5-4B SPEED-Bench launcher YAMLs from a
single-task qualitative-only smoke into a **3-task pipeline** that also
covers long-context throughput and verifies the S3-upload path
end-to-end. Two commits, cleanly cherry-picked from #1531's late branch
state — they were authored after the merge-commit was resolved against
an earlier rebased head and so didn't ride along with that merge.

### Pipeline shape (both YAMLs)

| Task | Split | Save dir |
|---|---|---|
| `task_0` | qualitative (existing quality / acceptance-rate signal) |
`/scratchspace/specdec_bench{,_mtp}/qualitative` |
| `task_1` | **throughput_32k** (new — long-context throughput) |
`/scratchspace/specdec_bench{,_mtp}/throughput_32k` |
| `task_2` | **upload to S3 in sweep layout** |
`s3://team-specdec-workgroup/results/specdec_bench{,_mtp}/<split>/` |

### New artifacts

* `tools/launcher/common/specdec_bench/upload_to_s3.sh` — thin wrapper
around `examples/specdec_bench/upload_to_s3.py` so it can be invoked as
a launcher task. Installs `boto3` from `requirements.txt` on cold
containers; warm pipelines pick it up from the prior `run.sh`.
*
`tools/launcher/common/specdec_bench/runtime_params_throughput_32k.yaml`
— pins `engine_args.max_model_len = 40,960` (32K input + 4K output + 4K
headroom) so vLLM doesn't silently auto-cap `max_model_len` below the
36K minimum needed for `throughput_32k` prompts on single-GPU runs.

### Why max_model_len matters

Without an explicit `max_model_len`, vLLM auto-derives it from the model
config (Qwen3.5-4B = 128K) **and from the GPU-memory budget**. On a
single GPU the second factor can cap effective `max_model_len` well
below 36K, silently truncating 32K-token prompts and producing wrong
throughput numbers. The qualitative split is not affected (its prompts
top out around 8K, well under any auto-derivation floor) so only
`task_1` carries the override.

### S3 credentials

`upload_to_s3.sh` reads `S3_ENDPOINT` / `S3_KEY_ID` / `S3_SECRET` from
the runtime environment (not hardcoded). `--skip-existing` +
`--allow-incomplete-provenance` are passed by default so re-runs land
alongside the prior upload, and runs lacking `CONTAINER_IMAGE` (Phase-2
harness work in OMNIML-4788 will populate it) still upload.

### Testing

Cluster smoke on cw_dfw via:

```
uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench.yaml --yes
```

is currently in-flight (jobs `12257378/79/80`, PD). Will update this PR
with timing/AR numbers + S3 upload confirmation once it lands.

### Before your PR is "Ready for review"

- Backward compatible: ✅ (additive — task_0 keeps the prior qualitative
behavior, just with `/qualitative` suffix in `save_dir`)
- New PIP dep: ✅ no (boto3 already in
`examples/specdec_bench/requirements.txt` from #1531)
- New tests: N/A (launcher YAML + shell wrapper; covered by cluster
smoke)
- Changelog: N/A (internal-facing tooling)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added a 32K-context runtime configuration (higher max model length) to
enable long-context throughput benchmarking and avoid silent prompt
truncation.
* Added a launcher helper to upload benchmark results to S3 with
incremental/retry-friendly options and pass/fail reporting.

* **Chores**
* Split Qwen3.5-4B benchmark into separate qualitative and 32K
throughput tasks and added coordinated S3 upload.
* Applied the same multi-task pipeline layout and clearer output
organization to the MTP speculative-decoding benchmark.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1564?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: chenhany <chenhany@nvidia.com>
Signed-off-by: Chenhan Yu <chenhany@nvidia.com>

fix(te-plugin): handle TE 2.15+ tuple return from `_Linear` / `_Group…

…edLinear`

TE 2.15+ changed `_Linear.forward` and `_GroupedLinear.forward` to return
`(out, new_workspace)` tuples instead of a single tensor. ModelOpt's
patched `te_quantized_linear_fn` / `te_grouped_quantized_linear_fn` still
passed the whole tuple into `self.output_quantizer`, crashing inside
`TensorQuantizer.forward` on `tuple.numel()`:

  AttributeError: 'tuple' object has no attribute 'numel'

Mirror the existing pattern from `_QuantTELayerNormLinear.forward`:
quantize only `output[0]` (activation) and pass auxiliary workspace
metadata through verbatim. TE <= 2.14 returns a single tensor and falls
through the isinstance branch unchanged.

This unblocks Megatron-Bridge's TE 2.15 path; the local
`patch_modelopt_te_linear_tuple_output` shim can be removed once this
ships in a tagged release.

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

fix(te-plugin): handle TE 2.15+ tuple return from `_Linear` / `_Group…

…edLinear`

TE 2.15+ changed `_Linear.forward` and `_GroupedLinear.forward` to return
`(out, new_workspace)` tuples instead of a single tensor. ModelOpt's
patched `te_quantized_linear_fn` / `te_grouped_quantized_linear_fn` still
passed the whole tuple into `self.output_quantizer`, crashing inside
`TensorQuantizer.forward` on `tuple.numel()`:

  AttributeError: 'tuple' object has no attribute 'numel'

Mirror the existing pattern from `_QuantTELayerNormLinear.forward`:
quantize only `output[0]` (activation) and pass auxiliary workspace
metadata through verbatim. TE <= 2.14 returns a single tensor and falls
through the isinstance branch unchanged.

This unblocks Megatron-Bridge's TE 2.15 path; the local
`patch_modelopt_te_linear_tuple_output` shim can be removed once this
ships in a tagged release.

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

fix(te-plugin): make _Linear arg indexing robust to TE signature chan…

…ges (#1473)

### What does this PR do?

Type of change: Bug fix

ModelOpt's `te_quantized_linear_fn` and `te_grouped_quantized_linear_fn`
read `weight` / `inp` from hard-coded positions in `args`. Two TE
signature changes broke this scheme:

- **TE 1.x → 2.0:** dropped the legacy `weight_fp8` slot between
`weight` and `inp`. ModelOpt handled this with an `if Version("2.0") <=
_TE_VERSION:` branch + a duplicate else branch.
- **TE 2.14 → 2.15:** inserted `weight_workspace` between `weight` and
`inp` at the `_Linear.forward` call site ([TE 2.15 linear.py
L1663](https://github.com/NVIDIA/TransformerEngine/blob/release_v2.15/transformer_engine/pytorch/module/linear.py#L1663)).
Unhandled by ModelOpt — `args[idx + 1]` resolved to `None` (workspace is
None outside FP8), which then crashed `TensorQuantizer.forward` on
`inputs.numel()` with `AttributeError: 'NoneType' object has no
attribute 'numel'`. Surfaced as a regression in Megatron-Bridge after
the TE 2.15 bump alongside ModelOpt 0.44.0rc3.
- **TE 2.10:** `_GroupedLinear.forward`'s second positional slot was
renamed `m_splits` → `non_tensor_args` (tuple wrapping). ModelOpt had a
separate `Version("2.10")` gate for this.

Replace all three version gates with **parameter-name introspection** of
the live `_Linear.forward` / `_GroupedLinear.forward` signature. The
parameter names (`weight`, `inp`, `m_splits`, `non_tensor_args`) have
been stable across TE 1.x, 2.x, and 2.15+; only their relative positions
shift. The new code reads the live signature via
`inspect.signature(...).parameters`, locates `weight`/`inp` by name, and
mutates only those positions in a list copy of `args` — everything
between (e.g. TE 2.15's `weight_workspace`) and after passes through
verbatim. The dual-branch code in `te_quantized_linear_fn` collapses to
a single path.

### Usage

No public API change. PTQ continues to work transparently across all
supported TE versions:

```python
import modelopt.torch.quantization as mtq
# Works on TE 1.x, 2.0-2.14, 2.15.x, and 2.16+ — no version flag needed.
mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop)
```

### Testing
<!-- Mention how have you tested your change if applicable. -->

Existing TE plugin tests
(`tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py`)
exercise both the `_forward` (no-grad calibration) and `_apply`
(grad-enabled training) paths of `te_quantized_linear_fn` for
`te.pytorch.Linear` — they would have caught the original TE 2.15
regression on a CI matrix entry pinned to TE 2.15. Verified trace
correctness across:

| TE version | `_Linear.forward` signature | `_te_linear` weight→inp gap
| `_GroupedLinear.forward` second slot |
|---|---|---|---|
| 1.x | `(ctx, weight, weight_fp8, inp, …)` | 1 | n/a |
| 2.0–2.14 | `(ctx, weight, inp, bias, …)` | 0 | `m_splits` |
| 2.15.x | `(ctx, weight, weight_workspace, inp, …)` | 1 |
`non_tensor_args` |
| 2.16+ (main) | `(ctx, weight, inp, bias, fwd_args)` | 0 |
`non_tensor_args` |

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ <!--- Public API unchanged;
broadens the range of TE versions that work (TE 2.15.x now supported, TE
1.x still supported via the same introspection path). -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A <!--- Only
adds a stdlib `inspect` import. -->
- Did you write any new necessary tests?: Existing tests sufficient
<!--- Bug fix is covered by existing `test_transformer_engine.py` for
whatever single TE version CI exercises. A multi-version TE matrix is
the right next step but is out of scope for this PR. -->

### Additional Information
<!-- E.g. related issue. -->
Triggered by Megatron-Bridge
NVIDIA-NeMo/Megatron-Bridge#3783 failing tests
after bumping ModelOpt 0.44.0rc2 → 0.44.0rc3 together with a Megatron-LM
bump that pulls TE 2.15. ModelOpt rc2 had the same latent bug — it just
wasn't exercised until TE 2.15 became the runtime version.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Refactor**
* Improved Transformer Engine quantization plugin robustness by using
runtime parameter inspection instead of version-based branching,
ensuring compatibility across TE versions without requiring manual
updates.

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1473)

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>