导出

export_config

模型导出任务依赖pipeline.config中的export_config 配置

export_config {
}
  • exporter_type: 导出类型, latest | best ,默认latest

    • latest 导出最新的模型

    • best 导出最好的模型

  • best_exporter_metric: 当exporter_type为best的时候,确定最优导出模型的metric,注意该metric要在对应任务的metrics设置了才行。对于多任务模型则需要设置 {metric_name}_{tower_name}。

  • metric_larger_is_better: 确定最优导出模型的metric是越大越好,还是越小越好,默认是越大越好

  • use_dense_ema: 是否使用稠密EMA参数导出模型。未配置时默认跟随train_config.dense_optimizer.ema是否启用;可显式配置为true或false覆盖默认行为,在线Dense导出同样使用该配置

  • dynamic_embedding: 分布式 embedding 导出时的 dynemb 读取块大小和可选 FeatureDB 全量上传配置,详见下文 dynemb 流式导出与 FeatureDB 全量上传

导出命令

torchrun --master_addr=localhost --master_port=32555 \
    --nnodes=1 --nproc-per-node=2 --node_rank=0 \
    -m tzrec.export \
    --pipeline_config_path experiments/multi_tower_din_taobao_local/pipeline.config \
    --export_dir experiments/multi_tower_din_taobao_local/export
  • –pipeline_config_path: 导出用的配置文件

  • –checkpoint_path: 指定要导出的checkpoint, 默认导出model_dir下面最新的checkpoint

  • –export_dir: 导出到的模型目录,对于向量召回模型,导出命令会自动进行切图,分别放在user和item的子目录下

环境变量

  • ODPS_ENDPOINT: 在PAI-DLC/PAI-DSW环境,数据为MaxCompute表的情况下需设置,详见文档的OdpsDataset章节

  • ODPS_CONFIG_FILE_PATH: 在本地环境,数据为MaxCompute表的情况下需设置为odps_conf的路径,详见文档的OdpsDataset章节

  • QUANT_EMB: 对EmebddingBagCollection(非序列特征)参数进行量化,默认开启,INT8量化在大部分场景中对AUC等指标基本无损,且能大幅提升推理性能

    • QUANT_EMB=INT8:启用量化,默认已启用,并且默认为INT8量化,可以支持FP32,FP16,INT8,INT4,INT2

    • QUANT_EMB=0:关闭量化

  • QUANT_EC_EMB: 对EmebddingCollection(序列特征)参数进行量化,默认关闭,INT8量化在大部分场景中对AUC等指标基本无损,且能大幅提升推理性能

    • QUANT_EC_EMB=INT8:启用量化,默认关闭,可以支持FP32,FP16,INT8,INT4,INT2

  • INPUT_TILE: 对User侧特征自动扩展,开启可减少请求大小、网络传输时间和计算时间,默认关闭。必须在TorchEasyRec Processor的fg_mode=normal下使用

    • INPUT_TILE=2:user侧特征fg仅计算一次

    • INPUT_TILE=3:user侧fg和embedding计算仅一次,适用于user侧特征比较多的情况

  • INPUT_TILE_3_ONLINE: 配合INPUT_TILE=3使用,对User侧序列特征使用在线推理模式,序列特征在线模型服务中推理性能更好,但导出的模型无法用于离线预测

    • INPUT_TILE_3_ONLINE=1:启用序列特征的在线推理模式

  • USE_DISTRIBUTED_EMBEDDING: 开启分布式 embedding 导出模式,导出 dense graph 与分片 sparse embedding 参数;该模式会自动使用 INPUT_TILE=3

    • USE_DISTRIBUTED_EMBEDDING=1:启用分布式 embedding 导出

    • 普通 sparse、ZCH 和原生 dynemb 模型使用统一导出流程;dynamic_embedding 仅调整 dynemb 读取和上传参数

    • 配置了 zch 的特征会自动转换为动态表导出,详见 零冲突Hash Embedding 的分布式 Embedding 导出章节

    • PYTORCH dense 图使用 CPU 导出,TRITON/CUTLASS dense 图需要 CUDA;导出目录保留 full/dense graph 调试文件,不再生成 graph/gm_sparse.graph

  • DIST_QUANT: 分布式 embedding 导出模式下的 sparse embedding 参数量化开关,默认关闭;当前仅支持 INT8,与普通导出的 QUANT_EMB / QUANT_EC_EMB 不同

    • USE_DISTRIBUTED_EMBEDDING=1 DIST_QUANT=INT8:启用分布式 sparse embedding INT8 量化

    • 未设置 DIST_QUANT 时,不启用分布式 sparse embedding 量化

  • ENABLE_AOT: 使用AOT(Ahead Of Time)编译优化导出的模型,可显著提升推理性能。AOT 编译产物与导出机器的 GPU 架构强绑定,在线服务的 GPU 类型必须与导出时使用的 GPU 类型完全一致,详见下文 在 PAI 上导出 AOT 模型 章节

    • ENABLE_AOT=1: 使用AOT编译优化导出模型(sparse 部分用 JIT,dense 部分用 AOTI)

    • ENABLE_AOT=2: 使用统一 AOTI 模型编译优化 (sparse + dense 融合为单一 .pt2) [experimental]

