环境搭建
- git 下载
bash
git clone https://github.com/yuantianyuan01/FastWAM.git- dockerfile 搭建
bash
ARG TARGETPLATFORM
ARG BASE_IMAGE=nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
FROM ${BASE_IMAGE}
ARG TARGETPLATFORM
ENV DEBIAN_FRONTEND=noninteractive
SHELL ["/bin/bash", "-lc"]
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
ca-certificates \
curl \
ffmpeg \
git \
libaio-dev \
libavcodec-dev \
libavformat-dev \
libavutil-dev \
libnuma-dev \
libswscale-dev \
pkg-config \
wget \
&& rm -rf /var/lib/apt/lists/*
ARG CONDA_DIR=/opt/conda
RUN if [[ "${TARGETPLATFORM:-linux/amd64}" == "linux/arm64" ]]; then \
installer_url="https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh"; \
else \
installer_url="https://repo.anaconda.com/miniconda/Miniconda3-py310_24.7.1-0-Linux-x86_64.sh"; \
fi && \
curl -fsSL "${installer_url}" -o /tmp/conda.sh && \
bash /tmp/conda.sh -b -p "${CONDA_DIR}" && \
rm -f /tmp/conda.sh && \
"${CONDA_DIR}/bin/conda" clean -afy
ENV PATH=/opt/conda/bin:$PATH
RUN conda create -n fastwam python=3.10 -y && conda clean -afy
ENV CONDA_DEFAULT_ENV=fastwam
ENV PATH=/opt/conda/envs/fastwam/bin:/opt/conda/bin:$PATH
RUN python -m pip install -U pip && \
python -m pip install torch==2.7.1+cu128 torchvision==0.22.1+cu128 --extra-index-url https://download.pytorch.org/whl/cu128
WORKDIR /workspace/FastWAM
COPY . .
RUN pip install -e .
CMD ["bash"]
bash
docker build -t fastwam .文件下载
- 下载模型
bash
docker run --gpus all -it --rm \
--network host \
--shm-size=64g \
-v $(pwd):/workspace \
fastwam
mkdir -p checkpoints
export DIFFSYNTH_MODEL_BASE_PATH="$(pwd)/checkpoints"
python scripts/preprocess_action_dit_backbone.py \
--model-config configs/model/fastwam.yaml \
--output checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt \
--device cuda \
--dtype bfloat16如果已经下载了，可以软链接：
bash
ln -s /workspace/checkpoints/Wan2.2-TI2V-5B /workspace/checkpoints/Wan-AI/Wan2.2-TI2V-5B- 下载数据集
bash
mkdir -p data/libero_mujoco3.3.2
cd data/libero_mujoco3.3.2
# Run after downloading all 4 tar.gz files
for f in *.tar.gz; do
tar -xzf "$f"
done
text
data/libero_mujoco3.3.2/
├── libero_10_no_noops_lerobot/
├── libero_goal_no_noops_lerobot/
├── libero_object_no_noops_lerobot/
└── libero_spatial_no_noops_lerobot/训练
1) Precompute T5 embedding cache
bash
# LIBERO
python scripts/precompute_text_embeds.py task=libero_uncond_2cam224_1e-4- 生成的目录结构
2) Training
- 运行命令
bash
bash scripts/train_zero1.sh 4 task=libero_uncond_2cam224_1e-4- 结果
bash
bash scripts/train_zero1.sh 4 task=libero_uncond_2cam224_1e-4
[launch] nproc_per_node=4 num_machines=1 machine_rank=0 run_id=2026-08-04_05-39-28
df: /root/.triton/autotune: No such file or directory
W0804 05:39:36.255000 670 site-packages/torch/distributed/run.py:766]
W0804 05:39:36.255000 670 site-packages/torch/distributed/run.py:766] *****************************************
W0804 05:39:36.255000 670 site-packages/torch/distributed/run.py:766] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W0804 05:39:36.255000 670 site-packages/torch/distributed/run.py:766] *****************************************
08/04 [05:39:40] INFO | >> Loading Wan2.2-TI2V-5B components... loader.py:152
08/04 [05:40:27] INFO | >> Skipping pretrained text encoder/tokenizer loader.py:206
load (`load_text_encoder=False`); training must
provide cached `context/context_mask`.
08/04 [05:40:32] INFO | >> Finished loading Wan2.2-TI2V-5B components in loader.py:211
52.39 seconds.
08/04 [05:40:40] INFO | >> ActionDiT backbone payload policy: action_dit.py:154
{'skip_prefixes': ['action_encoder.', 'head.'],
'alpha_scaling': True, 'interpolation':
'sequential_1d_linear_align_corners_true'}
INFO | >> Loaded ActionDiT backbone from action_dit.py:218
/workspace/FastWAM/checkpoints/ActionDiT_linear
_interp_Wan22_alphascale_1024hdim.pt (keys=820;
random_kept_prefixes=['action_encoder.',
'head.']).
INFO | >> Initialized MoT with experts: ['video', 'action'], mot.py:53
num_layers=30
INFO | >> Expert 'video': num_params=5.00 B mot.py:56
INFO | >> Expert 'action': num_params=1.02 B mot.py:56
08/04 [05:40:41] INFO | >> PyTorch version 2.7.1+cu128 available. config.py:54
08/04 [05:40:42] INFO | >> PyTorch version 2.7.1+cu128 available. config.py:54
08/04 [05:40:42] INFO | >> PyTorch version 2.7.1+cu128 available. config.py:54
08/04 [05:40:42] INFO | >> PyTorch version 2.7.1+cu128 available. config.py:54
Resolving data files: 100%|███████████████████████████| 433/433 [00:00<00:00, 29540.23it/s]
Resolving data files: 100%|███████████████████████████| 433/433 [00:00<00:00, 28070.51it/s]
Downloading data: 100%|████████████████████████████| 433/433 [00:00<00:00, 34913.56files/s]
Resolving data files: 100%|███████████████████████████| 433/433 [00:00<00:00, 50216.60it/s]
Resolving data files: 100%|███████████████████████████| 433/433 [00:00<00:00, 27117.40it/s]
Generating train split: 52895 examples [00:01, 44618.50 examples/s]
08/04 [05:40:45] INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmph8dpzzn9/test.c -o /tmp/tmph8dpzzn9/test.o
08/04 [05:40:46] INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmph8dpzzn9/test.o -laio -o
/tmp/tmph8dpzzn9/a.out
08/04 [05:40:46] INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpea71k0hp/test.c -o /tmp/tmpea71k0hp/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpea71k0hp/test.o -laio -o
/tmp/tmpea71k0hp/a.out
08/04 [05:40:46] INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmp7rb9ub6y/test.c -o /tmp/tmp7rb9ub6y/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmp7rb9ub6y/test.o -laio -o
/tmp/tmp7rb9ub6y/a.out
08/04 [05:40:46] INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmp3num8eul/test.c -o /tmp/tmp3num8eul/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmp3num8eul/test.o -laio -o
/tmp/tmp3num8eul/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpwn9ssry6/test.c -o /tmp/tmpwn9ssry6/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpwn9ssry6/test.o -L/usr/local/cuda
-L/usr/local/cuda/lib64 -lcufile -o
/tmp/tmpwn9ssry6/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpohaf4s4_/test.c -o /tmp/tmpohaf4s4_/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmplyuxh6ee/test.c -o /tmp/tmplyuxh6ee/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpohaf4s4_/test.o -L/usr/local/cuda
-L/usr/local/cuda/lib64 -lcufile -o
/tmp/tmpohaf4s4_/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmplyuxh6ee/test.o -laio -o
/tmp/tmplyuxh6ee/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpyji8_voi/test.c -o /tmp/tmpyji8_voi/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpr8knnynf/test.c -o /tmp/tmpr8knnynf/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpyji8_voi/test.o -L/usr/local/cuda
-L/usr/local/cuda/lib64 -lcufile -o
/tmp/tmpyji8_voi/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpq1fu50_7/test.c -o /tmp/tmpq1fu50_7/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpr8knnynf/test.o -laio -o
/tmp/tmpr8knnynf/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpq1fu50_7/test.o -L/usr/local/cuda
-L/usr/local/cuda/lib64 -lcufile -o
/tmp/tmpq1fu50_7/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpvj36ez9o/test.c -o /tmp/tmpvj36ez9o/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpvj36ez9o/test.o -laio -o
/tmp/tmpvj36ez9o/a.out
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
-Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv
-O2 -Wall -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -O2 -isystem
/opt/conda/envs/fastwam/include -fPIC -c
/tmp/tmpbo2_rvx9/test.c -o /tmp/tmpbo2_rvx9/test.o
INFO | >> gcc -pthread -B spawn.py:75
/opt/conda/envs/fastwam/compiler_compat
/tmp/tmpbo2_rvx9/test.o -laio -o
/tmp/tmpbo2_rvx9/a.out
08/04 [05:40:47] INFO | >> Calculating dataset stats for robot_video_dataset.py:92
normalization...
Iterating dataset to get normalization: 100%|███████████| 433/433 [00:04<00:00, 101.84it/s]
08/04 [05:40:52] INFO | >> Accelerate training: distributed_type=DistributedType.DEEPSPEED trainer.py:65
zero_stage=1 world_size=4 process_index=0 cfg_mixed_precision=bf16
accelerator_mixed_precision=bf16 grad_accum=1 grad_clip=1.0000
INFO | >> using accelerator.device=cuda:0 trainer.py:76
Before initializing optimizer states
MA 18.15 GB Max_MA 20.95 GB CA 20.98 GB Max_CA 21 GB
CPU Virtual Memory: used = 95.94 GB, percent = 9.5%
After initializing optimizer states
MA 18.15 GB Max_MA 23.76 GB CA 26.59 GB Max_CA 27 GB
CPU Virtual Memory: used = 96.91 GB, percent = 9.6%
After initializing ZeRO optimizer
MA 18.15 GB Max_MA 18.15 GB CA 26.59 GB Max_CA 27 GB
CPU Virtual Memory: used = 98.12 GB, percent = 9.7%
08/04 [05:41:10] INFO | >> Train/val dataset size: 52895/52895 trainer.py:129
INFO | >> Setting DiT to train mode and freezing other model components. trainer.py:282
INFO | >> Starting training with max_steps=8270. trainer.py:654
08/04 [05:41:48] INFO | >> epoch=0 step=10/8270 loss=1.7786 loss_action=1.3374 trainer.py:715
loss_video=0.4412 lr=2.66e-06 speed=0.26 step/s, 16.86 samples/s
eta=08:42:39
08/04 [05:42:09] INFO | >> epoch=0 step=20/8270 loss=1.1308 loss_action=0.7775 trainer.py:715