Instructions to use AIMadeSimpleResearch/AIMS-1-128M-base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use AIMadeSimpleResearch/AIMS-1-128M-base with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="AIMadeSimpleResearch/AIMS-1-128M-base", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("AIMadeSimpleResearch/AIMS-1-128M-base", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use AIMadeSimpleResearch/AIMS-1-128M-base with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "AIMadeSimpleResearch/AIMS-1-128M-base"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "AIMadeSimpleResearch/AIMS-1-128M-base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/AIMadeSimpleResearch/AIMS-1-128M-base

SGLang

How to use AIMadeSimpleResearch/AIMS-1-128M-base with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "AIMadeSimpleResearch/AIMS-1-128M-base" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "AIMadeSimpleResearch/AIMS-1-128M-base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "AIMadeSimpleResearch/AIMS-1-128M-base" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "AIMadeSimpleResearch/AIMS-1-128M-base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use AIMadeSimpleResearch/AIMS-1-128M-base with Docker Model Runner:
```
docker model run hf.co/AIMadeSimpleResearch/AIMS-1-128M-base
```

AIMS-1-128M-base / notebook_components.py

Samwise63

Upload verified custom BabyAIMS causal language model

afbd8e2 verified about 1 month ago

Raw History Blame Contribute Delete

11.8 kB

"""Original model computation from the supplied training notebook.
All computational definitions remain verbatim. AIMSModel has one lifecycle-only
addition to restore non-persistent buffers after loading; its forward is unchanged.
Hugging Face input/output handling lives in modeling_baby_aims.py.
"""

# Original notebook cell 17 (1-based).
import torch

def precompute_rope_params(head_dim, theta_base=10_000, context_length=4096, freq_config=None):
    assert head_dim % 2 == 0, "Embedding dimension must be even"

    # Compute the inverse frequencies
    inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2)[: (head_dim // 2)].float() / head_dim))

    ################################ NEW ###############################################
    # Frequency adjustments
    if freq_config is not None:
        low_freq_wavelen = freq_config["original_context_length"] / freq_config["low_freq_factor"]
        high_freq_wavelen = freq_config["original_context_length"] / freq_config["high_freq_factor"]

        wavelen = 2 * torch.pi / inv_freq

        inv_freq_llama = torch.where(
            wavelen > low_freq_wavelen, inv_freq / freq_config["factor"], inv_freq
        )

        smooth_factor = (freq_config["original_context_length"] / wavelen - freq_config["low_freq_factor"]) / (
            freq_config["high_freq_factor"] - freq_config["low_freq_factor"]
        )

        smoothed_inv_freq = (
            (1 - smooth_factor) * (inv_freq / freq_config["factor"]) + smooth_factor * inv_freq
        )

        is_medium_freq = (wavelen <= low_freq_wavelen) & (wavelen >= high_freq_wavelen)
        inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
        inv_freq = inv_freq_llama
    ####################################################################################


    # Generate position indices
    positions = torch.arange(context_length)

    # Compute the angles
    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)  # Shape: (context_length, head_dim // 2)

    # Expand angles to match the head_dim
    angles = torch.cat([angles, angles], dim=1)  # Shape: (context_length, head_dim)

    # Precompute sine and cosine
    cos = torch.cos(angles)
    sin = torch.sin(angles)

    return cos, sin


# Original notebook cell 18 (1-based).
def compute_rope(x, cos, sin):
    # x: (batch_size, num_heads, seq_len, head_dim)
    batch_size, num_heads, seq_len, head_dim = x.shape
    assert head_dim % 2 == 0, "Head dimension must be even"

    # Split x into first half and second half
    x1 = x[..., : head_dim // 2]  # First half
    x2 = x[..., head_dim // 2 :]  # Second half

    # Adjust sin and cos shapes
    cos = cos[:seq_len, :].unsqueeze(0).unsqueeze(0)  # Shape: (1, 1, seq_len, head_dim)
    sin = sin[:seq_len, :].unsqueeze(0).unsqueeze(0)

    # Apply the rotary transformation
    rotated = torch.cat((-x2, x1), dim=-1)
    x_rotated = (x * cos) + (rotated * sin)

    return x_rotated.to(dtype=x.dtype)


# Original notebook cell 19 (1-based).
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, emb_dim, eps=1e-5, dtype=None):
        super().__init__()
        self.eps = eps
        self.emb_dim = emb_dim
        self.weight = nn.Parameter(torch.ones(emb_dim, dtype=dtype))

    def forward(self, x):
        means = x.pow(2).mean(dim=-1, keepdim=True)
        x_normed = x * torch.rsqrt(means + self.eps)
        return (x_normed * self.weight).to(dtype=x.dtype)


# Original notebook cell 20 (1-based).
import torch.nn as nn


class GroupedQueryAttention(nn.Module):
    def __init__(
            self, d_in, d_out, num_heads,
            num_kv_groups,       # NEW
            dtype=None
        ):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
        assert num_heads % num_kv_groups == 0, "num_heads must be divisible by num_kv_groups"  # NEW

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads

        ############################# NEW  #############################
        # self.W_key = nn.Linear(d_in, d_out, bias=False, dtype=dtype)
        # self.W_value = nn.Linear(d_in, d_out, bias=False, dtype=dtype)
        self.W_key = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)
        self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)
        self.num_kv_groups = num_kv_groups
        self.group_size = num_heads // num_kv_groups
        ################################################################

        self.W_query = nn.Linear(d_in, d_out, bias=False, dtype=dtype)
        self.out_proj = nn.Linear(d_out, d_out, bias=False, dtype=dtype)


    def forward(self, x, cos=None, sin=None, is_causal=True):
        ##################### NEW  #####################
        # The forward method now accepts `mask` instead of accessing it via self.mask.
        # Also, we now have cos and sin as input for RoPE
        ################################################
        b, num_tokens, d_in = x.shape

        queries = self.W_query(x)  # Shape: (b, num_tokens, d_out)
        keys = self.W_key(x)  # Shape: (b, num_tokens, num_kv_groups * head_dim)
        values = self.W_value(x)  # Shape: (b, num_tokens, num_kv_groups * head_dim)

        # Reshape queries, keys, and values
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

        ##################### NEW  #####################
        # keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
        # values = values.view(b, num_tokens, self.num_heads, self.head_dim)
        keys = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim)
        values = values.view(b, num_tokens, self.num_kv_groups, self.head_dim)
        ################################################

        # Transpose keys, values, and queries
        keys = keys.transpose(1, 2)  # Shape: (b, num_kv_groups, num_tokens, head_dim)
        values = values.transpose(1, 2)  # Shape: (b, num_kv_groups, num_tokens, head_dim)
        queries = queries.transpose(1, 2)  # Shape: (b, num_heads, num_tokens, head_dim)

        ##################### NEW #####################
        # Apply RoPE
        if cos is not None:
            keys = compute_rope(keys, cos, sin)
            queries = compute_rope(queries, cos, sin)
        ################################################

        ##################### NEW  #####################
        # Expand keys and values to match the number of heads
        # Shape: (b, num_heads, num_tokens, head_dim)

        keys = keys.repeat_interleave(self.group_size, dim=1)  # Shape: (b, num_heads, num_tokens, head_dim)
        values = values.repeat_interleave(self.group_size, dim=1)  # Shape: (b, num_heads, num_tokens, head_dim)
        # For example, before repeat_interleave along dim=1 (query groups):
        #   [K1, K2]
        # After repeat_interleave (each query group is repeated group_size times):
        #   [K1, K1, K2, K2]
        # If we used regular repeat instead of repeat_interleave, we'd get:
        #   [K1, K2, K1, K2]
        ################################################

        # Memory-efficient causal scaled dot-product gqa attention
        context_vec = nn.functional.scaled_dot_product_attention(
            queries,
            keys,
            values,
            is_causal=is_causal,
            #enable_gqa=True
        )

        # Shape: (b, num_tokens, num_heads, head_dim)
        context_vec = context_vec.transpose(1, 2)

        # Combine heads, where self.d_out = self.num_heads * self.head_dim
        context_vec = context_vec.reshape(b, num_tokens, self.d_out)
        context_vec = self.out_proj(context_vec)  # optional projection

        return context_vec