dynemb 流式导出与 FeatureDB 全量上传

在 USE_DISTRIBUTED_EMBEDDING=1 下,dynemb 默认按块读取 checkpoint 并写入 NPZ,不需要将全部 dynemb 参数恢复到主机内存或显存。配置 export_config.dynamic_embedding.feature_store_config 后,同一批 dynemb 参数还会上传到 FeatureDB;不配置时只生成导出文件。该能力适用于原生 dynemb 参数。普通 sparse 和 ZCH 表仍按原有方式恢复到内存,其内存占用不受 dynemb 读取块预算限制。

配置 dynamic_embedding 要求 USE_DISTRIBUTED_EMBEDDING=1 且非 RTP 导出,否则导出启动时会报错。

配置

在 pipeline.config 中添加以下必填配置,其他参数使用默认值,仅在需要调整时显式设置:

export_config {
  dynamic_embedding {
    feature_store_config {
      region: "cn-beijing"
      project_name: "recommendation"
      feature_view_name: "model_embeddings"
      version: "full_ckpt100000_run1"
    }
  }
}

feature_store_config 复用 FeatureStoreConfig,必填参数如下:

配置项

导出时的含义

region

FeatureStore 地域;显式为空时从 ALIBABA_CLOUD_REGION 获取

project_name

已存在的 FeatureStore 项目

feature_view_name

DynamicEmbedding FeatureView;不存在时按配置创建

version

每次导出使用新的、尚未用于线上服务的版本,不能为 default

凭证沿用阿里云默认凭证链,以及运行环境中的 FEATUREDB_USERNAME、FEATUREDB_PASSWORD;不要将凭证写入 pipeline.config。

执行导出

从仓库根目录、依赖版本匹配的项目环境执行,显式指定一个已写入完成、导出期间不会再改变的 checkpoint,并使用新的导出目录:

PYTHONPATH=. USE_DISTRIBUTED_EMBEDDING=1 \
torchrun --standalone --nnodes=1 --nproc-per-node=1 \
    -m tzrec.export \
    --pipeline_config_path /mnt/model/pipeline.config \
    --checkpoint_path /mnt/model/model.ckpt-100000 \
    --export_dir /mnt/model/export/full_ckpt100000_run1

未设置 DIST_QUANT 时,NPZ 和 FeatureDB 使用 FP32 values,FeatureView 类型必须为 ARRAY<FLOAT>。在命令前增加 DIST_QUANT=INT8 后,两处使用同一份 UINT8 行量化结果,包含每行 4 字节的 scale/offset;FeatureView 类型必须为 ARRAY<UINT8>,导出时会校验已存在的表类型。量化维度必须为偶数。

在线学习 dense 模型热导出

