Internal Wiki

主题

RFC-0030: 统一 ling3-flash 任务提交方案（正交配置） ​

概述 ​

用「正交 YAML 配置 + 单一启动脚本 + 单一 JobSet 模板」取代当前 9 个近乎重复的 submit_ling3_flash_*chips*[_docker][_antdata]_job.sh 提交脚本与 13 个对应的 JobSet 模板，把任务提交方式固化下来。配置由一个模型 base 叠加几条互不相交、 彼此正交的 overlay 轴构成：硬件 × 上下文长度（卡数与 8k/32k/256k 共同决定 拓扑与并行度切分）、生产 vs 测试环境、算法实验覆盖。

背景与动机 ​

当前提交一个 ling3-flash 训练任务，需要从 9 个几乎一模一样的提交脚本 （每个 210–260 行）里挑一个，配上 .github/ci/ 下 13 个 JobSet 模板 （每个约 260 行）中的一个。问题：

  • ~90% 是重复样板。脚本之间真正不同的只有少数几个值：拓扑、并行度、 生产/测试开关、数据集。
  • 改一处要同步改一堆。新增一个环境变量、修一个 bug，都得在多个脚本和模板里 同步修改，极易漏改。
  • envsubst 白名单靠手维护。每个脚本里有约 30 行手写的 SUBST_VARS 白名单，随着字段增减很容易和模板漂移。

差异其实非常有限，且能干净地归到几个桶里：

差异维度具体内容
按卡数/拓扑gke-tpu-topology、completions/parallelism、ICI_DATA_PARALLELISM、GRADIENT_ACCUMULATION_STEPS（维持 global batch = 8192）、LAZY_LOADER_SCATTER
生产 vs 测试STEPS、EVAL_INTERVAL、checkpoint/goodput/cloud-monitoring/vertex-TB 开关、maxRestarts、GCS bucket、namespace、HF 转换 watcher
按实验数据集（DATASET_TYPE + DATASETS_YAML）、学习率/warmup、序列长度
正交开关docker（预装 venv 镜像）vs 源码 clone；antdata（lazy）vs opensource（grain）

上表是现状的差异盘点。新方案里它们分别归位：前两行 → topology 轴（含上下文）/ environment 轴；「按实验」的数据集与 LR → 数据集 YAML + experiment 文件；「正交 开关」不再是自由维度——docker/源码由 --env+--image/--branch 决定，antdata/grain 由所选数据集（dataset_type）决定。

唯一真正的条件分支（而非取值替换）是 Pod 内的命令块：源码模式 （git clone + uv venv + 安装依赖）vs docker 模式 （source /opt/maxtext_venv/bin/activate）。envsubst 无法分支，这点决定了模板 需要做的统一处理。

目标 ​

  • 消除重复：把 9 个提交脚本 + 13 个 JobSet 模板收敛为单一 launch.sh + 单一 jobset.yml，差异全部外移到分层 YAML overlay。
  • 配置正交：硬件×上下文、环境、实验三条轴互不相交，新增一个组合只需加一个 overlay 文件，不再复制整脚本。
  • 删冗余：删掉旧模板里 ~65 行逐键搬运的 env: 块；把 pretrain_*.sh 的 ~400 行 env 默认值/参数拼接收敛为只剩少量条件推导（D1-D5）的 run_pretrain.sh。
  • 零新增依赖：沿用现有 bash + yq + envsubst 工具链，不引入新的运行时框架。
  • 可验收：对任一 chips/context/env/experiment 组合，新方案渲染出的 jobset 与 run.yml 应与旧脚本逐键等价（含 D1-D5 推导结果），且 64 卡 test 实机冒烟能 自洽开跑、训练日志推进。

