Prhub

#36035 [Diffusion] Add component-scoped quantization overrides

原始 PR 作者 mickqian 合并时间 2026-08-25 09:20 文件变更 8 提交数 4 评论 2 代码增减 +124 / -9

执行摘要

新增组件级量化覆盖与忽略层配置

PR body 明确说明:串行化检查点的量化仍以 metadata 驱动,新的覆盖是为受支持的在线物化准备的,并不声称每个组件都支持每种格式,也不引入任何模型专属 CLI。实际诉求是让单个 diffusion 模型中的不同 DiT 组件(如 transformer 与 unconditional_transformer)能够选择不同的现有量化后端,并让原生的 MiniMax H3 Qwen3-VL 文本编码器也能使用现有 FP8 与 Kitchen INT8 在线量化能力。

值得 diffusion 与量化方向的工程师精读:核心设计是「串行化检查点保持 metadata 驱动、在线量化用显式覆盖、不支持则 fail closed」这三条边界;ServerArgs_extract_component_quantization_ignored_layers__post_init__ 的规范化/一致性校验模式也可作为新增组件级配置的参考范本。

讨论亮点

该 PR 没有可用的 code review 讨论线程,主要互动来自自动评论:mintlify bot 发布了文档预览链接;作者在 CI 状态更新后执行了 /tag-and-rerun-ci 触发重跑。需注意 Extra CI 与 AMD ROCm 7.2 CI 显示失败状态,而 NVIDIA 基础 CI 通过,说明跨平台 CI 稳定性仍是潜在关注点。

实现拆解

  1. 参数定义与规范化:在 python/sglang/multimodal_gen/runtime/server_args/server_args.pyServerArgs 中新增 component_quantization_ignored_layers: dict[str, list[str]] 字段,并在 __post_init__ 中做规范化:组件名去除首尾空白并把 - 归一为 _,层名统一为去空白的字符串列表;同时强制要求该配置必须搭配对应的 component_quantizations 条目,否则抛出 ValueError,保证配置自洽。
  2. CLI 参数提取:在 ServerArgs 中新增静态方法 _extract_component_quantization_ignored_layers,解析形如 --component-quantization-ignored-layers.<component> 的动态未知参数,支持 --key=value--key value 两种形式以及同一组件多个层名连续传入,并在 from_cli_args 中与已有的 _extract_component_quantizations 等逻辑串联,确保配置与 CLI 两路入口行为一致。
  3. 加载器透传:在 transformer_loader.py_server_args_for_transformer_component 中,把组件级 ignored layers 覆盖到副本 quantization_ignored_layers,与已有的 weights 路径、quantization 覆盖一起参与组件级 server args 构造;在 text_encoder_loader.py 中为 _configure_encoder_quantization_resolve_and_configure_encoder_quantization 增加 ignored_layers 参数,并在 load_customized 中从 server args 取出传给量化配置构造器(如 get_quantization_config(...)(ignored_layers=...))。
  4. 测试与文档配套test_server_args.py 覆盖 CLI 解析与配置保留路径,test_text_encoder_loader.py 验证 KitchenInt8Config.ignored_layers 被正确设置,test_ideogram4.py 补充组件级覆盖场景;docs/docs/sglang-diffusion/quantization.mdxapi/cli.mdx 补充新参数的说明。
文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/server_args/server_args.py 参数解析 modified 7.0
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py 加载器 modified 5.54
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py 加载器 modified 5.27
python/sglang/multimodal_gen/test/unit/test_ideogram4.py 测试 modified 4.06
python/sglang/multimodal_gen/test/unit/test_server_args.py 测试 modified 3.99
python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py 测试 modified 3.82
docs/docs/sglang-diffusion/quantization.mdx 文档 modified 3.19
docs/docs/sglang-diffusion/api/cli.mdx 文档 modified 2.45

关键符号

_extract_component_quantization_ignored_layers _server_args_for_transformer_component _configure_encoder_quantization _resolve_and_configure_encoder_quantization

关键源码片段

python/sglang/multimodal_gen/runtime/server_args/server_args.py core-logic

核心参数定义与 CLI 解析入口,新增 component_quantization_ignored_layers 字段、规范化校验逻辑及 _extract_component_quantization_ignored_layers 方法,是本次功能的配置源头。

