텐서 병렬화¶

목차¶

• 개념
• 현재 런타임 지원 범위
• 가중치 샤딩 전략
• 통신 패턴
• 아키텍처 인식 샤드 플랜
• MoE 전문가 병렬화
• 임베딩 및 LM Head 샤딩
• TP 스케줄러 및 실행기
• 설정
• 파이프라인 병렬화와의 조합 (2D PP × TP)

개념¶

텐서 병렬화(TP)는 개별 가중치 행렬을 GPU별로 수평 분할합니다. 각 GPU는 전체 행렬의 샤드 하나씩을 보유합니다. 모든 GPU는 동일한 입력 토큰을 처리하되, 각각 가중치의 서로 다른 부분을 계산합니다. 각 레이어가 끝나면 통신 집합 연산(all-reduce, all-gather)으로 부분 결과를 완전한 활성화로 결합한 뒤 다음 레이어로 진행합니다.

graph LR
    subgraph Input
        X[Input: shape B×S×D]
    end

    subgraph "GPU 0 (Rank 0)"
        W0["W_shard_0<br/>(D × D/2)"]
        P0["Partial output<br/>(B×S×D/2)"]
    end

    subgraph "GPU 1 (Rank 1)"
        W1["W_shard_1<br/>(D × D/2)"]
        P1["Partial output<br/>(B×S×D/2)"]
    end

    X --> W0
    X --> W1
    W0 --> P0
    W1 --> P1
    P0 -->|all-gather| OUT[Full output: B×S×D]
    P1 -->|all-gather| OUT

이것은 각 GPU가 서로 다른 레이어를 실행하는 파이프라인 병렬화와는 다릅니다. TP에서는 모든 GPU가 모든 레이어에 참여합니다.

현재 런타임 지원 범위¶

현재 mlxcel에 들어간 TP 런타임은 단일 프로세스 내 멀티 랭크(in-process multi-rank) 방식입니다. 하나의 mlxcel generate 또는 mlxcel serve 프로세스가 여러 랭크 샤드를 로드하고 한 프로세스 안에서 lockstep으로 실행합니다. 이 경로에서는 --node-role, --peers, 분산 클러스터 설정이 필요하지 않습니다.

현재 tp_size > 1에서 지원하는 런타임:

• dense Llama 계열
• dense Qwen2 / Qwen2.5 계열
• dense Qwen3 계열
• dense Qwen3.5 text 계열
• Gemma 3 텍스트 계열
• Gemma 4 텍스트 계열
• ERNIE 4.5 dense 계열
• Hunyuan v1 Dense 계열

현재 제약:

• --tp-embedding-mode는 replicated여야 합니다
• --tp-lm-head-mode는 replicated여야 합니다
• LoRA adapter는 TP와 함께 지원하지 않습니다
• dense Llama, Qwen2/Qwen2.5, Qwen3, Qwen3.5 text, Gemma 3 text, Gemma 4 text, ERNIE 4.5, Hunyuan v1 Dense에서는 TP server batching이 활성화됩니다
• TP 런타임은 배치된 행을 하나의 네이티브 배치 forward로 디코드하며(BatchedDecodeKind::Native), 서버는 이들에 대해 decode path: native batched forward를 로그로 남깁니다
• loaded-model 폴백이 필요한 Gemma 4 체크포인트(KV 공유 레이어가 있는 E2B 계열)는 server에서 sequential 실행으로 강제됩니다

로컬 CUDA에서 검증한 모델:

• models/llama-3.2-1b-4bit
• models/qwen2.5-0.5b-4bit (tp=2, greedy 24토큰 parity)
• models/qwen2.5-7b-4bit (tp=4)
• models/qwen3-0.6b-4bit
• models/qwen3.5-0.8b-4bit (tp=2)
• models/qwen3.5-4b-4bit (tp=4)
• models/qwen3.5-9b-4bit (tp=2, tp=4, stepwise CUDA parity)
• models/qwen3.5-27b-4bit (tp=4, stepwise CUDA parity)
• models/gemma3-1b-4bit
• models/ernie-4.5-0.3b-4bit
• models/hunyuan-1.8b-4bit

