nanochat Custom Distributed Optimizer

nanochat은 multi-GPU 학습을 하지만, 모델을 PyTorch DistributedDataParallel wrapper로 감싸는 방식만 쓰지는 않는다.

핵심은 다음이다.

data sharding:
  each rank reads a different slice of data

gradient synchronization:
  optimizer step에서 필요한 gradient/parameter shard를 통신

optimizer logic:
  AdamW parameter와 Muon parameter를 다르게 묶고 분산 update

일반적인 DDP는 backward 중 gradient bucket을 all-reduce한다. nanochat은 optimizer 쪽에서 communication과 update를 직접 통제한다. 특히 Muon은 matrix parameter를 shape별로 묶어 업데이트하는 성격이 강해서, 일반 DDP wrapper보다 optimizer 내부에서 분산 처리를 제어하는 편이 구조적으로 자연스럽다.

이 카드는 DDP를 대체하자는 이야기가 아니다. 기준선은 여전히 PyTorch DDP다. 다만 frontier-style training code를 읽을 때는 다음 질문을 봐야 한다.

model wrapper가 통신을 숨기는가?
optimizer가 통신을 직접 스케줄링하는가?
parameter/gradient/optimizer state 중 무엇이 shard되는가?

확인

• nanochat의 multi-GPU 학습은 data parallel인가?
• PyTorch DDP wrapper를 쓰지 않아도 multi-GPU data parallel이 될 수 있는 이유는 무엇인가?
• Muon optimizer가 communication 위치에 영향을 줄 수 있는 이유는 무엇인가?