TL;DR

본 논문은 USP(Unified Sequence Parallelism)를 제안한다. 이는 DeepSpeed-Ulysses와 Ring-Attention을 하나의 설정 가능한 프레임워크로 결합하여 긴 문맥 transformer 연산을 여러 GPU에 분산하는 하이브리드 접근법이다. 핵심 혁신은 SP 프로세스 그룹을 2D 메시로 구성하는 것이다. Ulysses는 행 방향으로 작동하고(AllToAll 사용), Ring은 열 방향으로 작동하여(P2P 사용), 하드웨어 토폴로지와 모델 아키텍처에 유연하게 조정할 수 있다. USP는 두 개의 8xA800 노드에서 LLAMA3-8B를 208K 시퀀스 길이로 학습할 때 47% MFU를 달성하며, 단일 방법만 사용할 때보다 최대 13% 성능이 향상된다. 중요한 통찰은 USP가 메모리 효율성 측면에서 tensor parallelism보다 무조건 우수한 것은 아니라는 점이다. TP-sp는 최대 시퀀스 길이를 위해 여전히 필요하며, 본 논문의 주요 가치는 SP가 DP, TP, ZeRO, PP와 4D 하이브리드 parallelism에서 어떻게 상호작용하는지에 대한 체계적 분석에 있다.

• Paper Link: https://arxiv.org/pdf/2405.07719

Related Papers

시퀀스 병렬화 기반:

• Ring Self-Attention - 링 기반 시퀀스 병렬화의 기초
• Blockwise RingAttention - USP에서 통합된 Ring 방법
• DeepSpeed Ulysses - USP에서 통합된 Ulysses 방법

관련 연구:

• DistCA: Efficient Long-Context Language Model Training by Core Attention Disaggregation - 분산 어텐션 분리
• Unlocking Agentic RL Training for GPT-OSS - 대규모 학습 시스템

하이브리드 병렬화:

• LoongTrain - 2D 어텐션을 활용한 하이브리드 접근법
• Tensor Parallelism - 텐서 병렬화와의 결합
• GPipe - 파이프라인 병렬화와의 통합

긴 시퀀스 처리:

• DISTFLASHATTN - 분산 어텐션 계산
• Striped Attention - 효율적인 시퀀스 분배
• Context Parallelism for Scalable Million-Token Inference - 추론 시 컨텍스트 병렬화

시스템 최적화:

• Reducing Activation Recomputation in Large Transformer Models - 메모리 효율적인 훈련
• Efficiently Scaling Transformer Inference - 효율적 추론 시스템

Takeaways

1. Contribution

생성형 AI 모델의 문맥 길이가 급속도로 증가하고 있다. Claude의 100K 토큰, GPT-4의 128K, Gemini 1.5 Pro의 1천만 토큰에 이르기까지, 극도로 긴 시퀀스의 연산을 여러 디바이스에 분산할 수 있는 parallelism 전략에 대한 긴급한 필요가 생겨났다. Sequence Parallelism(SP)은 입력 시퀀스를 시퀀스 차원을 따라 분할하여 이 문제를 해결한다. $Q$, $K$, $V$, $O$ 텐서를 디바이스에 분산시켜 각 디바이스가 전체 시퀀스의 일부만 처리하도록 한다. 2023년 말까지 두 가지 획기적인 접근법이 등장했다. DeepSpeed-Ulysses(SP-Ulysses)는 AllToAll collective communication을 사용하여 텐서를 시퀀스 분할 형태에서 헤드 분할 형태로 재분배하고, Ring-Attention(SP-Ring)은 링 토폴로지에서 peer-to-peer communication을 사용하여 KV 블록 전송과 attention 연산을 중첩시킨다. 두 접근법 모두 중복된 메모리 소비나 비효율적인 통신 패턴으로 어려움을 겪었던 초기 시도에 비해 큰 진전을 나타냈다.

