ferjorosa/tiny-lm-swallow-code-8k-ibis-16

Swallow Code Ibis-16 (LLaMA 3, 16 layers, 8k vocab)

This model was trained with the tiny-lm repository on the SwallowCode dataset.

The goal is educational: a compact pretraining run for studying tokenization, data pipelines, and transformer training end to end.

Data

The tokenizer is a custom 8k-token BPE trained on SwallowCode with Karpathy's rustbpe approach, then exported as a tiktoken encoding for inference.

The dataset was split into 99% train and 1% validation before tokenization. The resulting tokenized files contain 25.95B train tokens and 262M validation tokens. Training used 1,024-token contiguous windows over the token stream.

Training

The model was trained with PyTorch Lightning using bf16 mixed precision.

• Context window: 1,024 tokens
• Batch size: 64 sequences
• Gradient accumulation: 4
• Effective batch size: 262,144 tokens per optimizer step
• Training budget: about 25B tokens over 95,350 optimizer steps
• Optimizer: AdamW with cosine LR decay and 1% warmup
• Peak LR: 6e-4
• Weight decay: 0.1

Architecture

• LLaMA 3 style decoder-only transformer
• Parameters: 18,886,912 total; 16,789,760 non-embedding
• Layers: 16
• Vocab size: 8192
• Context length: 1024
• d_model: 256
• n_heads: 4
• n_kv_heads: 4
• ffn_hidden_dim: 1024
• RoPE theta: 10000.0
• Norm epsilon: 1e-05

Files

• model.safetensors: model weights (SafeTensors)
• model_config.yaml: tiny-lm model config
• tokenizer.pkl: tiktoken encoding
• tokenizer_config.yaml: tokenizer settings (BOS/EOS)
• checkpoint.ckpt: original Lightning checkpoint

Usage

This is a tiny-lm model (not Transformers-compatible). Load it with tiny-lm:

import pickle
import torch
from tiny_lm.model.llama3 import Llama3
from tiny_lm.model.config import Llama3Config

config = Llama3Config.from_yaml("model_config.yaml")
model = Llama3(
    vocab_size=config.vocab_size,
    d_model=config.d_model,
    n_layers=config.n_layers,
    n_heads=config.n_heads,
    context_length=config.context_length,
    n_kv_heads=config.n_kv_heads,
    ffn_hidden_dim=config.ffn_hidden_dim,
    multiple_of=config.multiple_of,
    rope_theta=config.rope_theta,
    norm_eps=config.norm_eps,
    emb_dropout=0.0,
    attn_dropout=0.0,
    resid_dropout=0.0,
    ffn_dropout=0.0,
    qkv_bias=config.qkv_bias,
    ffn_bias=config.ffn_bias,
    attn_backend=config.attn_backend,
)
from safetensors.torch import load_file as load_safetensors

state = load_safetensors("model.safetensors")
model.load_state_dict(state, strict=True)
model.eval()

with open("tokenizer.pkl", "rb") as f:
    tokenizer = pickle.load(f)