Prhub

#34391 [Diffusion] Support dynamic CPU offload components

原始 PR 作者 mickqian 合并时间 2026-08-12 11:38 文件变更 31 提交数 10 评论 0 代码增减 +631 / -173

执行摘要

--cpu-offload-components 支持 model_index 动态键与 all/none

PR body 明确说明 --cpu-offload-components 应当支持 Diffusers model_index.json 中的组件键,包括“model-specific components that are not known to SGLang in advance”。以往只有 dit/text_encoder/image_encoder/vae 四个预定义分组,新增 checkpoint 若带 transformer_2、audio_vae、connectors 等动态键则无法细粒度控制卸载,甚至被忽略导致显存占用超出预期。本 PR 引入统一选择器并保持 legacy 别名兼容,同时解决 single-DiT 常驻性能优化与显式 offload 策略冲突的历史问题。

值得精读。核心设计是三层:normalize_cpu_offload_components 做输入规范化、ServerArgs.should_cpu_offload_component 做统一决策、component_manager 的 _should_keep_single_dit 让性能快路径服从内存策略。建议重点对照 server_args.py 的 _adjust_cpu_offload_components 与 component_manager.py 的 finish_request 逻辑;如果团队在持续扩展新 diffusion 模型,这套“动态组件键 + 统一装载判定”模式可直接复用;若只维护现有模型,需关注 MPS 与 legacy flag 混用边界。

讨论亮点

该 PR 没有外部 review 评论(review_comments_count=0),核心权衡体现在 10 次作者自审提交中,从提交信息可读出关键决策:

  • a0b41c5 "honor offload policy in component residency":开始让驻留策略服从 offload 配置,这是优先级修正的起点。
  • 9ff7336 "restore causal DiT residency"、e2df3120 "keep Wan single DiTs resident on H100"、df96ed0 "keep distilled video DiTs resident on H100":在 offload 与 keep-resident 性能之间反复调优,最终把 single-DiT 常驻从“名称匹配”改为“实际策略为 ResidentStrategy 才保留”。
  • 94d5feb "preserve auto FSDP selection":调整 auto_tune 执行顺序,避免统一选择器干扰 FSDP 自动判定。
  • 结论:单一 DiT 常驻是性能快路径而非内存策略,必须服从显式/自动 CPU 与 layerwise offload;finish_request 的 preferred 预取也不得覆盖该策略。

实现拆解

  1. 新增统一参数与规范化:server_args.py 增加 cpu_offload_components 字段及 --cpu-offload-components CLI 参数;layerwise_offload_components.py 新增 normalize_cpu_offload_components(支持逗号分隔、连字符转下划线、大小写归一、none 单独出现)与 cpu_offload_component_matches(支持 all 通配、精确键、legacy 分组别名的匹配)。
  2. 统一决策入口:ServerArgs.should_cpu_offload_component 成为唯一判定入口,cpu_offload_components 优先,未设置时回退到 legacy flag 分组判断,并把 connectors、unconditional_transformer、vision_language_encoder 归入 DiT 组、sound_tokenizer 归入 VAE 组。_adjust_cpu_offload_components 在 _adjust_parameters 早期执行,将统一选择同步到 dit/text_encoder/image_encoder/vae 四个 legacy flag,且与显式 legacy flag 冲突时直接报 ValueError。
  3. 加载链路迁移:component_loader.py 基类 should_offload 增加 component_name 参数并默认走统一判断;text_encoder_loader.py、image_encoder_loader.py、sound_tokenizer_loader.py、vae_loader.py、vocoder_loader.py、upsampler_loader.py、vl_encoder_loader.py、adapter_loader.py、transformer_loader.py 等逐一改为调用 server_args.should_cpu_offload_component(component_name),component_loader.load() 中还增加对非 FSDP 模块的 CPU 兜底移动,覆盖通用组件加载路径。
  4. 驻留策略调整:component_manager.py 删除模块级 should_cpu_offload_component,把 MPS/FSDP/统一选择判断内联进 build_component_residency_strategy;新增 _is_single_dit_component,_should_keep_single_dit 改为依赖实际策略是否为 ResidentStrategy,并移除 lru_cache;finish_request 中 preferred 预取不再覆盖 single-DiT 的 CPU/layerwise offload。
  5. 内存分析与自动策略联动:gpu_worker.py 的 get_can_stay_resident_components 改为遍历 pipeline.memory_usages 全部键(含动态组件),用 should_cpu_offload_component 与 should_configure_layerwise_offload_for_lazy_component 判定,删除硬编码 offload flag 映射与 _format_offload_disable_suggestions;auto_tune.py 调整 FSDP 自动选择与组件驻留调整的顺序;wan.py、lingbot_world.py、zimage 等模型部署配置补充 keep_resident_min_available_gb 与 keep_resident_components 声明。
  6. 测试配套:test_server_args.py 新增 CLI 解析、model_index 键保留、all/none 通配、legacy flag 冲突、5090/H100 自动驻留策略等大量单元测试;test_causal_denoising.py 新增 DiT 前向调用前 prepare 顺序测试;最后提交刷新 B200 Ideogram4 基准 ground truth(+631/-173,覆盖 31 个文件)。