하지만 각 방법은 실용적 채택을 방해하는 중요한 한계를 가지고 있다. SP-Ulysses는 attention 헤드 수 $h_c$에 의해 근본적으로 제약받는다. AllToAll 연산이 헤드 차원을 따라 텐서를 재분배하기 때문에 parallelism 정도가 $h_c$를 초과할 수 없다. 이는 Grouped Query Attention(GQA)과 Multi-Query Attention(MQA) 같은 현대 아키텍처에서 심각한 문제가 된다. 예를 들어, LLAMA3-8B는 8개의 KV 헤드만 가진 GQA를 사용하므로 SP-Ulysses는 최대 parallelism 정도가 8로 제한된다. MQA(KV에 대해 $h_c = 1$)의 경우 SP-Ulysses는 전혀 작동할 수 없다. 게다가 Tensor Parallelism도 헤드 차원을 따라 분할하므로 SP-Ulysses와 TP는 직접적으로 충돌하며, 효과적으로 결합하는 것이 불가능하다. 반면 SP-Ring은 attention 연산을 더 작은 블록 수준 행렬 곱셈으로 세분화하기 때문에 연산 비효율성을 겪는다. 이는 FlashAttention 같은 융합 연산자의 효율성을 감소시킨다. 완벽한 통신-연산 중첩이 있더라도 SP-Ring의 총 실행 시간은 SP-Ulysses보다 느리다. 추가로, vanilla SP-Ring은 causal attention mask로 심각한 부하 불균형을 초래한다. 나중 시퀀스 세그먼트를 처리하는 디바이스가 초기 세그먼트를 처리하는 디바이스보다 최대 7배 더 많은 연산을 수행해야 한다.

USP의 핵심 기여는 Ulysses와 Ring이 상호 배타적인 경쟁 관계가 아니라 2D 프로세스 그룹 메시를 통해 통합될 수 있는 상호보완적 전략이라는 관찰이다. $N$개의 디바이스를 $N_1 \times N_2$ 크기의 2D 그리드로 구성하여($N_1$은 Ring 정도, $N_2$는 Ulysses 정도), USP는 Ulysses 그룹(메시의 행) 내에서 AllToAll communication을 적용하고 Ring 그룹(메시의 열)을 가로질러 P2P 링 통신을 적용한다. 이 하이브리드 공식은 두 방법을 일반화한다. $N_2 = N$으로 설정하면 순수 SP-Ulysses가 되고, $N_1 = N$으로 설정하면 순수 SP-Ring이 된다. 실용적 의의는 세 가지다. 첫째, USP는 SP-Ulysses의 헤드 수 제약을 제거한다. 16개 디바이스에서 8개의 KV 헤드를 가진 LLAMA3-8B의 경우, $N_2 = 8$(Ulysses)과 $N_1 = 2$(Ring)로 설정할 수 있다. 둘째, USP는 토폴로지 인식 통신을 가능하게 한다. AllToAll 연산은 고대역폭 노드 내 링크(예: NVLink)에 배치될 수 있고, P2P 링 통신은 저대역폭 노드 간 연결(예: RDMA 또는 Ethernet)에서 작동한다. 셋째, USP에 통합된 부하 균형 Ring-Attention 변형은 토큰 재정렬을 통해 causal mask 불균형을 해결한다.

통합 SP 방법 자체를 넘어, 본 논문은 SP와 Data Parallelism, Tensor Parallelism(Megatron-LM의 TP-sp 변형 포함), ZeRO-1/2/3, Pipeline Parallelism 간의 체계적 비교 분석을 제공한다. 이 분석은 4D 하이브리드 parallelism 시스템 설계를 위한 7가지 구체적인 “팁”을 도출하며, 통신 비용 공식과 메모리 풋프린트 비교로 뒷받침된다. 이 체계적 접근은 기존 문헌에서 이러한 parallelism 전략을 복잡한 상호작용을 다루지 않고 독립적으로 다루었던 중요한 공백을 채운다. 권장되는 프로세스 그룹 순서(가장 안쪽에서 바깥쪽으로 TP, SP-Ulysses, SP-Ring, ZeRO-DP, PP)는 대규모 학습 시스템을 구축하는 실무자에게 실행 가능한 지침을 제공한다.

2. Methodology

2.1 Core Intuition

Transformer에서 sequence parallelism의 근본적 과제는 self-attention 연산의 본질에 있다. 입력 $Q, K, V \in \mathbb{R}^{L \times d}$가 주어졌을 때, attention 연산은 $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V$를 계산한다. 여기서 $L$은 시퀀스 길이이고 $d = h_c \times h_s$는 은닉 차원(헤드 수와 헤드 크기의 곱)이다. 결정적 어려움은 시퀀스 차원 $L$이 행렬 곱셈 $QK^T$에서 공통(축약) 차원으로 작용한다는 것이다. 이는 $L$을 따라 단순하게 분할하면 softmax 정규화의 수학적 정확성이 깨진다는 것을 의미한다. 각 디바이스가 부분적인 내적에 대해 로컬 softmax를 계산하여 잘못된 attention 가중치를 생성하게 된다.

