BladeDISC初探

在大模型训练推理场景中，一个十分大的瓶颈是动态shape问题。比如nlp领域，处理的句子长短不一，tensor的shape是动态变化的，到runtime才能确定。这给机器学习编译器带来很大的困扰，以XLA为首的sota编译器均是静态shape的，在性能上会有一定损失。BladeDISC是阿里提出的针对动态shape的机器学习编译器，并且经过大量实验和实际生产检验。本文重点关注BladeDISC的构建，pytorch使用方式以及基础架构解读。后续文章会讲解优化流程和论文解读。

源码构建

Build from source

下载BladeDisc镜像

1	docker pull bladedisc/bladedisc:latest-devel-cu118

使用cu118版本

运行该镜像

1	docker run --gpus all -it -v $PWD:/disc bladedisc/bladedisc:latest-devel-cu118 bash

修改一下pytorch_blade/scripts/build_pytorch_blade.sh里面的TORCH_BLADE_CI_BUILD_TORCH_VERSION。修改为存在的requirements.txt即可。

构建过程中，onnx由于带宽等问题，可能会报error，添加-i https://pypi.tuna.tsinghua.edu.cn/simple指定pypi镜像即可。

pytorch版本构建

1
2
3

cd pytorch_blade && bash ./scripts/build_pytorch_blade.sh
python setup.py bdist_wheel
pip install ./pytorch_blade/dist/torch_blade-0.2.0+2.0.1.cu118-cp38-cp38-linux_x86_64.whl

错误处理

如果报错没有安全git，在docker中用：

1	git config --global --add safe.directory /disc

quick install

参考docker install

Pytorch部署BERT模型

Hugging Face模型下载

手动下载模型（适合服务器联网不稳定的情况使用）

找到Bert sentiment inference 模型，主要手动下载如下几个文件：

在python代码中使用离线下载的模型：
1
2
3
model_path = "./model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path).cuda().eval()

直接通过transformers 包下载，该下载方式通过huggingface对应模型网页的use this model获取

# Load model directly
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment")
model = AutoModelForSequenceClassification.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment")

确保环境有transformers包即可

通过huggingface-cli下载

1	huggingface-cli download nlptown/bert-base-multilingual-uncased-sentiment

做BERT Inference的testbench

我的测试codes如下：

import torch
import torch_blade
import time

from transformers import (
    pipeline,
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TextClassificationPipeline,
)

############################################# download model from huggingface #############################################
model_path = "./model"

tokenizer = AutoTokenizer.from_pretrained(model_path)

model = AutoModelForSequenceClassification.from_pretrained(model_path).cuda().eval()

def plain_tokenizer(inputs_str, return_tensors):
    inputs = tokenizer(inputs_str, return_tensors=return_tensors, padding=True)
    inputs = {key: value.cuda() for key, value in inputs.items()}
    
    # torch_blade.optimize 不支持 None 作为输入
    if "token_type_ids" in inputs and inputs["token_type_ids"] is None:
        del inputs["token_type_ids"]

    return (inputs['input_ids'], inputs['attention_mask'], inputs.get('token_type_ids', None))

class PlainTextClassificationPipeline(TextClassificationPipeline):
    def _forward(self, model_inputs):
        return self.model(*model_inputs)

classifier = pipeline(
    'sentiment-analysis',
    model=model,
    tokenizer=plain_tokenizer,
    pipeline_class=PlainTextClassificationPipeline,
    device=0
)

input_strs = [
    "We are very happy to show you the story.",
    "We hope you don't hate it."
]

results = classifier(input_strs)

for inp_str, result in zip(input_strs, results):
    print(inp_str)
    print(f" label: {result['label']}, with a score: {round(result['score'], 4)}")

############################################# Use BladeDISC for optimization #############################################
inputs_str = "Hey, the cat is cute."
inputs = plain_tokenizer(inputs_str, return_tensors="pt")

torch_config = torch_blade.config.Config()
torch_config.enable_mlir_amp = False  # disable mix-precision

# Ensure inputs are properly formatted for optimization
model_inputs = tuple(i for i in inputs if i is not None)

with torch.no_grad(), torch_config:
    optimized_ts = torch_blade.optimize(model, allow_tracing=True, model_inputs=model_inputs)

# Move optimized model to CUDA
optimized_ts = optimized_ts.cuda()

# Save the optimized TorchScript model
torch.jit.save(optimized_ts, "opt.disc.pt")

############################################# testbench #############################################
@torch.no_grad()
def benchmark(model, inputs, num_iters=1000):
    for _ in range(10):
        model(*inputs)
    torch.cuda.synchronize()

    start = time.time()
    for _ in range(num_iters):
        model(*inputs)
    torch.cuda.synchronize()
    end = time.time()
    return (end - start) / num_iters * 1000.0

def bench_and_report(input_strs):
    inputs = plain_tokenizer(input_strs, return_tensors="pt")
    model_inputs = tuple(i for i in inputs if i is not None)

    avg_latency_baseline = benchmark(model, model_inputs)
    avg_latency_bladedisc = benchmark(optimized_ts, model_inputs)

    print(f"Seqlen: {[len(s) for s in input_strs]}")
    print(f"Baseline: {avg_latency_baseline:.4f} ms")
    print(f"BladeDISC: {avg_latency_bladedisc:.4f} ms")
    print(f"BladeDISC speedup: {avg_latency_baseline / avg_latency_bladedisc:.4f}")

input_strs = [
    "We are very happy to show you the story.",
    "We hope you don't hate it."
]

bench_and_report(input_strs)