文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py 驻留管理 modified 7.96
python/sglang/multimodal_gen/runtime/server_args/server_args.py 参数解析 modified 7.81
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py 卸载选择 modified 7.36
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py 工作进程 modified 7.03
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py 加载器 modified 6.32
python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py 加载器 modified 5.99
python/sglang/multimodal_gen/test/unit/test_server_args.py 参数测试 modified 7.24

关键符号

should_cpu_offload_component _adjust_cpu_offload_components normalize_cpu_offload_components cpu_offload_component_matches build_component_residency_strategy _should_keep_single_dit _is_single_dit_component get_can_stay_resident_components should_offload load_customized

关键源码片段

python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py core-logic

驻留策略主路径:删除模块级 should_cpu_offload_component,内联 MPS/FSDP/ 统一选择判断;新增 _is_single_dit_component,_should_keep_single_dit 改为依据实际 ResidentStrategy,并调整 finish_request 中 preferred 预取不覆盖 offload 策略,是本 PR 优先级修正的核心。

# 源码:python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py
# 判断某个流水线组件是否应该做整模块 CPU offload(D2H)。
# 注意:MPS 与 FSDP 模块一律留在原设备,避免破坏已验证的 CUDA 行为。
def build_component_residency_strategy(
    component_name: str,
    module: nn.Module,
    server_args: ServerArgs,
) -> ComponentResidencyStrategy:
    if is_layerwise_offloaded_module(module):
        return LayerwiseOffloadStrategy()
    if (
        not current_platform.is_mps()
        and not server_args.use_fsdp_inference
        and not is_fsdp_managed_module(module)
        and server_args.should_cpu_offload_component(component_name)
    ):
        return VanillaD2HStrategy()
    return ResidentStrategy()
​
​
class ComponentResidencyManager:
    # 单一 DiT 的“留在显存”只是性能快路径,不是内存策略,
    # 因此只有真正拿到 ResidentStrategy 时才保留 DiT。
    # finish_request 里也禁止 preferred 预取覆盖 CPU/layerwise offload。
    def _should_keep_single_dit(self, component_name: str, module: nn.Module) -> bool:
        """Keep a single DiT resident only when its effective strategy is resident."""
        if not self._is_single_dit_component(component_name):
            return False
        return isinstance(self.strategy_for(component_name, module), ResidentStrategy)
​
    def _is_single_dit_component(self, component_name: str) -> bool:
        modules = self.pipeline.modules
        return (component_name == "transformer" and "transformer_2" not in modules) or (
            component_name == "video_dit" and "video_dit_2" not in modules
        )
python/sglang/multimodal_gen/runtime/server_args/server_args.py core-logic

新增 --cpu-offload-components 参数定义、_adjust_cpu_offload_components 与 should_cpu_offload_component 统一决策入口,并处理与 legacy flag 的显式冲突,是整个变更的配置契约层。

# 统一入口:--cpu-offload-components 会把 model_index.json 的任意组件键映射到
# 各 legacy 布尔 flag,保证下游 loader 不再需要感知动态组件名。
def _adjust_cpu_offload_components(self) -> None:
    if self.cpu_offload_components is None:
        return
