Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts

논문 링크

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

A :

논문 “Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts”는 MoE 모델의 핵심 병목인 GPU 간 통신 지연을 fine-grained 수준에서 컴퓨팅과 통신을 정교하게 겹치도록 설계함으로써 실행 성능을 크게 개선한 ByteDance의 시스템 최적화 논문입니다.


📌 결론 요약

Comet은 기존 MoE 시스템 대비 최대 1.96× 단일 레이어 속도 향상, 전체 모델 기준 평균 1.71× 속도 향상을 달성하였으며, ByteDance의 10,000+ GPU 클러스터에 상용 배포되어 수백만 GPU 시간 절감 효과를 보고함.

모델 비교 시스템 속도 향상 배율
단일 MoE Layer Megatron-Cutlass 외 3종 1.28× ~ 2.37×
End-to-End Megatron-Cutlass / TE / FasterMoE / Tutel 1.71× 평균 (최대 44.4%)

🧠 핵심 기여 및 독창성

기여 항목 설명
1. Shared Tensor 기반의 Data Dependency 분석 및 재구성 통신과 연산 간의 미세 단위 dependency를 추출하여 shared tensor를 M 또는 N 방향으로 나누고, computation tile 단위로 reschedule
2. Adaptive Thread Block Assignment GPU Thread Block을 통신/연산으로 분리한 후, workload에 따라 비율을 동적으로 조절하여 latency concealment 극대화
3. Fused Kernel을 통한 intra-kernel pipelining 별도 kernel 호출 대신 하나의 kernel 내에서 통신과 연산을 thread block 수준에서 분리해 fine-grained overlapping 실현
4. 실용성 높은 구현 및 실제 클러스터 적용 NVSHMEM 기반 구현, Megatron-LM 통합, 다양한 병렬 전략(Tensor/Expert) 및 플랫폼(H800, L20) 대응

🔍 핵심 알고리즘 설명 (예시 기반)

문제 상황:

MoE의 Layer0에서는 각 token이 여러 expert로 분산되어 GPU 간 통신 필요 → GEMM 실행은 tile 단위 (예: 128×128), 하지만 통신은 token 단위 → granularity mismatch 발생


예시 입력:

  • 총 토큰 수 ( M = 4096 ), 각 token은 ( N = 4096 ) 차원
  • top-k = 2 → token마다 2개의 expert로 라우팅됨
  • Expert 수 ( E = 8 ), GPU 수 = 8 (expert parallelism)

Comet 처리 흐름:

  1. Shared Tensor Decompose
    • Layer0에서는 GEMM의 입력이므로 ( M ) 방향(토큰 단위)으로 tensor 분해
    • Layer1에서는 TopK reduce 결과를 보낼 때 ( N ) 방향(embedding 차원)으로 분해
  2. Reschedule: 통신-연산 파이프라인 정렬
    • 각 expert의 GEMM tile이 가능한 한 local token으로 시작되도록 token 정렬
    • Remote token은 비동기로 가져오면서 다른 tile 연산 진행
  3. Adaptive Thread Block Assignment
    • GPU SM에 thread block 배정 시, 통신 전용 / 연산 전용 block 분리
    • Comet은 입력 길이, top-k, expert 수에 따라 최적 분할비 (e.g. 통신:연산 = 26:106)를 동적으로 선택
  4. Fine-grained Fused Kernel 실행
    • 통신 thread block은 async I/O (TMA 등)로 remote token 수신
    • 연산 thread block은 GroupGEMM tile 실행
    • 공유 메모리 경유로 통신-연산을 안전하게 동기화

🧪 성능 분석 (실험 결과)

실험 조건 Comet의 개선 효과
Input Token 수 증가 (2048→32768) 1.28× ~ 2.37× layer 속도 개선
top-k 및 expert 수 증가 Tutel, FasterMoE 대비 1.16× ~ 1.83× 안정적 개선
GPU 간 대역폭 제약 (L20 클러스터) 1.19× ~ 1.46× 개선, Comet의 유연성 입증
Load imbalance 상황 (token 분포 std=0.05) 여전히 Comet이 최저 latency 유지

⚠️ 한계점 및 향후 방향

한계 설명
NVSHMEM 의존성 저수준 통신 라이브러리 사용 → 유지보수 및 호환성 문제 가능
실행 전 profiling 필요 각 모델 config마다 최적 thread block 비율을 profiling 해야 함
아직은 MoE layer에 집중 attention, embedding 등 다른 연산과 통합된 end-to-end 전체 최적화는 미진

💡 요약

Comet은 기존 coarse-grained pipelining 방식의 한계를 극복하고, fine-grained 수준에서 연산과 통신을 정교하게 조율함으로써 MoE의 병목을 구조적으로 해결한 고성능 시스템입니다. 특히, CUDA 커널 레벨에서의 통합 스케줄링 및 동적 리소스 분배는 일반적인 kernel-level 병렬화에서 구현이 어려웠던 부분을 효과적으로 해결한 점에서 높은 독창성과 실용성을 지닙니다.

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

A :

좋습니다. 논문 Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts의 핵심 알고리즘은 크게 두 가지입니다:


🧩 핵심 알고리즘 구조 요약

구성 요소 역할 핵심 아이디어
① Shared Tensor 기반 Dependency Resolving 통신-연산 파이프라인 최적화 Tensor 분해 + 연산 재스케줄
② Adaptive Workload Assignment GPU 자원 효율화 Thread block 비율을 동적으로 조정

