Merge branch 'main' into kvcache3

Author: k223kim

Diff for: .github/workflows/cpu-tests.yml

@@ -88,11 +88,11 @@ jobs:
         continue-on-error: true
         with:
           path: .cache-HF
-          key: hf-cache-${{ runner.os }}-${{ matrix.python-version }}
+          key: hf-cache_${{ runner.os }}-py${{ matrix.python-version }}
           restore-keys: |
-            hf-cache-${{ runner.os }}-${{ matrix.python-version }}
-            hf-cache-${{ runner.os }}-
-            hf-cache-
+            hf-cache_${{ runner.os }}-py${{ matrix.python-version }}
+            hf-cache_${{ runner.os }}-
+            hf-cache_
 
       - name: Install dependencies
         run: |

Diff for: litgpt/config.py

@@ -1098,6 +1098,102 @@ def check_indicator_and_length(
 # Google Gemma 3
 ##################
 gemma3 = [
+    # https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json
+    dict(
+        name="Gemma-3-1b-it",
+        hf_config=dict(org="google", name="gemma-3-1b-it"),
+        scale_embeddings=True,
+        attention_scores_scalar=256,
+        vocab_size=262144,
+        block_size=131072,
+        sliding_window_size=512,
+        # 5 local layers for every global layer
+        sliding_window_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(26)],
+        intermediate_size=21504,
+        n_embd=1152,
+        n_layer=26,
+        n_head=4,
+        n_query_groups=1,
+        head_size=256,
+        rotary_percentage=1.0,
+        rope_adjustments=None,
+        parallel_residual=False,
+        bias=False,
+        norm_class_name="RMSNorm",
+        mlp_class_name="GemmaMLP",
+        gelu_approximate="tanh",
+        post_attention_norm=True,
+        post_mlp_norm=True,
+        norm_qk=True,
+        rope_base=1000000,
+        rope_local_base_freq=10000,
+        # 5 local layers for every global layer
+        rope_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(26)],
+    ),
+    # https://huggingface.co/google/gemma-3-4b-it/blob/main/config.json
+    dict(
+        name="Gemma-3-4b-it",
+        hf_config=dict(org="google", name="gemma-3-4b-it"),
+        scale_embeddings=True,
+        attention_scores_scalar=256,
+        vocab_size=262144,
+        block_size=131072,
+        sliding_window_size=1024,
+        # 5 local layers for every global layer
+        sliding_window_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(34)],
+        intermediate_size=10240,
+        n_embd=2560,
+        n_layer=34,
+        n_head=8,
+        n_query_groups=4,
+        head_size=256,
+        rotary_percentage=1.0,
+        rope_adjustments=dict(factor=8.0),
+        parallel_residual=False,
+        bias=False,
+        norm_class_name="RMSNorm",
+        mlp_class_name="GemmaMLP",
+        gelu_approximate="tanh",
+        post_attention_norm=True,
+        post_mlp_norm=True,
+        norm_qk=True,
+        rope_base=1000000,
+        rope_local_base_freq=10000,
+        # 5 local layers for every global layer
+        rope_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(34)],
+    ),
+    # https://huggingface.co/google/gemma-3-12b-it/blob/main/config.json
+    dict(
+        name="Gemma-3-12b-it",
+        hf_config=dict(org="google", name="gemma-3-12b-it"),
+        scale_embeddings=True,
+        attention_scores_scalar=256,
+        vocab_size=262144,
+        block_size=131072,
+        sliding_window_size=1024,
+        # 5 local layers for every global layer
+        sliding_window_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(48)],
+        intermediate_size=15360,
+        n_embd=3840,
+        n_layer=48,
+        n_head=16,
+        n_query_groups=8,
+        head_size=256,
+        rotary_percentage=1.0,
+        rope_adjustments=dict(factor=8.0),
+        parallel_residual=False,
+        bias=False,
+        norm_class_name="RMSNorm",
+        mlp_class_name="GemmaMLP",
+        gelu_approximate="tanh",
+        post_attention_norm=True,
+        post_mlp_norm=True,
+        norm_qk=True,
+        rope_base=1000000,
+        rope_local_base_freq=10000,
+        # 5 local layers for every global layer
+        rope_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(48)],
+    ),
     # https://huggingface.co/google/gemma-3-27b-it/blob/main/config.json
     dict(
         name="Gemma-3-27b-it",

Diff for: tests/test_model.py

@@ -812,7 +812,7 @@ def test_against_original_gemma_2(model_name, device, dtype):
 
 
 @torch.inference_mode()
-@pytest.mark.parametrize("model_name", ["gemma-3-27b-it"])
+@pytest.mark.parametrize("model_name", ["gemma-3-1b-it", "gemma-3-4b-it", "gemma-3-12b-it", "gemma-3-27b-it"])
 @pytest.mark.parametrize(
     ("device", "dtype"),
     [