[model, ops] feat: add Gemma 4 E2B support

docs/transformers_v5/gemma4.md

@@ -0,0 +1,54 @@
+# Gemma 4
+
+The `gemma4` and `gemma4_text` registrations provide generated modeling for
+Gemma 4 E2B against Transformers 5.16.1. They retain per-layer embeddings,
+shared KV states and the Hugging Face checkpoint layout. The multimodal
+registration includes image, video and audio towers.
+
+## Attention and loss
+
+Choose attention through `OpsImplementationConfig`, for example
+`attn_implementation="flex_attention"`. The normalized HF interface name is
+`veomni_flex_attention`. Packed full and sliding masks use the public
+`veomni.ops.mask` builders with `cu_seq_lens_q` boundaries. Packed attention
+does not support `past_key_values`; unpacked masks retain HF cache handling.
+With `use_bidirectional_attention="all"`, full attention remains bidirectional
+within each packed sample. Sliding attention preserves HF's inclusive, symmetric
+window on both sides of each token; document boundaries still apply.
+
+Each LM instance holds its own `VeomniOp("cross_entropy_loss", "standard", impl)`
+as `veomni_ce`. Training calls its `loss_function` wrapper so channel-loss
+callbacks can observe both eager and fused loss selections. An uncapped model
+retains the selected loss implementation. Logit softcapping forces eager loss
+on capped logits and materializes logits in memory. Capped models reject
+`return_log_probs=True`, including the top-k distillation path, since the
+auxiliary helpers would reproject uncapped logits.
+
+Audio uses a boolean local mask independently of the LM attention backend.
+Vision Ulysses is unsupported: use `ulysses_size=1` for models retaining the
+vision tower. Such models reject Ulysses on text-only ranks too, before peers
+could enter different collectives.
+
+Under FSDP2, missing modalities participate through config-derived dummy
+forwards in a fixed image, video, audio order, with live compute dtypes and
+zero-valued projected graph dependencies.
+
+## Validation
+
+CPU regressions cover output/loss/gradient parity, packed sample isolation,
+channel-loss interception, audio visibility and dummy graph participation:
+
+```bash
+pytest tests/models/test_gemma4.py -q
+pytest tests/models/base/test_auto_registry.py -k gemma4 -q
+```
+
+Actual asymmetric BF16 FSDP2 validation requires two CUDA GPUs:
+
+```bash
+pytest tests/distributed/test_dummy_forward.py -k gemma4 -q
+```
+
+The three cases cover image versus text, video/audio versus text, and split
+modalities across ranks. They skip on CPU. Local checks do not establish
+collective correctness or full-checkpoint memory requirements.

docs/transformers_v5/index.md

@@ -5,6 +5,8 @@ This section documents VeOmni's integration with HuggingFace
 
 ## Included Notes
 
+- [Gemma 4](gemma4.md): model scope, instance-local loss selection, packed masks, and validation limits.
+
 - [Attention interface](veomni_attention_interface.md): documents instance-local `VeomniOp` dispatch, config routing, native BlockMask contract, and Ulysses restrictions.
 - [Flash Attention custom-name handling](veomni_flash_attention_op_adapter.md): explains why `_lazy_imports` fails for VeOmni custom attention names and how the local hub-kernel loader adapter resolves it.
 - [MoE weight loading](transformers_v5_moe_weight_loading.md): explains how VeOmni expects MoE expert weights to be laid out and documents qwen3_moe handling.
@@ -20,5 +22,6 @@ veomni_flash_attention_op_adapter.md
 veomni_attention_interface.md
 transformers_v5_moe_weight_loading.md
 testing_new_model.md
+gemma4.md
 upgrade_5_9_to_5_16.md
 ```

tests/distributed/test_dummy_forward.py

@@ -178,6 +178,10 @@ def _asymmetric_forward_worker(model_type, config_path, batch_fn):
     loss.backward()
 
     for name, param in model.named_parameters():
+        if model_type == "gemma4" and name.startswith(
+            ("model.vision_tower.", "model.audio_tower.", "model.embed_vision.", "model.embed_audio.")
+        ):
+            assert param.grad is not None, f"[Rank {rank}] Missing modality gradient for {name}"
         if param.grad is not None:
             assert torch.isfinite(param.grad).all(), f"[Rank {rank}] Non-finite gradient in {name}"
 
@@ -257,3 +261,35 @@ def test_asymmetric_forward_omni(model_type: str, config_path: str, batch_fn):
         partial(_asymmetric_forward_worker, model_type, config_path, batch_fn),
         world_size=2,
     )
+
+
+def _gemma4_batch(*, rank, device, dtype, rank_zero_modalities, rank_one_modalities):
+    from transformers.models.gemma4.configuration_gemma4 import Gemma4Config
+
+    from tests.tools.gemma4_utils import make_gemma4_batch
+
+    config = Gemma4Config.from_pretrained("./tests/toy_config/gemma4_toy")
+    modalities = rank_zero_modalities if rank == 0 else rank_one_modalities
+    return make_gemma4_batch(config, modalities, device=device, dtype=dtype)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two CUDA GPUs for FSDP2 validation")
+@pytest.mark.parametrize(
+    "rank_zero_modalities,rank_one_modalities",
+    [
+        pytest.param(("image",), (), id="image_vs_text"),
+        pytest.param(("video", "audio"), (), id="video_audio_vs_text"),
+        pytest.param(("image",), ("video", "audio"), id="split_modalities"),
+    ],
+)
+def test_asymmetric_forward_gemma4(rank_zero_modalities, rank_one_modalities):
+    """Validate actual BF16 FSDP collectives with real modality data on different ranks."""
+    from ..tools.launch_utils import torchrun
+
+    batch_fn = partial(
+        _gemma4_batch, rank_zero_modalities=rank_zero_modalities, rank_one_modalities=rank_one_modalities
+    )
+    torchrun(
+        partial(_asymmetric_forward_worker, "gemma4", "./tests/toy_config/gemma4_toy", batch_fn),
+        world_size=2,
+    )

tests/models/base/test_auto_registry.py

@@ -36,6 +36,7 @@
 from tests.models.tiny_configs import (
     tiny_gemma3_text_config as _tiny_gemma3_text_config,
 )
+from tests.models.tiny_configs import tiny_gemma4_config, tiny_gemma4_text_config
 from tests.models.tiny_configs import (
     tiny_glm_moe_dsa_config as _tiny_glm_moe_dsa_config,
 )
@@ -180,6 +181,16 @@ def _eager_ops(model_case: _ModelCase) -> tuple[tuple[str, str], ...]:
 
 
 _MODEL_CASES = (
+    _ModelCase(
+        model_type="gemma4",
+        config_factory=tiny_gemma4_config,
+        architectures=("Gemma4ForConditionalGeneration", "Gemma4Model"),
+    ),
+    _ModelCase(
+        model_type="gemma4_text",
+        config_factory=tiny_gemma4_text_config,
+        architectures=("Gemma4ForCausalLM", "Gemma4TextModel"),
+    ),
     _ModelCase(
         model_type="deepseek_v3",
         eager_ops=(