SP-Ulysses는 시퀀스 차원에서 헤드 차원으로 분할을 변환하는 AllToAll communication을 수행하여 이 문제를 완전히 우회한다. AllToAll 후, 각 디바이스는 헤드의 부분집합에 대해 완전한 시퀀스를 보유하므로 표준 FlashAttention 연산이 가능하다. SP-Ring은 온라인 softmax 알고리즘에서 영감을 받은 다른 접근법을 취한다. KV 블록을 링 방식으로 반복 처리하며, 블록 간 올바른 점진적 softmax 계산을 가능하게 하는 실행 통계(log-sum-exp 값)를 유지한다. 이는 FlashAttention의 타일링 전략을 분산 환경으로 확장한 것과 수학적으로 동등하다.

USP의 통찰은 이 두 접근법이 직교하는 차원에서 작동한다는 것이다. Ulysses는 헤드 차원을 가로질러 재분배하고, Ring은 각 헤드의 연산 내에서 시퀀스 차원을 가로질러 재분배한다. 이들을 계층적으로 구성함으로써, USP는 먼저 AllToAll을 사용하여 Ulysses 하위그룹 내에서 헤드를 재분배하고, 그런 다음 Ring communication을 사용하여 Ring 하위그룹 간에 나머지 시퀀스 연산을 분배한다. 이 2단계 분해는 서로 다른 통신 링크가 서로 다른 대역폭을 가진 계층적 네트워크 토폴로지에 자연스럽게 매핑된다. 이는 현대 GPU 클러스터의 표준이다(노드 내 고속 NVLink, 노드 간 저속 RDMA/Ethernet).

Causal attention에서 SP-Ring의 부하 균형 메커니즘은 토큰 재정렬 전략을 기반으로 한다. Causal attention에서 $QK^T$ 행렬은 하삼각 행렬이며, 이는 시퀀스의 초기 토큰이 더 적은 연산을 필요로 하고(더 적은 유효 attention 대상) 나중 토큰이 더 많은 연산을 필요로 함을 의미한다. 시퀀스가 균등하게 분할되면, 4개 GPU에서 마지막 디바이스가 첫 번째 디바이스보다 거의 7배 더 많은 연산을 처리한다. 해결책은 각 디바이스가 초기 토큰과 후기 토큰의 혼합을 받도록 토큰을 재정렬하는 것이다. 구체적으로, 시퀀스는 $2 \times N_1$개의 청크로 분할되고, 디바이스 $r$은 위치 $r$과 $2N_1 - r - 1$의 청크를 받는다. 이는 각 디바이스가 동일한 양의 유효(마스크되지 않은) 연산을 처리하도록 보장하여 완벽한 부하 균형을 달성한다. 동일한 재정렬이 위치 인코딩 파라미터(예: RoPE)에 적용되며, 이는 길이 $bs \times L$의 정수 토큰 인덱스에서 작동하므로 오버헤드는 무시할 수 있다.

2.2 Model Architecture

USP는 새로운 모델 아키텍처가 아니라 표준 transformer 블록 내 attention 모듈을 위한 분산 연산 전략이다. 시스템 아키텍처는 2D 프로세스 메시를 가로지르는 데이터 흐름을 통해 이해할 수 있다.

입력: 시퀀스 차원을 따라 분할된 Q, K, V 텐서
      디바이스당 형태: (bs, L/N, h_c, h_s)

단계 1: Ulysses 그룹 내 AllToAll (행 방향)
        Scatter: 시퀀스 차원 (idx=1)
        Gather: 헤드 차원 (idx=2)
        출력 형태: (bs, L/N_2, h_c/N_2, h_s) -- 그러나 이제 헤드 부분집합에 대해 전체 L
        실제로는: reshape 후 (h_c/N_2, bs, L, d)

단계 2: Ring 그룹 내 부하 균형 Ring Attention (열 방향)
        링 토폴로지에서 K,V 블록의 P2P communication
        각 디바이스는 모든 K,V에 대해 로컬 Q 청크에 대한 attention 계산
        출력: 로컬 Q 청크에 대한 O 텐서

