전문가 병렬

Expert parallelism은 mixture-of-experts (MoE) 모델에 사용되는 병렬화 전략입니다. 각 expert의 feedforward 레이어는 서로 다른 하드웨어 가속기에 배치됩니다. 라우터가 토큰을 적절한 expert로 전달하고 결과를 수집합니다. 각 토큰이 소수의 expert만 활성화하기 때문에, 계산 비용 증가 없이 모델의 매개변수 수를 훨씬 더 크게 확장할 수 있습니다.

DistributedConfig

DistributedConfig 클래스와 ep_size 인수를 사용하여 전문가 병렬화를 활성화합니다. 대부분의 MoE 모델은 기본적으로 token dispatch를 사용하므로, ep_size는 tp_size와 독립적으로 설정할 수 있습니다. Llama 4와 Gemma 4와 같은 일부 모델은 여전히 기본적으로 masking과 all-reduce(model.ep_plan의 "ep_router" 및 "moe_tp_experts")를 사용합니다. ep_plan 오버라이드가 있는 모든 모델에서도 masking을 사용할 수 있으며, 전문가 그룹의 모든 rank가 동일한 token을 수신하도록 ep_size=tp_size가 필요합니다.

import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-120b",
    distributed_config=distributed_config,
)

각 MoE 모델은 config에서 두 가지 플랜을 정의합니다: dense 모듈을 위한 base_model_tp_plan과 experts를 위한 base_model_ep_plan입니다. 이들은 로드된 모델에서 model.tp_plan과 model.ep_plan으로 노출됩니다. tp_size > 1이고 ep_size > 1인 경우, 두 플랜이 모두 적용됩니다: tensor parallel 플랜은 attention과 dense MLP를 샤딩하고, expert parallel 플랜은 experts를 샤딩합니다. 동일한 모듈에 대해 EP 규칙이 TP 규칙보다 우선하므로, expert 가중치는 EP 플랜에 의해 한 번만 샤딩됩니다. EP 플랜에서 GroupedGemmParallel 스타일은 expert 차원 방향으로 expert 가중치를 분할하여 각 rank가 자신의 로컬 experts만 로드하도록 하고, ep_router는 all-reduce가 expert 출력을 결합하기 전에 다른 rank에 존재하는 experts를 마스킹합니다.

tp_plan은 tp_size > 1인 경우에만 적용되며, ep_plan은 ep_size > 1인 경우에만 적용됩니다. TP가 활성화되고 EP가 비활성화된 상태에서는 전문가 규칙을 포함하여 전체 TP 플랜이 적용됩니다.

ep_plan의 expert forward rule은 토큰이 expert에 도달하는 방식을 선택합니다:

ep_size 없이 enable_expert_parallel=True를 사용하는 것은 비추천되며 FutureWarning를 출력합니다. ep_size=tp_size, tp_plan={}와 마찬가지로 전문가만 샤딩합니다.

torchrun로 추론 스크립트를 실행합니다. 프로세스 수는 tp_size * fsdp_size * pp_size와 같아야 하며, ep_size는 전문가 수를 나누어 떨어뜨려야 합니다.

torchrun --nproc-per-node 8 your_script.py

플랜 오버라이드

tp_plan={...} 또는 ep_plan={...}을 DistributedConfig에 전달하여 사전 정의된 플랜의 개별 규칙을 재정의할 수 있습니다. 지정되지 않은 규칙은 유지되며, 병합된 플랜은 모델에 저장됩니다. 각 키는 모듈, 매개변수 또는 기존 플랜 항목과 일치해야 하며, 그렇지 않으면 샤딩이 시작되기 전에 ValueError가 발생합니다. 로드된 모델에서 확인되는 전체 경로를 사용해야 하므로, causal LM의 경우 model.layers.*, 해당 base model의 경우 layers.*를 사용합니다.

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    tp_plan={"model.layers.*.self_attn.q_proj": "colwise_rep"},
    ep_plan={"model.layers.*.mlp.experts.down_proj": "grouped_gemm"},
)

플랜을 제공해도 병렬 크기를 추론하지 않습니다: tp_size와 ep_size를 명시적으로 설정하세요.

대부분의 MoE 모델은 기본적으로 "ep_dispatch_experts"를 사용합니다. 대신 masking과 all-reduce를 사용하려면 ep_size=tp_size를 설정하고 라우터와 전문가의 forward 규칙을 모두 오버라이드해야 합니다(모델에 따라 라우터 모듈 이름이 다르며, 예를 들어 gpt-oss에서는 mlp.router입니다):

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)

반대로, masking을 사용하는 플랜을 가진 모델의 expert forward 규칙을 "ep_dispatch_experts"로 오버라이드하여 token dispatch를 사용할 수 있습니다. 이 경우 dispatch는 각 expert의 owner를 찾기 위해 전역 expert id가 필요하므로 router 규칙은 무시됩니다.

