밑바닥부터 언어 모델을 구현하는 스탠퍼드의 CS336 수업의 첫번째 과제에 대한 공식 저장소

9bow · 6월 23, 2026, 11:30오후

assignment1-basics 소개

대규모 언어 모델을 다룰 때 많은 학습자가 고수준 라이브러리의 추상화 뒤에서 내부 동작을 건너뛰게 됩니다. 스탠퍼드의 CS336 Language Modeling From Scratch 강의는 이 추상화를 걷어내고, 토크나이저부터 트랜스포머 언어 모델과 학습 루프까지 직접 구현하도록 설계된 과목입니다. assignment1-basics는 그 강의의 첫 번째 과제(2025년 봄학기 Assignment 1: Basics)를 학생용으로 공개한 저장소입니다.

이 저장소의 목표는 언어 모델의 핵심 구성요소를 밑바닥부터 작성하고, 작은 데이터셋으로 실제 학습까지 돌려보는 것입니다. 과제의 전체 설명은 저장소에 포함된 cs336_assignment1_basics.pdf 핸드아웃에 담겨 있고, 코드는 테스트 어댑터에 연결할 함수들의 골격만 제공됩니다. 처음 상태에서는 모든 테스트가 NotImplementedError 로 실패하며, 구현을 채워 넣어야 통과합니다.

학습 효과를 위해 정답 코드가 아니라 비어 있는 함수와 테스트가 주어진다는 점이 핵심입니다. 무엇을 구현해야 하는지는 테스트 어댑터 파일에 함수 시그니처로 명확히 드러나 있어, 학습자는 각 구성요소를 직접 작성한 뒤 동일한 테스트로 검증합니다. 본 게시물에서는 assignment1-basics가 구현하도록 요구하는 구성요소와 전체 작업 흐름, 그리고 시작하는 방법을 정리합니다.

assignment1-basics가 구현하는 구성요소

무엇을 만들어야 하는지는 저장소의 tests/adapters.py 에 정의된 함수들로 확인할 수 있습니다. 이 함수들은 크게 토크나이저, 모델 구성요소, 학습 유틸리티의 세 갈래로 나뉩니다.

토크나이저 쪽에서는 BPE(Byte Pair Encoding) 토크나이저를 직접 학습하고 사용하는 부분을 구현합니다(run_train_bpe, get_tokenizer). 모델 구성요소로는 선형층과 임베딩(run_linear, run_embedding), RMSNorm(run_rmsnorm), SiLU와 SwiGLU(run_silu, run_swiglu), 회전 위치 임베딩(RoPE, run_rope), 스케일드 닷-프로덕트 어텐션과 멀티헤드 셀프 어텐션(run_scaled_dot_product_attention, run_multihead_self_attention, run_multihead_self_attention_with_rope), 그리고 이들을 묶은 트랜스포머 블록과 전체 언어 모델(run_transformer_block, run_transformer_lm)을 작성합니다.

학습에 필요한 요소도 직접 구현합니다. 소프트맥스와 교차 엔트로피 손실(run_softmax, run_cross_entropy), 그라디언트 클리핑(run_gradient_clipping), AdamW 옵티마이저(get_adamw_cls), 코사인 학습률 스케줄(run_get_lr_cosine_schedule), 배치 구성(run_get_batch), 체크포인트 저장·복원(run_save_checkpoint, run_load_checkpoint)이 포함됩니다. 즉 모델 구조부터 옵티마이저와 학습 루프 보조 기능까지, 언어 모델 학습에 필요한 거의 모든 조각을 손으로 작성하게 됩니다.

assignment1-basics의 작업 흐름

전체 과정은 데이터 준비, 토크나이저 학습, 모델 구현, 학습과 평가의 순서로 이어집니다.

환경 관리에는 uv 를 사용합니다. 저장소는 재현성과 이식성을 위해 uv 기반으로 의존성을 관리하며, uv run <python_file> 형태로 코드를 실행하면 환경이 자동으로 구성·활성화됩니다. 테스트는 uv run pytest 로 실행하고, tests/adapters.py 의 함수를 채워 자신의 구현을 테스트에 연결합니다.

assignment1-basics 시작하기

uv 를 설치한 뒤, 저장소에서 테스트를 실행하면 구현해야 할 항목들이 실패로 표시됩니다.

uv run pytest
# 초기에는 모든 테스트가 NotImplementedError로 실패합니다.
# tests/adapters.py의 함수를 완성해 구현을 테스트에 연결합니다.

학습에 쓰이는 데이터는 TinyStories와 OpenWebText 서브샘플입니다. 저장소 안내에 따라 내려받습니다.

mkdir -p data && cd data

# TinyStories
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

# OpenWebText 서브샘플
wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_train.txt.gz
gunzip owt_train.txt.gz

작은 TinyStories로 먼저 동작을 확인한 뒤, 더 큰 OpenWebText 서브샘플로 확장하는 방식으로 진행할 수 있습니다.

assignment1-basics의 라이선스

assignment1-basics는 MIT 라이선스로 공개되어 있어 개인 및 상업적 목적으로 자유롭게 사용할 수 있습니다.

CS336 강의 홈페이지

assignment1-basics 프로젝트 GitHub 저장소

더 읽어보기

이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다.

파이토치 한국 사용자 모임이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일로 보내드립니다!
텔레그램(Telegram)이나 Slack/Discord/Teams/Dooray/GoogleChat 등으로도 새 글 알림을 받으실 수 있습니다.

아래쪽에 좋아요를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~