在 USE_DISTRIBUTED_EMBEDDING=1 的在线学习场景下,sparse embedding 通过增量 embedding dump(train_config.delta_embedding_dump_config)持续上传 FeatureStore,训练进程可以在训练过程中同步导出 dense 图,供推理 Processor 热切换。该能力不依赖 checkpoint:rank 0 在训练启动时一次性构建 serving 侧 dense 图,之后每次触发时全体 rank 从内存中的 DMP 模型收集 dense 权重(仅收集 dense 图实际携带的参数,不涉及 sparse/dynamicemb/MCH 状态),rank 0 在后台线程热替换权重、script 并构建新版本目录,训练主流程不阻塞在导出上。

dense 导出触发与增量 embedding dump 的粗粒度边界统一:每个 paired_dump_interval_steps / paired_dump_interval_minutes 边界执行一次 delta dump 并触发一次 dense 导出,同一训练步的 (dense, sparse) 严格成对。粗粒度参数为必选(steps 与 minutes 二选一)。在此基础上可选配置更细粒度的 sparse_dump_interval_steps / sparse_dump_interval_minutes(维度必须与粗粒度一致且数值严格更小):细粒度边界只执行 delta dump 与 FeatureStore 上传、不导出 dense,让线上 embedding 以更高频率更新,而 dense 版本仍按粗粒度节奏发布;两个边界落在同一训练步时粗粒度优先,只 dump 一次。训练正常结束时若存在尾部增量,也会成对补一次最终导出(尾部若恰好落在某个 dump 边界上,则不重复 dump,最终导出与该步配对)。dense 版本构建完成后不会立即生效:current.json 仅在该步的 sparse 增量在所有 rank 上都完成 FeatureStore 上传后才原子翻转(联合发布),保证推理侧绝不出现 dense 领先 sparse——sparse 领先 dense 是唯一允许的偏斜方向。

启用方式

训练配置中必须开启增量 embedding dump 并配置 FeatureStore 上传(dense 导出节奏完全由 dump 决定,发布水位来自上传完成进度,详见 DeltaEmbeddingDumpConfig):

train_config {
  ...
  delta_embedding_dump_config {
    paired_dump_interval_minutes: 10
    # 可选:每 2 分钟只做 delta dump + 上传,不导 dense
    sparse_dump_interval_minutes: 2
    feature_store_config {
      ...
    }
  }
}

训练命令前加上以下环境变量即可(ONLINE_DENSE_EXPORT_DIR 为必填):

ONLINE_DENSE_EXPORT=1 \
ONLINE_DENSE_EXPORT_DIR=/mnt/data/serving \
torchrun --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
    --nnodes=$WORLD_SIZE --nproc-per-node=$NPROC_PER_NODE --node_rank=$RANK \
    -m tzrec.train_eval \
    --pipeline_config_path experiments/online_learning/pipeline.config

环境变量

  • ONLINE_DENSE_EXPORT: 开启训练内在线 dense 导出,默认关闭

    • ONLINE_DENSE_EXPORT=1:启用(要求同时设置 USE_DISTRIBUTED_EMBEDDING=1,并在 train_config 中配置带 feature_store_config 的 delta_embedding_dump_config;仅本地落盘的 dump 没有上传水位可供门控,缺少任一项训练启动即报错)

  • ONLINE_DENSE_EXPORT_DIR: 在线服务读取的根目录(必填)。导出产物发布在 <ONLINE_DENSE_EXPORT_DIR>/dense_hot_export 下,必须为本地路径(发布依赖 os.rename 原子换版);model_dir 可以是远端(如 OSS)

  • ONLINE_DENSE_EXPORT_KEEP_VERSIONS: 保留的历史版本数,默认 0(不删除任何已导出的版本)。大于 0 时保留最新 K 个版本,且 K 最小为 3(serving 需要当前版本 + 上一版本用于原子切换)。current.json 指向的版本永远不被清理

  • ONLINE_DENSE_EXPORT_TIMEOUT: 单次导出的预算秒数,默认 3600。超时不会中断导出线程,仅打印告警,并用于训练结束时 drain 的等待上限

