support dreamzero

examples/dreamzero/README.md

@@ -0,0 +1,162 @@
+# DreamZero: Training
+
+This guide covers how to train DreamZero 14B using FlagScale with FSDP2 distributed training.
+
+## Overview
+
+DreamZero is a World Action Model — a Wan2.1 14B video DiT repurposed as a zero-shot robot policy. Architecture:
+- **T5-XXL** text encoder (~4.7B, frozen)
+- **CLIP** image encoder (~0.6B, frozen)
+- **VAE** video encoder (~0.1B, frozen)
+- **CausalWan DiT** video diffusion transformer (~14B, trainable)
+- Action/state encoder-decoder projectors (trainable)
+
+FlagScale uses FSDP2 (ZeRO-3 style sharding) to distribute the 14B trainable DiT across GPUs.
+
+## Installation
+
+### Clone Repository
+
+```sh
+git clone https://github.com/FlagOpen/FlagScale.git
+cd FlagScale/
+```
+
+### Setup Environment
+
+Create a conda environment with PyTorch 2.4+ and FSDP2 support:
+
+```sh
+conda create -n flagscale-dreamzero python=3.11
+conda activate flagscale-dreamzero
+pip install ".[cuda-train]" --verbose
+```
+
+Install additional dependencies:
+
+```sh
+pip install safetensors peft decord
+```
+
+DreamZero also requires the reference `groot` package for model architecture and data transforms:
+
+```sh
+git clone <dreamzero-reference-repo> /workspace/dreamzero
+```
+
+Add both FlagScale and the reference repo to your PYTHONPATH:
+
+```sh
+export PYTHONPATH=/path/to/FlagScale:/workspace/dreamzero:$PYTHONPATH
+```
+
+## Download Models
+
+Download the DreamZero-AgiBot checkpoint (contains `config.json` + safetensors shards):
+
+```sh
+# Place at /workspace/models/DreamZero-AgiBot/
+# Expected contents:
+#   config.json
+#   model.safetensors.index.json
+#   model-00001-of-00010.safetensors
+#   ...
+#   model-00010-of-00010.safetensors
+```
+
+Download the Wan2.1-I2V-14B-480P pretrained weights (for tokenizer, and optionally T5/CLIP/VAE):
+
+```sh
+# Place at /workspace/models/Wan2.1-I2V-14B-480P/
+# Required for tokenizer:
+#   google/umt5-xxl/tokenizer.json
+#   google/umt5-xxl/spiece.model
+```
+
+## Training
+
+### Prepare Dataset
+
+DreamZero uses a LeRobot-format dataset with `data/`, `meta/`, and `videos/` subdirectories.
+
+```sh
+# Place at /workspace/datasets/vla_arena_dreamzero/
+# Expected structure:
+#   data/chunk-000/, data/chunk-001/, ...
+#   meta/info.json, meta/episodes.jsonl, ...
+#   videos/chunk-000/, ...
+```
+
+### Edit Config
+
+FlagScale uses a two-level configuration system:
+
+1. **Experiment-level config** (`examples/dreamzero/conf/train.yaml`): Experiment settings, environment variables, and resource allocation
+2. **Task-level config** (`examples/dreamzero/conf/train/dreamzero_14b.yaml`): Model, dataset, and training hyperparameters
+
+#### Experiment-Level Config
+
+```sh
+vim examples/dreamzero/conf/train.yaml
+```
+
+Configure the following fields:
+
+- `experiment.exp_name` — Experiment name
+- `experiment.exp_dir` — Output directory for checkpoints and logs
+- `experiment.envs.CUDA_VISIBLE_DEVICES` — GPU devices to use (e.g., `"0,1,2,3,4,5,6,7"`)
+- `experiment.runner.nproc_per_node` — Number of GPUs
+
+#### Task-Level Config
+
+```sh
+vim examples/dreamzero/conf/train/dreamzero_14b.yaml
+```
+
+Configure the following fields:
+
+**System settings:**
+- `system.batch_size` — Per-GPU micro batch size (default: `1`)
+- `system.train_steps` — Total training steps (default: `5000`)
+- `system.checkpoint.save_freq` — Steps between checkpoints
+
+**Model settings:**
+- `model.pretrained_model_path` — Path to DreamZero-AgiBot checkpoint (e.g., `/workspace/models/DreamZero-AgiBot`)
+- `model.train_architecture` — `"full"` for full fine-tuning (recommended), `"lora"` for LoRA on DiT
+- `model.optimizer.lr` — Learning rate (default: `1.0e-5`)
+- `model.optimizer.betas` — Adam betas (default: `[0.95, 0.999]`)
+- `model.optimizer.scheduler.warmup_ratio` — Warmup ratio (default: `0.05`)
+
+**Data settings:**
+- `data.data_path` — Path to LeRobot dataset (e.g., `/workspace/datasets/vla_arena_dreamzero`)
+- `data.tokenizer_path` — Path to UMT5-XXL tokenizer (e.g., `/workspace/models/Wan2.1-I2V-14B-480P/google/umt5-xxl`)
+- `data.embodiment_tag` — Embodiment tag (e.g., `"libero"`)
+
+### Start Training
+
+```sh
+cd FlagScale/
+python flagscale/run.py --config-path examples/dreamzero/conf --config-name train
+```
+
+Training logs are saved to `outputs/dreamzero_train/logs/host_0_localhost.output` by default.
+
+### Stop Training
+
+```sh
+cd FlagScale/
+python flagscale/run.py --config-path examples/dreamzero/conf --config-name train action=stop
+```
+
+## Hardware Requirements
+
+With 8x H100/H800 80GB GPUs and `batch_size=1`:
+- Full fine-tuning (FSDP2): ~65 GB per GPU
+- Model loading takes ~30 minutes (14B params from 10 safetensors shards)
+- Training throughput: ~8.7s per step (effective batch size 8)
+
+## Known Issues
+
+1. **torch.compile + FSDP2**: Compiling attention sub-methods with `mode="reduce-overhead"` causes CUDA graph tensor deallocation mismatches during backward. Disabled by default via `system.disable_attention_compile: true`.
+2. **LoRA mode NaN**: `train_architecture: lora` can produce NaN due to bf16 backward overflow at certain timestep/data combinations. Use `train_architecture: full` to avoid this.
+3. **batch_size=2 OOM**: Full fine-tuning with `batch_size=2` exceeds 80GB GPU memory. Use `batch_size=1` with gradient accumulation if larger effective batches are needed.

