amd/gpt-oss-20b-BF16-da8w8-torchao-v0.17.0

import os
from collections import OrderedDict

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from torchao.quantization import (
    Int8DynamicActivationInt8WeightConfig,
    quantize_,
)
from torchao.quantization.granularity import PerRow
from torchao.quantization.quant_api import FqnToConfig
from torchao.quantization.quant_primitives import MappingType

MODEL_ID   = "unsloth/gpt-oss-20b-BF16"
OUTPUT_DIR = "amd/gpt-oss-20b-da8w8-torchao-v0.17.0"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Pass 1: standard dynamic-act / weight INT8 for nn.Linear layers.
# Skip lm_head and router (kept in BF16).
ao_config = Int8DynamicActivationInt8WeightConfig(
    version=2,
    act_mapping_type=MappingType.SYMMETRIC,
)
quantization_config = TorchAoConfig(
    ao_config,
    modules_to_not_convert=["lm_head", "router"],
)

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

# Pass 2: quantize MoE expert weights that the TorchAoConfig pass skipped
# because they live as nn.Parameter tensors, not nn.Linear modules.
# Use per-row granularity on activations and weights for expert tensors.
ao_config_experts = Int8DynamicActivationInt8WeightConfig(
    version=2,
    act_mapping_type=MappingType.SYMMETRIC,
    granularity=(PerRow(dim=-1), PerRow(dim=1)),
)

# Match the MoE expert parameter tensors via their fully-qualified names.
expert_fqn_config = FqnToConfig(
    fqn_to_config=OrderedDict({
        r"re:.*\.experts\.gate_up_proj$": ao_config_experts,
        r"re:.*\.experts\.down_proj$":    ao_config_experts,
    })
)

quantize_(quantized_model, expert_fqn_config, filter_fn=None)

quantized_model.save_pretrained(OUTPUT_DIR)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
tokenizer.save_pretrained(OUTPUT_DIR)

# Smoke test
inputs = tokenizer("What are we having for dinner?", return_tensors="pt")
out = quantized_model.generate(**inputs, max_new_tokens=30, cache_implementation="static")
print(tokenizer.decode(out[0], skip_special_tokens=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.22.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                    # required for gpt-oss-20b (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/gpt-oss-20b-da8w8-torchao-v0.17.0,tokenizer=unsloth/gpt-oss-20b-BF16,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 .