总览

本文档面向使用 LoongForge框架进行具身模型训练的用户，聚焦框架级通用能力，说明框架各模块支持的功能、配置入口及使用方式。具体模型的专用训练方案、数据预处理脚本、性能优化项目请参阅对应的 快速入门 文档。

1. 使用入口与配置约定

本章说明 embodied 的目录结构与启动脚本约定，并介绍 TrainingArgs / ModelConfig / DataConfig 三类配置对象及其覆盖优先级。

1.1 目录说明

路径	说明
examples/embodied/	模型级启动脚本目录，可通过脚本末尾的透传参数覆盖默认配置
configs/models/embodied/	模型默认 YAML 配置目录，包含 model: / data: 两个顶层配置段
loongforge/embodied/train.py	训练入口，负责解析配置、构建 Trainer 并启动训练
loongforge/embodied/train/training_args.py	通用训练参数定义文件，负责生成 shell CLI
loongforge/embodied/train/config_map.py	模型配置路由表，将 --model-name 绑定到 YAML、ModelConfig 与 DataConfig
loongforge/embodied/model/	模型组网、模型注册
loongforge/embodied/data/datasets/	数据处理相关功能

训练链路如下：

examples/embodied/<model>/run_*.sh
    ↓
loongforge/embodied/train.py
    ↓
parse_train_args()
    ↓
build_model_trainer()
    ↓
trainer.train()

1.2 启动脚本约定

启动脚本通常负责设置环境变量、路径、分布式参数与模型默认训练参数，并在命令末尾保留 "$@"，用于透传用户追加的 shell flag 或 YAML dotlist 覆盖项：

PYTHONPATH=$LOONGFORGE_PATH:${PYTHONPATH:-} \
torchrun "${DISTRIBUTED_ARGS[@]}" \
    "$LOONGFORGE_PATH/loongforge/embodied/train.py" \
    "${MODEL_CONFIG_ARGS[@]}" \
    "${DATA_ARGS[@]}" \
    "${TRAINING_ARGS[@]}" \
    "${DISTRIBUTED_TRAINING_ARGS[@]}" \
    "${LOGGING_ARGS[@]}" \
    "$@"

示例：

bash examples/embodied/pi05/run_pi05_ddp_finetune.sh \
    --train-iters 10000 \
    --per-device-batch-size 8 \
    model.action_horizon=64 \
    data.image_size=256

1.3 配置分层

将配置分为三类对象：

配置对象	配置入口	作用域	示例
TrainingArgs	Shell flag	训练流程参数	--train-iters, --lr-base
ModelConfig	YAML 或 model.xxx=...	模型结构与训练策略参数	model.action_horizon=64
DataConfig	YAML 或 data.xxx=...	数据加载与预处理参数	data.image_size=256

配置优先级为：

dataclass 默认值  <  YAML 配置  <  shell flag / dotlist 覆盖

2. 数据处理

数据处理模块负责：将数据集样本转换为模型可直接消费的 PreparedBatch，用户通过 TrainingArgs 指定数据集格式与 DataLoader 行为，通过 DataConfig 配置模型相关的数据处理逻辑。

数据处理链路：

Dataset
    ↓
sample-level transform
    ↓
preprocessor / collate_fn
    ↓
PreparedBatch
    ↓
batch.to(device)
    ↓
model.forward(batch)

2.1 数据集格式

通过 --dataset-format 选择数据读取格式，支持 LeRobot、HDF5 和 dummy 三种格式。dummy_datasets 可在无真实数据时生成随机样本，用于调试验证。

功能	配置项	默认值	取值 / 类型	说明
数据格式	--dataset-format	lerobot_datasets	lerobot_datasets, hdf5_datasets, dummy_datasets	选择数据集格式
数据路径	--dataset-path	None	本地路径或数据集 id	训练数据路径
数据 split	--split	train	字符串	数据集切分
dummy 样本数	--num-samples	100	正整数	dummy_datasets 下生成的样本数

2.2 LeRobot 数据策略

使用 LeRobot 格式时，可通过以下参数进一步控制数据加载行为：指定数据集格式版本（v2.0 / v2.1 / v3.0）、选择针对不同机器人或任务的构建策略、配置视频解码后端，以及指定 Robot型号以匹配 embodiment 的 action-state 布局。

