diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 27d44bd072d..f99d9020f74 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -28,7 +28,7 @@ from modelopt.torch.utils.distributed import distribution_plan from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax -from .moe_utils import _export_fused_experts +from .moe_utils import _export_fused_experts, _export_fused_experts_keep_fused from .quant_format import QUANTIZATION_NONE from .quant_utils import get_quantization_format from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry @@ -226,13 +226,37 @@ def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContext) -> None: """Split and quantize a fused-experts module with plural weight quantizers. - Under FSDP2 each rank holds only some experts, so it packs just those (kept fused) and the - split into per-expert keys is deferred until the gather brings all experts together. - Tied experts are packed independently and their duplicate keys are dropped by name in postprocess_state_dict; no per-module dedup cache is used. + + Distributed no-gather export (``ctx.keep_fused_experts``) folds the quantizers into the fused + weight IN PLACE, quantizing only this rank's local experts and keeping the result ``Shard(0)`` -- + the per-expert split happens later on the sharded weight (``split_fused_experts_state_dict``). + Otherwise (single process, or gather export) it splits into full per-expert submodules here. """ - _export_fused_experts(module, ctx.dtype) + attr = name.rpartition(".")[2] + if ctx.keep_fused_experts and attr != "experts": + # The split into per-expert keys finds a fused weight by its ``experts.`` key, as the + # writer does, so one named otherwise would be written fused and fail only at load. + raise NotImplementedError( + f"{name}: the keep-fused export splits fused experts into per-expert keys by their " + f"`experts` attribute name, and this module is named {attr!r}." + ) + if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: + # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands + # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its + # original shards the next time the layer reshards, and exported unquantized. The window + # unshards FSDP2's params, so under FSDP2 alone each rank folds every expert of the unit, one + # unit at a time as a forward holds it; expert weights FSDP2 does not shard, as expert + # parallelism keeps them, stay sharded over the experts, and each rank folds its own. + with _fsdp2_weight_window(ctx, module): + _export_fused_experts_keep_fused(module, ctx.dtype) + elif ctx.keep_fused_experts: + # No-gather distributed export: fold each rank's LOCAL experts and leave the weight + # fused + sharded; the per-expert split happens after consolidation. + _export_fused_experts_keep_fused(module, ctx.dtype) + else: + _export_fused_experts(module, ctx.dtype) @ExportModuleRegistry.register(predicate=is_quantlinear) @@ -286,9 +310,19 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) -> def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> None: """Export fused BMM-style expert weights (Llama4 / GPT-OSS). - Its weight quantizer has one amax covering all experts, so under FSDP2 each rank can pack - just the experts it owns and produce identical bytes -- no gather or unshard needed. + Their weight quantizer has one amax covering all experts. The keep-fused export refuses them, + since they store their experts transposed. """ + if ctx.keep_fused_experts: + # The keep-fused export leaves fused experts for the writer to split by their + # ``experts.`` key, as ``(num_experts, out_dim, in_dim)`` weights. These store theirs + # transposed (GPT-OSS also interleaves gate and up, and keeps biases), so the split would cut + # every expert wrong, quantized or not. + raise NotImplementedError( + f"{name}: the keep-fused export splits fused experts stored (num_experts, out_dim, " + f"in_dim), and {type(module).__name__} stores its experts transposed; export it through " + "the gather export instead." + ) if get_quantization_format(module) == QUANTIZATION_NONE: return # TODO: consolidate uncalibrated experts handling logic diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 02660142736..07e0ddd4129 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -21,6 +21,9 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate, Shard + +from modelopt.torch.export.quant_format import GGML_FORMATS def _delete_fused_moe_source_attrs(module: nn.Module) -> None: @@ -110,6 +113,295 @@ def _fallback_amax(weight: torch.Tensor, quantizer) -> torch.Tensor: return torch.maximum(groups.amax(dim=-1), -groups.amin(dim=-1)).float().reshape(like.shape) +def _rewrap_expert_rows(local: torch.Tensor, mesh, placements, num_experts: int): + """Wrap this rank's experts as a DTensor spanning ``num_experts`` rows on the expert axis.""" + shape = torch.Size((num_experts, *local.shape[1:])) + stride = torch.empty(shape, device="meta").stride() + return DTensor.from_local(local, mesh, placements, run_check=False, shape=shape, stride=stride) + + +class _LocalRows: + """``module``, with its fused ``proj`` replaced by this rank's rows of it.""" + + def __init__(self, module: nn.Module, proj: str, local: torch.Tensor): + self._module, self._proj, self._local = module, proj, local + + def __getattr__(self, name: str): + return self._local if name == self._proj else getattr(self._module, name) + + +def _local_expert_weight(module: nn.Module, proj: str, local: torch.Tensor, idx: int, dtype): + """``module.get_expert_weight(proj, idx, dtype)``, read from ``local`` rather than the module. + + The per-expert export packs each expert as that accessor returns it, which casts unpacked native + FP8 experts, held at the default dtype, to the export dtype. The accessor indexes the fused + parameter itself, a sharded DTensor here, so it is handed this rank's rows instead. Whatever + else it reads, such as how it picks this rank's experts out of a weight, is the module's own. + The rows stand in for ``self.`` only where the accessor reads it itself: a method it calls + is bound to the module, and would read the whole fused DTensor. + """ + return type(module).get_expert_weight(_LocalRows(module, proj, local), proj, idx, dtype) + + +def _stack_row(stack: torch.Tensor | None, idx: int, row: torch.Tensor, rows: int, device): + """Write ``row`` as row ``idx`` of a ``rows``-row stack, made on ``device`` for the first row. + + The fold packs one expert at a time: copied into one stack as each is packed, the projection is + held once, where a list of the packed experts and their ``torch.stack`` would hold it twice. + """ + if stack is None: + stack = torch.empty((rows, *row.shape), dtype=row.dtype, device=device) + elif row.shape != stack.shape[1:] or row.dtype != stack.dtype: + raise ValueError( + f"Expert {idx} packs to {row.dtype} {tuple(row.shape)}, the experts before it to " + f"{stack.dtype} {tuple(stack.shape[1:])}; they cannot be stacked." + ) + stack[idx].copy_(row) + return stack + + +def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantizers: int) -> None: + """Refuse a sharded fused expert weight the keep-fused fold cannot fold, on every rank alike. + + Every rank reaches the same verdict: the global shape and the mesh are the same on all of them, + and their quantizer counts differ only where each keeps its own experts' quantizers of an uneven + shard, which the checks below refuse on every rank. + """ + name = f"{type(module).__name__}.{proj}" + sharded = [placement for placement in weight.placements if isinstance(placement, Shard)] + if len(sharded) > 1 or any(placement.dim != 0 for placement in sharded): + # The fold shards the scales on the expert axis alone, and the ranks agree on their experts' + # quantization over the one mesh dim that shards it. + raise NotImplementedError( + f"{name}: placements {weight.placements}; the keep-fused fold needs the experts (dim 0) " + "sharded over one mesh dim and no other dim sharded." + ) + shards = 1 + for mdim, placement in enumerate(weight.placements): + if isinstance(placement, Shard) and placement.dim == 0: + size = weight.device_mesh.size(mdim) + # A rank holding no experts would stack none and register no input scale, so the ranks' + # state dicts would differ for the DCP write. The experts are chunked as torch does. + chunk = -(-weight.shape[0] // size) + if not chunk or -(-weight.shape[0] // chunk) < size: + raise NotImplementedError( + f"{name}: {weight.shape[0]} experts over {size} shards leaves a rank none; the " + "keep-fused fold needs every rank to hold an expert." + ) + shards *= size + # FSDP2 keeps every expert's quantizer on every rank; expert parallelism keeps only this rank's, + # in the order of its shard: transformers sets num_experts to the local count, so the quantized + # module makes one quantizer per expert of its local weight, which its forward picks by offset. + if num_quantizers != weight.shape[0] and num_quantizers * shards != weight.shape[0]: + raise NotImplementedError( + f"{name}: {num_quantizers} weight quantizers for {weight.shape[0]} experts in {shards} " + "shards; the keep-fused fold needs one per expert, or one per expert of this rank's shard." + ) + + +def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> None: + """Fold per-expert quantizers into the FUSED 3-D expert weights *in place*, staying sharded. + + Distributed (DCP) counterpart to :func:`_export_fused_experts`. That function slices the fused + weight into full per-expert tensors -- under FSDP2 it must first all-gather the ``Shard(0)`` + fused weight (via the export handler's ``unshard``) and then materializes every expert on every + rank, putting the whole model on each GPU (OOMs at 235B/480B). This keeps each projection a + single fused param: on a ``Shard(0)`` DTensor it quantizes ONLY this rank's local experts + (``to_local()``, each with its own weight quantizer: by global index where every rank keeps every + expert's, by position in the shard where a rank keeps only its own) and rewraps the result + ``Shard(0)``, so the model stays sharded and DCP writes per-rank shards. Where a rank keeps only + its own experts' quantizers, the ranks first gather each projection's quantizer states + (collective over the expert mesh dim), so every rank folds, leaves or refuses it alike. + On a plain tensor (single process) it quantizes all ``num_experts`` -- identical to the in-model + path. Quantizing the fused ``[i]`` slice with its per-expert quantizer is byte-identical to + the split-then-quantize, so ``_export_fused_experts_keep_fused`` + :func:`split_fused_experts_state_dict` + reproduce :func:`_export_fused_experts`'s output exactly. Registers, per quantized projection ``P``: + ``module.

