openmind/docs/zh/basic_tutorial/fused_ops.md

融合算子使能

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

开发者通过from openmind import apply_fused_kernel或者通过openmind-cli train即可使能融合算子，不需要修改模型代码。相较于不使能融合算子，性能会有较大提升（不同模型结构、不同超参效果有所不同）。

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

Model	npu_fusion_attention	rms_norm	npu_rotary_position_embedding	swiglu
Llama	✅	✅	✅	✅
Qwen2	✅	✅	✅	✅
Internlm2	✅	✅	✅	✅
Mistral	✅	✅	✅	✅

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

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

Flash Attention

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

attention\_out = Dropout(Softmax(Mask(scale*(pse + query*\mathbf{key}^\top), atten\_mask)), keep\_prob)* value

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

对于FA融合算子，当前openmind统一通过torch原生的sdpa接口调用，对于适配过的模型，使能后，sdpa走FA融合算子，不使能则会走transformers实现的eager模式；对于未适配的模型，其行为是默认行为，当前版本默认走sdpa接口，但sdpa的后端为小算子拼接，不保证性能。

RMSNorm

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

\text{RmsNorm}(x_i) = \frac{x_i}{\text{Rms}(\mathbf{x})} g_i,

where

\text{Rms}(\mathbf{x}) = \sqrt{\frac{1}{n} \sum_{i=1}^{n} x_i^2 + \epsilon},

用户可以在此处查询该融合算子的详细文档。

RoPE

RoPE (Rotary Positional Embedding，旋转式位置嵌入) 是一种位置编码技术，广泛应用于大型语言模型中，用于有效编码文本序列的位置信息。RoPE结合了绝对位置编码的稳定性与相对位置编码的灵活性，同时具备优秀的长度泛化能力。尽管RoPE已经在诸如LLaMA等多个前沿模型中得到采纳，但PyTorch框架目前尚未提供专门针对RoPE的实现与优化。因此，模型开发者通常需要通过自定义方式来实现RoPE，而这往往伴随着较高的计算和内存开销。

为了解决上述问题，openMind引入了针对Rotary Embedding的融合优化方案。通过将RoPE操作整合为单一算子，显著减少了数据传输次数和临时存储需求，进而优化了模型训练的性能，有效提升了RoPE在模型中的执行效率，该操作通过调用torch_npu侧的npu_rotary_mul接口实现。RoPE融合算子多被大模型用于Attention操作中，为了更好提升性能，建议在启用RoPE的同时，启用Flash Attention融合算子。

用户可以在此处查询该融合算子的详细文档。

SwiGLU

SwiGLU（Swish-Gated Linear Unit，Swish门控线性单元激活函数）常见于在LLaMA等大型语言模型中激活层。然而，由于PyTorch标准库中缺乏直接支持SwiGLU的算子接口，模型通常会以一系列基础算子组合的方式来实现SwiGLU功能，这种方式的执行效率并不理想。

为应对这一挑战，openMind对SwiGLU操作进行了融合优化，将其封装为一个高性能的融合算子，显著减少了数据在内存间的传输次数以及临时数据的存储需求，该操作通过调用torch_npu侧的npu_swiglu接口实现。

用户可以在此处查询该融合算子的详细文档。

使用示例

训练时使能融合算子

当通过openmind-cli train demo.yaml启动微调训练时，对于已适配融合算子的模型，openMind默认已开启所有已支持的融合算子，如需关闭，用户可在微调yaml配置文件中设置如下参数来单点关闭某个融合算子， 或者设置disable_fused_options: true以禁用融合算子功能（注：此参数只针对openmind-cli场景生效，对于通过SDK的外部使能，不支持该参数）：

### demo.yaml

disable_fused_options: true       # 默认值为false，设为true则禁用融合算子，且下面的融合算子使能开关失效

use_npu_fusion_attention: false   # 默认值为true，设为false则关闭使能flash attention融合算子
use_fused_rms_norm: false         # 默认值为true，设为false则关闭使能RMSNorm融合算子
use_fused_rope: false             # 默认值为true，设为false则关闭使能RoPE融合算子
use_fused_swiglu: false           # 默认值为true，设为false则关闭使能SwiGLU融合算子

# model
model_id: Qwen2.5-7B

# method
stage: pt
do_train: true
finetuning_type: full
deepspeed: examples/deepspeed/ds_z2_config.json

# dataset
dataset: text_zh_data
cutoff_len: 1024
packing: true

# output
output_dir: saves/qwen2.5_7b_full
logging_steps: 1
save_steps: 20000
overwrite_output_dir: true

# train
per_device_train_batch_size: 2
gradient_accumulation_steps: 2
learning_rate: 1.0e-5
lr_scheduler_type: cosine
warmup_ratio: 0.1
bf16: true
max_steps: 5000
seed: 1234

推理时使能融合算子

当通过openMind pipeline进行推理时，对于已适配融合算子的模型，只需要通过如下方式即可使能已支持的融合算子：

from openmind import pipeline, apply_fused_kernel

apply_fused_kernel()
# 调用apply_fused_kernel()即可使能所有已支持的融合算子
# 传入use_npu_fusion_attention=False可关闭flash attention使能
# 传入use_fused_rms_norm=False 可关闭RMSNorm使能

pipe = pipeline(
    task="text-generation",
    model="AI-Research/Qwen2.5-7B",
    framework="pt",
    backend="transformers",
    trust_remote_code=True,
    device="npu:0",
)
output = pipe("Give three tips for staying healthy.")
print(output)

# 输出
# Give three tips for staying healthy. 1. Eat a balanced diet: A healthy diet is essential for maintaining good health. It is important to consume a variety of foods from all the food groups, including fruits, vegetables, whole grains, lean proteins, and healthy fats. Avoid processed foods, sugary drinks, and excessive amounts of salt and sugar.
# 
# 2. Exercise regularly: Regular physical activity is important for maintaining good health. Aim for at least 30 minutes of moderate-intensity exercise most days of the week. This can include activities such as walking, jogging, cycling, swimming, or any other activity that gets your heart rate up and makes you breathe harder.
# 
# 3. Get enough sleep: Sleep is essential for good health. Aim for 7-9 hours of sleep each night. Establish a regular sleep schedule and create a relaxing bedtime routine to help you fall asleep more easily. Avoid using electronic devices before bed, as the blue light emitted by screens can interfere with your sleep.

注：由于transformers默认走sdpa，在外部不论有无使能apply_fused_kernel, 均会调用sdpa接口。但是使能后，openmind会对transformers的sdpa attention进行适配，适配后sdpa后端走npu FA融合算子，未适配的情况下，则是走小算子拼接。