#24429 Support NemotronHPuzzleForCausalLM

原始 PR 作者 netanel-haber 合并时间 2026-05-28 07:12 文件变更 7 提交数 6 评论 12 代码增减 +74 / -7

执行摘要

支持 NemotronHPuzzleForCausalLM 模型架构

支持NemotronHPuzzleForCausalLM模型变种，该变种使用逐层‘拼图’配置，可以针对每层定义不同的块类型（如Attention、MoE、Mamba）和专家数量，提供更灵活的模型定义。该PR主要从NVIDIA的模型需求出发，使SGLang能部署此类模型。

值得精读，特别是配置类如何支持逐层定制的设计模式，以及 MTP 集成方案。建议关注其中的 assert 与 ValueError 的取舍，在后续类似场景下优先使用异常。建议在合并后补充单元测试。

讨论亮点

Assertion 使用风险：Reviewer gemini-code-assist 建议将 max_n_routed_experts 中的 assert 改为 raise ValueError 以防止生产环境跳过验证。该建议未被采纳，代码保持使用 assert。
滑动窗口参数来源：Reviewer roikoren755 提议 sliding_window_size 应使用 layer_config.sliding_window 而非 config.sliding_window。作者回复此时 config 已是层配置，两者等价，无需修改。该点已解决。

实现拆解

新增配置类NemotronHPuzzleConfig：在configs/nemotron_h.py中基于NemotronHConfig派生，新增block_configs列表（每层配置字典）和mtp_block_configs。重写get_nemotron_h_config_for_layer方法对指定层返回配置副本并覆写独有属性；重写get_mtp_config返回MTP块配置；增加max_n_routed_experts属性遍历所有MoE块取最大值。
修改DecoderLayer初始化：在models/nemotron_h.py中，将NemotronHMoEDecoderLayer和NemotronHAttentionDecoderLayer等类的__init__改为调用config.get_nemotron_h_config_for_layer(layer_idx)获取该层配置，并传入Mixer模块。同时为NemotronHAttention的初始化补充sliding_window_size参数，并将权重加载中的num_experts来源改为config.max_n_routed_experts。
调整MTP子模块：在models/nemotron_h_mtp.py中，在MTP模块初始化前调用config = config.get_mtp_config()，使MTP子模块继承正确的块配置列表。
注册新架构：在configs/model_config.py的MTP路由白名单、server_args.py的模型特定调整判断、configs/__init__.py的导出列表和utils/hf_transformers/common.py中，将NemotronHPuzzleForCausalLM添加进对应集合，确保其被当作NemotronH系列模型统一处理。

文件	模块	状态	重要度
`python/sglang/srt/configs/nemotron_h.py`	配置层	modified	8.12
`python/sglang/srt/models/nemotron_h.py`	模型层	modified	6.77
`python/sglang/srt/configs/model_config.py`	模型配置	modified	5.34
`python/sglang/srt/configs/__init__.py`	配置导出	modified	4.59
`python/sglang/srt/models/nemotron_h_mtp.py`	MTP 模块	modified	4.96
`python/sglang/srt/server_args.py`	服务器参数	modified	4.49
`python/sglang/srt/utils/hf_transformers/common.py`	工具函数	modified	3.92

关键符号

get_nemotron_h_config_for_layer get_mtp_config max_n_routed_experts NemotronHPuzzleConfig.__init__ NemotronHPuzzleForCausalLM

关键源码片段

python/sglang/srt/configs/nemotron_h.py core-logic

核心新增 NemotronHPuzzleConfig 类，定义逐层配置机制；修改基类 __init__ 增加 * 强制关键字参数

# python/sglang/srt/configs/nemotron_h.py

