nanochat: tự train GPT-2 chỉ với khoảng $48

nanochat là gì

nanochat là một codebase viết từ đầu để train một mô hình ngôn ngữ nhỏ kiểu ChatGPT trên một node GPU duy nhất, bao gồm tokenization, pretraining, finetuning, evaluation và inference trong cùng một chỗ. Một núm vặn duy nhất, `--depth`, quy định kích thước transformer và tự suy ra phần còn lại của hyperparameter. Script tham chiếu `runs/speedrun.sh` chạy trọn pipeline đó trên một node 8xH100 và đưa cho bạn một CLI để chat khi xong.

nanochat phủ những gì từ đầu đến cuối

• ✓Pipeline đầy đủ trong một repo: train tokenizer (BPE kiểu GPT-4), pretraining, supervised finetuning, RL, evaluation và inference đều nằm trong `nanochat/` và `scripts/`.
• ✓Một núm vặn độ phức tạp duy nhất: đặt `--depth` cho transformer sẽ tự suy ra độ rộng, số attention head, lịch learning rate, thời lượng train và weight decay cho một mô hình compute-optimal ở kích thước đó.
• ✓Kiểm soát precision tường minh thay vì `torch.amp.autocast` — một biến global `COMPUTE_DTYPE` duy nhất (bfloat16, float32 hoặc float16) tự nhận diện phần cứng và có thể ghi đè bằng biến môi trường `NANOCHAT_DTYPE`.
• ✓Các tác vụ evaluation dựng sẵn — ARC, GSM8K, HumanEval, MMLU và SmolTalk — nằm trong `tasks/` để chấm điểm mô hình chat, cộng với một bộ đánh giá DCLM CORE score cho base model.
• ✓Một inference engine dùng KV cache (`nanochat/engine.py`) và một module thực thi công cụ (`nanochat/execution.py`) cho phép mô hình đã train chạy code Python như một tool.
• ✓Một bảng xếp hạng GPT-2 speedrun công khai, theo dõi thời gian wall-clock và DCLM CORE score để tái lập năng lực GPT-2, kèm commit và người đóng góp cho từng mục.

Cài đặt với uv

nanochat dùng uv để quản lý dependency thay vì một requirements.txt thuần hay môi trường Conda. Chạy `uv sync --extra gpu` trên máy CUDA (A100, H100, v.v.) hoặc `uv sync --extra cpu` cho CPU-only hay Apple Silicon, sau đó `source .venv/bin/activate`. Nếu định sửa code, thêm `--group dev` vào lệnh sync gpu để cài thêm pytest, matplotlib, ipykernel và transformers phục vụ development. Không có package pip riêng để cài — bạn clone repo và làm việc ngay bên trong đó.

Chạy script speedrun GPT-2

Toàn bộ workflow tái lập GPT-2 nằm trong một file, `runs/speedrun.sh`, thiết kế để chạy trên node 8xH100 — thuê một node GPU từ nhà cung cấp cloud, chạy `bash runs/speedrun.sh` trong một screen session, rồi chờ khoảng 1,5 giờ. Xong thì activate lại venv và chạy `python -m scripts.chat_cli` để chat với mô hình qua một CLI đơn giản. Cùng đoạn code đó chạy được trên một GPU đơn bằng cách bỏ `torchrun`, chuyển sang gradient accumulation, nhưng bạn sẽ phải chờ lâu gấp khoảng tám lần. Nếu GPU có dưới 80GB VRAM, bạn cần giảm `--device-batch-size` từ mặc định 32 xuống 16, 8, 4, 2 hoặc 1 cho tới khi hết out-of-memory.

Lựa chọn thay thế để train LLM quy mô nhỏ

Câu hỏi thường gặp

Vấn đề mà nanochat giải quyết

Phần lớn code train LLM công khai hiện nay hoặc là script đồ chơi train một mô hình tí hon trong notebook, hoặc là một framework production như những hệ thống đứng sau các sản phẩm ChatGPT thực thụ, được bọc trong nhiều lớp config dành cho cả team chứ không phải một người dùng GPU đi thuê. nanochat nằm ở khoảng giữa: nó chạy trọn pipeline từ tokenize đến chat, cho ra một mô hình GPT-2 class thật sự nói chuyện được, trên phần cứng mà một người có thể thuê trong vài giờ, mà không giấu giai đoạn nào sau lớp trừu tượng.

Trường hợp sử dụng tốt nhất

• •Học cách một mô hình kiểu ChatGPT thực sự được xây dựng, từng giai đoạn một, bằng cách đọc và chạy code train thật thay vì xem sơ đồ hay đọc blog post.
• •Chạy vòng lặp thử nghiệm nhanh tại chỗ — ví dụ trong README train một mô hình 12 layer, cỡ GPT-1, trong khoảng 5 phút để test một thay đổi code trước khi chạy full GPT-2 class.
• •Tái lập một điểm cụ thể trên đường cong compute-optimal scaling cho nghiên cứu, dùng `runs/scaling_laws.sh` hoặc `runs/miniseries.sh` để quét `--depth` qua một dải kích thước mô hình.
• •Có một mô hình ngôn ngữ nhỏ, chat được thật sự, với chi phí khoảng $48 tiền GPU cloud, mà không cần license một mô hình hosted hay xin API access.

Ai nên thử — và ai nên bỏ qua

nanochat hợp với developer và researcher muốn tự tay chạy mọi giai đoạn train LLM — tokenizer, pretraining, SFT, RL và inference — và không ngại thuê một node 8xH100 hoặc chờ một lượt train chậm hơn trên GPU đơn. Đây là lựa chọn tốt nếu bạn học bằng cách đọc code chạy thật, dày đặc, thay vì đọc documentation, và không phiền vì không có hệ thống config để dựa vào. Bỏ qua nếu bạn cần một mô hình chạy production ngay ngày mai, cần hỗ trợ multi-node hay phần cứng không phải NVIDIA mà README không hề nhắc tới, hoặc thà finetune một checkpoint có sẵn qua thư viện như transformers hay unsloth còn hơn train từ đầu.

Repo liên quan