Hidet: 开源深度学习编译器助力模型高效部署

Hidet简介

Hidet是一款强大的开源深度学习编译器,由Python编写而成。它的主要目标是简化在现代加速器(如NVIDIA GPU)上实现高性能深度学习算子的过程。Hidet的出现为深度学习从业者提供了一个新的选择,特别是对于那些需要实现高度优化的自定义算子的用户来说,Hidet无疑是一个极具吸引力的选项。

Hidet的主要特性

端到端编译:支持将PyTorch和ONNX模型编译为高效的CUDA内核。
图级优化:对模型的计算图进行优化,提升整体性能。
算子级优化:针对单个算子进行优化,进一步提升执行效率。
Python API:提供友好的Python接口,便于集成和使用。
自动调优:能够自动搜索最佳的算子调度方案。

Hidet的安装与使用

安装Hidet

Hidet的安装非常简单,只需要通过pip包管理器即可完成:

pip install hidet

需要注意的是,Hidet目前主要针对NVIDIA GPU上的推理工作负载进行优化,因此在使用时需要满足以下条件:

Linux操作系统
CUDA Toolkit 11.6+
Python 3.8+

对于想尝试最新功能的用户,也可以安装每日构建版本或从源代码构建。

使用Hidet优化PyTorch模型

以下是一个使用Hidet优化预训练ResNet18模型的简单示例:

import torch
import hidet

# 加载预训练的ResNet18模型
model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).cuda().eval()
x = torch.rand(1, 3, 224, 224).cuda()

# 配置Hidet优化选项
hidet.torch.dynamo_config.use_tensor_core(True)
hidet.torch.dynamo_config.search_space(2)

# 使用Hidet编译模型
model_opt = torch.compile(model, backend='hidet')

# 运行优化后的模型
y = model_opt(x)

在这个例子中,我们首先加载了一个预训练的ResNet18模型,然后使用Hidet对其进行编译优化。通过设置use_tensor_core(True),我们允许Hidet生成利用NVIDIA GPU上Tensor Core的CUDA内核。而search_space(2)则允许Hidet搜索最适合当前硬件和输入尺寸的算子调度方案。

Hidet的工作原理

Image 1: Eager vs Hidet latency

Hidet的工作原理可以概括为以下几个步骤:

模型转换:将PyTorch或ONNX模型转换为Hidet的内部图表示。
图级优化:对整个计算图进行优化,如算子融合、内存分配优化等。
算子级优化:针对单个算子进行优化,包括自动调优、张量核心利用等。
代码生成:生成优化后的CUDA内核代码。
运行时优化:使用CUDA Graph等技术减少框架级开销。

通过这一系列的优化步骤,Hidet能够显著提升模型的推理性能。如上图所示,在AWS g5.2xlarge实例上,Hidet优化后的ResNet50模型相比于PyTorch eager模式有明显的性能提升。

Hidet Script:自定义算子的利器

Hidet不仅提供了对现有模型的优化能力,还引入了Hidet Script,这是一种在Python中实现张量算子的方法。以下是使用Hidet Script实现简单矩阵乘法的示例:

import torch
import hidet

def matmul(m_size, n_size, k_size):
    from hidet.lang import f32, attr
    from hidet.lang.cuda import threadIdx, blockIdx, blockDim

    with hidet.script_module() as script_module:
        @hidet.script
        def matmul(
            a: f32[m_size, k_size],
            b: f32[k_size, n_size],
            c: f32[m_size, n_size]
        ):
            attr.cuda_grid_dim = ((m_size + 31) // 32, (n_size + 31) // 32)
            attr.cuda_block_dim = (32, 32)
            i = threadIdx.x + blockIdx.x * blockDim.x
            j = threadIdx.y + blockIdx.y * blockDim.y
            if i < m_size and j < n_size:
                c[i, j] = 0.0
                for k in range(k_size):
                    c[i, j] += a[i, k] * b[k, j]

    ir_module = script_module.ir_module()
    func = hidet.driver.build_ir_module(ir_module)
    return func

# 使用自定义算子
class NaiveMatmul(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a, b):
        m, k = a.shape
        k, n = b.shape
        c = torch.empty([m, n], dtype=a.dtype, device=a.device)
        func = matmul(m, n, k)
        func(a, b, c)
        return c

a = torch.randn([3, 4], device='cuda')
b = torch.randn([4, 5], device='cuda')
c = NaiveMatmul.apply(a, b)

这个例子展示了如何使用Hidet Script实现一个简单的矩阵乘法算子,并将其集成到PyTorch中。Hidet Script的强大之处在于,它在简化张量编程的同时,仍然允许用户以灵活的方式实现自己的优化。