Add Llama 3.2 1B

.gitignore

@@ -1,4 +1,5 @@
 tmp/
+models/
 checkpoints/
 
 *.bin

experiments/mini_transformer.py → experiments/02_mini_transformer.py

@@ -10,24 +10,26 @@
 # Reproduce https://sdbuchanan.com/blog/jax-2/
 
 from functools import partial
+
 import grain
 import jax
 import jax.numpy as jnp
 import numpy as np
 import optax
 from jax.nn.initializers import truncated_normal
+
 from julax.core import Learner, Trainer
 from julax.einops import Rearrange
 from julax.experiment import Experiment
 from julax.layers import (
     Chain,
-    Linear,
+    Embedding,
     LayerNorm,
+    Linear,
     Parallel,
-    Repeated,
+    Repeat,
+    Residual,
     RotaryEmbedding,
-    SkipConnection,
-    Embedding,
     Unembedding,
 )
 from julax.observers import default_observer
@@ -75,11 +77,11 @@ def main(
                         out_dim=dim,
                         w_init=truncated_normal(stddev=param_std),
                     ),
-                    blocks=Repeated(
+                    blocks=Repeat(
                         n=num_layers,
                         layer=Chain(
-                            attn=SkipConnection(
-                                layer=Chain(
+                            attn=Residual(
+                                processor=Chain(
                                     norm_attn=LayerNorm(dim=dim),
                                     attn=Chain(
                                         # qkv projection
@@ -130,8 +132,8 @@ def main(
                                     ),
                                 )
                             ),
-                            mlp=SkipConnection(
-                                layer=Chain(
+                            mlp=Residual(
+                                processor=Chain(
                                     norm_mlp=LayerNorm(dim=dim),
                                     mlp=Chain(
                                         up=Linear(