# Original notebook cell 21 (1-based).
import torch
import torch.nn as nn

# FeedForward from Qwen 3

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.fc1 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
        self.fc2 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
        self.fc3 = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], dtype=cfg["dtype"], bias=False)

    def forward(self, x):
        x_fc1 = self.fc1(x)
        x_fc2 = self.fc2(x)
        x = nn.functional.silu(x_fc1) * x_fc2
        return self.fc3(x)


# Original notebook cell 22 (1-based).
class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.att =  GroupedQueryAttention(  # MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            num_heads=cfg["n_heads"],
            num_kv_groups=cfg["n_kv_groups"],  # NEW
            dtype=cfg["dtype"]
        )

        self.ff = FeedForward(cfg)
        self.norm1 = RMSNorm(cfg["emb_dim"], eps=1e-5, dtype=cfg["dtype"])
        self.norm2 = RMSNorm(cfg["emb_dim"], eps=1e-5, dtype=cfg["dtype"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

        # Store cfg as an instance variable
        self.cfg = cfg

    # Add dtype
    def forward(self, x, cos=None, sin=None, is_causal=True):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x.to(self.cfg["dtype"]), cos, sin, is_causal=is_causal)   # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back


        # Shortcut connection for feed-forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x.to(self.cfg["dtype"]))
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x


# Original notebook cell 23 (1-based).
class AIMSModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"], dtype=cfg["dtype"])

        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

        ##################### NEW  #####################
        self.final_norm = RMSNorm(cfg["emb_dim"], dtype=cfg["dtype"])
        ################################################
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=cfg["dtype"]
        )

        self.cfg = cfg
        self.restore_runtime_buffers()

    @torch.no_grad()
    def restore_runtime_buffers(self):
        """Rebuild non-persistent RoPE tables after HF materializes model weights.
        This lifecycle-only addition leaves forward and the RoPE math unchanged.
        CPU construction guarantees real FP32 tables even in a meta-device context.
        """
        device = self.tok_emb.weight.device
        with torch.device("cpu"):
            cos, sin = precompute_rope_params(
                head_dim=self.cfg["emb_dim"] // self.cfg["n_heads"],
                theta_base=self.cfg["rope_base"],
                context_length=self.cfg["context_length"],
                freq_config=self.cfg["rope_freq"]
            )
        self.register_buffer("cos", cos.to(device=device), persistent=False)
        self.register_buffer("sin", sin.to(device=device), persistent=False)

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape

        tok_embeds = self.tok_emb(in_idx)
        x = tok_embeds

        num_tokens = x.shape[1]

        for block in self.trf_blocks:
            x = block(x, self.cos, self.sin, is_causal=True)

        x = self.final_norm(x)
        logits = self.out_head(x.to(self.cfg["dtype"]))
        return logits

	"""Original model computation from the supplied training notebook.

	All computational definitions remain verbatim. AIMSModel has one lifecycle-only
	addition to restore non-persistent buffers after loading; its forward is unchanged.
	Hugging Face input/output handling lives in modeling_baby_aims.py.
	"""

	# Original notebook cell 17 (1-based).
	import torch

	def precompute_rope_params(head_dim, theta_base=10_000, context_length=4096, freq_config=None):
	assert head_dim % 2 == 0, "Embedding dimension must be even"

	# Compute the inverse frequencies
	inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2)[: (head_dim // 2)].float() / head_dim))

	################################ NEW ###############################################
	# Frequency adjustments
	if freq_config is not None:
	low_freq_wavelen = freq_config["original_context_length"] / freq_config["low_freq_factor"]
	high_freq_wavelen = freq_config["original_context_length"] / freq_config["high_freq_factor"]

	wavelen = 2 * torch.pi / inv_freq

	inv_freq_llama = torch.where(
	wavelen > low_freq_wavelen, inv_freq / freq_config["factor"], inv_freq
	)

	smooth_factor = (freq_config["original_context_length"] / wavelen - freq_config["low_freq_factor"]) / (
	freq_config["high_freq_factor"] - freq_config["low_freq_factor"]
	)

	smoothed_inv_freq = (
	(1 - smooth_factor) * (inv_freq / freq_config["factor"]) + smooth_factor * inv_freq
	)

	is_medium_freq = (wavelen <= low_freq_wavelen) & (wavelen >= high_freq_wavelen)
	inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
	inv_freq = inv_freq_llama
	####################################################################################


	# Generate position indices
	positions = torch.arange(context_length)

	# Compute the angles
	angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) # Shape: (context_length, head_dim // 2)

	# Expand angles to match the head_dim
	angles = torch.cat([angles, angles], dim=1) # Shape: (context_length, head_dim)

	# Precompute sine and cosine
	cos = torch.cos(angles)
	sin = torch.sin(angles)

	return cos, sin


	# Original notebook cell 18 (1-based).
	def compute_rope(x, cos, sin):
	# x: (batch_size, num_heads, seq_len, head_dim)
	batch_size, num_heads, seq_len, head_dim = x.shape
	assert head_dim % 2 == 0, "Head dimension must be even"

	# Split x into first half and second half
	x1 = x[..., : head_dim // 2] # First half
	x2 = x[..., head_dim // 2 :] # Second half

	# Adjust sin and cos shapes
	cos = cos[:seq_len, :].unsqueeze(0).unsqueeze(0) # Shape: (1, 1, seq_len, head_dim)
	sin = sin[:seq_len, :].unsqueeze(0).unsqueeze(0)

	# Apply the rotary transformation
	rotated = torch.cat((-x2, x1), dim=-1)
	x_rotated = (x * cos) + (rotated * sin)

	return x_rotated.to(dtype=x.dtype)


	# Original notebook cell 19 (1-based).
	import torch.nn as nn

	class RMSNorm(nn.Module):
	def __init__(self, emb_dim, eps=1e-5, dtype=None):
	super().__init__()
	self.eps = eps
	self.emb_dim = emb_dim
	self.weight = nn.Parameter(torch.ones(emb_dim, dtype=dtype))

	def forward(self, x):
	means = x.pow(2).mean(dim=-1, keepdim=True)
	x_normed = x * torch.rsqrt(means + self.eps)
	return (x_normed * self.weight).to(dtype=x.dtype)


	# Original notebook cell 20 (1-based).
	import torch.nn as nn


	class GroupedQueryAttention(nn.Module):
	def __init__(
	self, d_in, d_out, num_heads,
	num_kv_groups, # NEW
	dtype=None
	):
	super().__init__()
	assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
	assert num_heads % num_kv_groups == 0, "num_heads must be divisible by num_kv_groups" # NEW

	self.d_out = d_out
	self.num_heads = num_heads
	self.head_dim = d_out // num_heads

	############################# NEW #############################
	# self.W_key = nn.Linear(d_in, d_out, bias=False, dtype=dtype)
	# self.W_value = nn.Linear(d_in, d_out, bias=False, dtype=dtype)
	self.W_key = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)
	self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)
	self.num_kv_groups = num_kv_groups
	self.group_size = num_heads // num_kv_groups
	################################################################

	self.W_query = nn.Linear(d_in, d_out, bias=False, dtype=dtype)
	self.out_proj = nn.Linear(d_out, d_out, bias=False, dtype=dtype)


	def forward(self, x, cos=None, sin=None, is_causal=True):
	##################### NEW #####################
	# The forward method now accepts `mask` instead of accessing it via self.mask.
	# Also, we now have cos and sin as input for RoPE
	################################################
	b, num_tokens, d_in = x.shape

	queries = self.W_query(x) # Shape: (b, num_tokens, d_out)
	keys = self.W_key(x) # Shape: (b, num_tokens, num_kv_groups * head_dim)
	values = self.W_value(x) # Shape: (b, num_tokens, num_kv_groups * head_dim)

	# Reshape queries, keys, and values
	queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

	##################### NEW #####################
	# keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
	# values = values.view(b, num_tokens, self.num_heads, self.head_dim)
	keys = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim)
	values = values.view(b, num_tokens, self.num_kv_groups, self.head_dim)
	################################################

	# Transpose keys, values, and queries
	keys = keys.transpose(1, 2) # Shape: (b, num_kv_groups, num_tokens, head_dim)
	values = values.transpose(1, 2) # Shape: (b, num_kv_groups, num_tokens, head_dim)
	queries = queries.transpose(1, 2) # Shape: (b, num_heads, num_tokens, head_dim)

	##################### NEW #####################
	# Apply RoPE
	if cos is not None:
	keys = compute_rope(keys, cos, sin)
	queries = compute_rope(queries, cos, sin)
	################################################

	##################### NEW #####################
	# Expand keys and values to match the number of heads
	# Shape: (b, num_heads, num_tokens, head_dim)

	keys = keys.repeat_interleave(self.group_size, dim=1) # Shape: (b, num_heads, num_tokens, head_dim)
	values = values.repeat_interleave(self.group_size, dim=1) # Shape: (b, num_heads, num_tokens, head_dim)
	# For example, before repeat_interleave along dim=1 (query groups):
	# [K1, K2]
	# After repeat_interleave (each query group is repeated group_size times):
	# [K1, K1, K2, K2]
	# If we used regular repeat instead of repeat_interleave, we'd get:
	# [K1, K2, K1, K2]
	################################################

	# Memory-efficient causal scaled dot-product gqa attention
	context_vec = nn.functional.scaled_dot_product_attention(
	queries,
	keys,
	values,
	is_causal=is_causal,
	#enable_gqa=True
	)

	# Shape: (b, num_tokens, num_heads, head_dim)
	context_vec = context_vec.transpose(1, 2)

	# Combine heads, where self.d_out = self.num_heads * self.head_dim
	context_vec = context_vec.reshape(b, num_tokens, self.d_out)
	context_vec = self.out_proj(context_vec) # optional projection

	return context_vec


	# Original notebook cell 21 (1-based).
	import torch
	import torch.nn as nn

	# FeedForward from Qwen 3

	class FeedForward(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.fc1 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
	self.fc2 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
	self.fc3 = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], dtype=cfg["dtype"], bias=False)

	def forward(self, x):
	x_fc1 = self.fc1(x)
	x_fc2 = self.fc2(x)
	x = nn.functional.silu(x_fc1) * x_fc2
	return self.fc3(x)


	# Original notebook cell 22 (1-based).
	class TransformerBlock(nn.Module):
	def __init__(self, cfg):
	super().__init__()

	self.att = GroupedQueryAttention( # MultiHeadAttention(
	d_in=cfg["emb_dim"],
	d_out=cfg["emb_dim"],
	num_heads=cfg["n_heads"],
	num_kv_groups=cfg["n_kv_groups"], # NEW
	dtype=cfg["dtype"]
	)

	self.ff = FeedForward(cfg)
	self.norm1 = RMSNorm(cfg["emb_dim"], eps=1e-5, dtype=cfg["dtype"])
	self.norm2 = RMSNorm(cfg["emb_dim"], eps=1e-5, dtype=cfg["dtype"])
	self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

	# Store cfg as an instance variable
	self.cfg = cfg

	# Add dtype
	def forward(self, x, cos=None, sin=None, is_causal=True):
	# Shortcut connection for attention block
	shortcut = x
	x = self.norm1(x)
	x = self.att(x.to(self.cfg["dtype"]), cos, sin, is_causal=is_causal) # Shape [batch_size, num_tokens, emb_size]
	x = self.drop_shortcut(x)
	x = x + shortcut # Add the original input back


	# Shortcut connection for feed-forward block
	shortcut = x
	x = self.norm2(x)
	x = self.ff(x.to(self.cfg["dtype"]))
	x = self.drop_shortcut(x)
	x = x + shortcut # Add the original input back

	return x


	# Original notebook cell 23 (1-based).
	class AIMSModel(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"], dtype=cfg["dtype"])

	self.trf_blocks = nn.Sequential(
	*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

	##################### NEW #####################
	self.final_norm = RMSNorm(cfg["emb_dim"], dtype=cfg["dtype"])
	################################################
	self.out_head = nn.Linear(
	cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=cfg["dtype"]
	)

	self.cfg = cfg
	self.restore_runtime_buffers()

	@torch.no_grad()
	def restore_runtime_buffers(self):
	"""Rebuild non-persistent RoPE tables after HF materializes model weights.

	This lifecycle-only addition leaves forward and the RoPE math unchanged.
	CPU construction guarantees real FP32 tables even in a meta-device context.
	"""
	device = self.tok_emb.weight.device
	with torch.device("cpu"):
	cos, sin = precompute_rope_params(
	head_dim=self.cfg["emb_dim"] // self.cfg["n_heads"],
	theta_base=self.cfg["rope_base"],
	context_length=self.cfg["context_length"],
	freq_config=self.cfg["rope_freq"]
	)
	self.register_buffer("cos", cos.to(device=device), persistent=False)
	self.register_buffer("sin", sin.to(device=device), persistent=False)

	def forward(self, in_idx):
	batch_size, seq_len = in_idx.shape

	tok_embeds = self.tok_emb(in_idx)
	x = tok_embeds

	num_tokens = x.shape[1]

	for block in self.trf_blocks:
	x = block(x, self.cos, self.sin, is_causal=True)

	x = self.final_norm(x)
	logits = self.out_head(x.to(self.cfg["dtype"]))
	return logits