단계 3: Ulysses 그룹 내 AllToAll (역방향)
        Scatter: 헤드 차원 (idx=2)
        Gather: 시퀀스 차원 (idx=1)
        출력: 시퀀스 차원을 따라 분할된 O
        디바이스당 형태: (bs, L/N, h_c, h_s)

2D 프로세스 메시 구성은 Data Parallelism과 Tensor Parallelism 프로세스 그룹이 기존에 배열되는 방식과 유사하다. 총 $N$개의 디바이스에 Ulysses 정도 $N_2$와 Ring 정도 $N_1$(여기서 $N = N_1 \times N_2$)을 가진 경우, SP 프로세스 그룹은 $N_1 \times N_2$ 메시로 간주된다. SP-Ring은 각 열(Ulysses 그룹 간 동일한 위치)을 가로질러 작동하고, SP-Ulysses는 각 행(각 Ulysses 그룹 내) 간에 실행된다. 이 구조는 4D 또는 5D 하이브리드 parallelism 프레임워크에서 다른 parallelism 차원과의 자연스러운 통합을 가능하게 한다.

2.3 Key Algorithms & Mechanisms

USP-Attention 알고리즘. 핵심 알고리즘(논문의 Algorithm 1)은 놀라울 정도로 간단하다. 프로세스 그룹 ulysses_pg와 ring_pg, 그리고 형태 $(bs, L/N, h_c, h_s)$의 입력 텐서 세그먼트 $Q$, $K$, $V$가 주어졌을 때, forward pass는 다음과 같이 진행된다. 첫째, AllToAll4D가 Ulysses 프로세스 그룹 내에서 $Q$, $K$, $V$에 적용되며, scatter_idx=1(시퀀스 차원)과 gather_idx=2(헤드 차원)를 사용한다. 이는 각 텐서를 시퀀스 차원을 따라 분할된 상태에서 헤드 차원을 따라 분할된 상태로 변환하므로, 각 디바이스는 이제 전체 시퀀스 길이를 보유하지만 $h_c / N_2$개의 헤드만 가진다. 둘째, 부하 균형 Ring Attention이 Ring 프로세스 그룹 내에서 실행되어 분산 시퀀스에 걸쳐 로컬 헤드 부분집합에 대한 전체 attention을 계산한다. 셋째, 출력 $O$는 Ulysses 그룹 내에서 역방향 AllToAll4D를 거치며, scatter_idx=2와 gather_idx=1을 사용하여 시퀀스 분할 레이아웃을 복원한다. 역전파 중에는 scatter와 gather 인덱스가 교환된다.

AllToAll4D 연산은 Ulysses 컴포넌트의 핵심 통신 프리미티브다. 데이터를 합산하는 AllReduce나 데이터를 연결하는 AllGather와 달리, AllToAll은 완전한 개인화된 교환을 수행한다. 각 디바이스가 데이터의 다른 부분을 다른 각 디바이스에 전송한다. USP의 맥락에서 AllToAll은 시퀀스 차원 $L$(모든 시퀀스 청크 수집)을 병합하고 헤드 차원 $h_c$(헤드 부분집합 분배)를 분할하여, 텐서 형태를 $(bs, L/N_2, h_c, h_s)$에서 $(bs, L, h_c/N_2, h_s)$로 변환한다. AllToAll 연산당 통신량은 $\frac{N_2 - 1}{N_2} \cdot bs \cdot \frac{L}{N_2} \cdot h_c \cdot h_s$ 요소이다. 총 8개의 AllToAll 연산이 있으므로(forward에서 Q, K, V, O에 대해 4개, backward에서 4개), 총 Ulysses 통신 비용은 $\frac{8}{N_2} O(bs \cdot L \cdot d)$이다.

부하 균형 Ring Attention. Ring Attention 컴포넌트는 attention 연산을 블록 단위로 처리한다. 각 디바이스는 로컬 $Q$ 블록을 보유하고 링 토폴로지에서 이웃 디바이스로부터 $K$, $V$ 블록을 반복적으로 수신한다. 핵심 수학적 과제는 분산 블록 간에 올바른 softmax 정규화를 계산하는 것이다. FlashAttention의 온라인 softmax 기법을 따라, 각 디바이스는 실행 중인 log-sum-exp(LSE) 누적기 $m_i$와 실행 중인 출력 누적기 $O_i$를 유지한다. 디바이스 $j$로부터 새로운 $K_j, V_j$ 블록을 처리할 때, 로컬 attention 점수 $S_{ij} = Q_i K_j^T / \sqrt{d}$가 계산되고, 누적기는 다음과 같이 업데이트된다.

