Skip to content

Commit 095f6d2

Browse files
[doc] PAI deployment docs: ENABLE_AOT export + serving processor params (#536)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent a1ef4cf commit 095f6d2

2 files changed

Lines changed: 139 additions & 26 deletions

File tree

docs/source/usage/export.md

Lines changed: 105 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ torchrun --master_addr=localhost --master_port=32555 \
2929
- --checkpoint_path: 指定要导出的checkpoint, 默认导出model_dir下面最新的checkpoint
3030
- --export_dir: 导出到的模型目录,对于向量召回模型,导出命令会自动进行切图,分别放在user和item的子目录下
3131

32+
(export-env-vars)=
33+
3234
### 环境变量
3335

3436
- ODPS_ENDPOINT: 在PAI-DLC/PAI-DSW环境,数据为MaxCompute表的情况下需设置,详见[文档](../feature/data.md)的OdpsDataset章节
@@ -43,6 +45,107 @@ torchrun --master_addr=localhost --master_port=32555 \
4345
- **INPUT_TILE=3**:user侧fg和embedding计算仅一次,适用于user侧特征比较多的情况
4446
- INPUT_TILE_3_ONLINE: 配合INPUT_TILE=3使用,对User侧序列特征使用在线推理模式,序列特征在线模型服务中推理性能更好,但导出的模型无法用于离线预测
4547
- **INPUT_TILE_3_ONLINE=1**:启用序列特征的在线推理模式
46-
- ENABLE_AOT:
47-
- **ENABLE_AOT=1**: 使用AOT(Ahead Of Time)编译优化导出优化的模型
48+
- ENABLE_AOT: 使用AOT(Ahead Of Time)编译优化导出的模型,可显著提升推理性能。**AOT 编译产物与导出机器的 GPU 架构强绑定,在线服务的 GPU 类型必须与导出时使用的 GPU 类型完全一致**,详见下文 [在 PAI 上导出 AOT 模型](#export-aot-on-pai) 章节
49+
- **ENABLE_AOT=1**: 使用AOT编译优化导出模型(sparse 部分用 JIT,dense 部分用 AOTI)
4850
- **ENABLE_AOT=2**: 使用统一 AOTI 模型编译优化 (sparse + dense 融合为单一 .pt2) [experimental]
51+
52+
(export-aot-on-pai)=
53+
54+
## 在 PAI 上导出 AOT 模型
55+
56+
### GPU 架构匹配要求
57+
58+
开启 `ENABLE_AOT` 后,导出产物(`.pt2` 中包含针对导出机器 GPU 算力/SM 版本编译的 cubin 与 Triton kernel)与 GPU 架构强绑定。因此:
59+
60+
- **PAI-EAS 在线服务使用的 GPU 类型,必须与导出时使用的 GPU 类型完全一致**,否则在线加载或推理会失败、或产生错误结果。
61+
- 更换 EAS 服务的 GPU 类型后,必须在对应 GPU 上重新导出模型。
62+
63+
根据 PAI-DLC 是否能提供与目标 EAS 服务相同的 GPU 类型,有以下两种导出方式。
64+
65+
### 场景一:PAI-DLC 有与 EAS 服务 GPU 类型相同的资源
66+
67+
直接在 PAI-DLC 上导出(推荐)。在 [DLC Tutorial 的导出章节](../quick_start/dlc_tutorial.md) 的命令前加上 `ENABLE_AOT=1` 即可,例如:
68+
69+
```bash
70+
ENABLE_AOT=1 CUDA_HOME=/usr/local/cuda-12 \
71+
torchrun --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
72+
--nnodes=$WORLD_SIZE --nproc-per-node=$NPROC_PER_NODE --node_rank=$RANK \
73+
-m tzrec.export \
74+
--pipeline_config_path /mnt/data/${MODEL_DIR}/pipeline.config \
75+
--export_dir /mnt/data/${MODEL_DIR}/export
76+
```
77+
78+
- `CUDA_HOME`: AOT 编译需要 `nvcc`,指向 CUDA 安装目录。
79+
- 其余量化、INPUT_TILE 等环境变量见上文 [环境变量](#export-env-vars) 章节,按需添加。
80+
81+
### 场景二:PAI-DLC 没有与 EAS 服务 GPU 类型匹配的资源
82+
83+
此时在 PAI-EAS 上以 `workload_type=elasticjob` 任务方式导出。由于该任务与在线服务运行在**相同的 EAS 资源组(同一 GPU 硬件)**,可保证导出产物与在线服务的 GPU 架构一致;导出产物写入挂载的数据集(NAS/OSS),随后即可用于部署在线服务(参见 [模型服务](serving.md))。
84+
85+
任务服务的 JSON 配置示例如下(各字段含义参见 [PAI-EAS 服务参数文档](https://help.aliyun.com/zh/pai/parameters-of-model-services),请将 `${...}` 占位符替换为自己的资源标识):
86+
87+
```json
88+
{
89+
"cloud": {
90+
"networking": {
91+
"security_group_id": "${SG_ID}",
92+
"vpc_id": "${VPC_ID}",
93+
"vswitch_id": "${VSWITCH_ID}"
94+
}
95+
},
96+
"containers": [
97+
{
98+
"image": "dsw-registry-vpc.{region}.cr.aliyuncs.com/pai/torcheasyrec:1.2.0-pytorch2.11.0-gpu-py311-cu129-ubuntu22.04",
99+
"port": 8000,
100+
"script": "set -e\npip install tzrec==${TZREC_NIGHTLY_VERSION} -f http://tzrec.oss-accelerate.aliyuncs.com/release/nightly/repo.html --trusted-host tzrec.oss-accelerate.aliyuncs.com\n\nCUDA_HOME=/usr/local/cuda-12 ODPS_ENDPOINT=http://service.{region}-vpc.maxcompute.aliyun-inc.com/api \\\nENABLE_AOT=1 QUANT_EC_EMB=INT8 \\\ntorchrun --master_addr=127.0.0.1 --master_port=12345 \\\n --nnodes=1 --nproc-per-node=1 --node_rank=0 \\\n -m tzrec.export \\\n --pipeline_config_path /mnt/data/${MODEL_DIR}/pipeline.config \\\n --export_dir /mnt/data/${MODEL_DIR}/export/ \\\n --additional_export_config '{\"cand_seq_pk\": \"cand_seq\"}'"
101+
}
102+
],
103+
"metadata": {
104+
"cpu": 60,
105+
"disk": "60Gi",
106+
"gpu": 1,
107+
"instance": 1,
108+
"memory": 383000,
109+
"name": "tzrec_aot_export",
110+
"resource": "${EAS_RESOURCE_ID}",
111+
"resource_burstable": false,
112+
"rpc": {
113+
"keepalive": 5000
114+
},
115+
"workload_type": "elasticjob",
116+
"workspace_id": "${WORKSPACE_ID}"
117+
},
118+
"options": {
119+
"enable_ram_role": true
120+
},
121+
"storage": [
122+
{
123+
"dataset": {
124+
"id": "${DATASET_ID}",
125+
"read_only": false,
126+
"version": "v1"
127+
},
128+
"mount_path": "/mnt/data/"
129+
}
130+
]
131+
}
132+
```
133+
134+
关键字段说明:
135+
136+
- `cloud.networking`: 任务运行所在的 VPC/交换机/安全组。由于脚本需从公网 `pip install` 安装 tzrec,而 EAS 默认不开放公网访问,需为该 VPC 配置 NAT 网关并绑定 EIP 以开启公网访问。详见 [配置网络连通](https://help.aliyun.com/zh/pai/configure-network-connectivity)
137+
- `metadata.workload_type`: 设为 `elasticjob`,即以一次性任务(而非常驻服务)方式运行,导出完成后任务自动退出。
138+
- `metadata.resource`: 填写**目标在线服务所用的 EAS 专属资源组 ID**,使导出运行在与服务相同的 GPU 上,从而保证 GPU 架构匹配。
139+
- `metadata.gpu` / `instance`: 单卡单实例即可完成导出。
140+
- `containers.image`: 运行任务的容器镜像,使用 PAI 官方 TorchEasyRec 镜像(已内置 CUDA、PyTorch 等依赖),`{region}` 替换为所在地域,镜像 tag 按需选择。
141+
- `containers.script`: 安装 tzrec 后执行 `tzrec.export`。其中 `ENABLE_AOT=1` 开启 AOT 编译,`QUANT_EC_EMB``CUDA_HOME``ODPS_ENDPOINT` 等见上文 [环境变量](#export-env-vars) 章节按需设置;HSTU 类模型通过 `--additional_export_config '{"cand_seq_pk": "cand_seq"}'` 指定候选序列特征名。
142+
- `storage`: 挂载数据集到 `/mnt/data/``--export_dir` 写入该挂载路径下,供后续在线服务加载。
143+
- `options.enable_ram_role`: **必须设为 `true`**,任务通过 RAM 角色访问 MaxCompute、OSS/NAS、镜像仓库等云资源,无需配置长期 AccessKey。需提前在 RAM 中创建受信主体为 PAI(`eas.pai.aliyuncs.com`)的角色并授予对应资源访问权限。详见 [配置 EAS 的 RAM 角色](https://help.aliyun.com/zh/pai/configure-the-eas-ram-role)
144+
145+
将上述 JSON 保存为 `aot_export.json`,可在 PAI-EAS 控制台创建任务服务,或通过 `eascmd` 提交:
146+
147+
```bash
148+
eascmd -i ${ACCESS_KEY_ID} -k ${ACCESS_KEY_SECRET} -e ${ENDPOINT} create aot_export.json
149+
```
150+
151+
任务运行结束后,`--export_dir` 指向的目录即为导出好的模型,将其作为在线服务的模型路径部署即可(部署方式参见 [模型服务](serving.md))。

0 commit comments

Comments
 (0)