i14, AWS 트레이니움2에서 선형 어텐션 모델 '라플라스 3B' 훈련 및 서비스 성공

새로운 모델 아키텍처와 성능

i14는 AWS 트레이니움2에서 선형 어텐션 기반의 라플라스 3B 모델을 네이티브 PyTorch(TorchNeuron Native)를 통해 성공적으로 훈련하고 서비스했다. 이 모델은 기존 트랜스포머 아키텍처의 소프트맥스 어텐션과 달리 Gated Linear Attention(GLA)을 사용하여 컨텍스트 길이에 관계없이 시퀀스당 29.5MB의 고정된 재귀 상태 메모리를 유지한다.

라플라스 3B는 AWS trn2.48xlarge 인스턴스에서 훈련 시 초당 약 12만 토큰의 처리량을 기록했다. 서비스 단계에서는 1,280개의 동시 시퀀스에서 초당 11,544개의 출력 토큰을 처리하는 성능을 보였다. 이는 기존 모델의 KV 캐시 증가로 인한 추론 메모리 문제를 해결하며, 예측 가능한 메모리 사용량을 제공한다.

선형 어텐션 모델의 필요성

대부분의 대규모 언어 모델(LLM)은 트랜스포머 아키텍처의 소프트맥스 어텐션을 사용하며, 이로 인해 추론 시 KV(Key-Value) 캐시가 대화나 문서 컨텍스트 길이에 비례하여 증가한다. 이러한 KV 캐시의 증가는 동시성 감소, 서비스 비용 증가, 그리고 에이전트 애플리케이션이 유용한 과거 정보를 버리도록 강제하는 문제점을 야기한다.

선형 어텐션 모델은 각 토큰에 대해 업데이트되는 고정 크기의 재귀 상태를 중심으로 계산을 재구성한다. 라플라스 3B는 GLA를 활용하여 이전 컨텍스트를 고정 크기 행렬 상태로 압축한다. 이 방식은 훈련 시 어텐션 계산 복잡도를 O(T²)에서 O(T)로 줄이고, 추론 시 메모리 상태를 시퀀스당 29.5MB로 고정하여 기존 트랜스포머 모델의 한계를 극복한다.

트레이니움2 맞춤형 구현 및 최적화

i14는 라플라스 3B를 AWS 트레이니움2에 최적화하기 위해 세 가지 주요 변경 사항을 적용했다. 첫째, GPU에 최적화된 Triton 커널 대신 표준 PyTorch 연산자를 사용하여 GLA 어텐션 커널을 재작성했다. 이 과정에서 16개의 128토큰 청크를 사용하여 NeuronCore 온칩 버퍼의 128개 파티션과 일치시켰으며, 이는 처리량을 14% 증가시켰다.

둘째, Muon 최적화기의 통신 방식을 재설계했다. MuonFSDP2에서 두 번의 all_to_all 집합 연산을 도입하여 랭크당 통신 트래픽을 24배 감소시키고 훈련 처리량을 1.63배 향상시켰다. 기존의 gather 및 broadcast 방식은 랭크당 10.87GB의 데이터를 이동시켰으나, all_to_all 방식은 0.45GB로 줄였다.

셋째, 순차적인 역전파 과정을 컴파일하기 위해 GLA 재귀의 16개 청크 단계를 정적으로 언롤링하여 컴파일러의 그래프 크기 및 온칩 메모리 제한에 맞췄다. 16개 청크(C=128) 구성은 스케줄러에 적합하고, 중간 데이터가 온칩에 유지되며, 타일이 128 SBUF 파티션과 정렬되는 최적의 조건으로 확인되었다.

AWS 트레이니움2는 16개의 트레이니움2 칩, 64개의 논리적 NeuronCore, 고대역폭 메모리, NeuronLink-v3 연결을 갖춘 trn2.48xlarge 인스턴스를 제공한다. 각 NeuronCore는 28MiB SBUF(128개 파티션)와 2MiB PSUM을 포함하며, SBUF는 소프트웨어로 관리되어 데이터 전송이 미리 스케줄링되어 효율성을 높인다.

라플라스 3B는 24개 디코더 레이어로 구성되며, 각 레이어는 소프트맥스 어텐션 대신 Simple GLA를 사용한다. 모델은 총 3.028B개의 파라미터를 가지며, 3,072의 히든 차원(D), 16개의 헤드, 2,048의 시퀀스 길이, 50,304의 어휘 크기를 지원한다. 훈련 데이터는 100B 토큰의 FineWeb-Edu를 사용했으며, Muon 최적화기와 AdamW를 혼합하여 적용했다.

개발자 및 사용자 영향

라플라스 3B의 고정된 추론 상태는 장기 실행 에이전트, 대규모 문서 분석, 그리고 예측 가능한 서비스 요구 사항을 가진 영구 메모리 애플리케이션에 큰 이점을 제공한다. 컨텍스트 길이에 따라 메모리 사용량이 증가하지 않으므로, 에이전트 시스템은 더 많은 컨텍스트를 유지하면서도 메모리 사용을 예측 가능하게 관리할 수 있다.

vLLM 엔진을 수정 없이 사용하여 모델을 서비스할 수 있었으며, 디코더 레이어의 토큰 믹서 속성만 변경하여 vLLM 풀에서 재귀 상태를 읽고 쓰도록 했다. 이는 vLLM이 모델을 어텐션 프리(attention-free)로 인식하여 시퀀스당 고정된 페이지를 할당하도록 한다.

추론 최적화 과정에서 대부분의 디코드 시간이 산술 연산이나 HBM 대역폭이 아닌 개별 연산자 디스패치에 소요됨을 확인했다. 컴파일된 디코드 단계는 이러한 오버헤드를 제거하여 처리량을 1.7배 증가시켰다. 또한, 실제 트래픽 패턴에서 새로운 디코드 배치 너비마다 발생하는 컴파일 비용(약 11초)을 고정된 너비로 패딩하여 제거할 수 있었다. 컴파일된 디코드는 지연 도착(staggered arrivals) 환경에서 3.52배의 처리량 개선을 보였다.

적용 조건 및 한계점

선형 어텐션은 모든 키와 값을 유지하는 대신 기록을 압축하므로, 특정 워크로드에 대한 검색 품질을 신중하게 평가해야 한다. 또한, BF16 재연결(reassociation)로 인해 컴파일된 출력은 이전에 실행된 결과와 약간 다를 수 있으므로, 정확한 재현성이 중요한 경우에는 이전에 실행된 디코드를 사용하는 것이 권장된다.

트레이니움2의 정적 컴파일러는 그래프 깊이와 온칩 메모리 제한을 가지므로, GLA 재귀의 청크 수와 크기를 신중하게 설정해야 한다. 16개 청크(C=128)가 이 제한을 만족하는 최적의 구성으로 확인되었다. Muon 최적화기의 Newton–Schulz 알고리즘은 훈련 단계 FLOPs의 약 2%에 불과하지만, 통신 비용이 병목 현상을 일으킬 수 있어 NeuronLink-v3와 같은 고속 인터커넥트의 역할이 중요하다.