최고의 AI 성능 엔지니어가 되기 위해 필요한 모든 자료 모음

가장 뛰어난 AI 성능 엔지니어가 되는 데 필요한 모든 자료를 모았습니다.

시리즈를 따라잡으려면 팔로우하고 저장하세요. 스레드에 링크가 있습니다 🧵

파트 7: 트랜스포머 추론 효율적으로 확장하기

Reiner Pope와 공동 저자들이 TPU v4에서 PaLM 추론을 연구했습니다. 이 논문은 성능 엔지니어들이 가중치(weights), 활성화 값(activations), KV 캐시가 어디에 위치해야 하는지, 매 단계마다 무엇이 이동해야 하는지, 그리고 어떤 전송이 지연 시간(latency)을 제한하는지에 대해 추론할 수 있는 방법을 제공합니다.

논문의 메커니즘은 추론 배포를 위해 다음과 같은 점검 사항들을 제안합니다:

• 프리필(prefill)과 디코드(decode)는 서로 다른 산술 강도(arithmetic intensity), 즉 전송된 바이트당 계산량을 가질 수 있습니다. 프리필은 하나의 패스로 많은 프롬프트 토큰을 처리할 수 있는 반면, 표준 자기회귀적 디코드는 단계마다 시퀀스당 새로운 토큰 하나를 생성합니다. 더 큰 디코드 배치(decode batches)는 가중치 읽기(weight reads)를 상각하지만, 시퀀스별 KV 캐시 트래픽을 추가합니다. 컨텍스트 길이와 함께 배치 크기를 스윕하고 선형 레이어(linear layers)와 어텐션(attention)을 프로파일링하여 대역폭 또는 계산 한계를 파악하세요.

• 텐서 병렬성(tensor parallelism)은 통신 병목 현상(communication bottleneck)을 제거하지 않으면서 로컬 계산량을 줄일 수 있습니다. 논문의 1D 레이아웃에서는 고정된 워크로드에 대해 칩 개수가 증가해도 활성화 값 집계 시간(activation-aggregation time)이 대략 일정하게 유지됩니다. 2D 파티셔닝은 가중치 차원(weight dimensions)을 모두 분할하여 각 칩이 더 작은 활성화 샤드(activation shards)를 통신하도록 합니다. 병렬성 정도를 늘리기 전에 실제 상호 연결(interconnect) 및 행렬 차원을 고려하여 이를 평가하세요.

• 최적의 이동 텐서는 토큰 배치에 따라 달라집니다. 가중치 정지 레이아웃(weight-stationary layouts)은 활성화 값을 교환하고, 가중치 수집 레이아웃(weight-gathered layouts)은 더 많은 가중치를 전송하여 활성화 값 교환을 줄입니다. 대규모 프리필은 그들의 활성화 텐서가 크기 때문에 이러한 트레이드오프를 정당화할 수 있습니다. 그러한 레이아웃들을 비교할 때는 프리필 중 프롬프트 토큰을 포함하여 순방향 패스(forward pass)마다 처리되는 토큰 수를 계산하세요.

• multiquery attention은 query head에 걸쳐 하나의 KV head를 공유하지만, head-sharded 레이아웃을 사용하면 이 캐시를 장치(device) 간에 복제할 수 있습니다. 해당 논문에서는 디코드 과정에서 어텐션을 배치(batch) 전체에 분할하고 현재 단계의 활성화 값(activations)을 교환하여 칩당 캐시 읽기(cache reads)를 줄입니다. vLLM의 grouped-query attention 배포 시, 텐서 병렬성(tensor parallelism)이 KV-head 개수를 초과하는지 확인하고, 그 전에 GPU당 용량을 추정할 때 캐시 복제가 발생하는지 점검해야 합니다.

• KV 캐시 용량과 프리필 어텐션 중간 값(intermediates)은 별개의 메모리 문제입니다. 고정된 배치 크기와 query head 개수에서 KV head를 공유한다고 해서 전체 어텐션 스코어 텐서가 줄어드는 것은 아닙니다. 밀집(dense)한 프리필의 경우, 이 텐서의 요소 수(element count)는 시퀀스 길이의 제곱에 따라 증가합니다. 부록 G에서는 임시 할당을 줄이기 위한 마이크로배칭(microbatching)과 전체 스코어 행렬(score matrix)을 생성하지 않도록 하는 FlashAttention에 대해 논하고 있습니다. 이러한 구분을 사용하여 메모리 실패가 유지되는 KV 상태에서 발생하는 것인지, 아니면 일시적인 어텐션 할당에서 발생하는 것인지를 파악해야 합니다.

• 통신 오버랩(communication overlap)은 실행 스케줄과 텐서 레이아웃에 따라 달라집니다. 해당 저자들은 Looped CollectiveEinsum을 사용하여 집합 연산(collectives)과 행렬 곱셈을 중첩시키고, 오버랩 기회를 노출하는 출력 분할(output sharding) 선택까지 포함합니다. GPU 타임라인을 검사하여 노출된 통신 및 종속적 컴퓨팅 지연(dependent compute stalls)이 있는지 확인해야 합니다. 비동기 집합 연산이 그 비용이 숨겨진다는 것을 의미하지는 않습니다.

• 양자화(quantization)는 무엇을 양자화하느냐에 따라 다른 병목 현상을 초래합니다. 해당 논문은 가중치(weights)를 int8로 저장하면서 bfloat16 행렬 연산을 유지합니다. 이는 가중치 로딩 트래픽을 줄이지만, 컴퓨팅 바운드 배치(compute-bound batches)에 대해 더 빠른 산술 연산을 제공하지는 않습니다. 측정된 병목 현상과 비교하여 가중치, 활성화 값, KV 캐시의 정밀도를 평가하고 두 단계 모두에 미치는 영향을 벤치마킹해야 합니다.

• 애플리케이션이 필요로 하는 동작 지점을 벤치마킹합니다. 논문에서는 토큰당 가속기 시간과 모델 FLOPS 활용도를 기준으로 레이턴시를 비교합니다. 서빙 시스템에서는 목표 부하 조건에서 첫 번째 토큰까지 걸리는 시간(time to first token), 토큰 간 레이턴시(inter-token latency), 처리량(throughput)을 추적해야 합니다. 큐잉 시간(queueing time)은 프리필 시간(prefill time)과 분리하여 측정합니다.

실질적인 활용 사례는 최적화 실험 범위를 좁히는 것입니다. 실제 워크로드에 대한 컴퓨팅 비용, HBM 트래픽, 집합 비용을 추정하고, 병목 지점을 검사하며, 관련 레이아웃이나 정밀도를 변경한 후 동일한 레이턴시 목표에서 재측정합니다.