在百舸上使用 WSA 加速镜像训练 VLA 模型
更新时间：2026-11-06
概述
WSA（TbotSA VLA）是视觉-语言-动作（VLA）模型：Qwen3-VL-2B 主干 + Cosmos tokenizer + DA3 depth teacher（WSA-Large 版另加 Wan2.2 / Wan2.1），在 RoboTwin-LeRobot 数据集上微调。提供 WSA-Base、WSA-Large 两个规模。本加速镜像针对 RTX Pro6000（sm_120）芯片做训练加速，已内置完整训练栈与优化后的训练脚本，支持单机多卡与多机扩展训练。
本示例以 RoboTwin-LeRobot-v3.0 仿真数据集为例，介绍在百舸平台上完成训练全流程。
适用场景
- 在百舸平台训练 WSA VLA 模型，用于机器人操作仿真（RoboTwin 等）任务
- 需要单机多卡训练，或双机扩展训练以获得更高吞吐的场景
- 需要断点续训、性能参考与常见问题指引的客户
环境与资源要求
| 项 | 内容 |
|---|---|
| 适配芯片 | 支持规格：aihc.lgn7ib |
| 源码 | 已烘入镜像 /root/workspace/WSA |
| 镜像环境 | 自带训练栈（无需 venv / conda） |
| 工作目录 | /root/workspace/WSA |
| 数据集 | RoboTwin-LeRobot-v3.0 |
| 模型权重 | WSA-Base、WSA-Large、Qwen3-VL-2B-Instruct、Cosmos-0.1-Tokenizer-CI8x8、DA3-LARGE-1.1、Wan2.2-TI2V-5B、Wan2.1-T2V-1.3B |
同一镜像通过开关切换即得 baseline / optimized：内置 run_wsa_base_optimized.sh（优化全开）/ run_wsa_base_baseline.sh（对照）+ 续训参数透传。
在百舸平台准备训练环境
- 登录百舸平台，创建或进入开发机，选择 RTX Pro6000 资源规格与上述加速镜像
- 打开开发机终端（或 WebIDE 终端），进入工作目录
/root/workspace/WSA - 按本章节顺序准备数据与权重，然后启动训练
准备数据与权重
数据集与权重不打包进镜像，训练前需下载/挂载到 /root/workspace 下（默认部署路径见下表，均可用环境变量整体覆盖到 PFS，见「启动训练 · 参数说明」）。HF 为公开参考地址，实际下载以 BOS 为准。
资产清单
| 物料 | HuggingFace 地址 | BOS 地址 | 部署路径（默认，可 env 覆盖） |
|---|---|---|---|
| RoboTwin-LeRobot-v3.0（数据集） | hxma/RoboTwin-LeRobot-v3.0 | bos://aihc-rdw-bj/huggingface.co/datasets/hxma/RoboTwin-LeRobot-v3.0/ |
/root/workspace/data/RoboTwin-LeRobot-v3.0（DATA_PATH） |
| WSA-Base（主权重） | zaleni/WSA-Base | bos://aihc-models-bj/zaleni/WSA-Base/ |
/root/workspace/models/WSA-Base（CHECKPOINT_PATH） |
| WSA-Large（主权重） | zaleni/WSA-Large | bos://aihc-models-bj/zaleni/WSA-Large/ |
/root/workspace/models/WSA-Large |
| Qwen3-VL-2B-Instruct | Qwen/Qwen3-VL-2B-Instruct | bos://aihc-models-bj/Qwen/Qwen3-VL-2B-Instruct/ |
/root/workspace/models/Qwen3-VL-2B-Instruct（QWEN_VL_PATH） |
| Cosmos-0.1-Tokenizer-CI8x8 | nvidia/Cosmos-0.1-Tokenizer-CI8x8 | bos://aihc-models-bj/nvidia/Cosmos-0.1-Tokenizer-CI8x8/ |
/root/workspace/models/Cosmos-0.1-Tokenizer-CI8x8（COSMOS_TOKENIZER_PATH） |
| DA3-LARGE-1.1 | depth-anything/DA3-LARGE-1.1 | bos://aihc-models-bj/depth-anything/DA3-LARGE-1.1/ |
/root/workspace/models/DA3-LARGE-1.1（DA3_PATH） |
| Wan2.2-TI2V-5B（WSA-Large 用） | Wan-AI/Wan2.2-TI2V-5B | bos://aihc-models-bj/Wan-AI/Wan2.2-TI2V-5B/ |
/root/workspace/models/Wan2.2-TI2V-5B |
| Wan2.1-T2V-1.3B（WSA-Large 用） | Wan-AI/Wan2.1-T2V-1.3B | bos://aihc-models-bj/Wan-AI/Wan2.1-T2V-1.3B/ |
/root/workspace/models/Wan2.1-T2V-1.3B |
方式一：HuggingFace 下载
使用 HuggingFace 下载前先设置镜像站：
Bash
1export HF_ENDPOINT=https://hf-mirror.com
2export HF_HUB_ENABLE_HF_TRANSFER=1
3export HF_HOME=/root/workspace/models/huggingface
Bash
1# 数据集
2huggingface-cli download hxma/RoboTwin-LeRobot-v3.0 --repo-type dataset --local-dir /root/workspace/data/RoboTwin-LeRobot-v3.0
3# 权重
4huggingface-cli download zaleni/WSA-Base --local-dir /root/workspace/models/WSA-Base
5huggingface-cli download zaleni/WSA-Large --local-dir /root/workspace/models/WSA-Large
6# ...（Qwen3-VL-2B-Instruct / Cosmos-0.1-Tokenizer-CI8x8 / DA3-LARGE-1.1 / Wan2.2-TI2V-5B / Wan2.1-T2V-1.3B 同理）方式二：BOS 下载
Bash
1mkdir -p /root/workspace/data/RoboTwin-LeRobot-v3.0
2bosutil cp -r bos://aihc-rdw-bj/huggingface.co/datasets/hxma/RoboTwin-LeRobot-v3.0/ /root/workspace/data/RoboTwin-LeRobot-v3.0/
3
4mkdir -p /root/workspace/models/WSA-Base
5bosutil cp -r bos://aihc-models-bj/zaleni/WSA-Base/ /root/workspace/models/WSA-Base/
6
7# ...（其余权重同理，路径见上方资产清单）训练前预处理（生成 repo 列表 + norm_stats，每套只需一次）
在 /root/workspace/WSA 下执行。
WSA-Base
Bash
1#!/bin/bash
2DATA_PATH="${DATA_PATH:-/root/workspace/data/RoboTwin-LeRobot-v3.0}"
3find -L "${DATA_PATH}" -path "*/meta/info.json" -print \
4 | while read -r info; do dirname "$(dirname "$info")"; done \
5 | sort -u > robotwin_repo_ids.txt
6python tools/compute_norm_stats_multi.py \
7 --repo_id_file robotwin_repo_ids.txt \
8 --action_mode delta --chunk_size 50 --num_workers 8 \
9 --output_path ./norm_stats/aloha/delta/stats.json产物落在相对路径
./norm_stats/aloha/delta/stats.json，与NORM_STATS_ROOT默认值对应。若把NORM_STATS_ROOT改到别处（如 PFS），须先把该文件复制到${NORM_STATS_ROOT}/aloha/delta/stats.json，否则训练在make_dataset报use_external_stats=True but no external stats file was found。
WSA-Large（另加 text_embeds 预计算）
Bash
1#!/bin/bash
2DATA_PATH="${DATA_PATH:-/root/workspace/data/RoboTwin-LeRobot-v3.0}"
3WAN_MODEL_PATH="${WAN_MODEL_PATH:-/root/workspace/models/Wan2.2-TI2V-5B}"
4WAN_TOKENIZER_PATH="${WAN_TOKENIZER_PATH:-/root/workspace/models/Wan2.1-T2V-1.3B}"
5python tools/discover_robotwin_repos.py \
6 --robotwin-root ${DATA_PATH} \
7 --output-file robotwin_repo_ids_large.txt --require-three-cameras true
8python tools/precompute_text_embeds.py \
9 --model-id ${WAN_MODEL_PATH} --tokenizer-model-id ${WAN_TOKENIZER_PATH} \
10 --repo-id-file robotwin_repo_ids_large.txt \
11 --text-embedding-cache-dir outputs/WSA_Large/text_embeds/robotwin \
12 --redirect-common-files false --device cuda
13python tools/compute_norm_stats_multi.py \
14 --repo_id_file robotwin_repo_ids_large.txt \
15 --action_mode delta --chunk_size 32 --num_workers 8 \
16 --output_path ./norm_stats_large/aloha/delta/stats.json实测总耗时约 25 分钟（text embeds 约 8.6 min、产物 24G；norm stats 约 16.6 min）。产物
outputs/WSA_Large/text_embeds/、norm_stats_large/为相对/root/workspace/WSA的路径、无环境变量入口，要落 PFS 需把这两个目录 bind mount 进容器。Large 的 stats（chunk_size 32）与 Base（chunk_size 50）不可混用。
启动训练
参数说明
优化版 launcher 已把数据/权重/输出/命名全部做成环境变量，默认落 /root/workspace，可整体覆盖到 PFS。
| 环境变量 | 默认值 | 说明 |
|---|---|---|
DATA_PATH |
/root/workspace/data/RoboTwin-LeRobot-v3.0 |
RoboTwin 数据集（底层 ROBOTWIN_ROOT） |
CHECKPOINT_PATH |
/root/workspace/models/WSA-Base |
WSA-Base 主权重（底层 POLICY_INIT_PATH） |
QWEN_VL_PATH |
/root/workspace/models/Qwen3-VL-2B-Instruct |
Qwen3-VL 主干 |
COSMOS_TOKENIZER_PATH |
/root/workspace/models/Cosmos-0.1-Tokenizer-CI8x8 |
Cosmos tokenizer |
DA3_PATH |
/root/workspace/models/DA3-LARGE-1.1 |
DA3 depth teacher |
OUTPUT_DIR |
/root/workspace/WSA/outputs |
输出根（含 checkpoints） |
NORM_STATS_ROOT |
/root/workspace/WSA/norm_stats |
norm_stats 目录 |
优化开关（run_wsa_base_optimized.sh 默认全开；run_wsa_base_baseline.sh 全关，用于对照/复现加速比）：COMPILE_MODEL / COMPILE_DA3_TEACHER_BACKBONE / OPTIMIZER_FUSED / FREEZE_PERMANENTLY_UNUSED_PARAMS（=true）、ATTENTION_BACKEND=sdpa、LEROBOT_DDP_BUCKET_CAP_MB=600（baseline=25）。多机由 NODE_COUNT / NODE_RANK / MASTER_ADDR / MASTER_PORT 注入。
覆盖到 PFS 示例：
Bash
1DATA_PATH=/pfs/pfs-OqLB2M/dataset/RoboTwin-LeRobot-v3.0 CHECKPOINT_PATH=/pfs/pfs-OqLB2M/models/zaleni/WSA-Base OUTPUT_DIR=/pfs/<你的目录>/outputs bash run_wsa_base_optimized.shWSA-Base 训练（单机 / 多机自适应）
镜像内置 /root/workspace/WSA/run_wsa_base_optimized.sh（优化全开）与 run_wsa_base_baseline.sh（全关，对照）：
Bash
1#!/usr/bin/env bash
2# WSA-Base 优化版 launcher（单机/多机自适应）—— 优化项全开
3# 打开：COMPILE_MODEL / COMPILE_DA3_TEACHER_BACKBONE / OPTIMIZER_FUSED /
4# FREEZE_PERMANENTLY_UNUSED_PARAMS / LEROBOT_DDP_BUCKET_CAP_MB=600（sdpa 同 baseline）。
5# 覆盖到 PFS 例：DATA_PATH=/pfs/... CHECKPOINT_PATH=/pfs/... ... bash run_wsa_base_optimized.sh
6# 续训（须 FREEZE=false）：FREEZE_PERMANENTLY_UNUSED_PARAMS=false bash run_wsa_base_optimized.sh \
7# --resume=true --config_path=<OUTPUT_DIR>/.../checkpoints/last/pretrained_model/train_config.json
8set -euo pipefail
9cd "$(dirname "${BASH_SOURCE[0]}")"
10
11DATA_PATH="${DATA_PATH:-/root/workspace/data/RoboTwin-LeRobot-v3.0}"
12CHECKPOINT_PATH="${CHECKPOINT_PATH:-/root/workspace/models/WSA-Base}"
13QWEN_VL_PATH="${QWEN_VL_PATH:-/root/workspace/models/Qwen3-VL-2B-Instruct}"
14COSMOS_TOKENIZER_PATH="${COSMOS_TOKENIZER_PATH:-/root/workspace/models/Cosmos-0.1-Tokenizer-CI8x8}"
15DA3_PATH="${DA3_PATH:-/root/workspace/models/DA3-LARGE-1.1}"
16OUTPUT_DIR="${OUTPUT_DIR:-/root/workspace/WSA/outputs}"
17NORM_STATS_ROOT="${NORM_STATS_ROOT:-/root/workspace/WSA/norm_stats}"
18
19export PROC_PER_NODE="${PROC_PER_NODE:-$(nvidia-smi -L | wc -l)}"
20export NODE_COUNT="${NODE_COUNT:-1}"
21export NODE_RANK="${NODE_RANK:-0}"
22export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
23export MASTER_PORT="${MASTER_PORT:-6379}"
24export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"
25export LEROBOT_DDP_BUCKET_CAP_MB="${LEROBOT_DDP_BUCKET_CAP_MB:-600}"
26
27USE_DIST_LOADING="${USE_DIST_LOADING:-true}" \
28BATCH_SIZE="${BATCH_SIZE:-8}" \
29GRADIENT_ACCUMULATION_STEPS="${GRADIENT_ACCUMULATION_STEPS:-4}" \
30COMPILE_MODEL="${COMPILE_MODEL:-true}" \
31COMPILE_MODE="${COMPILE_MODE:-default}" \
32COMPILE_DA3_TEACHER_BACKBONE="${COMPILE_DA3_TEACHER_BACKBONE:-true}" \
33COMPILE_REGIONAL_MODULES="${COMPILE_REGIONAL_MODULES:-false}" \
34OPTIMIZER_FUSED="${OPTIMIZER_FUSED:-true}" \
35FREEZE_PERMANENTLY_UNUSED_PARAMS="${FREEZE_PERMANENTLY_UNUSED_PARAMS:-true}" \
36ATTENTION_BACKEND="${ATTENTION_BACKEND:-sdpa}" \
37STEPS="${STEPS:-200000}" \
38WARMUP_STEPS="${WARMUP_STEPS:-20}" \
39LOG_FREQ="${LOG_FREQ:-10}" \
40SAVE_FREQ="${SAVE_FREQ:-10000}" \
41POLICY_INIT_PATH="${CHECKPOINT_PATH}" \
42QWEN3_VL_PRETRAINED_PATH="${QWEN_VL_PATH}" \
43COSMOS_TOKENIZER_PATH_OR_NAME="${COSMOS_TOKENIZER_PATH}" \
44DA3_MODEL_PATH_OR_NAME="${DA3_PATH}" \
45ROBOTWIN_ROOT="${DATA_PATH}" \
46BASE_OUTPUT_DIR="${OUTPUT_DIR}" \
47ACTION_TYPE=delta \
48USE_EXTERNAL_STATS=true \
49DATASET_EXTERNAL_STATS_ROOT="${NORM_STATS_ROOT}" \
50LOG_DA3_TEACHER_TIMING=true \
51bash launch/wsa_base_finetune_robotwin.sh "$@"baseline 版（run_wsa_base_baseline.sh）结构相同，仅 6 处开关不同：LEROBOT_DDP_BUCKET_CAP_MB=25、COMPILE_MODEL=false、COMPILE_DA3_TEACHER_BACKBONE=false、OPTIMIZER_FUSED=false、FREEZE_PERMANENTLY_UNUSED_PARAMS=false（ATTENTION_BACKEND=sdpa 与路径/多机段一致）。
调用：
Bash
1bash run_wsa_base_optimized.sh # 单机 8 卡
2NODE_COUNT=2 NODE_RANK=0 MASTER_ADDR=<rank0内网IP> MASTER_PORT=6379 bash run_wsa_base_optimized.sh # 双机 rank0
3NODE_COUNT=2 NODE_RANK=1 MASTER_ADDR=<rank0内网IP> MASTER_PORT=6379 bash run_wsa_base_optimized.sh # 双机 rank1（先起 rank1）WSA-Large 训练
训练脚本：
Bash
1#!/bin/bash
2
3# ==================== 路径配置 ====================
4DATA_PATH="${DATA_PATH:-/root/workspace/data/RoboTwin-LeRobot-v3.0}"
5CHECKPOINT_PATH="${CHECKPOINT_PATH:-/root/workspace/models/WSA-Large}"
6QWEN_VL_PATH="${QWEN_VL_PATH:-/root/workspace/models/Qwen3-VL-2B-Instruct}"
7COSMOS_TOKENIZER_PATH="${COSMOS_TOKENIZER_PATH:-/root/workspace/models/Cosmos-0.1-Tokenizer-CI8x8}"
8DA3_PATH="${DA3_PATH:-/root/workspace/models/DA3-LARGE-1.1}"
9WAN_MODEL_PATH="${WAN_MODEL_PATH:-/root/workspace/models/Wan2.2-TI2V-5B}"
10WAN_TOKENIZER_PATH="${WAN_TOKENIZER_PATH:-/root/workspace/models/Wan2.1-T2V-1.3B}"
11
12# ==================== 训练启动 ====================
13# 数据加载
14NUM_WORKERS="${NUM_WORKERS:-16}" \
15USE_DIST_LOADING="${USE_DIST_LOADING:-true}" \
16BATCH_SIZE="${BATCH_SIZE:-6}" \
17GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-2}" \
18\
19COMPILE_MODEL="${COMPILE_MODEL:-false}" \
20COMPILE_DA3_TEACHER_BACKBONE="${COMPILE_DA3_TEACHER_BACKBONE:-true}" \
21COMPILE_REGIONAL_MODULES="${COMPILE_REGIONAL_MODULES:-true}" \
22WSA_MOT_CKPT_LAYERS="${WSA_MOT_CKPT_LAYERS:-12}" \
23WSA_GC_FREEZE="${WSA_GC_FREEZE:-1}" \
24WSA_TS9="${WSA_TS9:-1}" \
25WSA_MOT_BLOCK_SPLIT_ATTN="${WSA_MOT_BLOCK_SPLIT_ATTN:-1}" \
26OPTIMIZER_FUSED="${OPTIMIZER_FUSED:-true}" \
27WSA_VAE_COMPILE_MODE="${WSA_VAE_COMPILE_MODE:-max-autotune-no-cudagraphs}" \
28ATTENTION_BACKEND="${ATTENTION_BACKEND:-sdpa}" \
29WARMUP_STEPS="${WARMUP_STEPS:-40}" \
30\
31POLICY_INIT_PATH="${CHECKPOINT_PATH}" \
32QWEN3_VL_PRETRAINED_PATH="${QWEN_VL_PATH}" \
33COSMOS_TOKENIZER_PATH_OR_NAME="${COSMOS_TOKENIZER_PATH}" \
34DA3_MODEL_PATH_OR_NAME="${DA3_PATH}" \
35ROBOTWIN_ROOT="${DATA_PATH}" \
36WAN_MODEL_ID="${WAN_MODEL_PATH}" \
37WAN_TOKENIZER_MODEL_ID="${WAN_TOKENIZER_PATH}" \
38\
39ACTION_TYPE=delta \
40NORMALIZATION_STATS_PATH=./norm_stats_large/aloha/delta/stats.json \
41TEXT_EMBED_CACHE_DIR=outputs/WSA_Large/text_embeds/robotwin \
42LOAD_TEXT_ENCODER=false \
43SKIP_DIT_LOAD_FROM_PRETRAIN=true \
44LOG_DA3_TEACHER_TIMING=true \
45bash launch/wsa_large_finetune_robotwin.sh注意：
- 步数用
TRAIN_MAX_STEPS=n控制。Large 下STEPS不生效，会被NUM_EPOCHS=5改写（实际约 31.6 万步） - 脚本中的
QWEN_VL_PATH/COSMOS_TOKENIZER_PATH底层 launcher 未引用，Large 不需要准备这两份权重
性能参考
以下数据为 RTX Pro6000 ×8 实测（STEPS=60，取 step20–60 稳态；同一镜像开关关闭 = baseline、打开 = optimized）。
WSA-Base 单机加速比
| 配置 | 稳态 update | 吞吐（全局） | 加速比 | da3 teacher |
|---|---|---|---|---|
| baseline（全关，sdpa） | ~3.18s | 80.5 samples/s | 1.00× | 0.132s |
| optimized（全开） | ~2.61s | ≈98.0 samples/s | ≈1.22× | 0.102s |
优化版峰值显存 ≈94.5 GB（近 97.9 GB 上限，主要来自 torch.compile inductor autotuning 首步瞬时占用）。
小显存卡/更大 batch 时建议关 COMPILE_MODEL 或设 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True。（sdpa 为两版共用功能基线，其收益不计入此 1.22×。）
WSA-Base 双机 2×8 扩展
双机 2×8（16 卡，effective batch = 8 × 16 × 4 = 512，取 step30–60 稳态）：稳态 ~2.96s/update，聚合吞吐 172.9 samples/s（约 10.8/卡），相对单机 optimized（98.0 samples/s）8→16 扩展比 1.76×、并行效率 88%。
每卡 batch 16 档（GA=4 不变，其余开关同 optimized，单机 batch 扫描 10/12/14/16 均不 OOM，batch 16 吞吐最高；多机训练建议 BATCH_SIZE=16）：
| 配置 | effective batch | 稳态 update | 吞吐（全局） | 峰值显存 |
|---|---|---|---|---|
| 单机 8 卡 | 8 × 16 × 4 = 512 | 4.733s | 108.2 samples/s | 91.3 GiB/卡 |
| 双机 2×8 | 16 × 16 × 4 = 1024 | 5.000s | 204.8 samples/s | 91.4 GiB/卡 |
8→16 扩展比 1.89×、并行效率 94.6%。
WSA-Large 单机
| 配置 | 稳态 update | 吞吐（全局） | 峰值显存 | loss（step10→60） |
|---|---|---|---|---|
train_large.sh 默认（8×6×2=96） |
1.700s | 56.47 samples/s | 86.3 GiB/卡 | 4.464 → 2.807 |
以上数据为特定测试环境下的参考值。实际性能会受到资源规格、数据读取方式、网络环境、任务参数等因素影响。
续训
FREEZE_PERMANENTLY_UNUSED_PARAMS=true（优化版默认）会硬拒 resume，续训须置为false。
WSA-Base 续训
Bash
1# 1) 出 checkpoint（FREEZE=false）
2FREEZE_PERMANENTLY_UNUSED_PARAMS=false bash run_wsa_base_optimized.sh
3# 2) 从断点续训
4FREEZE_PERMANENTLY_UNUSED_PARAMS=false bash run_wsa_base_optimized.sh --resume=true --config_path=<OUTPUT_DIR>/<POLICY>/<JOB>/checkpoints/last/pretrained_model/train_config.json续训时总步数需与首训一致：学习率调度按首训步数保存，调大步数后 lr 会停在调度终点。FREEZE=false 稳态 ≈91 samples/s（≈1.13×），比 FREEZE=true（98）降约 7%。
WSA-Large 续训
Bash
1RESUME=true RESUME_CHECKPOINT_DIR=<OUTPUT_DIR>/<JOB>/checkpoints/<step> TRAIN_MAX_STEPS=<与首训相同> bash train_large.sh续训时总步数须与首训一致：调大会使学习率调度沿原 cosine 回爬。
参考
评价此篇文章