README.md

---
language:
  - en
  - ja
tags:
  - llama3
  - swallow-code
  - tiny-lm
  - code
datasets:
  - tokyotech-llm/swallow-code
---

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

This model was trained with the
[tiny-lm](https://github.com/ferjorosa/tiny-lm) repository on the
[SwallowCode dataset](https://huggingface.co/datasets/tokyotech-llm/swallow-code).

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`](https://github.com/karpathy/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:

```python
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)
```