이는 블록 단위 분산에도 불구하고 수치적으로 안정적이고 수학적으로 정확한 attention 연산을 보장한다.

Causal masking에서의 부하 균형 스킴의 경우, 길이 $L$의 입력 시퀀스가 $2 \times N_1$개의 청크로 나뉜다. 링 랭크 $r$을 가진 디바이스는 위치 $r$과 $2N_1 - r - 1$의 청크를 받는다. 예를 들어, 4개의 Ring 디바이스와 16개의 토큰으로, 할당은 다음과 같다. GPU0은 토큰 {0,1,14,15}를 받고, GPU1은 {2,3,12,13}을, GPU2는 {4,5,10,11}을, GPU3은 {6,7,8,9}를 받는다. 각 디바이스의 청크 쌍은 시퀀스의 시작 부분(적은 causal attention 대상)에서 하나, 끝 부분(많은 대상)에서 하나가 선택되어 균형 잡힌 총 연산량을 얻는다. 링 랭크 $r$과 Ulysses 랭크 $u$를 가진 디바이스의 재정렬 공식은 로컬 시퀀스를 다음과 같이 추출한다.

이는 RoPE 위치 인코딩에 동일한 재정렬을 적용해야 하며, RoPE가 요소별 연산만 포함하므로 실현 가능하다.

통신 비용 분석. 본 논문은 parallelism 전략 간 상세한 비용 비교를 제공한다. 은닉 차원 $d$를 가진 표준 transformer 블록의 경우, 파라미터 통신(gradient 동기화)은 AllReduce(또는 ZeRO의 경우 동등한 AllGather + ReduceScatter)에 대해 $12 \cdot O(d^2)$ 요소를 포함한다. Activation 통신은 방법에 따라 다르다.

중요한 관찰은 SP-Ulysses 통신 비용이 parallelism 정도 증가에 따라 $1/N$으로 감소하는 반면, TP-sp 통신은 parallelism 정도와 관계없이 일정하게 유지된다는 것이다. 게다가, GQA는 SP 통신 비용을 감소시킨다. $K$와 $V$ 텐서가 더 적은 헤드를 가지기 때문이다. GQA 그룹 수 $G$에서, $K$와 $V$의 통신 비용은 $1/G$로 감소하지만, TP-sp 통신은 GQA의 영향을 받지 않는다.

메모리 분석은 미묘한 트레이드오프를 드러낸다. SP와 DP 모두 activation 메모리를 $A/N$으로 감소시킨다. 하지만 ZeRO 없는 SP는 전체 파라미터 및 gradient 복사본을 유지하며($P + G$), TP-sp는 파라미터를 $(P + G)/N$으로 분배한다. SP는 ZeRO-3와 결합될 때만 TP-sp와 메모리에서 동등해지며, 이는 파라미터 수집을 위한 추가 AllGather 통신을 추가한다. 이는 SP 단독으로 TP-sp의 최대 시퀀스 길이와 일치할 수 없는 이유를 설명한다. 파라미터 메모리 오버헤드가 activation에 사용 가능한 메모리를 제한하기 때문이다.

2.4 Implementation Details

구현은 Megatron-LM(커밋 2196398, 2024년 4월) 위에 구축되었으며, NeMo 24.03 Docker 이미지를 사용한다. SP-Ring은 Megatron-LM의 네이티브 Context Parallel 구현을 사용하고, SP-Ulysses는 저자의 커스텀 리포지토리에서 구현되었다. 주요 구현 세부사항은 다음과 같다.

혼합 정밀도 학습. 모든 실험은 bf16/fp16 형식을 사용한다. 메모리 모델은 파라미터에 대해 $P$ 바이트, gradient에 대해 $G$ 바이트, optimizer 상태에 대해 $6P$ 바이트(Adam의 fp32 파라미터 복사본 + 모멘텀 + 분산)를 고려한다.

ZeRO 통합. DP와 SP 프로세스 그룹 모두에 대해 기본적으로 ZeRO-1이 사용된다. ZeRO 프로세스 그룹은 크기 $N_{sp} \times N_{dp}$의 통합 SP+DP 그룹에 걸쳐 있으며, 이는 메모리 효율성에 중요하다.

