SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile

작성일 | In , , |

논문 링크

Q : 이 논문을 아주 자세하게 읽고 논문의 강점과 독창적인 지점을 설명해주고 핵심 알고리즘을 예시 입력을 들어서 전체적인 과정을 설명해줘 추가적으로 논문의 한계점에 대해서도 알려줘

A :

\

논문의 요약 및 강점

논문의 제목:
SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile

핵심 요약:
SimpleFSDP는 PyTorch-native compiler 기반의 Fully Sharded Data Parallel (FSDP) 프레임워크로, 대규모 모델 분산 학습에서 메모리 사용량을 최대 28.54% 줄이고 학습 처리량을 최대 68.67% 개선했습니다. 이 방법은 PyTorch의 기존 프리미티브를 활용하여 통합 구현되었으며, 새롭게 추가된 IR 노드 버킷화 및 재정렬을 통해 계산과 통신의 오버랩을 최적화합니다.

강점 및 독창적인 기여:

  1. 간결성: 기존 코드베이스를 크게 변경하지 않고 사용 가능하며, PyTorch의 torch.compile로 통합된 풀 그래프를 추적합니다.
  2. 구성 가능성: 텐서 병렬(Tensor Parallel), 파이프라인 병렬(Pipeline Parallel), 혼합 정밀도 학습(Mixed Precision Training)과 같은 최신 기술들과 원활히 통합 가능.
  3. 성능: IR 노드 버킷화와 재정렬 최적화를 통해 높은 학습 효율을 달성.
  4. 디버그 가능성: PyTorch 이거 모드(Eager Mode)에서 디버깅 기능을 유지하면서도 성능 향상.

핵심 알고리즘 설명

SimpleFSDP의 주요 과정:

  1. 모델 초기화 및 파라미터 샤딩
    • 모델의 파라미터, 그래디언트, 옵티마 상태를 샤딩하여 GPU 간 분산.
    • 파라미터는 PyTorch의 DTensor API를 사용해 샤드됨.
  2. IR 노드 최적화
    • 버킷화(Bucketing): 다수의 통신 노드를 하나의 버킷으로 병합해 통신 빈도를 줄임.
    • 재정렬(Reordering): 통신을 계산과 병렬적으로 수행하도록 재배치.
  3. 자동 래핑(Auto-Wrapping):
    • 사용자가 설정할 필요 없이, IR 노드를 자동으로 최적화하여 성능 향상.

예시 입력과 과정 설명

예시: 단순 모델 학습

입력:
Llama 3.1 모델, GPU 64개 사용, 배치 크기 16.

  1. 파라미터 샤딩:
    각 GPU에 샤딩된 모델 파라미터를 분배.
    예) 파라미터 ( W )가 총 4개라면, GPU 1에는 ( W_1 ), GPU 2에는 ( W_2 )가 저장.

  2. 버킷화 예시:
    • 원래 통신 순서: AG1 → AG2 → AG3
    • 버킷화 후: AG1+AG2 → AG3 (한 번에 처리).

  3. 재정렬:
    계산 ( C1 )과 통신 AG1을 병렬로 수행하여 대기 시간을 줄임.

  4. 자동 래핑:
    사용자는 코드에서 직접 정의하지 않아도, SimpleFSDP가 IR 노드를 최적화하여 통신과 계산을 오버랩.
결과:
  • 메모리 사용량 감소: 약 11.61%
  • 처리량 증가: 약 28.26%

한계점

  1. 자동 래핑의 비효율성:
    특정 상황에서 메모리 사용량이 증가하거나 성능이 감소할 수 있음. 예를 들어, 70B 모델의 경우 약 10.61GiB 메모리 오버헤드가 발생.

  2. 네트워크 모델링 한계:
    현재 통신 시간 추정은 단순히 데이터 크기에 의존하며, 네트워크 토폴로지와 같은 요인은 고려되지 않음.

  3. 비추적 가능한 코드:
    데이터 의존적 제어 흐름이 포함된 코드는 풀 그래프 트레이싱을 방해할 수 있음.

결론 및 미래 연구

SimpleFSDP의 의의:

  • 분산 학습에서 간결성과 성능을 모두 잡으면서도 PyTorch 생태계에 최적화된 혁신적인 접근 방식을 제시.
  • 추후 연구로, 네트워크 토폴로지 기반 최적화 및 보다 정교한 자동 래핑 알고리즘 개발이 제안됨.

이 논문은 대규모 모델의 분산 학습에서 최신 기술을 도입하며, 학계와 산업계 모두에 중요한 기여를 했습니다.

Q : 논문의 핵심알고리즘을 예시를 들어 아주 자세하게 설명해주고 정리해 줘

A :

SimpleFSDP의 핵심 알고리즘

SimpleFSDP는 대규모 모델의 Fully Sharded Data Parallel (FSDP) 학습을 최적화하기 위해 버킷화(Bucketing) 및 재정렬(Reordering)을 활용합니다. 아래는 핵심 알고리즘의 과정과 예시를 들어 상세히 설명합니다.

1. 파라미터 샤딩 및 통신 연산

