본문으로 건너뛰기

[sglang] sglang, 멀티모달 모델 인코더 병렬 처리 최적화: 전체 복제본 활용으로 성능 향상

PR 링크: sgl-project/sglang#30086 상태: Merged | 변경: +396 / -51

들어가며

최근 sglang 프로젝트의 PR([diffusion] perf: tp-shard every text/image encoder across the full DiT replica (any parallelism))은 멀티모달 생성 모델, 특히 DiT(Diffusion Transformer) 모델에서 텍스트 및 이미지 인코더의 성능을 획기적으로 개선하는 내용을 담고 있습니다. 기존에는 디노이징 단계 이전에 GPU 자원이 유휴 상태임에도 불구하고, 인코딩 단계에서 텍스트 인코더(T5)만이 특정 조건(순수 SP, tp==1 && sp>1) 하에서만 병렬 처리(TP-sharding)되었습니다. 이로 인해 다양한 병렬 처리 구성(CFG-parallel, 혼합 레이아웃 등)이나 T5 외의 다른 인코더(CLIP, Llama, Qwen 등)를 사용할 때 인코딩 단계가 병렬화되지 못하고 성능 저하의 원인이 되었습니다.

이번 PR은 이러한 문제를 해결하기 위해, 인코딩 단계에서 유휴 상태인 전체 DiT 복제본(replica)의 모든 GPU를 활용하여 텍스트 및 이미지 인코더를 TP-sharding함으로써 병렬 처리 효율을 극대화합니다. 이는 모델의 종류나 병렬 처리 구성에 관계없이 일관된 성능 향상을 기대할 수 있게 합니다.

코드 분석

이번 PR의 핵심 변경 사항은 크게 두 가지로 나눌 수 있습니다.

  1. 전체 DiT 복제본에 걸친 병렬 폴딩(Parallel Folding) 일반화
  2. T5에서 모든 텍스트/이미지 인코더로의 일반화

1. ServerArgs.adjust_pipeline_config를 통한 전체 복제본 활용

기존에는 parallel_folding 및 parallel_folding_mode 설정이 TextEncoderConfig에 국한되어 있었고, 특정 조건에서만 작동했습니다. 이번 PR에서는 이 설정을 EncoderConfig의 기본 클래스로 옮기고, ServerArgs.adjust_pipeline_config 함수를 통해 어떤 병렬 처리 구성에서도 전체 DiT 복제본(replica)에 걸쳐 인코더를 TP-sharding하도록 일반화했습니다.

replica_size = num_gpus // dp_size로 정의되는 전체 복제본 크기를 기준으로 삼아, tp × cfg × ulysses × ring × pp × any future axis와 같은 모든 병렬 처리 조합을 명시적으로 나열할 필요 없이 자동으로 커버합니다. parallel_folding_mode="world" 옵션은 get_world_group()을 사용하여 전체 월드 그룹을 폴딩 그룹으로 지정합니다.

Before:

# python/sglang/multimodal_gen/configs/models/encoders/base.py
# ...
class EncoderConfig(ModelConfig):
    # ...
    pass

@dataclass
class TextEncoderConfig(EncoderConfig):
    arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)

    # Use the SP Group of the transformer as the TP Group of T5.
    parallel_folding: bool = False
    # "sp" or "ulysses" or "ring"
    parallel_folding_mode: str = "sp"

After:

# python/sglang/multimodal_gen/configs/models/encoders/base.py
# ...
class EncoderConfig(ModelConfig):
    # ...
    # Parallel folding: during the encoding stage the whole DiT replica is idle,
    # so TP-shard the encoder across those otherwise-unused GPUs instead of
    # running it on a single rank
    parallel_folding_mode: str | None = None

@dataclass
class TextEncoderConfig(EncoderConfig):
    arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
    # ... (removed parallel_folding and parallel_folding_mode from here)

T5Config에서도 유사하게 기존의 parallel_folding 관련 설정을 제거하여 기본 EncoderConfig의 설정을 따르도록 변경되었습니다.

2. 모든 텍스트/이미지 인코더로의 일반화

이전에는 T5 인코더에만 병렬 폴딩이 적용되었지만, 이제는 모든 텍스트 및 이미지 인코더로 확장되었습니다.

  • parallel_folding_mode가 EncoderConfig의 기본 필드로 올라가면서 모든 인코더가 상속받게 되었습니다. 이전에는 AttributeError가 발생할 수 있었던 부분을 개선했습니다.
  • get_folding_tp_group() 함수가 공유되어, 텍스트 인코더 로더는 patch_tensor_parallel_group을 통해 인코더 생성 및 가중치 로드 시 TP 그룹을 패치합니다. 이는 모델별로 별도의 스레딩 없이 모든 병렬 선형 레이어(Linear) 및 임베딩이 생성 및 가중치 로드 시 폴딩 그룹에 맞춰 샤딩되도록 합니다.
  • adjust_pipeline_config 함수는 이제 인코더의 어텐션 헤드와 MLP 중간 차원이 폴딩 그룹의 크기로 나누어 떨어지는 경우(속성 이름의 차이에 강건함)에 대해 폴딩을 활성화합니다. 내부적으로 인코딩할 수 없는 인코더는 안전한 기본값으로 폴딩되지 않은 상태로 둡니다.

Gemma 3 인코더의 forward 메서드 수정:

Gemma 3와 같은 모델에서는 SiglipAttention이 순방향 실행 시점에 전역 TP 사이즈를 읽어왔습니다. 이는 construction-time 패치로 커버되지 않는 부분이었으나, 이번 PR에서는 해당 부분을 수정하여 런타임 시점에 per-rank 너비를 캐싱하도록 변경했습니다. 이를 통해 qkv.split 시 올바른 샤딩 크기를 사용하게 됩니다.

Before:

# python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py
# ...
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.hidden_size // get_tp_world_size()] * 3, dim=-1)

        batch_size, seq_len, _ = q.shape
        q = q.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)

        attn_output = self.attn(q, k, v)

        attn_output = attn_output.reshape(
            batch_size, seq_len, self.hidden_size // get_tp_world_size()
        )

        output, _ = self.o_proj(attn_output)
        return output

After:

# python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py
# ...
    def __init__(self, ..., hidden_size: int, num_heads: int, ...):
        # ...
        tp_size = get_tp_world_size()
        self.head_dim = hidden_size // num_heads
        self.num_heads_per_partition = num_heads // tp_size
        # Cache the per-rank projection width so forward() does not re-read the
        # global TP size (which is not patched to the folding group at run time).
        self.embed_dim_per_partition = self.num_heads_per_partition * self.head_dim
        self.scaling = self.head_dim**-0.5

        self.qkv_proj = QKVParallelLinear(
            hidden_size, hidden_size, bias=False, tp_size=tp_size
        )
        # ...

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.embed_dim_per_partition] * 3, dim=-1)

        batch_size, seq_len, _ = q.shape
        q = q.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)

        attn_output = self.attn(q, k, v)

        attn_output = attn_output.reshape(
            batch_size, seq_len, self.embed_dim_per_partition
        )

        output, _ = self.o_proj(attn_output)
        return output

또한, finalize_encoder_folding 함수가 추가되어, 인코더의 실제 차원이 결정된 후에 폴딩이 실제로 이득이 되는지 (예: hidden_size >= FOLD_MIN_HIDDEN_SIZE 및 헤드/MLP 차원이 그룹 크기로 나누어 떨어지는지) 검증하고, 그렇지 않으면 폴딩 모드를 해제하여 회귀를 방지합니다.

왜 이게 좋은가

이 PR의 가장 큰 장점은 인코딩 단계의 유휴 GPU 자원을 최대한 활용하여 전체 모델의 추론 속도를 향상시킨다는 점입니다.

  1. 성능 향상:

    • 기존에는 cfg-parallel 구성에서 텍스트 인코더가 단일 GPU에서 실행되어 나머지 CFG 복제본 GPU가 유휴 상태였습니다. 이 PR을 통해 해당 GPU들도 인코딩 작업에 투입되어 처리량이 증가합니다.
    • 다양한 병렬 처리 레이아웃(TP, SP, CFG, 혼합 등) 및 다양한 종류의 텍스트/이미지 인코더(CLIP, Llama, Qwen, Gemma 등)에 대해 일관되게 병렬 폴딩을 적용할 수 있게 되어, 특정 모델이나 구성에 국한되지 않는 성능 개선을 제공합니다.
    • 리뷰에서 언급된 바와 같이, 실제 H100 GPU 환경에서의 검증 결과, MAE(Mean Absolute Error)가 약 1.4e-8 수준으로 매우 낮아, 샤딩 및 재구성 과정에서 손실이 거의 없음을 확인했습니다. 이는 로직이 정확하며 성능 저하 없이 병렬화가 이루어졌음을 의미합니다.

  2. 코드의 일반성 및 유지보수성 향상:

    • 병렬 폴딩 로직을 기본 EncoderConfig로 추상화하고, patch_tensor_parallel_group과 같은 유틸리티를 사용하여 모델별 특수 처리를 최소화했습니다. 이는 새로운 인코더 모델을 추가하거나 기존 모델을 수정할 때 병렬 처리 관련 코드를 일관되게 관리할 수 있게 합니다.
    • finalize_encoder_folding 함수를 통해 실제 모델의 차원을 기반으로 폴딩의 유효성을 검증함으로써, 작은 모델에서 불필요한 오버헤드를 방지하고 회귀를 예방하는 안전장치를 마련했습니다.

  3. 일반적인 교훈:

    • 유휴 자원 활용: 딥러닝 모델 추론 파이프라인에서 특정 단계에 유휴 상태가 되는 GPU 자원이 있다면, 이를 다른 계산 집약적인 작업(여기서는 인코더 병렬화)에 활용하는 방안을 적극적으로 모색해야 합니다.
    • 추상화를 통한 일반화: 특정 모델이나 구성에 국한된 최적화 로직을 기본 클래스나 공통 유틸리티 함수로 추상화하면, 코드의 재사용성과 유지보수성을 크게 높일 수 있습니다.
    • 안전장치 마련: 최적화 로직이 모든 경우에 유효하지 않을 수 있으므로, 실제 환경에서의 검증(예: 모델 크기, 차원 등)을 통해 최적화 적용 여부를 동적으로 결정하는 메커니즘을 도입하는 것이 중요합니다.

리뷰 과정에서 transfer_weights 함수의 버그가 발견되어 샤딩된 파라미터 처리에 문제가 있었으나, 이는 테스트 코드 수정으로 해결되었고 실제 로직의 문제는 아니었습니다. 이는 복잡한 분산 시스템에서 테스트의 중요성을 다시 한번 강조합니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글