moe_utils

Utilities for Mixture-of-Experts (MoE) model export.

Functions

save_expert_token_count_table	Collect expert_token_count from all quantized MoE layers and save as an HTML table.
split_fused_experts_state_dict	Rewrite fused-experts entries in a (consolidated) HF state dict into per-expert entries.

save_expert_token_count_table(model, output_dir=None)

Collect expert_token_count from all quantized MoE layers and save as an HTML table.

The table has rows for each MoE layer and columns for each expert, with cell values showing the number of tokens routed to that expert during calibration.

Parameters:

• model (Module) – The model containing quantized MoE layers with expert_token_count attributes.

• output_dir (str | Path | None) – Directory to save the HTML file. Defaults to current directory.

split_fused_experts_state_dict(state_dict)

Rewrite fused-experts entries in a (consolidated) HF state dict into per-expert entries.

The distributed export writes experts fused (3-D experts.<proj> weights + fused scale buffers, kept sharded for DCP). This post-pass splits them into the same per-expert layout the in-model _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.

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)

Parameters:

state_dict (dict)

Return type:

dict