과정:

  • 모델 파라미터 샤딩: 모델의 파라미터는 각 GPU에 균등하게 분배됩니다.
    • 예를 들어, 파라미터 ( W )가 4개이고 GPU가 4개라면, GPU1은 ( W_1 ), GPU2는 ( W_2 ), … 로 샤딩됩니다.
  • 통신 연산: 각 GPU는 필요한 계산을 위해 다른 GPU로부터 일부 데이터를 수집해야 합니다. 이를 위해 All-Gather 및 Reduce-Scatter 연산이 수행됩니다.

2. IR 노드 버킷화(Bucketing)

목표:

통신 연산(All-Gather 및 Reduce-Scatter)의 개수를 줄이고, 병렬성을 극대화하여 계산 속도를 향상.

구현:

  • 버킷화는 여러 개의 All-Gather/Reduce-Scatter 연산을 하나로 묶습니다.
  • 큰 버퍼를 생성하여 데이터 전송을 한 번에 처리.

예시:

  • Before Bucketing:
    • GPU1: All-Gather ( W_1 )
    • GPU2: All-Gather ( W_2 )
    • GPU3: All-Gather ( W_3 )
  • After Bucketing:
    • GPU1~3: All-Gather ( [W_1, W_2, W_3] ) (한 번의 통신으로 처리).

3. IR 노드 재정렬(Reordering)

목표:

통신 연산과 계산 연산이 서로 대기하지 않도록 병렬 수행.

구현:

  • 계산 연산(C1, C2)과 통신 연산(All-Gather 및 Reduce-Scatter)을 CUDA 스트림 상에서 병렬 수행.
  • 선행 통신 작업이 후속 계산 작업과 겹치도록 조정.

예시:

  1. 원래 순서:
    • C1 → All-Gather → Wait → C2
  2. 재정렬 후:
    • All-Gather → C1 (병렬) → Wait → C2

이렇게 하면 계산(C1) 중에 통신이 완료되므로 전체 처리 속도가 빨라집니다.

4. 자동 래핑(Auto-Wrapping)

목표:

수작업 없이 사용자가 정의하지 않아도, IR 노드 버킷화 및 재정렬을 자동으로 최적화.

구현:

  • 그리디 알고리즘을 사용하여 통신 노드를 병합하고 계산 노드와 병렬로 수행.
  • 메모리 제약을 초과하지 않는 선에서 최적화 수행.

예시:

  • GPU에 필요한 파라미터 ( W_1, W_2, W_3 )가 있다면, Auto-Wrapping은 다음을 수행:
    1. ( W_1 )와 ( W_2 )를 하나의 버킷으로 묶음.
    2. ( W_3 )는 다른 버킷으로 처리.
    3. 각 버킷은 계산 ( C1, C2 )와 병렬 수행.

전체 과정 요약

단계 입력 데이터 처리 방식 결과
샤딩 모델 파라미터 ( W ) 각 GPU로 균등 분배 GPU1: ( W_1 ), GPU2: ( W_2 )
버킷화 All-Gather/Reduce-Scatter 연산 병합 및 통신 빈도 감소 통신 연산 수 감소
재정렬 IR 노드 계산과 통신이 병렬 수행되도록 최적화 대기 시간 단축
자동 래핑 모델 전체 최적 래핑을 자동 생성 사용자 입력 없이 최적화

알고리즘의 효과

  1. 메모리 절약: 버킷화와 재정렬을 통해 불필요한 중복 데이터를 최소화.
  2. 처리량 증가: 통신과 계산의 병렬화를 극대화.
  3. 자동화: 사용자의 개입 없이 최적화된 모델 래핑 제공.

결과

Llama 3.1 70B 모델에서:

  • 메모리 사용량: 약 11.61% 감소.
  • 처리량: 약 28.26% 증가.

이 알고리즘은 대규모 모델 학습에서 높은 효율성과 확장성을 제공하며, 분산 학습의 주요 문제를 해결합니다.

Q : 이 논문에서 제시한 결과를 자세하게 보고 다른 방법론에 비하여 특출난 점과 논문에서 제기하는 어떠한 방법이 이러한 결과를 도출하게 되었는지 논문에서 제시하는 이유와 너의 생각을 알려줘

A :

논문에서 제시한 결과와 특출난 점

제시된 주요 결과

  1. 성능 지표:
    • 메모리 절감: 최대 28.54% 절감 (Llama 3.1 405B 모델, FSDP2-Eager와 비교).
    • 처리량 증가: 최대 68.67% 증가 (405B 모델, FSDP2-Eager와 비교).
  2. 모델 크기에 따른 성능 변화:
    • 8B 모델 (FSDP):
      • 메모리: 27.72% 절감.
      • 처리량: 7.49% 증가.
    • 70B 모델 (FSDP + Tensor Parallel + Pipeline Parallel):
      • 메모리: 11.61% 절감.
      • 처리량: 28.26% 증가.
    • 405B 모델 (FSDP + Tensor Parallel + Pipeline Parallel):
      • 메모리: 16.26% 절감.
      • 처리량: 68.67% 증가.
  3. 자동 래핑 성능:
    • 8B 모델에서 수동 래핑 대비 약 7.34% 추가 처리량 증가.
    • 70B 모델에서 약간의 메모리 오버헤드 발생 (10.61GiB)에도 처리량 증가.

