Skip to content

quick_start_fastwam.md

Latest commit

 

History

History
80 lines (54 loc) · 2.62 KB

quick_start_fastwam.md

File metadata and controls

80 lines (54 loc) · 2.62 KB
Raw

Quick Start: FastWAM Training

This guide walks you through launching a FastWAM SFT (Supervised Fine-Tuning) job in the LoongForge framework.

0. Resource Preparation

0.1 Model Weights

0.1.1 Extract the Action DiT Backbone

Extract and linearly interpolate the Action DiT head weights from the Wan2.2 video DiT checkpoint, producing a .pt file that training can load directly:

LOCAL_MODEL_PATH=/data/models \
OUTPUT=$LOONGFORGE_PATH/checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt \
    bash examples/embodied/fastwam/preprocess_action_dit_backbone.sh

0.1.2 Wan-AI/Wan2.2-TI2V-5B Weights

hf download Wan-AI/Wan2.2-TI2V-5B --local-dir /workspace/huggingface.co/Wan-AI/Wan2.2-TI2V-5B

0.2 Tokenizer / Processor

Encode all task instructions in the dataset into text embeds in one pass, then read them on demand during training:

DATASET_PATH=/data/libero \
TEXT_EMBEDDING_CACHE_DIR=/data/cache/fastwam_text_embeds \
    bash examples/embodied/fastwam/precompute_text_embeds.sh

0.3 Dataset

hf download yuanty/LIBERO-fastwam --local-dir /workspace/data/LIBERO-fastwam
cd /workspace/data/LIBERO-fastwam
tar -zxvf libero_spatial_no_noops_lerobot.tar.gz

1. Data Configuration

Data-related fields are configured mainly through YAML and environment variables.

Model and data YAML:

model:
  action_dit_pretrained_path: checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt  # Path to the action dit generated above
  redirect_common_files: true
  dtype: bfloat16

data:
  text_embedding_cache_dir: data/text_embeds_cache/libero  # Path to the text embeds generated above

Environment variables:

export LOONGFORGE_PATH=/workspace/LoongForge                                                    # Repository root
export DATASET_PATH=/workspace/data/LIBERO-fastwam/libero_spatial_no_noops_lerobot              # LeRobot dataset root
export DIFFSYNTH_MODEL_BASE_PATH=/workspace/huggingface.co/Wan-AI/Wan2.2-TI2V-5B                # Wan2.2 5B base model path

2. Launch Training

2.1 Launch Script

Single-node DDP training example:

bash examples/embodied/fastwam/run_fastwam_sft_ddp.sh

2.2 Correctness Verification

To ensure that training accuracy is not affected by our optimizations, we ran a step-by-step action loss comparison between LoongForge's FastWAM and the official implementation under identical data, weights, and training configurations. The results show that LoongForge's performance optimizations are lossless with respect to training accuracy: alt text