이제 실제 예시를 기반으로 각 알고리즘이 어떻게 작동하는지 자세히 설명하겠습니다.


🔎 예시 설정

모델 설정

  • Token 수 ( M = 4096 ), 임베딩 차원 ( N = 4096 )
  • Expert 수 ( E = 4 ), top-k = 2
  • Expert parallelism: 각 GPU는 1개의 expert를 가지고 있음 (GPU0 ~ GPU3)
  • GEMM tile 크기 = 128 × 128

🧠 ① Shared Tensor 기반 Dependency Resolving

💡 목적

통신과 연산을 동시에 처리하기 위해, 연산에서 필요한 데이터가 준비되는 순간 바로 연산을 시작하게 만드는 것


📌 MoE Layer0의 예

✅ Step 1: Shared Tensor 분해 (Decompose)

  • 각 token은 top-2 expert로 라우팅됨 → GPU 간 All-to-All 통신 필요
  • GEMM 연산은 보통 128개의 token 단위로 tile 연산
  • 따라서, shared tensor를 M 방향(토큰 방향) 으로 분해
    • 예: tensor shape = (4096 × N), → 32개 블록 (4096/128 = 32)

✅ Step 2: 재스케줄 (Reschedule)

  • 각 expert는 local + remote token이 섞여 있음
  • 먼저 local token만 포함된 tile부터 GEMM 실행 시작
    • 예: GPU1의 Expert1은 local token 0~127을 먼저 처리
  • 그 사이 GPU1은 GPU0/GPU2로부터 필요한 remote token을 비동기로 수신

➡️ 결과적으로 연산과 통신이 자연스럽게 겹침 (overlap)


📌 MoE Layer1의 예

✅ Step 1: Shared Tensor 분해 (Decompose)

  • GEMM 후 결과를 top-k reduce하여 통신함 → GEMM 결과는 (M×N)
  • Reduce는 token끼리 합치므로 M 방향으로 분해 ❌, 대신 N 방향 분해 ✔️

✅ Step 2: 재스케줄

  • GEMM 결과는 column-wise로 처리되며
  • 앞부분 N 차원의 column부터 먼저 계산하고 바로 reduce+send
    • 예: column 0~127 완료 → 곧바로 All-to-All 시작

➡️ 연산 끝날 때까지 기다리지 않고, 앞부분부터 통신을 시작할 수 있게 됨


🧠 ② Adaptive Workload Assignment

💡 목적

연산 thread block과 통신 thread block의 수량 비율을 자동 조정하여 pipeline 효율 극대화


예시

  • GPU에 132개의 SM이 있음 → 132개의 thread block 배치 가능
  • 각 kernel 실행 시, nc개는 통신용 thread block, np개는 연산용 thread block 할당

문제: 어떤 비율이 최적일까?

해결:

  • 실행 전에 다양한 입력 길이와 모델 구성에 따라 profile 수행
  • 예: ( M=16384, TP=8 )이면 → 최적 nc=26, np=106
  • 이 값을 metadata로 저장 → 런타임에 동적으로 최적 커널 선택

➡️ workload마다 latency concealment 극대화 (예: 통신과 연산 latency 비슷하게 맞춤)


🎯 전체 실행 흐름 요약

[Input Tokens] →
① Shared Tensor 생성 →
①-1 M/N 방향으로 분해 →
①-2 Token 정렬 및 tile 단위 연산 재배열 →
② Adaptive thread block 분배 (e.g., nc=26, np=106) →
③ Fused kernel 내에서
   ├─ 통신 block: remote token 비동기 수신
   └─ 연산 block: local token부터 GEMM 시작 →
④ 연산 결과는 곧바로 column-wise로 reduce+send →
[Output Token]

📌 시각화 요약

  ┌────────────┐         ┌────────────┐
  │ Remote I/O │◄───────▶│ Local GEMM │
  └────────────┘         └────────────┘
        ▲                      ▲
        │ Shared Tensor (분해/재정렬)
        ▼                      ▼
  ┌─────────────────────────────────┐
  │    Fine-grained Overlapping     │
  └─────────────────────────────────┘
        ▲                      ▲
        │                      │
   [통신 Thread Block]    [연산 Thread Block]

✅ 요약 정리

알고리즘 설명 주요 효과
Shared Tensor Dependency Resolving 통신과 연산 사이의 데이터 의존성 해소 위해 tensor를 분해 및 재스케줄 미세 단위 연산-통신 overlap 실현
Adaptive Workload Assignment 통신/연산 thread block의 수를 동적으로 할당 SM 활용도 극대화 + bubble 최소화
Fused Kernel 실행 하나의 kernel에서 통신과 연산을 분리된 thread block으로 동시에 실행 커널 호출 오버헤드 제거

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

A :

논문 “Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts”는 기존 MoE 시스템들과 비교하여 성능 향상폭이 일관되고 크며, 그 원인을 아키텍처 차원의 정밀한 병렬성과 리소스 배분 전략에서 찾습니다. 아래에서는 Comet이 구체적으로 어떤 실험 조건에서 얼마나 뛰어난 성능을 보였는지, 그리고 어떤 메커니즘이 성능 향상을 유도했는지, 마지막으로 이에 대한 제 생각까지 종합적으로 정리하겠습니다.


✅ 실험 결과 요약: 수치로 보는 Comet의 강점

