본문으로 건너뛰기

[SGLang] Elastic Expert Parallelism: 동적 전문가 스케일링

들어가며

Expert Parallel 환경에서 GPU 장애가 발생하면 해당 GPU의 전문가들이 사라져 모델이 동작할 수 없다. Elastic EP는 이 문제를 해결하기 위해 전문가의 백업을 관리하고, 활성 랭크 상태를 추적하여 동적으로 전문가를 재배치하는 메커니즘이다.

관련 소스 경로:

구조도

┌─────────────────────────────────────────────┐
│            ElasticEPStateManager            │
│  (싱글턴, 클러스터 전체 활성 랭크 추적)       │
├─────────────────────────────────────────────┤
│  active_ranks: [1, 1, 1, 0, 1, 1, 1, 1]   │
│  last_active_ranks: (이전 스냅샷)            │
│  active_ranks_cpu: (CPU 복사본)             │
└──────────────┬──────────────────────────────┘
               │
               ▼
┌─────────────────────────────────────────────┐
│         ExpertBackupManager                 │
│  (별도 프로세스, DRAM에 전문가 가중치 보관)    │
├─────────────────────────────────────────────┤
│  continuous_buffer: [전문가 가중치 연속 버퍼]  │
│  weight_pointer_map: {이름 → 포인터 정보}    │
│  transfer_engine: Mooncake RDMA 전송        │
└─────────────────────────────────────────────┘

핵심 코드 분석

1. ElasticEPState: 활성 랭크 상태

ElasticEPState 데이터클래스는 클러스터 내 각 랭크의 활성/비활성 상태를 추적한다.

@dataclass
class ElasticEPState:
    active_ranks: Optional[torch.Tensor]
    last_active_ranks: Optional[torch.Tensor]
    active_ranks_cpu: Optional[torch.Tensor]

    def is_active_equal_last(self) -> bool:
        return torch.equal(self.active_ranks, self.last_active_ranks)

    def sync_active_to_cpu(self):
        if self.active_ranks is not None:
            self.active_ranks_cpu = self.active_ranks.detach().cpu().clone()

    def snapshot_active_to_last(self):
        if self.active_ranks is not None:
            self.last_active_ranks = self.active_ranks.clone()

active_ranks는 [1, 1, 1, 0, 1, ...] 형태의 텐서로, 0은 장애가 발생한 랭크를 의미한다. is_active_equal_last로 상태 변화를 감지하여 재배치 필요 여부를 판단한다.

2. ElasticEPStateManager: 싱글턴 관리

ElasticEPStateManager는 싱글턴 패턴으로 전역 상태를 관리한다.

class ElasticEPStateManager:
    _instance: Optional[ElasticEPState] = None

    @classmethod
    def init(cls, server_args: ServerArgs):
        if cls._instance is not None:
            return cls._instance
        if server_args.elastic_ep_backend is not None:
            cls._instance = cls._build_state(ep_size=None, device=None)
        return cls._instance

    @classmethod
    def healthy_rank_state(cls, *, ep_size=None, device=None) -> torch.Tensor:
        size = ep_size if ep_size is not None else torch.distributed.get_world_size()
        dev = device if device is not None else cls._select_device()
        return torch.ones(size, dtype=torch.int32, device=dev)

초기 상태는 모든 랭크가 1(건강)이다. elastic_ep_backend 설정이 있을 때만 활성화된다.

3. ExpertBackupManager: DRAM 가중치 백업

ExpertBackupManager는 별도 프로세스에서 실행되며, 디스크에서 전문가 가중치를 로드하여 CPU 메모리에 연속 버퍼로 보관한다.

class ExpertBackupManager:
    def backup_weights_from_disk(self):
        loader = get_model_loader(load_config, self.model_config)
        with set_default_torch_dtype(self.model_config.dtype):
            iter = loader._get_weights_iterator(...)
            for name, weight in iter:
                expert_id = extract_expert_id(name)
                if expert_id < self.idmx and expert_id >= self.idmn:
                    weight_info_dict[name] = {
                        "weight": weight, "numel": numel,
                        "shape": weight.shape, "dtype": weight.dtype, ...
                    }
                    total_bytes += byte_size

각 노드는 자신에게 할당된 전문가 범위(idmn ~ idmx)의 가중치만 백업한다.

4. 연속 버퍼 구성

백업 가중치는 하나의 연속 버퍼에 패킹되어 RDMA 전송 효율을 높인다.

self.continuous_buffer = torch.empty(total_bytes, dtype=torch.uint8, device="cpu")
buffer_base_ptr = self.continuous_buffer.data_ptr()

for name in sorted(weight_info_dict.keys()):
    weight_flat = weight.flatten().contiguous()
    weight_bytes = weight_flat.view(torch.uint8)
    self.continuous_buffer[start_byte:end_byte].copy_(weight_bytes)
    self.weight_pointer_map[name] = {
        "weight_ptr": buffer_base_ptr + current_byte_offset,
        "shape": weight_info["shape"], "dtype": weight_info["dtype"], ...
    }

weight_pointer_map은 각 가중치의 메모리 포인터, shape, dtype 정보를 저장한다.

5. Mooncake RDMA 전송 서버

백업 매니저는 Mooncake Transfer Engine을 사용하여 RDMA 기반 전송 서버를 시작한다.

def start_transfer_server(self):
    self.transfer_engine = get_mooncake_transfer_engine()
    self.session_id = self.transfer_engine.session_id
    server_ptr = self.continuous_buffer.data_ptr()
    server_len = self.continuous_buffer.numel() * self.continuous_buffer.element_size()
    ret_value = self.transfer_engine.engine.register_memory(server_ptr, server_len)

클라이언트(각 GPU 워커)는 장애 발생 시 이 서버에서 필요한 전문가 가중치를 RDMA로 빠르게 가져올 수 있다.

6. EPLB 연동: 탄력적 재균형

Elastic EP는 EPLB의 elasticity_aware 알고리즘과 연동된다. 비활성 랭크를 고려하여 전문가를 재배치한다.

# eplb_algorithms/__init__.py
if algorithm in [EplbAlgorithm.elasticity_aware, ...]:
    return elasticity_aware.rebalance_experts(
        weight=tokens_per_expert.sum(dim=0),
        active_ranks=(
            ElasticEPStateManager.instance().active_ranks
            if ElasticEPStateManager.instance() is not None
            else ElasticEPStateManager.healthy_rank_state()
        ),
    )

장애 복구 흐름

1. GPU 3 장애 감지
   active_ranks: [1,1,1,1,1,1,1,1] → [1,1,1,0,1,1,1,1]

2. is_active_equal_last() → False (변화 감지)

3. EPLB elasticity_aware 알고리즘 실행
   → GPU 3의 전문가를 다른 GPU에 복제 배치

4. ExpertBackupManager에서 가중치 RDMA 전송

5. 새 expert_location_metadata 적용

설계 근거

설계 결정 이유
별도 프로세스 백업 GPU 메모리 부담 없이 CPU DRAM에 보관
연속 버퍼 RDMA 전송 시 메모리 등록/전송 효율 극대화
싱글턴 상태 관리 클러스터 전체에서 일관된 랭크 상태 유지
Mooncake 전송 기존 인프라 재사용, 노드 간 고속 전송

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글