tiny-qwen3-moe-w4a8 / README.md
sourroy's picture
Add model card
0a5677b verified
|
Raw
History Blame Contribute Delete
5.94 kB
---
library_name: transformers
license: apache-2.0
pipeline_tag: text-generation
tags: [vllm_ci]
---
# Model Overview
- **Model Architecture:** Qwen3MoeForCausalLM (tiny, randomly initialized)
- **Input:** Text
- **Output:** Text
- **Supported Hardware Microarchitecture:** AMD MI300 / MI350 / MI355 (gfx942 / gfx950)
- **Inference Engine:** [vLLM](https://docs.vllm.ai/en/latest/)
- **Model Optimizer:** [AMD-Quark](https://quark.docs.amd.com/latest/index.html)
- **Weight quantization:** W4A8 — INT4 weights (per-channel, symmetric) produced
via a progressive FP8→INT4 spec, following the
[amd/Kimi-K2.5-W4A8](https://huggingface.co/amd/Kimi-K2.5-W4A8) recipe
- **Activation quantization:** FP8 E4M3, per-tensor, dynamic
- **Quantized layers:** routed MoE experts only (attention, router/gate, and
`lm_head` are kept in the original precision)
This is a **tiny, randomly-initialized** Qwen3-MoE model quantized to W4A8, used
purely as **vLLM CI coverage** for the Quark W4A8 fused-MoE path
(`QuarkW4A8Fp8MoEMethod`), which dispatches through the ROCm AITER fused MoE
kernel. It is not intended to produce meaningful text — it exists so CI can load
a real W4A8 checkpoint and run a forward pass on GPU.
The dimensions (hidden `2048`, MoE intermediate `1024`, `8` experts, top-`2`) are
multiples of 256 so the AITER W4A8 shuffle/GEMM tile constraints hold. The
`vocab_size` matches the tokenizer so token ids stay within the embedding table.
# Model Creation
Built and quantized with [AMD-Quark](https://quark.docs.amd.com/latest/index.html),
following the progressive FP8→INT4 weight spec from the
[amd/Kimi-K2.5-W4A8](https://huggingface.co/amd/Kimi-K2.5-W4A8) model card.
> Note: Quark quantizes `nn.Linear` modules. MoE experts are stored as individual
> `nn.Linear` layers in `transformers` ~4.57; quantize with that version so the
> routed experts are captured.
```python
import argparse
import torch
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from quark.torch import ModelQuantizer, export_safetensors
from quark.torch.quantization.config.config import (
FP8E4M3PerTensorSpec,
Int4PerChannelSpec,
ProgressiveSpec,
QConfig,
QLayerConfig,
)
def get_config() -> QConfig:
# Quantize the routed experts only.
exclude_layers = ["*self_attn*", "*mlp.gate", "*lm_head"]
input_spec = FP8E4M3PerTensorSpec(
observer_method="min_max", scale_type="float", is_dynamic=True
).to_quantization_spec()
# Progressive FP8 -> INT4 weight spec (Kimi-K2.5-W4A8 recipe).
weight_spec = ProgressiveSpec(
first_stage=FP8E4M3PerTensorSpec(
observer_method="min_max", scale_type="float", is_dynamic=False
),
second_stage=Int4PerChannelSpec(
symmetric=True,
scale_type="float",
round_method="half_even",
is_dynamic=False,
ch_axis=0,
),
).to_quantization_spec()
return QConfig(
global_quant_config=QLayerConfig(input_tensors=input_spec, weight=weight_spec),
exclude=exclude_layers,
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--export-path", required=True)
parser.add_argument("--tokenizer", default="Qwen/Qwen1.5-MoE-A2.7B-Chat")
parser.add_argument("--hidden", type=int, default=2048)
parser.add_argument("--moe-intermediate", type=int, default=1024)
parser.add_argument("--experts", type=int, default=8)
parser.add_argument("--topk", type=int, default=2)
parser.add_argument("--layers", type=int, default=2)
parser.add_argument("--seed", type=int, default=0)
args = parser.parse_args()
torch.manual_seed(args.seed)
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
# vocab_size MUST cover the tokenizer, else real prompts produce token ids
# beyond the embedding table -> out-of-bounds embedding lookup (GPU fault).
cfg = AutoConfig.for_model(
"qwen3_moe",
hidden_size=args.hidden,
intermediate_size=args.hidden,
moe_intermediate_size=args.moe_intermediate,
num_hidden_layers=args.layers,
num_attention_heads=16,
num_key_value_heads=2,
head_dim=128,
num_experts=args.experts,
num_experts_per_tok=args.topk,
vocab_size=len(tokenizer),
max_position_embeddings=2048,
)
model = AutoModelForCausalLM.from_config(cfg).to("cuda").eval().to(torch.bfloat16)
ds = load_dataset("mit-han-lab/pile-val-backup", split="validation")
samples = [
tokenizer(ds[i]["text"], return_tensors="pt", truncation=True,
max_length=64).input_ids.to("cuda")
for i in range(8)
]
dataloader = DataLoader(samples, batch_size=1)
quantizer = ModelQuantizer(get_config())
with torch.no_grad():
model = quantizer.quantize_model(model, dataloader)
export_safetensors(
model, args.export_path, custom_mode="quark",
weight_format="real_quantized", pack_method="reorder",
)
tokenizer.save_pretrained(args.export_path)
# Symmetric INT4 export emits all-zero `*_zero_point_2` tensors that vLLM's
# W4A8 loader does not expect; drop them so the checkpoint loads directly.
if __name__ == "__main__":
main()
```
# Usage in vLLM
W4A8 dispatches through the ROCm AITER fused MoE kernel, so run on gfx942/gfx950
with AITER enabled:
```bash
VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 \
vllm serve amd/tiny-qwen3-moe-w4a8 --enforce-eager
```
Because the weights are random, outputs are not meaningful — this model is a
structural / smoke-test fixture only.
# License
Apache-2.0. The tiny model is randomly initialized and derives no weights from
any base model.
Modifications Copyright(c) 2026 Advanced Micro Devices, Inc. All rights reserved.