1. 📈 End-to-End 모델 성능 비교

모델 비교 대상 Comet 성능 향상
Mixtral-8x7B, Qwen2-MoE, Phi-3.5-MoE Megatron-Cutlass 34.1% 감소
  Megatron-TE 42.6% 감소
  FasterMoE 44.4% 감소
  Tutel 31.8% 감소

2. 🔍 단일 MoE Layer 성능

Input Token 수 Comet vs. Baseline 향상 폭
M = 2048~32768 Megatron-TE 외 3종 1.28× ~ 2.37×

3. 🎯 통신 Overlap 효과

시스템 통신 latency 은닉률
Comet 86.5%
Tutel 68.6%
FasterMoE 29.2%
Megatron 계열 0% (overlap 없음)

🔍 Comet이 특출났던 이유 (논문에서 제시한 주장)

🎯 1. Fine-grained 수준에서의 연산/통신 병렬화

  • 기존 방법(FasterMoE, Tutel)은 kernel-level 수준에서 coarse하게 겹치거나 수동 pipelining만 제공
  • → Comet은 shared tensor를 token 단위로 분해 후, tile granularity로 연산 재스케줄
  • ✅ 효과: 전례 없는 수준의 latency hiding

🎯 2. Fused kernel 내 thread block specialization

  • 통신과 연산을 동일 kernel 내 thread block으로 완전히 분리하여 SM 간 리소스 충돌 방지
  • ✅ 효과: CUDA kernel 호출 오버헤드 제거 + 동시 실행 보장

🎯 3. Adaptive workload assignment

  • 입력 길이, expert 수, top-k 등 다양한 상황에 따라 thread block 비율 동적으로 조정
  • ✅ 효과: 다양한 config에서 consistently 우수한 성능

🎯 4. Host scheduling 제거

  • FasterMoE, Tutel은 전문가마다 kernel을 생성해야 해서 작은 expert 수에서는 오버헤드 심각
  • Comet은 fused kernel 내부에서 자체 스케줄링 수행
  • ✅ 효과: 작은 입력에서도 성능 저하 없음

🧠 내 생각: Comet이 왜 차별적인가?

관점 기존 방식 Comet
병렬 granularity 전문가별/배치별 tile/token 단위 (fine-grained)
통신 연산 통합 별도 kernel fused kernel 내부 병렬 thread block
runtime adaptiveness 없음 또는 고정 동적 thread block 조절
실질적 latency 은닉 미흡 86% 이상 은닉

📌 핵심 통찰:

기존 시스템들이 단순히 통신 시점을 당기거나 연산을 block-wise로 나누는 수준에 그쳤다면, Comet은 CUDA kernel 설계 차원에서 문제를 다시 정의함으로써 병목을 해결함. 이는 단순한 기법의 조합이 아니라, “연산 단위와 통신 단위 간 granularity mismatch를 해결하는 구조적 접근”으로 볼 수 있음.


💬 향후 확장 가능성

논문에서 언급된 것처럼, 이 구조는 TVM이나 Triton 같은 컴파일러 백엔드로 확장 가능하며, 실제 “token dispatch-aware kernel fusion” 같은 개념을 도입하면 더 일반적인 sparse 모델(예: sparse MLP, Switch Transformer 등)에도 적용 가능함.


✅ 요약 정리

항목 내용
성능 최대 2.37× 속도 향상, 통신 latency 86.5% 은닉
구조적 차별점 shared tensor 재구성 + fused kernel 병렬화
핵심 메커니즘 fine-grained dependency resolving + adaptive thread block assignment
내 견해 단순한 pipelining이 아닌, GPU 커널 스케줄링의 구조적 최적화

Q :이 논문에서 제시된 방법을 수도 코드로 설명을 해주었으면해 구체적인 주석이 추가가 되면 좋을것 같아

A :

아래는 Comet 논문의 핵심 알고리즘을 반영한 수도 코드 (pseudocode)입니다. 구조적으로는 두 핵심 구성인:

  1. Shared Tensor 기반 Dependency Resolving (분해 및 재스케줄링)
  2. Adaptive Workload Assignment (thread block 비율 조정)

을 중심으로 구성했으며, GPU의 SM 단위 실행을 추상화하여 launch_kernel()로 표현하고 CUDA 스타일의 thread block 역할도 구체적으로 주석으로 설명합니다.


📦 1. 주요 데이터 구조 정의

# SharedTensor: 통신-연산 사이의 공유 버퍼
class SharedTensor:
    def __init__(self, shape):  # shape = (M, N)
        self.data = zeros(shape)
        self.lock = Lock()  # fine-grained access 관리용

🔄 2. Dependency Resolving: shared tensor 분해 + 재스케줄

def decompose_and_reschedule(shared_tensor, layer_type):
    """
    shared_tensor: 통신과 연산 사이의 버퍼
    layer_type: "layer0" or "layer1" (통신-연산 or 연산-통신)
    """
    if layer_type == "layer0":
        # Layer0: GEMM이 consumer → token dimension(M)을 따라 분해
        chunks = split_tensor(shared_tensor, axis='M', chunk_size=128)  # GEMM tile 기준
        # local 우선 실행을 위한 토큰 정렬 (source GPU 기준)
        sorted_chunks = sort_by_source_rank(chunks)
        return sorted_chunks

    elif layer_type == "layer1":
        # Layer1: Top-K Reduce가 consumer → embedding dimension(N) 따라 분해
        chunks = split_tensor(shared_tensor, axis='N', chunk_size=128)
        return chunks