​
    legacy_flags = (
        "dit_cpu_offload",
        "text_encoder_cpu_offload",
        "image_encoder_cpu_offload",
        "vae_cpu_offload",
    )
    conflicting_flags = [
        flag_name
        for flag_name in legacy_flags
        if self.is_arg_explicitly_set(flag_name)
    ]
    if conflicting_flags:
        # 显式互斥,避免两个 selector 同时生效导致行为不可预测
        formatted_flags = ", ".join(
            "--" + flag_name.replace("_", "-") for flag_name in conflicting_flags
        )
        raise ValueError(
            "--cpu-offload-components cannot be combined with the legacy "
            f"CPU offload flags: {formatted_flags}"
        )
​
    selected_components = (
        normalize_cpu_offload_components(self.cpu_offload_components) or []
    )
    self.cpu_offload_components = selected_components
    # 把统一选择同步到 legacy flag,保持旧 loader 与旧诊断逻辑可读
    self.dit_cpu_offload = self.should_cpu_offload_component("transformer")
    self.text_encoder_cpu_offload = self.should_cpu_offload_component("text_encoder")
    self.image_encoder_cpu_offload = self.should_cpu_offload_component("image_encoder")
    self.vae_cpu_offload = self.should_cpu_offload_component("vae")
​
​
def should_cpu_offload_component(self, component_name: str) -> bool:
    # 统一选择器优先;未设置时回退到 legacy 分组判断,
    # 并把 connectors / sound_tokenizer 等归入既有组,保证动态组件可用。
    if self.cpu_offload_components is not None:
        return cpu_offload_component_matches(
            component_name, self.cpu_offload_components
        )
    if is_dit_component_name(component_name) or component_name in (
        "connectors",
        "unconditional_transformer",
        "vision_language_encoder",
    ):
        return bool(self.dit_cpu_offload)
    if is_text_encoder_component_name(component_name):
        return bool(self.text_encoder_cpu_offload)
    if is_image_encoder_component_name(component_name):
        return bool(self.image_encoder_cpu_offload)
    if is_vae_component_name(component_name) or component_name == "sound_tokenizer":
        return bool(self.vae_cpu_offload)
    return False
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py core-logic

新增 normalize_cpu_offload_components 与 cpu_offload_component_matches 两个核心辅助函数,负责参数规范化与动态组件匹配,是统一选择器的底层契约。

# 规范化 --cpu-offload-components 的输入:
# 支持逗号分隔、连字符转下划线、大小写归一;'none' 必须单独出现。
def normalize_cpu_offload_components(
    component_names: str | Sequence[str] | None,
) -> list[str] | None:
    if component_names is None:
        return None
    raw_components = (
        [component_names] if isinstance(component_names, str) else component_names
    )
    normalized_components: list[str] = []
    for raw_component in raw_components:
        if not isinstance(raw_component, str):
            raise ValueError(f"Invalid CPU offload component name: {raw_component}.")
        normalized_components.extend(
            component_name
            for value in raw_component.split(",")
            if (component_name := value.strip().replace("-", "_").lower())
        )
​
    unique_components = list(dict.fromkeys(normalized_components))
    if "none" in unique_components:
        if len(unique_components) != 1:
            raise ValueError("'none' cannot be combined with other components.")
        return [] # none 表示显式关闭所有组件 offload
    return unique_components or None
​
​
# 匹配一个运行时组件名是否落在用户选择中:
# 'all' 通配所有模块;显式键优先;其次兼容 legacy 分组别名。
def cpu_offload_component_matches(
    component_name: str,
    selected_component_names: Collection[str] | None,
) -> bool:
    if selected_component_names is None:
        return False
    if CPU_OFFLOAD_ALL_COMPONENTS in selected_component_names:
        return True
    if component_name in selected_component_names:
        return True
    if LAYERWISE_OFFLOAD_DIT_GROUP in selected_component_names:
        return is_dit_component_name(component_name)
    if LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP in selected_component_names:
        return is_text_encoder_component_name(component_name)
    if LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP in selected_component_names:
        return component_name in ("image_encoder", "condition_image_encoder")
    if LAYERWISE_OFFLOAD_VAE_GROUP in selected_component_names:
        return is_vae_component_name(component_name)
    return False

