1. FlagGems 简介

FlagGems 是使用 OpenAI 推出的 Triton编程语言 实现的高性能通用算子库，旨在为大语言模型提供一系列可应用于 PyTorch 框架的算子，加速模型的推理与训练。

FlagGems 通过对 PyTorch 的后端 aten 算子进行覆盖重写，默认支持 pytorch 的 eager 模式，实现算子库的无缝替换，使用户能够在不修改模型代码的情况下平稳地切换到 Triton 算子库。FlagGems 不会影响 aten 后端的正常使用，且可以带来良好的性能提升。Triton语言为算子库提供了更好的可读性和易用性，且针对多种算子可以带来良好的性能提升，因此开发者只需付出较低的学习成本，即可参与 FlagGems 的算子开发与建设。

2. 安装使用 FlagGems

2.1 FlagGems 所需依赖

Triton >= 2.2.0PyTorch >= 2.2.0Transformers >= 4.40.2

2.2 安装 FlagGems​​​​​​​

git clone https://github.com/FlagOpen/FlagGems.gitcd FlagGemspip install .

2.3 全局替换FlagGems算子​​​​​​​

import flag_gemsflag_gems.enable()

2.4 局部替换FlagGems算子​​​​​​​

import flag_gemswith flag_gems.use_gems():    pass

2.5 FlagGems 使用示例​​​​​​​

import torchimport flag_gemsM, N, K = 1024, 1024, 1024A = torch.randn((M, K), dtype=torch.float16, device="cuda")B = torch.randn((K, N), dtype=torch.float16, device="cuda")with flag_gems.use_gems():    C = torch.mm(A, B)

3. 参与贡献

• 新建一个 Issue 来反馈遇到的 Bug 或者提出新功能需求（https://github.com/FlagOpen/FlagGems/issues/new/choose）

• 提一个 Pull Request 来修复一个 Bug 或者实现一个新功能（https://github.com/FlagOpen/FlagGems/compare）

3.1 代码贡献流程

首先需要注册一个 GitHub 帐号，再执行以下步骤。如果你已经熟悉 GitHub 相关操作，建议重点浏览CI 相关的内容。

3.1.1 Fork 仓库并 clone 代码到本地

打开 FlagGems GitHub 首页（https://github.com/FlagOpen/FlagGems），单击 Fork 按钮创建一个仓库副本​​​​​​​

# Clone FlagGems 仓库到本地并进入 FlagGems 文件夹git clone https://github.com/FlagOpen/FlagGemscd FlagGems

3.1.2 创建本地分支

你需要从 FlagGems 的 master 分支上创建新分支来开发代码。​​​​​​​

# 创建并切换到一个名为 master 的分支 git checkout -b master

3.1.3 安装 pre-commit

FlagGems 使用 pre-commit（https://pre-commit.com） 的 git hooks 格式化源代码，在调用 git commit 命令时进行代码静态检查，并且 pre-commit 测试也是 CI 的一部分，不通过检查的 Pull Request 不能被提交到 FlagGems​​​​​​​

pip install pre-commitpre-commit installpre-commit

3.1.4 开发代码

在提交拉取请求时，贡献者应描述所做的更改以及原因。如果可以设计测试用例，请提供相应测试。拉取请求在合并前需要一位成员的批准，而且需要通过代码的持续集成检查。

目前持续集成检查设有四条流水线:

• 代码格式检查

• 算子单元测试

• 模型测试

• 代码覆盖率检查

随着 FlagGems 的代码量增大，会有更多的测试检查加入集成测试，可以根据贡献指南（https://github.com/FlagOpen/FlagGems/blob/master/CONTRIBUTING_cn.md）查看最新的说明。可通过 git status 查看当前目录的变化，通过 git diff 查看文件具体被修改的内容。

3.1.5 本地执行单元测试以及性能测试

1. 算子正确性测试​​​​​​​

cd testspytest test_xx_ops.py                # on CUDApytest test_xx_ops.py --device cpu   # on CPU

2. 模型正确性测试​​​​​​​

cd examplespytest model_xx_test.py

3.算子性能测试​​​​​​​

cd benchmarkpytest test_xx_perf.py -s            # kernelpytest test_xx_perf.py -s --mode cpu # e2e

4. 运行时打印日志信息（性能测试不建议打开）

pytest program.py --log-cli-level debug 

3.1.6 提交代码

1. 提交本地更改

git commit 前先使用 git add 添加变更或新增的文件。如果要撤销某个文件的变更，可使用 git checkout --。每次提交可通过 git commit -m 撰写提交说明

注意：git commit 会触发代码静态检测，检查有问题则代码提交失败

2. 保持本地仓库最新

发起 Pull Request 合入代码之前，需要同步原仓库最新的代码防止代码冲突

3. Push 本地的修改到远程仓库

4. 提交 Pull Request（PR）合入代码

打开 fork 的 FlagGems 页面，并切换到所建分支，然后单击 Compare & pull request 按钮

5. 确保通过 CI 测试

提交 Pull Request 后会触发 Github Action 进行 CI（Continuous Integration，持续集成）测试，并且之后每提交一次代码合入（git push）都会触发一次 CI 测试

6. Code Review

CI 测试通过后，请等待 Code Review。收到 Code Review 意见后，请回复评审人的意见，并根据意见修改代码。

7. PR Merge 后本次贡献结束

扫码回复“Triton”

加入Triton中文社区交流群