현재 비활성 상태:

• TP + LoRA

가중치 샤딩 전략¶

src/distributed/tensor_parallel/shard_strategy.rs에 정의되어 있습니다:

Column Parallel¶

가중치 행렬의 **출력 차원(열)**을 랭크에 걸쳐 분할합니다. 각 랭크는 출력 피처의 일부를 계산합니다.

block-beta
    columns 3
    W["W: [D_in × D_out]"]:3
    R0["Rank 0\nW[:, 0 : D_out/N]"] R1["Rank 1\nW[:, D_out/N : 2·D_out/N]"] Rn["...\nRank N-1"]

사용처: 어텐션 Q/K/V 프로젝션, FFN gate 및 up 프로젝션.

Column-parallel 연산 후 완전한 출력이 즉시 필요하면 all-gather를 수행합니다. 다음 연산이 row-parallel이라면 reduce-scatter를 뒤로 미룰 수 있습니다.

Row Parallel¶

가중치 행렬의 **입력 차원(행)**을 랭크에 걸쳐 분할합니다. 각 랭크는 부분 합에 기여하는 부분 활성화를 보유하고 이를 all-reduce로 합산합니다.

block-beta
    columns 3
    W["W: [D_in × D_out]"]:3
    R0["Rank 0\nW[0 : D_in/N, :]"] R1["Rank 1\nW[D_in/N : 2·D_in/N, :]"] Rn["...\nRank N-1"]

사용처: 어텐션 출력 프로젝션, FFN down 프로젝션.

Row-parallel 행렬곱 후에는 all-reduce(합산)를 적용해 올바른 결합 출력을 만듭니다.

Vocab Parallel¶

임베딩 테이블과 LM head의 어휘 차원을 랭크에 걸쳐 분할합니다. 각 랭크는 토큰 어휘의 연속적인 슬라이스를 담당합니다.

사용처: 임베딩 조회, LM head 출력 로짓.

Replicated¶

가중치가 모든 랭크에 완전히 복제됩니다. 샤딩이나 통신이 필요하지 않습니다. 레이어 정규화 파라미터 및 바이어스 벡터와 같은 소형 텐서에 사용됩니다.

통신 패턴¶

src/distributed/tensor_parallel/collective.rs에 정의되어 있습니다:

graph TD
    subgraph AllReduce["All-Reduce (Ring)"]
        AR0[Rank 0<br/>partial sum] -->|Reduce-scatter| RS[Each rank<br/>holds one reduced chunk]
        AR1[Rank 1<br/>partial sum] --> RS
        AR2[Rank 2<br/>partial sum] --> RS
        RS -->|All-gather| AG[All ranks<br/>hold full reduced tensor]
    end

    subgraph AllGather["All-Gather"]
        AG0[Rank 0 shard] -->|ring forwarding| FG[All ranks<br/>hold full tensor]
        AG1[Rank 1 shard] --> FG
    end

    subgraph ReduceScatter["Reduce-Scatter"]
        RS0[Rank 0 full tensor] -->|accumulate + scatter| RSO[Each rank<br/>holds 1/N reduced shard]
        RS1[Rank 1 full tensor] --> RSO
    end

All-reduce (ring):

TP에서 가장 일반적인 집합 연산입니다. Row-parallel 행렬곱 후 모든 랭크에 걸쳐 부분 결과를 합산하는 데 사용됩니다. 두 단계의 링 알고리즘을 사용합니다: reduce-scatter에 이어 all-gather. 랭크당 총 데이터 전송량: 2 × (N-1)/N × tensor_size.

All-gather:

샤딩된 출력을 완전한 텐서로 조합하는 데 사용됩니다. 각 랭크가 자신의 샤드를 기여하며 모든 랭크가 최종적으로 전체 텐서를 보유합니다.