class NemotronHPuzzleConfig(NemotronHConfig):
    """允许逐层指定 block 配置的 NemotronH 配置类 ."""
    model_type = "nemotron_h_puzzle"
    has_no_defaults_at_init = True

    def __init__(
        self,
        *,
        block_configs: list[dict[str, Any]],
        mtp_block_configs: list[dict[str, Any]] | None = None,
        **kwargs,
    ):
        # 基类 __init__ 已改为仅关键字参数 （星号确保所有参数为 keyword-only）
        super().__init__(**kwargs)
        self.block_configs = block_configs # 每层配置字典列表
        self.mtp_block_configs = mtp_block_configs # MTP 层的可选配置

    def get_nemotron_h_config_for_layer(self, layer_idx: int) -> NemotronHConfig:
        """为指定层返回一个配置副本，并覆写该层特定的属性。"""
        layer_config = copy.copy(self)
        for key, value in self.block_configs[layer_idx].items():
            setattr(layer_config, key, value)
        return layer_config

    def get_mtp_config(self) -> NemotronHConfig:
        """返回 MTP 子模块的配置（使用 mtp_block_configs）。"""
        assert self.mtp_block_configs
        mtp_config = copy.copy(self)
        mtp_config.block_configs = self.mtp_block_configs
        return mtp_config

    @property
    def max_n_routed_experts(self) -> int:
        """计算所有 MoE 块中最大的 n_routed_experts。"""
        block_n_routed_experts = [
            block["n_routed_experts"]
            for block in self.block_configs
            if block["block_type"] == "moe" # 只考虑 MoE 块
        ]
        max_experts = max(block_n_routed_experts)
        assert max_experts > 0 # 注意：使用 assert，生产环境可能被禁用 ，建议改用 raise ValueError
        return max_experts

python/sglang/srt/models/nemotron_h.py data-contract

修改 DecoderLayer 使用层特定配置；新增 NemotronHPuzzleForCausalLM 模型类；调整注意力滑动窗口参数

# python/sglang/srt/models/nemotron_h.py （关键改动片段）

