Llama-4 把模型结构里 expert 的参数组织搞得太恶心了。BF16 参数还好，如果想做一下量化，就会面临一堆麻烦。
先说一下 Llama-4 开源参数的问题。
HF BF16版本：3 维 experts 参数，gate_proj 和 up_proj 融合
Llama-4-Scout/Maverick 的 HuggingFace 版本参数 BF16 版本中，路由 expert 专家参数是按照 3 维存储的，而且 up_proj 和 gate_proj 放在了同一个 key 里。比如 Llama-4-Scout 的专家参数在 safetensors 文件中是这样存储的：
| Tensors | Shape | Precision |
|---|---|---|
| language_model.model.layers.n.feed_forward.experts.down_proj | [16,8192,5120] | BF16 |
| language_model.model.layers.n.feed_forward.experts.gate_up_proj | [16,5120,16384] | BF16 |
其中 n 是层数索引。Maverick 的 experts Shape 只是从 [16, , ] 变成了 [128, , ]，其余都一样。
这个形状是与 Transformers 库中 llama4 的模型代码保持一致的：
// https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama4/modeling_llama4.py
class Llama4TextExperts(nn.Module):
def __init__(self, config: Llama4Config):
...
self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, 2 * self.expert_dim))
self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size)))
...
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
...
gate_up = torch.bmm(hidden_states, self.gate_up_proj)
gate, up = gate_up.chunk(2, dim=-1) # not supported for DTensors
next_states = torch.bmm((up * self.act_fn(gate)), self.down_proj)
...
HF FP8版本：2 维 expert 参数，gate_proj 和 up_proj 分离
在 FP8 量化版本 Llama-4-Maverick-17B-128E-Instruct-FP8 中，expert 参数又被拆成了二维的，通过 key 中的索引标识属于哪个专家：
| Tensors | Shape | Precision |
|---|---|---|
| language_model.model.layers.n.feed_forward.experts.m.down_proj.weight | [5120, 8192] | F8_E4M3 |
| language_model.model.layers.n.feed_forward.experts.m.down_proj. weight_scale | [5120, 1] | BF16 |
| language_model.model.layers.n.feed_forward.experts.m.gate_proj.weight | [8192, 5120] | F8_E4M3 |
| language_model.model.layers.n.feed_forward.experts.m.gate_proj. weight_scale | [8192, 1] | BF16 |
| language_model.model.layers.n.feed_forward.experts.m.up_proj.weight | [8192, 5120] | F8_E4M3 |
| language_model.model.layers.n.feed_forward.experts.m.up_proj. weight_scale | [8192, 1] | BF16 |
其中 n 是层数索引，m 是层中的专家数索引。
之前说到，Transformers 库中 modeling_llama4.py 是只支持融合模型的，那这种格式的参数怎么加载？
哎，人家用了一个办法：如果读到模型配置里有量化配置，在加载模型前修改一下模型结构。硬是在 transformers/quantizers/base.py 中增加了一个 _convert_model_for_quantization 方法，如果模型有子 module 叫做 "Llama4TextExperts"，在量化 preprocess 的时候就给替换成 SequentialLlama4TextExperts 实现，不使用原始的 Llama4TextExperts 实现。
注意哦，这个替换对所有量化模型都生效。这种特化方法，要放在我的团队里，CR 估计都过不了。
怎么量化 llama4-scout ？
FB 官方提供了 128 专家的 FP8 模型，但是没有提供 16 专家的 FP8 量化模型。毕竟 16 专家模型也 200 多 G，如果想量化 16 专家模型，该怎么做呢？
meta 官方在 source/en/model_doc/llama4.md 里推荐的方法，是加载模型时使用 FbgemmFp8Config 进行 online 量化，这个我没跑成功，看着错误像是只支持单卡跑 fbgemm 量化，但是 H800 显存不够。如果这个问题可以解决，欢迎留言告诉我方法。
$ torchrun --nproc-per-node=4 test-16-fbgemm.py
self.pre_quantized False
Loading checkpoint shards: 100%|███████| 50/50 [00:45<00:00, 1.09it/s]
...
[rank0]: File "/workspace/venv-fbgemm/lib/python3.10/site-packages/transformers/integrations/fbgemm_fp8.py", line 52, in forward
[rank0]: x_quantized, x_scale = torch.ops.fbgemm.quantize_fp8_per_row(
[rank0]: File "/workspace/venv-fbgemm/lib/python3.10/site-packages/torch/_ops.py", line 1158, in __call__
...
[rank0]: File "/workspace/venv-fbgemm/lib/python3.10/site-packages/torch/distributed/tensor/_sharding_prop.py", line 486, in propagate_op_sharding_non_cached
[rank0]: raise NotImplementedError(
[rank0]: NotImplementedError: Operator fbgemm.quantize_fp8_per_row.default does not have a sharding strategy registered.
咱又没有那么大显存的卡，只能想别的办法，能不能仿照 Llama-4-Maverick-17B-128E-Instruct-FP8 来转换出来一个 16 专家的 FP8 模型呢？
Llama-4-Maverick-17B-128E-Instruct-FP8 怎么转出来的？
首先了解一下 llama4 的转换脚本：
- 在 github llama-models 代码库里，有一个 llama4/scripts/quantize.py 脚本，是用来将原始的 pytorch 模型参数，通过 fbgemm 转成 FP8 量化的 pytorch 模型参数。
- 在 transformers 代码库里，有一个 llama4/convert_llama4_weights_to_hf.py 脚本，是用来将原始的 pytorch 模型参数，通过映射表映射到 huggingface 的 safetensors 模型参数。
然后来看 llama4 发布的模型：
- Llama-4-Maverick-17B-128E-Instruct-Original：原始 pytorch 模型；
- Llama-4-Maverick-17B-128E-Instruct：HuggingFace 模型，通过 convert_llama4_weights_to_hf.py 转换原始模型而来。
- Llama-4-Maverick-17B-128E-Instruct-FP8-Original：FP8 量化的原始 pytorch 模型，通过 quantize.py 转换原始模型而来。
- Llama-4-Maverick-17B-128E-Instruct-FP8：FP8 量化的 HuggingFace 模型，哪来的？
它不是通过 convert_llama4_weights_to_hf.py 转换 Llama-4-Maverick-17B-128E-Instruct-FP8-Original 来的，因为：
1. FP8-Original 的量化是通过 fbgemm 做的，FP8 模型的量化是用 compressed-tensor 做的；
2. FP8-Original 是分 TP 做的量化，与 FP8 的量化方法也不同。
它也不是通过 llm-compressor(compressed-tensor) 转换 Llama-4-Maverick-17B-128E-Instruct 来的。因为 compressed-tensor 目前仅支持 Linear 算子的量化，但前面说过，加载原始 BF16 模型用的是 bmm 算子。
convert_llama4_weights_to_hf.py 中有一个 _OFFLINE_QUANT_COMPATIBLE 的参数，可以控制是否对专家进行融合。但是这样转出来的模型，modeling_llama4.py 是无法加载的。
我猜测 meta 线下可能替换了 modeling_llama4.py 中的专家层到 SequentialLlama4TextExperts，成功加载模型以后，然后再通过 llm-compressor 进行的模型转换。
然后就这样试了一下，先替换 transformers 库源码 modeling_llama4.py 中的 experts 组，并从源码安装 transformers：
--- a/src/transformers/models/llama4/modeling_llama4.py
+++ b/src/transformers/models/llama4/modeling_llama4.py
@@ -153 +153,2 @@ class Llama4TextMoe(nn.Module):
- self.experts = Llama4TextExperts(config)
+ from transformers.quantizers.base import SequentialLlama4TextExperts
+ self.experts = SequentialLlama4TextExperts(config)然后再执行：
export OFFLINE_QUANT_COMPATIBLE=1
python3 src/transformers/models/llama4/convert_llama4_weights_to_hf.py --instruct --convert_checkpoints --input_dir /workspace/Llama-4-Scout-17B-16E-Instruct-Original --output_dir /workspace/Llama-4-Scout-17B-16E-Instruct-Split-Experts
然后再使用 llm-compressor 库，用下面的 recipe 参数对模型进行量化，这样就成功了。转换出来的模型仓库内容基本与 Llama-4-Maverick-17B-128E-Instruct-FP8 相同。
recipe = QuantizationModifier(
targets="Linear",
scheme="FP8_DYNAMIC",
ignore=[
're:.*lm_head',
're:.*self_attn',
're:.*router',
're:.*vision_model',
're:.*multi_modal_projector',
're:.*shared_expert',
're:.*feed_forward.gate_proj',
're:.*feed_forward.up_proj',
're:.*feed_forward.down_proj'
],
)
为什么？
我是想不通 llama4 为啥要把简单的事情搞这么复杂，这样一改，很多开源库都要去适配这种格式的 MoE 参数，但功能上看起来又没有任何变化。做融合的 expert 参数有什么实际的收益吗？