⚙️ 3. Adaptive Thread Block Assignment

def select_optimal_thread_block_config(M, N, topk, TP, EP):
    """
    입력 크기 및 모델 설정에 따라 최적의 thread block 비율 선택
    """
    key = (M, topk, TP, EP)
    if key in PROFILED_CONFIG:
        return PROFILED_CONFIG[key]  # 사전 프로파일링된 메타데이터
    else:
        # 기본 fallback
        total_blocks = 132  # Hopper 기준
        nc = int(0.2 * total_blocks)  # 통신용
        np = total_blocks - nc       # 연산용
        return nc, np

🚀 4. Fused Kernel 실행 (커널 내부 동작 추상화)

def launch_comet_fused_kernel(shared_tensor, layer_type, nc, np):
    """
    하나의 CUDA 커널 내부에서 통신 thread block과 연산 thread block을 동시에 실행
    """
    # 각 thread block은 아래와 같은 동작을 함
    parallel_for block_id in range(nc + np):

        if block_id < nc:
            # 🔴 통신용 thread block
            while not all_required_chunks_ready(shared_tensor):
                data = remote_fetch_next_token()
                write_to_shared_tensor(shared_tensor, data)

        else:
            # 🔵 연산용 thread block (GEMM or Reduce)
            chunks = get_ready_chunks(shared_tensor, block_id)
            for chunk in chunks:
                if layer_type == "layer0":
                    result = group_gemm(chunk)
                else:
                    result = topk_reduce_and_scatter(chunk)
                write_output(result)

🧠 전체 실행 흐름 정리

def run_comet_pipeline(M, N, topk, TP, EP, layer_type):
    # Step 1: Shared tensor 준비
    shared_tensor = SharedTensor((M * topk, N))

    # Step 2: Shared tensor 분해 및 재정렬
    chunks = decompose_and_reschedule(shared_tensor, layer_type)

    # Step 3: 최적 thread block 수 계산
    nc, np = select_optimal_thread_block_config(M, N, topk, TP, EP)

    # Step 4: 커널 실행 (통신/연산 병렬 수행)
    launch_comet_fused_kernel(shared_tensor, layer_type, nc, np)

✅ 주석 요약 포인트

  • decompose_and_reschedule()는 통신과 연산 사이의 granularity mismatch를 해소
  • select_optimal_thread_block_config()는 workload 맞춤형 thread block 비율 제공
  • launch_comet_fused_kernel()은 하나의 fused kernel에서 thread block 단위로 역할 분리
    • block_id < nc: 통신 전용 thread block (NVSHMEM 기반 I/O)
    • block_id >= nc: 연산 전용 thread block (CUTLASS 기반 GEMM, Top-K Reduce 등)

이 구조는 CUDA GPU programming abstraction논문 설계 흐름을 정확히 대응시키며, 향후 Triton이나 TVM으로 구현 시에도 그대로 맵핑 가능한 수준으로 설계되었습니다.

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

A :

좋습니다. 이 논문인 Comet은 Mixture-of-Experts (MoE) 기반 대형 모델에서의 추론(inference) 또는 학습 과정에서 발생하는 GPU 간 통신 병목을 해소하기 위한 시스템적 최적화를 다룹니다. 질문하신 내용을 기준으로 다음과 같이 설명하겠습니다:


1. 📥 예시 기반 입력 데이터와 추론 과정

📌 예시: Mixtral-8x7B 모델 기준 (논문 실험 대상)

항목
Input token 수 ( M ) 8192
Embedding dimension ( N ) 4096
Hidden size ( K ) 14336
Expert 수 ( E ) 8
Top-k routing 2 (각 token은 2개의 expert에 전달)

🔁 추론 전처리 및 라우팅 (Token Routing)

  1. Input: 8192 tokens → [8192, 4096] 크기의 입력 텐서
  2. Gating:
    • 각 token은 top-2 expert로 라우팅됨 → 총 16384개 token이 expert에 전달됨
    • 실제 routing은 gate 네트워크(작은 MLP)가 softmax 후 top-2 선택
  3. All-to-All 통신:
    • Expert는 여러 GPU에 분산되어 있음 → 각 GPU는 자신에게 라우팅된 token만 받음

🧠 Expert Computation (MoE Layer 0)

각 expert는 일반적인 FFN 구조:

GEMM1: [token, 4096] × [4096, 14336] → [token, 14336]
Activation: GELU or ReLU
GEMM2: [token, 14336] × [14336, 4096] → [token, 4096]
  • token은 GPU별로 분산됨 (e.g., GPU0은 3000개, GPU1은 1500개 등)
  • 이 연산은 모든 expert에 대해 병렬적으로 수행됨 (Grouped GEMM)

📤 결과 병합 (MoE Layer 1)

  • 각 expert의 출력은 다시 All-to-All로 전송되어 원래 token 위치로 결합됨
  • Top-k reduction 수행:
    • 각 token의 두 expert 출력 → gating score 기반 weighted sum
    • 최종 output shape: [8192, 4096]

2. 🏗️ 모델 아키텍처 구성 (Comet 관점)

전체 구조 (1 Transformer block 기준)

Input Embedding
   ↓
Multi-Head Self Attention
   ↓
LayerNorm
   ↓
