NemoAI

NemoAI is a minimal deep learning framework that trains small transformer models (GPT) from scratch on an OpenCL device (or plain CPU). It is built for running on laptops / integrated GPUs such as AMD Vega (via Mesa rusticl).

The focus is on correctness (validated gradients, CPU/GPU symmetry, exact resume-from-checkpoint), reproducible training, and a small, auditable code base. The reference config is a ~5.32M parameter GPT (d_model 128, 2 heads, 6 blocks, sequence 128, vocab 16000 es) trained at ~1.1 s/step — 467 tok/s — on a Vega 8 iGPU. See docs/technical.md for the full benchmark and the config sweep.

Features

• Decoder-only GPT: token + positional embeddings, stacked transformer blocks, final LayerNorm, LM head. Configurable vocab, width, heads, blocks, context.
• Backends: OpenCL (GPU, async command queue) and CPU (reference), with a CPU/OpenCL symmetry test ensuring the GPU path is numerically consistent.
• Optimizers: SGD (momentum) and AdamW, running fully on-device (no host round-trips) on the GPU backend.
• Losses: CrossEntropy, MSE. Autograd over an explicit operation graph.
• Checkpoints: portable float32 binary format (CPU <-> OpenCL), weights and optimizer state (moments + step counter). Training can be resumed bit-exact.
• Inference: GPT::generate with greedy or top-k / temperature sampling.
• Full test suite (8 binaries) covering gradient checking, backend symmetry, convergence, checkpoints/resume, and end-to-end GPT training.

Requirements

• Linux, a C++20 compiler, libOpenCL.
• An OpenCL 1.2+ implementation (e.g. Mesa rusticl for AMD/Intel iGPUs, or the CPU OpenCL reference stack). If OpenCL is unavailable, examples fall back to CPU via --device cpu.

Build

make -f Makefile.core      # core unit test (core_test)
make -f Makefile.tests     # full automated test suite (8 test binaries)
make -f Makefile.examples  # reproducible examples (train_gpt, generate_gpt)

Run the test suite

make -f Makefile.tests
./test_gradient_checking    # analytic vs numeric gradients, all ops incl. causal attention
./test_backend_symmetry     # CPU vs OpenCL results for every op
./test_convergence_linear   # y = 2x+1 with AdamW
./test_convergence_mlp      # XOR 4/4
./test_checkpoints          # checkpoint round-trip
./test_gpt                  # end-to-end GPT training (loss must decrease)
./test_resume               # exact resume from a checkpoint + generation
./test_benchmark            # per-step profiling at 0.1M/1M/5M/10M params

Quick start

Train a tiny character-level GPT on the bundled corpus:

./examples/train_gpt --data examples/data/small.txt --steps 300 \
    --d_model 64 --heads 4 --blocks 2 --seq 16 --save /tmp/gpt.bin

Generate text from the trained checkpoint (pass the matching geometry for legacy v1 checkpoints; --tokenizer is the primary way to load a BPE model):

./examples/generate_gpt --model /tmp/gpt.bin --data examples/data/small.txt \
    --prompt "the cat" --tokens 32 --topk 1 \
    --d_model 64 --heads 4 --blocks 2 --seq 16

Resume interrupted training from a checkpoint:

./examples/train_gpt --data examples/data/small.txt --steps 600 \
    --d_model 64 --heads 4 --blocks 2 --seq 16 --save /tmp/gpt.bin --resume /tmp/gpt.bin

train_gpt writes a checkpoint every --save_every steps that includes both weights and the AdamW state, so resuming continues training exactly as if the run had never stopped.

Inference with a BPE model (no vocabulary reconstruction)

nlp/train_driver checkpoints embed the BPE tokenizer, so a single .model file is fully self-contained for inference — the vocabulary is never re-derived from a corpus:

./examples/generate_gpt --model runs/gpt-best.model --prompt "Hola" --tokens 32 --topk 1

The tokenizer is resolved in order of preference:

1. embedded in the checkpoint (checkpoints written by nlp/train_driver);
2. --tokenizer DIR → loads DIR/vocab.bin + DIR/merges.bin;
3. --vocab X --merges Y → explicit BPE files.

For example, the Spanish model trained on data/spanish/:

./examples/generate_gpt --model runs/gpt-best.model \
    --tokenizer data/spanish --prompt "Hola"

The BPE tokenizer itself lives permanently in the project (data/spanish/ has vocab.bin, merges.bin and the corpus.txt it was trained from), so any trained model remains usable at any time.

See docs/technical.md for the architecture, the OpenCL backend design, the checkpoint format, and the benchmark results.

Training pipeline (nlp/)

The nlp/ subproject adds the full training experience on top of the core framework: JSON configs, a checkpoint registry, resume, monitoring, and clean SIGINT/SIGTERM shutdown.

make -f Makefile.nlp all            # builds nlp/train_driver + nlp/corpus_driver
./nlp/train_driver --config runs/gpt.json            # train (resume latest)
./nlp/train_driver --config runs/gpt.json --run-steps 200   # short session
./nlp/train_driver --config runs/gpt.json --list-checkpoints
./nlp/train_driver --config runs/gpt.json --write-config /tmp/eff.json

runs/gpt.json is the reference config (d128, 2 heads, 6 blocks, seq 128, vocab 16000 es, ~5.32M params). Checkpoints land in runs/ as gpt-<step>.model with a checkpoints.json registry; Ctrl-C saves a final checkpoint. Per-step metrics go to runs/train_log.csv (CSV) and runs/metrics.jsonl (JSONL).

Pass the BPE tokenizer files with --vocab/--merges (or set identity.vocab_path / identity.merges_path in the config) so the tokenizer is embedded in every checkpoint — the resulting .model is self-contained:

./nlp/train_driver --config runs/gpt.json \
    --vocab data/spanish/vocab.bin --merges data/spanish/merges.bin

Corpus pipeline + permanent tokenizer

nlp/corpus_driver prepares the training corpus and stores the BPE tokenizer permanently inside the project (never in /tmp):

./nlp/corpus_driver --input data/raw_spanish \
    --train-tokenizer data/raw_spanish --tokenizer-dir data/spanish \
    --vocab-size 16000 --outdir data/spanish

This trains the BPE tokenizer (writing vocab.bin, merges.bin and the corpus.txt it was trained from into data/spanish/), then runs the two-pass cleaning pipeline to produce train.bin, val.bin and metadata.json. --train-tokenizer accepts a single text file or a directory (concatenated in sorted order); training is deterministic, so re-running it yields identical files. See docs/pipeline-fase2.md for the full flag reference.

Project layout

core/              framework sources (headers + .cpp, no external deps)
core/tests/        test programs
core/ops/          differentiable operations (attention, layernorm, ...)
core/layers/       modules (GPT, TransformerBlock, Linear, MHA)
examples/          train_gpt / generate_gpt + a small corpus
nlp/               train_driver + GPTTrainer (JSON config, checkpoints, logs)
runs/              gpt.json reference config; run artifacts
data/spanish/      Spanish BPE-16000 corpus (train/val .bin, vocab/merges tokenizer)
Makefile.core      core unit test
Makefile.tests     automated test suite
Makefile.examples  reproducible examples
Makefile.nlp       nlp/ training pipeline
docs/technical.md  technical documentation and benchmark data