How Deep Learning Computes

PyTorch Transformer 모델 직접 작성하기

이 경로는 PyTorch로 작은 decoder-only Transformer를 직접 작성하는 모델 구현 경로입니다.

목표는 거대한 모델을 학습하는 것이 아니라, Transformer가 PyTorch 코드에서 어떤 module과 tensor operation으로 나뉘는지 손으로 확인하는 것입니다. 이 경로가 끝나면 token embedding, RMSNorm, causal self-attention, MLP, decoder block, LM head가 어떻게 연결되는지 설명할 수 있어야 합니다.

이 path는 이후 두 경로의 기반입니다.

PyTorch Transformer 모델
  -> PyTorch 학습
  -> PyTorch profiling / kernel optimization
  1. PyTorch Tensor와 nn.Module 기본기 — Transformer를 직접 작성하기 전에 PyTorch tensor shape, device, dtype, nn.Module의 역할을 잡는다.
  2. PyTorch Token Embedding과 LM Head — token id를 hidden vector로 바꾸고, 마지막 hidden state를 vocab logits로 되돌리는 양끝 구조를 구현한다.
  3. PyTorch RMSNorm과 Residual Path — Transformer block에서 normalization과 residual connection이 activation 흐름을 어떻게 안정화하는지 코드로 구현한다.
  4. PyTorch Causal Self-Attention 구현 — Q/K/V projection, causal mask, softmax, attention output을 PyTorch tensor operation으로 직접 작성한다.
  5. PyTorch MLP와 Decoder Block 조립 — Attention, MLP, RMSNorm, residual path를 하나의 decoder block으로 묶는다.
  6. PyTorch Decoder-only Transformer 조립 — embedding, 여러 decoder block, final norm, LM head를 묶어 작은 language model을 완성한다.