다른 방법론 대비 특출난 점

방법론 메모리 절감 처리량 증가 특징
FSDP2-Eager 낮음 낮음 파라미터 샤딩만 지원, 통신 최적화 미흡.
FSDP2-Compile 중간 중간 계산 모듈만 컴파일, 통신 최적화 부족.
SimpleFSDP 최고 최고 통신과 계산의 전체 그래프 추적, IR 버킷화 및 재정렬 제공.

논문에서 제시하는 결과의 이유

1. IR 노드 버킷화(Bucketing)

  • 제안된 방식:
    • 통신 연산(All-Gather, Reduce-Scatter)을 병합하여 통신 빈도를 줄임.
    • 버킷화된 통신은 더 큰 데이터 덩어리를 한 번에 전송.
  • 효과: 통신 대기 시간(Base Latency)이 감소하여 GPU 리소스 활용이 극대화됨.

2. IR 노드 재정렬(Reordering)

  • 제안된 방식:
    • 통신과 계산 연산이 병렬로 수행될 수 있도록 CUDA 스트림 재배치.
    • 계산 작업 도중 통신을 미리 실행하여 전체 처리 시간 단축.
  • 효과: 통신과 계산 간 병렬성 증가로 처리량 대폭 개선.

3. 자동 래핑(Auto-Wrapping)

  • 제안된 방식:
    • 수작업 없이 IR 노드를 병합하고 재정렬하는 최적화.
    • 메모리 제한을 초과하지 않는 범위에서 그리디 알고리즘으로 최적화 수행.
  • 효과: 사용자 개입 없이 최적화 수행, 대규모 모델에서도 성능 유지.

논문에서 제시된 이유에 대한 내 의견

왜 이런 결과가 나왔는가?

  1. 통합적 접근 방식:
    • 기존 방법론(FSDP2)은 계산 모듈만 컴파일하거나 통신을 수동으로 최적화.
    • SimpleFSDP는 통신과 계산의 풀 그래프를 추적하고 최적화, 더욱 정교한 결과를 냄.
  2. IR 노드 기반 최적화의 효율성:
    • IR 노드 버킷화는 GPU 간 통신 오버헤드를 줄이면서 병렬성을 높임.
    • IR 노드 재정렬은 통신 대기 시간을 계산 시간과 겹치게 만들어 효율성을 극대화.
  3. PyTorch-native 구현:
    • PyTorch의 torch.compile과 DTensor API를 활용하여 기존 코드와 통합성이 뛰어남.
    • 기존 FSDP 방식보다 간단하면서도 강력한 성능 향상을 제공.

비판적 관점에서의 평가

  • 장점:
    • 통신 최적화와 계산의 동시 최적화라는 이중적 접근이 실제 학습 성능에 매우 긍정적 영향을 미침.
    • FSDP와 PyTorch-native 생태계의 강점을 최대한 활용.
  • 한계:
    • Auto-Wrapping은 여전히 특정 시나리오에서 메모리 오버헤드를 발생시킴.
    • 통신 시간 추정이 네트워크 토폴로지 등의 복잡한 요소를 반영하지 못해 최적 솔루션에서 벗어날 가능성 있음.

결론

SimpleFSDP는 대규모 모델 학습에 필요한 성능과 메모리 효율을 혁신적으로 개선한 방법론으로, IR 노드 최적화와 PyTorch-native 통합이 핵심 성공 요인입니다.
다만, 자동 래핑의 한계와 통신 시간 추정의 정확성을 보완하면 더욱 강력한 방법론이 될 것입니다.

Q : 이 논문의 입력데이터와 추론 과정에 대해서 예시를 들어 아주 자세하게 설명해 주겠니? 추가적으로 모델아키텍처의 구성 및 모델을 구성하는 연산과 메모리 요구량 컴퓨팅 요구량 대해서도 설명해줘

A :

논문의 입력 데이터 및 추론 과정

입력 데이터

  • 학습 데이터:
    C4 데이터셋 (Common Crawl 기반 텍스트 데이터셋).
    • 대규모 자연어 모델 학습에 적합.
    • Llama 3.1 모델 시리즈(8B, 70B, 405B)를 대상으로 학습.
  • 예시 입력 데이터:
    문장 텍스트, 예: 
    "The quick brown fox jumps over the lazy dog."

추론 과정