테스트된 하드웨어 구성. 실험은 세 가지 하드웨어 설정을 포괄한다. (1) 8xL20 PCIe 클러스터(이종 대역폭), (2) 8xA100-SXM4 NVLink 단일 노드(균일한 고대역폭), (3) 1.6 Tbps RDMA로 연결된 두 개의 8xA800 NVLink 노드(계층적 대역폭). MFU 계산은 causal masking을 고려하여(하삼각 attention 행렬의 유효 FLOP만 계산), Megatron-LM의 기본 계산보다 낮은 MFU를 보고한다.

하이퍼파라미터 선택. Ulysses 정도와 Ring 정도가 주요 조정 손잡이이다. 최적 구성은 하드웨어 토폴로지에 따라 다르다. NVLink 전용 설정에서는 Ulysses 정도를 최대화하는 것($h_c$까지)이 최적이고, PCIe 또는 다중 노드 설정에서는 균형 잡힌 혼합(예: 8xL20에서 Ulysses=4, Ring=2)이 대역폭 계층을 활용한다. 곱 $N_1 \times N_2$는 총 SP 정도와 같아야 하고, $N_2$는 $h_c$(또는 GQA/MQA의 KV 헤드 수)를 초과할 수 없다.

복잡도. USP attention의 연산 복잡도는 표준 attention과 동일하다. 디바이스당 $O(L^2 \cdot d)$가 SP로 $O(L^2 / N \cdot d)$가 된다. 통신 복잡도는 연산당 $O(bs \cdot L \cdot d / N_2)$의 AllToAll 연산이 지배하며, Ring 단계당 $O(bs \cdot L \cdot d / N)$의 P2P 통신이 추가된다(연산과 중첩됨). 부하 균형 토큰 재정렬은 정수 인덱스에서 $O(bs \cdot L)$이며, 무시할 수 있다.

3. Results

실험 평가는 5가지 차원으로 구조화되어 있다. 독립형 USP 성능, DP와의 비교, 하이브리드 SP+TP 구성, 상한 시퀀스 길이 탐색, 수렴 검증이다.

이종 네트워크에서의 독립형 USP 성능(표 3). LLAMA3-8B 파라미터로 8xL20 PCIe 클러스터에서, 32K 및 128K 시퀀스에 대한 최적 구성은 Ulysses 정도 4와 Ring 정도 2였으며, 각각 32.818과 4.235 iters/sec를 달성했다. 이는 토폴로지 인식 이점을 확인한다. 순수 SP-Ulysses(Ulysses=8, Ring=1)는 이 길이에서 각각 15.229와 2.563 iters/sec만 달성했고, 균형 잡힌 구성은 AllToAll을 고대역폭 PCIe 스위치 연결로 라우팅하고 P2P를 저대역폭 교차 스위치 링크로 라우팅한다. 부하 균형 Ring 변형(lb-ring)은 기본 Ring을 일관되게 능가했으며, 그 이점은 8K에서는 무시할 수 있었지만 128K에서는 상당했다(Ulysses=2, Ring=4의 경우 5.476 대 3.399). 이는 긴 시퀀스에서 causal mask 부하 균형의 중요성이 증가함을 보여준다.

동종 네트워크에서의 독립형 USP(표 4). 8xA100-SXM4 NVLink(균일한 고대역폭)에서 최적 구성은 순수 SP-Ulysses(Ulysses=8)였으며, 32K에서 136.375 iters/sec를 달성했다. NVLink가 균일하게 높은 대역폭을 제공하므로, AllToAll 통신(SP-Ulysses가 사용)이 최대한 효율적이기 때문에 이해가 된다. SP-Ring의 블록 세분화로 인한 연산 오버헤드가 이 설정에서 통신 은닉 이점보다 크다. 표 3과 4 간의 대조는 USP의 가치를 강력하게 보여준다. 동일한 알고리즘이 단일 파라미터 쌍을 조정하여 근본적으로 다른 하드웨어 토폴로지에 적응한다.

