Phoenix 训练

这一页回答什么

2026-08-13 起，phoenix/ 不再是演示模型。它包含训练和服务 For You 所用模型的代码，以及一套合成数据生成器（README.md:63、phoenix/README.md:9-21）。本页回答以下问题：

• 训练数据是什么格式，开源版用什么替代生产数据；
• 有哪些损失函数，包括 10-08 加入的停留时长损失和 10-13 加入的“停留达标”召回头；
• 一个训练步做哪些事，checkpoint 怎样写入、怎样恢复；
• Rust 服务层怎样加载 checkpoint 并对外提供 gRPC；
• 能否在本地跑一次概念验证训练，以及具体步骤。

核心结论

1. 上游发布的是生产训练与服务代码，不是样例。未发布的是生产数据流、集群编排、内部监控和训练好的 checkpoint（phoenix/README.md:9-21、phoenix/TRAINING.md:8-10）。
2. 排序与召回读取同一种训练数据：一行是一个用户会话快照，包含历史序列、候选序列和每条候选的动作标签。生产从 Kafka 流式读取；开源版用 reference/dump_gen.py 写出同格式的 parquet（phoenix/reference/README.md:100-109）。
3. 合成数据来自一个带种子的“世界”：1 万用户、2000 作者、5 万帖子、16 个话题。生成器按“用户话题偏好 × 作者质量 × 帖子质量”生成互动。模型能从中学到真实结构，但结果不代表线上效果（phoenix/reference/world.py:455-462、:301-311）。
4. 训练步由稠密优化器和稀疏逐行 AdaGrad 组成：排序旗舰与 nano 用 Muon，召回与旧排序预设用 AdamW。梯度非有限时，整步跳过（phoenix/TRAINING.md:12-36、phoenix/reference/train_step.py:101-148）。
5. 训练器在每次保存 checkpoint 时生成召回的候选索引，并写进 checkpoint。服务端直接加载该索引（phoenix/README.md:116-119、phoenix/xrex/train/trainer_recsys.py:4195-4250）。
6. 服务是 Python + Rust 混合进程：Rust 引擎 xai_recsys_engine（PyO3 扩展）负责 gRPC、请求排队、哈希与批次组装、大表加载。Python 侧用 JAX 执行模型前向（phoenix/xrex/inference/model_runner.py:3990、:4238、:4540、:5364）。
7. 可以在本地跑概念验证：需要 Linux 和 NVIDIA GPU。按 QUICKSTART.md 生成数据，训练 nano 排序与召回模型，启动三个服务，再发一次“召回 → 排序”请求。上游明确说明这只验证流程，不说明模型质量（phoenix/QUICKSTART.md:1-6、:169-170）。

整体流程

flowchart TB
    W["world.py<br/>带种子的合成世界"] --> SNAP["world_snapshots.py<br/>多模态嵌入、SID、发帖时间快照"]
    W --> DUMP["dump_gen.py<br/>会话快照 parquet"]
    SNAP -- "--sid 传入 SID 码" --> DUMP
    DUMP --> TRAIN["train_synth.py<br/>单卡训练 nano 配置"]
    TRAIN --> CKPT["checkpoint<br/>召回版含 post_embeddings 索引"]
    CKPT --> SERVE["launch_inference.py<br/>Rust gRPC + JAX 前向"]
    SNAP --> SIDSRV["sid_index_server.py<br/>SID 查询服务"]
    SIDSRV --> SERVE
    SERVE --> RTR["retrieve_then_rank.py<br/>召回再排序的客户端"]

数据

生产数据格式

• 生产训练从 Kafka 流式读取 Arrow 格式的用户会话快照。开源版的 aggregated_kafka 读取器从目录读取同格式的 parquet（phoenix/reference/README.md:100-104）。
• 读取器类为 PhoenixDataset（phoenix/xrex/data/parquet_recsys.py）。配置里的路径占位符为 /path/to/offline_kafka_dump（phoenix/xrex/configs/xrecsys_two_tower.py:232）。
• 目录结构为 partition={p}/<bucket>/batch_{b}.parquet，另有一个 .valid_batches.json 清单（phoenix/reference/dump_gen.py:52、:377）。
• 一行包含：用户 ID 与画像；历史序列（帖子、作者、动作多热向量、停留时长等连续值、产品面、时间）；候选序列（字段相同，动作作为标签）；语义 ID 列 semanticIdSeq（phoenix/reference/README.md:100-109、:132-136）。

