模型训练#

(暨scripts/train.sh 使用指南)

概述#

scripts/train.sh 是一个用于启动分布式训练的 Shell 脚本，它封装了 torchrun 命令，用于在多节点多 GPU 环境下训练模型。

基本用法#

bash scripts/train.sh [CONFIG] [WORK_DIR] [额外参数...]

预训练权重下载#

你可以使用预训练模型的全部或部分模块，以及将各个模型的不同模块像搭积木一样拼接，你可以从下面的列表中选择想要使用的模型权重下载到./checkpoints文件夹。

VLA Models#

Model	Size	Download Link
GR00T N1.5	3B	nvidia/GR00T-N1.5-3B
OpenVLA	7B	openvla/openvla-7b-finetuned-libero-10
PI0_base	3B	yinchimaoliang/pi0_base
PI05_base	3B	yinchimaoliang/pi05_base
PI05_libero	3B	yinchimaoliang/pi05_libero

Vision-Language Models (VLM)#

Model	Size	Download Link
Qwen2.5-VL	3B	Qwen/Qwen2.5-VL-3B-Instruct

Large Language Models (LLM)#

Model	Size	Download Link
Qwen 2.5	3B	Qwen/Qwen2.5-3B
Qwen 2.5	7B	Qwen/Qwen2.5-7B
Llama 2	7B	meta-llama/Llama-2-7b-hf

Vision Backbones#

Model	Download Link
ViT-Large (DINOv2)	timm/vit_large_patch14_reg4_dinov2.lvd142m
ViT-SO400M (SigLIP)	timm/ViT-SO400M-14-SigLIP
SigLIP2	google/siglip2-base-patch16-224
paligemma	google/paligemma-3b-pt-224

小Tip: 可以使用 huggingface-cli download <model-name> --local-dir ./checkpoints/<model-name> 下载想要的模型权重

脚本参数#

位置参数#

参数	位置	默认值	说明
`CONFIG`	$1	`configs/pi05/pi05_paligemma_libero10_full_finetune.py`	训练配置文件路径
`WORK_DIR`	$2	`work_dirs/pi05_paligemma_libero10_full_finetune`	保存日志和 checkpoint 的工作目录
额外参数	$3+	无	直接传递给 `train.py` 的额外参数

环境变量（分布式训练必需）#

这些环境变量用于配置 torchrun 的分布式训练参数：

环境变量	说明	示例
`MLP_WORKER_GPU`	每个节点使用的 GPU 数量	`8`
`MLP_WORKER_NUM`	参与训练的节点总数	`2`
`MLP_ROLE_INDEX`	当前节点的 rank 索引（从 0 开始）	`0`
`MLP_WORKER_0_HOST`	主节点（rank 0）的 IP 地址或主机名	`{MASTER_NODE_IP}`
`MLP_WORKER_0_PORT`	主节点用于通信的端口	`29500`

train.py 支持的额外参数#

以下参数可以作为额外参数传递给 train.py：

参数	类型	说明
`--cfg-options`	key=value 对	覆盖配置文件中的设置，格式为 `xxx=yyy`
`--eval-after-train`	flag	训练完成后自动运行评估
`--resume-from`	路径	从指定的 checkpoint 文件恢复训练

配置 Weights & Biases (wandb)#

Weights & Biases 用于实验追踪与可视化。配置方法如下：

1. 安装 wandb#

wandb 已包含在 requirements.txt 中，也可手动安装：

pip install wandb

2. 登录 wandb 账号#

wandb login

3. 设置环境变量#

export WANDB_PROJECT=fluxvla        # 项目名称（默认：fluxvla）
export WANDB_ENTITY=your-team-name  # 团队或用户名（默认：None）
export WANDB_MODE=online            # online、offline 或 disabled（默认：online）

4. 禁用 wandb 日志#

如果在训练时不需要 wandb 日志记录，可设置：