`` (quantized fused weight), ``module.

_weight_scale`` (stacked per-expert scales), + ``module.

_weight_scale_2`` (NVFP4 per-expert scalar), ``module.

_input_scale`` (shared). + """ + from modelopt.torch.export.unified_export_hf import _export_quantized_weight + + first = getattr(module, "_first_proj_attr", "gate_up_proj") + n = module.num_experts + # One fused tensor cannot hold quantized and unquantized experts side by side, so a projection + # is folded only when all its experts are quantized; one with none quantized stays as it is, + # which is what the per-expert export writes for it. Every refusal comes before either is + # folded, so it leaves the module untouched. + folded = [] + for proj in (first, "down_proj"): + weight_quantizers = getattr(module, f"{proj}_weight_quantizers") + enabled = [q.is_enabled for q in weight_quantizers] + wdata = getattr(module, proj).data + if wdata.element_size() == 1: + # get_expert_weight dequantizes these with their scale_inv, which the fold drops below + # whether it folds the projection or leaves it as it is. + raise NotImplementedError( + f"{type(module).__name__}.{proj}: native FP8 expert weights cannot be exported as " + "stored, since the fold drops their scale_inv; dequantize them first " + "(unpack_weight())." + ) + is_dt = isinstance(wdata, DTensor) + if is_dt: + _refuse_unfoldable_shards(module, proj, wdata, len(weight_quantizers)) + if len(weight_quantizers) != wdata.shape[0]: + # Expert parallelism leaves each rank only its own experts' quantizers, so the ranks + # decide on the projection together: alone, one rank could fold it while another + # leaves it as it is or refuses it. Joined by offset rather than by group rank. + mdim = next(i for i, p in enumerate(wdata.placements) if isinstance(p, Shard)) + gathered: list = [None] * wdata.device_mesh.size(mdim) + torch.distributed.all_gather_object( + gathered, + (_dtensor_dim0_offset(wdata), enabled), + group=wdata.device_mesh.get_group(mdim), + ) + enabled = [e for _, shard in sorted(gathered) for e in shard] + if any(enabled) and any( + getattr(q, "num_bits", None) in GGML_FORMATS for q in weight_quantizers + ): + # A GGML format packs each expert's weight into its blocks and registers no weight_scale + # for the fold to stack. Decided on the gathered states and on the format, which a recipe + # sets by module name and so is the same for every expert on every rank: every rank + # refuses alike. + raise NotImplementedError( + f"{type(module).__name__}.{proj}: GGML formats pack each expert's weight with no " + "weight_scale, which the keep-fused fold cannot stack; only the per-expert export " + "writes them." + ) + if all(enabled): + if not is_dt and not len(weight_quantizers) == wdata.shape[0] == n: + raise NotImplementedError( + f"{type(module).__name__}.{proj}: num_experts is {n}, but the unsharded weight " + f"holds {wdata.shape[0]} experts and there are {len(weight_quantizers)} weight " + "quantizers; the keep-fused fold needs the three equal." + ) + folded.append(proj) + elif any(enabled): + raise NotImplementedError( + f"{type(module).__name__}.{proj}: weight quantizers " + f"{[i for i, e in enumerate(enabled) if not e]} of {len(enabled)} are disabled and " + "the rest are not; a fused expert weight cannot mix quantized and unquantized experts." + ) + for proj in folded: + weight = getattr(module, proj) + weight_quantizers = getattr(module, f"{proj}_weight_quantizers") + input_quantizer = getattr(module, f"{proj}_input_quantizer", None) + + # A fused 3-D weight sharded on dim 0 (the expert axis) is a DTensor: quantize ONLY this + # rank's local experts and rewrap the result sharded identically -- indexing the DTensor + # (weight[i]) would all-gather the whole fused weight onto every rank, de-sharding the model + # and OOMing the GPU at scale. A plain tensor (a single process, or FSDP2's unsharded weight + # inside a weight window): all experts. + wdata = weight.data + is_dt = isinstance(wdata, DTensor) + quantizers_are_global = True + if is_dt: + mesh, placements = wdata.device_mesh, wdata.placements + local_w = wdata.to_local() + offset = _dtensor_dim0_offset(wdata) + scale_placements = [ + p if (isinstance(p, Shard) and p.dim == 0) else Replicate() for p in placements + ] + local_indices = list(range(offset, offset + local_w.shape[0])) + # Every expert's quantizers (FSDP2), or only this rank's, in shard order (expert + # parallelism): _refuse_unfoldable_shards allows no other count. + quantizers_are_global = len(weight_quantizers) == wdata.shape[0] + else: + local_w = wdata + local_indices = list(range(n)) + + local_fp8: torch.Tensor | None = None + local_ws: torch.Tensor | None = None + weight_scale_2s: list[torch.Tensor] = [] + input_scale: torch.Tensor | None = None + for local_idx, global_idx in enumerate(local_indices): + w_quantizer = weight_quantizers[global_idx if quantizers_are_global else local_idx] + w_slice = _local_expert_weight(module, proj, local_w, local_idx, dtype) + # Uncalibrated-expert fallback: derive amax from this expert's local weight slice + # (matches _export_fused_experts), so a never-routed expert still exports sane scales. + # Per stage, in the shape each expects (W4A8's SequentialQuantizer, block-wise INT4). + uncalibrated = [q for q in _tensor_quantizers(w_quantizer) if _is_uncalibrated(q)] + for q in uncalibrated: + q.amax = _fallback_amax(w_slice, q) + if uncalibrated: + warnings.warn( + f"Expert {global_idx} {proj} weight quantizer was not calibrated (amax missing " + f"or zero); using weight-derived amax. Increase calibration size to activate all " + f"experts.", + stacklevel=2, + ) + wrapper = nn.Module() + wrapper.weight = nn.Parameter(w_slice.contiguous(), requires_grad=False) + wrapper.weight_quantizer = w_quantizer + if input_quantizer is not None: + wrapper.input_quantizer = input_quantizer + _export_quantized_weight(wrapper, dtype) + rows = len(local_indices) + packed = wrapper.weight.data + local_fp8 = _stack_row(local_fp8, local_idx, packed, rows, packed.device) + # On the weight's device: per-tensor FP8 builds each expert's scale from a Python float, + # on the CPU, and a DTensor rewrap would label a CPU shard with the weight's device mesh. + local_ws = _stack_row(local_ws, local_idx, wrapper.weight_scale, rows, local_w.device) + # NVFP4 carries a SECOND per-tensor weight scale (weight_scale_2) that dequantizes the + # per-block weight_scale; FP8 has none. Keep it so the fused buffer (and the per-expert + # split) stay dequantizable. + if hasattr(wrapper, "weight_scale_2"): + weight_scale_2s.append(wrapper.weight_scale_2) + if hasattr(wrapper, "input_scale"): + input_scale = wrapper.input_scale + # Every rank holds an expert (_refuse_unfoldable_shards), so both stacks have a row. + assert local_fp8 is not None and local_ws is not None + if is_dt: + # Keep the global expert count: from_local would infer it as local x world, which is + # wrong for an uneven shard (3 experts on 2 ranks) and misnumbers the experts. + local_fp8 = _rewrap_expert_rows(local_fp8, mesh, placements, wdata.shape[0]) + local_ws = _rewrap_expert_rows(local_ws, mesh, scale_placements, wdata.shape[0]) + setattr(module, proj, nn.Parameter(local_fp8, requires_grad=False)) + module.register_buffer(f"{proj}_weight_scale", local_ws) + if weight_scale_2s: + # Per-expert scalar (E,) -- shards on the expert axis exactly like the weight scale. + local_ws2 = torch.stack(weight_scale_2s, dim=0).to(local_w.device) + if is_dt: + local_ws2 = _rewrap_expert_rows(local_ws2, mesh, scale_placements, wdata.shape[0]) + module.register_buffer(f"{proj}_weight_scale_2", local_ws2) + if input_scale is not None: + # Replicated, not sharded like the rest: calibration syncs the shared input quantizer's + # amax over the data-parallel ranks and, under expert parallelism, over the + # expert-parallel group (max_calibrate), so every rank registers the same value. + module.register_buffer(f"{proj}_input_scale", input_scale) + + # A projection left as it is goes out in the weight domain the per-expert export writes it in. + for proj in (first, "down_proj"): + if proj in folded: + continue + wdata = getattr(module, proj).data + is_dt = isinstance(wdata, DTensor) + local = wdata.to_local() if is_dt else wdata + domain = _local_expert_weight(module, proj, local, 0, dtype).dtype + if domain != local.dtype: + local = local.to(domain) + if is_dt: + local = _rewrap_expert_rows( + local, wdata.device_mesh, wdata.placements, wdata.shape[0] + ) + setattr(module, proj, nn.Parameter(local, requires_grad=False)) + + # Drop the quantizer modules -- their info now lives in the fused weight + scale buffers -- and + # the source checkpoint's scales, which described the weights the fold replaced, as the + # per-expert export does (_delete_fused_moe_source_attrs). + for attr in ( + f"{first}_weight_quantizers", + f"{first}_input_quantizer", + "down_proj_weight_quantizers", + "down_proj_input_quantizer", + f"{first}_scale_inv", + "down_proj_scale_inv", + f"{first}_activation_scale", + "down_proj_activation_scale", + ): + if hasattr(module, attr): + delattr(module, attr) + + def _export_fused_experts( module: nn.Module, dtype: torch.dtype, @@ -347,3 +639,165 @@ def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | Non output_path = Path(output_dir) / ".moe.html" output_path.write_text(html_content, encoding="utf-8") print(f"\033[1mExpert token count table saved to {output_path}\033[0m") + + +# ----- Distributed-export helpers (fused-experts split for the no-gather FSDP2 write) ----- +_FUSED_PROJ = ("gate_up_proj", "up_proj", "down_proj") + + +def _fused_experts_prefixes(state_dict: dict) -> list[str]: + """Return the prefixes of the modules holding a fused 3-D expert weight. + + Return the sorted set of ``...experts`` prefixes that carry a fused 3-D expert weight + (``.gate_up_proj`` / ``up_proj`` / ``down_proj``) in ``state_dict``. + """ + prefixes: set[str] = set() + for key, tensor in state_dict.items(): + for proj in _FUSED_PROJ: + if key.endswith(f".experts.{proj}") and hasattr(tensor, "dim") and tensor.dim() == 3: + prefixes.add(key[: -(len(proj) + 1)]) + break + return sorted(prefixes) + + +def split_fused_experts_state_dict(state_dict: dict) -> dict: + """Rewrite the fused-experts entries of a state dict into per-expert entries. + + The distributed export writes experts *fused* (3-D ``experts.`` weights + fused scale + buffers, kept sharded for DCP). This post-pass splits them into the same per-expert layout the + in-model :func:`_export_fused_experts` produces, so the two export paths yield identical + checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no + quantization. Non-fused-experts keys pass through unchanged. Every emitted tensor is a copy and + the input keeps its fused tensors, so the experts are held twice until it returns: pass one + module's tensors at a time, not a whole model's. Each module's experts are numbered from 0 by + their row in the tensors given, which must be plain: whole ones get their global numbers, and a + caller splitting a rank's local shard (``to_local()``) shifts each by the shard's global offset. + + Per fused-experts prefix ``P`` (``...experts``), inputs are (gated shown; ungated drops the gate split): + ``P.gate_up_proj`` ``(E, 2*I, H)`` + ``P.gate_up_proj_weight_scale`` + ``P.gate_up_proj_input_scale`` + ``P.down_proj`` ``(E, O, I)`` + ``P.down_proj_weight_scale`` + ``P.down_proj_input_scale`` + Outputs, for each expert ``i``: + ``P.{i}.gate_proj.weight`` ``(I, H)`` (+ ``.weight_scale``/``.input_scale``) [gated only] + ``P.{i}.up_proj.weight`` ``(I, H)`` (+ ``.weight_scale``/``.input_scale``) + ``P.{i}.down_proj.weight`` ``(O, I)`` (+ ``.weight_scale``/``.input_scale``) + """ + prefixes = _fused_experts_prefixes(state_dict) + if not prefixes: + return state_dict + + # Keys consumed (and thus removed) while emitting per-expert keys. + consumed: set[str] = set() + out: dict = {} + + for prefix in prefixes: + gated = f"{prefix}.gate_up_proj" in state_dict + first = "gate_up_proj" if gated else "up_proj" + first_w = state_dict[f"{prefix}.{first}"] + down_w = state_dict[f"{prefix}.down_proj"] + n_experts = first_w.shape[0] + if down_w.shape[0] != n_experts: + raise ValueError(f"{prefix}: {n_experts} {first} experts, {down_w.shape[0]} down_proj.") + + def _scales(proj_name): + ws = state_dict.get(f"{prefix}.{proj_name}_weight_scale") + ins = state_dict.get(f"{prefix}.{proj_name}_input_scale") + # NVFP4 second-level (per-tensor) weight scale; None for FP8. (E,) one scalar per expert. + ws2 = state_dict.get(f"{prefix}.{proj_name}_weight_scale_2") + return ws, ins, ws2 + + first_ws, first_in, first_ws2 = _scales(first) + down_ws, down_in, down_ws2 = _scales("down_proj") + + def _slice_weight_scale(ws, expert_idx, row_slice, rows): + # A per-expert scalar as is; else the expert's scale rows covering ``row_slice``, a cut of + # the fused gate|up's ``rows`` rows, in proportion (only gate and up are cut): an INT4 + # weight packs two rows a byte while its scales stay per row, an FP8 2-D block 128 rows. + if ws is None: + return None + if ws.dim() <= 1: # (E,) one scalar per expert + return ws[expert_idx] + sub = ws[expert_idx] + if row_slice is None: + return sub + start, stop, _ = row_slice.indices(rows) + n = sub.shape[0] + # Exact only for scale rows spread evenly over the weight rows, with the cut between + # two of them; otherwise a block of rows sharing one scale straddles gate and up. + if (rows % n and n % rows) or start * n % rows or stop * n % rows: + raise NotImplementedError( + f"A {n}-row weight scale over {rows} weight rows cannot be split at the gate|up " + f"boundary, row {start or stop}: a block of rows sharing one scale crosses it." + ) + return sub[start * n // rows : stop * n // rows] + + for expert_idx in range(n_experts): + if gated: + inter = first_w.shape[1] // 2 + # gate and up are halves of the fused gate_up: they share its single per-tensor + # weight_scale_2 (first_ws2), matching the in-model export + vLLM's W1/W3 fusion. + projections = [ + ( + "gate_proj", + first_w[expert_idx, :inter, :], + first_ws, + first_in, + first_ws2, + slice(0, inter), + ), + ( + "up_proj", + first_w[expert_idx, inter:, :], + first_ws, + first_in, + first_ws2, + slice(inter, None), + ), + ("down_proj", down_w[expert_idx], down_ws, down_in, down_ws2, None), + ] + else: + projections = [ + ("up_proj", first_w[expert_idx], first_ws, first_in, first_ws2, None), + ("down_proj", down_w[expert_idx], down_ws, down_in, down_ws2, None), + ] + for proj_name, weight, ws, ins, ws2, row_slice in projections: + base = f"{prefix}.{expert_idx}.{proj_name}" + # .clone() every emitted tensor so no two keys share storage. gate/up share the + # per-tensor scale object (and slices alias the fused parent); without cloning, + # safetensors/save_pretrained shared-tensor dedup silently drops the duplicate + # (e.g. up_proj.weight_scale/input_scale would go missing). + out[f"{base}.weight"] = weight.detach().clone().contiguous() + sliced_ws = _slice_weight_scale(ws, expert_idx, row_slice, first_w.shape[1]) + if sliced_ws is not None: + out[f"{base}.weight_scale"] = sliced_ws.detach().clone().contiguous() + # weight_scale_2 is a per-expert scalar (shared by gate/up): _slice_weight_scale's + # dim<=1 branch returns ws2[expert_idx] regardless of row_slice. NVFP4 only; None for FP8. + sliced_ws2 = _slice_weight_scale(ws2, expert_idx, row_slice, first_w.shape[1]) + if sliced_ws2 is not None: + out[f"{base}.weight_scale_2"] = sliced_ws2.detach().clone().contiguous() + if ins is not None: + out[f"{base}.input_scale"] = ins.detach().clone() + + # Mark the fused tensors + their scale buffers consumed. + for proj in (first, "down_proj"): + for suffix in ("", "_weight_scale", "_weight_scale_2", "_input_scale"): + consumed.add(f"{prefix}.{proj}{suffix}") + + out.update({key: tensor for key, tensor in state_dict.items() if key not in consumed}) + return out + + +def _dtensor_dim0_offset(dt) -> int: + """Global dim-0 start index of this rank's local shard of a ``Shard(0)`` DTensor.""" + mesh, placements = dt.device_mesh, dt.placements + try: + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + + _, global_offset = compute_local_shape_and_global_offset(tuple(dt.shape), mesh, placements) + return int(global_offset[0]) + except Exception: + n = dt.shape[0] + for mdim, p in enumerate(placements): + if isinstance(p, Shard) and p.dim == 0: + nshards = mesh.size(mdim) + return mesh.get_local_rank(mdim) * ((n + nshards - 1) // nshards) + return 0 diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index a0ad2c1deeb..29b517a4a66 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -66,7 +66,9 @@ class ExportContext: keep_fused_experts: bool = False """Marks the distributed no-gather export, whose handlers pack while the model is still sharded: - an FSDP2-owned weight is packed inside its unit's weight-update window.""" + an FSDP2-owned weight is packed inside its unit's weight-update window, and fused MoE experts are + folded in place and kept fused and sharded rather than split into whole per-expert tensors, + which :func:`~modelopt.torch.export.moe_utils.split_fused_experts_state_dict` does later.""" module_names: "ModuleNames | None" = None """``model``'s name maps, built once per export by the first handler that looks a module up diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index e6d75612869..7c57c3838f5 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -638,7 +638,7 @@ def _export_quantized_weight( An export handler packs through ``hf_export_handlers._export_weight`` rather than calling this directly: on the keep-fused export an FSDP2-owned weight must be packed inside its unit's - weight-update window, which that helper opens. + weight-update window, which that helper (or the fused-experts handler, around its fold) opens. Packed outside it, the unit's next reshard puts FSDP2's unquantized shard back, and the checkpoint is written unquantized. diff --git a/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py b/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py new file mode 100644 index 00000000000..83e2e974f3f --- /dev/null +++ b/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The keep-fused fold of sharded fused experts on every visible GPU, over NCCL.""" + +import copy +import re +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import ( + _dtensor_dim0_offset, + _export_fused_experts, + _export_fused_experts_keep_fused, + split_fused_experts_state_dict, +) + +HIDDEN_DIM = 32 +INTERMEDIATE_DIM = 16 +PREFIX = "model.layers.0.mlp.experts" + + +class _FusedExperts(nn.Module): + """Fused gate+up experts in the transformers 5.x layout (Qwen3-MoE / Mixtral).""" + + def __init__(self, num_experts): + super().__init__() + self.num_experts = num_experts + self.gate_up_proj = nn.Parameter( + torch.randn(num_experts, 2 * INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02 + ) + self.down_proj = nn.Parameter(torch.randn(num_experts, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02) + self.act_fn = nn.SiLU() + + def forward(self, hidden_states, top_k_index, top_k_weights): + out = torch.zeros_like(hidden_states) + mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) + for expert_idx in torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero()[:, 0]: + top_k_pos, token_idx = torch.where(mask[expert_idx]) + gate, up = F.linear(hidden_states[token_idx], self.gate_up_proj[expert_idx]).chunk( + 2, dim=-1 + ) + current = F.linear(self.act_fn(gate) * up, self.down_proj[expert_idx]) + out.index_add_(0, token_idx, current * top_k_weights[token_idx, top_k_pos, None]) + return out + + +def _quantized_experts(quant_cfg, num_experts): + """Fused experts calibrated on one token each, moved to this rank's GPU to be folded. + + Every rank builds and calibrates the same seeded module, so all hold identical quantizers. + """ + torch.manual_seed(0) + model = nn.Module() + model.experts = _FusedExperts(num_experts) + hidden = torch.randn(num_experts, HIDDEN_DIM) + tokens = (torch.arange(num_experts).unsqueeze(1), torch.ones(num_experts, 1)) + mtq.quantize(model, quant_cfg, lambda m: m.experts(hidden, *tokens)) + return model.experts.cuda() + + +def _state(experts): + return {f"{PREFIX}.{key}": value for key, value in experts.state_dict().items()} + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu() + + +def _whole(value): + """A sharded value gathered whole, one-byte values moved as bytes.""" + if not isinstance(value, DTensor): + return value + if value.element_size() != 1: + return value.full_tensor() + raw = DTensor.from_local( + value.to_local().view(torch.uint8), + value.device_mesh, + value.placements, + run_check=False, + shape=value.shape, + stride=value.stride(), + ) + return raw.full_tensor().view(value.dtype) + + +def _sharded(experts, mesh): + experts = copy.deepcopy(experts) + for proj in ("gate_up_proj", "down_proj"): + weight = distribute_tensor(getattr(experts, proj).data, mesh, [Shard(0)]) + setattr(experts, proj, nn.Parameter(weight, requires_grad=False)) + return experts + + +def _fold_job(rank, world_size, backend): + # Every collective below runs on this group, which gives up after a minute. + group = dist.new_group(backend=backend, timeout=timedelta(seconds=60)) + try: + # A mesh compares equal to any other over the same ranks and names, whatever its group, so + # DTensor's caches could hand a later job a mesh over this group once it is destroyed. + mesh = DeviceMesh.from_group(group, "cuda", mesh_dim_names=(f"ep{group.group_name}",)) + # Every rank keeps every expert's quantizers, as under FSDP2. Two experts a rank, then one + # expert fewer: the last rank holds one, and the folds must keep the global expert count. + cases = [(cfg, 2 * world_size) for cfg in (mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG)] + cases.append((mtq.NVFP4_DEFAULT_CFG, 2 * world_size - 1)) + for quant_cfg, num_experts in cases: + experts = _quantized_experts(quant_cfg, num_experts) + reference = copy.deepcopy(experts) + _export_fused_experts(reference, torch.bfloat16) + plain = copy.deepcopy(experts) + _export_fused_experts_keep_fused(plain, torch.bfloat16) + plain = _state(plain) + sharded = _sharded(experts, mesh) + + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + + fused = _state(sharded) + assert set(fused) == set(plain) + offset = _dtensor_dim0_offset(fused[f"{PREFIX}.gate_up_proj"]) + for key, value in fused.items(): + assert value.device.type == "cuda", key + if key.endswith("_input_scale"): # one scalar shared by every expert + assert not isinstance(value, DTensor), key + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + continue + assert isinstance(value, DTensor) and value.placements == (Shard(0),), key + assert value.shape == plain[key].shape, key + rows = plain[key][offset : offset + value.to_local().shape[0]] + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(rows)), key + # Gathered whole, the ranks' folds split into exactly what the per-expert export writes. + split = split_fused_experts_state_dict({k: _whole(v) for k, v in fused.items()}) + want = _state(reference) + assert set(split) == set(want) + for key, value in want.items(): + assert split[key].dtype == value.dtype and split[key].shape == value.shape, key + assert torch.equal(_raw_bytes(split[key]), _raw_bytes(value)), key + + # A layout that leaves a rank no expert is refused on every rank. + lone = _sharded(_quantized_experts(mtq.FP8_DEFAULT_CFG, world_size - 1), mesh) + with pytest.raises(NotImplementedError, match="leaves a rank none"): + _export_fused_experts_keep_fused(lone, torch.bfloat16) + + # Under expert parallelism a rank keeps only its own experts' quantizers, so the ranks gather + # each projection's quantizer states over the group before they fold it. + experts = _quantized_experts(mtq.NVFP4_DEFAULT_CFG, 2 * world_size) + plain = copy.deepcopy(experts) + _export_fused_experts_keep_fused(plain, torch.bfloat16) + plain = _state(plain) + mine = slice(2 * rank, 2 * rank + 2) + + def keeping_own_quantizers(): + sharded = _sharded(experts, mesh) + for proj in ("gate_up_proj", "down_proj"): + name = f"{proj}_weight_quantizers" + setattr(sharded, name, nn.ModuleList(list(getattr(sharded, name))[mine])) + return sharded + + sharded = keeping_own_quantizers() + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + assert set(_state(sharded)) == set(plain) + for key, value in _state(sharded).items(): + local = value.to_local() if isinstance(value, DTensor) else value + want = plain[key][mine] if isinstance(value, DTensor) else plain[key] + assert torch.equal(_raw_bytes(local), _raw_bytes(want)), key + + # The last rank disabling one or both of its experts' quantizers is refused on every rank, + # naming the global experts; every expert disabled leaves the projection as it is on all. + last = 2 * world_size - 1 + for disabled in ([last], [last - 1, last], list(range(2 * world_size))): + sharded = keeping_own_quantizers() + for expert, quantizer in enumerate(sharded.down_proj_weight_quantizers, mine.start): + if expert in disabled: + quantizer.disable() + if len(disabled) < 2 * world_size: + match = re.escape(f"quantizers {disabled} of {2 * world_size}") + with pytest.raises(NotImplementedError, match=match): + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + else: + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + assert hasattr(sharded, "gate_up_proj_weight_scale") + assert not hasattr(sharded, "down_proj_weight_scale") + finally: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("backend", [None, "nccl"], ids=["objects-over-gloo", "objects-over-nccl"]) +def test_ranks_fold_their_own_experts_over_nccl(need_2_gpus, dist_workers, backend): + """Each rank folds the experts it holds on its GPU, at their global offsets, and the folds, + gathered over NCCL, split into the per-expert export, an uneven shard included. Under expert + parallelism the ranks decide on each projection together: in the export's ``cpu:gloo,cuda:nccl`` + group their quantizer states go over Gloo, in a group of NCCL alone over NCCL.""" + dist_workers.run(_fold_job, backend=backend) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py new file mode 100644 index 00000000000..0a516db55aa --- /dev/null +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -0,0 +1,800 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Keep-fused MoE expert export used by the no-gather FSDP2 writer (CPU, no GPU needed). + +The no-gather writer never splits a fused 3-D expert weight while the model is sharded: +``_export_fused_experts_keep_fused`` folds every expert's quantizer into the fused weight in place, +and ``split_fused_experts_state_dict`` cuts the result into per-expert keys later. That pipeline +must write exactly what the single-process ``_export_fused_experts`` writes, so most tests compare +the two byte for byte on a tiny synthetic fused-experts module quantized with FP8 and NVFP4. +Sharded cases run on a gloo process group: one rank in process, and two spawned ranks where a +non-zero expert offset matters. +""" + +import contextlib +import copy +import math +import re +import warnings + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +pytest.importorskip("transformers") +pytest.importorskip("safetensors") +from safetensors.torch import load_file, save_file +from transformers import NemotronHConfig, PretrainedConfig +from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts + +import modelopt.torch.export.hf_export_handlers as handlers +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import ( + _dtensor_dim0_offset, + _export_fused_experts, + _export_fused_experts_keep_fused, + _fused_experts_prefixes, + split_fused_experts_state_dict, +) +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY +from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan + +NUM_EXPERTS = 4 +HIDDEN_DIM = 32 +INTERMEDIATE_DIM = 16 +PREFIX = "model.layers.0.mlp.experts" + +FP8_WEIGHT_ONLY_CFG = copy.deepcopy(mtq.FP8_DEFAULT_CFG) +FP8_WEIGHT_ONLY_CFG["quant_cfg"].append({"quantizer_name": "*input_quantizer", "enable": False}) + + +class _KeepFusedGatedExperts(nn.Module): + """Fused gate+up experts in the transformers 5.x layout (Qwen3-MoE / Mixtral).""" + + def __init__(self, num_experts=NUM_EXPERTS): + super().__init__() + self.num_experts = num_experts + self.gate_up_proj = nn.Parameter( + torch.randn(num_experts, 2 * INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02 + ) + self.down_proj = nn.Parameter(torch.randn(num_experts, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02) + self.act_fn = nn.SiLU() + + def forward(self, hidden_states, top_k_index, top_k_weights): + out = torch.zeros_like(hidden_states) + mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) + for expert_idx in torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero()[:, 0]: + top_k_pos, token_idx = torch.where(mask[expert_idx]) + gate, up = F.linear(hidden_states[token_idx], self.gate_up_proj[expert_idx]).chunk( + 2, dim=-1 + ) + current = F.linear(self.act_fn(gate) * up, self.down_proj[expert_idx]) + out.index_add_(0, token_idx, current * top_k_weights[token_idx, top_k_pos, None]) + return out + + +class _KeepFusedUngatedExperts(NemotronHExperts): + """Nemotron-H's ungated fused experts: a single ``up_proj`` and ``down_proj``. + + The real class, on its eager forward, so that whichever wrapper ModelOpt quantizes Nemotron-H's + experts with converts it. + """ + + def __init__(self, num_experts=NUM_EXPERTS): + config = NemotronHConfig( + n_routed_experts=num_experts, + hidden_size=HIDDEN_DIM, + moe_intermediate_size=INTERMEDIATE_DIM, + ) + config._experts_implementation = "eager" + super().__init__(config) + with torch.no_grad(): + self.up_proj.normal_(std=0.02) + self.down_proj.normal_(std=0.02) + + +LAYOUTS = pytest.mark.parametrize( + "experts_cls", + [_KeepFusedGatedExperts, _KeepFusedUngatedExperts], + ids=["gated", "ungated"], +) + + +@pytest.fixture(autouse=True) +def _unregister_synthetic_experts(): + """``mtq.quantize`` registers the synthetic classes on the fly; keep the registry clean. + + A class the registry already covers, through its base class's registration, stays as it is. + """ + on_the_fly = [ + cls + for cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts) + if cls not in QuantModuleRegistry + ] + yield + for cls in on_the_fly: + if QuantModuleRegistry.get(cls) is not None: + QuantModuleRegistry.unregister(cls) + + +def _quantized_experts(experts_cls, quant_cfg, num_experts=NUM_EXPERTS): + """Build, quantize and calibrate a fused-experts module; every expert sees one token.""" + torch.manual_seed(0) + model = nn.Module() + model.experts = experts_cls(num_experts) + + def forward_loop(m): + hidden = torch.randn(num_experts, HIDDEN_DIM) + m.experts(hidden, torch.arange(num_experts).unsqueeze(1), torch.ones(num_experts, 1)) + + mtq.quantize(model, quant_cfg, forward_loop) + return model.experts + + +def _prefixed(state_dict): + return {f"{PREFIX}.{key}": value for key, value in state_dict.items()} + + +def _per_expert_export(experts): + """What the single-process export writes: ``_export_fused_experts`` on a copy.""" + module = copy.deepcopy(experts) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _export_fused_experts(module, torch.bfloat16) + return _prefixed(module.state_dict()) + + +def _keep_fused_export(experts): + """The fused state ``_export_fused_experts_keep_fused`` leaves on a copy.""" + module = copy.deepcopy(experts) + _export_fused_experts_keep_fused(module, torch.bfloat16) + return _prefixed(module.state_dict()) + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8) + + +def _assert_byte_identical(got, want): + assert set(got) == set(want), (sorted(set(want) - set(got)), sorted(set(got) - set(want))) + for key, ref in want.items(): + assert got[key].dtype == ref.dtype, key + assert got[key].shape == ref.shape, key + assert torch.equal(_raw_bytes(got[key]), _raw_bytes(ref)), key + + +def _assert_quantized(state_dict): + """Guard against a config that quantized nothing, which would make any comparison trivial.""" + weights = [v for k, v in state_dict.items() if k.endswith(".weight")] + assert weights, "no per-expert weights" + assert all(v.dtype in (torch.float8_e4m3fn, torch.uint8) for v in weights) + + +# ------------------------------------------------------------------ keep-fused + split + + +def _uncalibrated(experts, projections): + """Expert 1's calibration never ran: zero its weight amax.""" + for proj in projections: + getattr(experts, f"{proj}_weight_quantizers")[1]._amax.zero_() + + +def _unquantized(experts, projections): + """A recipe left these projections' experts unquantized.""" + for proj in projections: + for quantizer in getattr(experts, f"{proj}_weight_quantizers"): + quantizer.disable() + + +def _with_source_scales(experts, projections): + """A native-FP8 checkpoint's experts keep their source block and activation scales.""" + for proj in projections: + experts.register_buffer(f"{proj}_scale_inv", torch.ones(NUM_EXPERTS, 1, 1)) + experts.register_buffer(f"{proj}_activation_scale", torch.ones(())) + + +def _export_cases(): + layouts = [(_KeepFusedGatedExperts, "gated"), (_KeepFusedUngatedExperts, "ungated")] + cfgs = [ + (mtq.FP8_DEFAULT_CFG, "fp8"), + (FP8_WEIGHT_ONLY_CFG, "fp8_weight_only"), + (mtq.NVFP4_DEFAULT_CFG, "nvfp4"), + ] + cases = [ + pytest.param(cls, cfg, None, (), id=f"{name}-{cfg_name}") + for cls, name in layouts + for cfg, cfg_name in cfgs + ] + cases += [ + pytest.param( + cls, + mtq.NVFP4_DEFAULT_CFG, + _uncalibrated, + ("first", "down_proj"), + id=f"{name}-uncalibrated", + ) + for cls, name in layouts + ] + cases += [ + pytest.param( + cls, + mtq.FP8_DEFAULT_CFG, + _with_source_scales, + ("first", "down_proj"), + id=f"{name}-source-scales", + ) + for cls, name in layouts + ] + cases += [ + pytest.param( + cls, mtq.FP8_DEFAULT_CFG, _unquantized, projs, id=f"{name}-unquantized-{which}" + ) + for cls, name in layouts + for which, projs in ( + ("first", ("first",)), + ("down", ("down_proj",)), + ("both", ("first", "down_proj")), + ) + ] + return cases + + +@pytest.mark.parametrize(("experts_cls", "quant_cfg", "change", "projections"), _export_cases()) +def test_keep_fused_then_split_writes_exactly_the_per_expert_export( + experts_cls, quant_cfg, change, projections +): + """Also when an expert calibration never reached falls back to an amax read off its own weight. + + For the gated layout that amax must come from the whole fused ``gate_up`` slice, so gate and up + keep one shared ``weight_scale_2``, as the per-expert export guarantees and vLLM's W1/W3 fusion + needs. And when a recipe leaves a projection's experts unquantized: it has nothing to fold. + """ + experts = _quantized_experts(experts_cls, quant_cfg) + if change is not None: + first = experts._first_proj_attr + change(experts, [first if proj == "first" else proj for proj in projections]) + want = _per_expert_export(experts) + if change is None: + _assert_quantized(want) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # an uncalibrated expert warns + got = split_fused_experts_state_dict(_keep_fused_export(experts)) + + _assert_byte_identical(got, want) + + +@pytest.mark.parametrize("unquantized", [(), ("down_proj",)], ids=["quantized", "down-unquantized"]) +def test_keep_fused_then_split_writes_unpacked_native_fp8_experts_as_the_per_expert_export( + unquantized, +): + """Native FP8 experts unpack at the default dtype, and the per-expert export reads each expert in + the export dtype (``get_expert_weight``): the fold reads them alike, for the scales it derives + and for a projection it leaves unquantized.""" + # transformers 5.5 imports triton in this module, which not every platform has. + fp8_experts_cls = pytest.importorskip("transformers.integrations.finegrained_fp8").FP8Experts + + torch.manual_seed(0) + config = PretrainedConfig( + hidden_size=HIDDEN_DIM, + intermediate_size=INTERMEDIATE_DIM, + num_experts=NUM_EXPERTS, + hidden_act="silu", + ) + experts = fp8_experts_cls(config, block_size=None) + for name in ("gate_up_proj", "down_proj"): + weight = getattr(experts, name) + with torch.no_grad(): + weight.copy_((torch.randn(weight.shape) * 4).to(weight.dtype)) + getattr(experts, f"{name}_scale_inv").fill_(0.01) + hidden = torch.randn(NUM_EXPERTS, HIDDEN_DIM) + route = (torch.arange(NUM_EXPERTS).unsqueeze(1), torch.ones(NUM_EXPERTS, 1)) + registered_before = QuantModuleRegistry.get(fp8_experts_cls) + try: + mtq.quantize(experts, mtq.NVFP4_DEFAULT_CFG, lambda m: m(hidden, *route)) + experts.unpack_weight() + _unquantized(experts, unquantized) + want = _per_expert_export(experts) + got = split_fused_experts_state_dict(_keep_fused_export(experts)) + finally: + if registered_before is None and QuantModuleRegistry.get(fp8_experts_cls) is not None: + QuantModuleRegistry.unregister(fp8_experts_cls) + _assert_byte_identical(got, want) + + +@pytest.mark.parametrize( + "quant_cfg", + [ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.W4A8_AWQ_BETA_CFG, + mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, + ], + ids=["int4-blockwise", "w4a8-awq", "fp8-2d-blockwise"], +) +def test_the_split_cuts_packed_and_block_scales_with_their_rows(monkeypatch, quant_cfg): + """An INT4 weight packs two rows a byte while its scales stay per row, and an FP8 2-D block + scale covers 128 rows: gate and up each take the scale rows of their own weight rows. Their + blocks need dimensions of whole 128-blocks.""" + monkeypatch.setitem(globals(), "HIDDEN_DIM", 128) + monkeypatch.setitem(globals(), "INTERMEDIATE_DIM", 128) + experts = _quantized_experts(_KeepFusedGatedExperts, quant_cfg) + want = _per_expert_export(experts) + + got = split_fused_experts_state_dict(_keep_fused_export(experts)) + + _assert_byte_identical(got, want) + + +@pytest.mark.parametrize("amax_state", ["calibrated", "zero"]) +def test_the_fold_derives_a_w4a8_experts_amax_per_quantizer_stage(monkeypatch, amax_state): + """W4A8's SequentialQuantizer has no amax of its own, and its INT4 block stage rejects a scalar: + an expert calibration never reached gets a per-block amax there and a scalar on the FP8 stage.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + block = 16 + rows = {"gate_up_proj": 2 * INTERMEDIATE_DIM, "down_proj": HIDDEN_DIM} + blocks = {"gate_up_proj": HIDDEN_DIM // block, "down_proj": INTERMEDIATE_DIM // block} + for proj in rows: + for i in range(NUM_EXPERTS): + int4 = TensorQuantizer( + QuantizerAttributeConfig(num_bits=4, block_sizes={-1: block, "type": "static"}) + ) + fp8 = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))) + fill = 0.0 if amax_state == "zero" else 1.0 + int4.amax = torch.full( + (rows[proj] * blocks[proj], 1), fill + ) # as block calibration records + fp8.amax = torch.tensor(fill) + getattr(experts, f"{proj}_weight_quantizers")[i] = SequentialQuantizer(int4, fp8) + + seen = [] + + def spy(wrapper, dtype, **_kwargs): + int4, fp8 = wrapper.weight_quantizer + seen.append((int4._amax.clone(), fp8._amax.clone())) + wrapper.weight_scale = torch.ones(1) + + monkeypatch.setattr("modelopt.torch.export.unified_export_hf._export_quantized_weight", spy) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + assert len(seen) == 2 * NUM_EXPERTS + assert {amax.shape for amax, _ in seen} == {(64, 1), (32, 1)} + assert all(torch.all(int4_amax > 0) and torch.all(fp8_amax > 0) for int4_amax, fp8_amax in seen) + assert any("was not calibrated" in str(w.message) for w in caught) == (amax_state == "zero") + + +def test_the_fold_puts_its_scales_on_the_weights_device(monkeypatch): + """Per-tensor FP8 makes each expert's scale from a Python float, on the CPU wherever the weight + is, and a sharded fold rewraps the scales on the weight's mesh. The weights here sit on meta, + the one other device a CPU runner has, so a stand-in packs them, its scales on the CPU.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + for proj in ("gate_up_proj", "down_proj"): + setattr(experts, proj, nn.Parameter(getattr(experts, proj).detach().to("meta"))) + + def pack(wrapper, dtype, **_kwargs): + wrapper.weight_scale = torch.tensor(1.0) + wrapper.weight_scale_2 = torch.tensor(1.0) + + monkeypatch.setattr("modelopt.torch.export.unified_export_hf._export_quantized_weight", pack) + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + for proj in ("gate_up_proj", "down_proj"): + for scale in ("weight_scale", "weight_scale_2"): + assert getattr(experts, f"{proj}_{scale}").is_meta, (proj, scale) + + +def _mixed(experts): + experts.down_proj_weight_quantizers[1].disable() + + +def _native_fp8(experts): + experts.down_proj = nn.Parameter( + experts.down_proj.data.to(torch.float8_e4m3fn), requires_grad=False + ) + + +def _native_fp8_unquantized(experts): + _native_fp8(experts) + _unquantized(experts, ["down_proj"]) + + +def _fewer_experts_counted(experts): + experts.num_experts = NUM_EXPERTS - 1 + + +def _ggml(experts): + for i in range(len(experts.down_proj_weight_quantizers)): + experts.down_proj_weight_quantizers[i] = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="q8_0", + block_sizes={-1: GGML_FORMAT_REGISTRY["q8_0"].block_size}, + backend="ggml", + ) + ) + + +@LAYOUTS +@pytest.mark.parametrize( + ("change", "match"), + [ + (_mixed, r"down_proj: weight quantizers \[1\] of 4"), + (_native_fp8, "native FP8"), + (_native_fp8_unquantized, "native FP8"), + (_fewer_experts_counted, "num_experts is 3, but the unsharded weight holds 4 experts"), + (_ggml, "GGML formats"), + ], + ids=["mixed", "native-fp8", "native-fp8-unquantized", "num-experts-short", "ggml"], +) +def test_a_projection_the_fold_cannot_write_is_refused_untouched(experts_cls, change, match): + """One fused tensor cannot mix quantized and unquantized experts, and native FP8 weights need the + scale_inv the fold drops, also where a recipe leaves them unquantized. A ``num_experts`` short of + the weight's rows would fold only the first experts, and a GGML format packs each expert's weight + with no scale to stack. The refusal comes before the other projection folds.""" + experts = _quantized_experts(experts_cls, mtq.FP8_DEFAULT_CFG) + change(experts) + before = {k: v.clone() for k, v in experts.state_dict().items()} + + with pytest.raises(NotImplementedError, match=match): + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + _assert_byte_identical(experts.state_dict(), before) + assert getattr(experts, f"{experts._first_proj_attr}_weight_quantizers")[0].is_enabled + + +# ------------------------------------------------------------------ split_fused_experts_state_dict + + +LOCAL_EXPERTS = 2 # experts in each hand-built fused entry below + + +def _fused_state_dict(*, gated, rows=4, cols=6): + """A fused-experts entry with position-coded values, so every slice is identifiable.""" + first = "gate_up_proj" if gated else "up_proj" + fused_rows = 2 * rows if gated else rows + n = LOCAL_EXPERTS + return { + f"{PREFIX}.{first}": torch.arange(n * fused_rows * cols, dtype=torch.float32) + .reshape(n, fused_rows, cols) + .to(torch.float8_e4m3fn), + # Per-row (per-channel) weight scale: its rows are split along with the weight's. + f"{PREFIX}.{first}_weight_scale": torch.arange(n * fused_rows).view(n, fused_rows) + 100.0, + # Per-expert scalars. + f"{PREFIX}.{first}_weight_scale_2": torch.arange(n) + 7.0, + f"{PREFIX}.{first}_input_scale": torch.tensor(0.5), + f"{PREFIX}.down_proj": torch.arange(n * cols * rows, dtype=torch.float32) + .reshape(n, cols, rows) + .to(torch.float8_e4m3fn), + f"{PREFIX}.down_proj_weight_scale": torch.arange(n) + 3.0, + f"{PREFIX}.down_proj_input_scale": torch.tensor(0.25), + } + + +def test_split_gives_every_tensor_its_own_storage(tmp_path): + """Slices alias the fused parent, and gate/up share one scale per expert. + + Aliasing keeps the whole fused tensor alive behind every per-expert key, and a shared scale is + refused by safetensors (``save_pretrained`` dedups it away instead). + """ + out = split_fused_experts_state_dict(_fused_state_dict(gated=True)) + + storages = [t.untyped_storage().data_ptr() for t in out.values()] + assert len(set(storages)) == len(storages) + path = tmp_path / "experts.safetensors" + save_file(out, str(path)) + reloaded = load_file(str(path)) + assert set(reloaded) == set(out) + for key, value in out.items(): + assert torch.equal(_raw_bytes(reloaded[key]), _raw_bytes(value)), key + + +def test_split_leaves_a_state_dict_without_fused_experts_unchanged(): + sd = { + "model.embed_tokens.weight": torch.ones(4, 2), + f"{PREFIX}.0.up_proj.weight": torch.ones(4, 2), # already per-expert + } + + out = split_fused_experts_state_dict(sd) + + assert out.keys() == sd.keys() + assert all(out[k] is sd[k] for k in sd) + + +@pytest.mark.parametrize( + ("rows", "scale_rows"), [(6, 3), (10, 4)], ids=["block-across-the-cut", "uneven-blocks"] +) +def test_split_refuses_a_scale_block_gate_and_up_would_share(rows, scale_rows): + """Two-row blocks over 6 rows put gate's last row and up's first under one scale; three-row + blocks over 10 rows end in a partial block, which no proportion places. Neither half can be given + scales of its own, so the split refuses rather than round.""" + sd = { + f"{PREFIX}.gate_up_proj": torch.zeros(1, rows, 4).to(torch.float8_e4m3fn), + f"{PREFIX}.gate_up_proj_weight_scale": torch.ones(1, scale_rows, 1), + f"{PREFIX}.down_proj": torch.zeros(1, 4, rows // 2).to(torch.float8_e4m3fn), + } + + with pytest.raises(NotImplementedError, match=r"cannot be split at the gate\|up boundary"): + split_fused_experts_state_dict(sd) + + +def test_split_refuses_projections_holding_different_expert_counts(): + """The split numbers experts by the first projection, so a down_proj expert beyond its count + would be dropped unwritten.""" + sd = _fused_state_dict(gated=True) + sd[f"{PREFIX}.down_proj"] = torch.zeros(LOCAL_EXPERTS + 1, 6, 4).to(torch.float8_e4m3fn) + + with pytest.raises(ValueError, match="down_proj"): + split_fused_experts_state_dict(sd) + + +def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): + fused = torch.zeros(2, 4, 4) + sd = { + "model.layers.1.mlp.experts.gate_up_proj": fused, + "model.layers.1.mlp.experts.down_proj": fused, + "model.layers.1.mlp.experts.gate_up_proj_weight_scale": fused, + "model.layers.0.mlp.experts.up_proj": fused, + "model.layers.0.mlp.experts.down_proj": fused, + # Not a fused-experts weight: already split, 2-D, outside an ``experts`` module, no tensor. + "model.layers.2.mlp.experts.0.down_proj": fused, + "model.layers.3.mlp.experts.down_proj": torch.zeros(4, 4), + "model.layers.4.mlp.gate_up_proj": fused, + "model.layers.5.mlp.experts.down_proj": "not a tensor", + } + assert _fused_experts_prefixes(sd) == [ + "model.layers.0.mlp.experts", + "model.layers.1.mlp.experts", + ] + + +# ------------------------------------------------------------------ sharded (DTensor) experts + + +def _shard_fused_weights(experts, mesh): + for proj in (experts._first_proj_attr, "down_proj"): + weight = getattr(experts, proj).data + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + setattr(experts, proj, nn.Parameter(sharded, requires_grad=False)) + return experts + + +def _whole(value): + """A sharded value gathered whole. gloo cannot gather FP8, so one-byte values move as bytes.""" + if not isinstance(value, DTensor): + return value + if value.element_size() != 1: + return value.full_tensor() + raw = value.to_local().view(torch.uint8) + raw = DTensor.from_local( + raw, + value.device_mesh, + value.placements, + run_check=False, + shape=value.shape, + stride=value.stride(), + ) + return raw.full_tensor().view(value.dtype) + + +def _two_rank_fold_job(rank, world_size): + mesh = init_device_mesh("cpu", (world_size,)) + cases = [ + (experts_cls, quant_cfg, NUM_EXPERTS) + for experts_cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts) + for quant_cfg in (mtq.FP8_DEFAULT_CFG, FP8_WEIGHT_ONLY_CFG, mtq.NVFP4_DEFAULT_CFG) + ] + # Three experts over two ranks: rank 0 holds two, rank 1 one. The folded tensors must keep the + # global expert count, not infer local x world rows. + cases.append((_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG, 3)) + for experts_cls, quant_cfg, num_experts in cases: + # Every rank calibrates the same seeded module, so all ranks hold identical quantizers. + experts = _quantized_experts(experts_cls, quant_cfg, num_experts) + want = _per_expert_export(experts) + plain = _keep_fused_export(experts) + sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + fused = _prefixed(sharded.state_dict()) + assert set(fused) == set(plain) + first = fused[f"{PREFIX}.{sharded._first_proj_attr}"] + n_local = first.to_local().shape[0] + offset = rank * math.ceil(num_experts / world_size) + assert _dtensor_dim0_offset(first) == offset + for key, value in fused.items(): + if key.endswith("_input_scale"): # one scalar shared by every expert + assert not isinstance(value, DTensor), key + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + else: # weights and per-expert scales stay sharded on the expert axis + assert isinstance(value, DTensor), key + assert value.placements == (Shard(0),), key + assert value.shape == plain[key].shape, key + local = plain[key][offset : offset + n_local] + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(local)), key + # Gathered whole, the ranks' folds split into exactly what the per-expert export writes. + whole = {key: _whole(value) for key, value in fused.items()} + _assert_byte_identical(split_fused_experts_state_dict(whole), want) + + # One expert over two ranks leaves rank 1 none: both ranks refuse, rather than rank 1 alone + # failing while rank 0 goes on to write. + lone = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG, num_experts=1) + with pytest.raises(NotImplementedError, match="leaves a rank none"): + _export_fused_experts_keep_fused(_shard_fused_weights(lone, mesh), torch.bfloat16) + + # Expert parallelism leaves a rank only its own experts' quantizers, in shard order, under a + # fused weight spanning every expert. Kept last, with the joint decision below: without that + # pairing rank 1 alone fails, and rank 0 then waits only in the collectives below, which the + # worker pool ends once a rank has failed. + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG) + plain = _keep_fused_export(experts) + n_local = NUM_EXPERTS // world_size + mine = slice(rank * n_local, (rank + 1) * n_local) + + def sharded_keeping_quantizers(kept): + sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) + for proj in (sharded._first_proj_attr, "down_proj"): + name = f"{proj}_weight_quantizers" + setattr(sharded, name, nn.ModuleList(list(getattr(sharded, name))[kept])) + return sharded + + with pytest.raises(NotImplementedError, match="3 weight quantizers for 4 experts in 2 shards"): + _export_fused_experts_keep_fused(sharded_keeping_quantizers(slice(0, 3)), torch.bfloat16) + sharded = sharded_keeping_quantizers(mine) + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + fused = _prefixed(sharded.state_dict()) + assert set(fused) == set(plain) + for key, value in fused.items(): + if isinstance(value, DTensor): + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(plain[key][mine])), key + else: + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + + # Each rank sees only its own experts' quantizers there, so the ranks decide on a projection + # together: rank 1 disabling some or all of its experts' is refused on both ranks, rather than + # folded on rank 0 alone, and every expert disabled leaves the projection as it is on both. + for disabled in ([3], [2, 3], [0, 1, 2, 3]): + sharded = sharded_keeping_quantizers(mine) + for expert, quantizer in enumerate(sharded.down_proj_weight_quantizers, start=mine.start): + if expert in disabled: + quantizer.disable() + if len(disabled) < NUM_EXPERTS: + with pytest.raises(NotImplementedError, match=re.escape(f"quantizers {disabled} of 4")): + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + else: + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + assert hasattr(sharded, "gate_up_proj_weight_scale") + assert not hasattr(sharded, "down_proj_weight_scale") + + +def test_two_ranks_fold_their_own_experts_under_global_indices(cpu_dist_workers): + """Rank 1 holds experts 2-3: each must be folded with ITS quantizer at its global offset, and the + fused tensors stay sharded on the expert axis at their global shape, an uneven shard included; + gathered, they split into the per-expert export. Its quantizers are every expert's (FSDP2) or, + under expert parallelism, its own two, and then the ranks decide on each projection's + quantization together. A layout that leaves a rank no expert, or quantizers that are neither, + is refused on every rank.""" + cpu_dist_workers.run(_two_rank_fold_job) + + +@pytest.mark.parametrize( + ("mesh_shape", "placements"), + [((1,), [Shard(2)]), ((1, 1), [Shard(0), Shard(0)])], + ids=["off-the-expert-axis", "over-two-mesh-dims"], +) +def test_a_fused_weight_sharded_other_than_by_expert_is_refused( + gloo_world_of_one, mesh_shape, placements +): + """The fold shards the scales on the expert axis alone, and the ranks agree on their experts + over one mesh dim: the scales of another layout would claim one rank's slice as the whole.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + mesh = init_device_mesh("cpu", mesh_shape) + weight = distribute_tensor(experts.down_proj.data, mesh, placements) + experts.down_proj = nn.Parameter(weight, requires_grad=False) + + with pytest.raises(NotImplementedError, match="sharded over one mesh dim"): + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + +# ------------------------------------------------------------------ export handlers + + +@pytest.mark.parametrize("keep_fused", [True, False], ids=["keep_fused", "split"]) +def test_fused_experts_handler_keeps_experts_fused_only_when_asked(keep_fused): + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + container = nn.Module() + container.experts = experts + assert ExportContext(model=container, dtype=torch.bfloat16).keep_fused_experts is False + ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=keep_fused) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) + + names = set(experts.state_dict()) + if keep_fused: + assert {"gate_up_proj", "gate_up_proj_weight_scale", "down_proj"} <= names + assert not any(n.startswith("0.") for n in names) + else: + assert "0.gate_proj.weight" in names + assert "gate_up_proj" not in names + + +def test_keep_fused_refuses_fused_experts_not_named_experts(): + """The split finds a fused weight by its ``experts.`` key, as the writer does: a container + named otherwise is refused before anything is folded, not written fused to fail at load.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + container = nn.Module() + container.moe = experts + before = copy.deepcopy(experts.state_dict()) + ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) + + with pytest.raises(NotImplementedError, match="named 'moe'"): + ExportModuleRegistry.match(experts)("model.layers.0.mlp.moe", experts, ctx) + + torch.testing.assert_close(experts.state_dict(), before, rtol=0, atol=0) + + +@pytest.mark.parametrize("dense", ["fsdp", "replicate"]) +def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_part( + monkeypatch, dense +): + """FSDP2 owns the fused weights when it shards the dense part around them: folded outside a + weight-update window, they would be overwritten with its original shards at the next reshard. + Outside held windows the handler closes the window itself.""" + events = [] + + @contextlib.contextmanager + def window(root, modules, names=None): + events.append("open") + yield + events.append("close") + + monkeypatch.setattr(handlers, "fsdp2_aware_weight_update", window) + monkeypatch.setattr( + handlers, "_export_fused_experts_keep_fused", lambda module, dtype: events.append("fold") + ) + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + container = nn.Module() + container.experts = experts + set_distribution_plan(container, DistributionPlan(2, ep_size=2, dense=dense)) + ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) + + ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) + + assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) + + +class GptOssExperts(nn.Module): + """Named as transformers' fused BMM experts, which store their weights transposed.""" + + +def test_keep_fused_refuses_bmm_experts(): + """The writer splits a keep-fused export's ``experts.`` weights as ``(num_experts, out_dim, + in_dim)``; BMM experts store theirs transposed, so the split would cut every expert wrong, and + the export refuses them before packing anything.""" + experts = GptOssExperts() + experts.gate_up_proj = nn.Parameter(torch.randn(NUM_EXPERTS, HIDDEN_DIM, 2 * INTERMEDIATE_DIM)) + before = copy.deepcopy(experts.state_dict()) + ctx = ExportContext(model=experts, dtype=torch.bfloat16, keep_fused_experts=True) + + with pytest.raises(NotImplementedError, match="stores its experts transposed"): + ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) + + torch.testing.assert_close(experts.state_dict(), before, rtol=0, atol=0)