I Pretrained a 360M LLaMA-Style Language Model from Scratch on 6B FineWeb Tokens (Single H100)

Pretraining an LLM from scratch usually sounds like “big-lab-only” territory. I wanted to test how far a smaller, practical setup can go while keeping the process transparent and reproducible.

This post documents an end-to-end run of training a ~360M parameter LLaMA-style model on a 6B-token FineWeb subset, using a single H100 GPU. The objective was not frontier SOTA, but a robust baseline: clear pipeline, stable training, and reproducible artifacts.

TL;DR

• Built and trained a 360M LLaMA-style decoder from scratch.

• Used FineWeb streaming (no full download), stopped at ~6B tokens.

• Trained on 1x H100, context length 2048, bf16, AdamW + warmup + cosine decay.

• Logged core dynamics (train/loss, train/lr, train/tps, val/loss, val/pplx).

• Missed some telemetry (grad norm, NaN counters, step-time histograms), which will be added next run.

1) Why This Project

I wanted to understand the full pretraining stack by doing it end-to-end:

• data construction under storage constraints,

• tokenizer and vocabulary decisions,

• model/optimizer scheduling,

• checkpoint/restart reliability,

• metrics and failure analysis.

360M was intentional: large enough to show a bit of meaningful behavior, small enough to iterate on with single-GPU constraints.

2) Setup and Constraints

This was a single-GPU training workflow designed for reliability over complexity.

• Hardware: 1x NVIDIA H100 (80GB)

• Precision: bf16 model training path

• Checkpointing: periodic saves + resume logic

• Constraint: FineWeb full corpus (~108TB scale) is not locally practical

• Constraint: early telemetry gaps (I didn’t log everything I should have)

Success criteria were simple:

1. finish target token budget without divergence,

2. maintain stable throughput,

3. produce progressively better qualitative outputs across checkpoints.

3) Architecture (Inspired by LLaMA + SmolLM-style Design)

This model follows a compact LLaMA-style decoder stack:

• RoPE positional encoding

• RMSNorm

• SwiGLU MLP

• Grouped Query Attention (GQA)

• FlashAttention2 path (with fallback)

Run-time configuration snapshot

4) Math Notes (LLaMA/SmolLM-style Core Equations)

Below are the core mechanics that guided this implementation style.

4.1 Causal language modeling objective

For token sequence \(x\_{1:T}\), training minimizes next-token negative log-likelihood:

\(\mathcal{L}*{\text{NLL}} = - \sum*{t=1}^{T} \log p\_\theta(x\_t \mid x\_{\lt t})\)

In code, this is standard cross-entropy over shifted logits/targets.

4.2 RMSNorm (used instead of LayerNorm)

For hidden vector \((h \in \mathbb{R}^d)\):

\(\text{RMSNorm}(h) = \gamma \odot \frac{h}{\sqrt{\frac{1}{d}\sum_{i=1}^d h_i^2 + \epsilon}}\)

No mean-centering term; scale-only normalization.

4.3 RoPE (rotary position embedding)

RoPE applies a position-dependent rotation to Q/K in each 2D subspace:

\(\begin{aligned} q'_{2i} &= q_{2i}\cos \theta_{p,i} - q_{2i+1}\sin \theta_{p,i} \\[0.5em] q'_{2i+1} &= q_{2i}\sin \theta_{p,i} + q_{2i+1}\cos \theta_{p,i} \end{aligned}\)

(and similarly for $k$), where \(\theta_{p,i}\) depends on position $p$ and frequency index $i$.

4.4 GQA attention shape idea

With \(n_q\) query heads and \(n_{kv}\) KV heads, each KV head is shared across groups of query heads:

\(\text{group size} = \frac{n_q}{n_{kv}}\)

This reduces KV memory/computation versus full MHA while keeping many query heads.

4.5 SwiGLU feed-forward block

A common LLaMA-family form:

\(\text{FFN}(x) = W_2\left(\text{SiLU}(W_1x) \odot (W_3x)\right)\)

4.6 Learning-rate schedule (warmup + cosine decay)

For step $s$, warmup \(s_w\), max steps \(s_{\max}\), peak LR \(\eta\), min ratio $r$:

\(\eta_s = \begin{cases} \eta \cdot \frac{s}{s_w}, & s < s_w \ \eta \left[r + (1-r)\cdot \frac{1+\cos\left(\pi\frac{s-s_w}{s_{\max}-s_w}\right)}{2}\right], & s \ge s_w \end{cases}\)

This is the schedule used in the run.

5) Data Methodology: FineWeb-6B Construction

I started from HuggingFaceFW/fineweb and used a streaming-first approach.

• Streamed row-by-row (no full dataset download).

• Preserved order as streamed (no shuffling).

• Tokenized on the fly and maintained a running token counter.

• Stopped ingestion at ~6B tokens.

• Wrote data frequently to disk instead of large in-memory buffering.

• Applied no extra filters beyond the 6B truncation boundary.

Tokenizer workflow was done via:

• 1-train-tokenizer.ipynb

6) Training Configuration

7) Engineering Workflow Choices

What kept this practical:

• Resume-first training loop with checkpoint loading.

• Regular checkpoint saves and artifact upload flow.

• Core metrics to monitor convergence + throughput.

• Simpler baseline architecture over novelty-heavy experimentation.

Known telemetry gaps in this run:

• grad norm not logged,

• NaN/inf counters not logged,

• step-time distribution not logged.

8) Training Dynamics and Stability

From logged metrics:

• train/loss decreased smoothly with expected slowing later.

• train/lr followed warmup then cosine decay as configured.

• train/tps stabilized after early/warmup region.

• validation checkpoints showed overall improvement with mild noise.

Interpretation: no divergence observed in tracked window, and training stayed operationally stable.
Caveat: stronger stability claims should include grad norm/NaN/step-time telemetry next run.

9) Next Steps / Future Work

• Add shuffle experiments (global/shard-level).

• Extend beyond 6B tokens.

• Add benchmark-style eval suite.

• Log grad norms, NaN counts, and step-time histograms.

• Sweep schedule/optimizer details more systematically.

• Try checkpoint averaging / EMA.

• Add instruction tuning (SFT).

• Explore distillation from stronger teachers.

10) Repro Appendix

Key files:

• Training entry: pretrain.py

• Model: model.py

• LR schedule: lr_schedule.py

Closing

This run shows that from-scratch pretraining at 360M scale is feasible for a small setup if you prioritize robust engineering, clear data methodology, and reproducibility.
The baseline is now in place; next iterations will focus on better data ordering, richer evals, and post-training quality improvements.

GitHub repo (pretraining code): https://github.com/weights-and-wires/smol-llama

Model weights (with inference code): https://huggingface.co/weights-and-wires/smol-llama

W&B report: https://api.wandb.ai/links/dotslasha/lydqsle8