功能	配置项	默认值	取值 / 类型	说明
LeRobotdataset 版本	--lerobotdataset-version	v3.0	v2.0, v2.1, v3.0	解析不同 LeRobot 磁盘格式
数据策略	--dataset-strategy	default	default, fastwam, groot_n1_7, cosmos3_droid, dreamzero	选择 LeRobot 构建策略
视频后端	--video-backend	torchcodec	torchcodec, decord, opencv, pyav, torchvision_av	视频解码实现
robot 类型	--robot-type	None	字符串	选择 embodiment / action-state layout

注意： --video-backend 用于指定 --lerobotdataset-version v2.x 系列及其变体的数据读取后端；v3.0 系列默认支持 torchcodec、pyav 两种后端。

2.3 DataLoader 行为

本节配置项控制 DataLoader 的 worker 并行度、多进程启动方式、分布式 index 切分与流式读取模式，覆盖从数据预取到超大规模数据集加载的常见需求。

功能	配置项	默认值	取值 / 类型	说明
worker 数	--num-workers	4	非负整数	每个 rank 的 DataLoader worker 数
worker seed	--dataloader-seed-workers	False	布尔开关	是否基于 --seed 设置 worker seed
多进程上下文	--dataloader-multiprocessing-context	None	fork, spawn, forkserver	DataLoader worker 启动方式
分布式采样方式	--distributed-sampler-mode	cyclic	cyclic, block	分布式 sampler 的 index 切分方式（可扩展）
流式读取	--streaming	False	布尔开关	使用 streaming / iterable dataset

3. 训练配置

训练配置模块负责解析 shell 参数和 YAML 文件，并生成 TrainingArgs、ModelConfig、DataConfig 三个类型化配置对象。通用训练能力均通过 TrainingArgs 暴露为 shell flag。

3.1 模型配置选择

支持以下能力：

• 通过 --model-name 选择预注册模型，自动绑定 YAML、ModelConfig 与 DataConfig

• 通过 --config-file 指定默认 YAML，使用自定义配置文件

• 通过 --tokenizer-path 指定 tokenizer 路径（本地路径或 HF repo id）

功能	配置项	默认值	取值 / 类型	说明
选择模型	--model-name	None	config_map.py 中注册的模型名	选择模型 schema、默认 YAML、ModelConfig 与 DataConfig
指定 YAML	--config-file	None	YAML 文件路径	覆盖 --model-name 对应的默认 YAML
指定 tokenizer	--tokenizer-path	None	本地路径或 HF repo id	设置 tokenizer 路径，并同步到 TOKENIZER_PATH 环境变量

注意： 即使使用 --config-file，仍需提供 --model-name，用于选择结构化配置类。

3.2 训练基础参数

本节配置项控制训练规模（迭代步数与 batch 大小）、可复现性（随机种子）与产物输出目录，通过梯度累积在不增加单卡显存的前提下弹性扩大全局 batch。

功能	配置项	默认值	取值 / 类型	说明
训练总步数	--train-iters	150000	正整数	optimizer update 步数
每设备 batch	--per-device-batch-size	4	正整数	单个 rank forward时的micro-batch
梯度累积	--gradient-accumulation-steps	1	正整数	每次 optimizer step 前累积的 micro-batch 数
随机种子	--seed	3047	整数	用于训练初始化与数据随机性控制
输出目录	--output-dir	outputs/default	路径	保存日志、checkpoint 与运行产物

全局 batch 计算方式：

global_batch_size = per_device_batch_size * world_size * gradient_accumulation_steps

3.3 学习率与优化器

学习率与优化器相关能力较多，本节按「基础调度 → 优化器实现 → 分组学习率」的顺序展开。

支持以下能力：

• 基础学习率与分模块独立学习率

• 10 种学习率调度策略，涵盖线性、cosine、polynomial、恒定及带 warmup / min_lr 变体

• 6 种优化器实现，包含标准 AdamW 及多种 CUDA 融合加速变体

• 梯度裁剪与权重衰减，支持对 bias / norm 参数单独分组

3.3.1 基础学习率与调度

