gemma-4-31B-it-MXFP4

Model Overview

• Base model: google/gemma-4-31B-it
• Model architecture: Gemma4ForConditionalGeneration (text-only eval uses Gemma4ForCausalLM override)
• Quantizer: AMD Quark
• Quantization scheme: mxfp4
• Weight quantization: FP4, per-group (MX format)
• Activation quantization: FP4, dynamic, per-group (MX format)
• Export format: Hugging Face safetensors (quant_method=quark)
• Excluded from quantization: vision tower, multimodal projector, lm_head

Runtime Environment

This workflow is intended to run on AMD Instinct MI355.

Start Docker Container

docker run -it -d \
  --name vllm-gemma4-openai \
  --ipc=host \
  --shm-size=64g \
  --network=host \
  --privileged \
  --cap-add=CAP_SYS_ADMIN \
  --device=/dev/kfd \
  --device=/dev/dri \
  --device=/dev/mem \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  --entrypoint /bin/bash \
  -v /shareddata:/shareddata -v /mnt:/mnt -v /data:/data -v /home/:/home -v /raid:/raid -v /docker:/docker -v /datasets:/datasets -v /scratch:/scratch -v /group:/group \
vllm/vllm-openai-rocm:gemma4

Quantization

Run quantization inside Docker image vllm/vllm-openai-rocm:gemma4.

Minimal quantization code

from quark.common.profiler import GlobalProfiler
from quark.torch import LLMTemplate, ModelQuantizer

ckpt_path = "google/gemma-4-31B-it"
output_dir = "/path/to/amd/gemma-4-31B-it-mxfp4"

exclude_layers = ["*vision_tower*", "*multi_modal_projector*", "*lm_head"]

if "gemma4" not in LLMTemplate.list_available():
    gemma4_template = LLMTemplate(
        model_type="gemma4",
        kv_layers_name=["*language_model.*k_proj", "*language_model.*v_proj"],
        q_layer_name="*language_model.*q_proj",
        exclude_layers_name=exclude_layers,
    )
    LLMTemplate.register_template(gemma4_template)

template = LLMTemplate.get("gemma4")
quant_config = template.get_config(
    scheme="mxfp4",
    exclude_layers=exclude_layers,
)

profiler = GlobalProfiler(output_path=f"{output_dir}/quark_profile.yaml")
quantizer = ModelQuantizer(quant_config)
quantizer.direct_quantize_checkpoint(
    pretrained_model_path=ckpt_path,
    save_path=output_dir,
)

Evaluation

Reproduction command

CUDA_VISIBLE_DEVICES=6,7 lm_eval run \
  --model vllm \
  --tasks gsm8k \
  --model_args '{"pretrained":"amd/gemma-4-31B-it-mxfp4","tensor_parallel_size":2,"max_model_len":16384,"gpu_memory_utilization":0.90,"max_gen_toks":2048,"trust_remote_code":true,"chat_template":"/scratch/kewang/workspace/vllm/examples/tool_chat_template_gemma4.jinja","reasoning_parser":"gemma4","hf_overrides":{"architectures":["Gemma4ForCausalLM"]}}' \
  --num_fewshot 5 \
  --batch_size auto \
  --apply_chat_template

Accuracy

The following GSM8K results are from local runs in this workspace.

License

This checkpoint is derived from google/gemma-4-31B-it. Please follow the upstream model license and usage terms.