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

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}
}