본문으로 건너뛰기

[sglang] SGLang, Helios 모델 통합으로 실시간 장편 비디오 생성의 새로운 지평을 열다

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

들어가며

최근 AI 기술의 발전은 텍스트, 이미지 생성을 넘어 이제는 동영상 생성 분야에서도 놀라운 성과를 보여주고 있습니다. 특히, 긴 길이의 비디오를 실시간에 가깝게 생성하는 것은 많은 연구자와 개발자들이 도전해 온 난제였습니다. SGLang 프로젝트는 이러한 흐름에 발맞춰, 최첨단 장편 비디오 생성 모델인 Helios를 성공적으로 통합하며 이 분야에 새로운 가능성을 제시했습니다.

이번 Pull Request (PR)는 SGLang 프레임워크에 Helios 모델을 도입하고, 관련 설정을 추가하는 것을 목표로 합니다. 이를 통해 SGLang 사용자는 Helios 모델을 활용하여 더욱 빠르고 효율적으로 고품질의 비디오를 생성할 수 있게 되었습니다. 본 글에서는 이 PR의 코드 변경 사항을 상세히 분석하고, Helios 모델 통합이 왜 SGLang의 성능과 기능 확장에 중요한 기여를 하는지 기술적인 관점에서 설명하고자 합니다.

코드 분석

이번 PR은 주로 SGLang의 설정(configuration) 관련 파일들을 수정하고 새로운 파일을 추가하는 방식으로 Helios 모델을 통합합니다. 핵심적인 변경 사항은 다음과 같습니다.

1. sglang/multimodal_gen/configs/models/dits/__init__.py

이 파일은 DiT(Diffusion Transformer) 계열 모델들의 설정을 관리하는 곳입니다. Helios 모델을 추가하기 위해 HeliosConfig__all__ 리스트에 포함시켜 외부에서 접근 가능하도록 합니다.

Before:

from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig
from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig

__all__ = [
    "HunyuanVideoConfig",
    "WanVideoConfig",
    "Hunyuan3DDiTConfig",
]

After:

from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig
from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig

__all__ = [
    "HeliosConfig",
    "HunyuanVideoConfig",
    "WanVideoConfig",
    "Hunyuan3DDiTConfig",
]

2. sglang/multimodal_gen/configs/models/dits/helios.py (신규 파일)

Helios 모델의 아키텍처 및 기본 설정을 정의하는 새로운 파일입니다. HeliosArchConfigHeliosConfig 두 개의 dataclass가 정의되어 있습니다.

  • HeliosArchConfig: DiT 아키텍처를 기반으로 하며, FSDP(Fully Sharded Data Parallel) 샤딩 조건을 위한 _fsdp_shard_conditions와 파라미터 이름 매핑(param_names_mapping)을 포함합니다. 특히, param_names_mapping은 기존 모델의 가중치 이름을 Helios 모델의 구조에 맞게 변환하는 역할을 하여, 사전 학습된 가중치를 효율적으로 로드할 수 있게 합니다. 리뷰어 yhyang201이 지적한 것처럼, m 파라미터는 FSDP 샤딩 조건 콜백 인터페이스에 필수적입니다.

  • HeliosConfig: HeliosArchConfig를 사용하여 전체 Helios 모델 구성을 정의합니다. prefix는 모델 식별에 사용됩니다.

# ... (생략) ...

@dataclass
class HeliosArchConfig(DiTArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])

    param_names_mapping: dict = field(
        default_factory=lambda: {
            # Patch embeddings
            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
            # Condition embedder: text
            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
            # Condition embedder: time
            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
            r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
            # Blocks: self-attention (keep attn1. prefix, drop .0. from to_out)
            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.attn1.to_out.\2",
            # Blocks: cross-attention output (drop .0. from to_out)
            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
            # Blocks: feed-forward
            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
            # Blocks: cross-attn residual norm
            r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.\2",
        }
    )

    # ... (기타 파라미터들) ...

    # Helios-specific
    rope_dim: tuple[int, int, int] = (44, 42, 42)
    rope_theta: float = 10000.0
    guidance_cross_attn: bool = True
    zero_history_timestep: bool = True
    has_multi_term_memory_patch: bool = True
    is_amplify_history: bool = False
    history_scale_mode: str = "per_head"

    def __post_init__(self):
        super().__post_init__()
        self.out_channels = self.out_channels or self.in_channels
        self.hidden_size = self.num_attention_heads * self.attention_head_dim
        self.num_channels_latents = self.out_channels


