모델을 확장하는 방법

LLM 확장은 계산, 메모리, 통신 병목을 함께 봐야 한다. TPU와 GPU의 한계를 이해하면 모델이 어디서 막히는지 예측하고, 훈련·추론 비용을 더 정확히 가늠할 수 있다.

핵심은 칩 수를 늘렸을 때 처리량이 선형에 가깝게 늘어나는 상태를 유지하는 것이다. 통신이 계산보다 느려지면 communication bound에 빠지고, 단일 칩 내부에서도 메모리 대역폭과 총 용량이 병목이 된다. 특히 행렬곱처럼 FLOPs/byte가 높은 연산은 TPU의 systolic array와 잘 맞고, GPU의 Tensor Cores도 같은 방향으로 진화하고 있다.

• Roofline analysis로 계산, 통신, 메모리 중 무엇이 한계를 만드는지 본다.
• TPU 구조와 칩 간 링크를 설명하고, sharding된 행렬 곱과 데이터 재배치를 다룬다.
• Transformer math로 파라미터 수, FLOPs, KV cache 크기를 계산한다.
• Training과 inference에서는 FSDP, Megatron sharding, pipeline parallelism, expert parallelism, disaggregated serving을 비교한다.
• Rematerialisation, ZeRO, host offload, gradient accumulation으로 메모리 요구량을 줄인다.

후반부는 LLaMA 3 적용 예제, JAX 구현과 프로파일링·디버깅, 그리고 새로운 GPU 장까지 이어진다. 결국 목표는 주어진 하드웨어에서 어떤 병렬화 조합이 가장 효율적인지 판단하고, 대략적인 학습 시간과 추론 비용을 추정할 수 있게 만드는 데 있다.