NVIDIA / NVIDIA/Model-Optimizer

[Bugs] Second batch of verified findings from the unit-coverage initiative (fold_weight crash on Llama-4/GPT-OSS MoE, KD-loss default footgun, and smaller items)

Open
#1,926 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
3.8k
Forks
604
Avg merge
2d 8h
Merged PRs (30d)
142

Description

Follow-up to #1902: writing direct unit suites for more modules (PRs #1913-#1916, #1922-#1925) surfaced another batch of verified defects. Everything below was confirmed by at least two independent passes (suite writer + adversarial reviewer, who re-derived math from source and reproduced the crashes). Ordered by priority; happy to send fix PRs for any of these.

1. fold_weight crashes on custom per-weight quantizers — reachable from public mtq.fold_weight() on Llama-4 / GPT-OSS (CRASH)

QuantModule.fold_weight (quant_module.py:141) derives the weight attribute as name[:-10], stripping "_quantizer" — while the adjacent comment claims it strips "_weight_quantizer". It works for the standard weight/weight_quantizer layout only by accident. For custom per-weight quantizers (quantizer_attr_names("gate_up_proj")gate_up_proj_weight_quantizer, weight attr gate_up_proj) it derives gate_up_proj_weight and trips the assertion. Reproduced: AssertionError: gate_up_proj_weight_quantizer doesn't have a corresponding gate_up_proj_weight. Neither _QuantLlama4TextExperts (huggingface.py:751) nor _QuantGptOssExperts (:1451) overrides fold_weight, and mtq.fold_weight() walks every QuantModule — so folding a quantized Llama-4 or GPT-OSS model crashes. Fix caveat: these experts quantize transposed (_transposed_quantize), so a correct fold needs more than name surgery. Documented in #1925.

2. LogitsDistillationLoss rides F.kl_div's deprecated "mean" default (SILENT FUTURE RESCALE)

losses.py:34/60 use the default reduction that PyTorch deprecates in favor of batchmean semantics; when torch flips it in a future major, the default KD loss silently rescales by the class-dim size. Pin reduction explicitly (or default to "batchmean" deliberately). Documented in #1924.

3. modelopt_state() returns aliases of the manager's live state

conversion.py:485 filters into a new list but keeps the same tuple/dict objects — caller mutation silently corrupts the model's recorded modelopt state. A deepcopy may be costly for large quantizer metadata, so this may be better fixed as a documented contract ("treat as read-only") or a shallow-copy of the per-mode dicts. Documented in #1915.

4. Loss-balancer validation gaps (small PR-sized pair)

StaticLossBalancer(1) crashes with TypeError (isinstance accepts only float, loss_balancers.py:95), and individually NEGATIVE weights pass the sum-only range check (:98-103), silently subtracting a loss term. Documented in #1924.

5. Bias-calibrator doc/API quartet (docs-only fix)

calib/bias.py: the axis docstring (and collect()'s comment block incl. its example shapes) says listed dims are KEPT, but the implementation REDUCES them (shipped behavior, locked in by test_affine_quant.py); axis: int | ... annotation contradicts the iterable-only implementation; compute_bias silently falls through to max_min on unknown methods while collect/compute_dynamic_bias raise; the mean/max_min comments in compute_dynamic_bias are swapped. Worst part: the bias field docstring examples in config.py (:505-514, e.g. {"enable": True, "axis": -1}) are all REJECTED by its own validate_bias validator — the documented schema is unusable as written. Documented in #1922.

6. Smaller items

  • DistillationModel.export() removes hooks but leaves stale _intermediate_output attributes pinning the last captured activations (incl. an autograd graph) on student and teacher layers (#1923).
  • MGDLoss does not detach teacher features, unlike its sibling losses — only standalone (non-DistillationModel) use is affected (#1924).
  • _ModeRegistryCls.__del__ is unreachable (the class-level _all_registries holds a strong reference) and unsafe if ever reached (unconditional list.remove raises after manual removal) — surfaced while fixing the vacuous assertion in #1921.
  • logging._disable_tqdm positional-args branch has a double off-by-one (drops args[index-1], inserts True one slot early); unreachable in practice (needs ≥10 positional tqdm args).
  • DistillationLossBalancer declares abstractmethod without ABCMeta, so it instantiates and only fails at call time (#1924).

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Treat the numbered findings as separate work items: start with the public mtq.fold_weight() path, then inspect QuantModule.fold_weight in quant_module.py:141 and the Llama-4/GPT-OSS expert definitions in huggingface.py:751 and :1451. Review losses.py, loss_balancers.py, conversion.py:485, calib/bias.py, config.py:505-514, and test_affine_quant.py for the other findings; done means each confirmed defect has a focused fix and regression coverage or documented contract.

Written by the indexing model from the issue text.

Assessment

Tech stack
huggingface, python, pytorch
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.