合成世界

world.py 一次性生成所有生成器共用的世界（generate_world，phoenix/reference/world.py:455-584）。开源默认值为 num_users=10_000、num_authors=2_000、num_posts=50_000、num_topics=16，种子为 20260721。

互动模型如下（world.py:301-340）：

• 亲和度 = 用户对帖子话题的偏好 × 作者质量 × 帖子质量（:301-305）。
• 门控：1 / (1 + exp(-14 × (亲和度 - 0.10)))（:307-308）。
• 各动作概率 = 基础率 × 门控。基础率为：点赞 0.14、点击 0.10、视频观看 0.08、转发 0.045、回复 0.035、收藏 0.02、引用 0.015、分享 0.012（:53-62、:310-311）。
• 停留达标：有互动且停留 ≥ DWELL_MARK_SECONDS = 10.0 秒。没有互动的记为“未停留”（:72、:337-338）。
• 帖子年龄在 1 到 36 小时之间（:36-37）。

多模态嵌入与语义 ID

生产环境有两个上游服务：多模态嵌入服务（“v5”）把帖子编码成 1024 维向量；语义 ID 服务把向量量化成 6 层 × 256 码的语义 ID（phoenix/reference/README.md:110-140）。开源版两条路径都提供：

world_snapshots.py 默认用合成向量训练一个 RQ-KMeans 码本，再分配 SID（phoenix/reference/world_snapshots.py:74-110）。产物放在 PHOENIX_INDEX_BASE 目录下：mm_snapshot/post_mm_v5.parquet、sid_snapshot/post_sid_v5_256x6.parquet、sid_snapshot/codebook_v5_256x6.npz、post_creation_snapshots/post_creation_1day.parquet（phoenix/reference/README.md:64-89）。

注意：运行 dump_gen.py 时必须传 --sid，否则数据里没有 SID 列。此时 SID 配置读到的全是“无 SID”；训练不报错，但 SID 表得不到梯度（phoenix/reference/README.md:132-136）。

损失函数

10-08 与 10-13 两次改动分别对应上游提交 35650fb 与 e1a0375。

训练步

可读版本是 recsys_train_step（phoenix/reference/train_step.py:101-148）。生产训练器 RecsysTrainer（phoenix/xrex/train/trainer_recsys.py:453）执行相同的步骤，另加分片与流水线：

1. 查嵌入：把本批所有哈希（用户、历史帖、历史作者、候选帖、候选作者、IP）拼接起来，一次从统一嵌入表取出对应行（train_step.py:33-55）。
2. 去重：compress_token_ids 计算本批用到的唯一行和反向索引（:119-122）。
3. 求梯度：同时对稠密参数和查出的嵌入求梯度（:124-126）。
4. 稠密更新：梯度范数非有限（或超过可选阈值）时，整步不更新（:87-97、:128-133）。
5. 稀疏更新：按唯一行累加嵌入梯度，交给逐行 AdaGrad，只更新被用到的行（:58-84、:135-147）。

优化器设置如下（开源默认值）：

combined 召回配置还在 checkpoint 里保存稠密参数的指数滑动平均（衰减 0.99）。保存时，训练器用平均值替换当前值（xrecsys_two_tower.py:325-326、trainer_recsys.py:3535-3547）。

checkpoint 与恢复

• 路径：<out>/<config_name>/elapsed_samples_<n>/<run_id>/（phoenix/TRAINING.md:86-92）。
• 恢复：用同一个 --out 和更大的 --steps 重跑。训练器找到最新的 checkpoint，恢复模型、优化器状态和数据读取位置，并检查校验和（phoenix/QUICKSTART.md:60-80）。
• 召回 checkpoint 额外包含 post_embeddings：训练器在保存时用候选塔处理整个语料。多头配置按语料类型选择头；混合头（ld）作为额外的列拼在后面（trainer_recsys.py:4195-4250）。
• reference/repack_checkpoint.py 把 checkpoint 精简成可发布的形态：只保留推理用的张量，去掉优化器状态，重新计算校验和（phoenix/reference/README.md:55）。

