🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

Ulysses Sequence Parallelism은 긴 시퀀스 학습에서 시퀀스와 어텐션 헤드를 함께 나누고 all-to-all 통신으로 재배치해, 단일 GPU 메모리 한계를 넘어 수십만~백만 토큰 문맥 학습을 가능하게 하는 방식이다.

📌 핵심 요약

• 긴 문맥 학습은 문서 분석, 코드 이해, 복잡한 추론, RAG처럼 많은 토큰을 한 번에 처리해야 하는 작업에서 중요해졌지만, 트랜스포머 어텐션의 계산량이 시퀀스 길이에 대해 제곱으로 증가해 단일 GPU 메모리 한계에 부딪힌다.
• FlashAttention은 어텐션 행렬을 직접 저장하지 않아 메모리를 선형 수준으로 줄이지만, 계산량 자체는 여전히 제곱으로 남기 때문에 32k 이상 같은 긴 시퀀스 학습에서는 추가적인 병렬화가 필요하다.
• Ulysses Sequence Parallelism은 입력 시퀀스를 GPU 여러 대에 나누고, QKV 투영 뒤 all-to-all 통신으로 각 GPU가 전체 시퀀스 위치의 일부 어텐션 헤드를 담당하게 만들어 로컬 어텐션을 수행한다.
• Accelerate, Transformers Trainer, TRL SFTTrainer에는 Ulysses 설정이 통합되어 있으며, 모델 준비, 데이터로더 래핑, 시퀀스 샤딩, shift labels 처리, 손실 집계 같은 세부 작업을 자동화한다.
• Ulysses는 Ring Attention과 비교해 GPU당 통신량이 더 작고 all-to-all 집합 통신을 활용할 수 있다는 장점이 있으며, 긴 문맥 학습에서는 position_ids를 사용해 거대한 4D attention mask를 피한다.

🧩 주요 포인트

1. 긴 문맥 학습은 문서 분석, 코드 이해, 복잡한 추론, RAG처럼 많은 토큰을 한 번에 처리해야 하는 작업에서 중요해졌지만, 트랜스포머 어텐션의 계산량이 시퀀스 길이에 대해 제곱으로 증가해 단일 GPU 메모리 한계에 부딪힌다.
2. FlashAttention은 어텐션 행렬을 직접 저장하지 않아 메모리를 선형 수준으로 줄이지만, 계산량 자체는 여전히 제곱으로 남기 때문에 32k 이상 같은 긴 시퀀스 학습에서는 추가적인 병렬화가 필요하다.
3. Ulysses Sequence Parallelism은 입력 시퀀스를 GPU 여러 대에 나누고, QKV 투영 뒤 all-to-all 통신으로 각 GPU가 전체 시퀀스 위치의 일부 어텐션 헤드를 담당하게 만들어 로컬 어텐션을 수행한다.
4. Accelerate, Transformers Trainer, TRL SFTTrainer에는 Ulysses 설정이 통합되어 있으며, 모델 준비, 데이터로더 래핑, 시퀀스 샤딩, shift labels 처리, 손실 집계 같은 세부 작업을 자동화한다.
5. Ulysses는 Ring Attention과 비교해 GPU당 통신량이 더 작고 all-to-all 집합 통신을 활용할 수 있다는 장점이 있으며, 긴 문맥 학습에서는 position_ids를 사용해 거대한 4D attention mask를 피한다.

🧠 상세 정리

1. 긴 문맥 학습이 필요한 배경

원문은 대규모 언어모델이 문서 분석, 코드 이해, 복잡한 추론, 검색 증강 생성 같은 작업에 쓰이면서 수십만 개 또는 수백만 개 토큰을 처리해야 하는 필요가 커졌다고 설명한다. 평균적인 책 한 권이 약 25만 토큰이라는 예시는, 여러 문서나 책 길이 입력을 학습에 사용하려면 단일 GPU에 담기 어려운 길이를 다뤄야 한다는 점을 보여준다. 특히 긴 문맥은 단순히 입력 길이를 늘리는 문제가 아니라, 모델이 전체 책, 법률 문서, 연구 논문, 대규모 코드베이스, 여러 검색 결과를 한 흐름 안에서 처리해야 하는 요구와 연결된다. 따라서 긴 시퀀스 학습은 선택적 최적화가 아니라 더 복잡한 실제 사용 사례를 지원하기 위한 핵심 조건으로 제시된다.

2. 표준 어텐션의 메모리와 계산 한계

