feat(quant): quantization of packed MoE experts

cosmos_framework/quantization/export.py

@@ -98,6 +98,11 @@ def remap_framework_state_dict(
             for suffix in _SCALE_SUFFIXES:
                 if not net_key.endswith(suffix):
                     continue
+                # Packed MoE expert scales sit beside the packed parameter itself.
+                param_diffusers_key = net_to_diffusers.get(net_key[: -len(suffix)])
+                if param_diffusers_key is not None:
+                    diffusers_key = param_diffusers_key + suffix
+                    break
                 weight_key = net_key[: -len(suffix)] + ".weight"
                 weight_diffusers_key = net_to_diffusers.get(weight_key)
                 if weight_diffusers_key is not None:
@@ -111,13 +116,40 @@ def remap_framework_state_dict(
     return remapped
 
 
+def _remap_packed_experts(
+    experts: dict[str, torch.Tensor], bf16_state_dict: dict[str, torch.Tensor]
+) -> dict[str, torch.Tensor]:
+    """:func:`remap_framework_state_dict` for packed experts, where every tensor must map.
+
+    A dropped expert would leave its BF16 source tensor without scales under an FP8 config.
+    """
+    remapped = remap_framework_state_dict(experts, bf16_state_dict)
+    if len(remapped) != len(experts):
+        unmapped = sorted(
+            key for key in experts if not remap_framework_state_dict({key: experts[key]}, bf16_state_dict)
+        )
+        raise ValueError(f"{len(unmapped)} packed expert tensor(s) have no Diffusers key: {unmapped[:10]}")
+    return remapped
+
+
 # Top-level projections produced by reverse_remap use the cosmos3_vfm names; the
 # bf16 diffusers base provides the equivalently-named proj_in/proj_out, so we drop
 # ours and keep the base copies.
 _DROP_FROM_QUANTIZED = ("vae2llm.", "llm2vae.")
 _FP8_DTYPES = {torch.float8_e4m3fn, getattr(torch, "float8_e5m2", torch.float8_e4m3fn)}
 FP8_MAX = 448.0
 _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale", ".input_scale_2", ".alpha", ".pre_quant_scale")
+_PACKED_EXPERT_KEY_SUFFIXES = tuple(
+    f"experts.{name}{suffix}"
+    for name in ("gate_up_proj", "down_proj")
+    for suffix in ("", ".weight_scale", ".input_scale")
+)
+
+
+def is_packed_expert_key(key: str) -> bool:
+    """True for a packed expert parameter or its scale, e.g. ``...experts.gate_up_proj.weight_scale``."""
+    return key.endswith(_PACKED_EXPERT_KEY_SUFFIXES)
+
 
 # Linears the diffusers checkpoint keeps in bf16 (never quantized), emitted in the
 # ModelOpt ``ignore`` list so a reasoning-path loader (vllm Cosmos3ForConditionalGeneration)
@@ -207,7 +239,12 @@ def _remap_framework_modelopt_state(
     def remap_module(name: str) -> str | None:
         return module_map.get(name.removeprefix("net."))
 
-    active_modules = {key.removesuffix(".weight_scale") for key in exported_state_dict if key.endswith(".weight_scale")}
+    # Packed MoE experts are not linears and have no component graph state.
+    active_modules = {
+        key.removesuffix(".weight_scale")
+        for key in exported_state_dict
+        if key.endswith(".weight_scale") and not is_packed_expert_key(key)
+    }
 
     real_quantize = modes["real_quantize"]
     real_metadata = real_quantize["metadata"]
@@ -381,7 +418,7 @@ def overlay(self, bf16: dict, quantized: dict) -> dict:
                 merged[k] = v
                 n_scale += 1
                 continue
-            if k.endswith(".weight") and v.dtype in _FP8_DTYPES:
+            if (k.endswith(".weight") or is_packed_expert_key(k)) and v.dtype in _FP8_DTYPES:
                 if k not in bf16:
                     raise KeyError(
                         f"quantized weight {k!r} has no bf16 counterpart — layer-key "
@@ -390,7 +427,9 @@ def overlay(self, bf16: dict, quantized: dict) -> dict:
                 if bf16[k].shape != v.shape:
                     raise ValueError(f"shape mismatch for {k}: bf16 {tuple(bf16[k].shape)} vs fp8 {tuple(v.shape)}")
                 merged[k] = v
-                quantized_modules.add(k[: -len(".weight")])
+                # Packed experts keep only their scales; the quantizer buffers below are per linear.
+                if not is_packed_expert_key(k):
+                    quantized_modules.add(k[: -len(".weight")])
                 n_weight += 1
                 continue
             # Remaining bf16 tensors (norms, embed_tokens, renamed projections) are
@@ -423,6 +462,7 @@ def export(
         modelopt_config: dict | None = None,
     ) -> None:
         from .checkpoint_io import load_sharded_safetensors, save_sharded_safetensors
+        from .moe import QuantPackedExperts, export_packed_experts
 
         net = mdl.net
         transformer_export_dir.mkdir(parents=True, exist_ok=True)
@@ -431,6 +471,14 @@ def export(
         if n_collapsed > 0:
             print(f"[export] collapsed per-channel input amax to scalar on {n_collapsed} quantizers")
 
+        # Packed MoE experts are exported directly; ModelOpt's compress and unified
+        # exporter only handle linears, so their quantizers are disabled afterwards.
+        experts = export_packed_experts(net)
+        for module in net.modules():
+            if isinstance(module, QuantPackedExperts):
+                for name in ("gate_up_proj", "down_proj"):
+                    getattr(module, f"{name}_input_quantizer").disable()
+
         # Record ModelOpt's real-quantized graph metadata before the unified
         # exporter unpacks those weights into plain FP8 checkpoint tensors.
         mtq.compress(mdl)
@@ -446,6 +494,7 @@ def export(
         with hide_quantizers_from_state_dict(net):
             net_state_dict = {k: v.detach().contiguous().cpu() for k, v in net.state_dict().items()}
         quantized = remap_framework_state_dict(net_state_dict, bf16)
+        quantized.update(_remap_packed_experts(experts, bf16))
         merged = self.overlay(bf16, quantized)
         n_shards = save_sharded_safetensors(merged, transformer_export_dir)
         print(f"[export] wrote {len(merged)} tensors across {n_shards} shard(s)")
@@ -505,12 +554,13 @@ def _write_root_hf_quant_config(root_dir: Path, transformer_dir: Path) -> None:
     print(f"[assemble] wrote {root_dir / 'hf_quant_config.json'} for upstream-vLLM discovery")
 
 
-def _write_root_weight_index(root_dir: Path) -> None:
-    """(Re)generate the root ``model.safetensors.index.json`` for whole-model loading."""
+def _write_root_weight_index(root_dir: Path) -> dict[str, str]:
+    """(Re)generate the root ``model.safetensors.index.json``; return its weight map."""
     index = build_root_index(root_dir)
     dst = root_dir / "model.safetensors.index.json"
     index.write(dst)
     print(f"[assemble] wrote {dst.name} ({len(index.weight_map)} tensors, {index.metadata['total_size']} bytes)")
+    return index.weight_map
 
 
 def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer_dir: Path):
@@ -520,6 +570,7 @@ def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer
     vision_encoder symlinked back); only ``transformer/`` is physically new (the
     quantized export). A fresh root weight index + a root ``hf_quant_config.json`` are
     written so both the diffusers and transformers loaders resolve the new weights.
+    The root Transformers ModelOpt state is skipped when the transformer has packed experts.
     """
     input_dir, output_dir = Path(input_dir), Path(output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
@@ -552,8 +603,12 @@ def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer
         rel_src = os.path.relpath(src, start=output_dir.resolve())
         print(f"[assemble] linking {entry.name} -> {rel_src}")
         dst.symlink_to(rel_src)
-    _write_root_weight_index(output_dir)
-    component_state = torch.load(target_transformer / "modelopt_state.pth", map_location="cpu", weights_only=False)
-    Fp8DiffusersExporter().write_transformers_modelopt_state(component_state, output_dir, output_dir)
+    weight_map = _write_root_weight_index(output_dir)
+    has_packed_experts = any(
+        is_packed_expert_key(key) for key, file in weight_map.items() if file.startswith("transformer/")
+    )
+    if not has_packed_experts:
+        component_state = torch.load(target_transformer / "modelopt_state.pth", map_location="cpu", weights_only=False)
+        Fp8DiffusersExporter().write_transformers_modelopt_state(component_state, output_dir, output_dir)
     _write_root_hf_quant_config(output_dir, target_transformer)
     print(f"[assemble] drop-in dir ready at {output_dir}")

cosmos_framework/quantization/moe.py

@@ -0,0 +1,167 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+"""ModelOpt FP8 quantization of the packed grouped-mm MoE experts.
+
+Calibrates and fake-quantizes the expert inputs through the wrapper's ``forward``, and quantizes
+the expert weights (with one scale per expert) directly in :func:`export_packed_experts`.
+This code does not support weight-only.
+
+How to run quantization: call :func:`register_packed_experts` before ``mtq.quantize`` (and before
+restoring a saved ModelOpt state), then export with
+:func:`~cosmos_framework.quantization.export.export_quantized_transformer`, which writes the
+experts via :func:`export_packed_experts`.
+
+Without the registration, ModelOpt's generic ``_QuantFusedExperts`` claims these experts and,
+since it only intercepts ``F.linear`` and they run ``torch._grouped_mm``, quantizes nothing.
+"""
+
+from __future__ import annotations
+
+import torch
+from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry, TensorQuantizer
+from torch import nn
+
+from cosmos_framework.model.generator.moe.interface import MoEExperts
+from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import Qwen3VLMoeTextExpertsGroupedMm
+from cosmos_framework.quantization.export import FP8_MAX
+
+_PACKED_EXPERT_TYPES = (Qwen3VLMoeTextExpertsGroupedMm, MoEExperts)
+
+
+class QuantPackedExperts(QuantModule):
+    """Calibrates and fake-quantizes expert inputs; per-expert FP8 weights are quantized on export."""
+
+    def _setup(self):
+        self.gate_up_proj_input_quantizer = TensorQuantizer()
+        self.down_proj_input_quantizer = TensorQuantizer()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # [num_tokens,hidden_size]
+        topk_scores: torch.Tensor,  # [num_tokens,top_k]
+        expert_indices: torch.Tensor,  # [num_tokens,top_k]
+        num_tokens_per_expert: torch.Tensor,  # [num_experts]
+        token_mask: torch.Tensor | None = None,  # [num_tokens]
+    ) -> torch.Tensor:  # [num_tokens,hidden_size]
+        # Group the real (token, slot) pairs by expert; ``num_tokens_per_expert`` counts exactly these.
+        top_k = expert_indices.shape[1]
+        slots = torch.arange(expert_indices.numel(), device=expert_indices.device)  # [num_tokens*top_k]
+        if token_mask is not None:
+            slots = slots[token_mask.repeat_interleave(top_k)]  # [num_routed]
+        if slots.numel() == 0:
+            return torch.zeros_like(hidden_states)  # [num_tokens,hidden_size]
+        slots = slots[torch.argsort(expert_indices.reshape(-1)[slots], stable=True)]  # [num_routed]
+        tokens = slots // top_k  # [num_routed]
+        scores = topk_scores.reshape(-1)[slots].to(hidden_states.dtype)  # [num_routed]
+        x = hidden_states[tokens]  # [num_routed,hidden_size]
+        routed = self._quantized_experts(x, num_tokens_per_expert, scores)  # [num_routed,hidden_size]
+        return torch.zeros_like(hidden_states).index_add_(0, tokens, routed)  # [num_tokens,hidden_size]
+
+    def _quantized_experts(
+        self,
+        x: torch.Tensor,  # [num_routed,hidden_size]  (grouped by expert)
+        num_tokens_per_expert: torch.Tensor,  # [num_experts]
+        scores: torch.Tensor,  # [num_routed]
+    ) -> torch.Tensor:  # [num_routed,hidden_size]
+        # a2 observes ``act(gate) * up`` before the routing scores, because vLLM applies the scores
+        # after the down GEMM. Per-expert weight scales are computed directly on export.
+        offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32)  # [num_experts]
+        x = self.gate_up_proj_input_quantizer(x)
+        h = torch._grouped_mm(x, self.gate_up_proj, offs=offsets)
+        g, u = torch.chunk(h, chunks=2, dim=-1)  # 2x [num_routed,moe_intermediate_size]
+        a = self.down_proj_input_quantizer(self.act_fn(g) * u)  # [num_routed,moe_intermediate_size]
+        return torch._grouped_mm(a * scores.unsqueeze(-1), self.down_proj, offs=offsets)  # [num_routed,hidden_size]
+
+
+def register_packed_experts() -> None:
+    """Register :class:`QuantPackedExperts` before ``mtq.quantize`` so ``_QuantFusedExperts`` skips it."""
+    pending: dict[type[nn.Module], str] = {}
+    for expert_cls in _PACKED_EXPERT_TYPES:
+        registered = QuantModuleRegistry.get(expert_cls)
+        if registered is not None and issubclass(registered, QuantPackedExperts):
+            continue
+        if registered is not None:
+            # An earlier ``mtq.quantize`` that saw this type unregistered has globally registered
+            # ModelOpt's generic ``_QuantFusedExperts`` for it, which would quantize nothing here.
+            QuantModuleRegistry.unregister(expert_cls)
+        pending[expert_cls] = expert_cls.__name__
+    if pending:
+        QuantModuleRegistry.register(pending)(QuantPackedExperts)
+
+
+def has_packed_experts(model: nn.Module) -> bool:
+    return any(isinstance(module, _PACKED_EXPERT_TYPES) for module in model.modules())
+
+
+def assert_quantizers_calibrated(model: nn.Module) -> None:
+    """Fail export instead of writing a missing or zero scale for an enabled quantizer."""
+    problems = []
+    for name, module in model.named_modules():
+        if isinstance(module, TensorQuantizer) and not module.is_enabled:
+            continue
+        if isinstance(module, TensorQuantizer):
+            amax = getattr(module, "amax", None)
+            if amax is None:
+                problems.append(f"{name}: no amax")
+            elif not torch.isfinite(amax).all():
+                problems.append(f"{name}: non-finite amax")
+            elif name.endswith("input_quantizer") and not (amax > 0).all():
+                problems.append(f"{name}: non-positive amax")
+        elif (
+            isinstance(getattr(module, "gate_up_proj", None), nn.Parameter)
+            and module.gate_up_proj.ndim == 3
+            and not isinstance(module, QuantPackedExperts)
+        ):
+            problems.append(f"{name}: packed experts {type(module).__name__} are not quantized")
+    if problems:
+        shown = "; ".join(problems[:10])
+        raise ValueError(f"{len(problems)} quantizer problem(s), first {min(len(problems), 10)}: {shown}")
+
+
+_PACKED_EXPERT_QUANTIZERS = (
+    "gate_up_proj_input_quantizer",
+    "down_proj_input_quantizer",
+)
+
+
+def _check_per_expert_fp8(name: str, module: QuantPackedExperts) -> None:
+    """Refuse any input quantizer layout but FP8 E4M3 with a scalar amax."""
+    for quantizer_name in _PACKED_EXPERT_QUANTIZERS:
+        quantizer = getattr(module, quantizer_name)
+        amax_shape = tuple(quantizer.amax.shape)
+        if quantizer.num_bits != (4, 3) or quantizer.block_sizes:
+            found = f"num_bits={quantizer.num_bits}, block_sizes={quantizer.block_sizes}"
+        elif quantizer.amax.numel() == 1:
+            continue
+        else:
+            found = f"amax shape {amax_shape}, expected a single element"
+        raise ValueError(f"{name}.{quantizer_name}: {found}; packed-expert export supports per-expert FP8 only")
+
+
+def export_packed_experts(net: nn.Module) -> dict[str, torch.Tensor]:
+    """FP8 packed experts in the source layout, with per-expert weight scales and scalar input scales."""
+    if not has_packed_experts(net):
+        return {}  # Dense models export exactly as before, without the calibration check.
+    assert_quantizers_calibrated(net)
+    exported = {}
+    for name, module in net.named_modules():
+        if not isinstance(module, QuantPackedExperts):
+            continue
+        enabled = [getattr(module, quantizer).is_enabled for quantizer in _PACKED_EXPERT_QUANTIZERS]
+        if not any(enabled):
+            continue  # Skipped by the coverage filter; the experts stay BF16 from the source.
+        if not all(enabled):
+            raise ValueError(f"{name or 'experts'}: packed-expert quantizers are only partly enabled: {enabled}")
+        _check_per_expert_fp8(name or "experts", module)
+        for param_name in ("gate_up_proj", "down_proj"):
+            # BF16 models keep BF16 amax; the scales are FP32 on disk.
+            param = getattr(module, param_name).detach()
+            weight_amax = param.abs().amax(dim=(1, 2), keepdim=True).float()
+            weight_scale = torch.where(weight_amax == 0, 1.0, weight_amax / FP8_MAX)  # [num_experts,1,1]
+            fp8 = (param.float() / weight_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+            input_amax = getattr(module, f"{param_name}_input_quantizer").amax.detach().float()
+            key = f"{name}.{param_name}" if name else param_name
+            exported[key] = fp8.contiguous().cpu()
+            exported[f"{key}.weight_scale"] = weight_scale.contiguous().cpu()
+            exported[f"{key}.input_scale"] = (input_amax / FP8_MAX).reshape(()).contiguous().cpu()
+    return exported

pyproject.toml

@@ -61,6 +61,7 @@ dependencies = [
 [project.optional-dependencies]
 automl = ["nvidia-automl-core==0.1.0"]
 automl-prompts = ["nvidia-automl-core==0.1.0", "gepa==0.1.2"]
+quantization = ["datasets", "nvidia-modelopt==0.44.0"]
 reasoner-eval = [
   "attrs",
   "evaluate",