[Paper] 다중 GPU에서 Tensor Parallelism으로 State‑Space Models 확장

Source: arXiv - 2602.21144v1

번역을 진행하려면 번역이 필요한 실제 텍스트(예: 초록, 본문, 섹션 등)를 제공해 주세요.
텍스트를 알려주시면 원본 형식과 마크다운을 유지하면서 한국어로 번역해 드리겠습니다.

개요

선택적 상태‑공간 모델(SSM)은 대형 언어 모델(LLM)에서 매우 긴 컨텍스트를 처리하기 위한 강력한 대안으로 부상하고 있습니다. 논문 *“Scaling State‑Space Models on Multiple GPUs with Tensor Parallelism”*은 실용적인 병목 현상을 해결합니다: 단일 GPU의 메모리 또는 대역폭이 부족할 때 SSM 기반 LLM을 효율적으로 실행하는 방법입니다. 이미 Transformer 확장에 널리 사용되고 있는 텐서 병렬화(TP) 기법을 적용함으로써, 저자들은 중요한 순환 경로를 빠르고 통신량을 최소화한 상태로 유지하면서 무거운 SSM 연산을 여러 GPU에 분산시키는 방법을 보여줍니다.

Key Contributions

• TP‑aware SSM 설계는 큰 투영 행렬을 분할하지만 토큰당 순환 상태 업데이트를 각 GPU에 로컬로 유지합니다.
• State‑cache 메커니즘은 프리필과 이후 디코딩 단계에서 SSM 은닉 상태를 재사용하여 “prefill‑to‑decode”(TTFT) 속도 향상을 가능하게 합니다.
• Quantized All‑Reduce는 TP 집계 단계에서 동기화 대역폭을 최대 약 18 %까지 감소시키면서 수치적 안정성을 유지합니다.
• Comprehensive evaluation은 NVIDIA A6000 및 A100 클러스터에서 세 가지 실제 SSM 기반 LLM 패밀리(Mamba, Falcon‑Mamba, Zamba)를 대상으로 수행했으며, GPU 수를 1개에서 4개로 확장할 때 1.6–4.0배의 처리량 향상을 보여줍니다.
• Open‑source implementation(또는 상세한 의사코드)은 최소한의 변경으로 기존 추론 스택에 쉽게 통합될 수 있습니다.

Methodology

1. Tensor Partitioning – 저자들은 SSM 믹서의 패킹된 가중치 텐서(긴 거리 투영과 로컬 믹싱 커널을 모두 포함)를 하나의 “큰 행렬”로 취급합니다. 이 행렬을 특징 차원에 따라 GPU들에 걸쳐 분할하는데, 이는 Transformer에 대한 고전적인 텐서 병렬(TP)과 유사합니다.
2. Local Recurrence – Transformer의 자체 주의(self‑attention)와 달리, SSM은 각 토큰에 대해 숨겨진 상태를 순차적으로 업데이트합니다. 설계는 각 GPU가 자체적인 숨겨진 상태 슬라이스를 유지하도록 하여, 중요한 경로에서 재귀적 업데이트가 GPU 간 통신을 필요로 하지 않게 합니다.
3. State Cache Across Prefill & Decode – 초기 “prefill”(긴 프롬프트 처리) 단계에서 숨겨진 상태를 캐시합니다. 모델이 토큰‑별 디코딩으로 전환될 때, 캐시된 상태를 재사용하여 이미 처리된 토큰에 대한 비용이 많이 드는 투영을 다시 계산하지 않도록 합니다.
4. Quantized All‑Reduce – 각 토큰 후에 모든 GPU의 부분 결과를 합산(All‑Reduce)하여 최종 숨겨진 표현을 만들어야 합니다. 저자들은 축소하기 전에 텐서를 8‑bit으로 양자화하고, 이후에 역양자화함으로써 인터커넥트를 통해 이동되는 데이터 양을 크게 줄입니다.
5. Benchmark Suite – 세 가지 모델 패밀리에 대해 추론을 수행하고, 컨텍스트 길이(2 K–64 K 토큰)와 배치 크기를 다양하게 바꾸어 원시 처리량(tokens/s)과 엔드‑투‑엔드 요청 지연 시간을 모두 측정합니다.

결과 및 발견

• 긴 컨텍스트 장점: 속도 향상은 컨텍스트 길이가 길어질수록 커지는데, 이는 SSM의 토큰당 비용이 지배적이며 TP 분할이 GPU당 메모리 압력을 감소시키기 때문입니다.
• TTFT 이점: 프리필과 디코드 사이에 상태를 캐시하면 프리필 지연 시간이 평균 약 30 % 감소합니다.
• 통신 효율성: 양자화된 All‑Reduce는 PCIe/NVLink 트래픽을 낮춰, 인터커넥트 대역폭이 제한된 클러스터에서도 이 접근 방식을 실현 가능하게 합니다.

Practical Implications

• Deployers can now run SSM‑based LLMs with 32 K+ context on 2‑4 GPU nodes without hitting memory limits, opening up use‑cases like document‑level QA, code‑base search, or long‑form generation.
• Cost‑effective scaling: Instead of buying a single massive GPU (e.g., H100 80 GB), teams can stitch together more affordable A6000/A100 cards and still achieve near‑linear speed‑up.
• Framework integration: The design maps cleanly onto existing TP libraries (e.g., Megatron‑LM, DeepSpeed), meaning developers can add SSM support with a few configuration changes rather than a full rewrite.
• Lower latency for real‑time apps: The TTFT cache reduces the “warm‑up” penalty when switching from prompt ingestion to token‑by‑token generation, which is crucial for chat‑style assistants.
• Quantized communication provides a template for other memory‑heavy models (e.g., retrieval‑augmented Transformers) where All‑Reduce becomes a bottleneck.

제한 사항 및 향후 연구

• 하드웨어 의존성: 가장 큰 성능 향상은 고속 NVLink를 갖춘 GPU에서 관찰되며, 느린 인터커넥트에서는 양자화된 All‑Reduce가 여전히 제한 요소가 될 수 있습니다.
• 모델‑특정 튜닝: 파티셔닝 전략은 특정 형태의 SSM 믹서를 전제로 합니다; 형태가 불규칙하거나 희소화된 믹서를 사용하는 모델은 맞춤형 슬라이싱 로직이 필요할 수 있습니다.
• 정밀도 트레이드‑오프: 8‑비트 양자화가 감소 단계에서는 작동하지만, 저자들은 일부 다운스트림 작업에서 (<0.2 BLEU) 작은 성능 저하를 언급합니다—향후 연구에서는 혼합 정밀도 또는 적응형 양자화를 탐색할 수 있습니다.
• 추론을 넘어 확장: 이 논문은 추론에 초점을 맞추고 있으며, 동일한 TP 방식을 훈련(특히 그래디언트 누적과 함께) 적용하는 것은 아직 해결되지 않은 과제입니다.

핵심 요약: 텐서 병렬성과 영리한 상태 캐싱 및 양자화된 통신을 결합함으로써, 이 작업은 오늘날의 다중 GPU 클러스터에서 대규모 장기 컨텍스트 SSM 추론을 실용화합니다—이는 차세대 LLM 기반 애플리케이션의 물결을 가속화할 수 있는 진전입니다.

저자

• Anurag Dutt
• Nimit Shah
• Hazem Masarani
• Anshul Gandhi

논문 정보

• arXiv ID: 2602.21144v1
• 분류: cs.DC, cs.LG
• 출판일: 2026년 2월 24일
• PDF: PDF 다운로드