amd/Mixtral-8x7B-Instruct-v0.1-da8w8-torchao-v0.17.0

import os
import json
from collections import OrderedDict

import torch
import torch.nn as nn
from huggingface_hub import split_torch_state_dict_into_shards
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from torchao.quantization import (
    Int8DynamicActivationInt8WeightConfig,
    quantize_,
)
from torchao.quantization.quant_api import FqnToConfig
from torchao.quantization.quant_primitives import MappingType

MODEL_ID   = "mistralai/Mixtral-8x7B-Instruct-v0.1"
OUTPUT_DIR = "mixtral-instruct-dynamic-quantized"
os.makedirs(OUTPUT_DIR, exist_ok=True)


class _LegacyExpert(nn.Module):
    """Legacy per-expert layout (w1=gate, w2=down, w3=up) for vLLM's mixtral.py loader."""

    def __init__(self, hidden: int, intermediate: int):
        super().__init__()
        self.w1 = nn.Linear(hidden, intermediate, bias=False)
        self.w2 = nn.Linear(intermediate, hidden, bias=False)
        self.w3 = nn.Linear(hidden, intermediate, bias=False)


# ---- Step 1: dense nn.Linear -> INT8 dyn-act / INT8 weight (symmetric) -----
ao_config = Int8DynamicActivationInt8WeightConfig(
    version=2, act_mapping_type=MappingType.SYMMETRIC,
)

# TorchAoConfig only touches nn.Linear; writing it to config.json is what tells
# vLLM to use TorchAOFusedMoEMethod.
quantization_config = TorchAoConfig(
    ao_config,
    modules_to_not_convert=["gate", "lm_head", "router"],
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    device_map="cpu",
    quantization_config=quantization_config,
    trust_remote_code=True,
)

# ---- Step 2: quantize packed MoE expert params (gate_up_proj, down_proj) ---
expert_fqn_config = FqnToConfig(
    fqn_to_config=OrderedDict({
        r"re:.*\.gate_up_proj$": ao_config,
        r"re:.*\.down_proj$":    ao_config,
    })
)
quantize_(model, expert_fqn_config, filter_fn=None)

# ---- Step 3: unpack MixtralExperts -> legacy per-expert nn.Linears ---------
# save_pretrained's auto-unpack hook doesn't fire for Int8Tensor, leaving
# vLLM's mixtral.py loader unable to map the packed keys.
print("=== Unpacking MixtralExperts -> legacy per-expert layout ===", flush=True)
for layer_idx, layer in enumerate(model.model.layers):
    moe_block = layer.mlp
    old_experts = moe_block.experts
    num_experts  = old_experts.num_experts
    intermediate = old_experts.intermediate_dim
    hidden       = old_experts.hidden_dim

    gup = old_experts.gate_up_proj.data  # Int8Tensor [E, 2*inter, hidden]
    dp  = old_experts.down_proj.data     # Int8Tensor [E, hidden, inter]

    new_experts = nn.ModuleList()
    for j in range(num_experts):
        e = _LegacyExpert(hidden, intermediate)
        # gate_up_proj packs [gate, up] along dim 1: first half=w1, second=w3.
        w1_data = gup[j, :intermediate, :].contiguous()
        w3_data = gup[j, intermediate:, :].contiguous()
        w2_data = dp[j].contiguous()
        e.w1.weight = nn.Parameter(w1_data, requires_grad=False)
        e.w3.weight = nn.Parameter(w3_data, requires_grad=False)
        e.w2.weight = nn.Parameter(w2_data, requires_grad=False)
        new_experts.append(e)
        if layer_idx == 0 and j == 0:
            for k in ("w1", "w2", "w3"):
                w = getattr(e, k).weight.data
                print(
                    f"  layer 0 expert 0 {k}: cls={type(w).__name__} "
                    f"shape={list(w.shape)} dtype={w.dtype}",
                    flush=True,
                )

    del moe_block.experts
    moe_block.experts = new_experts