MoE Layer
  ├─ Dispatch (token → expert)
  ├─ Expert FFNs (GEMM1 + Activation + GEMM2)
  └─ Combine (Top-k reduction)
   ↓
LayerNorm
   ↓
Output

이 중 Comet은 MoE Layer에만 집중하여 통신-연산 병렬화를 수행합니다.


3. 📊 연산량(Computational Cost) 분석

✅ MoE Layer 기준: 한 expert의 FFN

연산 형태 FLOPs 계산
GEMM1 [T, 4096] × [4096, 14336] ( 2 \times T \times 4096 \times 14336 )
GELU [T, 14336] 약 ( 10 \times T \times 14336 )
GEMM2 [T, 14336] × [14336, 4096] ( 2 \times T \times 14336 \times 4096 )
합계   약 ( 2 \times T \times 4096 \times 14336 \times 2 ) FLOPs
  • ( T ): 해당 expert에 라우팅된 token 수
  • 예: GPU당 평균 T = 2048이면, 약 480 GFLOPs per expert

4. 🧠 메모리 요구량

Shared Tensor 크기

shared tensor는 GEMM 입력/출력 및 통신을 위한 버퍼입니다.

  • MoE Layer0: ( M \times \text{topk} \times N )
    • 예: ( 8192 \times 2 \times 4096 \times 2 ) Bytes (BF16) ≈ 128MB
  • MoE Layer1: GEMM 출력 후 reduce 대상 → 동일 크기

→ 총 약 256MB의 버퍼 필요, 이는 NVSHMEM을 통해 모든 GPU에서 공유됨


5. 🧩 Comet에서의 최적화 요소 반영

처리 단계 Comet에서의 최적화 방법 효과
Dispatch Shared Tensor 기반 M 방향 분해 + 토큰 재배치 통신 시작과 연산 병렬화 가능
GEMM 실행 GroupGEMM → tile 우선순위 재스케줄 local token tile 우선 실행으로 overlap 증가
Combine N 방향 분해 → column-wise GEMM 처리 연산 일부 끝나자마자 통신 가능
통신 & 연산 통신/연산 thread block 분리 (fused kernel) CUDA scheduling 오버헤드 제거

✅ 정리 요약

항목 내용
입력 데이터 [8192, 4096] token embedding
라우팅 방식 top-2 expert routing (gate 기반)
주요 연산 GEMM1 + Activation + GEMM2 + Top-k Combine
FLOPs 약 2.4 TFLOPs (모든 expert 처리 시)
메모리 요구량 Shared Tensor 256MB (BF16 기준)
Comet 효과 통신-연산 fine-grained overlap + 커널 내 스케줄링으로 1.7×~2× 속도 향상

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

A :

좋습니다. 이 논문 Comet은 Mixture-of-Experts(MoE) 구조에서 fine-grained 수준의 통신-연산 overlap을 구현하여 기존 방법의 coarse-grained pipelining 한계를 극복했습니다. 하지만, 여전히 몇 가지 구조적, 시스템적 한계가 존재하며, 이러한 한계를 극복하려는 연구들은 다양한 방향으로 발전하고 있습니다.


✅ Comet의 주요 한계 요약

범주 한계점 설명
1. 통신 라이브러리 종속성 NVSHMEM 낮은 수준의 GPU 직접통신 API로, 플랫폼 독립성/유지보수 어려움
2. 컴파일러 최적화 부재 수작업 커널 튜닝 컴파일러 기반의 자동화된 tile 분해/스케줄링이 없음
3. attention 등 타 연산과의 통합 부재 범용성 제한 MoE 연산 외의 dense block에 대한 latency overlap은 고려되지 않음
4. 실행 전 profiling 필요 static optimal block config 입력/하드웨어가 달라지면 profiling 재필요 → 실시간 적응 어려움

🔍 한계 극복을 위한 연구 흐름별 정리

🧠 [1] 컴파일러 기반 자동화: TVM, Triton, Unity

접근법 설명
Triton-based Fused MoE Kernels 연산자 fusion + 메모리 관리 자동화
TVM dependency-aware scheduling + cost model 기반 autotuning
Unity (Meta, 2024) 모델 표현-컴파일-스케줄링 전체 자동화 파이프라인 구축

핵심 기여:
Comet의 shared tensor 분석 및 reschedule 과정을 자동으로 추론 및 커널 생성
→ 사람이 hand-tune한 커널보다 더 효율적인 layout & schedule을 제안 가능

🔗 대표 연구:

  • Unity: AI compiler for unified model optimization (Meta, 2024)
  • TrMoE: Triton-based sparse expert execution (ICLR 2025 under review)

🌐 [2] 네트워크/토폴로지 최적화: Hierarchical All-to-All & Topology-aware Scheduling

접근법 설명
2D Hierarchical All-to-All intra-node / inter-node 분리 후 효율적 통신
Topology-aware token dispatch bandwidth에 따라 routing 최적화 (Hotspot 방지)

핵심 기여:
NVSHMEM을 사용하는 fine-grained 통신의 한계를 통신 스케줄링 계층에서 완화
→ 더 큰 클러스터에 적응 가능

🔗 대표 연구:

  • Tutel (2023) – Hierarchical All2All + adaptive MoE
  • HetuMoE (2022) – 트래픽을 균형있게 분산시켜 통신 병목 완화

🔄 [3] Online/Runtime 재구성: Dynamic Expert Scheduling & Adaptive Routing

