How zhasion/cs336 works
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: inference.py Text Generation · train_loop.py Training Loop · chinchilla_isoflops.py IsoFLOP Analysis · quality_classifier.py Quality Filtering
Components 22
Model Basics
- BPE Training
assignment1-basics/cs336_basics/bpe.py - Tokenizer
assignment1-basics/cs336_basics/tokenizer.py - Transformer LM
assignment1-basics/cs336_basics/module.py - Training Loop
assignment1-basics/train_loop.py - Text Generation
assignment1-basics/inference.py
Systems
- Flash Attention
assignment2-systems/cs336_systems/flash_attention.py - Distributed Training
assignment2-systems/cs336_systems/naive_ddp.py - Gradient Overlap
assignment2-systems/cs336_systems/ddp_overlap_bucketed.py - Systems Benchmarks
assignment2-systems/cs336_systems/flash_benchmarking.py
Scaling
- Scaling Model
assignment3-scaling/cs336_scaling/model.py - IsoFLOP Analysis
assignment3-scaling/chinchilla_isoflops.py
Data Quality
- Quality Filtering
assignment4-data/cs336_data/quality_classifier.py - Data Utilities
assignment4-data/cs336_data/utils.py - Training Data Loader
assignment4-data/cs336-basics/cs336_basics/data.py
Alignment
- Supervised Fine-Tuning
assignment5-alignment/cs336_alignment/sft.py - Preference Optimization
assignment5-alignment/cs336_alignment/dpo.py - Group Relative Policy
assignment5-alignment/cs336_alignment/grpo.py - Math Answer Grader
assignment5-alignment/cs336_alignment/drgrpo_grader.py - Alignment Evaluation
assignment5-alignment/scripts/evaluate_safety.py
Other components
- Learner
- Model Checkpoint
- Common Crawl
Connections 19
- 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