Reduce-scatter:

Row-parallel 연산의 결과가 곧바로 column-parallel 연산으로 이어지는 최적화된 파이프라인에서 사용됩니다. 출력의 샤드만 필요한 경우 all-gather 단계를 피할 수 있습니다.

모든 집합 연산은 링 토폴로지를 사용합니다. 랭크 r은 각 단계에서 (r+1) % N으로 전송하고 (r+N-1) % N으로부터 수신합니다. N-1 단계에 걸쳐 모든 랭크는 다른 모든 랭크와 정확히 한 번씩 통신합니다.

아키텍처 인식 샤드 플랜¶

mlxcel은 src/distributed/tensor_parallel/plan_generator.rs를 통해 아키텍처별 샤드 플랜을 생성합니다. 플랜은 어느 가중치 텐서에 어떤 ShardStrategy를 적용하고 어떤 CommPattern이 이어지는지를 인코딩합니다.

표준 트랜스포머 (Llama, Qwen, Gemma 등)¶

GQA 인식 어텐션 샤딩¶

Grouped Query Attention(GQA)을 사용하는 모델은 KV 헤드 수가 Q 헤드 수보다 훨씬 적은 경우가 많습니다. 샤드 플랜은 이를 고려합니다:

• num_kv_heads >= tp_size인 경우: K/V 프로젝션이 일반적으로 샤딩됩니다.
• num_kv_heads < tp_size인 경우: K/V 프로젝션을 모든 랭크에 복제해 KV 헤드가 랭크당 하나 미만으로 단편화되는 일을 막습니다.

이는 src/distributed/tensor_parallel/parallel_attention.rs에서 계산됩니다.

graph LR
    subgraph "GQA Sharding (tp_size=4, 32 Q heads, 8 KV heads)"
        Q["Q heads: 32<br/>2 per rank"]
        K["K heads: 8<br/>2 per rank"]
        V["V heads: 8<br/>2 per rank"]
    end

    subgraph "GQA Replication (tp_size=8, 4 KV heads)"
        Q2["Q heads: 32<br/>4 per rank"]
        K2["K heads: 4<br/>replicated (< 1 per rank)"]
        V2["V heads: 4<br/>replicated (< 1 per rank)"]
    end

SSM/Mamba¶

Mamba 및 Mamba2 모델의 경우, 합성곱 및 SSM 상태 행렬이 채널 차원을 따라 샤딩되며 동일한 column/row parallel 패턴을 따릅니다.

MoE 전문가 병렬화¶

아래 mode는 shard-plan design을 설명합니다. 현재 production tp_size > 1 runtime은 위에 나열한 dense family로 제한되므로 이 mode들은 활성화되지 않습니다.

Mixture-of-Experts 모델(Mixtral, DeepSeek V2/V3, Qwen MoE 등)의 경우, mlxcel은 --tp-moe-mode로 설정 가능한 두 가지 MoE 샤딩 모드를 지원합니다:

Expert Parallel (기본값)¶

전체 전문가가 라운드 로빈으로 랭크에 할당됩니다: 전문가 i는 전적으로 랭크 i % tp_size에 위치합니다. 각 랭크의 전문가 가중치는 모두 로컬에 있으므로 전문가 내부 통신이 필요하지 않습니다.

flowchart LR
    subgraph "Rank 0"
        E0["Expert 0"]
        E2["Expert 2"]
    end
    subgraph "Rank 1"
        E1["Expert 1"]
        E3["Expert 3"]
    end
    Router --> E0 & E1 & E2 & E3

num_experts >= tp_size일 때 가장 효과적입니다. 라우터는 여전히 모든 랭크에서 실행되지만 각 랭크는 로컬에 할당된 전문가만 실행합니다.

Within-Expert Parallel¶

각 전문가의 내부 FFN 가중치는 일반 조밀 FFN과 동일한 column-parallel gate/up, row-parallel down 프로젝션 패턴으로 모든 랭크에 걸쳐 샤딩됩니다.