Rust 服务层

Cargo 工作区包含 9 个 crate（phoenix/Cargo.toml:3-14）：

加载与服务的过程

flowchart LR
    CLI["launch_inference.py<br/>--service_type ranking 或 retrieval"] --> RUN["RankingModelRunner<br/>或 RetrievalModelRunner"]
    RUN -- "恢复 checkpoint" --> P["JAX 参数"]
    RUN -- "load_tensor_no_resharding" --> PE["post_embeddings 等大表"]
    RUN --> SRV["Rust RecsysPredictorServer<br/>或 RecsysRetrievalPredictorServer"]
    REQ["gRPC 请求"] --> SRV
    SRV -- "哈希、组装批次" --> RUN
    RUN -- "JAX 前向" --> SRV
    SRV --> RESP["gRPC 响应"]

1. launch_inference.py 按 --service_type 选择 RankingModelRunner 或 RetrievalModelRunner（phoenix/xrex/inference/launch_inference.py:273、model_runner.py:3990、:4540）。
2. 运行器恢复 checkpoint。召回的 post_embeddings 由 Rust 函数 load_tensor_no_resharding 直接读入内存，再分发到各设备（model_runner.py:2313-2330，Rust 实现在 crates/serving/xai-recsys-engine/src/emb_table.rs:1095）。
3. 运行器构造 Rust 服务对象 RecsysPredictorServer / RecsysRetrievalPredictorServer（model_runner.py:4219-4245、:5355-5370）。配置了 sid_endpoint 时，服务在请求路径上调用 SID 查询服务，为帖子补齐语义 ID（model_runner.py:4228-4237，phoenix/README.md:123-125）。
4. Rust 侧接收 gRPC 请求，把请求排队并批量化，把编号哈希成表行号，填好输入数组（排序用 RankingBatchPrep，crates/serving/xai-recsys-engine/src/python.rs:1393）。然后 Rust 侧把批次交给 Python 执行前向，再返回响应。
5. ReloadModel RPC 支持在不重启的情况下更换模型（python.rs:1167-1200、recsys.proto:15）。

home-mixer 怎样调用这两个服务，见 Phoenix 排序模型 与 Phoenix 召回。

本地概念验证训练

能否运行

可以运行，但有以下硬件与软件前提：

• Linux + NVIDIA GPU。依赖固定为 jax[cuda12]==0.8.1（phoenix/pyproject.toml:12）。QUICKSTART.md 写的是“CUDA 12”；README.md 的完整验证环境是 nvidia/cuda:13.2.0-base-ubuntu22.04、GB300、aarch64（phoenix/QUICKSTART.md:12、phoenix/README.md:295-303）。
• 代码识别 A100 / H100 / H200 / GB200 / GB300。注意力核针对 Hopper 与 Blackwell 调优；其他 GPU 可能需要修改 attn_impl 与块大小（phoenix/README.md:298-303）。
• 需要 uv、Python ≥ 3.11、Rust 工具链、protoc ≥ 3.15，以及 RDMA verbs 头文件等系统包（phoenix/README.md:305-340）。
• 没有 GPU 的机器：源码未提供 CPU 训练路径的验证说明。

步骤

前提：在 phoenix/ 目录下执行；已按 phoenix/README.md 安装系统包、uv、Rust 和 protoc。

1. 安装依赖并编译 Rust 引擎（约 1 分钟）：

   uv sync --extra engine
   export PYTHONPATH=$PWD

2. 用随机权重自检。预期日志出现 Model warm up finished（phoenix/QUICKSTART.md:26-30）：

   uv run python xrex/inference/oss_bench/bench.py --smoke --service_type ranking