SimpleFSDP는 대규모 모델의 분산 추론에서 메모리와 통신 요구량을 줄이고 처리량을 증가시키기 위해 다음 과정을 수행합니다:

  1. 입력 텍스트 처리:
    • 입력 문장은 토큰화하여 정수 시퀀스로 변환: 
      [101, 2005, 2147, 3126, 5409, 2027, 3980, 1996, 3265, 3899, 102]
    • 각 정수는 임베딩 레이어에서 실수 벡터로 매핑.
  2. 파라미터 샤딩:
    • 모델의 파라미터를 각 GPU에 샤딩:
      • 예를 들어, ( W )가 4개라면, GPU1은 ( W_1 ), GPU2는 ( W_2 )만 보유.
    • 샤딩된 파라미터는 All-Gather 연산을 통해 필요한 계산 시에만 통합.
  3. 계산 및 통신 연산:
    • Self-Attention 연산:
      • 쿼리, 키, 값 벡터를 생성.
      • ( Q \cdot K^T ) 연산 후 소프트맥스.
      • 최종 출력은 ( \text{Softmax}(Q \cdot K^T) \cdot V ).
    • 통신 연산:
      • 각 GPU는 필요한 텐서를 All-Gather로 모으고, 계산 후 Reduce-Scatter로 결과를 분배.
  4. 출력 생성:
    • 최종 히든 상태는 출력 레이어에서 디코딩하여 텍스트 토큰 생성.
    • 예: 
      "A quick response."

모델 아키텍처 구성

1. Llama 3.1 모델 구조

  • 층 구성:
    • Transformer Encoder Layers:
      • 쿼리/키/값 생성 및 Self-Attention 연산.
      • 피드포워드 네트워크 (FFN): 활성화 함수와 선형 변환으로 구성.
    • 임베딩 레이어:
      입력 토큰을 고차원 벡터로 변환.
    • 출력 레이어:
      디코더가 최종 출력 확률 분포를 생성.
모델 크기 층 수 모델 차원 FFN 차원 헤드 수
8B 32 4,096 14,336 32
70B 80 8,192 28,672 64
405B 126 16,384 53,248 128

모델 연산 및 자원 요구량

1. 주요 연산

  1. Self-Attention 연산:
    • 복잡도: ( O(n^2 \cdot d) ),
      ( n ): 입력 시퀀스 길이, ( d ): 모델 차원.
    • 예: Llama 70B, ( n = 1024 ), ( d = 8192 ): ( 1024^2 \cdot 8192 = 8.59 \times 10^{10} ) FLOPs.
  2. 피드포워드 네트워크 (FFN):
    • 복잡도: ( O(n \cdot d \cdot d_\text{ff}) ),
      ( d_\text{ff} ): FFN 차원.
    • 예: ( n = 1024 ), ( d = 8192 ), ( d_\text{ff} = 28,672 ): ( 1024 \cdot 8192 \cdot 28,672 = 2.40 \times 10^{11} ) FLOPs.

2. 메모리 요구량

  • 파라미터 저장:
    • 모델 크기에 비례:
      • Llama 70B: ( 70 \times 10^9 ) 파라미터 × 4바이트(FP32) = 280GB.
  • 중간 활성화 값:
    • Self-Attention과 FFN 연산 중 생성.
    • Llama 70B의 경우, ( n = 1024 ), ( d = 8192 ):
      ( 1024 \cdot 8192 \cdot 4 \text{바이트} = 32MB) (단일 GPU).

3. 컴퓨팅 요구량

  • GPU 요구량:
    • Llama 70B 기준, 64개 GPU 필요.
    • 1 GPU당 11 TFLOPs/s 이상 요구.

IR 노드 최적화로 인한 요구량 감소

  • 버킷화 및 재정렬은 중간 활성화 값과 통신 데이터를 효율적으로 관리하여 메모리와 연산 요구량을 줄임.

정리

SimpleFSDP는 입력 텍스트 처리부터 대규모 모델 추론에 필요한 연산 및 메모리 요구량을 최적화하기 위해 설계되었습니다. 모델의 대규모 연산을 분산 처리하는 과정에서 버킷화와 재정렬 최적화가 특히 큰 효과를 발휘하며, 이는 메모리와 계산 자원의 요구량을 줄이고 처리 속도를 높이는 데 기여합니다.

Q : 이 논문의 한계를 극복하기 위한 방법으로 어떤 연구흐름이 있는지 정리해서 자세하게 설명해 줘

A :

논문의 한계 및 이를 극복하기 위한 연구 흐름

1. 논문의 주요 한계

1.1 자동 래핑(Auto-Wrapping)의 비효율성

  • SimpleFSDP는 통신과 계산의 최적화를 위해 자동 래핑(Auto-Wrapping)을 사용하지만, 특정 상황에서 비효율적일 수 있음.
  • 문제점:
    • 메모리와 처리량 간의 균형을 제대로 맞추지 못함.
    • 예: 70B 모델에서 메모리 오버헤드(10.61GiB) 발생.

1.2 통신 시간 추정의 한계

  • 현재 통신 시간 추정은 데이터 크기(전송량)에만 의존.
  • 문제점:
    • 네트워크 토폴로지, GPU 간 연결 대역폭, 병렬 통신 경합 등의 요인을 고려하지 않음.
    • 통신 시간 추정의 부정확성으로 최적화의 비효율 발생.

1.3 데이터 종속적 코드 처리

  • SimpleFSDP는 PyTorch의 torch.compile을 사용하지만, 데이터 의존적 제어 흐름이 있는 코드는 풀 그래프를 추적하지 못함.
  • 문제점:
    • 그래프가 여러 서브그래프로 분리되어 최적화 기회가 감소.

