MaxText의 1,400만 토큰 학습, 칩마다 무엇을 나눴나

1,400만 토큰짜리 시퀀스를 학습하려면 칩마다 무엇을 나눠 가져야 할까. 입력을 잘라 여러 칩에 배정하는 것만으로는 충분하지 않다. 각 토큰이 앞선 문맥을 참조하는 어텐션 단계에서는 다른 칩이 가진 정보까지 필요하기 때문이다.

구글은 2026년 11월 9일, 오픈소스 JAX 학습 프레임워크 MaxText에 긴 문맥 학습을 위한 네 가지 컨텍스트 병렬화 방식을 지원한다고 발표했다. 구글은 TPU 5p와 TPU 7에서 학습 길이를 1,400만 토큰까지 늘렸다고 보고했다. 이 수치는 학습 시퀀스의 길이다. 서비스에서 사용자가 입력할 수 있는 한도나, 그 길이의 문맥을 모델이 얼마나 잘 이해하는지에 대한 평가와는 구별해야 한다.

핵심은 긴 문맥을 나눈 뒤에도 남는 정보 교환을 어떻게 처리하느냐다. 모든 정보를 칩마다 모을 수도 있고, 계산의 담당 범위를 바꾸거나 필요한 조각을 차례로 전달할 수도 있다.

입력을 나눠도 KV가 다시 모이는 이유

긴 문맥 학습에는 서로 다른 부담이 있다. 일반적인 어텐션의 연산량은 시퀀스 길이의 제곱에 비례해 늘어난다. 한편 토큰마다 생기는 키와 값, 즉 KV를 비롯한 중간 데이터도 길이에 따라 커진다. 계산을 감당하는 문제와 데이터를 메모리에 담는 문제가 함께 생긴다.

컨텍스트 병렬화는 하나의 긴 시퀀스를 여러 칩에 나눠 각 칩이 전체 시퀀스의 중간 데이터를 보유하지 않게 한다. 하지만 앞선 토큰을 참조하는 인과적 어텐션에서는 각 토큰의 쿼리가 이전 토큰들의 KV에 접근해야 한다. 자신의 입력 조각만 들고 계산을 끝낼 수는 없다.

가장 단순한 All-gather는 어텐션을 계산할 때 다른 칩의 KV를 모두 모은다. 출발할 때 시퀀스를 나눠 가졌어도 이 단계에서는 각 칩이 전체 KV를 보유하게 된다. 칩을 더 붙여도 각 칩에 모이는 전체 KV의 크기는 줄지 않는다. 시퀀스가 충분히 길어지면 칩의 고대역폭 메모리인 HBM에 담기지 않는 이유다.

MaxText가 지원하는 다른 방식들은 이 지점에서 갈라진다. 아래 도식은 어텐션에 필요한 데이터를 각 칩이 어떻게 보유하고 주고받는지 비교한다.

헤드를 나누거나, KV를 돌리거나

A2A(All-to-All), 곧 Ulysses가 활용하는 성질은 어텐션을 헤드별로 독립 계산할 수 있다는 점이다. 처음에는 각 칩이 일부 토큰에 대한 모든 헤드를 맡는다. A2A는 어텐션 직전에 담당 대상을 바꿔, 각 칩이 일부 헤드에 대해서만 전체 토큰을 갖도록 재배치한다.

이렇게 하면 각 칩의 KV 저장량은 컨텍스트 병렬화 정도에 따라 줄어든다. 대신 나눠 맡길 수 있는 헤드 수가 확장의 한도가 된다. 텐서 병렬화도 같은 헤드 차원을 나누므로 두 방식이 이 자원을 함께 사용한다. 그룹 쿼리 어텐션에서는 KV 헤드 수가 제약이 된다. 또 데이터 교환이 끝나야 전체 시퀀스에 대한 계산을 시작할 수 있어, 이 교환 자체를 어텐션 계산 뒤에 숨길 수는 없다.

P2P, 곧 링 어텐션은 토큰 단위 분할을 유지한다. 각 칩은 자신의 쿼리와 현재 보유한 KV 조각으로 어텐션의 일부를 계산하고 결과를 합쳐 간다. 그동안 다음 KV 조각을 이웃 칩에서 전달받는다. 전체 KV를 한 번에 저장하는 대신, 필요한 조각이 칩 사이를 돌며 계산에 참여한다.

이 방식은 헤드 수에 따른 분할 한도를 피하고 통신과 계산을 겹칠 수 있다. 다만 계산 조각이 작아지면 칩 내부의 어텐션 커널 효율이 떨어질 수 있다. 뒤쪽 토큰일수록 참조할 앞선 문맥이 많아지는 인과적 어텐션의 특성 때문에 칩 사이의 계산량을 맞추는 작업도 필요하다.

네 번째 방식인 USP는 A2A와 링을 결합한다. A2A 그룹이 토큰을 먼저 모으면 링에서 한 번에 넘기는 조각이 커지고, 링을 도는 홉 수는 줄어든다. 이동 한 번에 수행하는 계산이 늘면 그 계산 동안 전송을 진행할 여지가 커진다. A2A만 쓸 때의 헤드 수 제약을 넘어 확장하면서, 긴 링에서 누적되는 이동 지연도 줄이려는 구성이다.

1,400만 토큰 결과가 보여 주는 범위

구글은 Llama 3.1 8B를 사용하는 64칩 비교를 제시하며, 1,400만 토큰 P2P 학습의 MFU를 다음과 같이 보고했다.

MFU(Model FLOPs Utilization)는 모델 학습 연산이 하드웨어의 이론적 연산 성능을 얼마나 활용하는지 나타내는 지표다. 표의 차이는 이 학습 구성에서의 연산 활용률 차이로 읽어야 한다. 모델의 정답률이나 서비스 응답 속도를 비교한 값은 아니다.

구글은 TPU 7 실험에서는 링 방식이 Ulysses와 USP보다 나았다고 밝혔다. 반면 일반적인 선택 지침에서는 KV 헤드 수가 허용하면 USP를 권한다. 특정 하드웨어에서 얻은 결과와 여러 구성을 위한 권고의 범위가 다르다. 전체 KV가 메모리에 들어가는지, 헤드를 얼마나 나눌 수 있는지, 각 칩에 충분한 계산량이 남는지에 따라 선택도 달라진다.

발표된 기능은 MaxText v0.2.4에서 제공된다. 이번 발표에서 1,400만 토큰 학습을 뒷받침하는 핵심은 각 칩이 전체 문맥을 복제하는 대신, 계산할 범위를 나누거나 KV 조각을 순환시키며 필요한 정보를 확보하도록 만든 데 있다.