From e95ddd168e9ff66a0d182ebf01db141aa0681474 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Fri, 15 May 2026 23:09:14 -0700 Subject: [PATCH 1/3] Clean up RoPE init + revert inference.py validator additions + add gemma4_31b CI - model.py: strip explanatory comment from materialize_runtime_buffers RoPE inv_freq block (keep hand-rolled formula as-is). - inference.py: revert all hf_validator + quant_compile_validator additions (--use-hf-api / --compare / --compare-quant / --prompts-file flags and their helpers); keep --bf16 HF checkpoint load path and existing prequantized / gguf flows. - .github/workflows/cuda.yml: add SocialLocalMobile/gemma-4-31B-it-HQQ-INT4 matrix entry (prequant tile-packed only) to export-model-cuda-artifact and test-model-cuda-e2e; pin to linux.aws.a100 like qwen3_5_moe. - .ci/scripts/export_model_artifact.sh: add gemma4_31b export branch mirroring qwen3_5_moe pattern. - .ci/scripts/test_model_e2e.sh: add gemma4_31b runner args + tokenizer handling. --- .ci/scripts/export_model_artifact.sh | 49 ++++++++++++++++++++++- .ci/scripts/test_model_e2e.sh | 19 ++++++++- .github/workflows/cuda.yml | 26 ++++++++++++- examples/models/gemma4_31b/inference.py | 52 +++++++++++++++---------- examples/models/gemma4_31b/model.py | 24 +++++++++++- 5 files changed, 143 insertions(+), 27 deletions(-) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 1f75d850e84..4bc8485dde8 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -195,9 +195,17 @@ case "$HF_MODEL" in PREPROCESSOR_FEATURE_SIZE="" PREPROCESSOR_OUTPUT="" ;; + SocialLocalMobile/gemma-4-31B-it-HQQ-INT4) + MODEL_NAME="gemma4_31b" + TASK="" + MAX_SEQ_LEN="" + EXTRA_PIP="" + PREPROCESSOR_FEATURE_SIZE="" + PREPROCESSOR_OUTPUT="" + ;; *) echo "Error: Unsupported model '$HF_MODEL'" - echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, openai/whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}, google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/diar_streaming_sortformer_4spk-v2, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4" + echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, openai/whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}, google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/diar_streaming_sortformer_4spk-v2, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4, SocialLocalMobile/gemma-4-31B-it-HQQ-INT4" exit 1 ;; esac @@ -459,6 +467,45 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then exit 0 fi +# Gemma 4 31B uses a prequantized checkpoint and custom export script +if [ "$MODEL_NAME" = "gemma4_31b" ]; then + pip install safetensors huggingface_hub gguf + + # Download prequantized model outside OUTPUT_DIR to avoid uploading on failure + LOCAL_MODEL_DIR=$(mktemp -d) + INDUCTOR_CACHE=$(mktemp -d) + trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE"' EXIT + + python -c "from huggingface_hub import snapshot_download; snapshot_download('${HF_MODEL}', local_dir='${LOCAL_MODEL_DIR}')" + + # Sanity check: run inference on the prequantized model + echo "::group::Inference sanity check" + python -m executorch.examples.models.gemma4_31b.inference \ + --prequantized "$LOCAL_MODEL_DIR" \ + --prompt "What is the capital of France?" \ + --max-new-tokens 32 \ + --temperature 0 \ + --no-compile + echo "::endgroup::" + + # Copy tokenizer for the runner + cp "$LOCAL_MODEL_DIR/tokenizer.json" "${OUTPUT_DIR}/tokenizer.json" + + # Export to .pte/.ptd (short cache dir avoids objcopy symbol length issues) + echo "::group::Export" + TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ + python -m executorch.examples.models.gemma4_31b.export \ + --prequantized "$LOCAL_MODEL_DIR" \ + --output-dir "${OUTPUT_DIR}" + echo "::endgroup::" + + test -f "${OUTPUT_DIR}/model.pte" + test -f "${OUTPUT_DIR}/aoti_cuda_blob.ptd" + ls -al "${OUTPUT_DIR}" + + exit 0 +fi + MAX_SEQ_LEN_ARG="" if [ -n "$MAX_SEQ_LEN" ]; then MAX_SEQ_LEN_ARG="--max_seq_len $MAX_SEQ_LEN" diff --git a/.ci/scripts/test_model_e2e.sh b/.ci/scripts/test_model_e2e.sh index 1678b0a4fbb..27b0dd9d597 100755 --- a/.ci/scripts/test_model_e2e.sh +++ b/.ci/scripts/test_model_e2e.sh @@ -228,9 +228,21 @@ case "$HF_MODEL" in AUDIO_FILE="" IMAGE_PATH="" ;; + SocialLocalMobile/gemma-4-31B-it-HQQ-INT4) + MODEL_NAME="gemma4_31b" + RUNNER_TARGET="gemma4_31b_runner" + RUNNER_PATH="gemma4_31b" + EXPECTED_OUTPUT="Paris" + PREPROCESSOR="" + TOKENIZER_URL="" + TOKENIZER_FILE="tokenizer.json" + AUDIO_URL="" + AUDIO_FILE="" + IMAGE_PATH="" + ;; *) echo "Error: Unsupported model '$HF_MODEL'" - echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, nvidia/diar_streaming_sortformer_4spk-v2, openai/whisper series (whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}), google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4" + echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, nvidia/diar_streaming_sortformer_4spk-v2, openai/whisper series (whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}), google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4, SocialLocalMobile/gemma-4-31B-it-HQQ-INT4" exit 1 ;; esac @@ -244,7 +256,7 @@ echo "::group::Prepare $MODEL_NAME Artifacts" # Download tokenizer files (skip for models that bundle tokenizer in export or do not use one) -if [ "$MODEL_NAME" != "parakeet" ] && [ "$MODEL_NAME" != "voxtral_realtime" ] && [ "$MODEL_NAME" != "sortformer" ] && [ "$MODEL_NAME" != "dinov2" ] && [ "$MODEL_NAME" != "qwen3_5_moe" ]; then +if [ "$MODEL_NAME" != "parakeet" ] && [ "$MODEL_NAME" != "voxtral_realtime" ] && [ "$MODEL_NAME" != "sortformer" ] && [ "$MODEL_NAME" != "dinov2" ] && [ "$MODEL_NAME" != "qwen3_5_moe" ] && [ "$MODEL_NAME" != "gemma4_31b" ]; then if [ "$TOKENIZER_FILE" != "" ]; then curl -L $TOKENIZER_URL/$TOKENIZER_FILE -o $MODEL_DIR/$TOKENIZER_FILE else @@ -368,6 +380,9 @@ EOF qwen3_5_moe) RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0 --cuda_graph" ;; + gemma4_31b) + RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0 --cuda_graph" + ;; voxtral_realtime) RUNNER_ARGS="--model_path ${MODEL_DIR}/model.pte --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --preprocessor_path ${MODEL_DIR}/$PREPROCESSOR --audio_path ${MODEL_DIR}/$AUDIO_FILE --temperature 0" # Add CUDA data path if present diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 087917c1116..eb7fc5a8939 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -185,6 +185,8 @@ jobs: name: "dinov2-small-imagenet1k-1-layer" - repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" + - repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" quant: - "non-quantized" - "quantized-int4-tile-packed" @@ -204,6 +206,15 @@ jobs: repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" quant: "quantized-int4-weight-only" + # Gemma 4 31B uses a prequantized checkpoint, only tile-packed + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "non-quantized" + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "quantized-int4-weight-only" # Voxtral Realtime only supports int4-tile-packed on CUDA - model: repo: "mistralai" @@ -258,7 +269,7 @@ jobs: with: timeout: 90 secrets-env: EXECUTORCH_HF_TOKEN - runner: ${{ matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} + runner: ${{ (matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' || matrix.model.name == 'gemma-4-31B-it-HQQ-INT4') && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} gpu-arch-type: cuda gpu-arch-version: 12.6 use-custom-docker-registry: false @@ -315,6 +326,8 @@ jobs: name: "dinov2-small-imagenet1k-1-layer" - repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" + - repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" quant: - "non-quantized" - "quantized-int4-tile-packed" @@ -334,6 +347,15 @@ jobs: repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" quant: "quantized-int4-weight-only" + # Gemma 4 31B uses a prequantized checkpoint, only tile-packed + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "non-quantized" + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "quantized-int4-weight-only" # Voxtral Realtime only supports int4-tile-packed on CUDA - model: repo: "mistralai" @@ -382,7 +404,7 @@ jobs: quant: "non-quantized" with: timeout: 90 - runner: ${{ matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} + runner: ${{ (matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' || matrix.model.name == 'gemma-4-31B-it-HQQ-INT4') && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} gpu-arch-type: cuda gpu-arch-version: 12.6 use-custom-docker-registry: false diff --git a/examples/models/gemma4_31b/inference.py b/examples/models/gemma4_31b/inference.py index 62dfe5956a7..788145fe791 100644 --- a/examples/models/gemma4_31b/inference.py +++ b/examples/models/gemma4_31b/inference.py @@ -6,12 +6,10 @@ """Eager inference on Gemma 4 31B-IT (CUDA + torch.compile). -Two input paths: +Three input paths: --prequantized Load a quantized checkpoint (from quantize_and_save.py). --gguf Load a GGUF file (e.g., Q4_K_M from the community). - -Packs for the target backend (--backend cuda), materializes runtime buffers, -optionally compiles with ``torch.compile``, and generates text autoregressively. + --bf16 Load the bf16 HF safetensors checkpoint via from_hf_checkpoint. Gemma 4 31B-IT is instruction-tuned and requires chat-template formatting. The ``--prompt`` is automatically wrapped with the Gemma 4 chat template @@ -38,7 +36,7 @@ import torch from executorch.examples.models.gemma4_31b.export import load_prequantized_model -from executorch.examples.models.gemma4_31b.model import materialize_runtime_buffers +from executorch.examples.models.gemma4_31b.model import Gemma4_31B, materialize_runtime_buffers def _move_to_cuda(model, config) -> None: @@ -147,6 +145,11 @@ def main() -> None: default=None, help="Path to a GGUF file (e.g., gemma-4-31B-it-Q4_K_M.gguf).", ) + src.add_argument( + "--bf16", + default=None, + help="Path to a bf16 hf directory (e.g., gemma-4-31B).", + ) parser.add_argument( "--tokenizer-path", default=None, @@ -192,12 +195,34 @@ def main() -> None: if args.backend == "cuda" and not torch.cuda.is_available(): parser.error("CUDA is required for the cuda backend.") + # ---- Tokenizer ---- + if args.tokenizer_path: + tokenizer_path = args.tokenizer_path + elif args.prequantized: + tokenizer_path = os.path.join(args.prequantized, "tokenizer.json") + elif args.bf16: + tokenizer_path = os.path.join(args.bf16, "tokenizer.json") + else: + parser.error("--tokenizer-path is required with --gguf.") + from tokenizers import Tokenizer + + tokenizer = Tokenizer.from_file(tokenizer_path) + + prompt_str = args.prompt if args.raw_prompt else apply_chat_template(args.prompt) + + # Gemma 4 EOS tokens (from generation_config.json: ids 1, 50, 106). + eos_token_ids = {1, 50, 106} + if args.gguf: from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model model, config = load_gguf_model( args.gguf, args.max_seq_len, backend=args.backend ) + elif args.bf16: + model, config = Gemma4_31B.from_hf_checkpoint( + args.bf16, max_seq_len=args.max_seq_len + ) else: print(f"Loading prequantized model from {args.prequantized}...") model, config = load_prequantized_model( @@ -212,21 +237,6 @@ def main() -> None: print("Compiling model with torch.compile...") model = torch.compile(model, mode="default") - if args.tokenizer_path: - tokenizer_path = args.tokenizer_path - elif args.prequantized: - tokenizer_path = os.path.join(args.prequantized, "tokenizer.json") - else: - parser.error("--tokenizer-path is required with --gguf.") - from tokenizers import Tokenizer - - tokenizer = Tokenizer.from_file(tokenizer_path) - - # Gemma 4 EOS tokens (from generation_config.json: ids 1, 50, 106). - eos_token_ids = {1, 50, 106} - - prompt = args.prompt if args.raw_prompt else apply_chat_template(args.prompt) - print(f"\nPrompt: {args.prompt}") print("-" * 40) @@ -234,7 +244,7 @@ def main() -> None: output = generate( model, tokenizer, - prompt, + prompt_str, max_new_tokens=args.max_new_tokens, temperature=args.temperature, eos_token_ids=eos_token_ids, diff --git a/examples/models/gemma4_31b/model.py b/examples/models/gemma4_31b/model.py index b0eb4004c52..b457c8807ca 100644 --- a/examples/models/gemma4_31b/model.py +++ b/examples/models/gemma4_31b/model.py @@ -675,7 +675,29 @@ def materialize_runtime_buffers( for layer in model.layers: attn = layer.self_attn - attn.inv_freq = attn.inv_freq.to(device) + if attn.is_sliding: + rotary_dim = attn.head_dim + else: + rotary_dim = int(attn.head_dim * attn.partial_rotary) + rope_angles = rotary_dim // 2 + inv_freq_rotated = 1.0 / ( + attn.rope_theta + ** ( + torch.arange(0, rotary_dim, 2, device=device, dtype=torch.float32) + / attn.head_dim + ) + ) + nope_angles = attn.head_dim // 2 - rope_angles + if nope_angles > 0: + inv_freq = torch.cat( + [ + inv_freq_rotated, + torch.zeros(nope_angles, device=device, dtype=torch.float32), + ] + ) + else: + inv_freq = inv_freq_rotated + attn.register_buffer("inv_freq", inv_freq, persistent=False) model.register_buffer( "embed_normalizer", From 7cffb3d036e7029b3e27cd6ce2912a880cdf8078 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 17 May 2026 23:45:12 -0700 Subject: [PATCH 2/3] lint --- .github/workflows/cuda.yml | 4 ---- examples/models/gemma4_31b/inference.py | 5 ++++- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index eb7fc5a8939..52bd1943cd6 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -148,10 +148,6 @@ jobs: # Run Qwen 3.5 MoE tests (quantize roundtrip + TurboQuant KV cache + sampler) python -m pytest examples/models/qwen3_5_moe/test_quantize_roundtrip.py examples/models/qwen3_5_moe/test_turboquant.py examples/models/qwen3_5_moe/test_sampler.py -v -o "addopts=" - # Run Gemma 4 31B tests (quant unit tests + pipeline integration tests) - pip install gguf - python -m pytest examples/models/gemma4_31b/quant/tests/ examples/models/gemma4_31b/tests/ -v -o "addopts=" - export-model-cuda-artifact: name: export-model-cuda-artifact # Skip this job if the pull request is from a fork (HuggingFace secrets are not available) diff --git a/examples/models/gemma4_31b/inference.py b/examples/models/gemma4_31b/inference.py index 788145fe791..535473161e1 100644 --- a/examples/models/gemma4_31b/inference.py +++ b/examples/models/gemma4_31b/inference.py @@ -36,7 +36,10 @@ import torch from executorch.examples.models.gemma4_31b.export import load_prequantized_model -from executorch.examples.models.gemma4_31b.model import Gemma4_31B, materialize_runtime_buffers +from executorch.examples.models.gemma4_31b.model import ( + Gemma4_31B, + materialize_runtime_buffers, +) def _move_to_cuda(model, config) -> None: From 86ba97b3d13df967ce80a943a1562b83e68f78f4 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 18 May 2026 00:02:31 -0700 Subject: [PATCH 3/3] lint --- examples/models/gemma4_31b/inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/models/gemma4_31b/inference.py b/examples/models/gemma4_31b/inference.py index 535473161e1..e1563c04ff6 100644 --- a/examples/models/gemma4_31b/inference.py +++ b/examples/models/gemma4_31b/inference.py @@ -75,7 +75,8 @@ def apply_chat_template(prompt: str) -> str: Does not include BOS — ``generate()`` prepends it at the token-ID level. """ return ( - "<|turn>user\n" + prompt + "<|turn>user\n" + + prompt + "\n<|turn>model\n<|channel>thought\n" )