Context Parallelism v.s. Sequence Parallelism (Megatron-LM 기준)

참고 코드: https://github.com/nvidia/megatron-lm

GitHub - NVIDIA/Megatron-LM: Ongoing research training transformer models at scale

참고 글: Introducing Context Parallelism · Better Tomorrow with Computer Science - https://insujang.github.io/2024-09-20/introducing-context-parallelism/

Sequence Parallelism (SP)

기준 논문 [MLSys'23]: REDUCING ACTIVATION RECOMPUTATION IN LARGE TRANSFORMER MODELS

SP(Sequence Parallelism) 는 Megatron-LM 기준으로 TP에 붙어서 동작하는 보조 기법에 가깝다. TP에서는 linear layer의 output이 hidden dimension 기준으로 all-reduce 이후 TP rank마다 동일한 activation이 생긴다. 그러면 LayerNorm, Dropout처럼 TP가 직접 적용되지 않는 구간 (i.e., weight가 hidden dim으로 나뉘지 않는 구간)에서는 rank마다 연산이 불필요하게 중복된다. 중복 연산으로 인한 연산량 자체는 뭐 크지 않더라도, activation 저장을 할 때 중복 저장하는 memory overhead를 줄이고 싶은 것이다. 

따라서 SP는 activation을 sequence dimension으로 다시 나누어, TP rank들이 동일한 activation을 중복 저장하지 않도록 한다. 즉, SP는 주로 TP 내부에서 이미 존재하던 통신 패턴을 all-reduce 대신 reduce-scatter + all-gather 형태로 재구성하는 방식이다. all-reduce = reduce-scatter + all-gather로 볼 수 있기 때문에, 통신량 자체는 크게 늘지 않으면서 activation memory를 줄일 수 있다. 이 때문에 SP는 보통 추가 computation/communication overhead가 거의 없고, 시간 overhead도 작다고 볼 수 있다.

 

Context Parallelism (CP)

Megatron-LM / Megatron-Core 기준으로, Context Parallelism(CP)은 긴 sequence를 여러 GPU에 나누어 저장하고 계산하기 위한 parallelism이다. 즉, batch나 hidden dimension이 아니라 sequence length / context dimension을 shard한다.

`context_parallel_size` > 1일 때 켜지고, 기본값은 1이라서 CP가 꺼진 상태다. Megatron-Core 문서 기준으로 CP는 TP, PP, DP와 함께 조합 가능하며, 전체 GPU 수(`world_size`)는 기본적으로 TP × CP × PP × DP 구조로 잡힌다. `DP x CP` 수만큼 model weight가 replicate되고, `TP x PP` 수만큼 model weight shard된다. 각 `DP` rank의 model replica에 대해서 서로 다른 data batch가 들어오고, 그 안에서 `CP` rank만큼 이 batch가 sequence length가 잘라진다. 

context_parallel_size > 1 일 때 CP가 켜지고, 기본값은 1이므로 기본적으로는 CP가 꺼져 있다. CP는 TP, PP, DP와 함께 조합 가능하며, 전체 GPU 수는 보통 다음과 같이 구성된다.

$$ \text{world\_size} = TP \times CP \times PP \times DP $$

CP는 weight를 shard하지 않는다. Sequence activation만 나누기 때문에, 같은 CP group 안의 GPU들은 model weight를 복제해서 가진다. Megatron-LM 코드 주석에서도 CP는 sequence length를 partition하므로 weight에는 영향을 주지 않고, CP group 내에서는 weight가 duplicated된다고 설명한다.

4D parallelism (출처: Scaling Llama 3 Training with Efficient Parallelism Strategies - https://dl.acm.org/doi/epdf/10.1145/3695053.3731410)

 

따라서 직관적으로는 다음과 같이 볼 수 있다.

• TP × PP 방향으로는 model weight가 shard된다.
• DP × CP 방향으로는 model weight가 replicate된다.
• 각 DP rank에는 서로 다른 data batch가 들어간다.
• 그 batch 안에서 CP rank들이 sequence length를 나누어 가진다.

핵심 목적은 long-context training에서 activation memory를 줄이는 것이다. Sequence length가 길어질수록 transformer activation memory가 커지기 때문에 (기본 attention일 경우 $O(L^2)$) OOM이 발생하기 쉽다. 물론 activation recomputation을 쓰면 activation을 저장하지 않고 backward 때 다시 계산해주기 때문에 메모리 용량이 줄지만, 그러면 사실상 forward pass를 한번씩 더 돌아야 하는거라 forward-backward 연산 시간이 약 1.33배 (forward:backward time=1:2) 증가한다. CP는 각 GPU가 전체 sequence가 아니라 자기에게 할당된 sequence chunk만 저장/처리하게 해서 per-GPU activation memory를 CP size에 비례해 줄인다.

 

구현 특이점: Symmetric chunk assignment

좌: 기본 naive partitioning / 우: megatron의 head-tail partitioning. (출처: USP - https://arxiv.org/abs/2405.07719)

 

Megatron-LM의 CP는 sequence를 단순히 앞에서부터 `cp_size`개로 균등하게 자르지 않는다. 특히 causal attention에서는 attention의 유효 영역이 lower triangular 형태이므로, causal mask를 실제로 활용해 불필요한 upper-triangle 계산을 생략하는 attention kernel (e.g., FlashAttention)에서는 뒤쪽 query chunk일수록 attend해야 하는 key/value 범위가 길어지고 연산량이 커진다. 즉, 같은 길이의 chunk라도 앞쪽 chunk는 연산량이 작고, 뒤쪽 chunk는 연산량이 크다. 따라서 sequence를 contiguous하게 나누면 뒤쪽 chunk를 맡은 rank에 load가 몰리는 imbalance가 생긴다.

이를 완화하기 위해 Megatron-LM은 sequence를 먼저 `2 × cp_size`개 chunk로 나눈 뒤, 앞쪽 chunk와 뒤쪽 chunk를 하나의 CP rank에 pair로 배정한다. 예를 들어 cp_size = 4라면 sequence를 8개 chunk로 나누고, 각 rank는 다음과 같이 두 chunk를 가진다.

rank 0: chunk 0 + chunk 7
rank 1: chunk 1 + chunk 6
rank 2: chunk 2 + chunk 5
rank 3: chunk 3 + chunk 4

 

이 구조 때문에 실제 indexing도 일반적인 slicing과 다르다. 각 CP rank의 local sequence는 원래 sequence에서 연속된 한 구간이 아니므로, attention input, position id, attention mask, packed sequence metadata가 모두 이 앞/뒤 pairing 구조를 따라야 한다. 특히 RoPE는 token의 position 정보를 사용하므로, local rank 안에서 position을 단순히 0, 1, 2, ...로 다시 매기면 안 된다. 예를 들어 rank 0이 chunk 0 + chunk 7을 들고 있다면, RoPE position도 chunk 0의 global positions와 chunk 7의 global positions를 함께 가져와야 한다.

또한 이 방식에서는 sequence를 `cp_size`개 chunk가 아닌 `2 × cp_size`개 chunk로 균등하게 나눌 수 있어야 하므로, sequence length가 `2 × cp_size`의 배수가 되도록 padding을 맞춰야 한다. 그렇지 않으면 rank별 chunk 크기나 앞/뒤 pairing index가 맞지 않아 CP layout이 깨질 수 있다.

CP vs SP

SP는 TP에 합쳐져 있다고 보면 되고, RS/AG 와 AG/RS 사이 Dropout, LN 부분이 SP가 활성화되는 구간이라고 보면 된다. 저 Ring Attention 부분이 CP로 인한 통신 overhead. (출처: https://docs.nvidia.com/megatron-core/developer-guide/0.16.0/user-guide/features/context_parallel.html)

 

공통점

• 둘 다 sequence length를 sharding한다. 
• activation memory가 병목이 되는 long-context training 상황에서 sequence 축으로 activation을 shard해서 per-GPU memory 사용량을 줄이는 것을 목표로 한다. 
• 둘 다 보통 TP/DP/PP와 함께 조합되는 보조 parallelism으로 쓰인다. 특히 Megatron-LM 계열에서는 TP, PP, DP 위에 SP 또는 CP를 추가해서 메모리 병목을 완화한다.

차이점

적용 범위:

• SP: 주로 LayerNorm, Dropout 등 TP가 적용되지 않는 일부 activation을 sequence dimension으로 shard.
• CP: input batch와 대부분의 activation 전체를 sequence dimension으로 partition.

Communication 방식: 

• SP: TP와 SP 사이 boundary에서 reduce-scatter (hidden dim 축 shard를 reduce - sequence dim 축 shard로 scatter) / all-gather. (근데 사실상 identity/all-reduce -> reduce-scatter/all-gather 라서 통신량 자체는 같다.)
• CP: attention 연산 때 K, V head를 ring-style P2P 방식 (또는 all-gather도 가능)으로 통신 & 연산. 통신량이 $O(LH_{kv})$ ($H_{kv}$: kv chunk의 head 수)만큼 증가. 통신 지연 시간을 고려하면 순수 

Memory 이점:

• SP: TP에서 중복되는 activation memory 감소. LayerNorm 전에 TP에서 all-reduce를 하고 나면 LayerNorm에서는 모든 TP rank들이 동일한 activation 값을 가지게 되는데 굳이 이렇게 중복 값을 가져서 쓸데없는 메모리 소모를 해야 할까? -> 이 때 바로 sequence 축으로 sharding을 해줌으로써 per-GPU activation memory를 늘 $1/TP$만큼으로 유지.
• CP: sequence 축이 있는 모든 input, activation에 대해 $1/CP$만큼으로 감소시킴.