
零基础上手 Inferentia:从开机器到跑出第一个推理
第一次碰 AWS 自研芯片,别去搞训练,先用 Inferentia 跑一次推理。训练要调的东西太多,你还没摸到芯片,先被环境劝退了。
我手头有一批文本要给因子打分,以前都租 GPU,账单看着肉疼。这几天在试 Inferentia,素材里说第一代 Inf1 实例比同等 GPU 实例大概 2.3 倍吞吐、成本低七成。这个数我自己验证不了,一台机器一个模型,说明不了什么。但方向是我要的,推理这种重复劳动,本来就该找条便宜路。
AWS 把两块芯片分得很清楚。Inferentia 管推理,追的是高吞吐低延迟;Trainium 管训练,吃的是计算密集型任务。名字容易记混,记住推理是 I 开头就行。
下面是从零到出结果的完整路径,我按控制台的实际顺序写,照着点。
先准备一个 AWS 账号,一张能付美元的卡。区域选美东,Inferentia 不是每个区都有货。
1. 登录控制台,右上角把区域切到 us-east-1。
2. 顶部搜索框输 EC2,进 实例 页面,点橙色的 启动实例。
3. 名称随便填,我叫它 test-inferentia。
4. 镜像这一项最容易翻车。点 Amazon Machine Image,搜 Neuron,选 Deep Learning AMI Neuron 那个 Ubuntu 版本。普通 DLAMI 里没有 Neuron 驱动,装了也是白装。
5. 实例类型选 inf1.xlarge。老一代,便宜,新手练手够用。想新一点就 inf2.xlarge。
6. 密钥对新建一个,.pem 文件会下载到本地,别删。
7. 网络、存储全默认,直接 启动实例。
等一两分钟,实例状态变成 运行中。点进去,右上角 连接,选 SSH 客户端,把那段命令复制到本地终端。看到 ubuntu@ip-xxx 的提示符,就进来了。
第二步,确认芯片在。敲 neuron-ls。这是 Neuron SDK 自带的小工具,作用就一个,告诉你这台机器上有几颗加速芯片、驱动版本对不对。正常会打出几行,能看到设备数量和型号。什么都没有,说明镜像选错了,别在原地折腾,重开一台比修环境快。
第三步,装依赖。pip install torch-neuronx transformers。这里注意别手贱升级 PyTorch,torch-neuronx 对版本挑得厉害,AMI 里配好的那套就是能跑的。
第四步,写一个最小脚本。从 HuggingFace 拉一个小的情感分类模型,给一句话打分。代码二十来行,关键只有一处,模型加载时加一个 torch_neuron 的后端参数,等于告诉它把模型往芯片上放。
第五步,编译。第一次跑,终端会一行行刷 Compiler status,这是 Neuron 把模型图翻译成芯片看得懂的格式。我这边等了几分钟。编完会在本地落一个缓存目录,第二次跑直接读缓存,快得多。
第六步,看结果。输出是一个正负分。到这一步,整条流程就通了。
踩坑的部分,我这边踩过,也看别人踩过,主要是这几个。
Trainium 和 Inferentia 记混,是最常见的。开一台 trn1 跑推理没必要,还贵。训练归训练,推理归推理。
neuron-ls 报找不到命令,九成是镜像选错,回头改第一步。
改了模型结构或者 batch 大小,缓存失效,会重新编译。终端看着像卡住,多等一会。
inf1.xlarge 有时提示容量不足。换个可用区,或者直接上 inf2 试试。
跑完记得 停止实例。加速实例按小时算钱,关终端不等于关机,这个坑最费钱。
学完这个,下一步可以试三件事。一是把 batch 慢慢加大,看吞吐怎么变,这是后面做容量规划的基础。二是拿 vLLM 部署一个小语言模型,Inferentia 对这条路支持得还行。三是如果你确实有训练需求,再去碰 Trainium 和 trn1 实例,那是另一套东西,好在 Neuron SDK 是通的,不用重新学一遍。
我这边就一台机器、一个模型,跑得顺不顺,样本量太小,不敢下结论。流程先跑通,别的以后再说。
物界前沿