實(shí)現(xiàn)PyTorch模型高效部署的實(shí)戰(zhàn)教程)
在AI模型落地過(guò)程中推理性能往往比訓(xùn)練階段更容易成為瓶頸。同樣是跑一個(gè)模型離線(xiàn)訓(xùn)練能接受分鐘級(jí)耗時(shí)線(xiàn)上服務(wù)卻要求幾十毫秒內(nèi)返回結(jié)果顯存占用、吞吐量、批量調(diào)度都需要進(jìn)一步優(yōu)化。我們團(tuán)隊(duì)在開(kāi)發(fā)AI系統(tǒng)時(shí)就遇到這類(lèi)問(wèn)題模型換了好幾種推理后端還是無(wú)法同時(shí)滿(mǎn)足延遲和吞吐要求最后決定自己設(shè)計(jì)并部署一個(gè)輕量級(jí)推理加速器項(xiàng)目代號(hào) Redwood。整個(gè)原型設(shè)計(jì)、編碼、部署和驗(yàn)證過(guò)程壓縮在兩周內(nèi)完成這里把完整思路整理成一套可參考的實(shí)操教程。需要先說(shuō)明的是本文提到的“加速器”指深度學(xué)習(xí)推理加速器負(fù)責(zé)算子融合、內(nèi)存調(diào)度、推理服務(wù)優(yōu)化并不是網(wǎng)絡(luò)連接類(lèi)工具。Redwood 定位為一個(gè)偏業(yè)務(wù)側(cè)的推理加速組件它不追求替代 TensorRT 這類(lèi)大型引擎而是把模型優(yōu)化、自定義算子、服務(wù)部署串聯(lián)起來(lái)讓AI系統(tǒng)可以在短時(shí)間內(nèi)獲得穩(wěn)定可測(cè)的加速收益。1. Redwood是什么AI場(chǎng)景下的推理加速器1.1 從AI系統(tǒng)部署痛點(diǎn)到加速器AI系統(tǒng)從訓(xùn)練走向生產(chǎn)時(shí)通常會(huì)面對(duì)三個(gè)層面的問(wèn)題。第一層是模型推理性能。訓(xùn)練時(shí)用的 PyTorch 模型直接跑推理默認(rèn)可能存在很多冗余計(jì)算比如卷積后緊跟 BatchNorm 和 ReLU這三個(gè)算子分別讀寫(xiě)了多遍中間張量增加了顯存帶寬壓力。第二層是服務(wù)化瓶頸。如果直接用 Python 循環(huán)處理請(qǐng)求沒(méi)有批量調(diào)度和并發(fā)控制GPU 利用率會(huì)很低。第三層是硬件差異。開(kāi)發(fā)機(jī)器可能是 NVIDIA 顯卡生產(chǎn)環(huán)境可能是另一套GPU也可能只有 CPU不同設(shè)備上的算子實(shí)現(xiàn)效率差別很大。Redwood 解決的問(wèn)題可以概括為在不改變?cè)寄Ｐ蜆I(yè)務(wù)邏輯的前提下通過(guò)靜態(tài)圖優(yōu)化、算子融合、內(nèi)存復(fù)用和統(tǒng)一的推理運(yùn)行時(shí)讓模型在目標(biāo)設(shè)備上跑得更快、更穩(wěn)。它不是一個(gè)從零寫(xiě)出來(lái)的深度學(xué)習(xí)框架而是一個(gè)中間優(yōu)化層。1.2 Redwood加速器解決的問(wèn)題從實(shí)際需求出發(fā)Redwood 需要回答幾個(gè)具體問(wèn)題。如何把一個(gè) PyTorch 模型轉(zhuǎn)換成更適合推理的靜態(tài)計(jì)算圖如何在計(jì)算圖上自動(dòng)完成常見(jiàn)融合減少算子往返如何為缺少內(nèi)置實(shí)現(xiàn)的算子提供自定義實(shí)現(xiàn)并接入運(yùn)行時(shí)如何通過(guò)批量調(diào)度提高 GPU 吞吐如何暴露成對(duì)業(yè)務(wù)方友好的 HTTP 推理接口這些點(diǎn)聽(tīng)起來(lái)很多但拆開(kāi)后每個(gè)模塊的邊界都很清楚。項(xiàng)目第一天我們把需求收斂成三塊模型優(yōu)化器、算子運(yùn)行時(shí)、服務(wù)入口。Redwood 的所有代碼都圍繞這三塊展開(kāi)。1.3 與常見(jiàn)推理引擎的區(qū)別提到推理加速很多人會(huì)想到 TensorRT、ONNX Runtime、Triton Inference Server。Redwood 和這些方案不是競(jìng)爭(zhēng)關(guān)系而是互補(bǔ)關(guān)系。TensorRT 很強(qiáng)但圖優(yōu)化過(guò)程相對(duì)封閉自定義算子擴(kuò)展成本較高。ONNX Runtime 提供了豐富的 EP適合通用場(chǎng)景但要讓業(yè)務(wù)業(yè)務(wù)側(cè)完全擁抱 ONNX轉(zhuǎn)換和調(diào)優(yōu)也需要不少時(shí)間。Triton Inference Server 更偏服務(wù)化調(diào)度本身不負(fù)責(zé)模型內(nèi)部的算子改寫(xiě)。Redwood 的切入點(diǎn)是先用 PyTorch 的 torch.fx 拿到計(jì)算圖做一層輕量級(jí)優(yōu)化再把高性能算子通過(guò) Triton 或 CUDA 寫(xiě)成插件掛載到運(yùn)行時(shí)。這樣業(yè)務(wù)側(cè)可以繼續(xù)使用 PyTorch 生態(tài)又能在關(guān)鍵路徑上獲得接近手寫(xiě)算子的性能。2. 環(huán)境準(zhǔn)備與版本說(shuō)明2.1 硬件與操作系統(tǒng)Redwood 的驗(yàn)證環(huán)境以 Linux 為主推薦使用 Ubuntu 20.04 或 22.04。如果只做 CPU 推理驗(yàn)證不強(qiáng)制要求 GPU但本文示例中的自定義算子需要 NVIDIA GPU 和 CUDA環(huán)境。版本需要根據(jù)你的項(xiàng)目實(shí)際情況調(diào)整本文示例以常見(jiàn)環(huán)境為例重點(diǎn)是演示配置思路。操作系統(tǒng)Ubuntu 22.04GPUNVIDIA Tesla T4 / V100 / RTX 3090 均可驅(qū)動(dòng)建議 NVIDIA Driver 535 或更新版本內(nèi)存32GB 以上磁盤(pán)20GB 以上可用空間2.2 軟件環(huán)境Redwood 基于 Python 和 PyTorch 搭建自定義算子使用 Triton 編寫(xiě)。Triton 是一種面向GPU編程的語(yǔ)言可以在不使用復(fù)雜 CUDA 代碼的情況下編寫(xiě)高性能算子。下面是一個(gè)常見(jiàn)環(huán)境組合。Python 3.10PyTorch 1.13 或 2.xTriton 2.xonnx 1.14fastapi 0.100uvicorn 0.23docker / docker-composenvidia-container-toolkit先創(chuàng)建項(xiàng)目虛擬環(huán)境conda create -n redwood python3.10 -y conda activate redwood pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install triton onnx fastapi uvicorn pydantic這里沒(méi)有寫(xiě)死 PyTorch 小版本因?yàn)?PyTorch 升級(jí)相對(duì)頻繁。只要你的 Python 和 CUDA 版本與安裝包匹配即可。安裝完成后可以檢查環(huán)境python -c import torch; print(torch.__version__, torch.cuda.is_available()) python -c import triton; print(triton.__version__)如果輸出torch.cuda.is_available()為T(mén)rue說(shuō)明 GPU 環(huán)境正常。2.3 項(xiàng)目結(jié)構(gòu)Redwood 項(xiàng)目代碼結(jié)構(gòu)如下后續(xù)核心代碼都會(huì)落到這些目錄中。redwood/ ├── redwood/ │ ├── __init__.py │ ├── graph_optimizer.py │ ├── memory_pool.py │ ├── runtime.py │ ├── kernels/ │ │ ├── __init__.py │ │ └── fused_matmul_relu.py │ └── server.py ├── models/ │ └── example_model.py ├── scripts/ │ └── benchmark.py ├── requirements.txt ├── Dockerfile └── docker-compose.yml這個(gè)結(jié)構(gòu)保持簡(jiǎn)單方便兩周內(nèi)迭代。實(shí)際項(xiàng)目中可以根據(jù)團(tuán)隊(duì)分工繼續(xù)拆分。3. Redwood整體架構(gòu)與設(shè)計(jì)思路3.1 模塊劃分Redwood 的運(yùn)行時(shí)模塊分為四層。圖優(yōu)化層負(fù)責(zé)加載 PyTorch 模型使用 torch.fx 抓取計(jì)算圖執(zhí)行常量折疊、算子融合、無(wú)用節(jié)點(diǎn)刪除。算子層提供內(nèi)置融合算子也支持注冊(cè)外部自定義算子例如通過(guò) Triton 編寫(xiě)的 kernel。調(diào)度層負(fù)責(zé)推理請(qǐng)求的批量聚合支持動(dòng)態(tài) batch減少 GPU 空轉(zhuǎn)。服務(wù)層基于 FastAPI 提供 HTTP 接口把模型推理封裝成可調(diào)用的 API。每一層都可以單獨(dú)測(cè)試也可以組合使用。這樣設(shè)計(jì)的好處是如果某個(gè)層出現(xiàn)問(wèn)題不需要改動(dòng)其他層代碼。3.2 為什么選擇“業(yè)務(wù)無(wú)關(guān) 算子插件”模式在設(shè)計(jì)初期我們考慮過(guò)直接引入一套完整推理引擎。但后來(lái)發(fā)現(xiàn)團(tuán)隊(duì)的模型迭代非?？旖裉煊镁矸e網(wǎng)絡(luò)下周可能換成 Transformer 結(jié)構(gòu)。如果加速器只支持固定算子新模型上線(xiàn)時(shí)又要重新做適配。Redwood 選擇“業(yè)務(wù)無(wú)關(guān) 算子插件”模式。圖優(yōu)化層只處理通用的圖改寫(xiě)規(guī)則不關(guān)心模型是 CV 還是 NLP算子層通過(guò)字典注冊(cè)算子實(shí)現(xiàn)新算子只需實(shí)現(xiàn)統(tǒng)一接口即可被運(yùn)行時(shí)調(diào)用。這樣業(yè)務(wù)模型可以像插拔插件一樣接入 Redwood。下面是一個(gè)算子注冊(cè)表示例用來(lái)理解插件模式# redwood/kernels/__init__.py _OP_REGISTRY {} def register_op(name): def decorator(func): _OP_REGISTRY[name] func return func return decorator def get_op(name): if name not in _OP_REGISTRY: raise ValueError(fUnsupported op: {name}) return _OP_REGISTRY[name]這種注冊(cè)模式在框架集成中很常見(jiàn)。后續(xù)如果接入新的自定義算子只需要在導(dǎo)入路徑中調(diào)用register_op運(yùn)行時(shí)就能發(fā)現(xiàn)它。3.3 兩周迭代計(jì)劃兩周時(shí)間看起來(lái)短但把任務(wù)拆細(xì)后完全能夠完成一個(gè)可演示的加速器原型。第1-2天完成環(huán)境搭建、需求拆解、架構(gòu)設(shè)計(jì)。第3-6天實(shí)現(xiàn)圖優(yōu)化層支持 ConvBNReLU 融合。第7-10天實(shí)現(xiàn) Triton 自定義算子并通過(guò)單元測(cè)試。第11-13天實(shí)現(xiàn) FastAPI 推理服務(wù)和動(dòng)態(tài) batch 邏輯。第14天容器部署、性能基準(zhǔn)測(cè)試、輸出總結(jié)文檔。這個(gè)計(jì)劃不是所有團(tuán)隊(duì)都必須遵循但它強(qiáng)調(diào)了核心路徑先讓圖優(yōu)化跑通再優(yōu)化單算子再接入服務(wù)。這樣即使時(shí)間緊張至少能完成第一個(gè)階段的演示。4. 核心代碼實(shí)現(xiàn)4.1 模型優(yōu)化PassConvBNReLU融合圖優(yōu)化層的第一個(gè)功能是實(shí)現(xiàn)算子融合。以卷積網(wǎng)絡(luò)中非常常見(jiàn)的Conv2d - BatchNorm2d - ReLU結(jié)構(gòu)為例三個(gè)算子單獨(dú)執(zhí)行時(shí)每次都會(huì)讀取并寫(xiě)回中間張量。融合后只需要一次數(shù)據(jù)讀取和一次寫(xiě)入可以有效降低顯存帶寬消耗。Redwood 使用torch.fx對(duì)模型進(jìn)行符號(hào)化追蹤然后遍歷計(jì)算圖找到滿(mǎn)足條件的卷積和 BN 節(jié)點(diǎn)將 BN 參數(shù)折算到卷積權(quán)重中最后去掉 BN 節(jié)點(diǎn)并保留 ReLU。這里給出一個(gè)完整可運(yùn)行的優(yōu)化示例# redwood/graph_optimizer.py import torch import torch.nn as nn from torch.fx import GraphModule, symbolic_trace def fuse_conv_bn(conv: nn.Conv2d, bn: nn.BatchNorm2d): 將 Conv2d 與 BatchNorm2d 融合為一個(gè) Conv2d。 計(jì)算方式將 BN 的縮放和偏移折算到卷積權(quán)重和偏置中。 conv conv.eval() bn bn.eval() bn_mean bn.running_mean bn_var bn.running_var bn_eps bn.eps bn_gamma bn.weight bn_beta bn.bias scale bn_gamma / torch.sqrt(bn_var bn_eps) new_weight conv.weight * scale.view(-1, 1, 1, 1) if conv.bias is not None: new_bias (conv.bias - bn_mean) * scale bn_beta else: new_bias -bn_mean * scale bn_beta fused_conv nn.Conv2d( conv.in_channels, conv.out_channels, conv.kernel_size, conv.stride, conv.padding, conv.dilation, conv.groups, biasTrue, ) fused_conv.load_state_dict({weight: new_weight, bias: new_bias}) return fused_conv def optimize_graph(model: nn.Module, example_input: torch.Tensor) - GraphModule: 通過(guò) torch.fx 追蹤模型計(jì)算圖并執(zhí)行 ConvBN 融合。 注意這里只演示核心邏輯實(shí)際還需要處理 BN 后的 ReLU 等節(jié)點(diǎn)。 model.eval() traced symbolic_trace(model) nodes list(traced.graph.nodes) for node in nodes: if node.op call_module and isinstance(traced.get_submodule(node.target), nn.BatchNorm2d): prev_node node.args[0] if prev_node.op call_module and isinstance(traced.get_submodule(prev_node.target), nn.Conv2d): fused_conv fuse_conv_bn( traced.get_submodule(prev_node.target), traced.get_submodule(node.target), ) # 將原 conv 模塊替換為融合后的 conv traced.delete_submodule(prev_node.target) traced.add_submodule(prev_node.target, fused_conv) # 讓所有指向 BN 的節(jié)點(diǎn)直接指向原 conv node.replace_all_uses_with(prev_node) traced.graph.erase_node(node) traced.recompile() return traced這段代碼的核心思路是把 BN 的縮放系數(shù)和偏移量折算進(jìn)卷積的 weight 和 bias然后在計(jì)算圖中刪除 BN 節(jié)點(diǎn)。由于 ReLU 是逐元素操作一般可以放在融合的最后一步或者在后續(xù)算子中直接包含激活函數(shù)。注意這里只是核心片段要放入redwood/graph_optimizer.py實(shí)際使用還需考慮模型中的不同模塊命名情況。4.2 自定義Triton算子MatMul ReLU融合圖優(yōu)化完成后還需要讓關(guān)鍵算子跑得更快。這里給出一個(gè)使用 Triton 編寫(xiě)的MatMul ReLU融合算子示例。它在一個(gè) kernel 內(nèi)完成矩陣乘法和激活計(jì)算避免中間結(jié)果的顯存讀寫(xiě)。# redwood/kernels/fused_matmul_relu.py import torch import triton import triton.language as tl triton.jit def fused_matmul_relu_kernel( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): pid_m tl.program_id(0) pid_n tl.program_id(1) offs_m pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) offs_k tl.arange(0, BLOCK_K) a_ptrs a_ptr (offs_m[:, None] * stride_am offs_k[None, :] * stride_ak) b_ptrs b_ptr (offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn) acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for k in range(0, K, BLOCK_K): a tl.load(a_ptrs) b tl.load(b_ptrs) acc tl.dot(a, b) a_ptrs BLOCK_K * stride_ak b_ptrs BLOCK_K * stride_bk # 融合 ReLU 激活 acc tl.maximum(acc, 0) offs_cm pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_cn pid_n * BLOCK_N tl.arange(0, BLOCK_N) c_ptrs c_ptr (offs_cm[:, None] * stride_cm offs_cn[None, :] * stride_cn) tl.store(c_ptrs, acc) def matmul_relu(a: torch.Tensor, b: torch.Tensor) - torch.Tensor: assert a.is_cuda and b.is_cuda M, K a.shape K2, N b.shape assert K K2 c torch.empty((M, N), devicea.device, dtypetorch.float32) BLOCK_M, BLOCK_N, BLOCK_K 16, 16, 16 grid (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) fused_matmul_relu_kernel[grid]( a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_MBLOCK_M, BLOCK_NBLOCK_N, BLOCK_KBLOCK_K, ) return c這個(gè) kernel 中每個(gè)線(xiàn)程塊負(fù)責(zé)計(jì)算結(jié)果矩陣中的一個(gè)BLOCK_M x BLOCK_N分塊內(nèi)層循環(huán)按BLOCK_K累加。與 PyTorch 自帶的torch.matmul不同這里沒(méi)有產(chǎn)生M x N x K的中間張量而是直接在累加后執(zhí)行ReLU然后寫(xiě)回結(jié)果。實(shí)際使用時(shí)應(yīng)根據(jù) GPU 型號(hào)調(diào)整BLOCK大小。4.3 調(diào)度器與內(nèi)存池推理服務(wù)端通常需要處理大量請(qǐng)求。如果每個(gè)請(qǐng)求都單獨(dú)調(diào)用一次 GPU 算子GPU 的利用率很低。Redwood 的調(diào)度層負(fù)責(zé)收集短時(shí)間內(nèi)到達(dá)的請(qǐng)求把它們拼成一個(gè) batch 后統(tǒng)一推理。# redwood/runtime.py import threading import time import torch class RedwoodInferenceRuntime: def __init__(self, model, max_batch_size8, wait_time0.01): self.model model self.max_batch_size max_batch_size self.wait_time wait_time self._queue [] self._lock threading.Lock() def submit(self, tensor): with self._lock: self._queue.append(tensor) time.sleep(self.wait_time) with self._lock: if len(self._queue) self.max_batch_size: return self._flush_locked() return None def _flush_locked(self): batch torch.cat(self._queue, dim0) self._queue.clear() with torch.no_grad(): return self.model(batch) def flush(self): with self._lock: if self._queue: return self._flush_locked() return None這段代碼是一個(gè)簡(jiǎn)化版動(dòng)態(tài) batch 調(diào)度器。請(qǐng)求到達(dá)后不立即推理而是等待一小段時(shí)間盡量積攢成一個(gè) batch。當(dāng) batch 大小達(dá)到上限后立即處理。實(shí)際生產(chǎn)環(huán)境還需要處理隊(duì)列積壓、超時(shí)、顯存池復(fù)用等問(wèn)題但核心思路可以作為入門(mén)實(shí)現(xiàn)。內(nèi)存池方面Redwood 可以復(fù)用 PyTorch 的 caching allocator同時(shí)在包含可變長(zhǎng)度輸入的業(yè)務(wù)中通過(guò)對(duì)輸入做 padding 和緩存輸出緩沖區(qū)減少反復(fù)申請(qǐng)顯存。真正實(shí)現(xiàn)時(shí)可以?xún)?yōu)先觀察torch.cuda.memory_stats()判斷是否存在頻繁顯存分配。4.4 FastAPI推理服務(wù)最后把 Redwood 的圖優(yōu)化和運(yùn)行時(shí)封裝成一個(gè) HTTP 服務(wù)。這里使用 FastAPI 提供/predict接口。# redwood/server.py import io import torch from fastapi import FastAPI from pydantic import BaseModel from redwood.graph_optimizer import optimize_graph from redwood.runtime import RedwoodInferenceRuntime app FastAPI(titleRedwood Inference Service) class PredictRequest(BaseModel): data: list class PredictResponse(BaseModel): result: list cost_ms: float _model None _runtime None def load_model(): global _model, _runtime from models.example_model import ExampleModel model ExampleModel() checkpoint torch.load(models/example_model.pt, map_locationcpu) model.load_state_dict(checkpoint) model.eval() example_input torch.randn(1, 3, 224, 224) optimized_model optimize_graph(model, example_input) _runtime RedwoodInferenceRuntime(optimized_model) app.on_event(startup) def startup(): load_model() app.post(/predict, response_modelPredictResponse) def predict(req: PredictRequest): import time tensor torch.tensor(req.data, dtypetorch.float32) start time.time() with torch.no_grad(): output _runtime.submit(tensor) cost_ms (time.time() - start) * 1000 return PredictResponse(resultoutput.tolist(), cost_mscost_ms)這里省略了復(fù)雜的模型加載和批處理細(xì)節(jié)重點(diǎn)展示服務(wù)層如何調(diào)用底層 Runtime。通過(guò)app.on_event(startup)在服務(wù)啟動(dòng)時(shí)加載和優(yōu)化模型避免每次預(yù)測(cè)都重復(fù)做圖優(yōu)化。5. 部署與驗(yàn)證5.1 容器鏡像構(gòu)建Redwood 使用 Docker 部署可以保證運(yùn)行時(shí)環(huán)境一致。以下是一個(gè)基礎(chǔ) Dockerfile 示例。FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 ENV DEBIAN_FRONTENDnoninteractive RUN apt-get update apt-get install -y \ python3.10 python3.10-dev python3-pip \ rm -rf /var/lib/apt/lists/* WORKDIR /app COPY requirements.txt . RUN pip3 install --no-cache-dir -r requirements.txt COPY . . EXPOSE 8000 CMD [uvicorn, redwood.server:app, --host, 0.0.0.0, --port, 8000]部署到測(cè)試環(huán)境前務(wù)必在測(cè)試環(huán)境驗(yàn)證流程。5.2 啟動(dòng)服務(wù)本地開(kāi)發(fā)環(huán)境啟動(dòng)服務(wù)可以直接使用 Uvicornuvicorn redwood.server:app --host 0.0.0.0 --port 8000使用 Docker Compose 啟動(dòng)時(shí)可以參考下面的配置# docker-compose.yml services: redwood: build: . ports: - 8000:8000 volumes: - ./models:/app/models environment: - CUDA_VISIBLE_DEVICES0 deploy: resources: reservations: devices: - driver: nvidia count: 1 capabilities: [gpu]啟動(dòng)命令docker-compose up --build注意NVIDIA Container Toolkit 需要提前安裝否則容器內(nèi)無(wú)法識(shí)別 GPU。5.3 性能對(duì)比Redwood 上線(xiàn)后的性能驗(yàn)證不能只關(guān)注單個(gè)算子的耗時(shí)還要關(guān)注端到端推理延遲和吞吐量。我們寫(xiě)一個(gè)簡(jiǎn)單的基準(zhǔn)腳本# scripts/benchmark.py import time import torch from redwood.graph_optimizer import optimize_graph from models.example_model import ExampleModel model ExampleModel().eval() example_input torch.randn(1, 3, 224, 224) # 原始模型 start time.time() for _ in range(100): with torch.no_grad(): model(example_input) original_cost (time.time() - start) / 100 * 1000 # 優(yōu)化后的模型 optimized optimize_graph(model, example_input) start time.time() for _ in range(100): with torch.no_grad(): optimized(example_input) optimized_cost (time.time() - start) / 100 * 1000 print(fOriginal: {original_cost:.3f} ms) print(fOptimized: {optimized_cost:.3f} ms) print(fSpeedup: {original_cost / optimized_cost:.2f}x)在真實(shí)項(xiàng)目中建議至少跑 1000 次并先 warmup 數(shù)次避免首次調(diào)用包含 CUDA kernel 加載耗時(shí)。性能數(shù)據(jù)可能因模型和 GPU 不同而差異很大關(guān)鍵是觀察優(yōu)化前后同一環(huán)境下的相對(duì)提升。5.4 準(zhǔn)確性驗(yàn)證加速器不能只快不準(zhǔn)。每次圖優(yōu)化后都要對(duì)比模型輸出與原始輸出的誤差。一般情況下ConvBN 融合是數(shù)值等價(jià)變換誤差應(yīng)保持在非常小的范圍。def verify_consistency(original, optimized, inputs): with torch.no_grad(): y1 original(inputs) y2 optimized(inputs) max_diff (y1 - y2).abs().max().item() rel_diff ((y1 - y2).abs() / (y1.abs() 1e-6)).max().item() print(fmax absolute diff: {max_diff:.6e}) print(fmax relative diff: {rel_diff:.6e})如果誤差明顯偏大優(yōu)先檢查融合過(guò)程中的均值、方差取的是訓(xùn)練階段的統(tǒng)計(jì)量還是當(dāng)前 batch 的統(tǒng)計(jì)量。推理階段必須使用running_mean和running_var否則輸出會(huì)有偏差。6. 常見(jiàn)問(wèn)題與排查思路在 Redwood 的開(kāi)發(fā)與部署過(guò)程中我們遇到了不少問(wèn)題。這里整理成表格方便后續(xù)讀者按圖索驥。問(wèn)題現(xiàn)象常見(jiàn)原因解決思路torch.cuda.is_available() 為 falseNVIDIA 驅(qū)動(dòng)或 CUDA 版本不匹配用nvidia-smi檢查驅(qū)動(dòng)重新安裝匹配的 PyTorchTriton kernel 編譯報(bào)錯(cuò)BLOCK 參數(shù)設(shè)置不適合 GPU調(diào)整 BLOCK_M/N/K或檢查是否使用 GPU 運(yùn)行推理結(jié)果與原始模型不一致BN 融合時(shí)使用了訓(xùn)練模式參數(shù)確保模型處于 eval 狀態(tài)使用 running_mean/varFastAPI 服務(wù)啟動(dòng)慢圖優(yōu)化在啟動(dòng)階段執(zhí)行提前離線(xiàn)保存優(yōu)化后的模型啟動(dòng)時(shí)直接加載并發(fā)請(qǐng)求時(shí) GPU 利用率低沒(méi)有批量調(diào)度啟用 Redwood 的 batch 調(diào)度調(diào)整 wait_time顯存占用持續(xù)上漲內(nèi)存池沒(méi)有復(fù)用或存在緩存泄漏使用torch.cuda.memory_stats()定位限制緩存清理策略容器內(nèi)無(wú)法識(shí)別 GPU未安裝 nvidia-container-toolkit安裝相應(yīng)版本并重啟容器服務(wù)下面展開(kāi)幾個(gè)高頻問(wèn)題的排查過(guò)程。6.1 Triton Kernel 編譯失敗Triton 在首次執(zhí)行 kernel 時(shí)會(huì)有 JIT 編譯過(guò)程。如果設(shè)置不當(dāng)編譯時(shí)間會(huì)很長(zhǎng)甚至直接報(bào)錯(cuò)。常見(jiàn)原因是BLOCK_SIZE與 GPU 寄存器限制不匹配。排查思路先查看完整報(bào)錯(cuò)棧確認(rèn)是否在tl.dot或tl.load附近。調(diào)小BLOCK_M、BLOCK_N例如從 32 改為 16。確認(rèn)輸入張量是 GPU 上的連續(xù)內(nèi)存。在代碼中設(shè)置TRITON_KERNEL_OVERRIDE1或檢查 Triton 緩存目錄權(quán)限。6.2 圖優(yōu)化后模型輸出不一致這是優(yōu)化器最容易踩的坑。ConvBN 融合時(shí)如果模型處于 training 模式BN 層會(huì)使用當(dāng)前 batch 的均值和方差導(dǎo)致融合結(jié)果錯(cuò)誤。因此執(zhí)行optimize_graph之前必須調(diào)用model.eval()并確認(rèn)所有 BN 層的running_mean和running_var已經(jīng)完成更新。如果在加載 checkpoint 后立即做融合需要先跑一次驗(yàn)證模式的前向或者在保存模型時(shí)就確保 BN 統(tǒng)計(jì)量正確。最簡(jiǎn)單的方式是加載完權(quán)重后先執(zhí)行model.eval()再做 graph trace。6.3 動(dòng)態(tài) batch 導(dǎo)致請(qǐng)求阻塞Redwood 的簡(jiǎn)單調(diào)度器會(huì)等待wait_time來(lái)聚合請(qǐng)求。如果業(yè)務(wù)流量本身很小等待時(shí)間會(huì)導(dǎo)致額外延遲。針對(duì)這種情況可以把wait_time設(shè)置得很小或者采用“達(dá)到最小 batch 立即推理”的策略。更復(fù)雜但穩(wěn)定的方案是使用定時(shí)器 flush 隊(duì)列比如每 5ms 檢查一次。7. 最佳實(shí)踐與工程建議7.1 用AI編程工具提升開(kāi)發(fā)效率Redwood 能在兩周內(nèi)完成很大程度上依賴(lài) AI 編程工具和 AI Agent 的輔助。寫(xiě) Triton kernel、調(diào)試 torch.fx 圖改寫(xiě)、生成 Dockerfile 等大量重復(fù)性工作都可以借助 AI 編程助手快速產(chǎn)出初稿。需要注意的是AI 生成代碼必須經(jīng)過(guò)測(cè)試驗(yàn)證不能直接信任尤其涉及 CUDA 算子和圖優(yōu)化邏輯時(shí)要重點(diǎn)審查數(shù)值正確性。推薦工作流是先自己梳理模塊邊界讓 AI 輔助生成函數(shù)骨架再用單元測(cè)試鎖定行為最后根據(jù)編譯報(bào)錯(cuò)迭代。這樣既能提升效率又能保證代碼質(zhì)量。7.2 安全與生產(chǎn)注意事項(xiàng)Redwood 涉及模型部署和推理服務(wù)生產(chǎn)環(huán)境變更時(shí)必須遵守幾個(gè)原則。在測(cè)試環(huán)境驗(yàn)證后再發(fā)布不要在業(yè)務(wù)高峰期直接改動(dòng)模型服務(wù)。模型文件和配置統(tǒng)一放在只讀目錄避免運(yùn)行時(shí)被篡改。對(duì)外服務(wù)接口需要鑒權(quán)預(yù)測(cè)接口不能裸奔暴露在公網(wǎng)。對(duì)請(qǐng)求體大小做限制避免超大輸入導(dǎo)致顯存溢出。日志記錄不要輸出完整的模型輸入輸出防止數(shù)據(jù)泄露。推理服務(wù)盡量使用最小權(quán)限容器賬號(hào)運(yùn)行避免使用 root。部署前還要做備份。如果使用模型版本管理工具上線(xiàn)新版本時(shí)可以快速回滾到舊版本。7.3 后續(xù)演進(jìn)方向Redwood 還只是一個(gè)淺層加速器原型后續(xù)可以從幾個(gè)方向繼續(xù)完善。支持更多融合模式例如 LayerNorm 激活、Attention 融合。引入動(dòng)態(tài)形狀處理解決 NLP 場(chǎng)景下變長(zhǎng)序列問(wèn)題。增加多模型管理不同模型使用不同的優(yōu)化策略。接入 Prometheus 監(jiān)控將延遲、吞吐、顯存指標(biāo)標(biāo)準(zhǔn)化。擴(kuò)展 CPU 后端讓 Redwood 在無(wú) GPU 環(huán)境也能運(yùn)行。如果團(tuán)隊(duì)業(yè)務(wù)逐漸穩(wěn)定可以考慮把 Redwood 的一部分能力逐漸貢獻(xiàn)到開(kāi)源社區(qū)或者對(duì)接成熟的推理引擎減少自研維護(hù)成本。最后想說(shuō)的是自研加速器并不一定適合所有團(tuán)隊(duì)。如果只是快速上線(xiàn)一個(gè)模型服務(wù)直接使用 ONNX Runtime 或 TensorRT 會(huì)更穩(wěn)妥。Redwood 的意義在于讓我們理解推理鏈路中哪些優(yōu)化真正有效并且保留了對(duì)業(yè)務(wù)模型的深度定制能力。希望這份兩周落地的經(jīng)驗(yàn)?zāi)芙o你提供參考也歡迎在實(shí)際項(xiàng)目中根據(jù)自身場(chǎng)景進(jìn)行調(diào)整。 自研推理加速器Redwood：兩周內(nèi)實(shí)現(xiàn)PyTorch模型高效部署的實(shí)戰(zhàn)教程
在AI模型落地過(guò)程中推理性能往往比訓(xùn)練階段更容易成為瓶頸。同樣是跑一個(gè)模型離線(xiàn)訓(xùn)練能接受分鐘級(jí)耗時(shí)線(xiàn)上服務(wù)卻要求幾十毫秒內(nèi)返回結(jié)果顯存占用、吞吐量、批量調(diào)度都需要進(jìn)一步優(yōu)化。我們團(tuán)隊(duì)在開(kāi)發(fā)AI系統(tǒng)時(shí)就遇到這類(lèi)問(wèn)題模型換了好幾種推理后端還是無(wú)法同時(shí)滿(mǎn)足延遲和吞吐要求最后決定自己設(shè)計(jì)并部署一個(gè)輕量級(jí)推理加速器項(xiàng)目代號(hào) Redwood。整個(gè)原型設(shè)計(jì)、編碼、部署和驗(yàn)證過(guò)程壓縮在兩周內(nèi)完成這里把完整思路整理成一套可參考的實(shí)操教程。需要先說(shuō)明的是本文提到的“加速器”指深度學(xué)習(xí)推理加速器負(fù)責(zé)算子融合、內(nèi)存調(diào)度、推理服務(wù)優(yōu)化并不是網(wǎng)絡(luò)連接類(lèi)工具。Redwood 定位為一個(gè)偏業(yè)務(wù)側(cè)的推理加速組件它不追求替代 TensorRT 這類(lèi)大型引擎而是把模型優(yōu)化、自定義算子、服務(wù)部署串聯(lián)起來(lái)讓AI系統(tǒng)可以在短時(shí)間內(nèi)獲得穩(wěn)定可測(cè)的加速收益。1. Redwood是什么AI場(chǎng)景下的推理加速器1.1 從AI系統(tǒng)部署痛點(diǎn)到加速器AI系統(tǒng)從訓(xùn)練走向生產(chǎn)時(shí)通常會(huì)面對(duì)三個(gè)層面的問(wèn)題。第一層是模型推理性能。訓(xùn)練時(shí)用的 PyTorch 模型直接跑推理默認(rèn)可能存在很多冗余計(jì)算比如卷積后緊跟 BatchNorm 和 ReLU這三個(gè)算子分別讀寫(xiě)了多遍中間張量增加了顯存帶寬壓力。第二層是服務(wù)化瓶頸。如果直接用 Python 循環(huán)處理請(qǐng)求沒(méi)有批量調(diào)度和并發(fā)控制GPU 利用率會(huì)很低。第三層是硬件差異。開(kāi)發(fā)機(jī)器可能是 NVIDIA 顯卡生產(chǎn)環(huán)境可能是另一套GPU也可能只有 CPU不同設(shè)備上的算子實(shí)現(xiàn)效率差別很大。Redwood 解決的問(wèn)題可以概括為在不改變?cè)寄Ｐ蜆I(yè)務(wù)邏輯的前提下通過(guò)靜態(tài)圖優(yōu)化、算子融合、內(nèi)存復(fù)用和統(tǒng)一的推理運(yùn)行時(shí)讓模型在目標(biāo)設(shè)備上跑得更快、更穩(wěn)。它不是一個(gè)從零寫(xiě)出來(lái)的深度學(xué)習(xí)框架而是一個(gè)中間優(yōu)化層。1.2 Redwood加速器解決的問(wèn)題從實(shí)際需求出發(fā)Redwood 需要回答幾個(gè)具體問(wèn)題。如何把一個(gè) PyTorch 模型轉(zhuǎn)換成更適合推理的靜態(tài)計(jì)算圖如何在計(jì)算圖上自動(dòng)完成常見(jiàn)融合減少算子往返如何為缺少內(nèi)置實(shí)現(xiàn)的算子提供自定義實(shí)現(xiàn)并接入運(yùn)行時(shí)如何通過(guò)批量調(diào)度提高 GPU 吞吐如何暴露成對(duì)業(yè)務(wù)方友好的 HTTP 推理接口這些點(diǎn)聽(tīng)起來(lái)很多但拆開(kāi)后每個(gè)模塊的邊界都很清楚。項(xiàng)目第一天我們把需求收斂成三塊模型優(yōu)化器、算子運(yùn)行時(shí)、服務(wù)入口。Redwood 的所有代碼都圍繞這三塊展開(kāi)。1.3 與常見(jiàn)推理引擎的區(qū)別提到推理加速很多人會(huì)想到 TensorRT、ONNX Runtime、Triton Inference Server。Redwood 和這些方案不是競(jìng)爭(zhēng)關(guān)系而是互補(bǔ)關(guān)系。TensorRT 很強(qiáng)但圖優(yōu)化過(guò)程相對(duì)封閉自定義算子擴(kuò)展成本較高。ONNX Runtime 提供了豐富的 EP適合通用場(chǎng)景但要讓業(yè)務(wù)業(yè)務(wù)側(cè)完全擁抱 ONNX轉(zhuǎn)換和調(diào)優(yōu)也需要不少時(shí)間。Triton Inference Server 更偏服務(wù)化調(diào)度本身不負(fù)責(zé)模型內(nèi)部的算子改寫(xiě)。Redwood 的切入點(diǎn)是先用 PyTorch 的 torch.fx 拿到計(jì)算圖做一層輕量級(jí)優(yōu)化再把高性能算子通過(guò) Triton 或 CUDA 寫(xiě)成插件掛載到運(yùn)行時(shí)。這樣業(yè)務(wù)側(cè)可以繼續(xù)使用 PyTorch 生態(tài)又能在關(guān)鍵路徑上獲得接近手寫(xiě)算子的性能。2. 環(huán)境準(zhǔn)備與版本說(shuō)明2.1 硬件與操作系統(tǒng)Redwood 的驗(yàn)證環(huán)境以 Linux 為主推薦使用 Ubuntu 20.04 或 22.04。如果只做 CPU 推理驗(yàn)證不強(qiáng)制要求 GPU但本文示例中的自定義算子需要 NVIDIA GPU 和 CUDA環(huán)境。版本需要根據(jù)你的項(xiàng)目實(shí)際情況調(diào)整本文示例以常見(jiàn)環(huán)境為例重點(diǎn)是演示配置思路。操作系統(tǒng)Ubuntu 22.04GPUNVIDIA Tesla T4 / V100 / RTX 3090 均可驅(qū)動(dòng)建議 NVIDIA Driver 535 或更新版本內(nèi)存32GB 以上磁盤(pán)20GB 以上可用空間2.2 軟件環(huán)境Redwood 基于 Python 和 PyTorch 搭建自定義算子使用 Triton 編寫(xiě)。Triton 是一種面向GPU編程的語(yǔ)言可以在不使用復(fù)雜 CUDA 代碼的情況下編寫(xiě)高性能算子。下面是一個(gè)常見(jiàn)環(huán)境組合。Python 3.10PyTorch 1.13 或 2.xTriton 2.xonnx 1.14fastapi 0.100uvicorn 0.23docker / docker-composenvidia-container-toolkit先創(chuàng)建項(xiàng)目虛擬環(huán)境conda create -n redwood python3.10 -y conda activate redwood pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install triton onnx fastapi uvicorn pydantic這里沒(méi)有寫(xiě)死 PyTorch 小版本因?yàn)?PyTorch 升級(jí)相對(duì)頻繁。只要你的 Python 和 CUDA 版本與安裝包匹配即可。安裝完成后可以檢查環(huán)境python -c import torch; print(torch.__version__, torch.cuda.is_available()) python -c import triton; print(triton.__version__)如果輸出torch.cuda.is_available()為T(mén)rue說(shuō)明 GPU 環(huán)境正常。2.3 項(xiàng)目結(jié)構(gòu)Redwood 項(xiàng)目代碼結(jié)構(gòu)如下后續(xù)核心代碼都會(huì)落到這些目錄中。redwood/ ├── redwood/ │ ├── __init__.py │ ├── graph_optimizer.py │ ├── memory_pool.py │ ├── runtime.py │ ├── kernels/ │ │ ├── __init__.py │ │ └── fused_matmul_relu.py │ └── server.py ├── models/ │ └── example_model.py ├── scripts/ │ └── benchmark.py ├── requirements.txt ├── Dockerfile └── docker-compose.yml這個(gè)結(jié)構(gòu)保持簡(jiǎn)單方便兩周內(nèi)迭代。實(shí)際項(xiàng)目中可以根據(jù)團(tuán)隊(duì)分工繼續(xù)拆分。3. Redwood整體架構(gòu)與設(shè)計(jì)思路3.1 模塊劃分Redwood 的運(yùn)行時(shí)模塊分為四層。圖優(yōu)化層負(fù)責(zé)加載 PyTorch 模型使用 torch.fx 抓取計(jì)算圖執(zhí)行常量折疊、算子融合、無(wú)用節(jié)點(diǎn)刪除。算子層提供內(nèi)置融合算子也支持注冊(cè)外部自定義算子例如通過(guò) Triton 編寫(xiě)的 kernel。調(diào)度層負(fù)責(zé)推理請(qǐng)求的批量聚合支持動(dòng)態(tài) batch減少 GPU 空轉(zhuǎn)。服務(wù)層基于 FastAPI 提供 HTTP 接口把模型推理封裝成可調(diào)用的 API。每一層都可以單獨(dú)測(cè)試也可以組合使用。這樣設(shè)計(jì)的好處是如果某個(gè)層出現(xiàn)問(wèn)題不需要改動(dòng)其他層代碼。3.2 為什么選擇“業(yè)務(wù)無(wú)關(guān) 算子插件”模式在設(shè)計(jì)初期我們考慮過(guò)直接引入一套完整推理引擎。但后來(lái)發(fā)現(xiàn)團(tuán)隊(duì)的模型迭代非?？旖裉煊镁矸e網(wǎng)絡(luò)下周可能換成 Transformer 結(jié)構(gòu)。如果加速器只支持固定算子新模型上線(xiàn)時(shí)又要重新做適配。Redwood 選擇“業(yè)務(wù)無(wú)關(guān) 算子插件”模式。圖優(yōu)化層只處理通用的圖改寫(xiě)規(guī)則不關(guān)心模型是 CV 還是 NLP算子層通過(guò)字典注冊(cè)算子實(shí)現(xiàn)新算子只需實(shí)現(xiàn)統(tǒng)一接口即可被運(yùn)行時(shí)調(diào)用。這樣業(yè)務(wù)模型可以像插拔插件一樣接入 Redwood。下面是一個(gè)算子注冊(cè)表示例用來(lái)理解插件模式# redwood/kernels/__init__.py _OP_REGISTRY {} def register_op(name): def decorator(func): _OP_REGISTRY[name] func return func return decorator def get_op(name): if name not in _OP_REGISTRY: raise ValueError(fUnsupported op: {name}) return _OP_REGISTRY[name]這種注冊(cè)模式在框架集成中很常見(jiàn)。后續(xù)如果接入新的自定義算子只需要在導(dǎo)入路徑中調(diào)用register_op運(yùn)行時(shí)就能發(fā)現(xiàn)它。3.3 兩周迭代計(jì)劃兩周時(shí)間看起來(lái)短但把任務(wù)拆細(xì)后完全能夠完成一個(gè)可演示的加速器原型。第1-2天完成環(huán)境搭建、需求拆解、架構(gòu)設(shè)計(jì)。第3-6天實(shí)現(xiàn)圖優(yōu)化層支持 ConvBNReLU 融合。第7-10天實(shí)現(xiàn) Triton 自定義算子并通過(guò)單元測(cè)試。第11-13天實(shí)現(xiàn) FastAPI 推理服務(wù)和動(dòng)態(tài) batch 邏輯。第14天容器部署、性能基準(zhǔn)測(cè)試、輸出總結(jié)文檔。這個(gè)計(jì)劃不是所有團(tuán)隊(duì)都必須遵循但它強(qiáng)調(diào)了核心路徑先讓圖優(yōu)化跑通再優(yōu)化單算子再接入服務(wù)。這樣即使時(shí)間緊張至少能完成第一個(gè)階段的演示。4. 核心代碼實(shí)現(xiàn)4.1 模型優(yōu)化PassConvBNReLU融合圖優(yōu)化層的第一個(gè)功能是實(shí)現(xiàn)算子融合。以卷積網(wǎng)絡(luò)中非常常見(jiàn)的Conv2d - BatchNorm2d - ReLU結(jié)構(gòu)為例三個(gè)算子單獨(dú)執(zhí)行時(shí)每次都會(huì)讀取并寫(xiě)回中間張量。融合后只需要一次數(shù)據(jù)讀取和一次寫(xiě)入可以有效降低顯存帶寬消耗。Redwood 使用torch.fx對(duì)模型進(jìn)行符號(hào)化追蹤然后遍歷計(jì)算圖找到滿(mǎn)足條件的卷積和 BN 節(jié)點(diǎn)將 BN 參數(shù)折算到卷積權(quán)重中最后去掉 BN 節(jié)點(diǎn)并保留 ReLU。這里給出一個(gè)完整可運(yùn)行的優(yōu)化示例# redwood/graph_optimizer.py import torch import torch.nn as nn from torch.fx import GraphModule, symbolic_trace def fuse_conv_bn(conv: nn.Conv2d, bn: nn.BatchNorm2d): 將 Conv2d 與 BatchNorm2d 融合為一個(gè) Conv2d。 計(jì)算方式將 BN 的縮放和偏移折算到卷積權(quán)重和偏置中。 conv conv.eval() bn bn.eval() bn_mean bn.running_mean bn_var bn.running_var bn_eps bn.eps bn_gamma bn.weight bn_beta bn.bias scale bn_gamma / torch.sqrt(bn_var bn_eps) new_weight conv.weight * scale.view(-1, 1, 1, 1) if conv.bias is not None: new_bias (conv.bias - bn_mean) * scale bn_beta else: new_bias -bn_mean * scale bn_beta fused_conv nn.Conv2d( conv.in_channels, conv.out_channels, conv.kernel_size, conv.stride, conv.padding, conv.dilation, conv.groups, biasTrue, ) fused_conv.load_state_dict({weight: new_weight, bias: new_bias}) return fused_conv def optimize_graph(model: nn.Module, example_input: torch.Tensor) - GraphModule: 通過(guò) torch.fx 追蹤模型計(jì)算圖并執(zhí)行 ConvBN 融合。 注意這里只演示核心邏輯實(shí)際還需要處理 BN 后的 ReLU 等節(jié)點(diǎn)。 model.eval() traced symbolic_trace(model) nodes list(traced.graph.nodes) for node in nodes: if node.op call_module and isinstance(traced.get_submodule(node.target), nn.BatchNorm2d): prev_node node.args[0] if prev_node.op call_module and isinstance(traced.get_submodule(prev_node.target), nn.Conv2d): fused_conv fuse_conv_bn( traced.get_submodule(prev_node.target), traced.get_submodule(node.target), ) # 將原 conv 模塊替換為融合后的 conv traced.delete_submodule(prev_node.target) traced.add_submodule(prev_node.target, fused_conv) # 讓所有指向 BN 的節(jié)點(diǎn)直接指向原 conv node.replace_all_uses_with(prev_node) traced.graph.erase_node(node) traced.recompile() return traced這段代碼的核心思路是把 BN 的縮放系數(shù)和偏移量折算進(jìn)卷積的 weight 和 bias然后在計(jì)算圖中刪除 BN 節(jié)點(diǎn)。由于 ReLU 是逐元素操作一般可以放在融合的最后一步或者在后續(xù)算子中直接包含激活函數(shù)。注意這里只是核心片段要放入redwood/graph_optimizer.py實(shí)際使用還需考慮模型中的不同模塊命名情況。4.2 自定義Triton算子MatMul ReLU融合圖優(yōu)化完成后還需要讓關(guān)鍵算子跑得更快。這里給出一個(gè)使用 Triton 編寫(xiě)的MatMul ReLU融合算子示例。它在一個(gè) kernel 內(nèi)完成矩陣乘法和激活計(jì)算避免中間結(jié)果的顯存讀寫(xiě)。# redwood/kernels/fused_matmul_relu.py import torch import triton import triton.language as tl triton.jit def fused_matmul_relu_kernel( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): pid_m tl.program_id(0) pid_n tl.program_id(1) offs_m pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) offs_k tl.arange(0, BLOCK_K) a_ptrs a_ptr (offs_m[:, None] * stride_am offs_k[None, :] * stride_ak) b_ptrs b_ptr (offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn) acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for k in range(0, K, BLOCK_K): a tl.load(a_ptrs) b tl.load(b_ptrs) acc tl.dot(a, b) a_ptrs BLOCK_K * stride_ak b_ptrs BLOCK_K * stride_bk # 融合 ReLU 激活 acc tl.maximum(acc, 0) offs_cm pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_cn pid_n * BLOCK_N tl.arange(0, BLOCK_N) c_ptrs c_ptr (offs_cm[:, None] * stride_cm offs_cn[None, :] * stride_cn) tl.store(c_ptrs, acc) def matmul_relu(a: torch.Tensor, b: torch.Tensor) - torch.Tensor: assert a.is_cuda and b.is_cuda M, K a.shape K2, N b.shape assert K K2 c torch.empty((M, N), devicea.device, dtypetorch.float32) BLOCK_M, BLOCK_N, BLOCK_K 16, 16, 16 grid (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) fused_matmul_relu_kernel[grid]( a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_MBLOCK_M, BLOCK_NBLOCK_N, BLOCK_KBLOCK_K, ) return c這個(gè) kernel 中每個(gè)線(xiàn)程塊負(fù)責(zé)計(jì)算結(jié)果矩陣中的一個(gè)BLOCK_M x BLOCK_N分塊內(nèi)層循環(huán)按BLOCK_K累加。與 PyTorch 自帶的torch.matmul不同這里沒(méi)有產(chǎn)生M x N x K的中間張量而是直接在累加后執(zhí)行ReLU然后寫(xiě)回結(jié)果。實(shí)際使用時(shí)應(yīng)根據(jù) GPU 型號(hào)調(diào)整BLOCK大小。4.3 調(diào)度器與內(nèi)存池推理服務(wù)端通常需要處理大量請(qǐng)求。如果每個(gè)請(qǐng)求都單獨(dú)調(diào)用一次 GPU 算子GPU 的利用率很低。Redwood 的調(diào)度層負(fù)責(zé)收集短時(shí)間內(nèi)到達(dá)的請(qǐng)求把它們拼成一個(gè) batch 后統(tǒng)一推理。# redwood/runtime.py import threading import time import torch class RedwoodInferenceRuntime: def __init__(self, model, max_batch_size8, wait_time0.01): self.model model self.max_batch_size max_batch_size self.wait_time wait_time self._queue [] self._lock threading.Lock() def submit(self, tensor): with self._lock: self._queue.append(tensor) time.sleep(self.wait_time) with self._lock: if len(self._queue) self.max_batch_size: return self._flush_locked() return None def _flush_locked(self): batch torch.cat(self._queue, dim0) self._queue.clear() with torch.no_grad(): return self.model(batch) def flush(self): with self._lock: if self._queue: return self._flush_locked() return None這段代碼是一個(gè)簡(jiǎn)化版動(dòng)態(tài) batch 調(diào)度器。請(qǐng)求到達(dá)后不立即推理而是等待一小段時(shí)間盡量積攢成一個(gè) batch。當(dāng) batch 大小達(dá)到上限后立即處理。實(shí)際生產(chǎn)環(huán)境還需要處理隊(duì)列積壓、超時(shí)、顯存池復(fù)用等問(wèn)題但核心思路可以作為入門(mén)實(shí)現(xiàn)。內(nèi)存池方面Redwood 可以復(fù)用 PyTorch 的 caching allocator同時(shí)在包含可變長(zhǎng)度輸入的業(yè)務(wù)中通過(guò)對(duì)輸入做 padding 和緩存輸出緩沖區(qū)減少反復(fù)申請(qǐng)顯存。真正實(shí)現(xiàn)時(shí)可以?xún)?yōu)先觀察torch.cuda.memory_stats()判斷是否存在頻繁顯存分配。4.4 FastAPI推理服務(wù)最后把 Redwood 的圖優(yōu)化和運(yùn)行時(shí)封裝成一個(gè) HTTP 服務(wù)。這里使用 FastAPI 提供/predict接口。# redwood/server.py import io import torch from fastapi import FastAPI from pydantic import BaseModel from redwood.graph_optimizer import optimize_graph from redwood.runtime import RedwoodInferenceRuntime app FastAPI(titleRedwood Inference Service) class PredictRequest(BaseModel): data: list class PredictResponse(BaseModel): result: list cost_ms: float _model None _runtime None def load_model(): global _model, _runtime from models.example_model import ExampleModel model ExampleModel() checkpoint torch.load(models/example_model.pt, map_locationcpu) model.load_state_dict(checkpoint) model.eval() example_input torch.randn(1, 3, 224, 224) optimized_model optimize_graph(model, example_input) _runtime RedwoodInferenceRuntime(optimized_model) app.on_event(startup) def startup(): load_model() app.post(/predict, response_modelPredictResponse) def predict(req: PredictRequest): import time tensor torch.tensor(req.data, dtypetorch.float32) start time.time() with torch.no_grad(): output _runtime.submit(tensor) cost_ms (time.time() - start) * 1000 return PredictResponse(resultoutput.tolist(), cost_mscost_ms)這里省略了復(fù)雜的模型加載和批處理細(xì)節(jié)重點(diǎn)展示服務(wù)層如何調(diào)用底層 Runtime。通過(guò)app.on_event(startup)在服務(wù)啟動(dòng)時(shí)加載和優(yōu)化模型避免每次預(yù)測(cè)都重復(fù)做圖優(yōu)化。5. 部署與驗(yàn)證5.1 容器鏡像構(gòu)建Redwood 使用 Docker 部署可以保證運(yùn)行時(shí)環(huán)境一致。以下是一個(gè)基礎(chǔ) Dockerfile 示例。FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 ENV DEBIAN_FRONTENDnoninteractive RUN apt-get update apt-get install -y \ python3.10 python3.10-dev python3-pip \ rm -rf /var/lib/apt/lists/* WORKDIR /app COPY requirements.txt . RUN pip3 install --no-cache-dir -r requirements.txt COPY . . EXPOSE 8000 CMD [uvicorn, redwood.server:app, --host, 0.0.0.0, --port, 8000]部署到測(cè)試環(huán)境前務(wù)必在測(cè)試環(huán)境驗(yàn)證流程。5.2 啟動(dòng)服務(wù)本地開(kāi)發(fā)環(huán)境啟動(dòng)服務(wù)可以直接使用 Uvicornuvicorn redwood.server:app --host 0.0.0.0 --port 8000使用 Docker Compose 啟動(dòng)時(shí)可以參考下面的配置# docker-compose.yml services: redwood: build: . ports: - 8000:8000 volumes: - ./models:/app/models environment: - CUDA_VISIBLE_DEVICES0 deploy: resources: reservations: devices: - driver: nvidia count: 1 capabilities: [gpu]啟動(dòng)命令docker-compose up --build注意NVIDIA Container Toolkit 需要提前安裝否則容器內(nèi)無(wú)法識(shí)別 GPU。5.3 性能對(duì)比Redwood 上線(xiàn)后的性能驗(yàn)證不能只關(guān)注單個(gè)算子的耗時(shí)還要關(guān)注端到端推理延遲和吞吐量。我們寫(xiě)一個(gè)簡(jiǎn)單的基準(zhǔn)腳本# scripts/benchmark.py import time import torch from redwood.graph_optimizer import optimize_graph from models.example_model import ExampleModel model ExampleModel().eval() example_input torch.randn(1, 3, 224, 224) # 原始模型 start time.time() for _ in range(100): with torch.no_grad(): model(example_input) original_cost (time.time() - start) / 100 * 1000 # 優(yōu)化后的模型 optimized optimize_graph(model, example_input) start time.time() for _ in range(100): with torch.no_grad(): optimized(example_input) optimized_cost (time.time() - start) / 100 * 1000 print(fOriginal: {original_cost:.3f} ms) print(fOptimized: {optimized_cost:.3f} ms) print(fSpeedup: {original_cost / optimized_cost:.2f}x)在真實(shí)項(xiàng)目中建議至少跑 1000 次并先 warmup 數(shù)次避免首次調(diào)用包含 CUDA kernel 加載耗時(shí)。性能數(shù)據(jù)可能因模型和 GPU 不同而差異很大關(guān)鍵是觀察優(yōu)化前后同一環(huán)境下的相對(duì)提升。5.4 準(zhǔn)確性驗(yàn)證加速器不能只快不準(zhǔn)。每次圖優(yōu)化后都要對(duì)比模型輸出與原始輸出的誤差。一般情況下ConvBN 融合是數(shù)值等價(jià)變換誤差應(yīng)保持在非常小的范圍。def verify_consistency(original, optimized, inputs): with torch.no_grad(): y1 original(inputs) y2 optimized(inputs) max_diff (y1 - y2).abs().max().item() rel_diff ((y1 - y2).abs() / (y1.abs() 1e-6)).max().item() print(fmax absolute diff: {max_diff:.6e}) print(fmax relative diff: {rel_diff:.6e})如果誤差明顯偏大優(yōu)先檢查融合過(guò)程中的均值、方差取的是訓(xùn)練階段的統(tǒng)計(jì)量還是當(dāng)前 batch 的統(tǒng)計(jì)量。推理階段必須使用running_mean和running_var否則輸出會(huì)有偏差。6. 常見(jiàn)問(wèn)題與排查思路在 Redwood 的開(kāi)發(fā)與部署過(guò)程中我們遇到了不少問(wèn)題。這里整理成表格方便后續(xù)讀者按圖索驥。問(wèn)題現(xiàn)象常見(jiàn)原因解決思路torch.cuda.is_available() 為 falseNVIDIA 驅(qū)動(dòng)或 CUDA 版本不匹配用nvidia-smi檢查驅(qū)動(dòng)重新安裝匹配的 PyTorchTriton kernel 編譯報(bào)錯(cuò)BLOCK 參數(shù)設(shè)置不適合 GPU調(diào)整 BLOCK_M/N/K或檢查是否使用 GPU 運(yùn)行推理結(jié)果與原始模型不一致BN 融合時(shí)使用了訓(xùn)練模式參數(shù)確保模型處于 eval 狀態(tài)使用 running_mean/varFastAPI 服務(wù)啟動(dòng)慢圖優(yōu)化在啟動(dòng)階段執(zhí)行提前離線(xiàn)保存優(yōu)化后的模型啟動(dòng)時(shí)直接加載并發(fā)請(qǐng)求時(shí) GPU 利用率低沒(méi)有批量調(diào)度啟用 Redwood 的 batch 調(diào)度調(diào)整 wait_time顯存占用持續(xù)上漲內(nèi)存池沒(méi)有復(fù)用或存在緩存泄漏使用torch.cuda.memory_stats()定位限制緩存清理策略容器內(nèi)無(wú)法識(shí)別 GPU未安裝 nvidia-container-toolkit安裝相應(yīng)版本并重啟容器服務(wù)下面展開(kāi)幾個(gè)高頻問(wèn)題的排查過(guò)程。6.1 Triton Kernel 編譯失敗Triton 在首次執(zhí)行 kernel 時(shí)會(huì)有 JIT 編譯過(guò)程。如果設(shè)置不當(dāng)編譯時(shí)間會(huì)很長(zhǎng)甚至直接報(bào)錯(cuò)。常見(jiàn)原因是BLOCK_SIZE與 GPU 寄存器限制不匹配。排查思路先查看完整報(bào)錯(cuò)棧確認(rèn)是否在tl.dot或tl.load附近。調(diào)小BLOCK_M、BLOCK_N例如從 32 改為 16。確認(rèn)輸入張量是 GPU 上的連續(xù)內(nèi)存。在代碼中設(shè)置TRITON_KERNEL_OVERRIDE1或檢查 Triton 緩存目錄權(quán)限。6.2 圖優(yōu)化后模型輸出不一致這是優(yōu)化器最容易踩的坑。ConvBN 融合時(shí)如果模型處于 training 模式BN 層會(huì)使用當(dāng)前 batch 的均值和方差導(dǎo)致融合結(jié)果錯(cuò)誤。因此執(zhí)行optimize_graph之前必須調(diào)用model.eval()并確認(rèn)所有 BN 層的running_mean和running_var已經(jīng)完成更新。如果在加載 checkpoint 后立即做融合需要先跑一次驗(yàn)證模式的前向或者在保存模型時(shí)就確保 BN 統(tǒng)計(jì)量正確。最簡(jiǎn)單的方式是加載完權(quán)重后先執(zhí)行model.eval()再做 graph trace。6.3 動(dòng)態(tài) batch 導(dǎo)致請(qǐng)求阻塞Redwood 的簡(jiǎn)單調(diào)度器會(huì)等待wait_time來(lái)聚合請(qǐng)求。如果業(yè)務(wù)流量本身很小等待時(shí)間會(huì)導(dǎo)致額外延遲。針對(duì)這種情況可以把wait_time設(shè)置得很小或者采用“達(dá)到最小 batch 立即推理”的策略。更復(fù)雜但穩(wěn)定的方案是使用定時(shí)器 flush 隊(duì)列比如每 5ms 檢查一次。7. 最佳實(shí)踐與工程建議7.1 用AI編程工具提升開(kāi)發(fā)效率Redwood 能在兩周內(nèi)完成很大程度上依賴(lài) AI 編程工具和 AI Agent 的輔助。寫(xiě) Triton kernel、調(diào)試 torch.fx 圖改寫(xiě)、生成 Dockerfile 等大量重復(fù)性工作都可以借助 AI 編程助手快速產(chǎn)出初稿。需要注意的是AI 生成代碼必須經(jīng)過(guò)測(cè)試驗(yàn)證不能直接信任尤其涉及 CUDA 算子和圖優(yōu)化邏輯時(shí)要重點(diǎn)審查數(shù)值正確性。推薦工作流是先自己梳理模塊邊界讓 AI 輔助生成函數(shù)骨架再用單元測(cè)試鎖定行為最后根據(jù)編譯報(bào)錯(cuò)迭代。這樣既能提升效率又能保證代碼質(zhì)量。7.2 安全與生產(chǎn)注意事項(xiàng)Redwood 涉及模型部署和推理服務(wù)生產(chǎn)環(huán)境變更時(shí)必須遵守幾個(gè)原則。在測(cè)試環(huán)境驗(yàn)證后再發(fā)布不要在業(yè)務(wù)高峰期直接改動(dòng)模型服務(wù)。模型文件和配置統(tǒng)一放在只讀目錄避免運(yùn)行時(shí)被篡改。對(duì)外服務(wù)接口需要鑒權(quán)預(yù)測(cè)接口不能裸奔暴露在公網(wǎng)。對(duì)請(qǐng)求體大小做限制避免超大輸入導(dǎo)致顯存溢出。日志記錄不要輸出完整的模型輸入輸出防止數(shù)據(jù)泄露。推理服務(wù)盡量使用最小權(quán)限容器賬號(hào)運(yùn)行避免使用 root。部署前還要做備份。如果使用模型版本管理工具上線(xiàn)新版本時(shí)可以快速回滾到舊版本。7.3 后續(xù)演進(jìn)方向Redwood 還只是一個(gè)淺層加速器原型后續(xù)可以從幾個(gè)方向繼續(xù)完善。支持更多融合模式例如 LayerNorm 激活、Attention 融合。引入動(dòng)態(tài)形狀處理解決 NLP 場(chǎng)景下變長(zhǎng)序列問(wèn)題。增加多模型管理不同模型使用不同的優(yōu)化策略。接入 Prometheus 監(jiān)控將延遲、吞吐、顯存指標(biāo)標(biāo)準(zhǔn)化。擴(kuò)展 CPU 后端讓 Redwood 在無(wú) GPU 環(huán)境也能運(yùn)行。如果團(tuán)隊(duì)業(yè)務(wù)逐漸穩(wěn)定可以考慮把 Redwood 的一部分能力逐漸貢獻(xiàn)到開(kāi)源社區(qū)或者對(duì)接成熟的推理引擎減少自研維護(hù)成本。最后想說(shuō)的是自研加速器并不一定適合所有團(tuán)隊(duì)。如果只是快速上線(xiàn)一個(gè)模型服務(wù)直接使用 ONNX Runtime 或 TensorRT 會(huì)更穩(wěn)妥。Redwood 的意義在于讓我們理解推理鏈路中哪些優(yōu)化真正有效并且保留了對(duì)業(yè)務(wù)模型的深度定制能力。希望這份兩周落地的經(jīng)驗(yàn)?zāi)芙o你提供參考也歡迎在實(shí)際項(xiàng)目中根據(jù)自身場(chǎng)景進(jìn)行調(diào)整。