Skip to content

fix(qwen3_moe): export W4A16 NVFP4 checkpoints via NVFP4 MoE path #176

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
potoior wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(qwen3_moe): export W4A16 NVFP4 checkpoints via NVFP4 MoE path #176

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions tensorrt_edgellm/checkpoint/repacking.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1320,7 +1320,8 @@ def repack_nvfp4_gated_moe_experts(
:data:`NVFP4_MOE_INTERMEDIATE_SIZE_ALIGNMENT` for ``"concat"``.
"""
from ..models.linear import \
is_nvfp4_linear # local import to avoid circular dep
NVFP4A16MarlinLinear # local import to avoid circular dep
from ..models.linear import is_nvfp4_linear

if fc1_layout == "interleave":
build_fc1_dense = _interleave_gated_moe_fc1
Expand Down Expand Up @@ -1364,8 +1365,11 @@ def repack_nvfp4_gated_moe_experts(
gate = expert.gate_proj
up = expert.up_proj
down = expert.down_proj
if not (is_nvfp4_linear(gate) and is_nvfp4_linear(up)
and is_nvfp4_linear(down)):

def _is_nvfp4(m):
return is_nvfp4_linear(m) or isinstance(m, NVFP4A16MarlinLinear)

if not (_is_nvfp4(gate) and _is_nvfp4(up) and _is_nvfp4(down)):
raise TypeError("Gated NVFP4 MoE experts must use NVFP4 quant")

gate_dense = decode_modelopt_nvfp4(gate.weight, gate.weight_scale,
Expand Down
5 changes: 3 additions & 2 deletions tensorrt_edgellm/models/qwen3_moe/modeling_qwen3_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@
import torch
import torch.nn as nn

from ...config import QUANT_FP16, QUANT_NVFP4, ModelConfig
from ...config import QUANT_FP16, QUANT_NVFP4, QUANT_NVFP4_A16, ModelConfig
from ..default.modeling_default import (MLP, Attention, OnnxSpec, RMSNorm,
_make_flat_wrapper)
from ..linear import FP16Linear, make_linear
Expand Down Expand Up @@ -223,7 +223,8 @@ def __init__(self, config: ModelConfig, module_prefix: str = "") -> None:
self.hidden_size = config.hidden_size
self.group_size = config.quant.group_size
self.zero_point_offset = config.quant.gptq_zero_point_offset
self._use_nvfp4_moe = config.quant.quant_type == QUANT_NVFP4
self._use_nvfp4_moe = config.quant.quant_type in (QUANT_NVFP4,
QUANT_NVFP4_A16)
self._use_fp16_moe = config.quant.quant_type == QUANT_FP16
# All paths compute the same SwiGLU expert FFN; the integer is just
# each plugin's own enum for it. Int4MoePlugin names the elementwise
Expand Down