Goldfish

Goldfish is a small PyTorch framework for sequential modelling. Its first vertical slice trains character-level causal language models, with reusable training, dataset-locking, registry, and experiment-management infrastructure for future numeric and multimodal sequence tasks.

For the architecture and data contracts, see:

• docs/ARCHITECTURE.md
• docs/DATASETS.md
• docs/TASKS.md
• docs/MODELS.md
• docs/EXPERIMENTS.md
• docs/OBSERVABILITY.md

Quick start

Install the locked project environment:

uv sync

Prepare the included text dataset. This validates manifest.yaml, fits the tokenizer on the train split only, and writes the dataset and tokenizer locks:

uv run goldfish prepare data/alphabet
# Equivalent: uv run python prepare.py data/alphabet

Then train a new GRU experiment:

uv run goldfish train data/alphabet \
  --name alphabet-gru \
  --sequence-length 13 \
  --batch-size 32 \
  --epochs 40 \
  --model-profile model-profiles/language/gru-small.yaml \
  --device cuda \
  --optimizer adamw \
  --weight-decay 0.0001 \
  --gradient-clip-norm 1.0 \
  --prompt "cdefg" \
  --max-new-tokens 100

The dispatcher forwards all arguments after train directly to train.py, so this is equivalent:

uv run python train.py data/alphabet --name alphabet-gru --epochs 40 \
  --model-profile model-profiles/language/gru-small.yaml

The run is created under runs/, for example:

runs/exp1-alphabet-gru/
├── config.yaml
├── data.json
├── environment.json
├── metrics.jsonl
├── summary.json
├── run.log
├── checkpoints/
│   ├── latest.pt
│   ├── best.pt
│   └── final.pt
└── artifacts/
    ├── samples/
    └── probes/          # Present when --observability is enabled
        ├── manifest.json
        ├── mixer-state.jsonl
        └── activation-stats.jsonl

Dataset workflow

Goldfish trains from a self-contained, manifest-driven dataset root:

data/<dataset-name>/
├── manifest.yaml
├── dataset-lock.json              # Generated by `prepare`; do not edit manually
├── train/
├── val/
├── test/
└── tokenizer/
    ├── tokenizer.json             # Generated from train only
    └── tokenizer-lock.json        # Generated by `prepare`

manifest.yaml lists the exact split files and their order. The current text builder treats each listed file as one document and appends an EOS token between documents.

data/alphabet/ is the checked-in text example dataset bundle, including its manifest, source splits, generated tokenizer, and lock artifacts. data/fourier/ is the checked-in numeric forecasting example: 28,000 deterministic observations from a 12-component, multi-frequency Fourier series, split across two ordered training shards, validation, and test, with frozen train-only normalization and lock artifacts. data/fourier-lb256/ has identical raw rows and horizons but a lookback of 256, making it the paired long-history A/B dataset. Other datasets under data/ remain ignored by default unless explicitly unignored.

Prepare before training

# Text language-model example
uv run goldfish prepare data/alphabet

# Numeric Fourier-series forecasting example
uv run goldfish prepare data/fourier
uv run goldfish train data/fourier --name fourier --epochs 10 --batch-size 32 \
  --model-profile model-profiles/forecast/gru-small.yaml
uv run goldfish forecast runs/exp1-fourier --checkpoint best --split test \
  --plot runs/exp1-fourier/artifacts/forecasts/test-best.png

Preparation performs:

validate manifest
→ fit tokenizer from train files only
→ write tokenizer/tokenizer.json
→ write dataset-lock.json
→ write tokenizer/tokenizer-lock.json

Training validates both locks before constructing a DataLoader or optimizing. If a manifest-listed raw file, file order, tokenizer artifact, or tokenizer configuration changes, run prepare again explicitly.

Training

The command form is:

uv run goldfish train <dataset-root> [options]

Models

Model architectures are selected with a required YAML profile:

--model-profile model-profiles/language/gru-small.yaml
--model-profile model-profiles/forecast/lstm-small.yaml
--model-profile model-profiles/forecast/multihead-lstm-small.yaml
--model-profile model-profiles/forecast/deltanet-small.yaml  # DeltaNet fast-weight memory encoder

Profiles supply architecture-owned parameters and model registry identity. Goldfish injects data-derived dimensions (vocabulary size or numeric feature/target/horizon counts), then saves the complete resolved model configuration in the run. This keeps new model-specific hyperparameters out of the training CLI.

Compilation

Use --compile to run model forward passes through torch.compile. Goldfish keeps checkpoint state dictionaries in the uncompiled model format, so checkpoints remain compatible with infer and forecast without compilation.

