[sglang] SGLang, 레이어별 오프로딩 기본값 설정을 통한 인코더/VAE 성능 최적화

PR 링크: sgl-project/sglang#25517 상태: Merged | 변경: +0 / -0

들어가며

최근 딥러닝 모델들은 점점 더 거대해지고 복잡해지면서, 제한된 GPU 메모리 환경에서 효율적으로 모델을 실행하는 것이 중요한 과제가 되었습니다. 특히 이미지 생성 모델과 같이 멀티모달 데이터를 처리하는 경우, 텍스트 인코더, 이미지 인코더, VAE(Variational Autoencoder) 등 다양한 컴포넌트들이 GPU 메모리를 많이 차지하게 됩니다. SGLang은 이러한 문제를 해결하기 위해 모델의 일부 레이어를 CPU 메모리로 옮기는 '오프로딩(offloading)' 기법을 지원해왔습니다.

이번 PR (#PR_NUMBER)은 기존의 CPU 오프로딩 방식보다 더 효율적인 '레이어별 오프로딩(layer-wise offloading)'을 텍스트 인코더, 이미지 인코더, VAE와 같은 주요 컴포넌트들에 대해 기본값으로 설정함으로써, 모델의 추론 속도를 크게 향상시키는 것을 목표로 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 분석하고, 왜 이러한 변경이 성능 향상으로 이어지는지, 그리고 이 최적화가 가지는 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py 및 관련 파일들에서 이루어졌습니다. 주요 변경점을 파일별로 살펴보겠습니다.

1. component_loader.py 및 layerwise_offload.py: 레이어별 오프로딩 로직 개선

이 PR은 layerwise_offload_components 설정의 동작 방식을 개선했습니다. 기존에는 특정 컴포넌트 그룹(LAYERWISE_OFFLOAD_DEFAULT_COMPONENTS)이 명시적으로 지정되지 않은 경우에만 기본값으로 레이어별 오프로딩을 적용했습니다. 하지만 이제는 LAYERWISE_OFFLOAD_DIT_GROUP이라는 새로운 그룹을 도입하고, normalize_layerwise_offload_components 함수를 통해 설정을 더 유연하게 처리합니다.

Before:

-from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
-    LAYERWISE_OFFLOAD_ALL_COMPONENTS,
-    LAYERWISE_OFFLOAD_DEFAULT_COMPONENTS,
-    layerwise_component_matches_selection,
-)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
+    LAYERWISE_OFFLOAD_ALL_COMPONENTS,
+    LAYERWISE_OFFLOAD_DIT_GROUP,
+    layerwise_component_matches_any_selection,
+    normalize_layerwise_offload_components,
+)

After:

-        selected_component_names = server_args.layerwise_offload_components
+        selected_component_names = normalize_layerwise_offload_components(
+            server_args.layerwise_offload_components
+        )
         if selected_component_names is None:
             return False
         selected_component_names = set(selected_component_names)
         if LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_component_names:
             return True
         explicit_component_names = selected_component_names - {
-            LAYERWISE_OFFLOAD_DEFAULT_COMPONENTS
+            LAYERWISE_OFFLOAD_DIT_GROUP
         }
