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

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

Model Overview

• Model Architecture: MixtralForCausalLM (Mixture-of-Experts)
  • Input: Text
  • Output: Text
• Source Model: Mixtral-8x7B-Instruct-v0.1
• Supported Hardware: AMD EPYC (CPU inference)
• Preferred Operating System: Linux
• Inference Engine: vLLM v0.23.0
• Quantization Framework: TorchAO v0.17.0
• Quantization Method: 8-bit Dynamic Activation, 8-bit Weight Quantization, Symmetric
  • Dense nn.Linear layers quantized via TorchAoConfig
  • MoE expert weights (gate_up_proj, down_proj) quantized via a second FqnToConfig pass
  • Skipped (kept in BF16): lm_head, router, gate
• On-disk Layout: Sharded pytorch_model-*.bin (legacy Mixtral per-expert w1/w2/w3 layout under block_sparse_moe.experts.*) — required so vLLM's mixtral.py weight loader can map the keys
• Compatible Stack:
  • ZenDNN v6.0.0
  • zentorch v2.11.0.2
  • PyTorch v2.11.0
  • TorchAO v0.17.0
  • vLLM v0.23.0

[!NOTE] zentorch v2.11.0.2 for PyTorch v2.11.0 has to be built from source.

Quantization

The model is produced in four steps:

1. Load in BF16 with a TorchAoConfig so all dense nn.Linear layers are converted to INT8 dynamic-activation / INT8-weight (symmetric). lm_head, router, and gate are kept in BF16.
2. Run a second quantize_ pass with FqnToConfig (regex-based, requires torchao >= 0.17) to quantize the packed MoE expert parameters (gate_up_proj, down_proj) that the TorchAoConfig pass skips because they are nn.Parameter, not nn.Linear.
3. Unpack each MixtralExperts block (3D packed Int8Tensor parameters) into the legacy per-expert w1/w2/w3 nn.Linear layout that vLLM's mixtral.py weight loader expects. The gate_up_proj packs [gate, up] along dim 1 — first half becomes w1, second half becomes w3. down_proj becomes w2.
4. Save with torch.save as sharded pytorch_model-*.bin (not safetensors). This is required because:
   • safetensors splits each Int8Tensor into three underscore-prefixed sub-keys (_data, _scale, _zero_point) that vLLM's loader drops.
   • transformers >= 5.7 forces safetensors from save_pretrained, so we bypass it and write shards directly via huggingface_hub.split_torch_state_dict_into_shards.
   • In-memory keys are renamed .mlp. → .block_sparse_moe. to match the legacy Mixtral on-disk layout.

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)

[!NOTE] Weights are saved as pytorch_model-*.bin (not safetensors) because torchao's Int8Tensor subclass cannot currently be serialized via safetensors without losing its identity (it would split into three underscore-prefixed sub-tensors that vLLM's loader drops).

[!NOTE] FqnToConfig with regex-based FQN matching requires torchao >= 0.17. Older versions will silently leave the MoE expert weights in BF16.

Quick Start

Requirements

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

CPU runtime libraries (only needed if not already present):

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

Recommended environment variables

# 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}"

Locate the libraries with find / -name 'libtcmalloc_minimal.so.4' and find / -name 'libiomp5.so', then substitute the resulting directory for <path to lib>.

Evaluation

The model was evaluated against the BF16 (unquantized) baseline using lm-evaluation-harness with the vLLM engine.

Evaluation Command

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 .

Limitations

• Version Lock: This model is quantized with TorchAO v0.17.0 and is compatible only with PyTorch v2.11.0 / ZenDNN v6.0.0. It will not load correctly on other PyTorch versions.
• Serialization Format: Weights are stored as sharded pytorch_model-*.bin rather than safetensors, because torchao's Int8Tensor subclass cannot currently be round-tripped through safetensors in a form vLLM's Mixtral loader accepts.
• Loader Layout: Expert weights are stored in the legacy block_sparse_moe.experts.{i}.w1/w2/w3 layout (not the packed experts.gate_up_proj / experts.down_proj layout) to remain compatible with vLLM's mixtral.py weight loader.
• CPU Only: This model is optimized for AMD EPYC CPU inference via ZenDNN. It is not intended for GPU inference.

License

This model is distributed under the same license as the source model. See the LICENSE file for details.

Modifications copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.