# zhasion/cs336 architecture

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

- Interactive diagram: https://gitdiagram.com/zhasion/cs336
- Repository: https://github.com/zhasion/cs336

## Overview

This repository is a set of personal implementations for five Stanford CS336 assignments, not one connected application. Its learning stages span tokenizer and transformer basics, systems optimization, scaling-law analysis, web-data quality filtering, and model alignment. The graph keeps these branches separate; only relationships supported by the sampled source or README are shown. Several capabilities are documented but their implementation was not sampled, so those nodes have no speculative call edges.

Main entry points: [`assignment1-basics/inference.py`](https://github.com/zhasion/cs336/blob/main/assignment1-basics/inference.py) (Text Generation), [`assignment1-basics/train_loop.py`](https://github.com/zhasion/cs336/blob/main/assignment1-basics/train_loop.py) (Training Loop), [`assignment3-scaling/chinchilla_isoflops.py`](https://github.com/zhasion/cs336/blob/main/assignment3-scaling/chinchilla_isoflops.py) (IsoFLOP Analysis), [`assignment4-data/cs336_data/quality_classifier.py`](https://github.com/zhasion/cs336/blob/main/assignment4-data/cs336_data/quality_classifier.py) (Quality Filtering)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_basics["Model Basics"]
  node_bpe["BPE Training<br/>[bpe.py]"]
  node_tokenizer["Tokenizer<br/>[tokenizer.py]"]
  node_transformer["Transformer LM<br/>[module.py]"]
  node_training["Training Loop<br/>[train_loop.py]"]
  node_inference["Text Generation<br/>[inference.py]"]
end

subgraph group_systems["Systems"]
  node_flash["Flash Attention<br/>[flash_attention.py]"]
  node_distributed["Distributed Training<br/>[naive_ddp.py]"]
  node_overlap["Gradient Overlap"]
  node_systems_bench["Systems Benchmarks"]
end

subgraph group_scaling["Scaling"]
  node_scaling_model["Scaling Model<br/>[model.py]"]
  node_scaling_fit["IsoFLOP Analysis"]
end

subgraph group_data["Data Quality"]
  node_quality["Quality Filtering"]
  node_data_utils["Data Utilities<br/>[utils.py]"]
  node_data_loader["Training Data Loader<br/>[data.py]"]
end

subgraph group_alignment["Alignment"]
  node_sft["Supervised Fine-Tuning<br/>[sft.py]"]
  node_dpo["Preference Optimization<br/>[dpo.py]"]
  node_grpo["Group Relative Policy<br/>[grpo.py]"]
  node_grader["Math Answer Grader<br/>[drgrpo_grader.py]"]
  node_evaluation["Alignment Evaluation<br/>[evaluate_safety.py]"]
end

node_learner(("Learner"))
node_checkpoint["Model Checkpoint"]
node_crawl["Common Crawl"]

node_learner -->|"starts training"| node_training
node_training -->|"trains"| node_transformer
node_learner -->|"requests generation"| node_inference
node_inference -->|"encodes prompt"| node_tokenizer
node_inference -->|"generates with"| node_transformer
node_inference -->|"loads weights"| node_checkpoint
node_inference -->|"returns text"| node_learner
node_learner -->|"runs tokenizer work"| node_bpe
node_learner -->|"runs benchmarks"| node_systems_bench
node_learner -->|"analyzes results"| node_scaling_fit
node_scaling_fit -.->|"fits scaling context"| node_scaling_model
node_crawl -->|"provides web records"| node_quality
node_quality -->|"uses filtering helpers"| node_data_utils
node_learner -->|"loads training data"| node_data_loader
node_learner -->|"runs fine-tuning"| node_sft
node_learner -->|"runs preference training"| node_dpo
node_learner -->|"runs policy training"| node_grpo
node_grpo -.->|"grades answers"| node_grader
node_learner -->|"evaluates safety"| node_evaluation

click node_bpe "https://github.com/zhasion/cs336/blob/main/assignment1-basics/cs336_basics/bpe.py"
click node_tokenizer "https://github.com/zhasion/cs336/blob/main/assignment1-basics/cs336_basics/tokenizer.py"
click node_transformer "https://github.com/zhasion/cs336/blob/main/assignment1-basics/cs336_basics/module.py"
click node_training "https://github.com/zhasion/cs336/blob/main/assignment1-basics/train_loop.py"
click node_inference "https://github.com/zhasion/cs336/blob/main/assignment1-basics/inference.py"
click node_flash "https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/flash_attention.py"
click node_distributed "https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/naive_ddp.py"
click node_overlap "https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/ddp_overlap_bucketed.py"
click node_systems_bench "https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/flash_benchmarking.py"
click node_scaling_model "https://github.com/zhasion/cs336/blob/main/assignment3-scaling/cs336_scaling/model.py"
click node_scaling_fit "https://github.com/zhasion/cs336/blob/main/assignment3-scaling/chinchilla_isoflops.py"
click node_quality "https://github.com/zhasion/cs336/blob/main/assignment4-data/cs336_data/quality_classifier.py"
click node_data_utils "https://github.com/zhasion/cs336/blob/main/assignment4-data/cs336_data/utils.py"
click node_data_loader "https://github.com/zhasion/cs336/blob/main/assignment4-data/cs336-basics/cs336_basics/data.py"
click node_sft "https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/sft.py"
click node_dpo "https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/dpo.py"
click node_grpo "https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/grpo.py"
click node_grader "https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/drgrpo_grader.py"
click node_evaluation "https://github.com/zhasion/cs336/blob/main/assignment5-alignment/scripts/evaluate_safety.py"

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_bpe,node_tokenizer,node_transformer,node_training,node_inference toneBlue
class node_flash,node_distributed,node_overlap,node_systems_bench toneAmber
class node_scaling_model,node_scaling_fit toneMint
class node_quality,node_data_utils,node_data_loader toneRose
class node_sft,node_dpo,node_grpo,node_grader,node_evaluation,node_learner,node_checkpoint,node_crawl toneIndigo
```

## Components

### Model Basics

- [BPE Training](https://github.com/zhasion/cs336/blob/main/assignment1-basics/cs336_basics/bpe.py): `assignment1-basics/cs336_basics/bpe.py`
- [Tokenizer](https://github.com/zhasion/cs336/blob/main/assignment1-basics/cs336_basics/tokenizer.py): `assignment1-basics/cs336_basics/tokenizer.py`
- [Transformer LM](https://github.com/zhasion/cs336/blob/main/assignment1-basics/cs336_basics/module.py): `assignment1-basics/cs336_basics/module.py`
- [Training Loop](https://github.com/zhasion/cs336/blob/main/assignment1-basics/train_loop.py): `assignment1-basics/train_loop.py`
- [Text Generation](https://github.com/zhasion/cs336/blob/main/assignment1-basics/inference.py): `assignment1-basics/inference.py`

### Systems

- [Flash Attention](https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/flash_attention.py): `assignment2-systems/cs336_systems/flash_attention.py`
- [Distributed Training](https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/naive_ddp.py): `assignment2-systems/cs336_systems/naive_ddp.py`
- [Gradient Overlap](https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/ddp_overlap_bucketed.py): `assignment2-systems/cs336_systems/ddp_overlap_bucketed.py`
- [Systems Benchmarks](https://github.com/zhasion/cs336/blob/main/assignment2-systems/cs336_systems/flash_benchmarking.py): `assignment2-systems/cs336_systems/flash_benchmarking.py`

### Scaling

- [Scaling Model](https://github.com/zhasion/cs336/blob/main/assignment3-scaling/cs336_scaling/model.py): `assignment3-scaling/cs336_scaling/model.py`
- [IsoFLOP Analysis](https://github.com/zhasion/cs336/blob/main/assignment3-scaling/chinchilla_isoflops.py): `assignment3-scaling/chinchilla_isoflops.py`

### Data Quality

- [Quality Filtering](https://github.com/zhasion/cs336/blob/main/assignment4-data/cs336_data/quality_classifier.py): `assignment4-data/cs336_data/quality_classifier.py`
- [Data Utilities](https://github.com/zhasion/cs336/blob/main/assignment4-data/cs336_data/utils.py): `assignment4-data/cs336_data/utils.py`
- [Training Data Loader](https://github.com/zhasion/cs336/blob/main/assignment4-data/cs336-basics/cs336_basics/data.py): `assignment4-data/cs336-basics/cs336_basics/data.py`

### Alignment

- [Supervised Fine-Tuning](https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/sft.py): `assignment5-alignment/cs336_alignment/sft.py`
- [Preference Optimization](https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/dpo.py): `assignment5-alignment/cs336_alignment/dpo.py`
- [Group Relative Policy](https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/grpo.py): `assignment5-alignment/cs336_alignment/grpo.py`
- [Math Answer Grader](https://github.com/zhasion/cs336/blob/main/assignment5-alignment/cs336_alignment/drgrpo_grader.py): `assignment5-alignment/cs336_alignment/drgrpo_grader.py`
- [Alignment Evaluation](https://github.com/zhasion/cs336/blob/main/assignment5-alignment/scripts/evaluate_safety.py): `assignment5-alignment/scripts/evaluate_safety.py`

### Other components

- **Learner**
- **Model Checkpoint**
- **Common Crawl**

## Connections

- Learner → Training Loop: starts training
- Training Loop → Transformer LM: trains
- Learner → Text Generation: requests generation
- Text Generation → Tokenizer: encodes prompt
- Text Generation → Transformer LM: generates with
- Text Generation → Model Checkpoint: loads weights
- Text Generation → Learner: returns text
- Learner → BPE Training: runs tokenizer work
- Learner → Systems Benchmarks: runs benchmarks
- Learner → IsoFLOP Analysis: analyzes results
- IsoFLOP Analysis → Scaling Model: fits scaling context
- Common Crawl → Quality Filtering: provides web records
- Quality Filtering → Data Utilities: uses filtering helpers
- Learner → Training Data Loader: loads training data
- Learner → Supervised Fine-Tuning: runs fine-tuning
- Learner → Preference Optimization: runs preference training
- Learner → Group Relative Policy: runs policy training
- Group Relative Policy → Math Answer Grader: grades answers
- Learner → Alignment Evaluation: evaluates safety