功能	配置项	默认值	取值 / 类型	说明
基础学习率	--lr-base	2.5e-5	float	默认参数组学习率
学习率策略	--lr-decay-style	cosine_with_min_lr	见下表	scheduler 类型
warmup 步数	--lr-warmup-iters	2000	非负整数	线性 warmup 步数
decay 步数	--lr-decay-iters	None	正整数或不传	不传时使用 --train-iters
最小学习率	--min-lr	1e-6	float	decay 下限
梯度裁剪	--clip-grad	1.0	float，<=0 表示关闭	最大梯度范数
权重衰减	--weight-decay	0.01	float	decoupled weight decay 系数
权重衰减分组	--weight-decay-grouping	all	all, bias_norm	是否对 bias / norm 参数禁用 weight decay

--lr-decay-style 支持 10 种调度策略，涵盖linear、cosine、polynomial等。

取值	说明
linear	warmup 后线性衰减
cosine	warmup 后 cosine 衰减
cosine_with_restarts	cosine 衰减并带 hard restart
polynomial	polynomial decay
constant	恒定学习率
constant_with_warmup	warmup 后保持恒定
inverse_sqrt	inverse sqrt decay
cosine_with_min_lr	cosine 衰减到 --min-lr
cosine_warmup_with_min_lr	带最小 LR 的 cosine warmup
lambda_linear	框架自定义 cycle linear scheduler

3.3.2 优化器实现

通过 --optimizer 指定优化器实现，支持标准 AdamW、多种融合加速实现（PyTorch / TransformerEngine / Apex）及 Adam、SGD：

取值	说明
AdamW（默认）	标准 AdamW
TorchFusedAdamW	PyTorch fused AdamW
TEFusedAdamW	TransformerEngine FusedAdam
ApexFusedAdamW	Apex FusedAdam
Adam	torch Adam
SGD	torch SGD

Fused Adam 加速实现说明：

• TEFusedAdamW：将参数更新融合为单次 CUDA kernel，显著降低显存带宽压力，依赖 Transformer-engine

• ApexFusedAdamW：多参数组融合更新，在大模型场景下优化器步骤耗时更低，依赖 apex

3.3.3 模块级分组学习率

通过 --lr-group 为不同模块配置独立学习率，常用于微调时对 backbone 使用较小学习率、对动作头使用较大学习率，未匹配的参数使用 --lr-base。

• 配置项：--lr-group

• 默认值：None（不启用，所有参数统一使用 --lr-base）

• 格式：module.path=lr,module.path=lr

示例：

bash examples/embodied/pi05/run_pi05_ddp_finetune.sh \
    --lr-base 1.0e-4 \
    --lr-group "model.backbone=1.0e-5,model.action_head=1.0e-4"

配置规则：

• 路径匹配顺序敏感，子模块路径应位于父模块路径之前

• 未匹配参数使用 --lr-base

• 模块路径以模型实现中的实际属性路径为准

3.4 Checkpoint

Checkpoint 模块提供以下能力：

• 权重保存：支持 safetensors、pt、dcp 三种格式，其中 DCP 可启用异步保存

• 训练状态保存：持久化 optimizer、scheduler、RNG 与 DataLoader 状态，用于断点续训

• 预训练加载：可指定外部 checkpoint 初始化模型参数

功能	配置项	默认值	取值 / 类型	说明
加载预训练权重	--pretrained-checkpoint	None	checkpoint 路径	用于初始化模型参数
续训	--resume	False	布尔开关	从 output_dir/checkpoints 查找最新 checkpoint 并恢复
保存间隔	--save-interval	10000	非负整数，0 表示关闭	每 N 个 update step 保存一次 checkpoint
保存格式	--save-format	safetensors	safetensors, pt, dcp	checkpoint 文件格式
保存训练状态	--save-training-state	True	布尔开关	保存 optimizer、scheduler、RNG 与 DataLoader 状态
异步保存	--async-save	False	布尔开关	DCP 格式下可启用异步保存

续训示例：

bash examples/embodied/pi05/run_pi05_ddp_finetune.sh \
    --output-dir /path/to/previous_run \
    --resume

3.5 日志与监控