-        return any(
-            layerwise_component_matches_selection(component_name, selected_component)
-            for selected_component in explicit_component_names
+        return layerwise_component_matches_any_selection(
+            component_name, explicit_component_names
         )

또한, configure_layerwise_offload_modules 함수는 이제 LAYERWISE_OFFLOAD_DIT_GROUP을 포함한 컴포넌트들을 기본값으로 설정할 수 있게 되었습니다. 이는 사용자가 명시적으로 layerwise_offload_components를 설정하지 않았을 때, DIT(Diffusion Image Transformer) 관련 컴포넌트들을 포함하여 레이어별 오프로딩을 적용하도록 변경된 것입니다. 이는 텍스트 인코더, 이미지 인코더, VAE와 같이 메모리를 많이 차지하는 컴포넌트들이 기본적으로 더 효율적인 오프로딩 방식을 사용하게 함을 의미합니다.

Before:

-    # Legacy --dit-layerwise-offload configures these modules when no component is named.
-    layerwise_offload_default_enabled: bool = True
+    # whether the current module is selected by the `dit` group
+    layerwise_offload_dit_group_enabled: bool = True

After:

-    layerwise_offload_default_enabled: bool = True
+    layerwise_offload_dit_group_enabled: bool = True

2. gpu_worker.py: layerwise_offload_components 설정 시 경고 로직 강화

gpu_worker.py 파일에서는 init_device_and_model 함수 내에서 layerwise_offload_components를 설정할 때, 사용자가 해당 설정을 명시적으로 지정했는지 여부를 warn_missing 인자를 통해 전달하도록 변경되었습니다. 이는 사용자가 의도하지 않은 컴포넌트에 대해 레이어별 오프로딩이 적용될 경우 경고를 발생시켜, 설정 오류를 방지하는 데 도움을 줍니다.

Before:

                 component_names=self.server_args.layerwise_offload_components,
+            )
 
         logger.info(

After:

                 component_names=self.server_args.layerwise_offload_components,
+                warn_missing=(
+                    self.server_args.is_arg_explicitly_set(
+                        "layerwise_offload_components"
+                    )
+                    or self.server_args.is_arg_explicitly_set("dit_layerwise_offload")
+                ),
+            )
 
         logger.info(

3. transformer_loader.py: 불필요한 로깅 제거

transformer_loader.py 파일에서는 모델 로딩 후 총 파라미터 수를 로깅하는 부분이 제거되었습니다. 이는 직접적인 성능 개선과는 관련이 없지만, 불필요한 로깅을 줄여 코드의 간결성을 높이는 변경입니다.

Before:

         for post_load_hook in quant_spec.post_load_hooks:
             post_load_hook(model)
 
-        total_params = sum(p.numel() for p in model.parameters())
-        logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
-
         # considering the existent of mixed-precision models (e.g., nunchaku)
         if (

After: (해당 로깅 라인 제거)

왜 이게 좋은가: 성능 향상과 메모리 효율성

이번 PR의 가장 큰 장점은 성능 향상과 메모리 효율성 증대입니다. PR 설명에 포함된 벤치마크 결과는 이를 명확하게 보여줍니다.

Component	GPU	cpu_offload	layerwise_offload	resident	layerwise vs cpu_offload
FLUX.1-dev text encoders, CLIP + T5	RTX 4090	552.67 ms	485.35 ms	57.15 ms	12.2% faster
Qwen-Image text encoder	RTX 4090	926.31 ms	664.80 ms	105.83 ms	28.2% faster
FLUX.1-dev text encoders, CLIP + T5	H100	225.56 ms	196.22 ms	24.11 ms	13.0% faster
FLUX.1-dev CLIP-only text encoder	RTX 4090	23.35 ms	14.73 ms	8.46 ms	36.9% faster
FLUX.1-dev CLIP-only text encoder	H100	12.82 ms	7.66 ms	5.78 ms	40.3% faster
Wan2.1 I2V image encoder, CLIPVision	H100	106.58 ms	50.69 ms	9.71 ms	52.4% faster
FLUX.1-dev VAE decode, AutoencoderKL	H100	151.99 ms	130.05 ms	130.01 ms	14.4% faster

위 표에서 볼 수 있듯이, 레이어별 오프로딩(layerwise_offload)은 기존의 CPU 오프로딩(cpu_offload) 방식에 비해 모든 경우에서 더 빠른 속도를 보여줍니다. 특히 Qwen-Image 텍스트 인코더의 경우 28.2%, FLUX.1-dev CLIP-only 텍스트 인코더는 40.3%까지 성능 향상이 있었습니다. 이는 레이어별 오프로딩이 GPU와 CPU 간의 데이터 전송을 최소화하고, 필요한 레이어만 동적으로 GPU로 가져오는 방식으로 동작하기 때문입니다.

또한, 레이어별 오프로딩은 GPU 메모리 사용량을 크게 줄여줍니다. 예를 들어 FLUX.1-dev VAE 디코딩 시, CPU 오프로딩은 5975 MiB, 레이어별 오프로딩은 5795 MiB, 완전 GPU 상주 시 5973 MiB의 메모리를 사용했습니다. 이는 완전 GPU 상주 방식과 거의 동일한 메모리 사용량으로 더 빠른 속도를 달성할 수 있음을 의미합니다. 즉, GPU 메모리가 부족한 환경에서도 더 큰 모델을 실행할 수 있는 가능성을 열어줍니다.

일반적인 교훈:

1. 세분화된 오프로딩의 중요성: 모델의 전체를 한 번에 옮기는 것보다, 필요한 부분만 동적으로 옮기는 레이어별 오프로딩이 훨씬 효율적입니다. 이는 GPU 메모리 제약이 있는 환경에서 모델의 성능을 극대화하는 핵심 전략입니다.
2. 기본값 설정의 힘: 사용자가 복잡한 설정을 직접 하지 않아도, 기본적으로 성능이 좋은 옵션을 제공하는 것은 사용자 경험과 전반적인 시스템 성능 향상에 크게 기여합니다. 이번 PR은 텍스트/이미지 인코더, VAE와 같이 성능에 민감한 컴포넌트들에 대해 레이어별 오프로딩을 기본값으로 설정함으로써 이러한 이점을 제공합니다.
3. 벤치마킹의 중요성: 다양한 하드웨어와 모델 구성에 대한 벤치마킹 결과는 변경의 효과를 정량적으로 입증하고, 최적화 방향을 설정하는 데 필수적입니다. PR에 제시된 표는 레이어별 오프로딩의 우수성을 명확히 보여줍니다.

References

• torch.compile - PyTorch의 컴파일 기능으로, 코드 실행 속도를 향상시키는 데 사용될 수 있습니다. (이 PR에서 직접적으로 사용되지는 않았지만, SGLang의 전반적인 성능 최적화 맥락에서 관련될 수 있습니다.)
• Layer-wise Offloading in Deep Learning - Hugging Face 블로그에서 레이어별 오프로딩에 대한 개념적 설명 (이 PR의 직접적인 참고 자료는 아니지만, 관련 기술에 대한 이해를 돕습니다.)
• SGLang GitHub Repository - SGLang 프로젝트의 소스 코드 및 개발 정보

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