Support nested Gemma 4 checkpoint layout

[ousamabenyounes]

Summary

• route Gemma4ForConditionalGeneration through its nested language-model layout
• keep Gemma 4 vision and optional per-layer embedding modules resident while streaming decoder layers
• preserve nested fused MoE expert tensors during checkpoint splitting and loading

Fixes #345

Test verification (RED -> GREEN)

On current upstream main (430adb1), the regression test routes Gemma 4 through the generic layout, the prerequisite for the reported nested-prefix parser failure:

AssertionError: ('airllm', 'AirLLMBaseModel') != ('airllm', 'AirLLMGemma4')
using generic AirLLM streaming model for architecture: Gemma4ForConditionalGeneration

On this rebased branch (b3a3c31), a tiny real Gemma4ForConditionalGeneration model splits, round-trips bit-exactly, and reloads with no missing or unexpected keys. The fixture covers nested decoder layers, fused MoE expert tensors, the vision tower, and PLE modules:

PYTHONPATH=air_llm pytest -q air_llm/tests/test_gemma4_split.py
1 passed

The ignored local CI runner is iso-baseline: upstream runs 46 tests with 2 pre-existing collection/import errors; this branch runs 47 tests with the same 2 errors. Package sdist/wheel builds and twine check pass, changed executable production lines have 14/14 coverage, and the visible GitHub check (GitGuardian Security Checks) passes on b3a3c31.

The official 51.6 GB checkpoint was not downloaded and CUDA 4-bit generation was not run; validation covers its published architecture and nested/fused module layout without making performance claims.

[ousamabenyounes]

Pushed follow-up 89d0bda: rebased onto current main, kept both Qwen4Exp and Gemma4 registry entries, and added Gemma4 PLE modules to the resident split/load map.

Validation: reproduced the PLE omission with air_llm/tests/test_gemma4_split.py RED (missing embed_tokens_per_layer, per_layer_model_projection, per_layer_projection_norm), then reran the focused test GREEN. Local run-ci.sh is iso-baseline: origin/main has 32 unittest cases with the same 2 existing import errors, final has 33 with the same 2 errors; package build and twine check pass; diff coverage is 11/11. GitGuardian passes on the pushed SHA.

[ousamabenyounes]

Rebased onto current main and resolved the lazy-export conflict while preserving the newer LoRA registrations.

Validation is RED -> GREEN on the current upstream base: Gemma 4 routes through the incorrect generic model before this patch, while the rebased branch passes a real tiny-model split, bit-exact round-trip, and reload test covering nested decoder layers, fused MoE experts, vision modules, and PLE modules. Local CI is iso-baseline (46 -> 47 tests with the same two pre-existing import errors); package build, twine check, 14/14 changed-line coverage, and the exact-SHA GitGuardian check pass.