트랜스포머의 어텐션은 시퀀스 길이 n에 대해 계산량과 어텐션 점수 행렬 저장 비용이 모두 O(n²)로 증가한다. FlashAttention 같은 최적화 구현은 타일링을 이용해 전체 어텐션 행렬을 물질화하지 않으므로 메모리 사용량을 O(n) 수준으로 줄일 수 있지만, 원문은 계산량 자체가 여전히 O(n²)로 남는다고 강조한다. 그 결과 32k 토큰 이상처럼 매우 긴 시퀀스를 학습할 때는 메모리 최적화만으로 충분하지 않고, 어텐션 블록 내부에서 시퀀스 처리를 여러 장치로 나누는 방식이 필요해진다. 전통적인 데이터 병렬화는 각 GPU가 여전히 전체 시퀀스를 처리해야 하므로 이 문제를 해결하지 못한다.

3. Ulysses Sequence Parallelism의 핵심 아이디어

Ulysses Sequence Parallelism은 DeepSpeed Ulysses 논문에서 소개된 방식으로, 시퀀스 차원을 나누는 것에 더해 어텐션 헤드도 GPU들 사이에 분할한다. 먼저 입력 시퀀스를 P개의 GPU에 나누어 각 GPU가 전체 토큰 중 일부 구간을 보유하게 하고, 각 장치가 자신의 로컬 시퀀스 청크에 대해 query, key, value 투영을 계산한다. 이후 all-to-all 통신을 통해 데이터를 재배치하여 각 GPU가 모든 시퀀스 위치를 갖되 어텐션 헤드의 일부만 담당하도록 만든다. 핵심 통찰은 어텐션 헤드들이 서로 독립적으로 계산될 수 있기 때문에, 시퀀스 지역성을 헤드 지역성으로 바꾸면 긴 시퀀스 어텐션을 여러 GPU에 효율적으로 분산할 수 있다는 점이다.

4. Ulysses의 처리 단계와 통신 구조

원문이 설명하는 Ulysses의 흐름은 시퀀스 샤딩, QKV 투영, 첫 번째 all-to-all 통신, 로컬 어텐션, 두 번째 all-to-all 통신, 출력 투영 순서로 진행된다. 첫 번째 all-to-all 이후 각 GPU는 전체 시퀀스 위치를 보지만 자신에게 할당된 어텐션 헤드만 계산하며, 이때 FlashAttention 또는 SDPA 같은 표준 어텐션 구현을 사용할 수 있다. 로컬 어텐션이 끝나면 두 번째 all-to-all 통신으로 데이터를 다시 시퀀스 샤딩 형태로 되돌리고, 각 GPU는 자신의 로컬 시퀀스 청크에 대한 출력 투영을 수행한다. 이 구조는 긴 시퀀스 전체를 한 GPU가 모두 감당하지 않도록 하면서도 어텐션 계산의 의미를 유지하는 병렬화 방식으로 제시된다.

5. Ring Attention과의 통신 복잡도 비교

원문은 Ulysses가 각 어텐션 레이어마다 두 번의 all-to-all 연산을 필요로 하며, GPU당 통신량은 O(n·d/P)라고 설명한다. 여기서 n은 시퀀스 길이, d는 hidden dimension, P는 병렬화 정도를 의미한다. 반면 Ring Attention은 GPU당 O(n·d)를 통신하며, P-1번의 순차적인 point-to-point 전송을 링 구조로 수행하기 때문에 Ulysses보다 P배 많은 통신량을 갖는 것으로 비교된다. 또한 Ulysses는 all-to-all 집합 연산이 한 번의 collective step에서 전체 bisection bandwidth를 활용할 수 있어, P-1개의 홉을 직렬화하는 Ring Attention보다 지연시간 측면에서도 유리하다고 설명한다.

6. Accelerate 통합과 사용자 설정

Accelerate는 ParallelismConfig와 DeepSpeedSequenceParallelConfig를 통해 Ulysses sequence parallelism의 기반을 제공한다. 설정에서는 sp_backend를 deepspeed로 지정하고, sp_size로 sequence parallelism에 사용할 GPU 수를 정하며, sp_seq_length_is_variable을 True로 두면 배치마다 시퀀스 길이가 달라지는 상황을 다룰 수 있다. sp_attn_implementation에는 flash_attention_2, flash_attention_3, sdpa 같은 어텐션 구현을 지정할 수 있다고 원문은 설명한다. accelerator.prepare()를 호출하면 모델이 DeepSpeed의 UlyssesSPAttentionHF에 등록되고, 데이터로더는 UlyssesSPDataLoaderAdapter로 감싸지며, 올바른 손실 계산을 위해 shift_labels도 자동으로 주입된다.