SP 대 DP(그림 5). LLAMA2-7B와 전역 배치 크기 8로 단일 8xA800 노드에서, DP가 모든 시퀀스 길이에 걸쳐 SP-Unified를 일관되게 능가한다(4K에서 161 대 149 TFLOPS, 512에서 55 대 53 TFLOPS). 이 5-10% 격차는 예상된다. SP가 DP가 완전히 피하는 attention 모듈에 대한 추가 AllToAll 및 P2P 통신 오버헤드를 도입하기 때문이다. 이는 논문의 팁 2를 검증한다. 배치 크기가 충분할 때 DP가 선호되어야 하며, SP는 배치 크기를 더 이상 분할할 수 없을 때만 사용되어야 한다.

단일 노드에서의 하이브리드 SP+TP(표 5). LLAMA2-7B로 8xA800에서 달성 가능한 최대 시퀀스 길이는 64K였다. 이 길이에서 최고 성능(154.49 TFLOPS, 0.50 MFU)은 TP=4, Ulysses=2, Ring=1로 달성되었으며, TP 전용(141.85 TFLOPS, 0.45 MFU)을 10% 능가했다. 중요하게도, SP 전용(TP 없음)은 64K에서 OOM을 발생시켜 TP-sp가 SP 단독보다 메모리 효율적임을 확인했다. 더 큰 배치 크기(16)로 더 짧은 시퀀스(30K)에서, 순수 SP-Ulysses(Ulysses=8, TP 없음)가 163.42 TFLOPS(0.52 MFU)로 최고 성능을 달성하여 TP 전용을 26% 능가했다. 이는 중요한 발견을 드러낸다. 메모리가 병목이 아닐 때, SP-Ulysses의 낮은 통신 오버헤드가 실제로 TP-sp보다 우수하게 만든다.

다중 노드 학습(표 6). LLAMA3-8B로 두 개의 8xA800 노드에서, 64K 및 80K 시퀀스에 대한 최적 구성은 모두 Ulysses=4, Ring=4(SP 전용, TP 없음)였으며, 각각 137.48과 148.90 TFLOPS를 달성했다. USP는 64K에서 13%, 80K에서 12% 순수 SP-Ring을 능가하여 하이브리드 이점을 직접 보여준다. 120K 시퀀스에서 SP 전용은 OOM을 발생시켜 최적 성능(152.51 TFLOPS, 0.49 MFU)을 위해 TP=4, Ulysses=2, Ring=2가 필요했다. 이 패턴(메모리가 허용할 때 SP 전용이 더 빠르지만, 최대 시퀀스 길이를 위해서는 TP가 필요함)은 핵심 실용적 발견이다.

상한 탐색(표 7). 두 노드에서 달성된 최대 시퀀스 길이는 TP=8, Ring=2로 208K 토큰이었으며, 147.26 TFLOPS(0.47 MFU)를 달성했다. 160K에서 최고 MFU는 TP=4/Ulysses=2/Ring=2 또는 TP=4/Ring=4로 0.51이었다. 이러한 결과는 가능한 가장 긴 시퀀스를 달성하려면 모든 헤드 차원 parallelism을 TP-sp에 할당해야 함(메모리 효율성을 위해)을 보여주며, Ring이 추가 시퀀스 분배를 제공한다.

수렴 검증(그림 6). 2K 시퀀스 길이로 10K 반복에 걸쳐 LLAMA2-7B에 대한 DP 및 USP(Ulysses=2, Ring=2)의 손실 곡선이 완전히 겹쳐 수학적 동등성을 확인한다. 이는 부하 균형을 위한 토큰 재정렬과 RoPE 수정이 학습 정확성을 보존함을 검증한다.

4. Critical Assessment

Strengths

1. 상호보완적 방법의 우아한 통합. 2D 메시 분해는 개념적으로 간단하면서도 강력하며, 단일 파라미터 쌍으로 순수 Ulysses와 순수 Ring 사이의 연속적인 스펙트럼을 가능하게 한다. 이는 실무자가 두 개의 불완전한 솔루션 중 하나를 선택하도록 강요하지 않는다.
2. 포괄적인 체계적 분석. SP-Ulysses, SP-Ring, DP, ZeRO-1/2/3, TP, TP-sp를 통신 비용, 메모리 비용, 분할 차원 전반에 걸쳐 다루는 비교 표(표 2)는 중요한 참고 기여이다. 7가지 실용적 팁이 이 분석을 실행 가능한 지침으로 정리한다.
3. 하드웨어 토폴로지 인식. 세 가지 뚜렷한 하드웨어 구성(PCIe, NVLink 단일 노드, 다중 노드 NVLink+RDMA)에 걸친 실험 검증은 USP가 이종 대역폭 계층에 적응함을 설득력 있게 보여준다. 이는 현대 GPU 클러스터의 현실이다.
4. 부하 균형 causal attention. Causal masking에서 균형 잡힌 연산을 위한 토큰 재정렬 전략은 간단하고 효과적이며 무시할 수 있는 오버헤드를 추가한다. 통합 프레임워크에 통합된 것은 실제 학습 워크로드에 실용적 필요성이다.
5. 오픈소스 구현. 코드가 공개적으로 제공되고 Megatron-LM과 통합되어 채택 장벽을 낮추고 결과 재현을 가능하게 한다.

