不装PyTorch,推理引擎从6.82GB瘦到1.48GB

通过更聚焦的阅读结构呈现文章主线、时间信息与延展入口,帮助访问者更轻松地完成完整阅读与继续浏览。

不装PyTorch,推理引擎从6.82GB瘦到1.48GB
READ ARTICLE DETAIL

不装PyTorch,推理引擎从6.82GB瘦到1.48GB

EBET易博官网|赛程比分与球队动态 337 阅读 约 5.5 分钟

想用Python在GPU上跑一个现代模型,默认答案往往是"装PyTorch"。训练场景下这没问题,但放到推理上,它带来的东西有点多:庞大的安装体积、以CUDA为默认视角的世界观,以及一套内存行为你得靠摸索而不是靠设定去掌握的运行时。 nbxtensor是NeuroBrix给出的推理侧答案:一个带Triton内核、内存可规划、调优可认证的张量库,而且不需要PyTorch。安装只要一行 pip install nbxtensor ,凡是Triton能跑的地方它都能跑——NVIDIA、AMD、Intel和苹果GPU,Linux、Windows、macOS都覆盖。 "不需要PyTorch"到底指什么? 不是"PyTorch可选"。nbxtensor从不导入torch——不作为依赖,不作为可选附加项,也不藏在某个回退分支里。它的内核、测试,以及NeuroBrix在其上运行的模型,全部能在没装torch的环境里工作。 实际效果体现在基于它构建的NeuroBrix引擎上:在干净机器上实测,轻量安装(nbxtensor上的Triton引擎)为 1.48 GB ;同一个引擎加上PyTorch后是 6.82 GB 。如果你本来就活在PyTorch代码库里,双方通过DLPack交换数据,哪个方向都不会被锁死。 四个组成部分 张量 :设备内存、形状、步长、视图、数据类型及其提升规则——跑一张计算图所需的管道,底下不垫一个框架。 Triton内核 :矩阵乘、注意力、卷积、归一化、逐元素运算和归约,都用Triton写,一份源码服务Triton支持的所有GPU厂商。一个设备由一份事实文件描述——tile尺寸、共享内存、矩阵单元——而不是由代码描述。 可规划内存 :每设备的分配器藏在统一接口后面,另有一个规划器在首次运行前就把缓冲区尺寸算好。每个内核都声明自己的开销(输出、工作区和临时空间),规划器因此能在运行前给一个模型"报价"。在NeuroBrix里,这个规划器叫Prism:它根据你的硬件配置决定模型是放进单卡、拆到多卡、流式加载,还是把最大的算子切块。 认证调优 :一个Triton内核有block尺寸、warps、stages等设置,合适的取值取决于GPU和形状。常见做法是在用户机器上首次运行时自动调优。nbxtensor改为提供按GPU类别一次性认证好的设置:每个候选先对照float64基准校验,通过的再计时,胜出的连同其证明一起存下来。目前苹果GPU上已认证 18,552 项设置。 如果你的GPU类别还没有证书,缺失的设置会在你机器上调优一次并保留;你可以把它认证后通过pull request分享出去,下一个用同款GPU的人就完全不用再调。 快在该快的地方 两个例子。在NeuroBrix 2026年11月5日于V100 16 GB GPU上的全局基准中,nbxtensor上的Triton引擎在49个模型里有18个快过PyTorch引擎,Janus-Pro-7B上最多快5倍(212.6秒降到41.8秒)。 在Volta架构上,Triton会把矩阵乘编译成FMA,V100的张量核心因此闲置。nbxtensor用手写的m8n8k4内核把它们重新用起来,按算子逐个实测选择,而不是按规则决定。 苹果芯片 PyPI上没有面向macOS的Triton。nbxtensor用一条命令装上NeuroBrix为苹果GPU自建的Triton版本(NeuroBrix/triton-ext): pip install nbxtensor python -m nbxtensor install-metal python -m nbxtensor doctor,输出backend in force: metal doctor在所有平台都能用:它会报出当前生效的后端,并在缺少GPU运行时或Triton后端时,给出对应的安装命令。 用在哪里 nbxtensor是NeuroBrix的Triton引擎的基础。该引擎运行9个家族共49个模型——语言、代码、视觉语言、全模态、语音识别、语音理解、文本转语音、图像、视频和超分——统一使用一种.nbx容器格式和一个入口。NeuroBrix每次发版都会通过这个入口在两个引擎上跑完49个模型,并把输出与判定输出逐字节比对。 这一版以推理为先,训练不在其中。项目以Apache 2.0协议开源,仓库在github.com/NeuroBrix/nbxtensor。 特别

about image
觉得有用?分享给朋友

让更多人看到这篇内容

推荐阅读