uv run goldfish train data/fourier-lb256 \
  --model-profile model-profiles/forecast/lstm-lb256-128x2.yaml \
  --compile

Compilation has a one-time startup cost; it is usually most useful for longer runs. The resolved value is saved in config.yaml and cannot be changed while resuming a run.

Device selection

Both training and inference accept:

--device cpu
--device cuda
--device mps

When --device is omitted, Goldfish chooses the best available platform device in this order:

CUDA → MPS → CPU

An explicit unavailable accelerator is an error: --device cuda does not silently fall back to CPU. The resolved device is recorded in the run configuration and environment metadata.

Startup brief and DataLoader resources

Before the first training batch, Goldfish prints a run brief with the resolved device, dataset/window dimensions, model shape, DataLoader settings, optimizer/scheduler, reproducibility mode, and lock fingerprint. For example, numeric runs include:

Device:     cuda
Loader:     train_workers=9, val_workers=7, pin_memory=True, prefetch=2, persistent=True
Reproduce:  deterministic=False, seed=None

--num-workers auto is the default. It reserves roughly 20% of logical CPUs for the operating system, then allocates the remaining worker budget 60:40 between training and validation/test. Training and validation execute sequentially, so this is a phase-specific limit rather than simultaneous CPU consumption.

# Default automatic allocation
uv run goldfish train data/fourier-lb256 --num-workers auto

# Explicit total budget or phase-specific overrides
uv run goldfish train data/fourier-lb256 \
  --num-workers 16 --train-workers 9 --val-workers 7 --prefetch-factor 4

# Disable worker processes
uv run goldfish train data/fourier-lb256 --num-workers 0

On CUDA, Goldfish enables pinned-memory DataLoaders and non-blocking batch transfer. For worker counts above zero, it also enables persistent workers and uses --prefetch-factor (default 2). The resolved settings are saved in config.yaml and reused by strict resume. On successful completion, Goldfish also writes artifacts/plots/training-curves.png, with separate train and validation trajectories for every scalar metric in metrics.jsonl.

Reproducibility

Deterministic execution is opt-in; it is not enabled by default because deterministic algorithms can reduce performance or reject unavailable kernels. Enable it only with an explicit seed:

uv run goldfish train data/fourier-lb256 \
  --deterministic --seed 7

--deterministic without --seed is rejected. In deterministic mode, Goldfish seeds Python, NumPy, and PyTorch; enables deterministic PyTorch/cuDNN algorithms; and configures deterministic cuBLAS behavior before CUDA initialization. A seed may also be supplied without --deterministic when only repeatable random initialization is desired. The resolved seed and deterministic values are recorded in the run config and restored on resume.

Optimizers

--optimizer adamw  # default
--optimizer adam
--optimizer sgd

Common optimizer options:

Option	Default	Meaning
--learning-rate, --lr	0.001	Base learning rate.
--weight-decay	Optimizer-specific	L2/decoupled weight-decay setting.
--momentum	Optimizer-specific	SGD momentum only.

The resolved run config records all effective optimizer defaults, including Adam/AdamW betas and epsilon, even when they were not supplied on the CLI.

Learning-rate schedulers

--scheduler none         # default
--scheduler cosine
--scheduler step
--scheduler exponential
--scheduler plateau

Scheduler	Required/typical options	Step timing
none	None	None
cosine	--scheduler-t-max, optional --scheduler-eta-min	Epoch by default; --scheduler-step-timing batch is also valid.
step	--scheduler-step-size, --scheduler-gamma	Epoch.
exponential	--scheduler-gamma	Epoch.
plateau	--scheduler-factor, --scheduler-patience	Validation.

Example cosine run:

uv run goldfish train data/alphabet \
  --name alphabet-cosine \
  --epochs 40 \
  --scheduler cosine \
  --scheduler-t-max 40 \
  --scheduler-eta-min 0.000001

Invalid scheduler options for the selected scheduler are rejected rather than silently ignored.

Checkpoints, metrics, and samples

For each completed epoch, Goldfish:

• appends train/validation metrics and effective learning rate to metrics.jsonl;
• writes checkpoints/latest.pt;
• updates checkpoints/best.pt when validation/loss improves;
• writes a generated text sample at --sample-frequency;
• optionally writes checkpoints/epoch-NNNN.pt using --checkpoint-frequency.

At the end of a successful run, Goldfish writes:

checkpoints/final.pt
artifacts/samples/final.txt
summary.json

Useful options:

