BladeDISC 开源项目教程
项目介绍
BladeDISC 是一个端到端的动态形状编译器项目,专门为机器学习工作负载设计。它是阿里巴巴 PAI-Blade 的关键组件之一,提供对 TensorFlow 和 PyTorch 工作负载在 GPGPU 和 CPU 后端上的通用、透明且易于使用的性能优化。BladeDISC 原生支持动态形状工作负载,并在性能方面进行了多方面的考虑。
项目快速启动
对于 TensorFlow 用户
以下是启用 BladeDISC 的 TensorFlow 程序示例:
import blade_disc_tf as disc
disc.enable()
# 构建 TensorFlow 图并运行
g = tf.Graph()
with g.as_default():
with tf.Session() as sess:
sess.run()
更多信息请参考 QuickStart for TensorFlow Users。
对于 PyTorch 用户
以下是启用 BladeDISC 的 PyTorch 程序示例:
import torch_blade
from torch import nn
# 构建 PyTorch 模块
class MyModule(nn.Module):
pass
module = MyModule()
module.eval()
with torch.no_grad():
# blade_module 是 BladeDISC 优化后的模块
blade_module = torch_blade.optimize(module, allow_tracing=True, model_inputs=(x, y))
# 运行优化后的模块
blade_module(x, y)
更多信息请参考 QuickStart for PyTorch Users。
应用案例和最佳实践
BladeDISC 在生产环境中显示出高达 6.95 倍的加速比,相比于 PyTorch。与静态优化编译器(如 XLA 和 TensorRT)相比,BladeDISC 显示出可比或更好的性能。特别是在动态形状工作负载中,例如 BERT large 推理在 T4 GPU 上,BladeDISC 显示出 1.75 倍的加速。
典型生态项目
BladeDISC 与多个生态项目紧密结合,包括 TensorFlow、PyTorch 以及 MLIR。这些项目共同构成了一个强大的机器学习优化生态系统,支持从模型开发到部署的全过程优化。
- TensorFlow: BladeDISC 提供了对 TensorFlow 工作负载的优化支持。
- PyTorch: BladeDISC 提供了对 PyTorch 工作负载的优化支持。
- MLIR: BladeDISC 利用 MLIR 进行中间表示的优化和转换。
通过这些生态项目的结合,BladeDISC 能够为各种机器学习工作负载提供高效的性能优化解决方案。