运行与排障说明

  • 仅 rank 0 执行建图与发布;启动时会对整条建图 + script 链路做一次试运行,任何 trace/script 失败会在训练开始前就暴露(fail-fast)。

  • 启动时还会校验 dense 图的每个参数都能在训练模型 state dict 中找到来源(INPUT_TILE=3 的 user 侧孪生模块自动回落到非 user 侧权重),校验失败训练不会启动。

  • 导出在后台线程执行,dense 版本构建失败(日志中 online dense export task failed)只跳过该版本、不翻转 current.json,训练继续,由后续版本取代;FeatureStore 上传失败则会使训练报错退出,指针同样不翻转。两种失败下推理侧都停留在上一个自洽版本。

  • 两次导出间隔小于单次导出耗时时,排队中的旧任务被最新版本顶替(latest-wins),不产生积压;水位到达后也只发布最新的就绪版本,被顶替的未发布版本目录会被直接删除,不占用 ONLINE_DENSE_EXPORT_KEEP_VERSIONS 的保留配额。

  • 新鲜度预算:dump/发布间隔必须大于单次增量上传耗时的 p99(可通过监控上传耗时占发布间隔的比例确认),否则全局水位持续落后于配对步,current.json 长期不翻转。

  • 建议周期性做全量再基线(anti-entropy):停训后用离线全量导出 + FeatureStore 全量导入,再用下述 bootstrap 命令发布配对的 dense 版本,以清理 crash 重启死时间线残留的旧值、僵尸 tombstone 等漂移。全量导入必须以重启式切换进行,不能在增量流写入的同时在线叠写。

  • MatchModel(向量召回)与 TDM 模型的完整导出按 user/item(或按模块)分目录,与单体 dense 热导出的布局不兼容,启用会直接报错,请使用完整的 tzrec.export。

  • 手动/离线从某个 checkpoint 导出一个版本(bootstrap,与训练内导出发布到同一目录结构)。该命令不经过水位门控、直接翻转 current.json,manifest 中 sparse_step 等于 checkpoint_step(--checkpoint_step 省略时从 model.ckpt-<step> 路径解析)、sparse_probes 为空(Processor 跳过探针验证),前提是 FeatureStore 中的全量 embedding 与该 checkpoint 同源:

USE_DISTRIBUTED_EMBEDDING=1 ONLINE_DENSE_EXPORT_DIR=/mnt/data/serving \
python -m tzrec.tools.online_dense_export \
    --pipeline_config_path experiments/online_learning/pipeline.config \
    --checkpoint_path /mnt/data/model/model.ckpt-1200 \
    --model_dir /mnt/data/model \
    --checkpoint_step 1200

在 PAI 上导出 AOT 模型

GPU 架构匹配要求

开启 ENABLE_AOT 后,导出产物(.pt2 中包含针对导出机器 GPU 算力/SM 版本编译的 cubin 与 Triton kernel)与 GPU 架构强绑定。因此:

  • PAI-EAS 在线服务使用的 GPU 类型,必须与导出时使用的 GPU 类型完全一致,否则在线加载或推理会失败、或产生错误结果。

  • 更换 EAS 服务的 GPU 类型后,必须在对应 GPU 上重新导出模型。

根据 PAI-DLC 是否能提供与目标 EAS 服务相同的 GPU 类型,有以下两种导出方式。

场景一:PAI-DLC 有与 EAS 服务 GPU 类型相同的资源

直接在 PAI-DLC 上导出(推荐)。在 DLC Tutorial 的导出章节 的命令前加上 ENABLE_AOT=1 即可,例如:

ENABLE_AOT=1 CUDA_HOME=/usr/local/cuda-12 \
torchrun --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
    --nnodes=$WORLD_SIZE --nproc-per-node=$NPROC_PER_NODE --node_rank=$RANK \
    -m tzrec.export \
    --pipeline_config_path /mnt/data/${MODEL_DIR}/pipeline.config \
    --export_dir /mnt/data/${MODEL_DIR}/export
  • CUDA_HOME: AOT 编译需要 nvcc,指向 CUDA 安装目录。

  • 其余量化、INPUT_TILE 等环境变量见上文 环境变量 章节,按需添加。