토큰 디스패치

token dispatch를 사용하면 각 rank는 배치의 자기 몫만 학습합니다. 모든 MoE 레이어에서 rank는 자신의 토큰을 라우팅하고, all-to-all을 통해 각 (token, expert) 쌍을 해당 expert를 소유한 rank로 전송한 후, 수신한 데이터에 대해 로컬 expert를 실행합니다. 이후 두 번째 all-to-all로 결과를 받아와서 라우팅 가중치와 결합합니다. 이 과정에서 전송되는 것은 라우팅된 activations와 expert 출력뿐이며, 어떤 rank도 자신이 소유하지 않은 토큰에 대해 expert를 계산하지 않습니다.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=1,
    fsdp_size=8,
    ep_size=4,
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

tp_size=1인 경우, ep_size는 fsdp_size와 전문가 수를 나누어야 하며, 어텐션은 num_key_value_heads에 의해 제한되지 않습니다. 모델의 나머지 부분에 대해:

• experts 외부의 파라미터는 fsdp에 걸쳐 FSDP2로 샤딩되며, 이를 통해 그라디언트가 감소합니다.
• 전문가들은 ep에 걸쳐 샤드되며, efsdp_size = fsdp_size * tp_size // ep_size가 1보다 클 경우 efsdp에 걸쳐 추가로 FSDP 샤딩됩니다. 이들은 항상 FSDP로 래핑되므로 fsdp_mixed_precision와 fsdp_cpu_offload도 적용되며, save_pretrained()는 다른 매개변수와 마찬가지로 이들을 수집합니다.
• Expert parallel 그룹은 ep_size / tp_size개의 배치로 구성되므로, 각 expert의 그래디언트는 해당 수만큼의 배치에 대한 합입니다. efsdp 리덕션은 그룹 크기 대신 fsdp_size로 나누며, 이는 FSDP2가 dense 모듈에 대해 사용하는 배치별 평균과 동일한 값을 제공합니다.
• 각 로컬 전문가(expert)는 레이어당 하나의 제로 패딩(zero pad) 행을 처리합니다. 토큰을 받지 못한 전문가를 가진 랭크(rank)도 reverse all-to-all 및 전문가(expert) 그래디언트 감소(reduction)에 참여합니다.

크기가 결합되는 방식

• tp_size * fsdp_size는 프로세스 수입니다. ep_size는 아무것도 추가하지 않습니다: 전문가 가중치에 대해서만 동일한 랭크를 다시 그룹화합니다.
• ep_size는 expert 목록을 ep_size개의 블록으로 나눕니다. 각 rank는 num_experts / ep_size개의 expert를 계산하며, ep_size개의 연속된 rank가 하나의 완전한 expert 세트를 보유합니다. 이 rank의 집합은 all-to-all이 실행되는 그룹입니다.
• efsdp_size = fsdp_size * tp_size / ep_size는 expert set의 완전한 사본이 몇 개 존재하는지를 나타냅니다. 서로 다른 사본에서 동일한 위치에 있는 rank들은 FSDP가 dense 모듈에 대해 수행하는 방식과 유사하게, 메모리를 위해 해당 expert를 샤딩하고 그라디언트를 평균화합니다.
• 랭크가 보유하는 배치의 구성은 fsdp에만 의존합니다. 연속된 랭크들은 TP 그룹을 형성하며 동일한 배치를 받습니다. fsdp_size개의 그룹은 서로 다른 배치를 받습니다.

8개의 프로세스, tp_size=2, fsdp_size=4, 8개의 전문가:

rank         0      1       2      3       4      5       6      7
batch       [====B0====]   [====B1====]   [====B2====]   [====B3====]     one batch per TP pair
tp           0      1       0      1       0      1       0      1

ep_size=2   E0-3   E4-7    E0-3   E4-7    E0-3   E4-7    E0-3   E4-7     group = a TP pair,   efsdp_size=4
ep_size=4   E0E1   E2E3    E4E5   E6E7    E0E1   E2E3    E4E5   E6E7     group = two pairs,   efsdp_size=2
ep_size=8   E0     E1      E2     E3      E4     E5      E6     E7       group = all ranks,   efsdp_size=1

그림에서 두 가지 수치가 도출됩니다:

• EP 그룹 내에서 각 토큰은 해당 배치 보유 쌍의 각 랭크마다 한 번씩, 총 tp_size번 존재합니다. 이는 ep_size에 의존하지 않습니다.
• EP 그룹은 ep_size / tp_size개의 서로 다른 배치(batch)를 보유합니다. 이는 각 전문가(expert)의 그래디언트가 합산되는 횟수입니다.