print("=== Unpacking complete; saving as sharded .bin (torch.save) ===", flush=True)

# ---- Step 4: write sharded pytorch_model-*.bin (bypass safetensors) --------
# Remove any stale shards so the index points only at the new .bin files.
for stale in os.listdir(OUTPUT_DIR):
    if stale.endswith((".safetensors", ".safetensors.index.json",
                       ".bin", ".bin.index.json")):
        try:
            os.remove(os.path.join(OUTPUT_DIR, stale))
        except OSError:
            pass

model.config.save_pretrained(OUTPUT_DIR)
if model.can_generate():
    model.generation_config.save_pretrained(OUTPUT_DIR)

state_dict = model.state_dict()
print(f"  state_dict: {len(state_dict)} tensors", flush=True)

# Rename `.mlp.` -> `.block_sparse_moe.` to match legacy on-disk Mixtral layout.
state_dict = {
    (k.replace(".mlp.", ".block_sparse_moe.") if ".mlp." in k else k): v
    for k, v in state_dict.items()
}
print(f"  state_dict after .mlp.->.block_sparse_moe. rename: {len(state_dict)} tensors", flush=True)

WEIGHTS_NAME       = "pytorch_model.bin"
WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
filename_pattern   = WEIGHTS_NAME.replace(".bin", "{suffix}.bin")
state_dict_split = split_torch_state_dict_into_shards(
    state_dict, filename_pattern=filename_pattern, max_shard_size="5GB",
)

for shard_file, tensor_names in state_dict_split.filename_to_tensors.items():
    shard_state_dict = {name: state_dict[name] for name in tensor_names}
    shard_path = os.path.join(OUTPUT_DIR, shard_file)
    print(f"  saving shard {shard_file} ({len(shard_state_dict)} tensors)", flush=True)
    torch.save(shard_state_dict, shard_path)

if state_dict_split.is_sharded:
    index = {
        "metadata": {
            "total_parameters": sum(p.numel() for p in state_dict.values()),
            **state_dict_split.metadata,
        },
        "weight_map": state_dict_split.tensor_to_filename,
    }
    with open(os.path.join(OUTPUT_DIR, WEIGHTS_INDEX_NAME), "w") as f:
        json.dump(index, f, indent=2, sort_keys=True)
    print(f"  wrote {WEIGHTS_INDEX_NAME}", flush=True)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
tokenizer.save_pretrained(OUTPUT_DIR)
print("=== Save complete ===", flush=True)

pip install --extra-index-url https://download.pytorch.org/whl/cpu \
            --extra-index-url https://wheels.vllm.ai/cpu/ \
    torch==2.11.0+cpu \
    vllm==0.23.0 \
    torchao==0.17.0 \
    "lm-eval[vllm]==0.4.12" \
    huggingface_hub

conda install -c conda-forge gperftools=2.17.2 llvm-openmp=18.1.8 --no-deps -y

# TorchInductor + zentorch
export TORCHINDUCTOR_FREEZING=1
export TORCHINDUCTOR_AUTOGRAD_CACHE=0
export VLLM_USE_AOT_COMPILE=0
export ZENDNNL_MATMUL_ALGO=1
export ZENTORCH_FUSED_MOE=1                    # recommended for Mixtral-8x7B (MoE)

# Required CPU runtime libraries
export LD_PRELOAD="<path to lib>/libtcmalloc_minimal.so.4:<path to lib>/libiomp5.so${LD_PRELOAD:+:$LD_PRELOAD}"

lm_eval \
    --model vllm \
    --model_args pretrained=amd/Mixtral-8x7B-Instruct-v0.1-da8w8-torchao-v0.17.0,tokenizer=mistralai/Mixtral-8x7B-Instruct-v0.1,dtype=bfloat16 \
    --tasks gsm8k \
    --batch_size auto \
    --trust_remote_code \
    --num_fewshot 5 \
    --log_samples \
    --gen_kwargs "max_gen_toks=2048" \
    --apply_chat_template \
    --output_path .