접근법 설명
Dynamic Expert Pruning / Dropout 실행 시점에 부하 기반으로 expert 활성화
Expert Affinity Routing 통신량 최소화 방향으로 routing 동적 조절
Latency-aware Kernel Dispatcher 입력 token 수에 따라 fused kernel 동적 선택

핵심 기여:
Comet의 static profiling 의존성을 해결하고 실시간 workload에 적응
→ workload imbalance나 token skew에도 강건

🔗 대표 연구:

  • Gating Dropout (ICML 2022) – MoE 통신량 감소
  • SparTA (Sparse Transformer Autotuner) – online tile tuning

🔀 [4] 다른 연산자와의 통합: Attention-MoE Fusion, Cross-block Scheduling

접근법 설명
Fused Attention + MoE Block 연속된 attention → MoE block을 fused kernel로 실행
Pipeline-aware compiler 연산자 간 dep graph를 global scheduling

핵심 기여:
현재 MoE layer만 target으로 하는 Comet을 모델 전체 수준으로 확장

🔗 대표 연구:

  • PipeMoE (INFOCOM 2023) – MoE operator 간 full-path pipeline 구성
  • Flux (ByteDance, 2024) – MoE 외 연산까지 fused kernel 구성

📌 비교 요약 테이블: 연구 흐름별 특징

연구 흐름 해결 대상 대표 기술 기대 효과
Compiler 기반 자동화 수작업 커널 최적화 Triton, TVM, Unity 구조적 재현성과 확장성
통신 최적화 통신량/latency Hierarchical All2All 더 큰 클러스터 대응
Runtime 적응 Static tuning Gating Dropout, SparTA 입력 변화에 유연
전체 모델 통합 연산자 간 경계 PipeMoE, Flux 모델 수준 latency 최적화

✅ 정리

  • Comet은 MoE 구조 내 통신-연산 병목을 fine-grained pipelining + fused kernel 방식으로 해결한 최첨단 시스템이지만, 컴파일러 지원 부재, 통신 추상화 제한, 전체 모델 통합 부족 등의 구조적 한계가 있습니다.
  • 이를 극복하기 위한 연구는 아래와 같은 흐름으로 확장 중입니다:
    1. 컴파일러 기반 자동화 (Triton/TVM/Unity)
    2. 네트워크 aware 통신 최적화 (Tutel, HetuMoE)
    3. 동적 실행 적응 (SparTA, Gating Dropout)
    4. 모델 전체 연산 통합 (PipeMoE, Flux)

각 흐름은 Comet의 장점을 유지하면서도 확장성과 실시간성, 범용성 측면에서 더 큰 발전 가능성을 열어주고 있습니다.

Q : Comet이 기존 MoE 시스템(FasterMoE, Tutel 등) 대비 fine-grained overlapping을 실현하기 위해 어떤 granularity-level 최적화를 적용했으며, 이로 인해 어떠한 구조적 trade-off가 발생하는가?

A :

좋은 질문입니다. Comet이 기존 MoE 시스템(FasterMoE, Tutel 등)과 비교해 가장 근본적으로 차별화되는 점은 granularity 수준에서의 정교한 최적화입니다. 아래에서 구체적으로 어떤 수준에서 fine-grained overlapping을 적용했는지, 그리고 그에 따라 생기는 구조적 trade-off는 무엇인지 정리해드립니다.


✅ Comet의 granularity-level 최적화

1. 기존 시스템의 한계 (Coarse-Grained Overlap)

| 시스템 | 겹침 수준 | 단점 | |——–|————|——| | FasterMoE | expert 단위 pipelining | 작은 expert 처리 시 kernel launch 오버헤드 커짐 | | Tutel | token group 단위 pipelining | Heuristic scheduling → 실제 종속성 반영 어려움 |

이들은 통신과 연산을 별도의 kernel로 나누어 coarse하게 병렬 실행하며, kernel 간 오버헤드, 종속성 미스매치 등으로 latency hiding이 제한적


2. Comet의 fine-grained overlapping 방식

항목 Comet의 접근
Tensor 분해 단위 token dimension(M), embedding dimension(N) 방향 분해
Execution granularity GroupGEMM tile (e.g., 128×128) 수준에서 재스케줄링
Kernel 실행 단위 하나의 fused kernel 내부에서 통신/연산 thread block 분리
연산 순서 token 소스 rank 기준 정렬 → local token tile 우선 실행

➡️ 이렇게 하여 통신이 완료된 일부 데이터가 즉시 연산에 들어갈 수 있는 path 확보


⚠️ 구조적 trade-off

1. 커널 복잡도 증가 및 구현 부담

  • Fused kernel은 통신/연산 thread block 간의 동기화, shared memory 할당, warp 경쟁 등을 세밀하게 제어해야 함
  • → 일반적인 CUTLASS GEMM이나 기존 NCCL 사용보다 low-level CUDA 설계 부담 증가

2. 하드웨어 종속성 증가

  • Hopper GPU에서는 TMA/Tensor Core 등을 효율적으로 활용하지만, 다른 GPU 아키텍처(Ampere, Volta)에서는 이점이 제한될 수 있음
  • 특히 NVSHMEM은 모든 클러스터 환경에서 지원되는 것이 아니므로 이식성 문제 발생

