nanochat

The nanochat repository is organized into several functional directories that separate core logic, training scripts, evaluation tasks, and execution scripts:

• nanochat/: The core library containing the model architecture (gpt.py), inference engine (engine.py), optimizer (optim.py), tokenizer (tokenizer.py), and data utilities (dataloader.py, dataset.py).
• scripts/: High-level entry points for users, including chat_cli.py for interactive chat, base_train.py for pretraining, chat_sft.py for supervised fine-tuning, and infer_bench.py for performance benchmarking.
• tasks/: Definitions for evaluation datasets and task mixtures (e.g., arc.py, gsm8k.py, mmlu.py).
• runs/: Shell scripts for reproducible experiments, such as speedrun.sh for fast training or runcpu.sh for CPU/MPS execution.
• tests/: Unit tests for core components like the inference engine, optimizer, and tokenizer.

The project uses a BOS-aligned dataloader to ensure every sequence in a batch starts with a Beginning-of-Sequence (BOS) token. This improves training quality by providing proper context for every document.

Two implementation strategies exist:

1. Original (Simple): High efficiency and 100% token utilization, but some documents may start mid-stream without context.
2. _bos_bestfit (BestFit-Crop): The new default. It uses a bin-packing algorithm to pick the largest documents that fit into a row, reducing the number of 'unlucky' crops. While it achieves 100% utilization (no padding), it discards approximately 34% of tokens when documents do not fit perfectly into the sequence length.

Note on Loss Scaling: Switching to a BOS-aligned dataloader will result in lower validation loss compared to non-BOS loaders. This is because the model is no longer being trained on 'confusing' tokens that lack proper context. Absolute loss values are not directly comparable across different dataloader implementations.

The chat_sft.py script is used for Supervised Fine-Tuning (SFT). It has been updated to match the capabilities of base_train.py and includes several specific tuning options:

• Optimizer Warm-start: Use --load-optimizer=1 (default) to load pretrained momentum buffers via load_optimizer_state() in checkpoint_manager.py. Note that Learning Rates (LRs) are reset to fresh SFT values after loading.
• LR Schedule: Uses a warmup/constant/warmdown schedule. Key flags include --warmup-ratio, --warmdown-ratio, --init-lr-frac, and --final-lr-frac. A warmdown_ratio of 0.5 is recommended.
• Data Mixture Tuning: You can configure the number of epochs for specific datasets using --mmlu-epochs and --gsm8k-epochs.
• Evaluation: Periodic evaluation across 6 tasks is performed using --chatcore-every=N (e.g., --chatcore-every=200), with results logged to wandb.
• Hyperparameter Inheritance: By default, SFT inherits batch sizes and LRs from the pretrained checkpoint metadata, though CLI overrides remain functional.

# Example SFT command with custom epoch counts and evaluation frequency
python chat_sft.py --mmlu-epochs 3 --gsm8k-epochs 4 --chatcore-every 200 --load-optimizer=1

To improve model performance by offloading static N-gram patterns to a lookup table, implement an 'engram-lite' module. This approach uses a simple hash of bigrams and adds the resulting embeddings directly to the residual stream at every layer, gated by a per-layer learnable lambda ($\lambda$).

Implementation Details:

• Hash Function: (36313 * curr) XOR (27191 * prev) mod table_size
• Embedding Table: Use a zero-initialized embedding table to ensure it starts as an identity mapping.
• Injection: Add to the residual stream at every layer using a learned lambda (initial value of 0.1 is recommended over 0.0).
• Optimal Hyperparameters:
  • Table Size: vocab_size * 5 (e.g., ~164K entries for a 32K vocab).
  • Optimizer: Use AdamW with the same learning rate as the token embeddings.

Note: Avoid complex context-aware gating (like sigmoid dot-product gates) or restricting injection to early layers only, as simple direct addition to the residual stream has been shown to perform better.

class BigramEmbed(nn.Module):
    def __init__(self, vocab_size, embed_dim, table_multiplier=5):
        self.embed = nn.Embedding(vocab_size * table_multiplier, embed_dim)

    def forward(self, idx):
        # idx is the token indices
        # h is the hashed bigram index
        h = (36313 * idx[:, 1:]) ^ (27191 * idx[:, :-1]) % (table_size - 1)
        return self.embed(h)

Nanochat includes learnable per-layer residual connections to improve model performance with minimal compute overhead. This is controlled via the --scalar_lr CLI argument (default 0.5).

Mechanisms

1. x0_lambdas (x0 residual connections): Blends the initial normalized embedding (x0) back into each layer: x = resid_lambdas[i] * x + x0_lambdas[i] * x0. This provides a direct path from the embedding to deep layers.
2. resid_lambdas (residual stream scaling): A multiplicative scaling of the residual stream at each layer, initialized to 1.0.

Learning Rate Sensitivity

These two scalar types require different learning rates to train effectively:

• x0_lambdas: Use the standard --scalar_lr (e.g., 0.5).
• resid_lambdas: Require a much smaller learning rate, approximately 100x smaller than the x0_lambdas (e.g., scalar_lr * 0.01).

Implementation Warning

When implementing custom layers, remember that __init__ runs in a meta device context. Any tensor values set during __init__ are fake; actual values must be initialized in the init_weights() method.

# Example CLI usage
python train.py --scalar_lr 0.6