텐서 병렬성

tp_size > 1을 설정하면 TP 플랜에 따라 dense 모듈을 샤딩하고, experts는 dispatch를 사용합니다. 8개 프로세스에서:

distributed_config = DistributedConfig(
    tp_size=2,
    fsdp_size=4,
    ep_size=4,
)

TP rank의 각 쌍은 동일한 배치(batch)를 수신합니다. 이는 텐서 병렬성(tensor parallelism)이 쌍 내부에서 활성화(activations)를 복제하기 때문입니다. 두 rank가 모두 토큰을 전부 디스패치(dispatch)했다면, 특정 expert를 소유한 rank는 모든 토큰을 두 번 수신하고 두 번 계산하게 되며, 해당 weight gradient가 두 배가 됩니다. Expert는 하나의 rank에 온전히 존재하므로 중복을 weight로 분할할 수 없으며, 소유자는 일반적으로 다른 rank이므로 masking처럼 소유권으로 해결할 수도 없습니다. 따라서 이 쌍은 행(rows)을 분할합니다. 각 TP rank는 토큰의 서로 겹치지 않는 1 / tp_size만 디스패치하고, 결과는 보낸 rank로 돌아오며, zero-padded된 절반에 대해 쌍을 대상으로 all-reduce를 수행하여 다음 레이어가 기대하는 복제된 출력을 복원합니다. 이 분할은 tp_size에 기반하며 ep_size에 기반하지 않습니다. 배치를 보유하여 보낼 수 있는 rank가 해당 rank이기 때문입니다. Expert 그룹은 4개의 rank에 걸쳐 있으며, 각 expert는 efsdp_size = 2개의 rank에 걸쳐 FSDP-sharding됩니다. 한편 trunk의 FSDP 그룹은 4개의 rank에 걸쳐 있습니다. 모델의 일반적인 TP 제약 조건, 예를 들어 attention-head의 나누어 떨어지는 조건 등은 dense 모듈에 여전히 적용됩니다. 토큰 슬라이스는 균등하지 않거나 비어 있을 수 있으며, 단일 토큰 디코딩 중에도 마찬가지입니다.

Token dispatch는 아직 pipeline parallelism과 함께 사용할 수 없으므로 pp_size=1을 사용하세요(아직 테스트되지 않음) 이 구성들은 각각 8개의 GPU를 사용합니다:

Trainer는 아직 토큰 디스패치를 지원하지 않습니다. 배치 및 토큰 카운팅은 all-reduce 레이아웃을 가정하며, 이 레이아웃에서는 TP 그룹의 랭크들이 배치를 공유하고 fsdp_size개의 데이터 병렬 샤드가 존재합니다. Trainer의 디스패치 지원은 후속 업데이트에서 제공됩니다.

FSDP2와 결합

Tensor parallelism과 expert parallelism은 가중치를 tp에 걸쳐 샤딩하지만, 옵티마이저 상태와 규칙이 없는 모듈은 여전히 그룹의 모든 랭크에 복제되므로 학습할 수 있는 모델 크기가 제한됩니다. fsdp_size를 사용하여 두 번째 메시 차원에 FSDP2를 추가하세요. 마스킹과 all-reduce를 사용하여 ep_size=tp_size를 유지하고 `ep_plan={“layers..mlp.gate”: “ep_router”, “layers..mlp.experts”: “moe_tp_experts”}“`를 전달하세요.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,  # expert parallel size, must match tp_size with masking and all-reduce
    fsdp_size=2,  # data parallel shards
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

모델을 평소와 같이 로드한 후 Trainer로 학습합니다. 두 메시 전체에 걸린 그래디언트 노름을 계산하고 각 메시에 고유한 옵티마이저 파라미터 그룹을 할당합니다. save_model()은 샤드된 가중치를 일반 체크포인트로 수집합니다. 이를 위해 accelerate>=1.12가 필요하며, Trainer가 tp_size와 fsdp_size를 ~Accelerate.ParallelismConfig에 반영할 수 있어야 합니다.

아래 표는 8xH100 GPU에서 EP 전용 학습과 2D EP+FSDP2를 비교합니다. 워크로드는 시퀀스 길이 2048에서 bf16로 Qwen3-30B-A3B를 완전 파인튜닝하는 것입니다. FSDP 샤드가 많을수록 피크 메모리가 줄고, FSDP2가 fsdp에 걸쳐 전문가를 올게더와 리듀스스캐터하기 때문에 토큰/초가 일부 감소합니다.

로드 시 샤딩된 모델에 대해 체크포인트에서 재개하는 기능은 아직 지원되지 않으므로, Trainer는 해당 모델에 대해 save_only_model=True 또는 save_strategy="no"만 허용합니다.

API 참조