@dataclass
class HeliosConfig(DiTConfig):
    arch_config: DiTArchConfig = field(default_factory=HeliosArchConfig)

    prefix: str = "Helios"

3. sglang/multimodal_gen/configs/pipeline_configs/__init__.py

파이프라인 설정을 관리하는 파일에 Helios 관련 파이프라인 설정을 추가합니다. HeliosT2VConfig, HeliosMidConfig, HeliosDistilledConfig__all__에 포함됩니다.

4. sglang/multimodal_gen/configs/pipeline_configs/helios.py (신규 파일)

Helios 모델을 위한 다양한 파이프라인 설정을 정의합니다. 텍스트-비디오(T2V) 생성, 중간 해상도 강화(Mid), 증류(Distilled) 버전 등 여러 시나리오에 맞는 설정을 제공합니다.

  • HELIOS_MAX_SEQUENCE_LENGTH: UMT5 텍스트 인코더의 최대 시퀀스 길이를 정의합니다.
  • umt5_postprocess_text: UMT5 인코더의 출력을 지정된 최대 시퀀스 길이로 패딩하는 함수입니다. 이는 모델 입력의 일관성을 유지하는 데 중요합니다.
  • Helios...Config 클래스는 PipelineConfig를 상속받아 모델, VAE, 텍스트 인코더, 정밀도 설정 등을 정의합니다. 특히, HeliosT2VConfig는 기본적인 텍스트-비디오 생성 설정을, HeliosMidConfig는 스테이지 2 피라미드 SR을, HeliosDistilledConfig는 스테이지 2 및 스테이지 3 DMD를 활성화하는 설정을 포함합니다.
# ... (생략) ...

# Helios UMT5 max sequence length (used for both tokenizer and post-processing padding)
HELIOS_MAX_SEQUENCE_LENGTH = 226

def umt5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
    """Post-process UMT5 text encoder outputs, padding to HELIOS_MAX_SEQUENCE_LENGTH tokens."""
    max_seq_len = HELIOS_MAX_SEQUENCE_LENGTH
    mask: torch.Tensor = outputs.attention_mask
    hidden_state: torch.Tensor = outputs.last_hidden_state
    seq_lens = mask.gt(0).sum(dim=1).long()
    assert torch.isnan(hidden_state).sum() == 0
    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
    prompt_embeds_tensor: torch.Tensor = torch.stack(
        [
            torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))])
            for u in prompt_embeds
        ],
        dim=0,
    )
    return prompt_embeds_tensor

# ... (HeliosT2VConfig, HeliosMidConfig, HeliosDistilledConfig 정의) ...

5. sglang/multimodal_gen/configs/sample/helios.py (신규 파일)

Helios 모델 사용 시 필요한 샘플링 파라미터들을 정의합니다. 비디오 해상도, 프레임 수, FPS, 가이던스 스케일, 네거티브 프롬프트 등이 포함됩니다. 각 파이프라인 설정(HeliosT2VConfig, HeliosMidConfig, HeliosDistilledConfig)에 맞는 샘플링 파라미터 클래스(HeliosT2VSamplingParams, HeliosMidSamplingParams, HeliosDistilledSamplingParams)가 제공됩니다.

# ... (생략) ...

@dataclass
class HeliosT2VSamplingParams(SamplingParams):
    # Video parameters
    height: int = 384
    width: int = 640
    num_frames: int = 99
    fps: int = 24

    # Denoising stage
    guidance_scale: float = 5.0
    negative_prompt: str = (
        "Bright tones, overexposed, static, blurred details, subtitles, style, "
        "works, paintings, images, static, overall gray, worst quality, low quality, "
        "JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, "
        "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, "
        "still picture, messy background, three legs, many people in the background, "
        "walking backwards"
    )
    num_inference_steps: int = 50

    # Helios T2V supported resolutions
    supported_resolutions: list[tuple[int, int]] | None = field(
        default_factory=lambda: [
            (640, 384),  # ~5:3
            (384, 640),  # ~3:5
            (832, 480),  # ~16:9-ish
            (480, 832),  # ~9:16-ish
        ]
    )

# ... (HeliosMidSamplingParams, HeliosDistilledSamplingParams 정의) ...

6. sglang/multimodal_gen/registry.py

리뷰어 mickqian이 언급한 model_detectors 관련 부분은 이 PR에서 직접적으로 수정되지 않았지만, 향후 model_detectors 기능의 유지 또는 폐기 여부에 대한 논의가 필요함을 시사합니다. 현재로서는 Helios 모델을 위한 별도의 탐지기(detector) 로직이 추가되지는 않았습니다.

