transformer_impl.py — Mindspeed（华为 NPU）引擎

文件概述

基于 Mindspeed 适配器的训练引擎，继承自 MegatronEngineWithLMHead，专为华为 NPU 设备优化。文件非常简短（约 49 行）。

核心类

MindspeedEngineWithLMHead

@EngineRegistry.register(model_type="language_model", backend="megatron", device="npu")
class MindspeedEngineWithLMHead(MegatronEngineWithLMHead):
    def __init__(self, model_config, engine_config, ...):
        super().__init__(model_config, engine_config, ...)

        # 应用 Mindspeed 的 repatch（为 NPU 优化 Flash Attention 等）
        repatch_config = {"use_flash_attn": True}
        if self.engine_config.context_parallel_size > 1:
            repatch_config["context_parallel_size"] = ...
        repatch(repatch_config)

设计思路

注册时使用 device="npu"，当用户在 NPU 环境下使用 megatron 后端时，EngineRegistry 会自动选择这个引擎而非标准的 Megatron 引擎：

# 在 NPU 上自动选择 Mindspeed
engine_cls = EngineRegistry.get(model_type="language_model", backend="megatron", device="npu")
# → MindspeedEngineWithLMHead

与其他模块的关系

• 继承 engine/megatron/transformer_impl.py 的 MegatronEngineWithLMHead
• 使用 mindspeed 库的 repatch 函数进行 NPU 适配
• 通过 EngineRegistry 在 NPU 设备上自动启用

小结

Mindspeed 引擎是 Megatron 引擎的 NPU 适配层，通过最小化的代码改动实现了 NPU 兼容性。体现了引擎注册表机制的扩展性。