Skip to content

fix(gemma4): load mixed compressed-tensors nvfp4 - #493

Open
Poyo-Yon wants to merge 1 commit into
FlashML-org:mainfrom
Poyo-Yon:codex/fix-gemma4-compressed-tensors-nvfp4
Open

Poyo-Yon wants to merge 1 commit into
FlashML-org:mainfrom
Poyo-Yon:codex/fix-gemma4-compressed-tensors-nvfp4

Conversation

@Poyo-Yon

Copy link
Copy Markdown

Fixes #492

Summary

Teach the Gemma 4 checkpoint reader to load the mixed compressed-tensors layout used by unsloth/gemma-4-26B-A4B-it-NVFP4:

  • fuse per-channel FP8 q/k/v weight_scale roles alongside their weights and normalize [out_features, 1] bf16 scales to FreeToken's flat fp32 buffers;
  • map compressed-tensors NVFP4 shared-MLP roles through the existing nvfp4_parts_ct / ct_nvfp4_fuse helpers;
  • build the NVFP4 routed-expert source spec from the installed checkpoint dialect, including reciprocal compressed-tensors global scales;
  • ignore static KV quantizer metadata which is not part of the runtime KV cache.

The existing ModelOpt Gemma 4 path is retained unchanged.

Tests

Weight-free regression coverage:

pytest -q \
  tests/models/test_gemma4_compressed_tensors.py \
  tests/models/test_gemma4_vision.py \
  tests/models/test_gemma4_gguf_rope.py

6 passed, 4 skipped in 20.93s

Real-hardware validation:

  • OS: Windows 11 + Ubuntu 24.04 under WSL2
  • GPU: NVIDIA GeForce RTX 5060 Ti 16 GB, driver 616.92
  • Other GPU: NVIDIA GeForce RTX 5070 Ti 16 GB
  • FreeToken base: cac247a (0.1.3)
  • Checkpoint: unsloth/gemma-4-26B-A4B-it-NVFP4
  • Mode: text-only, 4096 context, NVFP4 expert offload, 40% expert cache
  • Result: model load, serial expert-bank build, 52,597-token KV-cache allocation, CUDA graph capture, /v1/models, and /v1/chat/completions all completed
  • Three post-warm-up requests: 34.48, 59.44, and 60.65 output tok/s (median 59.44 tok/s)

Exact serve command is included in the linked issue.

AI assistance was used to prepare the patch and tests. The contributor reviewed the behavior and ran the real-checkpoint validation described above.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Gemma 4 reader cannot load mixed compressed-tensors NVFP4 checkpoints

1 participant