非目标 ​

  • 不引入新的模板引擎/渲染框架（Jinja2 等），见「备选方案」。
  • 本轮只覆盖 ling3-flash；ling2 / ling3-tiny 等其他模型不在范围内（目录 结构已为其预留 models/<name>/，后续复用）。
  • 不强行迁移现有 scripts/datasets/ant_datasets_*.yml（17 个）：仍可被路径 引用，照常工作，后续按需逐步迁往外部 GitOps repo。
  • 不改动诊断/CI 专用 YAML 与 scripts/comm/*（详见「影响范围」）。
  • 不改动 maxtext 训练逻辑本身：只重构「如何把配置喂给 train.py」，不动 train.py 内部行为。

方案 ​

沿用现有 bash + yq + envsubst 工具链（零新增依赖）。用「正交 YAML 配置 + 单一启动脚本 + 单一 JobSet 模板」替换 9 个脚本及其模板。配置用 deploy:/maxtext: 两个命名空间区分去向：deploy: 经 envsubst 进 JobSet 模板，maxtext: 汇成单一 run.yml（声明 base_config: base.yml）挂进 Pod 直接交给 train.py。借此删掉旧 模板的 env: 块，并把 pretrain_*.sh 收敛为只剩条件推导的 run_pretrain.sh。

本次范围仅限 ling3-flash；目录结构按模型隔离，后续 ling2 / ling3-tiny 只需新增 models/<name>/ 目录即可复用。

用户故事 ​

提交一次训练，用户只需选「跑多大、多长上下文、什么环境、哪个实验」：

bash
# 例 1：算法同学在 1024 卡 v5p、8k 上下文、生产环境跑自己的实验
scripts/launch/launch.sh \
    --model ling3-flash --chips 1024 --context 8k --env prod \
    --experiment ~/my-exps/ling3-flash/lr3e4-fineweb.yml \
    --image main-abc1234-20260609          # prod 必填

# 例 2：test 环境源码模式，指定分支
scripts/launch/launch.sh \
    --model ling3-flash --chips 64 --context 8k --env test \
    --branch feature/new-lr-schedule       # 不给则用当前仓库分支

launch.sh 内部三步，全程可 --dry-run 预览：

text
第 1 步 合并配置              第 2 步 按命名空间劈两路              第 3 步 提交
─────────────────          ─────────────────────────          ────────────
model.yml                ┐                    ┌ .deploy ─envsubst─▶ jobset.yml ┐
topology/1024chips/8k.yml│                    │                                │
environment/prod.yml     ├─ yq deep-merge ──▶─┤                                ├▶ kubectl apply
~/my-exps/.../lr3e4.yml* │  (后者覆盖前者)     │                                │
KEY=VALUE (CLI 覆盖)     ┘                    └ .maxtext ──▶ run.yml ─ConfigMap─┘
                                                (base_config: base.yml)
                                                     ↓
                                              /etc/maxtext/run.yml → train.py

launch.sh 还会：推导 JOB_NAME / IMAGE / BRANCH / ConfigMap 名、并行度乘积自检

合并是「后者覆盖前者」：平台预设打底，用户的 experiment 文件（★，可来自另一个 repo）覆盖其上任意键，命令行 KEY=VALUE 最后兜底。deploy: 走 envsubst 进模板， maxtext: 汇成单一 run.yml 交给 train.py（详见「设计细节」）。

--dry-run 输出示例（节选，便于核对每个键的最终来源与去向）：

yaml
# === deploy （→ envsubst 进 jobset.yml）===
TPU_TOPOLOGY: 8x8x16            # ← topology/1024chips/8k.yml
COMPLETIONS: 256               # ← topology/1024chips/8k.yml
K8S_NAMESPACE: ling3-prod      # ← environment/prod.yml
DATASETS_YAML: <外部 repo>/datasets/fin_35p.yml  # 选择器：选哪个数据集并入 run.yml
# === maxtext （→ run.yml，train.py 的唯一位置 YAML）===
base_config: base.yml
model_name: ling3-flash        # ← --model 注入
ici_data_parallelism: 16       # ← topology/1024chips/8k.yml
max_target_length: 8192        # ← 数据集 YAML 并入
lazy_train_files: 0.66 /models_prod/...  # ← 数据集 YAML 并入
enable_checkpointing: true     # ← environment/prod.yml
learning_rate: 3.0e-4          # ← 实验文件 (覆盖)
# === rendered jobset.yml ===
metadata: { name: ling3-fl-1024-..., namespace: ling3-prod }
...   # completions: 256, gke-tpu-topology: 8x8x16, image: <指定镜像>,
      # 仅挂载 /etc/maxtext/run.yml（单个 ConfigMap）；无 65 行 env: 块

备选方案 ​

  1. Python/Jinja 渲染器 — 用 Jinja2 模板原生支持条件分支（docker/源码/多 slice 一个模板搞定），表达力更强。但引入新依赖与新的维护面，偏离现有 bash/envsubst 习惯。未采用：现有 yq + envsubst 已能覆盖需求，零新增依赖 更易被团队接受。
  2. 把生产/测试折叠进 topology 层（沿用「小卡=测试、大卡=生产」惯例，不单独 抽 environment 层）。文件更少，但 64 卡跑生产、大卡跑测试时不灵活。未采用： 用户明确要求把生产/测试作为独立维度。
  3. 环境轴命名用 profile — 已否决，与本仓库 profiling（性能采集）概念冲突， 改用 environment。
  4. context 单列为第四条正交轴 — 已否决。上下文长度会改变最优并行度切分（长 上下文用 CP 顶替 DP），与卡数耦合，故并入 topology 轴做成 <chips>/<context> 二维组合，而非与 topology 平行的独立轴。
  5. DP/GAS 由 launch.sh 纯算术推导 — 已否决。最优切分常需实测手调，写死在每个 <chips>/<context>.yml 里更可控；launch.sh 只做 乘积==设备数 的自检兜底。

约束与注意事项 ​

方案不引入新的运行时框架，但依赖一组命令行工具与集群资源。分三处列明。

A. 提交端（本地 MacBook 即可） ​

launch.sh 只做 YAML 合并、模板渲染与 kubectl apply，不导入 JAX、不碰 TPU/CUDA，因此在本地 MacBook（或任意配好 kubectl 的开发机 / CI runner）直接运行 即可——与「训练/测试必须在 TPU Pod 上跑」是两回事，提交这一步留在本地。所需工具都能 brew install：

工具用途备注
bash ≥ 4脚本宿主macOS 自带 3.2 偏旧，建议 brew install bash
yq（mikefarah/yq，Go 版，≥ v4）多层 YAML 深合并、.deploy/.maxtext 取子树、生成 run.yml见下方「yq 版本陷阱」——本方案统一锁定 mikefarah/yq，不再兼容 kislyuk/yq
envsubst（gettext）用 deploy 变量渲染 jobset 模板brew install gettext / apt install gettext-base
kubectlapply JobSet、建/查 ConfigMap、查 Pod需预先配好目标 GKE 集群凭证
git--branch 未指定时推导当前分支（rev-parse）；源码模式在 Pod 内克隆
sha256sum 或 shasumConfigMap 内容哈希（数据集 + run.yml）二者择一，_lib/datasets_configmap.sh 已兼容
gcloud（可选）下载产物、prod 镜像仓库鉴权仅 prod / 取数据时需要
python3（可选）复用旧脚本里 warmup_steps_fraction 的浮点计算（也可改用纯 bash/awk）见 run_pretrain.sh

yq 版本陷阱：现有 pretrain_ling3_flash.sh:86-91 为兼容两种 yq 写了分支 （mikefarah：yq eval-all 'select(fileIndex==0)*select(fileIndex==1)'； kislyuk：yq -y -s '.[0]*.[1]'）。本方案的合并逻辑更重（4 层 + 命名空间取子树）， 为避免分支维护与行为差异，统一要求 mikefarah/yq ≥ v4，并在 launch.sh 启动时 检测版本、不符即报错给出安装指引。

B. 集群前置（GKE，一次性 / 由平台准备） ​

  • GKE 集群 + TPU 节点池：tpu7x，且存在与各 --chips/--context 对应的拓扑 节点池（4x4x4 … 16x8x16）。
  • JobSet controller：集群已安装 jobset.x-k8s.io/v1alpha2 CRD。
  • GCS Fuse CSI driver：模板用 gke-gcsfuse/volumes 注解挂载。
  • PVC：ant-data-pvc、ant-data-prod-pvc（数据集与 checkpoint 挂载点）。
  • 命名空间：ling3-prod（prod）、default（test）。
  • Secret perf-64chips-token：源码模式克隆私有仓库需要；docker 模式标记 optional: true，缺失也能启动。
  • HF 转换 watcher 资源（仅 prod 且 ENABLE_HF_CONVERSION=true）： _lib/eval_conv_gate.sh 所需的 watcher Deployment/权限。

C. Pod 端（容器内，由模板的引导块准备，无需手动装） ​

  • docker 模式：镜像已预装 venv（/opt/maxtext_venv）、仓库、全部依赖，零安装。
  • 源码模式（--env test 不带 --image，通过 --branch 指定分支或默认当前分支）：引导块在 Pod 内现装 uv → venv（python 3.12）→ uv pip install -e '.[tpu]' → install_maxtext_tpu_github_deps → jq + mikefarah/yq（run_pretrain.sh 读 run.yml 用，须与提交端同款）→ 从源码装 optax（Muon 需要）。docker 镜像同样 需预装 mikefarah/yq。这段大体沿用现有模板逻辑，仅把 yq 来源统一为 mikefarah 版。

设计细节 ​

配置结构：base + 2 条仓库内正交预设轴 + 1 个用户覆盖文件 + CLI ​

配置由四部分叠加而成，优先级从低到高。前三部分（base + 两条预设轴）由本仓库 维护、是经过验证的「平台预设」；第四部分（experiment）是用户自己填写的覆盖 文件，优先级最高（仅次于 CLI 临时覆盖），可以覆盖前面任意键，并且可以放在另一个 repo 里做 GitOps 管理。

每个配置文件用两个顶层命名空间区分一个键最终流向哪里（详见后文「配置如何流向 train.py」）：

yaml
# 任意一层文件的形态
deploy:                     # → envsubst 进 JobSet 模板的编排键（约 10 个）
  TPU_TOPOLOGY: 8x8x16
  COMPLETIONS: 256
  K8S_NAMESPACE: ling3-prod
maxtext:                    # → 合并成 run.yml 交给 train.py 的 maxtext 配置键（其余几十个）
  per_device_batch_size: 2
  ici_data_parallelism: 16
  learning_rate: 3e-4

maxtext: 下用 maxtext 原生小写配置键（与 base.yml 同名），不再是 ICI_DATA_PARALLELISM 这种 env 变量名——这样它们能直接进 run.yml，无需中间映射。

数据集 YAML 也是 maxtext 配置：数据集 YAML 的内容本身就是一批 maxtext 键 （lazy_train_files、max_target_length、eod_mask_loss 等，实测每个 ≤135 KB）。 因此它作为合并链的一员并入 run.yml（位置见下「合并优先级」第 ②.5 项），而非 单独挂一个文件——由此 消掉原本独立的 datasets ConfigMap，全部配置汇成单一 run.yml。哪个数据集 YAML 参与合并，由 deploy.DATASETS_YAML 指定（仅作 launch.sh 的「读取哪个文件」选择器，不进 run.yml）。

数据集 YAML 的两类来源（与 experiment 一样支持外置）：

  • 仓库内少量默认集：scripts/launch/models/ling3-flash/datasets/ 下保留几个 用于日常精度对齐的稳定数据集，随 maxtext 仓库版本化。
  • 外部 GitOps 数据集：生产/实验用的大量、频繁迭代的数据集，与 experiment 文件 一同住在另一个 GitOps repo。deploy.DATASETS_YAML 既接受仓库内名字，也接受 外部绝对/相对路径——与 --experiment 同样的「名字 or 路径」解析规则。
text
scripts/launch/
  launch.sh
  templates/
    jobset.yml                       # 唯一 JobSet 模板（原 13 个）
  models/
    ling3-flash/
      model.yml                      # ① base：模型默认值（平台维护）
      topology/                      # ② 预设轴 A：硬件 × 上下文（平台维护）
        64chips/{8k,32k,256k}.yml
        256chips/{8k,...}.yml
        1024chips/{8k,...}.yml
      environment/                   # ③ 预设轴 B：生产 vs 测试（平台维护，正交于 A）
        prod.yml
        test.yml
      datasets/                      # 仓库内少量默认数据集（精度对齐用，稳定）
        align_smoke.yml
      experiment/
        example.yml                  # 仅作模板/示例；真实实验文件由用户拥有，可外置

datasets/ 与 experiment/ 在仓库内都只放少量稳定样例；生产/实验用的大量数据集 与 experiment 文件随外部 GitOps repo 迭代，提交时由 deploy.DATASETS_YAML / --experiment 以路径形式引入。

合并优先级（yq deep-merge，后者覆盖前者，键碰撞时高优先级胜出）：

text
① model.yml
      ↓ 被覆盖
② topology/<chips>/<context>.yml          （由 --chips + --context 选中）
      ↓ 被覆盖
②.5 数据集 YAML（ant_datasets_*.yml）      （由 deploy.DATASETS_YAML 选中，整份并入 maxtext.）
      ↓ 被覆盖
③ environment/<env>.yml                    （由 --env 选中）
      ↓ 被覆盖
④ 用户 experiment 文件（可选）              （由 --experiment <path|name> 指定）★ 用户拥有
      ↓ 被覆盖
⑤ CLI 上的 KEY=VALUE                        （一次性临时覆盖，最高）

数据集 YAML 排在 topology 之后、environment 之前：它带的 max_target_length、 lazy_loader_scatter 等可被 environment / experiment / CLI 进一步覆盖。其顶层 键直接归入 maxtext. 命名空间（它们本就是 maxtext 配置键）。

①②③ 是「选预设」：--chips/--context/--env 是选择器，决定加载哪个预设 文件，文件内容由平台维护、不轻易改动。这三层本身即一套自洽可跑的完整配置—— 不给 --experiment 时就跑它们。 ④是「填覆盖」（可选）：experiment 文件是用户的实验声明，可覆盖以上任意键 （包括并行度、环境开关、模型旋钮），是日常调参的主战场；不给则跳过。 ⑤是「临时试」：命令行 KEY=VALUE 用于不落盘的一次性试验。

命名说明：轴 B 刻意叫 environment 而非 profile。本仓库里 profiling 是 一等概念（PROFILER、ENABLE_TPU_PROFILING_OPTIONS、PROFILER_STEPS 等一组 性能采集环境变量），再叫 profile 会与之语义混淆。

--experiment 可选，接受路径或名字，支持外置 GitOps：

  • 不给 → 跳过第④层，跑①②③三层预设的默认配置。
  • 传名字（如 --experiment dev）→ 解析为仓库内 experiment/dev.yml（仅用于 示例/快速试）。
  • 传路径（含 / 或以 .yml 结尾，如 --experiment ~/exp-configs/run42.yml 或 --experiment /gitops/ling3-flash/prod-0609.yml）→ 按原样读取。这样实验配置可以 住在另一个 repo 里，由该 repo 的 PR / CI 流程管理版本与审批，launch.sh 只在 提交时把它读进来合并。

为什么 context 折进 topology 轴而不是单列一轴：上下文长度不只是改 MAX_SEQ_LEN，它会改变最优并行度切分——8k 用 FSDP+DP，而 256k 往往要拿 context parallelism（ICI_CONTEXT_PARALLELISM）顶替 DP，同时联动 PER_DEVICE_BATCH_SIZE、remat、packing。并行度由 chips 和 context 共同决定， 且常需实测手调（不是纯算术能得出），所以把它们做成一个二维组合 <chips>/<context>， 每个组合一套经过验证的最优切分，比自动推导更诚实、更可控。

各部分职责（从现有脚本提炼）：

  • ① model.yml（base，平台维护）— 与规模、上下文都无关的共享旋钮：OPT_TYPE、 全部 MUON_*、REMAT_POLICY 默认、ICI_EXPERT/TENSOR_PARALLELISM、 SHARD_EXP_ON_FSDP、WARMUP_ITERS、monitor_* 与 checkpoint 周期的 extra-args 默认、LOAD_PARAMETERS_PATH（scan 检查点）、profiler 变量（空）、 MODEL_NAME=ling3-flash，以及默认数据集（dataset_type + 默认 DATASETS_YAML） ——使不带 --experiment 时三层预设即可自洽开跑。

  • ② topology/<chips>/<context>.yml（预设轴 A，平台维护）— 由「卡数 × 上下文」 共同决定的硬件与切分预设：TPU_TOPOLOGY、COMPLETIONS（= 主机数）、并行度 （ICI_FSDP/DATA/CONTEXT_PARALLELISM，乘积需 = 设备数）、 GRADIENT_ACCUMULATION_STEPS（凑目标 global batch）、PER_DEVICE_BATCH_SIZE （随 context 变）、LAZY_LOADER_SCATTER、MAX_SEQ_LEN、ROPE_MAX_TIMESCALE、 以及随 context 调整的 remat 旋钮。8k 参考值：

    chipstopologycompletionsFSDPDPCPGAS
    644x4x416128111
    2564x8x864128418
    5128x8x8128128814
    10248x8x162561281612
    204816x8x165121283211

    长上下文（如 256k）的同卡数文件会以 CP>1 顶替部分 DP，并下调 PDB；具体切分按 实测填入对应 <chips>/256k.yml。

  • ③ environment/{prod,test}.yml（预设轴 B，平台维护）— 生产 vs 测试：STEPS、 EVAL_INTERVAL、ENABLE_CHECKPOINTING、ASYNC_CHECKPOINTING、ENABLE_GOODPUT、 ENABLE_CLOUD_MONITORING、USE_VERTEX_TENSORBOARD、K8S_MAX_RESTARTS、 K8S_NAMESPACE、GCS_BUCKET、ENABLE_HF_CONVERSION。 （prod = 监控/检查点/HF-watcher 全开、ling3-prod 命名空间、ant-data-prod bucket、maxRestarts 10；test = 全关、default 命名空间、data_ant bucket、 maxRestarts 3、小 STEPS。）

  • ④ 用户 experiment 文件（用户拥有，最高优先级，可外置 GitOps）— 一次实验的 完整声明。典型内容是 maxtext: 下的算法旋钮：dataset_type（grain/lazy）、 learning_rate、warmup 相关、steps，以及 deploy: 下的 DATASETS_YAML （指向 scripts/datasets/ant_datasets_*.yml）。但不受限：用户可以覆盖①②③的 任意键——比如 maxtext.ici_data_parallelism 调小、maxtext.enable_checkpointing 打开。并行度被覆盖后仍会过 launch.sh 的 乘积==设备数 自检（见下），写错会被拦下。

    yaml
    # 一个 experiment 文件示例（可住在外部 GitOps repo）
maxtext:
  learning_rate: 3.0e-4
  dataset_type: lazy
  steps: 50000
deploy:
  DATASETS_YAML: scripts/datasets/ant_datasets_fin_35p_0512.yml

配置如何流向 train.py：双命名空间 + 单一 run.yml ​

这一节回答「合并后这么多配置项怎么传给 maxtext」，以及为什么旧模板里那 ~65 行 env: 块可以整体删除。

关键约束（已核对 pyconfig.py）：train.py 命令行只接受一个位置 YAML， 其余参数一律当作 key=value（pyconfig.py:82 仅取 argv[1]）。因此不能用 train.py base.yml run.yml 传两个文件。但 maxtext 的 OmegaConf 合并与本仓库现有 yq 合并都是深合并、后者/子配置覆盖前者（同名标量后者胜出；logical_axis_rules 按规则名合并，override_logical_axis_rules=true 时整体替换）——这正好支撑下面的方案。

做法：launch.sh 把合并后的扁平配置按命名空间劈成两路：

text
合并后的配置（含 deploy: 与 maxtext: 两个命名空间，
              maxtext: 已并入所选数据集 YAML 的全部键）
        │
        ├── yq '.deploy'  → 导出为 shell 变量 ──envsubst──▶ JobSet 模板
        │                    （~12 个结构性变量：JOB_NAME / IMAGE /
        │                     TPU_TOPOLOGY / COMPLETIONS / NAMESPACE …）
        │
        └── yq '.maxtext' → 写出 run.yml（顶部声明 base_config: base.yml）
                             → 单个内容哈希 ConfigMap → 挂载到 /etc/maxtext/run.yml
                             → train.py /etc/maxtext/run.yml

run.yml 顶部写 base_config: base.yml，于是 pyconfig.py:140-151 会先加载 base.yml 再把 run.yml 合并上去（run.yml 胜出），等价于「base + 本次实验全部覆盖」 一次到位。模型选择仍靠 model_name=ling3-flash（在 run.yml 里），其 configs/models/ling3-flash.yml 由 loader 自动加载。

收益：

  • 旧模板里那 ~65 行逐键搬运的 env: 块整体删除——maxtext 配置改由挂载的 run.yml 承载，不再经 Pod 环境变量中转。
  • Pod 脚本里 ~60 行 key=value 拼接收敛为「train.py /etc/maxtext/run.yml + 少量推导出的覆盖 + "$@"」。
  • run.yml 是一份可 dry-run 打印、可 kubectl exec 查看的「本次运行的完整配置」， 排障时一眼看到最终值。

与「推导留 Pod」的衔接（见下「Pod 入口脚本」）：Pod 脚本继任者只需读 run.yml 里 它要分支的那几个键（per_device_batch_size、scan_layers、steps、warmup 等）， 算出 D1-D5 推导，作为少量 key=value 追加在 train.py /etc/maxtext/run.yml 之后即可（CLI 覆盖优先级最高，正好覆盖 run.yml 里的占位）。

launch.sh（单一入口） ​

bash
scripts/launch/launch.sh --model ling3-flash --chips 1024 --context 8k --env prod \
  [--experiment <name|path.yml>] [--image <tag>] [--branch <branch>] [--name <name>] [KEY=VALUE ...]

--experiment 可选：不给就只用 base + topology + environment 三层预设跑默认 配置（这三层本身即一套自洽可跑的完整配置，数据集等默认值下沉在 base/environment 里）。它只在要换数据集、调 LR 等实验性覆盖时才需要。

运行模式由 --env 与 --image/--branch 共同决定（不再有独立的 docker/source 开关）：

--env--image--branch运行模式
prod必填忽略docker（预装 venv 镜像）。缺 --image 即报错——生产要可复现，不在 Pod 内现 build
test不给可选，默认当前仓库分支源码（git clone 指定分支 + uv venv 安装，python:3.12）
test给定忽略docker（用指定镜像测试）

--image 与 --branch 互斥：给了 --image 则进入 docker 模式，镜像内已含代码， --branch 被忽略（若同时给出，打印警告）。源码模式下 --branch 不给则默认取 当前仓库的分支（git rev-parse --abbrev-ref HEAD），方便本地开发时"提交当前 正在改的分支去跑"的场景。

职责：

  1. 解析配置文件：base model.yml + 预设 topology/<chips>/<context>.yml + environment/<env>.yml（由 --chips/--context/--env 选中）+ 可选的用户 experiment 文件。--experiment 既接受名字（解析为仓库内 experiment/<name>.yml） 也接受路径（含 / 或 .yml 结尾，按原样读取，可来自外部 GitOps repo）； 不给则跳过这一层，直接用三层预设的默认配置。被选中的预设文件缺失则清晰报错。
  2. 校验 env/image/branch 约束：--env prod 必须有 --image，否则报错退出。 --image 与 --branch 同时给出时打印警告（--branch 被忽略）。
  3. 用 yq 按 ①→②→②.5→③→④ 顺序深度合并（后者覆盖前者，experiment 可覆盖任意 预设键）：第 ②.5 步把 deploy.DATASETS_YAML 选中的数据集 YAML 整份并入 maxtext.。同时把 --model 注入 maxtext.model_name。再应用末尾的 KEY=VALUE 命令行覆盖（最高优先级），按是否为 maxtext 原生键决定落到 maxtext. 还是 deploy. 命名空间。
  4. 劈成两路（见「配置如何流向 train.py」）：
    • yq '.deploy' → export 为 shell 变量；由 deploy 键名程序化生成 envsubst 白名单（取代手维护的 SUBST_VARS）。
    • yq '.maxtext' → 写出 run.yml（顶部加 base_config: base.yml）→ 经 _lib/datasets_configmap.sh 同款内容哈希逻辑建单个 ConfigMap 挂进 Pod （数据集已并入 run.yml，不再有独立 datasets ConfigMap）。
  5. 推导少数 deploy 侧计算值：
    • BRANCH：给了 --branch 则用其值；否则取当前仓库分支 （git rev-parse --abbrev-ref HEAD）。BRANCH_LABEL 由 BRANCH 转化。
    • JOB_NAME（未给 --name 时用 branch+时间戳；转小写、≤63 字符）、USER。
    • IMAGE：给了 --image 时为镜像仓库 …/maxtext-tpu-pretrain:<tag>（docker 模式）；未给（仅 --env test 允许）时默认 python:3.12（源码模式）。
    • RUN_CONFIG_CONFIGMAP_NAME：run.yml 的内容哈希 ConfigMap 名。
    • 注：CUSTOM_REMAT_FLAGS（PDB→remat）等条件推导（D1-D5）留在 Pod 脚本，不在 launch.sh 计算（见下「Pod 入口脚本」）。
  6. 并行度自检：从 maxtext. 读 ici_*，断言 FSDP×DP×TP×EP×CP == 设备数， 不等则报错退出（防止 topology×context 组合写错切分）。
  7. prod 环境下，用现有 _lib/eval_conv_gate.sh （hf_validate_env / hf_pre_apply_check / hf_apply_watcher）执行 HF-watcher gate，由 ENABLE_HF_CONVERSION 控制。
  8. envsubst "$SUBST_VARS" < templates/jobset.yml | kubectl apply -f -。
  9. 打印与现有脚本一致的 monitor/tail/exec/cleanup 帮助块（带命名空间）。

单一 JobSet 模板 ​

把 13 个模板收敛为 scripts/launch/templates/jobset.yml。除了消除 docker-vs-源码的 结构差异，最大的简化是删掉原模板里那 ~65 行逐键搬运的 env: 块——maxtext 配置 改由挂载的 run.yml 承载，模板只保留少量编排用 env（GCS_BUCKET、RUN_NAME、 GITHUB_TOKEN 等）。

  • image: ${IMAGE}（docker 模式为指定镜像；源码模式 launch.sh 默认 python:3.12）。

  • Pod 内运行时引导探测（取代「一个 clone 模板 + 一个 docker 模板」）：

bash
if [ -f /opt/maxtext_venv/bin/activate ]; then
  source /opt/maxtext_venv/bin/activate          # docker：镜像内预装 venv + 仓库
else
  git config --global url."https://x-access-token:${GITHUB_TOKEN}@github.com/".insteadOf "https://github.com/"
  git clone --depth=1 --branch="${BRANCH}" https://github.com/primatrix/maxtext.git /root/maxtext
  cd /root/maxtext
  pip install uv && uv venv --python 3.12 --seed maxtext_venv && source maxtext_venv/bin/activate
  uv pip install -e '.[tpu]' --resolution=lowest && install_maxtext_tpu_github_deps
  apt-get update -qq && apt-get install -y -qq jq >/dev/null && pip install yq
  pip uninstall optax -y 2>/dev/null || true && pip install git+https://github.com/google-deepmind/optax
fi
cd /root/maxtext
# ... index 0 上 job_log_sync ...
bash scripts/launch/run_pretrain.sh /etc/maxtext/run.yml   # 新 Pod 入口（见下节）
  • 单个内容哈希 ConfigMap 挂进 Pod（源码/docker 都适用）：合并后的 run.yml （已并入数据集 YAML 的全部键）→ /etc/maxtext/run.yml，即 train.py 的唯一位置 YAML。不再有独立的 datasets ConfigMap / 挂载。GITHUB_TOKEN 的 secretKeyRef 标记 optional: true，缺 perf-64chips-token 的 docker 命名空间也能启动。

模板需暴露的结构性变量收敛为约 11 个： ${JOB_NAME} ${K8S_NAMESPACE} ${K8S_MAX_RESTARTS} ${IMAGE} ${COMPLETIONS} ${TPU_TOPOLOGY} ${RUN_CONFIG_CONFIGMAP_NAME} ${GCS_BUCKET} ${BRANCH} ${BRANCH_LABEL} ${USER}。仅支持单 slice。

Pod 入口脚本：run_pretrain.sh（承接条件推导） ​

按「推导留 Pod」的决定，新建 scripts/launch/run_pretrain.sh 接替 pretrain_ling3_flash.sh 的角色，但只保留无法静态表达的条件推导，不再有几十行 env 默认值与 key=value 拼接。它：

  1. 读入参 run.yml（已含本次实验的全部 maxtext 配置）。
  2. yq 从中读出它要分支的几个键：per_device_batch_size、scan_layers、steps、 warmup 相关、load_full_state_path、compiled_trainstep_file。
  3. 计算这几处条件推导（从旧 pretrain_ling3_flash.sh 逐字移植，编号本地重排）：
    • D1 remat_policy=custom 且 per_device_batch_size<=2 → 追加 activation-saver （mlpwi_0/1=device moe_topk_indices=device）。
    • D2 steps<=warmup → 放大 learning_rate_schedule_steps，再算 warmup_steps_fraction。
    • D3 load_parameters_path 按 scan_layers 真假二选一。
    • D4 load_full_state_path 非空 → 清空 params 路径并强制开 checkpoint。
    • D5 compiled_trainstep_file 非空（AOT）→ eval_interval=0。
  4. source comm/env_check.sh、comm/xla_flags_common.sh、comm/libtpu_diag_common.sh （沿用，负责 vm.max_map_count、LIBTPU_INIT_ARGS、诊断）。
  5. 调起 python3 -m maxtext.trainers.pre_train.train /etc/maxtext/run.yml <上面推导出的少量 key=value>——CLI 覆盖优先级最高，正好覆盖 run.yml 占位。

这样模型训练逻辑仍跟着仓库脚本走、kubectl exec 进去可手改 run.yml 重跑，同时把 「静态配置」与「条件推导」彻底分开：前者在 run.yml 一眼可见，后者集中在一个短脚本里。

测试计划 ​

验证策略是「先逐项核对渲染产物、再实机冒烟」，确保新方案与旧脚本行为等价后才 删除旧脚本。

  • 渲染校验（本地可做，无需 TPU）：给 launch.sh 加 --dry-run，打印 deploy/maxtext 两段 + 渲染后的 jobset.yml + run.yml。对每个 chips/context/env/experiment 组合确认：jobset 无残留 ${...}、 completions/拓扑/命名空间/bucket 正确、运行时变量（GITHUB_TOKEN、 JOB_COMPLETION_INDEX、LOG_FILE、RUN_NAME）未被替换；run.yml 含 base_config: base.yml 且并行度乘积 = 设备数。对 jobset 跑 kubectl apply --dry-run=client，对 run.yml 跑一次 train_compile 或 pyconfig 加载，校验其能被 maxtext 接受（兼测类型：env 走管道时是字符串，run.yml 里是 原生 YAML 类型）。
  • 与旧产物逐键 diff：渲染 64 卡 test 与 1024 卡 prod。jobset 与旧模板 diff 应只 剩有意的统一项（image 变量、引导 if-block、删 env: 块、挂 run.yml、optional secret）；run.yml 的最终 maxtext 配置应与旧 pretrain_ling3_flash.sh 实际拼出的 key=value 集合逐键等价（含 D1-D5 推导结果）。
  • 实机冒烟（64 卡 test，走默认配置不带 experiment）：确认三层预设即可自洽开跑、 Pod 进入 Running、训练日志推进若干 step；再补一次带 --experiment <path> 的提交， 验证覆盖生效。
  • Docker 生产路径冒烟：--chips 1024 --env prod --experiment dev --image <tag> dry-run +（有 Pod 时）短暂实机提交，确认预装 venv 激活与 HF-watcher gate 生效。
  • CI 断言：加一条「所有 maxtext 原生键不得出现在 deploy: 下」的断言，防止键 归错命名空间（难以靠单测覆盖的是「键漂移」，故用渲染期断言兜底）。

影响范围 ​

  • maxtext 仓库 scripts/ 与 .github/ci/ 目录。
  • 提交 ling3-flash 任务的所有人：命令从 scripts/submit_ling3_flash_<chips>chips[_docker][_antdata]_job.sh 改为统一的 scripts/launch/launch.sh --model ling3-flash --chips <N> ...。
  • 引用了被删脚本/模板的 CI workflow（迁移前需 grep .github/workflows/ 与 scripts/ 并更新调用方）。
  • 实验配置与数据集的归属转移：experiment 文件与生产/实验用数据集 YAML 均由用户 拥有、随外部 GitOps repo 迭代，可不再住本仓库。本仓库 experiment/ 仅留 example.yml 模板、datasets/ 仅留少量精度对齐用的默认集。--experiment <path> / deploy.DATASETS_YAML <path> 把外部文件读进来合并，其版本/审批走该 repo 的 PR 流程，与平台预设（model/topology/environment）解耦。

不在本次范围（保持不动）：

  • 诊断/CI 专用 YAML：tpu-1x64chips-ling3-flash-dump-job.yaml、 tpu-2x64chips-ling3-flash-multislice-job.yaml、 tpu-loss-validation-ling3-flash-64chips-jobset.yaml、 tpu-long-running-job-64chips-ling3-flash.yaml（被 CI workflow 直接引用 / 一次性用途）。
  • scripts/comm/*（env_check.sh / xla_flags_common.sh / libtpu_diag_common.sh， 由新 Pod 入口 run_pretrain.sh 继续 source）。
  • 现有 scripts/datasets/ant_datasets_*.yml（17 个）本轮不强行迁移：仍可被 deploy.DATASETS_YAML 以路径引用，照常工作；后续按需逐步迁往外部 GitOps repo， 仓库内只长期保留少量精度对齐用的默认集。

被取代：scripts/pretrain_ling3_flash.sh（~400 行 env 默认值 + key=value 拼接） 由 scripts/launch/run_pretrain.sh（仅 D1-D5 条件推导 + 调 train.py）取代；其静态 配置职责上移到分层 YAML / run.yml。验证通过后删除。

实施计划 ​

  1. 搭骨架：新建 scripts/launch/ 目录、launch.sh、run_pretrain.sh、 templates/jobset.yml，以及 models/ling3-flash/ 下的 model（base）/ topology（<chips>/<context>）/ environment / experiment overlay 文件（均按 deploy:/maxtext: 双命名空间组织）。
  2. 移植逻辑：把 JOB_NAME、ConfigMap、HF-watcher gate、env/image 约束、 并行度自检移植到 launch.sh；把条件推导（D1-D5）逐字移植到 run_pretrain.sh；复用现有 _lib/* 与 comm/*。
  3. 渲染校验与逐键 diff：按「测试计划」对各组合做 --dry-run 校验，并与旧产物 逐键 diff（64 卡 test、1024 卡 prod）。
  4. 实机冒烟：先跑 64 卡 test 默认配置 + 带 --experiment 覆盖各一次；再走 1024 卡 prod docker 路径，确认预装 venv 激活与 HF-watcher gate 生效。
  5. 删除与迁移：以上通过后，删除 9 个 submit_ling3_flash_*chips*_job.sh 及其 对应的 9 个 .github/ci/tpu-*chips-ling3-flash-*-job.yaml，在 CLAUDE.md 与新增 scripts/launch/README.md 中给出新旧命令映射表。

风险与缓解 ​

  • 渲染产物与旧脚本行为不一致 → 通过「测试计划」的逐项 diff + 实机冒烟把关， 先验证再删除旧脚本。
  • envsubst 误替换运行时变量 → 白名单由合并键程序化生成，运行时变量天然不在 键集合内；dry-run 中专门断言 GITHUB_TOKEN 等保持未替换。
  • docker 命名空间缺 perf-64chips-token → secret 标记 optional: true， 缺失也能启动（docker 模式不需要 clone）。
  • CI workflow 仍引用被删文件 → 删除前 grep .github/workflows/ 与 scripts/， 更新所有调用方。
  • 统一 ConfigMap 改变源码模式行为（旧 64 卡源码脚本直接用 DATASETS_YAML 路径） → 内容哈希 ConfigMap 去重，行为等价；在逐键 diff 中确认挂载路径与内容一致。
  • run.yml 漏键 / 类型不符导致 train.py 行为漂移（env 走管道时是字符串，run.yml 里是原生 YAML 类型） → 「逐键等价」diff 专门比对旧 key=value 集合与 run.yml 合并后的最终值；用 pyconfig 加载校验类型；保留 model_name + base_config: base.yml 确保 configs/models/ling3-flash.yml 仍自动加载。
  • deploy:/maxtext: 键归错命名空间（如把 ici_* 误放 deploy，导致并行度没进 run.yml） → launch.sh 的并行度自检从 maxtext. 读取，缺键即报错；CI 加一条 断言：所有 maxtext 原生键不得出现在 deploy: 下。
  • 条件推导散落两处难追溯（launch.sh 不再算 remat，全部在 run_pretrain.sh） → D1-D5 集中在一个 ~50 行脚本、各带注释与来源行号，dry-run 也打印推导后的覆盖。