# stanford-cs336/assignment1-basics architecture

> Architecture diagram of the GitHub repository stanford-cs336/assignment1-basics, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-13). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/stanford-cs336/assignment1-basics
- Repository: https://github.com/stanford-cs336/assignment1-basics

## Overview

#### 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: [`tests/test_model.py`](https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/test_model.py) (Component contract tests), [`cs336_basics/__init__.py`](https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/__init__.py) (Student implementation package), [`tests/fixtures/ts_tests/model.pt`](https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/fixtures/ts_tests/model.pt) (PyTorch reference state), [`tests/adapters.py`](https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/adapters.py) (Test adapter boundary)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_runtime["Runtime & package"]
  node_assignment["Assignment specification"]
  node_uv["uv Python environment<br/>[pyproject.toml]"]
  node_package["Student implementation package<br/>Python package<br/>[__init__.py]"]
end

subgraph group_tokenization["Tokenization"]
  node_pretoken_example["Pretokenization example<br/>runnable example"]
  node_bpe["BPE tokenizer<br/>tokenization API"]
  node_corpora[("Text corpora<br/>external data")]
  node_bpe_artifacts["Vocabulary &amp; merge rules<br/>BPE artifacts"]
end

subgraph group_training["Model &amp; training"]
  node_data_loader["Sequence batching<br/>data API"]
  node_transformer_lm["Transformer language model<br/>model API"]
  node_transformer_blocks["Transformer blocks<br/>model layer"]
  node_embeddings_rope["Embeddings &amp; RoPE<br/>model layer"]
  node_attention["Causal self-attention<br/>model layer"]
  node_loss_utils["Loss tensor utilities<br/>NN utility API"]
  node_adamw["AdamW optimizer<br/>optimizer API"]
  node_serialization["Model serialization<br/>persistence API"]
end

subgraph group_validation["Validation &amp; references"]
  node_test_adapter{{"Test adapter boundary<br/>integration adapter<br/>[adapters.py]"}}
  node_contract_tests["Component contract tests<br/>pytest suite<br/>[test_model.py]"]
  node_snapshots["Numerical snapshots<br/>reference artifacts"]
  node_pytorch_reference["PyTorch reference state<br/>external compatibility fixture<br/>[model.pt]"]
end

node_assignment -.->|"specifies"| node_package
node_uv -->|"runs"| node_package
node_pretoken_example -->|"demonstrates"| node_bpe
node_corpora -->|"training text"| node_bpe
node_bpe -->|"trains"| node_bpe_artifacts
node_bpe -->|"token IDs"| node_data_loader
node_data_loader -->|"batches"| node_transformer_lm
node_embeddings_rope -->|"input representations"| node_transformer_blocks
node_attention -->|"causal mixing"| node_transformer_blocks
node_transformer_blocks -->|"stacked decoder"| node_transformer_lm
node_transformer_lm -->|"logits"| node_loss_utils
node_loss_utils -->|"gradients"| node_adamw
node_adamw -->|"updates parameters"| node_transformer_lm
node_transformer_lm -->|"state"| node_serialization
node_pytorch_reference -.->|"compatibility target"| node_serialization
node_test_adapter -->|"binds APIs"| node_package
node_contract_tests -->|"calls through"| node_test_adapter
node_snapshots -.->|"numerical references"| node_contract_tests
node_pytorch_reference -.->|"serialization fixture"| node_contract_tests

click node_assignment "https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_assignment1_basics.pdf"
click node_uv "https://github.com/stanford-cs336/assignment1-basics/blob/main/pyproject.toml"
click node_package "https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/__init__.py"
click node_pretoken_example "https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/pretokenization_example.py"
click node_test_adapter "https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/adapters.py"
click node_contract_tests "https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/test_model.py"
click node_snapshots "https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/_snapshots/test_transformer_lm.npz"
click node_pytorch_reference "https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/fixtures/ts_tests/model.pt"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_assignment,node_uv,node_package toneBlue
class node_pretoken_example,node_bpe,node_corpora,node_bpe_artifacts toneAmber
class node_data_loader,node_transformer_lm,node_transformer_blocks,node_embeddings_rope,node_attention,node_loss_utils,node_adamw,node_serialization toneMint
class node_test_adapter,node_contract_tests,node_snapshots,node_pytorch_reference toneRose
```

## Components

### Runtime & package

- [Assignment specification](https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_assignment1_basics.pdf): specification · Defines the educational implementation requirements. · `cs336_assignment1_basics.pdf`
- [uv Python environment](https://github.com/stanford-cs336/assignment1-basics/blob/main/pyproject.toml): environment · `pyproject.toml`
- [Student implementation package](https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/__init__.py): Python package · Implementations are added beneath this namespace. · `cs336_basics/__init__.py`

### Tokenization

- [Pretokenization example](https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/pretokenization_example.py): runnable example · `cs336_basics/pretokenization_example.py`
- **BPE tokenizer**: tokenization API · Pre-tokenizes text, trains BPE, and encodes or decodes special tokens.
- **Text corpora**: external data · Optional TinyStories and OpenWebText sample downloads.
- **Vocabulary & merge rules**: BPE artifacts · Ordered merges and vocabulary produced by BPE training.

### Model & training

- **Sequence batching**: data API · Converts token IDs into training sequences and batches.
- **Transformer language model**: model API · Decoder-only language model with truncated-input support.
- **Transformer blocks**: model layer · Repeated causal attention, SwiGLU, normalization, and residual layers.
- **Embeddings & RoPE**: model layer
- **Causal self-attention**: model layer
- **Loss tensor utilities**: NN utility API · Cross-entropy-related tensor operations.
- **AdamW optimizer**: optimizer API · Updates model parameters from gradients.
- **Model serialization**: persistence API · Reads or writes compatible model and parameter state.

### Validation & references

- [Test adapter boundary](https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/adapters.py): integration adapter · Wires student module layout to expected public APIs. · `tests/adapters.py`
- [Component contract tests](https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/test_model.py): pytest suite · Covers tokenizer, BPE training, data, NN utilities, optimizer, model, and serialization. · `tests/test_model.py`
- [Numerical snapshots](https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/_snapshots/test_transformer_lm.npz): reference artifacts · Fixed numerical parity references for model and optimizer behavior. · `tests/_snapshots/test_transformer_lm.npz`
- [PyTorch reference state](https://github.com/stanford-cs336/assignment1-basics/blob/main/tests/fixtures/ts_tests/model.pt): external compatibility fixture · `tests/fixtures/ts_tests/model.pt`

## Connections

- 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