3. 컴파일러 최적화와의 통합 어려움

  • shared tensor의 분해/재배치는 현재 수작업으로 최적화됨 → 컴파일러가 인지하고 최적화하기 어려움
  • TVM, Triton 등과 연동 시 의존성 추론이 까다로움

🔁 요약

항목 기존 방법 Comet
Overlap 단위 expert/kernel token/tile
실행 방식 병렬 커널 실행 fused kernel 내부 thread block 병렬
장점 구현 단순, 일부 overlap 정밀한 latency hiding (86% 은닉)
단점 overlap 제한적 설계 복잡성, 이식성, 컴파일러 통합성 문제

결론적으로, Comet은 granularity를 token-level 및 tile-level로 끌어올린 덕분에 MoE의 통신 병목을 근본적으로 해소하지만, 이와 맞물려 커널 설계 및 시스템 통합 측면에서 복잡성과 이식성에 대한 trade-off를 감수해야 하는 구조입니다. 이 점을 후속 연구에서는 컴파일러 통합이나 runtime adaptation 방식으로 해결하려는 흐름이 있습니다.

Q : Shared Tensor 기반 dependency resolving 과정에서 M과 N 방향 분해를 사용하는 기준은 무엇이며, 이를 다른 연산자(예: Attention, LayerNorm)에도 일반화하려면 어떤 제약이 발생하는가?

A :

매우 중요한 질문입니다. Comet의 핵심 기술인 Shared Tensor 기반 dependency resolving“어떤 방향으로 tensor를 분해할 것인가?”에 따라 overlapping 성능이 크게 좌우되며, 이를 다른 연산자(예: Attention, LayerNorm 등)로 확장하려면 연산자의 입출력 형태 및 데이터 의존성 구조에 대한 깊은 이해가 필요합니다.


✅ 1. Comet에서 M/N 방향 분해 기준

Shared Tensor의 역할

  • 통신과 연산 사이에서 데이터를 연결하는 producer-consumer buffer
  • overlapping을 위해선 두 연산이 독립적으로 접근 가능한 방향으로 분해되어야 함

MoE Layer 0 (All-to-All → GEMM)

항목 설명
Producer 통신 (All-to-All)
Consumer GEMM1 (FFN의 첫 번째 연산)
Shared Tensor shape ((M \times \text{topk}, N))
분해 방향 M 방향 (토큰 단위)

이유:

  • GEMM은 row-wise 병렬화 → 각 token은 독립적으로 계산 가능
  • 따라서, M 차원을 따라 분해하면 consumer(GEMM)가 부분적으로 먼저 실행 가능
  • 이로 인해 통신된 일부 token만 도착해도 연산이 시작됨

MoE Layer 1 (GEMM2 → Combine)

항목 설명
Producer GEMM2 (FFN의 두 번째 연산)
Consumer Top-k Combine + All-to-All
Shared Tensor shape 동일: ((M \times \text{topk}, N))
분해 방향 N 방향 (임베딩 차원)

이유:

  • Top-k Combine은 token 단위로 M 방향 reduce 연산
  • 이 경우 M 방향은 inter-dependent → 분해 불가
  • 대신, N 차원은 token 내부 차원이므로 독립적 → 연산 일부 완료 시 곧바로 통신 가능

✅ 2. 일반화 시 기준

다른 연산자에 이 방식을 적용하려면 다음 기준을 따져야 합니다:

판단 기준 의미 예시
연산자의 access pattern 어떤 차원을 기준으로 병렬화하는가? Attention → QK^T: head-wise, seq-wise
의존성 방향 어느 차원이 독립적이고, 어느 차원이 종속적인가? LayerNorm → feature(N) 종속, token(M) 독립
output consumption timing 출력이 언제부터 다음 연산에 소비 가능한가? Residual add: 전체 완료 필요 → overlap 어려움

✅ 3. 연산자별 적용 가능성과 제약

🧠 Self-Attention (QK^T, softmax, V 곱)

연산 가능 여부 제약
QK^T (dot-product) 부분적으로 가능 seq_iseq_j 간 의존성 존재 → M과 M의 상호참조
softmax 전체 query sequence 필요 → non-overlappable
V 곱 가능 (tile 단위) softmax(QK)가 완료되어야 시작 가능

➡️ Attention은 token 간 의존성이 있어 M 분해가 어려우며, V 곱 단계만 일부 적용 가능


🧪 LayerNorm

연산 가능 여부 제약
Mean, Var 계산 불가능 (전체 feature 필요) N 차원 reduction 필요
정규화 계산 부분 적용 가능 N 방향 병렬 가능하나, 의미 있는 overlapping은 어려움

➡️ LayerNorm은 output을 생성하기 위해 전체 feature vector(N)가 필요하기 때문에 overlapping 자체가 어려움


✅ 요약 정리

연산자 가능 분해 방향 가능한 overlapping 단계 제약 사항
GEMM M (row-wise) Early computation by tile 없음
Combine (TopK) N (column-wise) Early reduce + send M 종속성
Attention 일부 M, 일부 head V 곱 정도 query-key 간 상호참조
LayerNorm 거의 불가 없음 feature 전체 필요

📌 결론

Comet의 M/N 방향 분해 기준은 다음과 같은 원칙적 판단에 기반합니다:

  1. Consumer 연산자가 어떤 차원을 병렬화하는가?
  2. Shared tensor 내에서 독립적으로 접근 가능한 차원이 존재하는가?
  3. 그 차원으로 tensor를 분해했을 때 결과가 일관성을 유지하는가?

