---
license: mit
language:
- en
tags:
- llama
- transformers
- text-generation
- from-scratch
- pretraining
- fineweb
- smollm
- gqa
- rope
library_name: transformers
pipeline_tag: text-generation
datasets:
- weights-and-wires/fineweb-6b
---

# smol-llama 🦙 (360M) — from-scratch LLaMA-style pretraining

This repository contains **smol-llama**, a **~360M parameter LLaMA-architecture causal LM** trained **from scratch** for **next-token prediction**.  
This project was primarily an educational + engineering effort to reproduce a **SmolLM-like training setup** at a smaller scale.

> TL;DR: Wanted to see if it's possible to *actually pretrain an LLM from scratch* end-to-end — data pipeline → tokenizer → training → checkpoint → Hub.

### Model Description

**smol-llama** is a compact implementation of the LLaMA architecture, featuring modern techniques like Grouped Query Attention (GQA), RoPE embeddings, and SwiGLU activations. It was trained on the [weights-and-wires/fineweb-6b](https://huggingface.co/datasets/weights-and-wires/fineweb-6b) dataset.

> NOTE: This is an early checkpoint / research-y base model. Expect imperfect generations.

---

### Model Architecture

| Component | Value |
|-----------|-------|
| **Parameters** | 360M |
| **Hidden Dimension** | 960 |
| **Layers** | 32 |
| **Attention Heads** | 15 (Query) / 5 (KV) |
| **Head Dimension** | 64 |
| **Context Length** | 2048 |
| **Vocabulary Size** | 49,152 |
| **Architecture** | LLaMA-style decoder-only |

**Key Features:**
- **Grouped Query Attention (GQA)**: 3:1 ratio for efficient inference
- **RoPE**: Rotary Position Embeddings for better length generalization
- **RMSNorm**: Root Mean Square Layer Normalization
- **SwiGLU**: Gated linear unit activation in FFN
- **Flash Attention 2**: Memory-efficient attention computation
- **Gradient Checkpointing**: Enables training larger batches

## Training Details

### Dataset

Trained on [weights-and-wires/fineweb-6b](https://huggingface.co/datasets/weights-and-wires/fineweb-6b), a curated subset of the FineWeb dataset containing ~6 billion high-quality web tokens.

### Training Hyperparameters

| Hyperparameter | Value |
|----------------|-------|
| **Optimizer** | AdamW (fused) |
| **Learning Rate** | 3e-4 (peak) |
| **LR Schedule** | Cosine with linear warmup |
| **Warmup Steps** | 900 |
| **Total Steps** | 5,725 (~1 epoch) |
| **Batch Size** | 64 |
| **Gradient Accumulation** | 8 |
| **Effective Batch Size** | 512 sequences |
| **Context Length** | 2048 tokens |
| **Tokens per Step** | ~1M |
| **Total Tokens** | ~6B |
| **Precision** | bfloat16 |
| **Gradient Clipping** | 1.0 |

### Infrastructure

| Resource | Specification |
|----------|--------------|
| **GPU** | 1× NVIDIA H100 (80GB PCIe) |
| **Training Time** | ~22 hours |
| **Throughput** | ~75,000 tokens/sec |
| **Cloud Provider** | RunPod |
| **Cost** | ~$53 (total) |

### Training Loss

The model was trained for one full epoch over the dataset with checkpoints saved every 200 steps. Final training loss: **~2.8** (see training checkpoints for intermediate metrics).

## Quick Start

### Installation

```bash
uv add torch transformers accelerate
```

### Basic Usage

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model and tokenizer
model_name = "weights-and-wires/smol-llama"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
  model_name,
  torch_dtype=torch.bfloat16,
  device_map="auto"
)

# Generate text
prompt = "The future of artificial intelligence is"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# Remove token_type_ids if present (not used by LLaMA models)
if 'token_type_ids' in inputs:
  del inputs['token_type_ids']

outputs = model.generate(
  **inputs,
  max_new_tokens=100,
  temperature=0.7,
  top_p=0.9,
  do_sample=True,
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```

### Advanced Generation

```python
# More controlled generation
outputs = model.generate(
  **inputs,
  max_new_tokens=200,
  temperature=0.8,
  top_k=50,
  top_p=0.95,
  repetition_penalty=1.1,
  do_sample=True,
  pad_token_id=tokenizer.eos_token_id,
)

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)
```

### Batch Generation

```python
prompts = [
  "Once upon a time",
  "The key to success is",
  "In the year 2050,",
]

inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)

outputs = model.generate(
  **inputs,
  max_new_tokens=50,
  temperature=0.7,
  do_sample=True,
  pad_token_id=tokenizer.eos_token_id,
)

for i, output in enumerate(outputs):
  print(f"\nPrompt {i+1}: {prompts[i]}")
  print(f"Generated: {tokenizer.decode(output, skip_special_tokens=True)}")
```

## Loading from Custom Checkpoint Format

If you want to load the original training checkpoints:

```python
import torch
from transformers import PreTrainedTokenizerFast