왜 이게 좋은가?

1. 최신 모델 통합을 통한 성능 향상

Helios 모델은 장편 비디오 생성 분야에서 최신 기술을 집약한 모델입니다. SGLang에 Helios를 통합함으로써, 사용자는 SGLang 프레임워크의 효율적인 추론 엔진 위에서 Helios의 강력한 비디오 생성 성능을 활용할 수 있게 됩니다. 이는 특히 실시간에 가까운 속도로 긴 비디오를 생성해야 하는 애플리케이션에서 큰 이점을 제공합니다. 비록 이 PR 자체에 구체적인 성능 수치(예: FPS 향상률)가 포함되어 있지는 않지만, 모델 통합은 일반적으로 다음과 같은 최적화를 가능하게 합니다:

  • 효율적인 메모리 관리: SGLang의 최적화된 메모리 관리 기법이 Helios 모델에 적용될 수 있습니다.
  • 커널 퓨전 및 연산 최적화: SGLang의 추론 엔진은 모델의 연산 그래프를 최적화하여 불필요한 연산을 줄이고 GPU 활용률을 높입니다. Helios 모델의 복잡한 트랜스포머 구조에 이러한 최적화가 적용되면 상당한 속도 향상을 기대할 수 있습니다.
  • 병렬 처리 강화: 분산 학습 및 추론 환경에서의 최적화(_fsdp_shard_conditions 등)는 대규모 모델을 효율적으로 실행하는 데 필수적입니다.

2. 유연하고 확장 가능한 설정 관리

이 PR은 dataclasses를 활용하여 모델 및 파이프라인 설정을 매우 유연하고 체계적으로 관리합니다. 각 설정 파일은 특정 모델이나 파이프라인의 특성을 명확하게 정의하며, 필요한 경우 쉽게 수정하거나 확장할 수 있습니다. 예를 들어, HeliosT2VConfig, HeliosMidConfig, HeliosDistilledConfig와 같이 동일한 기본 모델을 사용하더라도 다른 파이프라인 구성 요소를 조합하여 다양한 결과물을 얻을 수 있도록 설계되었습니다.

3. 재현성 및 일관성 확보

명확하게 정의된 설정 파일과 샘플링 파라미터는 모델의 동작을 재현 가능하게 만듭니다. 사용자는 동일한 설정을 사용하여 언제든지 동일한 결과를 얻을 수 있으며, 이는 연구 및 개발 과정에서 매우 중요합니다. 또한, umt5_postprocess_text와 같은 함수는 텍스트 인코더 출력의 일관성을 보장하여 후속 비디오 생성 단계의 안정성을 높입니다.

일반적인 교훈

  • 최신 AI 모델의 신속한 통합: SGLang과 같이 빠르게 발전하는 프레임워크는 최신 연구 성과를 신속하게 통합하는 것이 경쟁력의 핵심입니다. Helios와 같은 최첨단 모델의 통합은 프레임워크의 가치를 높입니다.
  • 체계적인 설정 관리: 복잡한 AI 모델을 다룰 때는 dataclasses와 같은 구조화된 방식을 사용하여 설정을 관리하는 것이 유지보수성과 확장성을 크게 향상시킵니다.
  • 명확한 역할 분담: 모델 아키텍처 설정, 파이프라인 설정, 샘플링 파라미터 설정을 분리하여 관리함으로써 각 구성 요소의 역할을 명확히 하고 코드의 가독성을 높일 수 있습니다.

References

  • torch.compile: SGLang은 PyTorch의 torch.compile과 같은 최적화 도구를 활용하여 모델 추론 속도를 향상시킬 수 있습니다. (이 PR에서 직접적으로 사용되지는 않았으나, SGLang의 전반적인 성능 최적화 전략에 포함될 수 있습니다.)
  • FSDP (Fully Sharded Data Parallel): 대규모 모델 학습 및 추론 시 메모리 효율성을 높이는 PyTorch의 기능입니다. Helios 모델의 _fsdp_shard_conditions 설정은 FSDP 활용을 염두에 둔 것입니다.
  • DiT (Diffusion Transformer): 비디오 생성 모델 Helios의 기반이 되는 아키텍처입니다. (논문 링크)
  • UMT5: Helios 모델에서 텍스트 인코딩에 사용되는 모델입니다. (Hugging Face 모델 카드 링크)

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글