在百舸上使用 lingbot-va 加速镜像训练视觉-动作模型

概述

lingbot-va 是基于 Wan2.2（TI2V-5B 系）的视觉-动作（Vision-Action，VA）模型，在 fixed176（LIBERO pick-n-place 系）数据集上做监督微调（SFT）训练。

本镜像面向多机多卡训练场景进行了针对性优化，通过 FSDP 重分片（resharding）、block 级分片粒度、HSDP 二维 mesh 与数据加载调优，消除跨节点 FSDP AllGather 瓶颈，在单机 8 卡下训练吞吐较基线提升约 24.9%，双机 16 卡时吞吐几乎线性翻倍（多机扩展效率约 98.9%）。

适用场景

• 基于 Wan2.2 系列基座模型进行视觉-动作模型微调训练。
• 需要多机多卡分布式训练、对训练扩展效率有要求的场景。
• 使用 fixed176（LIBERO pick-n-place 系）等标准数据集开箱训练，或替换为自有数据集。

环境与资源要求

镜像内置软件栈：

准备工作

创建训练任务

1. 登录百舸平台，进入训练任务创建页。
2. 选择推荐 GPU 资源规格（单机 8 卡或双机 16 卡）。
3. 工作目录设置为 /root/workspace/lingbot-va。
4. 启动命令配置见下文「启动训练」。

准备模型权重

训练权重为 lingbot-va-base（Wan2.2 系），从百舸 BOS 同步到工作目录：

1bcecmd bos sync "bos:/aihc-models-bj/robbyant/lingbot-va-base/" \
2  "/root/workspace/models/lingbot-va-base/" --concurrency 16

路径可改：train.sh 中 BASE_MODEL_PATH 与 DATASET_PATH 均为 ${VAR:-默认值} 形式，权重与数据集可以放到任意路径（含挂载盘），只需在启动前覆盖对应环境变量，例如：

1export BASE_MODEL_PATH=/root/workspace/models/lingbot-va-base
2export DATASET_PATH=/root/workspace/datasets/pick-n-place-sq-lerobot-v21-fixed176
3bcecmd bos sync "bos:/aihc-models-bj/robbyant/lingbot-va-base/" "${BASE_MODEL_PATH}/" --concurrency 16
4bash /root/workspace/lingbot-va/train.sh

准备数据集

数据集 fixed176（pick-n-place-sq-lerobot-v21-fixed176，约 794M）已随镜像内置到 /root/workspace/datasets/pick-n-place-sq-lerobot-v21-fixed176，无需额外下载。

如需换用自有数据集，通过 DATASET_PATH 环境变量指向实际路径即可，内置副本不会被读取。

启动训练

train.sh 已内置全部优化配置与默认参数，直接执行即可：

1# 单机 8 卡
2bash /root/workspace/lingbot-va/train.sh
3
4# 双机 16 卡：两节点各执行一次，拓扑由平台注入环境变量决定
5bash /root/workspace/lingbot-va/train.sh

单机与多机使用同一套脚本，集群拓扑由平台注入的环境变量决定（NPROC_PER_NODE / WORLD_SIZE=节点数 / RANK=节点序号 / MASTER_ADDR / MASTER_PORT）；单机时 WORLD_SIZE 缺省为 1。

启动脚本 train.sh 完整内容如下（可直接复制为训练任务启动命令）：

1#!/bin/bash
2set -o pipefail
3
4CODE_DIR=${CODE_DIR:-"/root/workspace/lingbot-va"}
5DATASET_PATH=${DATASET_PATH:-"/root/workspace/datasets/pick-n-place-sq-lerobot-v21-fixed176"}
6BASE_MODEL_PATH=${BASE_MODEL_PATH:-"/root/workspace/models/lingbot-va-base"}
7LOG_DIR=${LOG_DIR:-"/root/workspace/logs"}
8RUN_TAG=${RUN_TAG:-"opt_66edf60"}
9
10GPUS_PER_NODE=${NPROC_PER_NODE:-8}
11NNODES=${WORLD_SIZE:-1}
12NODE_RANK=${RANK:-0}
13MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
14MASTER_PORT=${MASTER_PORT:-"29700"}
15
16export LINGBOT_FSDP_GRANULARITY=block
17export LINGBOT_FSDP_RESHARD_AFTER_FORWARD=0
18if (( NNODES > 1 )); then
19  export LINGBOT_FSDP_SHARDING=hsdp
20  export LOCAL_WORLD_SIZE="${GPUS_PER_NODE}"
21fi
22export PYTHONPATH="${CODE_DIR}:${PYTHONPATH:-}"
23export WANDB_DISABLED=true
24export WANDB_MODE=disabled
25export HF_HUB_OFFLINE=1
26export TRANSFORMERS_OFFLINE=1
27export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
28export OMP_NUM_THREADS=8
29export DATASET_PATH BASE_MODEL_PATH
30
31LOG_FILE="${LOG_DIR}/${RUN_TAG}_$(hostname).log"
32mkdir -p "${LOG_DIR}"
33cd "${CODE_DIR}" || exit 1
34
35echo "[lingbot-va][topology] nnodes=${NNODES} node_rank=${NODE_RANK} gpus_per_node=${GPUS_PER_NODE} master=${MASTER_ADDR}:${MASTER_PORT} sharding=${LINGBOT_FSDP_SHARDING:-fsdp}"
36
37python -m torch.distributed.run \
38  --nnodes="${NNODES}" \
39  --node_rank="${NODE_RANK}" \
40  --master_addr="${MASTER_ADDR}" \
41  --master_port="${MASTER_PORT}" \
42  --nproc_per_node="${GPUS_PER_NODE}" \
43  -m wan_va.train --config-name=libero_train \
44  2>&1 | tee "${LOG_FILE}"