场景二:PAI-DLC 没有与 EAS 服务 GPU 类型匹配的资源

此时在 PAI-EAS 上以 workload_type=elasticjob 任务方式导出。由于该任务与在线服务运行在相同的 EAS 资源组(同一 GPU 硬件),可保证导出产物与在线服务的 GPU 架构一致;导出产物写入挂载的数据集(NAS/OSS),随后即可用于部署在线服务(参见 模型服务)。

任务服务的 JSON 配置示例如下(各字段含义参见 PAI-EAS 服务参数文档,请将 ${...} 占位符替换为自己的资源标识):

{
    "cloud": {
        "networking": {
            "security_group_id": "${SG_ID}",
            "vpc_id": "${VPC_ID}",
            "vswitch_id": "${VSWITCH_ID}"
        }
    },
    "containers": [
        {
            "image": "dsw-registry-vpc.{region}.cr.aliyuncs.com/pai/torcheasyrec:1.5.0-pytorch2.14.0-gpu-py311-cu126-ubuntu22.04",
            "port": 8000,
            "script": "set -e\npip install tzrec[cu126]==${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\"}'"
        }
    ],
    "metadata": {
        "cpu": 60,
        "disk": "60Gi",
        "gpu": 1,
        "instance": 1,
        "memory": 383000,
        "name": "tzrec_aot_export",
        "resource": "${EAS_RESOURCE_ID}",
        "resource_burstable": false,
        "rpc": {
            "keepalive": 5000
        },
        "workload_type": "elasticjob",
        "workspace_id": "${WORKSPACE_ID}"
    },
    "options": {
        "enable_ram_role": true
    },
    "storage": [
        {
            "dataset": {
                "id": "${DATASET_ID}",
                "read_only": false,
                "version": "v1"
            },
            "mount_path": "/mnt/data/"
        }
    ]
}

关键字段说明:

  • cloud.networking: 任务运行所在的 VPC/交换机/安全组。由于脚本需从公网 pip install 安装 tzrec,而 EAS 默认不开放公网访问,需为该 VPC 配置 NAT 网关并绑定 EIP 以开启公网访问。详见 配置网络连通。

  • metadata.workload_type: 设为 elasticjob,即以一次性任务(而非常驻服务)方式运行,导出完成后任务自动退出。

  • metadata.resource: 填写目标在线服务所用的 EAS 专属资源组 ID,使导出运行在与服务相同的 GPU 上,从而保证 GPU 架构匹配。

  • metadata.gpu / instance: 单卡单实例即可完成导出。

  • containers.image: 运行任务的容器镜像,使用 PAI 官方 TorchEasyRec 镜像(已内置 CUDA、PyTorch 等依赖),{region} 替换为所在地域,镜像 tag 按需选择。

  • containers.script: 安装 tzrec 后执行 tzrec.export。其中 ENABLE_AOT=1 开启 AOT 编译,QUANT_EC_EMB、CUDA_HOME、ODPS_ENDPOINT 等见上文 环境变量 章节按需设置;HSTU 类模型通过 --additional_export_config '{"cand_seq_pk": "cand_seq"}' 指定候选序列特征名。

  • storage: 挂载数据集到 /mnt/data/,--export_dir 写入该挂载路径下,供后续在线服务加载。

  • options.enable_ram_role: 必须设为 true,任务通过 RAM 角色访问 MaxCompute、OSS/NAS、镜像仓库等云资源,无需配置长期 AccessKey。需提前在 RAM 中创建受信主体为 PAI(eas.pai.aliyuncs.com)的角色并授予对应资源访问权限。详见 配置 EAS 的 RAM 角色。

将上述 JSON 保存为 aot_export.json,可在 PAI-EAS 控制台创建任务服务,或通过 eascmd 提交:

eascmd -i ${ACCESS_KEY_ID} -k ${ACCESS_KEY_SECRET} -e ${ENDPOINT} create aot_export.json

任务运行结束后,--export_dir 指向的目录即为导出好的模型,将其作为在线服务的模型路径部署即可(部署方式参见 模型服务)。