How stanford-cs336/assignment1-basics works
Purpose and execution
This is a Python educational repository for implementing foundational language-model infrastructure: BPE tokenization, tensor/Transformer components, data loading, optimization, and checkpoint serialization. The assignment specification is cs336_assignment1_basics.pdf. Code is run through the uv-managed Python environment (pyproject.toml, uv.lock); the included runnable example is cs336_basics/pretokenization_example.py.
Implementation boundary
The package namespace is cs336_basics/__init__.py; student implementations are expected to be added beneath this package. The test suite is intentionally decoupled from student module layout through tests/adapters.py: implementations must be wired there so tests can call the expected tokenizer, model, optimizer, data, and serialization APIs. This adapter is the principal integration boundary.
Tokenization pipeline
The tokenizer stage covers corpus pre-tokenization, BPE vocabulary/merge training, special-token behavior, and encode/decode compatibility. The only shipped pipeline example is cs336_basics/pretokenization_example.py; expected behavior and reference artifacts are exercised through tests/test_tokenizer.py, tests/test_train_bpe.py, and fixtures such as tests/fixtures/gpt2_vocab.json and tests/fixtures/gpt2_merges.txt. BPE training consumes text corpora and produces vocabulary plus ordered merge rules.
Neural-model stack
The central model path is a decoder-style Transformer language model: token embeddings feed positional encoding/RoPE, repeated Transformer blocks, multi-head causal self-attention, feed-forward/SwiGLU layers, normalization, and output projection. Architectural expectations are defined by tests/test_model.py and numerical reference snapshots including tests/_snapshots/test_transformer_lm.npz, tests/_snapshots/test_multihead_self_attention_with_rope.npz, and tests/_snapshots/test_transformer_block.npz. The model must support truncated inputs as indicated by tests/_snapshots/test_transformer_lm_truncated_input.npz.
Training primitives and data flow
Training-oriented components include batched/tokenized data access, cross-entropy-related tensor utilities, and AdamW parameter updates. Their contracts are separated into tests/test_data.py, tests/test_nn_utils.py, and tests/test_optimizer.py; optimizer numerical parity is checked against tests/_snapshots/test_adamw.npz. The intended flow is raw corpus text → tokenizer IDs → sequence/batch construction → Transformer forward pass → loss/gradients → AdamW update.
Persistence, validation, and external inputs
Serialization must preserve model/parameter state compatibly with external PyTorch reference material in tests/fixtures/ts_tests/model.pt and tests/fixtures/ts_tests/model_config.json; its boundary is validated by tests/test_serialization.py. Tests are run with uv run pytest and use fixed snapshots under tests/_snapshots. Optional training corpora are downloaded outside the package into a local data/ directory from Hugging Face: TinyStories and Stanford CS336’s OpenWebText sample.
Main entry points: test_model.py Component contract tests · __init__.py Student implementation package · model.pt PyTorch reference state · adapters.py Test adapter boundary
Components 19
Runtime & package
- Assignment specification · Defines the educational implementation requirements.
cs336_assignment1_basics.pdf - uv Python environment
pyproject.toml - Student implementation package · Implementations are added beneath this namespace.
cs336_basics/__init__.py
Tokenization
- Pretokenization example
cs336_basics/pretokenization_example.py - BPE tokenizer · Pre-tokenizes text, trains BPE, and encodes or decodes special tokens.
- Text corpora · Optional TinyStories and OpenWebText sample downloads.
- Vocabulary & merge rules · Ordered merges and vocabulary produced by BPE training.
Model & training
- Sequence batching · Converts token IDs into training sequences and batches.
- Transformer language model · Decoder-only language model with truncated-input support.
- Transformer blocks · Repeated causal attention, SwiGLU, normalization, and residual layers.
- Embeddings & RoPE
- Causal self-attention
- Loss tensor utilities · Cross-entropy-related tensor operations.
- AdamW optimizer · Updates model parameters from gradients.
- Model serialization · Reads or writes compatible model and parameter state.
Validation & references
- Test adapter boundary · Wires student module layout to expected public APIs.
tests/adapters.py - Component contract tests · Covers tokenizer, BPE training, data, NN utilities, optimizer, model, and serialization.
tests/test_model.py - Numerical snapshots · Fixed numerical parity references for model and optimizer behavior.
tests/_snapshots/test_transformer_lm.npz - PyTorch reference state
tests/fixtures/ts_tests/model.pt
Connections 19
- Assignment specification → Student implementation package: specifies
- uv Python environment → Student implementation package: runs
- Pretokenization example → BPE tokenizer: demonstrates
- Text corpora → BPE tokenizer: training text
- BPE tokenizer → Vocabulary & merge rules: trains
- BPE tokenizer → Sequence batching: token IDs
- Sequence batching → Transformer language model: batches
- Embeddings & RoPE → Transformer blocks: input representations
- Causal self-attention → Transformer blocks: causal mixing
- Transformer blocks → Transformer language model: stacked decoder
- Transformer language model → Loss tensor utilities: logits
- Loss tensor utilities → AdamW optimizer: gradients
- AdamW optimizer → Transformer language model: updates parameters
- Transformer language model → Model serialization: state
- PyTorch reference state → Model serialization: compatibility target
- Test adapter boundary → Student implementation package: binds APIs
- Component contract tests → Test adapter boundary: calls through
- Numerical snapshots → Component contract tests: numerical references
- PyTorch reference state → Component contract tests: serialization fixture