Motivation
Gemma 4 26B MoE checkpoints are not yet covered by the current Gemma 4 multimodal adapter. A TP8 GSPO smoke test with synchronized video and audio reaches multimodal feature construction but fails before rollout with:
ValueError: Gemma4 audio token/feature count mismatch: tokens=2000 features=0
The adapter should support the 26B MoE architecture without weakening the existing media-token validation.
Proposed feature
Add Gemma 4 26B MoE support for multimodal serving and rollout-based training, including:
- image, audio, video, and combined-media inputs supported by the checkpoint processor;
- TP8 serving and GSPO training with Adam 8-bit;
- media feature extraction and token replacement for the MoE model layout;
- tower and projector/merger frozen by default;
--unfreeze-mm-tower and --unfreeze-mm-projector/--unfreeze-mm-merger;
- independent tower and projector LR, minimum LR, decay steps, and decay style;
- policy synchronization for unfrozen media parameters;
- complete native checkpoint save/load for text, tower, and projector parameters.
Representative reproduction:
areno train \
--ckpt /path/to/gemma4-26b-moe \
--dataset-path /path/to/ave.jsonl \
--dataset-loader-fn examples/multimodal/ave_event_recognition/dataset_loader.py \
--reward-fn-path examples/multimodal/ave_event_recognition/reward.py \
--agent-fn examples/multimodal/ave_event_recognition/run_agent.py \
--algo gspo --world-size 8 --tp-size 8 \
--batch-size 1 --n-samples 8 --mini-bs 1 \
--max-running-prompts 8 --adam-8bit --max-steps 2
Acceptance criteria
- A combined video+audio request completes in
areno serve and GSPO without media token/feature mismatches.
- Frozen-mode tests verify tower/projector tensors receive no gradients and remain byte-equivalent after save/reload.
- Unfrozen-mode tests verify nonzero tower/projector gradients and independent LR schedules.
- A checkpoint round trip verifies modified tower/projector tensors are preserved.
- TP8 text-only behavior remains unchanged.
Related work
Motivation
Gemma 4 26B MoE checkpoints are not yet covered by the current Gemma 4 multimodal adapter. A TP8 GSPO smoke test with synchronized video and audio reaches multimodal feature construction but fails before rollout with:
The adapter should support the 26B MoE architecture without weakening the existing media-token validation.
Proposed feature
Add Gemma 4 26B MoE support for multimodal serving and rollout-based training, including:
--unfreeze-mm-towerand--unfreeze-mm-projector/--unfreeze-mm-merger;Representative reproduction:
Acceptance criteria
areno serveand GSPO without media token/feature mismatches.Related work