| |
|
|
| |
| import os |
|
|
| import torch |
| from transformers import AutoConfig |
|
|
|
|
| def load_hf_mixtral_config(): |
| hf_model = os.getenv("HF_MODEL") |
| assert hf_model is not None, "Please set HF_MODEL to a HuggingFace name e.g. meta-llama/Llama-3.1-8B-Instruct" |
| config = AutoConfig.from_pretrained(hf_model, local_files_only=os.getenv("CI") == "true") |
| return config |
|
|
|
|
| def fuse_mixtral_experts(state_dict, num_experts, src_prefix, dst_prefix): |
| """Convert per-expert w1/w2/w3 weights to the transformers 5.x fused MixtralExperts format. |
| |
| transformers 5.x replaced the per-expert MixtralBlockSparseTop2MLP modules with a single fused |
| ``MixtralExperts`` module whose weights are applied via ``F.linear`` (same orientation as the old |
| per-expert ``nn.Linear`` weights): |
| * ``{dst_prefix}.gate_up_proj``: (num_experts, 2*intermediate, hidden) = cat([w1, w3], dim=0) |
| per expert (gate first, up second; split via ``chunk(2, dim=-1)`` in the forward). |
| * ``{dst_prefix}.down_proj``: (num_experts, hidden, intermediate) = w2 per expert. |
| |
| ``src_prefix`` is where the per-expert weights live (e.g. ``experts`` for a standalone |
| MixtralSparseMoeBlock, ``block_sparse_moe.experts`` inside a decoder layer). The matching |
| per-expert source keys are consumed; all other keys pass through unchanged. |
| """ |
| consumed = set() |
| gate_up, down = [], [] |
| for e in range(num_experts): |
| w1 = state_dict[f"{src_prefix}.{e}.w1.weight"] |
| w2 = state_dict[f"{src_prefix}.{e}.w2.weight"] |
| w3 = state_dict[f"{src_prefix}.{e}.w3.weight"] |
| gate_up.append(torch.cat([w1, w3], dim=0)) |
| down.append(w2) |
| consumed.update(f"{src_prefix}.{e}.w{i}.weight" for i in (1, 2, 3)) |
| out = {k: v for k, v in state_dict.items() if k not in consumed} |
| out[f"{dst_prefix}.gate_up_proj"] = torch.stack(gate_up, dim=0) |
| out[f"{dst_prefix}.down_proj"] = torch.stack(down, dim=0) |
| return out |
|
|
