基于 Triton 的大模型优化

发表于 分类于 本文字数： 445

矩阵乘，flash attention和应用

镜像收藏：mlu370_ubuntu22.04-student-1.22.1:v5.2

Triton 入门:矩阵乘的实现

实验目的：

掌握使用智能编程语言 BCL 进行算子开发、编译扩展高性能库算子,并集成到 PyTorch 框架中的方法和流程。能够用 BCL 实现 Sigmoid 算子,并集成进 PyTorch 的推理网络高效 地运行在 MLU 硬件上。

代码文件组织：

代码补全

Triton 进阶:Flash Attention 的实现

Triton 实战:模型应用的实现

需要挂载模型卷：模型收藏中的large-scale-models:model-v1

运行环境：

  • 实验目录: /opt/code_chap5_new/exp_5_3_application/
  • 模型目录: /workspace/model/favorite/large-scale-models/model-v1/

三个应用的代码分别存 放在/opt/code_chap5/exp_5_1_deepseek/、/opt/code_chap5/exp_5_2_stable_diffusion/和/opt/code_chap5/exp_5 录下

以 Llama 3.1 为基座模型的 DeepSeek 蒸馏模型 DeepSeek-R1-Distill-Llama-8B 的算子替换需要修改/opt/tools/native/transformers_mlu/src/transformers/models/llama /modeling_llama.py 文件;文生图模型 Stable Diffusion 的算子替换需要 ldm/modules/attention.py 文件;以 Qwen 2.5 为基座模型的 Verilog 代码生成模型 CodeV-QC 的算子替换需要修改/opt/tools /native/transformers_mlu/src/transformers/models/qwen2/modeling_qwen2.py 文件

代码补全

flash_attention_triton_opt.py

文件位于三个目录下，内容相同：

  • exp_5_1_deepseek/flash_attention_triton_opt.py
  • exp_5_2_stable_diffusion/flash_attention_triton_opt.py
  • exp_5_3_codev/flash_attention_triton_opt.py

DeepSeek: modeling_llama.py

路径: /opt/tools/native/transformers_mlu/src/transformers/models/llama/modeling_llama.py

CodeV: modeling_qwen2.py

路径: /opt/tools/native/transformers_mlu/src/transformers/models/qwen2/modeling_qwen2.py

Stable Diffusion: attention.py

路径: /opt/code_chap5_new/exp_5_3_application/exp_5_2_stable_diffusion/ldm/modules/attention.py

实验运行

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 实现单算子测试
python flash_attention_triton_opt.py

# 运行 DeepSeek 聊天机器人程序
cd /opt/code_chap5/exp_5_1_deepseek/
python testdeepseek.py

# 运行图生图推理程序
cd /opt/code_chap5/exp_5_2_stable_diffusion/
bash run-stable-diffusion-img2img.sh

# 运行文生图推理程序
bash run-stable-diffusion-txt2img.sh

# 运行图像修复推理程序
bash run-stable-diffusion-painting.sh

# 运行 Verilog 代码生成程序
cd /opt/code_chap5/exp_5_3_codev/
bash run_cli.sh