Option	Default	Meaning
--runs-dir	runs	Base directory for new experiment runs.
--name	generated	Human-readable suffix in exp<N>-<name>.
--sample-frequency	1	Write a generated sample every N epochs.
--checkpoint-frequency	unset	Write periodic epoch-NNNN.pt checkpoints.
--gradient-clip-norm	unset	Clip gradient norm before each optimizer update.
--prompt	first train document	Prompt used for generated samples.
--max-new-tokens	100	Number of tokens generated after the prompt.

best.pt currently monitors validation/loss with min mode. The full monitor policy is saved in each run's resolved config.yaml.

Observability

--observability enables probe-based state trajectories for a run. Probes are declared by the model profile (its observability.probes block), so a run only enables the system and configures the reference input set:

uv run goldfish train data/fourier-lb256 \
  --model-profile model-profiles/forecast/multihead-lstm-small.yaml \
  --observability --observability-batches 8

Option	Default	Meaning
--observability	off	Enable the probes declared by the model profile.
--observability-batches	8	Reference batches captured from the validation split for activation probes.

If the validation split yields fewer batches than --observability-batches, training fails at start with an error naming the available count; reduce the value accordingly (the reference set must be identical across runs, so Goldfish never silently captures fewer batches than configured).

Three probe kinds exist (see docs/OBSERVABILITY.md for the full specification):

• mixer-state — doubly stochastic / unconstrained mixing matrices, logits, identity distances, and gradient-norm snapshots at every sampled epoch;
• communication-state — dense inter-layer blocks and gated latent communication (routing, gates, block norms);
• activation-stats — statistics of intermediate tensors (norms, mean-abs, std, max, p95) on a fixed reference forward pass, including per-head reductions and composite quantities such as the injected message magnitude ||gate ⊙ D(m)|| / ||h||.

Records are written as append-only JSONL under artifacts/probes/, one file per probe, with a manifest.json that records the resolved configuration, matched modules, and the reference split fingerprint for reproducibility. Schedules are configurable per probe (every-N-epochs or explicit epoch lists, plus initial/final records); dense sampling of the early phase is recommended for trajectory questions. The observability block is persisted in config.yaml and restored on resume.

For example, the multihead-lstm-small*.yaml profiles declare the probes appropriate to their model family, so the same command works for the mixer, dense-communication, and latent-communication variants.

Strict resume

Resume from the run's checkpoints/latest.pt:

uv run goldfish train data/alphabet \
  --resume runs/exp1-alphabet-gru \
  --epochs 10

On resume, --epochs means additional epochs. Goldfish appends to metrics.jsonl and retains the same run directory.

Resume is deliberately strict. It verifies:

• dataset and tokenizer fingerprints;
• run ID and checkpoint format/provenance;
• model family, model name, and architecture configuration;
• optimizer and scheduler configuration.

Model/data/optimizer/scheduler changes require a new run. Forking experiment runs is planned but not yet implemented.

Inference

Generate from a managed experiment without creating a new run:

uv run goldfish infer runs/exp1-alphabet-gru \
  --checkpoint best \
  --prompt "cdefg" \
  --max-new-tokens 100

uv run python infer.py runs/exp1-alphabet-gru --checkpoint best --prompt "cdefg" --device cpu

infer.py loads the run's resolved model configuration and the selected managed checkpoint (best, latest, or final). It validates the current dataset and tokenizer locks before generation, so inference does not silently run with a changed vocabulary or dataset bundle.

Forecast visualization

goldfish forecast can save a raw-unit matplotlib PNG for any exported forecast window:

uv run goldfish forecast <run-dir> --checkpoint best --split test \
  --plot forecast.png --plot-window 0

Each target receives its own subplot. It shows the target’s model input history in blue, actual future values in green, and forecast values in dashed orange. The cutoff is row offset 0; horizon positions follow the manifest-declared row offsets. Use --plot-window to choose a different exported window.

Lookback A/B datasets

data/fourier/ uses lookback: 32; data/fourier-lb256/ uses lookback: 256. They contain byte-identical raw CSV rows, features, targets, and horizons, so they can isolate the effect of historical context. Both bundles have independently generated locks and normalizer artifacts. Recreate the long-history bundle with uv run goldfish prepare data/fourier-lb256 after modifying its manifest or raw files.

Commands, help, and tests

Dispatcher command	Direct entry point	Purpose
goldfish prepare ...	prepare.py ...	Build text tokenizer or numeric normalizer and lock artifacts.
goldfish train ...	train.py ...	Create or strictly resume a managed training run.
goldfish infer ...	infer.py ...	Generate text from a managed checkpoint.
goldfish forecast ...	forecast.py ...	Export and optionally plot numeric forecasts from a managed run.

uv run goldfish --help
uv run python train.py --help
uv run python infer.py --help
uv run python forecast.py --help

uv run pytest -q