[sglang] Qwen3.5 MoE 모델 로딩 및 Mamba 캐시 PP 모드 샤딩 수정

PR 링크: sgl-project/sglang#21448 상태: Merged | 변경: +78 / -8

들어가며

Pipeline Parallelism(PP)은 모델의 레이어를 여러 GPU에 분산하여 처리하는 병렬화 기법입니다. 각 GPU는 start_layer부터 end_layer까지의 레이어만 로딩해야 합니다. Qwen3.5 MoE 모델과 Mamba 하이브리드 모델에서 이 범위 필터링이 누락되어, PP 모드에서 로딩 오류가 발생하는 문제를 수정한 PR입니다.

핵심 코드 분석

1. Qwen3.5 MoE weight 로딩 필터링

Before (qwen3_5.py):

def load_weights(self, weights):
    for name, loaded_weight in weights:
        name = name.replace(r"model.language_model.", r"model.")
        if ".self_attn." in name:
            name = name.replace(".self_attn", "")
        # 모든 레이어의 weight를 로딩 시도 -> PP 모드에서 오류
        for param_name, weight_name, shard_id in stacked_params_mapping:
            # ...

After:

def load_weights(self, weights):
    for name, loaded_weight in weights:
        name = name.replace(r"model.language_model.", r"model.")
        if ".self_attn." in name:
            name = name.replace(".self_attn", "")
        layer_id = get_layer_id(name)
        if (layer_id is not None
                and hasattr(self, "start_layer")
                and (layer_id < self.start_layer or layer_id >= self.end_layer)):
            continue  # 이 GPU의 레이어 범위 밖이면 스킵

get_layer_id(name) 유틸리티로 weight 이름에서 레이어 번호를 추출하고, PP 범위 밖이면 건너뜁니다. 이 패턴이 load_weights와 load_fused_expert_weights 양쪽 모두에 추가되었습니다.

2. Mamba 캐시 레이어 범위 수정

Before:

class MambaReqToTokenPool:
    def __init__(self, ...):
        self.start_layer = 0  # TODO: Support PP
        num_mamba_layers = len(cache_params.layers)

After:

class MambaReqToTokenPool:
    def __init__(self, ..., mamba_layer_ids: List[int], start_layer: int = None):
        self.start_layer = start_layer if start_layer is not None else 0
        num_mamba_layers = len(mamba_layer_ids)  # PP 범위 내 레이어만

기존에는 cache_params.layers(전체 Mamba 레이어)의 길이를 사용했으나, PP 모드에서는 이 GPU가 담당하는 레이어만 포함해야 합니다. 호출 측에서 필터링된 mamba_layer_ids를 전달합니다:

mamba_layer_ids=[
    i for i in config.mamba2_cache_params.layers
    if self.start_layer <= i < self.end_layer
]

3. Qwen3 VL의 PP 속성 전달

class Qwen3VLForConditionalGeneration:
    @property
    def start_layer(self) -> int:
        return getattr(getattr(self, "model", None), "start_layer", 0)

    @property
    def end_layer(self) -> int:
        model = getattr(self, "model", None)
        end_layer = getattr(model, "end_layer", None)
        if end_layer is not None:
            return end_layer
        return int(getattr(cfg, "num_hidden_layers", 0))

왜 이게 좋은가

1. PP 모드 지원 완성: Qwen3.5 MoE와 Mamba 하이브리드 모델이 PP 모드에서 정상 동작합니다.
2. 메모리 효율: 각 GPU가 자신의 레이어 weight만 로딩하여 메모리 사용량이 줄어듭니다.
3. 캐시 정확성: Mamba 캐시가 정확한 레이어 수만큼만 할당되어, 과다 할당을 방지합니다.

정리

PP 모드에서 레이어 범위 필터링은 모델 로딩과 캐시 할당 양쪽 모두에 적용해야 합니다. TODO: Support PP로 남겨둔 코드를 실제로 구현한 PR입니다.

참고 자료

• sgl-project/sglang#21448
• python/sglang/srt/models/qwen3_5.py
• python/sglang/srt/mem_cache/memory_pool.py

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.