1.4 대규모 분산 환경의 병목

  • 다수의 GPU 노드에서 발생하는 통신 병목은 완전히 해결되지 않음.
  • 문제점:
    • 멀티노드 환경에서 통신과 계산의 동기화가 부족.

2. 한계를 극복하기 위한 연구 흐름

2.1 자동 래핑의 개선

목표: 메모리 효율성과 처리량 간의 균형을 향상.

  1. 강화 학습 기반 최적화:
    • 강화 학습(RL)을 통해 IR 노드의 래핑 및 재정렬을 학습.
    • 메모리 제약과 처리량 목표를 만족하는 최적의 래핑 전략을 생성.
  2. 통합 래핑 알고리즘:
    • 메모리, 처리량, 통신 시간을 동시에 고려하는 다목적 최적화 알고리즘 도입.
    • 예: NSGA-II 같은 다목적 진화 알고리즘.
  3. 동적 래핑:
    • 학습 과정에서 실시간으로 래핑 전략을 조정.
    • 예: GPU 사용량이나 통신 대기 시간이 특정 임계값에 도달하면 자동 재구성.

2.2 통신 시간 추정의 개선

목표: 통신 시간 추정 정확성을 높여 최적화 효율성을 개선.

  1. 네트워크 토폴로지 고려:
    • GPU 간 연결 대역폭과 NVLink/NIC 구조를 모델링하여 통신 시간 추정에 반영.
    • 예: NVIDIA NCCL 라이브러리에서 제공하는 네트워크 프로파일링 데이터 활용.
  2. 통신 경합 모델링:
    • GPU 간 동시 통신의 경합을 고려.
    • 예: 다수의 All-Gather 연산이 경합할 경우, 병렬 수행 효율성 감소를 반영.
  3. 머신 러닝 기반 추정:
    • 통신 로그 데이터를 기반으로 통신 시간을 학습하고 추정.
    • 예: 그래프 신경망(GNN)을 사용하여 통신 연산의 상호작용과 병목 예측.

2.3 데이터 종속적 코드 지원

목표: 데이터 의존적 제어 흐름이 포함된 코드의 그래프 추적 가능성을 향상.

  1. 제어 흐름 인라인 확장:
    • 데이터 의존적 분기문을 그래프 내에서 명시적으로 표현.
    • 예: TorchDynamo의 데이터 흐름 추적을 확장.
  2. 조건부 그래프 최적화:
    • 데이터 의존성에 따라 다른 최적화를 적용.
    • 예: 조건문별로 별도의 최적화 그래프 생성.
  3. Mixed-Mode Compilation:
    • 추적 불가능한 부분은 이거 모드(Eager Mode)로 처리하고, 나머지 부분은 컴파일.

2.4 대규모 분산 환경 최적화

목표: 멀티노드 환경에서 통신 병목과 계산 병렬화를 개선.

  1. 비동기 통신 최적화:
    • 통신 연산을 완전히 비동기로 수행.
    • 예: NCCL의 비동기 API를 활용하여 통신과 계산의 중첩 증가.
  2. 파이프라인 통신 스케줄링:
    • 통신과 계산을 노드 간 파이프라인 방식으로 스케줄링.
    • 예: GPipe와 유사한 방식으로 파라미터 전송을 분할.
  3. 호모지니어스(Homogeneous) 스케줄링:
    • 각 GPU 노드의 작업 분배를 균일하게 조정.
    • 예: GPU 간 로드 밸런싱 알고리즘 도입.

3. 결론

SimpleFSDP의 한계는 주로 자동화 최적화와 통신 모델링의 부족에서 비롯되며, 이를 극복하기 위해 다음과 같은 연구 흐름이 유망합니다:

  1. 자동 래핑 최적화: 강화 학습 및 다목적 알고리즘을 활용한 래핑 전략 개선.
  2. 통신 시간 추정 모델링: 네트워크 토폴로지와 경합을 반영한 ML 기반 추정.
  3. 데이터 종속 코드 처리: 혼합 모드 컴파일 및 조건부 최적화 지원.
  4. 대규모 분산 최적화: 비동기 통신 및 파이프라인 스케줄링 강화.

이러한 연구 흐름은 SimpleFSDP가 대규모 모델 학습에서 더욱 강력한 성능을 발휘하도록 기여할 것입니다.

Q : SimpleFSDP에서 IR 노드의 버킷화(Bucketing)와 재정렬(Reordering) 과정이 기존 FSDP 대비 어떤 구조적 이점을 제공하며, 실제 학습에서 통신과 계산의 병렬화를 어떻게 개선하는가?

A :

SimpleFSDP의 IR 노드 버킷화(Bucketing)와 재정렬(Reordering)의 구조적 이점 및 개선점

1. 기존 FSDP의 구조적 한계

기존 FSDP(특히 PyTorch의 FSDP2)는 각 파라미터에 대해 통신(All-Gather)과 계산(Forward/Backward Pass)을 독립적으로 처리:

  1. 통신 연산 빈도 증가:
    • 파라미터별로 All-Gather와 Reduce-Scatter가 수행되어 통신이 자주 발생.
  2. 통신과 계산의 병렬성 부족:
    • 통신 연산이 완료될 때까지 계산 연산이 대기.
    • CUDA 스트림에서 통신과 계산이 명확히 분리되어 비효율적.

