[2/4] add gemma 3 1b (#2000)

Author: k223kim

litgpt/config.py

@@ -1058,6 +1058,38 @@ def norm_class(self) -> Type:
 # 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-27b-it/blob/main/config.json
     dict(
         name="Gemma-3-27b-it",

tests/test_model.py

@@ -802,7 +802,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-27b-it"])
 @pytest.mark.parametrize(
     ("device", "dtype"),
     [