class NemotronHMoEDecoderLayer(nn.Module):
    def __init__(self, config, layer_idx, quant_config=None, prefix=""):
        super().__init__()
        # 获取该层特定配置
        layer_config = config.get_nemotron_h_config_for_layer(layer_idx)
        self.mixer = NemotronHMoE(
            layer_config, # 传入层特定配置
            layer_idx=layer_idx,
            quant_config=quant_config,
            prefix=f"{prefix}.mixer",
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

class NemotronHAttentionDecoderLayer(nn.Module):
    def __init__(self, config, layer_idx, quant_config=None, prefix=""):
        super().__init__()
        # 同样获取层特定配置
        layer_config = config.get_nemotron_h_config_for_layer(layer_idx)
        self.mixer = NemotronHAttention(
            layer_config,
            layer_idx,
            quant_config,
            prefix=f"{prefix}.mixer",
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

# 新增的模型类，继承所有前向逻辑
class NemotronHPuzzleForCausalLM(NemotronHForCausalLM):
    pass

# EntryClass 注册新增架构
EntryClass = [NemotronHForCausalLM, NemotronHPuzzleForCausalLM]

评论区精华

Assertion in max_n_routed_experts 正确性

Reviewer gemini-code-assist 指出使用 assert 进行验证可能在生产环境因 -O 标志被跳过，建议改为 raise ValueError。

结论：作者未修改，代码保持 assert（见 head 版本） · unresolved

sliding_window_size should use layer_config 正确性

Reviewer roikoren755 建议滑动窗口参数应使用 layer_config.sliding_window 而非 config.sliding_window。

结论：作者回复此时 config 已通过 get_nemotron_h_config_for_layer 返回 layer_config，两者等价，无需修改。PR 保持原样。 · 已解决

风险与影响

配置向后兼容风险：NemotronHConfig.__init__ 增加了 * 强制关键字参数，原有位置参数调用方式将报错，需确保所有调用方改为关键字参数。
断言验证缺陷：max_n_routed_experts 使用 assert 进行验证，断言在 -O 优化模式下被跳过，可能导致生产环境使用错误的专家数。
测试覆盖不足：未提供针对 NemotronHPuzzleConfig 或 NemotronHPuzzleForCausalLM 的单元测试或集成测试，新路径边缘情况未覆盖。
性能开销：get_nemotron_h_config_for_layer 每层通过 copy.copy 复制配置并设置属性，增加少量初始化开销，但属一次性成本。

用户：能够部署和推理 NemotronHPuzzle 模型，获得更灵活的分层配置能力；现有 NemotronH 模型不受影响。
系统：模型注册和配置解析路径增加新分支；MTP 模块新增 get_mtp_config() 调用路径，不影响已有 MTP 使用。
团队：维护成本小幅增加，需要跟踪新架构的演进，建议后续补充测试和验证。

配置向后兼容风险断言可被禁用缺少测试覆盖新模型路径未经验证

关联 Issue

未识别关联 Issue

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

完整报告

执行摘要

一句话：支持 NemotronHPuzzleForCausalLM 模型架构
推荐动作：值得精读，特别是配置类如何支持逐层定制的设计模式，以及 MTP 集成方案。建议关注其中的 assert 与 ValueError 的取舍，在后续类似场景下优先使用异常。建议在合并后补充单元测试。

功能与动机

实现拆解

新增配置类NemotronHPuzzleConfig：在configs/nemotron_h.py中基于NemotronHConfig派生，新增block_configs列表（每层配置字典）和mtp_block_configs。重写get_nemotron_h_config_for_layer方法对指定层返回配置副本并覆写独有属性；重写get_mtp_config返回MTP块配置；增加max_n_routed_experts属性遍历所有MoE块取最大值。
修改DecoderLayer初始化：在models/nemotron_h.py中，将NemotronHMoEDecoderLayer和NemotronHAttentionDecoderLayer等类的__init__改为调用config.get_nemotron_h_config_for_layer(layer_idx)获取该层配置，并传入Mixer模块。同时为NemotronHAttention的初始化补充sliding_window_size参数，并将权重加载中的num_experts来源改为config.max_n_routed_experts。
调整MTP子模块：在models/nemotron_h_mtp.py中，在MTP模块初始化前调用config = config.get_mtp_config()，使MTP子模块继承正确的块配置列表。
注册新架构：在configs/model_config.py的MTP路由白名单、server_args.py的模型特定调整判断、configs/__init__.py的导出列表和utils/hf_transformers/common.py中，将NemotronHPuzzleForCausalLM添加进对应集合，确保其被当作NemotronH系列模型统一处理。

关键文件：

python/sglang/srt/configs/nemotron_h.py（模块配置层；类别 source；类型 core-logic；符号 get_nemotron_h_config_for_layer, get_mtp_config, max_n_routed_experts, NemotronHPuzzleConfig）: 核心新增 NemotronHPuzzleConfig 类，定义逐层配置机制；修改基类 init 增加 * 强制关键字参数
python/sglang/srt/models/nemotron_h.py（模块模型层；类别 source；类型 data-contract；符号 NemotronHPuzzleForCausalLM）: 修改 DecoderLayer 使用层特定配置；新增 NemotronHPuzzleForCausalLM 模型类；调整注意力滑动窗口参数
python/sglang/srt/configs/model_config.py（模块模型配置；类别 source；类型 data-contract）: 在 MTP 路由白名单中加入 NemotronHPuzzleForCausalLM，确保 draft 模型映射正确
python/sglang/srt/configs/__init__.py（模块配置导出；类别 source；类型 dependency-wiring）: 导出 NemotronHPuzzleConfig，使配置类可被外部引用
python/sglang/srt/models/nemotron_h_mtp.py（模块 MTP模块；类别 source；类型 data-contract）: 在 MTP 父模块初始化前调用 get_mtp_config()，确保 MTP 层使用正确的块配置
python/sglang/srt/server_args.py（模块服务器参数；类别 source；类型 core-logic）: 在模型特定调整判断中加入 NemotronHPuzzleForCausalLM，应用 NemotronH 默认参数
python/sglang/srt/utils/hf_transformers/common.py（模块工具函数；类别 source；类型 core-logic）: 添加识别 NemotronHPuzzleForCausalLM 架构的映射，确保模型加载正确

关键符号：get_nemotron_h_config_for_layer, get_mtp_config, max_n_routed_experts, NemotronHPuzzleConfig.init, NemotronHPuzzleForCausalLM

关键源码片段

`python/sglang/srt/configs/nemotron_h.py`

核心新增 NemotronHPuzzleConfig 类，定义逐层配置机制；修改基类 init 增加 * 强制关键字参数

# python/sglang/srt/configs/nemotron_h.py

class NemotronHPuzzleConfig(NemotronHConfig):
    """允许逐层指定 block 配置的 NemotronH 配置类 ."""
    model_type = "nemotron_h_puzzle"
    has_no_defaults_at_init = True

    def __init__(
        self,
        *,
        block_configs: list[dict[str, Any]],
        mtp_block_configs: list[dict[str, Any]] | None = None,
        **kwargs,
    ):
        # 基类 __init__ 已改为仅关键字参数 （星号确保所有参数为 keyword-only）
        super().__init__(**kwargs)
        self.block_configs = block_configs # 每层配置字典列表
        self.mtp_block_configs = mtp_block_configs # MTP 层的可选配置

    def get_nemotron_h_config_for_layer(self, layer_idx: int) -> NemotronHConfig:
        """为指定层返回一个配置副本，并覆写该层特定的属性。"""
        layer_config = copy.copy(self)
        for key, value in self.block_configs[layer_idx].items():
            setattr(layer_config, key, value)
        return layer_config

    def get_mtp_config(self) -> NemotronHConfig:
        """返回 MTP 子模块的配置（使用 mtp_block_configs）。"""
        assert self.mtp_block_configs
        mtp_config = copy.copy(self)
        mtp_config.block_configs = self.mtp_block_configs
        return mtp_config

    @property
    def max_n_routed_experts(self) -> int:
        """计算所有 MoE 块中最大的 n_routed_experts。"""
        block_n_routed_experts = [
            block["n_routed_experts"]
            for block in self.block_configs
            if block["block_type"] == "moe" # 只考虑 MoE 块
        ]
        max_experts = max(block_n_routed_experts)
        assert max_experts > 0 # 注意：使用 assert，生产环境可能被禁用 ，建议改用 raise ValueError
        return max_experts

`python/sglang/srt/models/nemotron_h.py`

修改 DecoderLayer 使用层特定配置；新增 NemotronHPuzzleForCausalLM 模型类；调整注意力滑动窗口参数

# python/sglang/srt/models/nemotron_h.py （关键改动片段）

class NemotronHMoEDecoderLayer(nn.Module):
    def __init__(self, config, layer_idx, quant_config=None, prefix=""):
        super().__init__()
        # 获取该层特定配置
        layer_config = config.get_nemotron_h_config_for_layer(layer_idx)
        self.mixer = NemotronHMoE(
            layer_config, # 传入层特定配置
            layer_idx=layer_idx,
            quant_config=quant_config,
            prefix=f"{prefix}.mixer",
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

class NemotronHAttentionDecoderLayer(nn.Module):
    def __init__(self, config, layer_idx, quant_config=None, prefix=""):
        super().__init__()
        # 同样获取层特定配置
        layer_config = config.get_nemotron_h_config_for_layer(layer_idx)
        self.mixer = NemotronHAttention(
            layer_config,
            layer_idx,
            quant_config,
            prefix=f"{prefix}.mixer",
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

# 新增的模型类，继承所有前向逻辑
class NemotronHPuzzleForCausalLM(NemotronHForCausalLM):
    pass

# EntryClass 注册新增架构
EntryClass = [NemotronHForCausalLM, NemotronHPuzzleForCausalLM]

评论区精华

Assertion 使用风险：Reviewer gemini-code-assist 建议将 max_n_routed_experts 中的 assert 改为 raise ValueError 以防止生产环境跳过验证。该建议未被采纳，代码保持使用 assert。
滑动窗口参数来源：Reviewer roikoren755 提议 sliding_window_size 应使用 layer_config.sliding_window 而非 config.sliding_window。作者回复此时 config 已是层配置，两者等价，无需修改。该点已解决。
- Assertion in max_n_routed_experts (correctness): 作者未修改，代码保持 assert（见 head 版本）
- sliding_window_size should use layer_config (correctness): 作者回复此时 config 已通过 get_nemotron_h_config_for_layer 返回 layer_config，两者等价，无需修改。PR 保持原样。

风险与影响

风险：
- 配置向后兼容风险：NemotronHConfig.__init__ 增加了 * 强制关键字参数，原有位置参数调用方式将报错，需确保所有调用方改为关键字参数。
- 断言验证缺陷：max_n_routed_experts 使用 assert 进行验证，断言在 -O 优化模式下被跳过，可能导致生产环境使用错误的专家数。
- 测试覆盖不足：未提供针对 NemotronHPuzzleConfig 或 NemotronHPuzzleForCausalLM 的单元测试或集成测试，新路径边缘情况未覆盖。
- 性能开销：get_nemotron_h_config_for_layer 每层通过 copy.copy 复制配置并设置属性，增加少量初始化开销，但属一次性成本。
影响：
- 用户：能够部署和推理 NemotronHPuzzle 模型，获得更灵活的分层配置能力；现有 NemotronH 模型不受影响。
- 系统：模型注册和配置解析路径增加新分支；MTP 模块新增 get_mtp_config() 调用路径，不影响已有 MTP 使用。
- 团队：维护成本小幅增加，需要跟踪新架构的演进，建议后续补充测试和验证。
- 风险标记：配置向后兼容风险, 断言可被禁用, 缺少测试覆盖, 新模型路径未经验证

关联脉络

暂无明显关联 PR

#24429 Support NemotronHPuzzleForCausalLM

执行摘要

支持 NemotronHPuzzleForCausalLM 模型架构

实现拆解

评论区精华

风险与影响

关联 Issue

未识别关联 Issue

完整报告

参与讨论