Limitations

1. 제한된 실험 규모. 모든 실험은 최대 16개 GPU(두 개의 8-GPU 노드)를 사용한다. 대규모 이점(10K+ GPU)에 대한 논문의 주장은 검증되지 않았다. 규모에서 USP와 통신 패브릭 간 상호작용(혼잡, 대역폭 공유)은 탐색되지 않았다.
2. 동시 작업과의 비교 없음. 본 논문은 Striped Attention, LightSeq, 또는 다른 신흥 SP 방법과 벤치마크하지 않아 Ulysses/Ring 베이스라인 너머의 상대적 장점을 평가하기 어렵다.
3. 메모리 분석 정밀도 부족. 메모리 비용 비교(표 2)는 정확한 바이트 수가 아닌 점근적 표기법($O$)을 사용하며, TP activation 메모리에 대한 $\alpha$ 팩터가 불명확하게 남아 있다. 실무자는 용량 계획을 위해 정확한 숫자가 필요하다.
4. GQA/MQA 영향 미탐색. 논문이 GQA가 SP 통신 비용을 감소시킨다고 언급하지만, 실험 평가는 이 이점을 정량화하기 위해 GQA 그룹 수를 체계적으로 변화시키지 않는다. 모든 LLAMA3-8B 실험은 다른 구성을 탐색하지 않고 $h_c = 8$개의 KV 헤드를 사용한다.
5. 통신 병목의 프로파일링 없음. 본 논문은 종단 간 처리량을 보고하지만 AllToAll, P2P, 연산 단계에 소요된 시간을 분해하지 않는다. 이는 실무자가 통신이 언제 병목이 되는지 이해하고 추가 최적화 방법을 파악하는 데 도움이 될 것이다.
6. 최소한의 수렴 검증. 단 하나의 수렴 실험(LLAMA2-7B, 2K 시퀀스, 10K 반복)만 표시된다. 더 긴 시퀀스와 더 큰 모델은 마이크로 배치 내 attention 패턴에 영향을 미치는 부하 균형 토큰 재정렬로 인해 다른 수렴 동작을 보일 수 있다.

Future Directions

1. 대규모 검증. 프로덕션 수준 모델로 1000+ GPU 규모에서 USP를 시연하면 이론적 장점을 검증하고 통신 스케줄링에서 새로운 과제를 드러낼 것이다.
2. MoE 아키텍처와의 통합. 논문이 언급하듯, attention 모듈의 SP는 FFN 모듈과 분리되어 MoE 호환성을 실현 가능하게 하지만, attention SP 그룹과 전문가 parallelism 그룹 간 AllToAll 전환에 대한 신중한 설계가 필요하다.
3. 자동화된 구성 튜닝. 큰 구성 공간(Ulysses 정도, Ring 정도, TP 정도, ZeRO 단계)을 고려할 때, 하드웨어 토폴로지와 모델 아키텍처를 프로파일링하여 최적 설정을 권장하는 자동 튜닝 메커니즘을 개발하면 사용성이 크게 향상될 것이다.
4. Megatron-LM에서 SP+ZeRO-3. 본 논문은 이를 SP가 TP-sp의 메모리 효율성과 일치할 수 있게 하는 누락된 기능으로 식별한다. 이 조합을 구현하고 벤치마크하는 것이 자연스러운 다음 단계이다.
5. 추론 워크로드로의 확장. 본 논문이 학습에 초점을 맞추지만, 긴 문맥 추론(특히 prefill 단계)도 유사한 과제에 직면한다. 다른 배치 크기와 시퀀스 길이 트레이드오프가 있는 추론 시나리오에 USP를 적응시키는 것이 중요한 방향이다.