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)
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_outputattributes pinning the last captured activations (incl. an autograd graph) on student and teacher layers (#1923).MGDLossdoes not detach teacher features, unlike its sibling losses — only standalone (non-DistillationModel) use is affected (#1924)._ModeRegistryCls.__del__is unreachable (the class-level_all_registriesholds a strong reference) and unsafe if ever reached (unconditionallist.removeraises after manual removal) — surfaced while fixing the vacuous assertion in #1921.logging._disable_tqdmpositional-args branch has a double off-by-one (dropsargs[index-1], inserts True one slot early); unreachable in practice (needs ≥10 positional tqdm args).DistillationLossBalancerdeclaresabstractmethodwithout ABCMeta, so it instantiates and only fails at call time (#1924).
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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