2. SimpleFSDP의 버킷화(Bucketing)와 재정렬(Reordering)의 구조적 이점

2.1 버킷화(Bucketing)의 구조적 이점

  1. 연산 병합:
    • 여러 IR 노드(All-Gather 및 Reduce-Scatter)를 하나의 버킷으로 묶어 한 번의 통신 연산으로 처리.
    • 통신의 고정 오버헤드(Base Latency)를 줄임.
  2. 큰 데이터 덩어리 전송:
    • 개별 파라미터 대신 병합된 데이터를 한 번에 전송하여 네트워크 대역폭 활용도 극대화.
  3. 샘플 사례:
    • 기존 방식: GPU1이 (W_1), GPU2가 (W_2)에 대해 독립적으로 All-Gather 수행.
    • 버킷화 방식: ( [W_1, W_2] )를 하나의 버킷으로 병합해 단일 All-Gather 수행.
  4. 결과: 통신 빈도가 줄어들고 각 통신 연산의 효율성이 증가.

2.2 재정렬(Reordering)의 구조적 이점

  1. 통신과 계산의 병렬화:
    • 통신(All-Gather, Reduce-Scatter)을 CUDA 스트림에서 계산 연산과 겹치도록 재배치.
    • 계산 도중 통신을 미리 수행하여 대기 시간 감소.
  2. 선행 통신 작업 배치:
    • 다음 계산 단계에 필요한 파라미터를 미리 All-Gather하여 현재 계산 단계와 병렬 수행.
  3. 샘플 사례:
    • 기존 방식: (C1 \to All-Gather \to Wait \to C2).
    • 재정렬 방식: (All-Gather \to C1) (병렬) ( \to Wait \to C2).
  4. 결과: CUDA 스트림 상에서 통신과 계산이 병렬적으로 처리되어 GPU 활용률 증가.

3. 실제 학습에서 통신과 계산 병렬화의 개선점

3.1 통신 오버헤드 감소

  • 기존 방식에서는 각 파라미터마다 통신 오버헤드가 발생.
  • 버킷화로 인해 고정 오버헤드(Base Latency)가 병합되어 전체 통신 시간이 감소.

3.2 GPU 사용률 향상

  • 재정렬로 통신과 계산이 병렬적으로 수행되면서 GPU 대기 시간이 감소.
  • 전체 학습 처리량(Throughput)이 증가.

3.3 메모리 관리 최적화

  • 버킷화와 재정렬은 중간 활성화 데이터와 통신 데이터를 적시에 해제하거나 로드할 수 있도록 관리:
    • 메모리 사용량 감소.
    • 통신과 계산의 효율적 자원 할당.

4. 실제 성능 향상

실험 결과 (Llama 3.1 시리즈)

모델 버킷화 & 재정렬 적용 전 버킷화 & 재정렬 적용 후 개선율
8B 처리량: 2.2M TPS 처리량: 2.4M TPS +7.49% 처리량 증가
70B 메모리: 40.8 GiB 메모리: 36.1 GiB -11.61% 메모리 절감
405B 처리량: 4.8M TPS 처리량: 8.1M TPS +68.67% 처리량 증가

5. 한계 및 발전 가능성

  1. 통신 경합 문제:
    • 다수의 GPU 노드가 병렬로 통신을 수행할 때 경합 발생 가능.
    • 해결 방향: 네트워크 토폴로지와 GPU 통신 대역폭을 모델링한 최적화 연구 필요.
  2. 자동화 한계:
    • Auto-Wrapping은 현재 그리디 알고리즘에 의존, 글로벌 최적화 부족.
    • 해결 방향: 강화 학습 기반 또는 다목적 최적화 알고리즘 도입.
  3. 메모리 사용량 최적화:
    • 일부 상황에서 재정렬이 메모리 오버헤드를 유발.
    • 해결 방향: 동적 메모리 관리 및 실시간 리밸런싱 도입.

결론

SimpleFSDP의 버킷화와 재정렬은 기존 FSDP 대비 통신 연산의 빈도를 줄이고 계산과 병렬화하여 GPU 활용 효율을 극대화합니다. 이를 통해 처리량(Throughput)과 메모리 사용량 모두에서 눈에 띄는 개선을 이끌어냈으며, 대규모 모델 학습에서의 분산 시스템 최적화를 선도하는 중요한 기여를 했습니다.

Q : 자동 래핑(Auto-Wrapping) 알고리즘의 설계에서 메모리 사용량과 처리량 간의 균형을 맞추기 위한 최적화 방식이 어떤 기준과 조건에 기반하여 작동하며, 이를 개선하기 위해 추가적인 접근법이 필요할 가능성은 무엇인가?

A :

자동 래핑(Auto-Wrapping) 알고리즘의 설계와 최적화 방식

1. 자동 래핑(Auto-Wrapping)의 목적

SimpleFSDP의 자동 래핑은 IR 노드의 버킷화와 재정렬을 자동으로 처리하여, 수작업 없이 메모리 사용량과 처리량 간의 최적 균형을 찾는 것을 목표로 합니다. 이를 통해 대규모 모델 학습의 복잡성을 줄이고 성능을 최적화합니다.