7. 손실 집계와 position_ids 사용

Ulysses에서는 각 GPU가 시퀀스의 서로 다른 부분에 대한 손실을 계산하므로, 전체 손실을 단순 평균이 아니라 유효 토큰 수로 가중해 집계해야 한다. 원문은 일부 rank가 padding만 포함하거나 prompt 토큰이 마스킹된 경우처럼 유효 토큰 수가 불균형할 수 있으므로, 각 rank의 loss와 valid token count를 모아 가중 aggregation을 수행해야 올바른 gradient를 얻는다고 설명한다. Transformers Trainer나 TRL SFTTrainer를 사용하면 이 처리는 자동으로 수행되며, 사용자 정의 Accelerate 학습 루프를 작성할 때만 별도 코드가 필요하다. 또한 긴 시퀀스에서는 4D attention mask가 어텐션 점수 행렬만큼 부담스러울 수 있으므로, Ulysses와 Ring Attention은 학습 중 causal masking에 attention_mask 대신 position_ids를 사용한다고 설명한다.

8. Transformers Trainer와 TRL SFTTrainer 통합

Transformers Trainer는 TrainingArguments.parallelism_config를 통해 Ulysses를 통합하며, 데이터로더 래핑, 시퀀스 샤딩, 손실 집계 같은 sequence parallelism 관련 세부 사항을 자동 처리한다. Trainer의 compute_loss는 SP 모드를 감지해 DeepSpeed SP 전용 손실 계산 경로로 보내고, SP rank별 손실과 유효 토큰 수를 모아 가중 손실을 계산한다. TRL의 SFTTrainer는 Transformers Trainer 위에 구축되어 있으며, 긴 시퀀스 supervised fine-tuning을 위해 max_length, pad_to_multiple_of, packing 같은 설정을 함께 사용한다. 원문에서는 pad_to_multiple_of가 sp_size와 같아야 시퀀스 divisibility를 보장한다고 설명하며, Ulysses가 활성화된 경우 SFTTrainer가 pre-shifted labels 처리를 자동으로 다룬다고 제시한다.

🧾 핵심 주장 / 시사점

• Ulysses의 장점은 어텐션 계산 자체를 바꾸기보다, 독립적인 어텐션 헤드라는 구조적 특성을 이용해 긴 시퀀스 학습의 병목을 여러 GPU로 재배치한다는 데 있다.
• 긴 문맥 학습에서는 메모리 절감 기법만으로 충분하지 않으며, 데이터 병렬화가 해결하지 못하는 어텐션 내부의 시퀀스 길이 문제를 직접 나누는 병렬화가 필요하다.
• Accelerate, Transformers Trainer, TRL SFTTrainer에 통합된 덕분에 사용자는 복잡한 손실 집계와 데이터로더 샤딩을 직접 구현하지 않고도 Ulysses 기반 긴 문맥 학습을 구성할 수 있다.

✅ 액션 아이템

• 긴 문맥 학습에서 32k 이상의 시퀀스는 계산량 제곱 증가가 크므로 Ulysses 적용 구간을 문맥 길이 기준으로 선별해 우선 반영한다.
• QKV 투영 뒤 all-to-all로 시퀀스와 헤드를 분할했을 때 GPU당 통신량·메모리 사용량·학습 속도를 실제 측정해 설정값을 정량 비교한다.
• Accelerate, Transformers, TRL 경로에서 모델 준비·데이터로더 래핑·시퀀스 샤딩·shift labels·손실 집계 자동화가 정상 동작하는지 학습 설정을 점검한다.

❓ 열린 질문

• 긴 문맥 100k에서 1M 토큰 구간에서 all-to-all 기반 Ulysses가 Ring Attention 대비 통신 병목을 줄이는 임계 토큰 수와 GPU 구성은 어떤가?
• position_ids로 거대한 4D attention mask를 피할 때 메모리 절감과 계산 비용 완화 효과는 기존 방식 대비 얼마나 달라지는가?
• Ulysses가 통합된 Accelerate/Transformers/TRL 조합에서 시퀀스 샤딩과 손실 집계 동기화 신뢰성은 어떤 지표로 검증해 판단할 것인가?