3. 生成合成快照和训练数据。默认生成 12,288 行、4 个分区、每个文件 1024 行（phoenix/reference/dump_gen.py:687-691）：

   uv run python reference/world_snapshots.py --out ./synth_index --seed 20260721
   export PHOENIX_INDEX_BASE=./synth_index
   uv run python reference/dump_gen.py --out ./synth_dump --seed 20260721 \
     --num-rows 12288 --partitions 4 --rows-per-file 1024 \
     --sid ./synth_index/sid_snapshot/post_sid_v5_256x6.parquet --self-check

4. 训练 nano 排序模型 6 步，并把每步损失记录到 checkpoints/run/metrics.jsonl：

   uv run python reference/train_synth.py --data ./synth_dump --steps 6 --out "$PWD/checkpoints" --metrics

5. 训练 nano 召回模型 6 步。产出的 checkpoint 带有候选索引：

   uv run python reference/train_synth.py --config xrecsys_two_tower_nano_offline_kafka_dump \
     --data ./synth_dump --steps 6 --out "$PWD/checkpoints"

6. 按 QUICKSTART.md 第 5 节启动 SID 服务（端口 50061）、召回服务（9990）和排序服务（9988）。等待两个模型服务都输出 Server ready to serve（phoenix/QUICKSTART.md:106-151）。
7. 发送三次“召回 → 排序”请求。成功时，最后一行是 retrieve_then_rank: 3 session(s) completed the full loop.（phoenix/QUICKSTART.md:153-167）：

   uv run python reference/retrieve_then_rank.py --data ./synth_dump --sessions 3 --topk 16 \
     --retrieval-port 9990 --ranking-port 9988

启动器的行为

train_synth.py（phoenix/reference/train_synth.py）的行为如下：

• 默认配置为 home_direct_packed_nano_offline_kafka_dump，默认 500 步（:15-16）。
• 只用 1 张卡（:110），并把配置中的数据路径指向 --data。
• 步数上限 = 数据行数 ÷ 每卡批大小 − 分区数。nano 排序每卡批大小为 64；按默认数据量，最多约 188 步，超过时自动截断（:126-139）。要训练更多步，用 dump_gen.py --num-rows 生成更多数据。
• 每 steps // 2 步保存一次 checkpoint，结束时再保存一次；评估关闭（:141-146）。

这次训练能证明什么

• 能证明：代码能在单卡上完整执行“生成数据 → 训练 → 保存 checkpoint → 恢复 → 服务 → 召回再排序”，且损失在合成数据上下降（phoenix/TRAINING.md:83-84）。
• 不能证明：模型质量、线上效果、吞吐或规模。上游原话是“不是质量或收敛性的声明”（phoenix/TRAINING.md:83-84、:102-112）。
• 也无法看到线上参数：上游没有发布训练好的 checkpoint（phoenix/TRAINING.md:96-100）。

设计决策

源码锚点

• phoenix/TRAINING.md、phoenix/QUICKSTART.md、phoenix/reference/README.md —— 上游训练与快速上手说明
• phoenix/reference/world.py:301-340、:455-584 —— 合成世界与互动模型
• phoenix/reference/dump_gen.py:683-737 —— 训练数据生成入口
• phoenix/reference/train_step.py:101-148 —— 可读训练步
• phoenix/reference/train_synth.py:94-160 —— 单卡启动器
• phoenix/xrex/train/trainer_recsys.py:453、:3535-3560、:4195-4250 —— 训练器、保存、召回索引
• phoenix/xrex/inference/model_runner.py:3990-4245、:4540-5370 —— 排序与召回服务运行器
• phoenix/crates/serving/xai-recsys-engine/src/python.rs:1393、:1167 —— 批次组装、模型热加载

相关页面

• Phoenix 排序模型 —— 排序模型结构与损失
• Phoenix 召回：双塔检索架构 —— 召回模型结构与损失
• Phoenix 召回索引：rankall 与事件层 —— 生产环境中召回语料快照的来源
• 哈希嵌入 —— 统一嵌入表与稀疏更新
• 端到端推理与训练入口 —— 旧 run_pipeline.py 的去向与现在的入口
• 多模态帖子嵌入器 —— 帖子多模态嵌入