评论区精华

单一 DiT 驻留快路径与 offload 策略的优先级 设计

无 review 评论;从提交历史可看出 a0b41c5 先实现 offload 策略生效,随后 9ff7336、e2df3120、df96ed0 分别恢复 causal/Wan/ 蒸馏视频 DiT 的常驻,最后通过“仅当实际策略为 ResidentStrategy 才保留”收敛矛盾。

结论:_should_keep_single_dit 从名称匹配改为策略匹配,finish_request 中 preferred 预取不覆盖 CPU/layerwise offload。 · 已解决

统一选择器与 auto FSDP 自动调整的执行顺序 设计

94d5feb 提交专门提出 preserve auto FSDP selection;_adjust_parameters 中 _adjust_cpu_offload_components 先执行,随后 auto_tuner 的 maybe_adjust_auto_fsdp_with_offload_enabled 再调整,避免统一选择器覆盖 FSDP 判定。

结论:调整 auto_tune 顺序:先 FSDP 相关调整,再做组件驻留 after-offload 调整。 · 已解决

MPS 平台下 --cpu-offload-components 与强制禁用 legacy flag 的交互 question

代码推断:_adjust_platform_specific 在 MPS 上把 legacy flags 全部置 False 并告警,但 should_cpu_offload_component 在 cpu_offload_components 非 None 时优先返回匹配结果,MPS 上仍可能把组件放到 CPU,与“仅在 CUDA 上验证”的注释冲突。

结论:未解决;现有测试未覆盖 MPS 与统一选择器组合,建议补充短路或测试。 · 待处理

风险与影响

  1. MPS 平台组合未覆盖:_adjust_platform_specific 对 MPS 会强制把所有 legacy flag 置 False 并告警,但 should_cpu_offload_component 在 cpu_offload_components 非 None 时直接返回匹配结果,MPS 上可能仍把组件移到 CPU,与“仅 CUDA 验证”的注释相矛盾,存在设备行为不一致风险。
  2. 兼容性:--cpu-offload-components 与任一显式 legacy flag 同时设置会直接 ValueError 导致启动失败,虽然错误信息清晰,但既有启动脚本若混用会中断。
  3. 加载峰值内存:component_loader.load() 通用路径对非 FSDP 模块执行 to("cpu"),若后续又参与 layerwise offload 配置可能发生重复搬运,对低显存场景峰值内存影响需要 GPU/CI 实测(PR body 也注明请求 GPU/CI 验证)。
  4. lru_cache 生命周期:_should_keep_single_dit 移除缓存后每次请求多一次 strategy_for 查找与 isinstance 判断,开销可接受;但 strategy_for 本身仍以 module 为键做 lru_cache,若 pipeline 模块动态重建可能累积旧引用。
  5. auto_tune 顺序敏感:_adjust_cpu_offload_components 在 auto_tuner 之前执行,若用户在 performance_mode=auto 下同时传 --cpu-offload-components,部分 auto 默认逻辑仍可能改写 legacy flag,需关注测试未覆盖组合。

影响范围集中在 sglang.multimodal_gen 模块:所有 Diffusion 组件加载路径(native、FSDP、transformer、VAE、encoder、bridge、adapter、vocoder、sound-tokenizer、upsampler、generic)与内存驻留管理、内存分析诊断均改为统一入口。用户侧获得可扩展的统一 offload 参数,对现有模型完全向后兼容 legacy flags;系统侧把“组件名 -> 是否卸载 -> 驻留策略”判定收敛到单点,后续新增 diffusion 模型无需改动四个旧 flag。团队侧需要理解 should_cpu_offload_component 单入口与 auto_tune 顺序,诊断日志不再给出具体 --xxx-cpu-offload 建议,改为提示调整 --cpu-offload-components 或 --layerwise-offload-components。DB/部署无外部依赖变更。

核心路径变更 需要 GPU/CI 验证 MPS 平台组合未覆盖 与 legacy flag 混用会启动失败 移除 lru_cache 后 module 引用生命周期

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论