Apodex-1.1-mini-NVFP4 fails: dense_quant hardcoded to NVFP4 when experts are NVFP4, but shared_expert is FP8 (modelopt MIXED_PRECISION)

[zangetsu787]

Apodex-1.1-mini-NVFP4 fails: dense_quant hardcoded to NVFP4 when experts are NVFP4, but Apodex's shared_expert is FP8 (modelopt MIXED_PRECISION)

Environment

• FreeToken 0.1.1+cu130 (freetoken-0.1.1+g30aa89115)
• PyTorch 2.11.0+cu130
• GPU: RTX 4070 Ti 12 GB
• Model: apodex/Apodex-1.1-mini-NVFP4 (huggingface)
• Conversion: ft checkpoint --model ... --out ... --dtype bfloat16 --moe-backend offload succeeds; ft serve crashes at weight load

Model profile

• Architecture: Qwen3_5MoeForConditionalGeneration (Qwen3.5 MoE A3B, 256 experts, 40 layers)
• Experts: NVFP4 (the "NVFP4" in the name)
• Shared expert: FP8 (float8_e4m3fn), per hf_quant_config.json — every mlp.shared_expert.{gate,up,down}_proj is tagged quant_algo: "FP8"
• quantized_layers covers linear_attn.* and self_attn.{q,k,v,o}_proj as FP8 too
• So this is a modelopt MIXED_PRECISION checkpoint, not pure NVFP4

Root cause

In freetoken/models/qwen3_5_moe/config.py, line ~194:

dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config)

This unconditionally sets dense_quant="nvfp4" whenever experts are NVFP4. For Qwen3.5 MoE NVFP4 checkpoints where the shared_expert is actually FP8 (or bf16), this is wrong. FreeToken then routes the shared_expert through Nvfp4DenseColMerged (in qwen3_5_moe/moe.py), which calls nvfp4_linear.py:866 looking for weight/weight_scale/weight_global — keys that don't exist on an FP8 shared_expert (which has weight + scalar weight_scale + input_scale, and no weight_global).

Result:

KeyError: 'model.layers.0.mlp.shared_expert.gate_up_proj.weight'
  at freetoken/kernel/triton/nvfp4_linear.py:866 in load_state_dict
  called from freetoken/models/qwen3_5_moe/moe.py:Nvfp4DenseColMerged.load_state_dict

_dense_mlp_quant is documented to not match shared_expert.* (only bare .mlp.gate_proj), so it never overrides the bad default.

Suggested fix

Either:

1. Probe hf_quant_config.json for the actual shared_expert quant algo (similar to _expert_quant) and use it instead of forcing NVFP4. e.g. add a _shared_expert_quant(hf_config) helper that inspects quantized_layers["...mlp.shared_expert.gate_proj"].
2. Default to dense_quant="none" when _dense_mlp_quant returns "none" — i.e. remove the unconditional dense_quant = "nvfp4" if expert_quant == "nvfp4" line and just use the function's return value. The current comment claims NVFP4 dense weights are "independent of attention quant" but that's only true for pure NVFP4 checkpoints, not the mixed-precision ones modelopt produces.

Repro:

# Download (23 GB on D:)
huggingface-cli download apodex/Apodex-1.1-mini-NVFP4 --local-dir D:/freetoken-models/Apodex-1.1-mini-NVFP4

# Convert (works fine)
ft checkpoint --model "D:/freetoken-models/Apodex-1.1-mini-NVFP4" \
              --out "D:/freetoken-models/Apodex-1.1-mini-NVFP4-FTW" \
              --dtype bfloat16 --moe-backend offload --shard-gib 8
# OK: 20.19 GiB FTW written

# Serve (crashes at weight load)
ft serve --model "D:/freetoken-models/Apodex-1.1-mini-NVFP4-FTW" --port 1922 --memory-ratio 0.85
# KeyError: 'model.layers.0.mlp.shared_expert.gate_up_proj.weight'

The model itself is fine — apodex.ai demo and vLLM/SGLang work; the Apodex README even provides an SGLang launch command using the same modelopt_mixed quantization. So this is purely FreeToken's mixed-precision inference path.