[sglang] SGLang LTX-2 최적화: LoRA 병합 오버헤드 제거를 통한 추론 성능 향상

PR 링크: sgl-project/sglang#28594 상태: Merged | 변경: +264 / -3

들어가며

SGLang의 LTX-2 모델은 2단계(two-stage) 추론 방식을 사용합니다. 기존의 original 모드에서는 stage-1 distilled LoRA를 요청마다 베이스 모델에 병합(merge)하고, stage-2로 넘어갈 때 다시 해제(unmerge)하는 과정을 반복했습니다. 이 과정은 GPU 연산 자원을 소모하며 추론 지연 시간(latency)을 증가시키는 병목이었습니다. 본 PR은 이 문제를 해결하기 위해 stage-1 LoRA를 초기화 시점에 베이스 모델에 영구적으로 병합하고, stage-2는 동적 델타(dynamic delta)로 처리하는 최적화 기법을 도입했습니다.

코드 분석

1. BaseLayerWithLoRA.commit_merged_as_base 추가

linear.py 파일에 추가된 commit_merged_as_base 메서드는 현재 병합된 가중치를 새로운 베이스로 승격시킵니다.

# python/sglang/multimodal_gen/runtime/layers/lora/linear.py
@torch.no_grad()
def commit_merged_as_base(self) -> None:
    weight = self.base_layer.weight
    if isinstance(weight, DTensor):
        weight = weight.to_local()
    self.cpu_weight = weight.detach().to("cpu").clone()
    self.merged = False
    self.disable_lora = True
    self.lora_weights_list = []
    # ... (상태 초기화)

이 메서드는 기존의 cpu_weight 스냅샷을 업데이트하고, LoRA 관련 상태를 초기화하여 시스템이 이 가중치를 '기본값'으로 인식하게 만듭니다.

2. LTX2Pipeline의 초기화 및 상태 관리

ltx_2_pipeline.py에서는 파이프라인 초기화 시점에 _maybe_merge_stage1_distilled_into_base를 호출하여 최적화를 수행합니다.

# python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
def _switch_lora_phase_base_merged(self, phase: str, distilled_lora_strength: float) -> bool:
    if phase == "stage2":
        delta = distilled_lora_strength - float(self._stage1_distilled_base_strength)
        self.set_lora(..., strength=delta, merge_weights=False)
        return True

이제 stage-2 전환 시 전체를 다시 병합하는 대신, 이미 병합된 베이스에 필요한 델타값만 적용하여 연산량을 획기적으로 줄였습니다.

왜 이게 좋은가

이 최적화의 핵심은 '반복적인 메모리 복사 및 가중치 연산 제거'입니다. 기존에는 요청마다 merge와 unmerge가 발생하여 LTX2LoRASwitchStage에서 약 180ms가 소요되었으나, 최적화 이후 약 6ms로 단축되었습니다. 전체 E2E 지연 시간 또한 26.6초에서 15.9초로 약 40% 이상 개선되었습니다.

교훈:

1. 상태 관리의 최적화: 빈번하게 발생하는 모델 가중치 변경은 초기화 시점에 미리 계산(Pre-compute)하여 고정 비용으로 전환하는 것이 유리합니다.
2. 동적 델타 활용: 전체 모델을 교체하는 대신, 차이값(delta)만 적용하는 방식은 멀티 스테이지 파이프라인에서 매우 효과적인 전략입니다.

결론

이번 PR은 LTX-2 모델의 추론 파이프라인에서 불필요한 오버헤드를 제거함으로써 실질적인 성능 향상을 이끌어냈습니다. 특히 commit_merged_as_base와 같은 유연한 설계는 향후 다른 LoRA 기반 모델 최적화에도 훌륭한 레퍼런스가 될 것입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.no_grad.html
• https://github.com/sgl-project/sglang

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