Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
My Solution and Notes for the Stanford CS336: LLM from scratch
| Date | Stars |
|---|---|
| 2026-07-31 | 264 |
| 2026-08-06 | 264 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
<h1 align="center">My SOLUTION to <br/>
CS336: Language Modeling from Scratch <br/>
(Spring 2025 Version)</h1>
- [Assignment 01: Tokenization \& Language Modeling](#assignment-01-tokenization--language-modeling)
- [Part 0: Environment Setup \& Data Download](#part-0-environment-setup--data-download)
- [Part 1: BPE Tokenizer](#part-1-bpe-tokenizer)
- [Part 2: Language Model \& Optimizer](#part-2-language-model--optimizer)
- [Part 3: Training Model](#part-3-training-model)
- [Configuration](#configuration)
- [Learning Curves](#learning-curves)
- [Genrated Sample](#genrated-sample)
- [Bonus: MoE Layer](#bonus-moe-layer)
- [Assignment 02: Flash Attention \& Parallelism](#assignment-02-flash-attention--parallelism)
- [Flash Attention](#flash-attention)
- [Parallelism](#parallelism)
- [Assignment 03: Scaling Laws](#assignment-03-scaling-laws)
- [Assignment 04: Data](#assignment-04-data)
- [Assignment 05: Alignment \& RLHF (GRPO)](#assignment-05-alignment--rlhf-grpo)
- [SFT](#sft)
- [Expert Iteration](#expert-iteration)
- [GRPO](#grpo)
This repository contains my notes and solutions for the [Stanford CS336: Language Modeling from Scratch](https://stanford-cs336.github.io/spring2025).
> [!note]
>
> This is the NEW version of my solution, for those who have see my solutions before, here is the [OLD version](https://github.com/YYZhang2025/Stanford-CS336/tree/old-main), if you are interested.
# Assignment 01: Tokenization & Language Modeling
## Part 0: Environment Setup & Data Download
We first need install the virtual environment manager `uv` to ensure reproducibility, portability, and ease of use.
```sh
pip install uv
# or
brew install uv
```
After installing `uv`, we can run any code in the repo using
```sh
uv run <python_file_path>
```
and the environment will be automatically solved and activated when necessary.
Or create and activate the environment manually using:
```sh
uv sync
source .venv/bin/activate
```
It will install all the dependencies specified in `pyproject.toml`.
Than we can download the TinyStories data and a subsample of OpenWebText
``` sh
mkdir -p data
cd data
wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt
wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt
wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_train.txt.gz
gunzip owt_train.txt.gz
wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_valid.txt.gz
gunzip owt_valid.txt.gz
cd ..
```
It will create a `data` folder in the current directory and download the required datasets into it.
## Part 1: BPE Tokenizer
With my implementation, it took about 1 min to train the BPE Tokenizer for Tinystories
```bash
Identified 10 chunks for pre-tokenization.
Pre-tokenization processes completed. Aggregating results...
Completed pre-tokenization. Vocabulary size: 59933 unique tokens.
100%|████████████████████████████████████████████████████████████████████████| 9743/9743 [00:53<00:00, 181.61it/s]
[TIME] train_bpe took 85.40s
```
And token **30mins** to pre-tokenize the whole file and save as `.bin`.
> [!note]
>
> For those who want to re-produce my results or just want to focus on the model training, I highly recommend to download my version of tokenization, you can download it through following commend:
>
> ```bash
> pip install -U huggingface_hub
> hf download YuYangZhang/TinyStory-Tokenized --repo-type dataset --local-dir datasets/tiny_stories
> ```
>
> It will download the `merges.txt`, `vocab.json`, `special_tokens.txt`,`train.bin`, `eval.bin` to local-directory.
## Part 2: Language Model & Optimizer
The Language Model is as following:
Here is the model configuration:
```Python
@dataclass
class ModelConfig:
vocab_size: int = 10000
mExcerpt of 14,073 characters
Read on GitHubWould you bet a product on this? Bounded 0–100 and slow moving.
matched fp:7ff60fe5386465b3, topic:course