이를 다른 연산자에 일반화하기 위해서는 연산자의 내부 데이터 의존성 모델링tile-level scheduler의 보편화가 필수이며, 이러한 generalization은 앞으로의 컴파일러 기반 접근(TVM, Unity 등)에서 중요한 방향이 될 것입니다.

Q : Adaptive thread block assignment에서 profiling 기반으로 최적 block 비율을 선택하는 방식이 Comet의 실시간성이나 범용성에 미치는 영향을 줄이고자 할 때, 어떠한 online scheduling 기법이 보완적으로 적용될 수 있을까?

A :

이 질문은 Comet의 Adaptive Thread Block Assignment의 핵심 한계를 지적하는 매우 예리한 질문입니다. Comet은 현재 사전 profiling 기반(static metadata)으로 최적 nc:nc (통신:연산 thread block 비율)를 선택하기 때문에, 실시간성(real-time adaptability)범용성(generalizability) 측면에서 제약이 있습니다.


✅ 문제 요약: 왜 문제가 되는가?

항목 문제점
실시간성 결여 새로운 입력 크기(M), top-k, 하드웨어 환경 등 바뀔 때마다 profiling 재수행 필요
범용성 부족 다른 모델이나 dynamic routing 정책에 대해 성능 보장 어려움
모델 일반화 한계 Comet은 MoE의 FFN 구조에 최적화 → 다른 구조(예: depth-wise conv)에서 어려움

✅ 해결 방향: Online Scheduling 기법 적용

아래는 실시간 적응성을 높이기 위해 적용 가능한 Online Scheduling 기법 3가지와 그 특징입니다.


1. 🔁 Reinforcement Learning 기반 Thread Block 튜너 (Runtime Auto-Tuner)

개념 환경 상태(state)를 입력으로 받고, thread block 비율을 action으로 선택 → reward는 latency
입력 현재 입력 token 수(M), expert 수(E), top-k, device ID, previous latency
방법 Bandit / PPO / DQN 기반 정책 학습 (심지어 lightweight model도 가능)
장점 실시간 적응 + 환경 변화에 강건
단점 초기 학습 비용, cold-start 문제

예시:

  • 첫 번째 100 step 동안 탐색 (explore)
  • 이후 빠르게 best config 선택
  • latency에 따라 reward 조정

🔗 관련 사례:

  • AutoTVM (TVM), SparTA (Sparse Transformer Auto-tuner)

2. 📊 Lightweight Heuristic Rule + 피드백 보정 (Latency Ratio Tuning)

| 개념 | 통신 연산 latency를 runtime에서 실측 → 일정 비율 차이만큼 thread block 수 재조정 | | 수식 예시 |

if latency_comm > latency_comp:
    nc += Δ
    np -= Δ
else:
    np += Δ
    nc -= Δ

| 장점 | 매우 경량 + 빠른 적용 가능 | | 단점 | feedback loop tuning이 안정적이지 않을 수 있음

적용 방식:

  • 각 MoE layer 시작 시 latency 예측
  • 커널 launch 직전 shared memory/warp 수 고려해 thread block 조정

🔗 관련 사례:

  • Flux (ByteDance internal)
  • Optimus-CC: Heuristic communication partitioning

3. 📈 Predictive Performance Modeling (Dynamic Cost Model)

개념 학습된 모델을 통해 입력 파라미터 기반으로 block 비율 예측
학습 대상 입력 길이 M, expert 부하 분포 std, top-k, 하드웨어 정보
모델 종류 Linear Regression / Gradient Boosted Trees / Lightweight NN
장점 매우 빠르고 정확도 높음 (한 번만 학습하면 됨)
단점 모델 구조 바뀔 경우 재학습 필요

구현 방식:

  • 다양한 config에 대해 오프라인 latency log 수집
  • 이후 실행 시, 해당 input에 맞는 nc/np 예측

🔗 관련 사례:

  • GShard / Unity compiler의 auto-partitioning cost model

📌 요약 비교표

방식 실시간성 정확도 구현 난이도 보완 가능성
RL 기반 튜닝 ★★★★☆ ★★★★☆ ★★★★☆ 최고 수준 적응성
Heuristic + Feedback ★★★★★ ★★☆☆☆ ★★☆☆☆ 빠른 적용, 간단
Predictive 모델링 ★★★★☆ ★★★★☆ ★★★☆☆ 적당한 트레이드오프

✅ 추천: Hybrid 방식

  • 초기 실행에는 Predictive 모델 + Heuristic 보정
  • 이후 자주 쓰이는 config에는 RL 기반 튜너를 학습

→ 초기 latency는 낮추고, 장기적으로 실행 환경 변화에 적응 가능한 구조 확보


✅ 결론

Comet의 static profiling 기반 thread block assignment는 고정된 config에서는 매우 강력하지만, 동적 환경/모델 구조/입력 변동에는 민감합니다. 이를 극복하기 위해서는:

  1. RL 기반 auto-tuner (강력하지만 복잡)
  2. Heuristic feedback loop (간단하고 빠름)
  3. Cost model 기반 prediction (성능 예측 정확)

을 조합한 실시간 hybrid thread block scheduler가 가장 현실적이며 효과적인 대안입니다. 이러한 방식은 Comet을 더 넓은 범위의 모델과 상황에 적용 가능한 범용 MoE 시스템으로 확장하는 데 핵심 기여를 할 수 있습니다.