支持以下监控方式：

• 控制台 metrics 日志，可配置记录间隔与阶段计时详细程度

• W&B 集成，支持 online / offline / disabled 三种模式

• TensorBoard 集成，指定目录即可启用

• 可按 rank 粒度控制 loss 聚合与输出来源

功能	配置项	默认值	取值 / 类型	说明
日志间隔	--log-interval	1	正整数	每 N 步记录 metrics
详细计时间隔	--detail-log-interval	20	非负整数	每 N 步记录阶段耗时
计时日志级别	--timing-log-level	0	0, 1	阶段耗时日志详细程度
W&B 项目	--wandb-project	loongforge	字符串	W&B project 名称
W&B 模式	--wandb-mode	disabled	online, offline, disabled	W&B 启用模式
TensorBoard 目录	--tensorboard-dir	None	路径	不传表示关闭 TensorBoard
loss 日志 rank	--loss-log-rank	[-1]	rank 列表，-1 表示全局平均	控制 loss 聚合与输出来源

3.6 冻结训练

通过 --freeze-modules 冻结指定模块参数，常用于微调时固定视觉编码器或语言模型主干，仅更新动作头等目标模块。模块路径以模型实现中的 named_modules() 为准，具体模型的常用冻结路径在对应 Quick Start 中说明。

功能	配置项	默认值	取值 / 类型	说明
冻结模块	--freeze-modules	空字符串（不冻结任何模块）	逗号分隔模块路径	将匹配模块参数设置为 requires_grad=False

4. 分布式 Trainer

Trainer 模块负责训练生命周期编排，包括分布式上下文初始化、模型构建、权重加载、模型包装、优化器与 scheduler 构建、DataLoader 构建、训练循环、日志、checkpoint 与资源清理。

4.1 Trainer 选择

通过 --trainer-type 选择训练器，可选值为 trainer_builder.py 中注册的 Trainer 类名。默认的 FinetuneTrainer 适用于标准单数据流的监督微调；若涉及多数据流、特殊 loss 组合或非标准 step 调度，可在 trainer_builder.py 中注册自定义 Trainer 类。

4.2 分布式策略

支持两种分布式并行策略：

• DDP：标准数据并行，适用于模型与 optimizer state 均可放入单卡的场景；可叠加 ZeRO-1 分片 optimizer state 以节省显存

• FSDP：全参数分片，适用于模型、梯度或 optimizer state 超出单卡显存的场景；支持 HSDP（二维 mesh 分片）

训练精度支持 bfloat16（默认）、float16、float32。

功能	配置项	默认值	取值 / 类型	说明
分布式策略	--distributed-strategy	fsdp	ddp, fsdp	选择 DDP 或 FSDP
训练精度	--dtype	bfloat16	bfloat16, float16, float32	模型训练 dtype
DDP ZeRO-1	--zero-optimizer	False	布尔开关	DDP 下分片 optimizer state
HSDP shard size	--hsdp-shard-size	None	正整数	FSDP 下启用 HSDP

DDP 示例：

bash examples/embodied/pi05/run_pi05_ddp_finetune.sh \
    --distributed-strategy ddp \
    --dtype bfloat16

FSDP 示例：

bash examples/embodied/pi05/run_pi05_fsdp_finetune.sh \
    --distributed-strategy fsdp \
    --dtype bfloat16

DDP + ZeRO-1 示例：

bash examples/embodied/pi05/run_pi05_ddp_finetune.sh \
    --distributed-strategy ddp \
    --zero-optimizer

策略选择建议：

场景	建议
模型参数与 optimizer state 可完整放入单卡	DDP
模型可放入单卡，但 optimizer state 显存占用较高	DDP + ZeRO-1
模型、梯度或 optimizer state 难以放入单卡	FSDP
多节点训练，希望参数分片限制在 shard group 内、减少跨节点 FSDP 通信	FSDP + HSDP

4.2.1 DDP / ZeRO 通用参数

以下参数提供对 DDP 策略的通信行为和 ZeRO-1 的精细控制，仅在 --distributed-strategy ddp 时生效：

• DDP 行为：可调节未使用参数检测、静态图优化、bucket 大小与 bucket view，用于减少通信开销或节省内存

