Support per expert weight quantizer in TEGroupedMLP #1550
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -65,6 +65,7 @@ | |
| from ..functional import normalized_hadamard_transform | ||
|
|
||
| __all__ = [ | ||
| "GroupedQuantizer", | ||
| "HardDisabledTensorQuantizer", | ||
| "NVFP4StaticQuantizer", | ||
| "SequentialQuantizer", | ||
|
|
@@ -1641,3 +1642,71 @@ def convert_to_single_quantizer(model, indx: int = 0): | |
| ) in original_sequential_quantizers.items(): | ||
| for name, sequential_quantizer in sequential_quantizers_list: | ||
| setattr(parent_module, name, sequential_quantizer) | ||
|
|
||
|
|
||
| class GroupedQuantizer(nn.ModuleList): | ||
|
Contributor
There was a problem hiding this comment. nit:
Contributor
There was a problem hiding this comment. RB: Can you create a plan on how to do this and share in this thread?
Contributor
There was a problem hiding this comment.
🐝 I’m reviewing the overlapping behavior and will share a concrete refactoring plan here shortly.
Contributor
There was a problem hiding this comment.
Plan:
This keeps the refactor local to
Contributor
There was a problem hiding this comment. @jenchen13 could you take a look at this comment? |
||
| """A container for per-group :class:`TensorQuantizer` modules. | ||
|
|
||
| Used when a single linear holds several independently-quantized weights — e.g. the | ||
| fused experts of a TEGroupedLinear, where each of the ``num_gemms`` weights needs its | ||
| own ``amax``. Unlike :class:`SequentialQuantizer` (an ``nn.Sequential`` that *chains* | ||
| quantizers over one tensor), the contained quantizers act on *different* tensors, so | ||
| there is no inherent forward path: index in with ``grouped[i](weight_i)``. | ||
|
|
||
| Property reads (``amax``, ``is_enabled``) delegate to the first quantizer — all members | ||
| share one config, so the first is representative for "is this calibrated/enabled" | ||
| checks; the real per-group values live on the members and are used via indexing. | ||
| Lifecycle/config methods broadcast to every member. | ||
| """ | ||
|
|
||
| _delegated_properties = ["fake_quant", "is_enabled", "amax"] | ||
| _delegated_methods = [ | ||
| "reset_amax", | ||
| "disable", | ||
| "enable", | ||
| "load_calib_amax", | ||
| "load_calib_bias", | ||
| ] | ||
|
|
||
| def __init__(self, *quantizers: "TensorQuantizer | SequentialQuantizer"): | ||
| """Initialize GroupedQuantizer module.""" | ||
| super().__init__(quantizers) | ||
| assert all(isinstance(q, (TensorQuantizer, SequentialQuantizer)) for q in self), ( | ||
| "All quantizers must be a TensorQuantizer or SequentialQuantizer." | ||
| ) | ||
|
|
||
| def forward(self, inputs): | ||
| """Apply the representative quantizer for single-weight compatibility paths.""" | ||
| return self[0](inputs) | ||
|
|
||
| def __getattr__(self, name): | ||
| """Delegate property reads to the first member and method calls to all members.""" | ||
| if name in self._delegated_properties: | ||
| return getattr(self[0], name) | ||
|
|
||
| if name in self._delegated_methods: | ||
|
|
||
| def method_wrapper(*args, **kwargs): | ||
| return [getattr(quantizer, name)(*args, **kwargs) for quantizer in self] | ||
|
|
||
| return method_wrapper | ||
|
|
||
| return super().__getattr__(name) | ||
|
|
||
| def __setattr__(self, name, value): | ||
| if name in self._delegated_properties: | ||
| for quantizer in self: | ||
| setattr(quantizer, name, value) | ||
| else: | ||
| super().__setattr__(name, value) | ||
|
|
||
| def set_from_attribute_config(self, attributes): | ||
| """Set the attributes of contained quantizers; a single config broadcasts to all.""" | ||
| if not isinstance(attributes, (list, tuple)): | ||
| attributes = [attributes] * len(self) | ||
| for attribute, quantizer in zip(attributes, self): | ||
| quantizer.set_from_attribute_config(attribute) | ||
|
|
||
| def get_modelopt_state(self) -> dict[str, Any]: | ||
| """Get meta state to be saved in checkpoint.""" | ||
| return {"num_quantizers": len(self), "is_grouped_quantizer": True} | ||
|
There was a problem hiding this comment. [SUGGESTION] |
||
|
Contributor
There was a problem hiding this comment. Design looks good to me overall, left a comment here - https://github.com/NVIDIA/Model-Optimizer/pull/1550/changes#r3468773327 |
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -15,7 +15,9 @@ | |||||
|
|
||||||
| """Support quantization for Transformer Engine layers.""" | ||||||
|
|
||||||
| import copy | ||||||
| import inspect | ||||||
| import os | ||||||
| import warnings | ||||||
|
|
||||||
| import torch | ||||||
|
|
@@ -27,11 +29,13 @@ | |||||
|
|
||||||
| from modelopt.torch.quantization.utils import replace_function | ||||||
|
|
||||||
| from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer | ||||||
| from .custom import _ParallelLinear | ||||||
|
|
||||||
| _TE_VERSION = Version(te.__version__) | ||||||
|
|
||||||
| _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV = "MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP" | ||||||
|
|
||||||
|
|
||||||
| def _assert_te_fp8_enabled(): | ||||||
| """Check if Transformer Engine FP8 autocast is enabled and raise error if so.""" | ||||||
|
|
@@ -48,6 +52,13 @@ def _assert_te_fp8_enabled(): | |||||
| pass # Older TE versions may not have this API | ||||||
|
|
||||||
|
|
||||||
| def _is_calibrating(quantizer): | ||||||
| """Return whether a tensor or sequential quantizer is collecting calibration stats.""" | ||||||
| if isinstance(quantizer, SequentialQuantizer): | ||||||
| return any(getattr(q, "_if_calib", False) for q in quantizer) | ||||||
| return getattr(quantizer, "_if_calib", False) | ||||||
|
|
||||||
|
|
||||||
| @QuantModuleRegistry.register({te.pytorch.Linear: "te_Linear"}) | ||||||
| class _QuantTELinear(_ParallelLinear): | ||||||
| @property | ||||||
|
|
@@ -137,8 +148,27 @@ def _setup(self): | |||||
| # Remove self.weight after setup. | ||||||
| delattr(self, "weight") | ||||||
|
|
||||||
| # Each fused expert gets its own weight quantizer (independent amax). Storing them in a | ||||||
| # GroupedQuantizer (an nn.ModuleList) surfaces them as ``weight_quantizer.{i}``, which the | ||||||
| # fused-experts name normalizer maps to ``*weight_quantizer`` so the stock configs apply. | ||||||
| self.weight_quantizer = GroupedQuantizer( | ||||||
| *(copy.deepcopy(self.weight_quantizer) for _ in range(self.num_gemms)) | ||||||
|
Contributor
There was a problem hiding this comment. We dont need this deep copy here.
Suggested change
The TensorQuantizers are always initialized with default args; they are configured after
Contributor
There was a problem hiding this comment. RB: Can you please explain concisely "The TensorQuantizers are always initialized with default args; they are configured after _setup -> so we can safely initialize new TensorQuantizers here. deepcopy seems unwarranted to me." with code pointers ?
Contributor
There was a problem hiding this comment.
🐝 I’m tracing the initialization and post-
Contributor
There was a problem hiding this comment.
Yes—the key is the conversion order:
So |
||||||
| ) | ||||||
|
|
||||||
| # Compile only the per-expert quantizer loop. The surrounding TE grouped GEMM remains | ||||||
| # eager, and the opt-in flag leaves the default execution path unchanged. | ||||||
| if os.getenv(_COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, "0") == "1": | ||||||
| quantizers = tuple(self.weight_quantizer) | ||||||
|
|
||||||
| def quantize_weights(*weights): | ||||||
| return tuple(quantizer(weight) for quantizer, weight in zip(quantizers, weights)) | ||||||
|
|
||||||
| self._compiled_weight_quantizer_loop = torch.compile( | ||||||
|
realAsma marked this conversation as resolved.
|
||||||
| quantize_weights, | ||||||
| backend="inductor", | ||||||
| fullgraph=False, | ||||||
| mode="reduce-overhead", | ||||||
| ) | ||||||
|
|
||||||
| def modelopt_post_restore(self, prefix: str = ""): | ||||||
| # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to | ||||||
|
|
@@ -150,12 +180,29 @@ def modelopt_post_restore(self, prefix: str = ""): | |||||
| # Remove self.weight after post_restore. | ||||||
| delattr(self, "weight") | ||||||
|
|
||||||
| # Base post_restore only re-calibrates the first member; the other per-expert | ||||||
| # quantizers also need re-calibration so a TP/EP change between save and restore | ||||||
| # produces correctly shaped per-channel _amax. Mirror base behavior: only | ||||||
| # re-calibrate quantizers whose loaded state had _amax (skip unused ones). | ||||||
| from modelopt.torch.quantization.model_calib import max_calibrate | ||||||
|
|
||||||
| for i in range(self.num_gemms): | ||||||
| weight_i = getattr(self, f"weight{i}", None) | ||||||
| if weight_i is None: | ||||||
| continue | ||||||
| wq_i = self.weight_quantizer[i] | ||||||
| q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i | ||||||
| if not hasattr(q, "_amax"): | ||||||
|
There was a problem hiding this comment. [CRITICAL ModeState] Per-expert restore loses What happens: self.weight_quantizer.reset_amax() # GroupedQuantizer delegated method -> broadcasts to ALL members
max_calibrate(self.weight_quantizer, lambda wq: wq(self.weight), distributed_sync=False)
This loop is then supposed to recalibrate the rest, but the guard Impact: After Fix: determine which members had saved had_amax = []
for i in range(self.num_gemms):
wq_i = self.weight_quantizer[i]
q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i
had_amax.append(hasattr(q, "_amax"))
super().modelopt_post_restore(prefix=prefix)
delattr(self, "weight")
from modelopt.torch.quantization.model_calib import max_calibrate
for i in range(self.num_gemms):
weight_i = getattr(self, f"weight{i}", None)
if weight_i is None or not had_amax[i]:
continue
wq_i = self.weight_quantizer[i]
wq_i.reset_amax()
max_calibrate(wq_i, lambda wq, w=weight_i: wq(w), distributed_sync=False)(Note |
||||||
| continue | ||||||
| wq_i.reset_amax() | ||||||
| max_calibrate(wq_i, lambda wq, w=weight_i: wq(w), distributed_sync=False) | ||||||
|
|
||||||
| def iter_weights_for_calibration(self): | ||||||
| """Yield ``(weight_i, weight_quantizer)`` for each of the ``num_gemms`` grouped weights.""" | ||||||
| for i in range(self.num_gemms): | ||||||
| weight_i = getattr(self, f"weight{i}", None) | ||||||
| if weight_i is not None: | ||||||
| yield weight_i, self.weight_quantizer[i] | ||||||
|
|
||||||
| @staticmethod | ||||||
| def te_grouped_quantized_linear_fn(package, func_name, self, *args): | ||||||
|
|
@@ -184,8 +231,19 @@ def te_grouped_quantized_linear_fn(package, func_name, self, *args): | |||||
|
|
||||||
| new_args = list(args) | ||||||
| new_args[inp_pos] = self.input_quantizer(args[inp_pos]) | ||||||
| weights = tuple(args[weights_start : weights_start + num_gemms]) | ||||||
| # Calibration mutates collector state and must stay outside Inductor/CUDAGraph capture. | ||||||
| use_compiled_loop = hasattr(self, "_compiled_weight_quantizer_loop") and not any( | ||||||
| _is_calibrating(quantizer) for quantizer in self.weight_quantizer | ||||||
| ) | ||||||
| if use_compiled_loop: | ||||||
| quantized_weights = self._compiled_weight_quantizer_loop(*weights) | ||||||
| else: | ||||||
| quantized_weights = tuple( | ||||||
| self.weight_quantizer[gemm_idx](weight) for gemm_idx, weight in enumerate(weights) | ||||||
| ) | ||||||
| for gemm_idx, quantized_weight in enumerate(quantized_weights): | ||||||
| new_args[weights_start + gemm_idx] = quantized_weight | ||||||
| output = getattr(package, func_name)(*new_args) | ||||||
| # TE 2.15+ returns `(out, new_workspaces)`; TE <= 2.14 returns just `out`. | ||||||
| # Only the activation tensor participates in output quantization. | ||||||
|
|
||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit:
Create QuantizerContainerBase base class and use that here to avoid explicitly listing like this.
see https://github.com/NVIDIA/Model-Optimizer/pull/1550/changes#r3553373837