融合算子使能 ​

openMind Library 已支持Ascend Extension for PyTorch插件torch_npu 提供的融合算子特性，让使用PyTorch框架的开发者更充分地释放昇腾AI处理器的算力。

开发者修改from_pretrained入参或者修改模型config.json即可使能融合算子，不需要修改模型代码。相较于不使能融合算子，性能会有较大提升（不同模型结构、不同超参效果有所不同）。

openMind Library当前支持的融合算子如下：

更多的模型和融合算子正在持续适配开发中

下文将分别介绍上述融合算子的使用方法。

npu_fusion_attention ​

该融合算子实现了“Transformer Attention Score”的融合计算，计算公式：

用户可以在此处查询该融合算子详细文档，在固定shape的场景中，可以较大幅度提升性能。

使用示例 ​

用户可以通过以下两种方式使用该融合算子

• 方法一：通过openMind Auto Classes加载，例如使用AutoModelForCausalLM实例化模型，传入_attn_implementation="npu_fusion_attention"参数即可。

  python

  from openmind import AutoModelForCausalLM
  
  model_id = "/your/model/path"
  
  model = AutoModelForCausalLM.from_pretrained(
      model_id,
      trust_remote_code=True,
      _attn_implementation="npu_fusion_attention"
      )
  

• 方法二：通过模型config.json文件传入"_attn_implementation": "npu_fusion_attention"。

  以模型AI-Research/internlm2-base-7b为例，将该模型下载到本地后，在模型文件夹中config.json文件中，添加"_attn_implementation": "npu_fusion_attention"字段。

  json

  {
  "architectures": [
      "InternLM2ForCausalLM"
  ],
  "auto_map": {
      "AutoConfig": "configuration_internlm2.InternLM2Config",
      "AutoModelForCausalLM": "modeling_internlm2.InternLM2ForCausalLM",
      "AutoModel": "modeling_internlm2.InternLM2ForCausalLM"
  },
  "bias": false,
  "bos_token_id": 1,
  "eos_token_id": 2,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 32768,
  "model_type": "internlm2",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pad_token_id": 2,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 1000000,
  "tie_word_embeddings": false,
  "_attn_implementation": "npu_fusion_attention",
  "torch_dtype": "bfloat16",
  "transformers_version": "4.41.0",
  "use_cache": true,
  "vocab_size": 92544,
  "pretraining_tp": 1
  }
  

  修改完成后单击保存，使用Auto Classes加载模型即可。此时from_pretrained方法不再需要传入_attn_implementation参数，参考如下：

  python

  from openmind import AutoModelForCausalLM
  
  model_id = "/your/local/model/path"
  
  model = AutoModelForCausalLM.from_pretrained(
      model_id,
      trust_remote_code=True,
      )
  

npu_rms_norm ​

RmsNorm算子是大模型常用的归一化操作，相比LayerNorm算子，其去掉了减去均值的部分 ,其计算公式为:

用户可以在此处查询该融合算子的详细文档。目前，该融合算子支持llama和Qwen2等系列模型。

使用示例 ​

用户可以通过如下方式使用该融合算子：

• 通过模型config.json文件传入"use_npu_rms_norm": true。

  以模型AI-Research/Llama-3.2-3B-Instruct为例，将该模型下载到本地后，在模型文件夹中config.json文件中，添加"use_npu_rms_norm": true字段。

{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": [
    128001,
    128008,
    128009
  ],
  "head_dim": 128,
  "hidden_act": "silu",
  "hidden_size": 3072,
  "initializer_range": 0.02,
  "intermediate_size": 8192,
  "max_position_embeddings": 131072,
  "mlp_bias": false,
  "model_type": "llama",
  "num_attention_heads": 24,
  "num_hidden_layers": 28,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": {
    "factor": 32.0,
    "high_freq_factor": 4.0,
    "low_freq_factor": 1.0,
    "original_max_position_embeddings": 8192,
    "rope_type": "llama3"
  },
  "rope_theta": 500000.0,
  "tie_word_embeddings": true,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.45.0.dev0",
  "use_cache": true,
  "use_npu_rms_norm": true,
  "vocab_size": 128256
}