2. 알고리즘의 작동 기준 및 조건

2.1 주요 입력 변수

  • 통신 시간 (( T_\text{AG}, T_\text{RS} )): All-Gather 및 Reduce-Scatter 연산에 필요한 예상 시간.
  • 계산 시간 (( T_c )): 각 계산 노드의 실행 시간.
  • 메모리 사용량 (( M_c )): 계산 및 통신에 필요한 메모리 사용량.
  • 메모리 제한 (( M_\text{max} )): 각 GPU의 가용 메모리.

2.2 최적화 조건

  1. 시간 제약(Time Constraint):
    • 통신 시간(( T_\text{AG} ))과 계산 시간(( T_c ))이 겹칠 수 있어야 함: [ T_\text{AG} + T_\text{RS} \leq T_c ]
  2. 메모리 제약(Memory Constraint):
    • 통신 및 계산에 필요한 메모리 사용량이 가용 메모리를 초과하지 않아야 함: [ M_c + M_{c(i)} \leq M_\text{max} ]

2.3 알고리즘의 동작

  • 버킷화 결정:
    • IR 노드(All-Gather, Reduce-Scatter)를 메모리 제약과 시간 제약 조건 하에서 병합.
    • 조건을 만족하지 못하면 새 버킷을 생성.
  • 재정렬(Reordering):
    • 병합된 IR 노드가 계산 노드와 병렬로 수행될 수 있도록 CUDA 스트림 재배치.
  • 그리디 알고리즘 사용:
    • 현재 상태에서 메모리 제약과 시간 제약을 만족하는 가장 최적의 버킷화/재정렬을 선택.
    • 글로벌 최적화 대신 로컬 최적화를 반복적으로 수행.

3. 메모리와 처리량 균형의 어려움

3.1 균형 조정의 한계

  1. 그리디 알고리즘의 제약:
    • 로컬 최적화에만 의존하므로 글로벌 최적화에 도달하지 못할 가능성.
    • 예: 메모리 효율성을 극대화하려다 처리량이 감소하거나 반대의 상황 발생.
  2. 모델 크기에 따른 비선형적 메모리 사용량:
    • 대규모 모델(예: Llama 3.1 70B, 405B)의 경우, 중간 활성화 값 및 통신 데이터가 비선형적으로 증가.
  3. 통신 경합 및 네트워크 병목:
    • 멀티노드 환경에서 통신 경합이 발생하면 처리량 저하.

3.2 실제 사례

  • 70B 모델: 자동 래핑에서 메모리 오버헤드(10.61 GiB)가 발생.
    이는 메모리 제한 조건을 완화한 결과 처리량이 증가했으나, 균형 유지 실패를 보여줌.

4. 추가적인 접근법

4.1 강화 학습 기반 최적화

  • 방식: 강화 학습(RL) 알고리즘으로 메모리와 처리량의 균형을 학습.
  • 장점: 글로벌 최적화를 목표로 설정 가능.
  • 적용: RL 에이전트가 메모리와 처리량의 보상 함수(reward function)를 기반으로 IR 노드 래핑 전략을 학습.

4.2 다목적 최적화 알고리즘

  • 예: NSGA-II (Non-dominated Sorting Genetic Algorithm II)
  • 장점: 메모리 사용량 최소화와 처리량 최대화를 동시에 고려.
  • 적용: IR 노드의 래핑과 재정렬을 유전 알고리즘 방식으로 최적화.

4.3 동적 메모리 관리

  • 방식: 학습 중 실시간으로 메모리 사용량을 모니터링하고 필요 시 IR 노드 병합/분할.
  • 장점: 메모리와 처리량 간의 균형을 동적으로 유지.

4.4 네트워크 토폴로지 모델링

  • 방식: 네트워크 병목과 경합을 포함한 통신 시간 예측 모델 도입.
  • 장점: 다수 GPU 환경에서 처리량 최적화 가능.

5. 결론

자동 래핑은 SimpleFSDP의 효율성을 높이는 중요한 구성 요소지만, 메모리와 처리량 간의 균형 조정에 제한이 있습니다. 이를 개선하기 위해 강화 학습 기반 최적화, 다목적 최적화 알고리즘, 동적 메모리 관리, 네트워크 토폴로지 모델링과 같은 접근법이 필요하며, 이러한 개선은 대규모 분산 학습의 성능을 더욱 향상시킬 것입니다.

Q : SimpleFSDP가 PyTorch-native 통합 방식을 사용하여 다른 최신 분산 학습 기술(Tensor Parallel, Pipeline Parallel 등)과의 호환성을 유지하는 구체적인 방법은 무엇이며, 이 과정에서 발생할 수 있는 병목 현상은 어떻게 해결되었는가?

A :

SimpleFSDP의 PyTorch-native 통합 및 최신 분산 학습 기술과의 호환성

1. SimpleFSDP의 PyTorch-native 통합 방식