# Load tokenizer
tokenizer = PreTrainedTokenizerFast.from_pretrained("weights-and-wires/smol-llama")

# Load custom checkpoint
checkpoint_path = "training_checkpoints/checkpoint_step_5000.pt"
ckpt = torch.load(checkpoint_path, map_location="cuda")

# Create model from scratch (you'll need the model definition)
from utils.model import Llama, ModelArgs
model = Llama(ModelArgs()).cuda().to(torch.bfloat16)

# Handle torch.compile prefix if present
state_dict = {k.replace("_orig_mod.", ""): v for k, v in ckpt['model'].items()}
model.load_state_dict(state_dict)
model.eval()

# Generate
def generate(prompt, max_tokens=50):
  input_ids = tokenizer.encode(prompt, return_tensors="pt").cuda()
  
  with torch.no_grad():
    for _ in range(max_tokens):
      logits, _ = model(input_ids[:, -2048:])
      next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
      input_ids = torch.cat([input_ids, next_token], dim=1)
      if next_token.item() == tokenizer.eos_token_id:
        break
  
  return tokenizer.decode(input_ids[0])

print(generate("The meaning of life is"))
```

## Training Checkpoints

Intermediate training checkpoints are available in the `training_checkpoints/` folder:

| Checkpoint | Steps | Tokens Seen | Loss |
|------------|-------|-------------|------|
| `checkpoint_step_200.pt` | 200 | ~200M | - |
| `checkpoint_step_400.pt` | 400 | ~400M | - |
| ... | ... | ... | - |
| `checkpoint_step_4800.pt` | 4,800 | ~4.8B | - |
| `checkpoint_step_5000.pt` | 5,000 | ~5B | - |

These checkpoints include full training state (model, optimizer, step, loss) and can be used to resume training or analyze training dynamics.

## Limitations

This is a small model trained on a limited dataset (~6B tokens) for demonstration purposes. As such, it has several limitations:

- **Limited Knowledge**: The model has only seen 6B tokens, compared to 100B+ for larger models
- **Generalization**: May not perform well on out-of-distribution tasks
- **Factual Accuracy**: Should not be relied upon for factual information
- **Biases**: Inherits biases present in the web-scraped training data
- **No Instruction Tuning**: This is a base model without instruction following or chat capabilities
- **No Safety Alignment**: Has not undergone safety training or RLHF

## Intended Use

This model is intended for:
- Research and experimentation with small language models
- Educational purposes and learning about LLM pre-training
- Fine-tuning on downstream tasks
- Exploring efficient training techniques
- Prototyping and proof-of-concept projects

This model is **NOT** intended for:
- Production deployments without further fine-tuning
- Safety-critical applications
- Generating factual information without verification
- Applications requiring instruction following (use an instruction-tuned variant)

## Training Code

The complete pre-training code is available in the model repository. Key components:

```bash
# Clone the repository
git clone https://github.com/weights-and-wires/smol-llama
cd smol-llama

# Install dependencies
uv add sync

# Run training (requires GPU)
uv run pretrain.py
```

See the repository files for complete implementation details including:
- Custom LLaMA architecture (`utils/model.py`)
- Rotary embeddings (`utils/rotary.py`)
- Data loading utilities (`utils/data.py`)
- Checkpoint management (`utils/checkpoint.py`)
- Learning rate scheduling (`utils/lr_schedule.py`)

## Citation

If you use this model in your research, please cite:

```bibtex
@misc{smol-llama-2026,
  author = {Kashif, Ananya},
  title = {smol-llama: A 360M Parameter LLaMA Model Trained From Scratch},
  year = {2026},
  publisher = {HuggingFace},
  url = {https://huggingface.co/weights-and-wires/smol-llama}
}
```

Also consider citing the FineWeb dataset:

```bibtex
@inproceedings{
  penedo2024the,
  title={The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale},
  author={Guilherme Penedo and Hynek Kydl{\'\i}{\v{c}}ek and Loubna Ben allal and Anton Lozhkov and Margaret Mitchell and Colin Raffel and Leandro Von Werra and Thomas Wolf},
  booktitle={The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track},
  year={2024},
  url={https://openreview.net/forum?id=n6SCkn2QaG}
}
```

## Resources

- **Model Repository**: [weights-and-wires/smol-llama](https://huggingface.co/weights-and-wires/smol-llama)
- **Training Dataset**: [weights-and-wires/fineweb-6b](https://huggingface.co/datasets/weights-and-wires/fineweb-6b)
- **Reference Implementation**: [HuggingFaceTB/SmolLM-360M](https://huggingface.co/HuggingFaceTB/SmolLM-360M)

## License

This model is released under the MIT License. See the LICENSE file for details.

## Acknowledgments

- Inspired by [HuggingFaceTB/SmolLM-360M](https://huggingface.co/HuggingFaceTB/SmolLM-360M)
- Trained on [FineWeb](https://huggingface.co/datasets/HuggingFaceFW/fineweb) data
- Built with [PyTorch](https://pytorch.org/) and [Transformers](https://huggingface.co/transformers/)