num_experts < tp_size일 때 가장 효과적입니다(전문가 수가 적고 GPU가 많은 드문 경우).

src/distributed/tensor_parallel/parallel_moe.rs에 정의되어 있습니다.

임베딩 및 LM Head 샤딩¶

Configuration type은 이 strategy를 모델링하지만 런타임은 embedding과 LM-head mode 모두 replicated를 요구합니다. 아래 sharded 대안은 production에서 활성화되지 않습니다.

--tp-embedding-mode 및 --tp-lm-head-mode로 설정 가능한 두 가지 모드:

어휘가 매우 큰(10만 토큰 초과) 모델에서 메모리가 주된 제약이라면 vocab-parallel 모드를 권장합니다.

TP 스케줄러 및 실행기¶

TP 실행은 두 가지 컴포넌트로 조율됩니다:

TP 스케줄러 (랭크 0 전용)¶

src/distributed/tensor_parallel/tp_scheduler.rs

스케줄러는 랭크 0에서만 실행됩니다. 스케줄러의 역할:

• 새 시퀀스를 수용하고 프리필 큐를 관리합니다.
• 배치를 구성합니다 (각 단계에서 프리필 또는 디코드할 시퀀스를 선택).
• 각 추론 단계 전에 다른 모든 랭크에 StepDecision 메시지를 브로드캐스트합니다.
• 토큰 샘플링을 조율합니다: 랭크 0이 다음 토큰을 샘플링하고 결과를 브로드캐스트합니다.
• 연속 배치: GPU 활용률을 극대화하기 위해 프리필과 디코드 단계를 교차합니다.

다른 랭크는 독립적인 스케줄링 결정을 내리지 않습니다. 모든 랭크는 각 단계를 실행하기 전에 스케줄러의 브로드캐스트를 기다립니다.

시퀀스 생명주기:

stateDiagram-v2
    [*] --> Waiting: request submitted
    Waiting --> PendingPrefill: admitted to batch
    PendingPrefill --> Decoding: prefill complete
    Decoding --> Completed: EOS or max_tokens
    Decoding --> Evicted: preempted
    Evicted --> Waiting: re-admitted
    Completed --> [*]

TP 실행기 (전체 랭크)¶

src/distributed/tensor_parallel/tp_executor.rs

모든 랭크(랭크 0 포함)는 TPExecutor를 실행합니다. 실행기의 역할:

• 랭크 0 브로드캐스트에서 StepDecision 값을 수신합니다.
• 모든 랭크에서 동일한 시퀀스 상태를 유지합니다 (잠금 단계 실행을 위해).
• 각 단계 전에 배리어 동기화에 참여합니다.
• 디코드 단계에서 랭크 0이 다음 토큰을 샘플링하고 다른 랭크는 브로드캐스트로 샘플링된 토큰을 받습니다.
• 랭크 헬스 상태를 보고하며 오류가 발생하면 전파해 전체 TP 그룹을 중지합니다.

sequenceDiagram
    participant S as Scheduler (rank 0)
    participant E0 as Executor rank 0
    participant E1 as Executor rank 1
    participant E2 as Executor rank 2

    S->>E0: broadcast StepDecision
    S->>E1: broadcast StepDecision
    S->>E2: broadcast StepDecision
    Note over E0,E2: All ranks execute forward pass in lockstep
    E0->>E1: all-reduce (per layer)
    E1->>E2: all-reduce (per layer)
    E2->>E0: all-reduce (per layer)
    Note over E0: Rank 0 samples next token
    E0->>E1: broadcast sampled token
    E0->>E2: broadcast sampled token
    Note over S: Update sequence state

설정¶

현재 런타임에서는 임베딩 관련 두 플래그를 모두 replicated로 두어야 합니다.

예시: 한 머신에서 2랭크 TP 텍스트 생성:

mlxcel generate \
  --model ./models/llama-3.2-1b-4bit \
  --prompt "텐서 병렬화를 한 문장으로 설명해 주세요." \
  --tp-size 2

예시: 한 머신에서 2랭크 TP 서버:

mlxcel serve \
  --model ./models/qwen3-0.6b-4bit \
  --tp-size 2

파이프라인 병렬화와의 조합 (2D PP × TP)¶

텐서 병렬화는 파이프라인 병렬화와 조합해 2D (pp_stage, tp_rank) 메시를 구성할 수 있습니다. PP-only 토폴로지에서 한 스테이지가 여전히 한 노드의 메모리를 초과할 때 사용합니다. 대표 사례는 Llama-70B 급 모델을 중형 GPU 4대 위에 2 PP 스테이지 × 2 TP 랭크 메시로 서빙하는 경우입니다.

참고: 동일한 메시를 PP 관점에서 서술한 내용은 파이프라인 병렬화 → 텐서 병렬화와의 조합 (2D PP × TP) 에서 확인할 수 있습니다.

스테이지 내부의 TP 동작¶

한 파이프라인 스테이지 내부의 TP 컬렉티브 레이어는 위의 1D 설명과 동일합니다. 랭크 0이 admission과 step 브로드캐스트를 주도하고 해당 스테이지의 모든 랭크가 그 스테이지 레이어에 대한 링 토폴로지 all-reduce / reduce-scatter / all-gather에 참여합니다. TP 피어 검색은 스테이지 단위로 한정됩니다: all_reduce_sum, all_gather, reduce_scatter는 같은 pp_stage를 공유하는 TP 랭크들하고만 통신합니다.

스테이지 경계의 PP 동작¶

PP 스테이지 경계에서, 각 tp_rank는 자기 활성화 샤드를 다음 스테이지의 같은 tp_rank를 소유하는 노드에 전달합니다. 즉, 스테이지 0 랭크 0은 스테이지 1 랭크 0과 통신하고 스테이지 0 랭크 1은 스테이지 1 랭크 1과 통신합니다. TP 샤드 레이아웃이 스테이지 간에 균일하므로 경계에서 크로스-랭크 교환은 발생하지 않습니다.

트랜스포트 트래픽 클래스¶

2D 라우팅 레이어는 메시의 모든 노드에 걸쳐 단일 연결 풀을 공유하지만 메시지마다 올바른 핸들러가 처리하도록 두 트래픽 클래스를 구분합니다:

두 클래스 모두 PipelineTensorParallelRouter를 통해 선택됩니다. 각 클래스의 전체 피어 집합은 컬렉티브 레이어가 사용하는 레지스트리 조회 (NodeRegistry::nodes_at_stage)와 활성화 전송 레이어(NodeRegistry::nodes_at_rank)에서 얻습니다.

CLI 진입점¶

mlxcel generate \
  --model ./models/llama-70b-4bit \
  --pp-size 2 \
  --tp-size 2 \
  -p "Hello" -n 50

src/commands/generate.rs의 검증기는 이 조합을 받아들이고 일관성 없는 토폴로지에 대해 가독성 있는 에러를 방출합니다.

클러스터 TOML과 캐시 조정¶

클러스터 TOML은 pipeline_tensor_parallel 역할(pp_tp, 2d 별칭)을 사용하며 정확히 pp_size × tp_size개의 해당 노드를 요구합니다. 각 노드는 stage와 rank를 모두 지정해야 합니다. 2×2 TOML 예시와 캐시 조정 규칙 (coordinated_2d_admission, broadcast_2d_eviction)의 전체 내용은 파이프라인 병렬화 → 텐서 병렬화와의 조합 (2D PP × TP) 를 참조하십시오.

제약¶

• 스테이지별 TP 폭은 균일: 이종 폭(예: 스테이지 0이 tp=4, 스테이지 1이 tp=2)은 지원하지 않습니다.
• 3D 병렬화는 범위 밖: Expert 병렬화(PP × TP × EP)는 지원하지 않습니다.