SimpleFSDP는 PyTorch의 기존 프리미티브(Primitive)를 활용하여 설계되었으며, 이를 통해 Tensor Parallel, Pipeline Parallel 등 최신 분산 학습 기술과 자연스럽게 통합됩니다.

1.1 주요 PyTorch 프리미티브

  1. DTensor API:
    • 데이터 분산 및 통신 관리.
    • 모델 파라미터를 다양한 병렬화 방식(Tensor Parallel, Data Parallel 등)으로 샤딩 가능.
  2. Parametrization 모듈:
    • 파라미터를 샤딩하거나 통합(All-Gather)하는 작업을 자동화.
    • 통신과 계산을 동일한 방식으로 처리하여 추적 가능.
  3. Selective Activation Checkpointing:
    • 특정 통신/계산 연산에서 활성화를 선택적으로 저장 및 해제.
    • 메모리 사용량 감소.
  4. TorchInductor 컴파일러:
    • 통신 및 계산 연산의 IR 노드를 추적하여 최적화.

2. 최신 분산 학습 기술과의 호환성

2.1 Tensor Parallel

  • 구현 방식:
    • DTensor를 활용하여 각 레이어의 파라미터를 Tensor Parallel 서브그룹으로 분할.
    • 계산 단계에서는 All-Gather를 통해 필요한 데이터를 통합.
  • 특징:
    • SimpleFSDP는 Tensor Parallel 연산에 필요한 통신(All-Gather 및 Reduce-Scatter)을 IR 노드로 추적하여 최적화.
  • 결과:
    • 데이터 분산 및 샤딩 작업이 병렬로 수행됨.
    • Tensor Parallel과의 호환성을 유지하며 처리량 최적화.

2.2 Pipeline Parallel

  • 구현 방식:
    • 모델을 여러 서브모듈로 분할하여 GPU 간 파이프라인 실행.
    • SimpleFSDP는 각 파이프라인 단계에서 통신과 계산의 IR 노드를 추적.
  • 특징:
    • 파라미터를 각 파이프라인 단계로 제한하여 메모리 사용량 감소.
    • 통신과 계산의 순차적 병목을 최소화.
  • 결과:
    • Pipeline Parallel의 데이터 흐름과 자연스럽게 결합.

2.3 혼합 정밀도 학습(Mixed Precision Training)

  • 구현 방식:
    • DTensor API를 사용하여 모델 파라미터를 16비트(float16)로 변환.
    • 그래디언트 업데이트 시에는 32비트(float32)로 변환.
  • 특징:
    • 단일 IR 그래프에서 메모리 효율성과 학습 안정성을 모두 제공.

2.4 Meta Initialization

  • 구현 방식:
    • 초기화 시 파라미터를 Meta 디바이스(추상 디바이스)에 로드.
    • 실제 GPU/CPU에 로드하기 전에 모델 구조를 분석.
  • 특징:
    • 모델 초기화 속도 증가 및 메모리 절약.

3. 병목 현상 및 해결 방법

3.1 병목 현상

  1. 통신 병목:
    • Tensor Parallel 및 Pipeline Parallel 환경에서 GPU 간 통신 시간이 증가.
    • 멀티노드 환경에서는 네트워크 대역폭 경합 발생.
  2. 계산과 통신의 비효율적 병렬화:
    • 계산 작업이 통신 완료를 대기하는 경우 처리량 감소.
  3. 메모리 사용량 증가:
    • 통신 데이터와 계산 데이터가 동시에 메모리를 차지할 때 메모리 오버헤드 발생.

3.2 해결 방법

  1. 통신 병목 해결:
    • 버킷화(Bucketing): 다수의 All-Gather 및 Reduce-Scatter 연산을 병합하여 통신 빈도 감소.
    • 비동기 통신: 통신 연산을 비동기로 수행하여 계산과 겹치도록 최적화.
  2. 계산과 통신 병렬화:
    • 재정렬(Reordering): 통신 작업을 계산 작업과 병렬로 수행하도록 CUDA 스트림 조정.
    • 선행 통신(Prefetching): 다음 계산 단계에 필요한 데이터를 미리 All-Gather.
  3. 메모리 사용량 최적화:
    • Selective Activation Checkpointing: 특정 활성화 값만 저장하여 메모리 절약.
    • 동적 메모리 관리: 학습 과정에서 메모리 사용량을 실시간 조정.

4. 실제 성능 개선

  • Llama 3.1 모델에서 Tensor Parallel 및 Pipeline Parallel과 결합 시:
    • 8B 모델: 처리량 29.35% 증가.
    • 70B 모델: 메모리 사용량 11.61% 감소.
    • 405B 모델: 처리량 68.67% 증가.

5. 결론

SimpleFSDP는 PyTorch-native 통합 방식을 통해 Tensor Parallel, Pipeline Parallel 등 최신 분산 학습 기술과 완벽히 호환되며, IR 노드 추적 및 최적화를 활용하여 통신과 계산의 병렬화를 극대화합니다. 이 과정에서 발생하는 병목 현상은 버킷화, 재정렬, 비동기 통신, 동적 메모리 관리 등을 통해 효과적으로 해결되었으며, 대규모 모델 학습의 성능을 크게 개선하였습니다.