examples/dreamzero/conf/train.yaml

@@ -0,0 +1,39 @@
+defaults:
+  - _self_
+  - train: dreamzero_14b
+
+experiment:
+  exp_name: dreamzero_train
+  seed: 42
+  save_steps: 500
+  load: null
+  exp_dir: outputs/${experiment.exp_name}
+  ckpt_format: torch
+  task:
+    type: train
+    backend: native
+    entrypoint: flagscale/train/train_dreamzero.py
+  runner:
+    per_node_task: false
+    no_shared_fs: false
+    rdzv_backend: static
+    hostfile: null
+    nproc_per_node: 8
+  cmds:
+    before_start: echo "Starting DreamZero Training"
+  envs:
+    LOGLEVEL: "INFO"
+    CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7"
+    CUDA_DEVICE_MAX_CONNECTIONS: 1
+    WANDB_MODE: offline
+    OTEL_SDK_DISABLED: true
+    TOKENIZERS_PARALLELISM: false
+    NCCL_TIMEOUT: "1800"
+    TORCH_NCCL_ASYNC_ERROR_HANDLING: "1"
+    PYTHONUNBUFFERED: "1"
+
+action: run
+
+hydra:
+  run:
+    dir: ${experiment.exp_dir}/hydra

examples/dreamzero/conf/train/dreamzero_14b.yaml

@@ -0,0 +1,92 @@
+# DreamZero 14B Training Config for FlagScale
+# LIBERO fine-tuning with full DIT fine-tuning via FSDP2 on 8xH100/H800
+
+system:
+  train_steps: 5000
+  log_freq: 10
+  grad_clip_norm: 1.0
+  use_amp: true
+  shuffle: true
+  num_workers: 4
+  batch_size: 1
+  # Disable torch.compile on attention sub-methods — CUDA graph replay
+  # crashes with FSDP2 reshard-after-forward (tensor deallocation mismatch).
+  disable_attention_compile: true
+
+  checkpoint:
+    save_freq: 500
+    output_directory: ${experiment.exp_dir}
+    save_total_limit: 10
+
+model:
+  model_name: dreamzero
+  checkpoint_dir: null
+  # Path to DreamZero-AgiBot checkpoint (contains config.json + safetensors shards)
+  pretrained_model_path: /workspace/models/DreamZero-AgiBot
+  # Optional: explicit paths for T5/CLIP/VAE (null = loaded from checkpoint)
+  text_encoder_pretrained_path: null
+  image_encoder_pretrained_path: null
+  vae_pretrained_path: null
+  compute_dtype: bfloat16
+  use_gradient_checkpointing: true
+
+  # Training mode: "full" (full fine-tuning) or "lora" (LoRA on DIT)
+  # "full" recommended — avoids LoRA-on-fine-tuned-weights NaN issue.
+  train_architecture: full
+  # LoRA config (only used if train_architecture=lora)
+  lora_rank: 4
+  lora_alpha: 4
+  lora_target_modules: "q,k,v,o,ffn.0,ffn.2"
+
+  # Dimensions matching LIBERO data
+  num_frames: 33
+  action_horizon: 24
+  state_horizon: 1
+  action_dim: 32
+  max_state_dim: 64
+
+  # Frame/block settings
+  num_frame_per_block: 2
+  num_action_per_block: 24
+  num_state_per_block: 1
+  # frame_seqlen: 2 views tiled to 352x640 → VAE 8x → 44x80 → patch(1,2,2) → 22x40 = 880
+  frame_seqlen: 880
+  num_views: 2
+
+  optimizer:
+    name: AdamW
+    lr: 1.0e-5
+    weight_decay: 1.0e-5
+    betas: [0.95, 0.999]
+    eps: 1.0e-8
+    scheduler:
+      name: cosine
+      warmup_ratio: 0.05
+      scheduler_kwargs: null
+
+data:
+  # Path to LeRobot-format dataset (with data/, meta/, videos/ subdirs)
+  data_path: /workspace/datasets/vla_arena_dreamzero
+  # Path to UMT5-XXL tokenizer directory
+  tokenizer_path: /workspace/models/Wan2.1-I2V-14B-480P/google/umt5-xxl
+  # Per-view resolution: 176x320 (two views tiled vertically: 352x640 before VAE encoding)
+  image_size: [176, 320]
+  max_text_length: 512
+  embodiment_tag: libero
+  # Native multi-chunk dataloader
+  use_aligned_dataloader: false
+  max_chunk_size: 4
+  macro_stride: 24
+  relative_action: true
+  crop_ratio: 0.95
+  color_jitter: true
+  # Embodiment ID mapping
+  embodiment_tag_mapping:
+    agibot: 26
+    oxe_droid: 17
+    bridge: 1
+    libero: 17
+    gr1_unified: 24
+    mecka_hands: 27
+    xdof: 22
+    yam: 31