Ultra-FineWeb-12L-llama3-1B-Tokens
Trained on ~1,000,000,000 tokens across 1 dataset(s) using MaxText on a Kaggle TPU v5e-8 accelerator.
Architecture (Llama2/Llama3 decoder block)
| Key | Value |
|---|---|
| Decoder block | llama2 (RMSNorm + RoPE + SwiGLU, causal masking) |
| Hidden layers | 12 |
| Hidden size | 640 |
| Query heads | 5 (head_dim=128) |
| KV heads | 5 |
| Intermediate (MLP) | 1920 (SwiGLU) |
| Vocab size | 32,000 |
| Max position embeddings | 2048 |
| Tie word embeddings | false |
| Approx parameters | 104,858,240 (~104.9M) |
Training
Datasets (1 total, trained sequentially)
- [0] openbmb/Ultra-FineWeb
- Subset:
default - Train/eval split:
en/ `` - Text columns: ['content']
- Target tokens: 1,000,000,000 -> 3814 steps
- Subset:
Hyperparameters
- Tokenizer:
hf-internal-testing/llama-tokenizer - Optimizer: AdamW (b1=0.9, b2=0.95, eps=1e-08, weight_decay=0.1, decoupled)
- LR schedule: cosine, peak 0.0003, warmup 10%, final 0.1x peak
- Gradient clip: global-norm 1.0
- Z-loss: 1e-5 (auxiliary softmax z-loss for stability)
- Skip-step-on-spike: enabled
- Mixed precision: bf16 forward/backward, fp32 master weights & Adam state
- Batch: 8 sequences x 2048 tokens x 8 TPU chips x 2 grad_accum = 262,144 tokens/step
- Total steps: 3,814 (~1,000,000,000 tokens trained)
- Log period: every 10 steps
- Hardware: Kaggle TPU v5e-8 (single host, 8 chips, 16 GB HBM/chip)
- Framework: MaxText (JAX + Flax NNX + Optax + Orbax)
Tokenizer
The tokenizer from hf-internal-testing/llama-tokenizer is included in this repo.
Load with:
from transformers import LlamaTokenizer
tok = LlamaTokenizer.from_pretrained("Cion-lab/Ultra-FineWeb-12L-llama3-1B-Tokens")
Usage
from transformers import AutoModelForCausalLM, LlamaTokenizer
import torch
REPO = "Cion-lab/Ultra-FineWeb-12L-llama3-1B-Tokens"
m = AutoModelForCausalLM.from_pretrained(REPO, torch_dtype=torch.bfloat16)
tok = LlamaTokenizer.from_pretrained(REPO)
prompt = "Once upon a time, "
ids = tok(prompt, return_tensors="pt").input_ids
out = m.generate(ids, max_new_tokens=64, do_sample=True, temperature=0.8, top_p=0.9)
print(tok.decode(out[0], skip_special_tokens=True))
License
The model weights are released under CDLA-Sharing-1.0. The tokenizer follows its upstream licensing.
Citation
@misc{ultra-fineweb},
title = {Ultra-FineWeb},
howpublished = {https://huggingface.co/datasets/openbmb/Ultra-FineWeb}
@misc{maxtext,
title = {MaxText: A Simple, Performant and User-Friendly OpenSource LLM training Codebase},
author = {Google},
howpublished = {https://github.com/AI-Hypercomputer/maxtext}
}