训练参数说明

所有训练参数均通过环境变量传入，不支持命令行传参（如 save_interval=20 会被 argparse 直接拒绝并退出）。

性能参考

以下数据为特定测试环境下的参考值（8 卡单机 / 16 卡双机，num_steps=50，batch_size=1、gradient_accumulation_steps=10）。实际性能受资源规格、数据读取方式、网络环境、任务参数等因素影响。

• 加速比（单机 8 卡，优化版 vs 基线）：1.249x（+19.9%）。
• 多机扩展效率（双机 vs 单机，均值窗口 step 10-49）：98.9%，跨机扩展近乎无损。
• 训练中 loss 正常收敛。

梯度累积步数（GA）标定参考

由于 batch_size 恒为 1，训练批次仅能通过 gradient_accumulation_steps 调节。6 个档位实测流量（参考值）：

注：GA=1 与 GA≥2 之间存在约 19 GiB 的峰值显存台阶，该差异与 GA 取值无关（GA≥2 后显存不再随 GA 增长）；显存敏感场景可优先考虑 GA=1，追求训练效率建议 GA≥2。

常见问题

Q：如何调整训练步数、保存间隔等参数？ 通过环境变量传入，例如：SAVE_INTERVAL / NUM_STEPS / GRADIENT_ACCUMULATION_STEPS / SAVE_ROOT / RESUME_FROM。命令行传参会被直接拒绝（报 unrecognized arguments，训练不启动）。

Q：如何断点续训？ 通过 RESUME_FROM 环境变量指向已保存的 checkpoint 目录（含 transformer/ 与 training_state.pt），续训从断点步数继续、lr 不回 warmup。

1# 首训：每 20 步保存一次
2SAVE_INTERVAL=20 NUM_STEPS=20 SAVE_ROOT=<空目录> bash /root/workspace/lingbot-va/train.sh
3
4# 续训：从 checkpoint 继续
5RESUME_FROM=$SAVE_ROOT/checkpoints/checkpoint_step_20 NUM_STEPS=23 SAVE_ROOT=<另一个目录> bash /root/workspace/lingbot-va/train.sh

当前版本为 step 级断点续训：恢复模型权重、全局步数与学习率调度器，不恢复优化器动量（优化器从零重建）。

Q：权重/数据集可以放在非默认目录吗？ 可以。train.sh 中 BASE_MODEL_PATH、DATASET_PATH 均为环境变量（含默认值），启动前用环境变量覆盖即可，支持任意挂载路径（含 PFS 等）。

Q：多机训练需要注意什么？ 两节点各执行一次 train.sh 即可。脚本将平台注入的 WORLD_SIZE 当作节点数使用（NNODES=${WORLD_SIZE:-1}），多机时自动开启 HSDP 分片，无需额外配置。

Q：显存不足如何优化？ 确认脚本已生效的 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True；必要时可通过降低 gradient_accumulation_steps 调节（注意 GA=1 与 GA≥2 之间存在约 19 GiB 显存台阶）。

后续操作

1. 训练日志输出到 ${LOG_DIR}/${RUN_TAG}_$(hostname).log，可通过日志观察训练进度与 STEP_TIMING 指标。
2. 训练完成前确认 $SAVE_ROOT/checkpoints/ 下已正常落盘 transformer/ 权重与 training_state.pt（后者应远小于 1 MB）。
3. 如需替换数据集或权重，通过 DATASET_PATH / BASE_MODEL_PATH 环境变量重新指向后重新启动训练。