export WANDB_MODE=disabled

提示：所有 wandb 配置均通过环境变量读取，无需在配置文件中进行额外设置。

使用示例#

1. 单节点多 GPU 训练（8 GPU）#

export MLP_WORKER_GPU=8
export MLP_WORKER_NUM=1
export MLP_ROLE_INDEX=0
export MLP_WORKER_0_HOST=localhost
export MLP_WORKER_0_PORT=29500

bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/pi05_paligemma_libero10_full_finetune

2. 火山云平台训练#

在火山云 MLP 平台上，环境变量 MLP_WORKER_GPU、MLP_WORKER_NUM、MLP_ROLE_INDEX、MLP_WORKER_0_HOST、MLP_WORKER_0_PORT 由平台自动注入，无需手动设置，直接运行即可。

bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/pi05_paligemma_libero10_full_finetune

3. 多节点分布式训练（手动配置）#

如果不在火山云平台上，需要手动设置环境变量：

节点 0（主节点）：

export MLP_WORKER_GPU=8
export MLP_WORKER_NUM=2
export MLP_ROLE_INDEX=0
export MLP_WORKER_0_HOST={MASTER_NODE_IP}
export MLP_WORKER_0_PORT=29500

bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/pi05_paligemma_libero10_full_finetune

节点 1：

export MLP_WORKER_GPU=8
export MLP_WORKER_NUM=2
export MLP_ROLE_INDEX=1
export MLP_WORKER_0_HOST={MASTER_NODE_IP}
export MLP_WORKER_0_PORT=29500

bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/pi05_paligemma_libero10_full_finetune

4. 使用额外参数#

# 覆盖配置参数
bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/custom_run \
    --cfg-options runner.max_epochs=100 runner.save_interval=500

# 从 checkpoint 恢复训练
bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/pi05_resume \
    --resume-from work_dirs/previous_run/checkpoint_epoch_5.pt

# 训练后自动评估
bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/pi05_with_eval \
    --eval-after-train

5. 组合使用多个额外参数#

bash scripts/train.sh \
    configs/pi05/pi05_paligemma_libero10_full_finetune.py \
    work_dirs/pi05_full \
    --resume-from work_dirs/previous_run/latest.pt \
    --eval-after-train \
    --cfg-options runner.max_epochs=50 train_dataloader.batch_size=16

输出内容#

训练完成后，WORK_DIR 目录下会生成以下文件：

config.yaml / config.json - 训练配置的备份
checkpoint_*.pt - 模型检查点文件
dataset_statistics.json - 数据集统计信息（如果有）
{config_name}_{timestamp}.jsonl - 训练日志文件，记录了 loss 变化等训练过程信息
tokenizer 相关文件 - 如果使用了 tokenizer

后端引擎#

FluxVLA 支持多种后端引擎，用于训练和推理 LIBERO、Isaac Sim 等环境下的 VLA 模型。

引擎	说明
FSDP (Fully Sharded Data Parallel)	通过模型权重分片存储减少显存占用，适用于大规模模型训练
DDP (Distributed Data Parallel)	在多 GPU 上并行处理数据，提供高效的分布式训练能力

⚠️ 注意：FSDP 需要使用 两张或以上的 GPU 才能完整使用其分布式训练功能。

注意事项#

环境变量必须设置：在运行脚本前，必须正确设置所有 MLP_* 环境变量，否则 torchrun 无法正确启动分布式训练。
配置文件格式：配置文件使用 MMEngine 的 Config 格式（Python 文件），需要遵循项目的配置规范。
多节点同步：在多节点训练时，所有节点必须使用相同的 CONFIG 和 WORK_DIR 参数。
端口冲突：确保 MLP_WORKER_0_PORT 指定的端口未被占用，且在防火墙中开放。
NCCL 后端：默认使用 NCCL 作为分布式后端，确保 CUDA 和 NCCL 版本兼容。