@staticmethod
def _extract_component_quantization_ignored_layers(
    unknown_args: list[str],
) -> tuple[dict[str, list[str]], list[str]]:
    ignored_layers: dict[str, list[str]] = {}
    remaining: list[str] = []
    i = 0
    prefixes = (
        "--component-quantization-ignored-layers.",
        "--component_quantization_ignored_layers.",
    )
    # 使用 while 循环逐个扫描 CLI 剩余参数,因为单个组件可能带多个层名,
    # 需要把连续的不以 "-" 开头的参数都收集为该组件的忽略层列表。
    while i < len(unknown_args):
        arg = unknown_args[i]
        # 同时支持 "--key=value" 与 "--key value" 两种书写形式。
        key_part = arg.split("=", 1)[0] if "=" in arg else arg
        prefix = next(
            (candidate for candidate in prefixes if key_part.startswith(candidate)),
            None,
        )
        if prefix is None:
            remaining.append(arg)
            i += 1
            continue
        # 组件名统一把 "-" 换成 "_",与 component_quantizations 的规范化保持一致。
        component = key_part[len(prefix):].replace("-", "_")
        if "=" in arg:
            values = [arg.split("=", 1)[1]]
        else:
            values = []
            # 支持多个层名连续传入,直到遇到下一个参数为止。
            while i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith("-"):
                i += 1
                values.append(unknown_args[i])
        if component and values:
            ignored_layers[component] = values
        else:
            remaining.append(arg)
        i += 1
    return ignored_layers, remaining
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py core-logic

Transformer 组件加载路径中透传组件级 ignored layers,确保覆盖能进入 quantization_ignored_layers 并影响在线量化行为。

def _server_args_for_transformer_component(
    server_args: ServerArgs, component_name: str
) -> ServerArgs:
    """为次级 transformer 组件屏蔽全局量化覆盖标志。"""
    component_weights_path = server_args.component_weights_paths.get(component_name)
    component_quantization = server_args.component_quantizations.get(component_name)
    component_ignored_layers = server_args.component_quantization_ignored_layers.get(
        component_name
    )
    # 只要存在任一组件级量化覆盖,就为当前组件构造独立副本,
    # 避免全局覆盖泄漏到其它组件(如 unconditional_transformer)。
    if (
        component_weights_path is not None
        or component_quantization is not None
        or component_ignored_layers is not None
    ):
        component_server_args = copy.copy(server_args)
        if component_weights_path is not None:
            component_server_args.transformer_weights_path = component_weights_path
            component_server_args.nunchaku_config = None
            logger.info(
                "Using transformer_weights_path override for %s: %s",
                component_name,
                component_weights_path,
            )
        if component_quantization is not None:
            component_server_args.quantization = component_quantization
            logger.info(
                "Using quantization override %s for %s",
                component_quantization,
                component_name,
            )
        # 组件级忽略层直接映射到全局 quantization_ignored_layers 字段,
        // 由下游量化配置构造器消费
        if component_ignored_layers is not None:
            component_server_args.quantization_ignored_layers = component_ignored_layers
        return component_server_args
​
    # 非特殊组件且无覆盖时直接复用原 server args。
    if component_name not in ("transformer_2", "unconditional_transformer"):
        return server_args
    ...

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

  1. 参数解析兼容性_extract_component_quantization_ignored_layers 依赖未知参数扫描,若与其他 component 前缀参数顺序冲突或参数带引号,可能产生残留参数进入 remaining,导致 CLI 解析错误。
  2. fail-closed 校验过严__post_init__ 要求 ignored layers 必须匹配 quantization override,对于只配置 weights path 而想沿用全局 quant 的场景可能误伤,需要用户明确同时提供 quant override。
  3. 透传链完整性:ignored layers 需要从 server args 一路传到 KitchenInt8Config 等量化配置,任何一层缺失都会静默丢失设置;当前测试覆盖了 transformer 与 text encoder 两条主链,但尚未覆盖所有组件加载器。
  4. CI 失败:Extra 与 AMD ROCm 7.2 CI 未通过,可能暴露跨平台问题或仅为基础设施不稳定,合并前未完全澄清。

影响范围集中在 sglang-diffusion 运行路径:多组件扩散模型(如 Ideogram 4 的 transformer/unconditional_transformer)可以分别配置量化后端与忽略层,原生文本编码器(如 MiniMax H3 的 Qwen3-VL 编码器)可获得 FP8/Kitchen INT8 在线量化能力。对既有用户,全局 --quantization--quantization-ignored-layers 作为 legacy 接口保持不变,行为向后兼容。对团队而言,该 PR 扩展了 ServerArgs 的组件级配置家族(weights/quantization/attention backend 之后又加 ignored layers),后续新增组件级参数可复用同一套提取与校验模式。

新增参数校验为 fail-closed CI 扩展与 AMD 测试未通过 参数透传链较长

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论