• ZeRO-1：开启后分片 optimizer state，可进一步配置 bucket view 和 fp32 master 参数维护

功能	配置项	默认值	取值 / 类型	说明
未使用参数检测	--ddp-find-unused-parameters	True	布尔开关	模型存在条件分支时通常需要保持开启
静态图优化	--ddp-static-graph	False	布尔开关	计算图每步稳定时可开启
bucket view 梯度	--ddp-gradient-as-bucket-view	False	布尔开关	复用 DDP bucket 内存
DDP bucket 大小	--ddp-bucket-cap-mb	None	整数 MB	控制 DDP all-reduce bucket 大小
ZeRO-1	--zero-optimizer	False	布尔开关	分片 optimizer state
ZeRO bucket view	--zero-parameters-as-bucket-view	False	布尔开关	ZeRO 下复用 bucket 内存
ZeRO master 参数	--zero-master-param-dtype	none	none, fp32	是否维护 fp32 master 参数

示例：

bash examples/embodied/pi05/run_pi05_ddp_finetune.sh \
    --distributed-strategy ddp \
    --no-ddp-find-unused-parameters \
    --ddp-static-graph

4.2.2 FSDP 通用参数

以下参数提供对 FSDP 分片、wrap 策略和 dtype 的精细控制，仅在 --distributed-strategy fsdp 时生效：

• 分片与 reshard：控制 forward 后是否立即 reshard，可针对 root group 单独配置

• wrap 策略：可手动指定或排除 FSDP unit 类，也可按参数量阈值自动包装

• dtype 控制：分片前参数 dtype、all-gather 后 dtype 与梯度 reduce dtype 均可独立配置

功能	配置项	默认值	取值 / 类型	说明
HSDP	--hsdp-shard-size	None	正整数	启用二维 mesh 的 shard 维度
默认 reshard 策略	--fsdp-reshard-default	None	true, false, none, 大于 1 的整数	控制 forward 后参数 reshard
root reshard 策略	--fsdp-reshard-root	False	true, false, none, 大于 1 的整数	root FSDP group 的 reshard 策略
指定 wrap 类	--fsdp-wrap-modules	None	逗号分隔模块类名	指定 FSDP unit
排除 wrap 类	--fsdp-no-wrap-modules	None	逗号分隔模块类名	排除指定模块类
自动 wrap 阈值	--fsdp-min-num-params	1000000	非负整数	自动包装重复层的参数阈值
leftover wrap 阈值	--fsdp-leftover-min-num-params	1000000	非负整数	自动包装剩余模块的参数阈值
原始参数 dtype	--fsdp-original-param-dtype	None	fp32, bf16, fp16	FSDP 分片前参数 dtype
unsharded 参数 dtype	--fsdp-unsharded-param-dtype	None	fp32, bf16, fp16	all-gather 后前向/反向 dtype
reduce dtype	--fsdp-reduce-dtype	fp32	fp32, bf16, fp16	梯度 reduce dtype
cast forward inputs	--fsdp-cast-forward-inputs	True	布尔开关	是否将输入 cast 到参数 dtype

4.3 稳定性与运行时控制

训练过程中如遇 loss 异常波动、数值不稳定或显存 OOM，可通过以下手段介入：

• loss spike 防护与 NaN/Inf 检查，保障训练健壮性

• 确定性模式与 TF32 禁用，用于结果复现对齐

• 手动 GC 模式，减少 GC 引入的步间抖动

功能	配置项	默认值	取值 / 类型	说明
loss spike 防护	--loss-spike-threshold	100.0	float	loss 超过阈值或为 NaN/Inf 时，该次 loss 贡献置零
NaN / Inf 检查	--check-for-nan-in-loss-and-grad	True	布尔开关	检查 loss 与 gradient 中的异常值
确定性模式	--deterministic-mode	False	布尔开关	启用确定性算法
禁用 TF32	--disable-tf32	False	布尔开关	禁用 CUDA TF32
手动 GC	--manual-gc	False	布尔开关	关闭自动 GC 并改为显式触发
手动 GC 间隔	--manual-gc-interval	0	非负整数	--manual-gc 启用后每 N 步执行 GC