[Train] Support FSDP2 checkpoint resume (#1149)

### PR Category
<!-- One of [ Train | Inference | Compress | Serve | RL | Core |
Hardware | CICD | Tools | Others ] -->
Train
### PR Types
<!-- One of [ User Experience | New Features | Bug Fixes | Improvements
| Performance | Breaking Change| Deprecations | Test Case | Docs |
Others ] -->
Improvements
### PR Description
<!-- Describe what you’ve done -->
  - Refactor to use standard `from_pretrained`, `save_pretrained` api
- Reorganize qwen_gr00t into proper subpackage (qwen_gr00t/,
vlm/qwenvl_backbone)
- Add FSDP2 checkpoint resume support (model, optimizer, scheduler, RNG
state)
  - vla serving script now supports arbitrary observation key names

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: legitnull

.github/workflows/unit_tests_common.yml

@@ -168,7 +168,7 @@ jobs:
             --retry-count 3
 
           # TODO: temp solution to install newly added deps, remove once the new image is built
-          pip install qwen_vl_utils==0.0.14 diffusers==0.36.0 websocket-client==1.8.0 websocket==0.2.1 websockets==15.0.1 msgpack==1.1.0
+          pip install qwen_vl_utils==0.0.14 diffusers==0.36.0 websocket-client==1.8.0 websocket==0.2.1 websockets==15.0.1 msgpack==1.1.0 datasets==4.5.0
 
           # Copy test data (keep existing logic)
           mkdir -p /opt/data

examples/gr00t_n1_5/conf/train/gr00t_n1_5.yaml

@@ -11,13 +11,15 @@ system:
     output_directory: ${experiment.exp_dir}
     save_checkpoint: true
     save_freq: 1000
+    # Path to a checkpoint directory to resume training from (e.g. outputs/gr00t_n1_5_train/checkpoints/001000)
+    # resume_from:
 
 model:
   model_name: gr00t_n1_5
   # Path or HuggingFace model ID for the pretrained GR00T N1.5 model
   checkpoint_dir: /workspace/models/nvidia/GR00T-N1.5-3B
 
-  # Fine-tuning control (from groot_finetune_script.py)
+  # Fine-tuning control
   tune_llm: true
   tune_visual: true
   tune_projector: true
@@ -27,6 +29,39 @@ model:
   # Embodiment tag to use for training (e.g. 'new_embodiment', 'gr1')
   embodiment_tag: new_embodiment
 
+  # Number of future action steps predicted per forward pass.
+  # Determines action_delta_indices = [0, 1, ..., chunk_size - 1].
+  chunk_size: 16
+
+  # Padding dimensions — shorter state/action sequences are zero-padded to these sizes
+  max_state_dim: 64
+  max_action_dim: 32
+
+  normalization_mapping:
+    VISUAL: IDENTITY
+    STATE: MIN_MAX
+    ACTION: MIN_MAX
+
+  # LoRA fine-tuning (lora_rank: 0 disables LoRA)
+  # lora_rank: 0
+  # lora_alpha: 16
+  # lora_dropout: 0.1
+  # lora_full_model: false
+
+  # ============================================================
+  # Module Freezing Configuration
+  # ============================================================
+  # Freezing logic: freeze_patterns are applied first, then keep_patterns override.
+  # Patterns are regex matched against full parameter names.
+  #
+  # Common patterns for GR00T N1.5:
+  #   - "_groot_model\\.backbone\\..*"         # Entire backbone (VLM + vision)
+  #   - "_groot_model\\.action_head\\..*"       # Action diffusion head
+  #
+  # freeze:
+  #   freeze_patterns:
+  #     - "_groot_model\\.backbone\\..*"
+
   optimizer:
     name: AdamW
     lr: 1.0e-4
@@ -42,10 +77,8 @@ model:
 
 data:
   dataset_type: lerobot
-  data_path: /workspace/IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot/
+  data_path: /workspace/datasets/IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot/
   tolerance_s: 0.0001
-  observation_delta_indices: null
-  action_delta_indices: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
   preprocessor:
     name: policy_preprocessor
     steps:

examples/pi0/conf/train/pi0.yaml

@@ -22,7 +22,6 @@ model:
   # Path to paligemma tokenizer
   tokenizer_path: /workspace/models/google/paligemma-3b-pt-224
   tokenizer_max_length: 48
-  action_steps: 50
 
   optimizer:
     name: AdamW

examples/pi0_5/conf/serve.yaml

@@ -7,11 +7,11 @@ experiment:
   exp_dir: outputs/${experiment.exp_name}
   task:
     type: serve
-    entrypoint: flagscale/serve/run_serve_pi.py
+    entrypoint: flagscale/serve/run_serve_qwen_gr00t.py
   runner:
     hostfile: null
     deploy:
-      use_fs_serve: false
+      use_fs_serve: true
   envs:
     CUDA_VISIBLE_DEVICES: 0
     CUDA_DEVICE_MAX_CONNECTIONS: 1

examples/pi0_5/conf/serve/pi0_5.yaml

@@ -2,18 +2,13 @@
   engine_args:
     host: 0.0.0.0
     port: 5000
-    # Model variant: "pi0" or "pi0.5"
     model_variant: "pi0.5"
-    model: /workspace/models/lerobot/pi05_base
-    tokenizer: /workspace/models/google/paligemma-3b-pt-224
-    stat_path: /workspace/datasets/lerobot/aloha_mobile_cabinet/meta/stats.json
+    model: /workspace/models/pi0_5_train/checkpoints/last/pretrained_model
     device: "cuda"
-    # By default, Pi0.5 uses quantiles for state and action normalization, if false, it uses mean and std instead
-    use_quantiles: false
-    images_keys:
-      - observation.images.base_0_rgb
-      - observation.images.left_wrist_0_rgb
-      - observation.images.right_wrist_0_rgb
-    # Only used for warmup
-    images_shape: [3, 480, 640]
-    state_key: observation.state
+    # Maps client-sent observation keys to the keys the model was trained with.
+    # Format: {key_from_client: key_expected_by_model}
+    # rename_map:
+    #   "observation/image": "observation.images.image"
+    #   "observation/wrist_image": "observation.images.wrist_image"
+    #   "observation/state": "observation.state"
+    #   "prompt": "task"

examples/pi0_5/conf/train.yaml

@@ -22,7 +22,7 @@ experiment:
     before_start: echo "Starting PI0.5 Training"
   envs:
     LOGLEVEL: "INFO"
-    CUDA_VISIBLE_DEVICES: "0,1"
+    CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7"
     CUDA_DEVICE_MAX_CONNECTIONS: 1
     WANDB_MODE: offline
     OTEL_SDK_DISABLED: true

examples/pi0_5/conf/train/pi0_5.yaml

@@ -1,10 +1,10 @@
 system:
-  batch_size: 1
-  train_steps: 100000
+  batch_size: 16
+  train_steps: 30000
   log_freq: 10
   grad_clip_norm: 1.0
-  use_amp: false
-  shuffle: false
+  use_amp: true
+  shuffle: true
   num_workers: 4
 
   checkpoint:
@@ -18,11 +18,12 @@ system:
 model:
   model_name: pi0.5
   # Path to the pretrained pi05_base model checkpoint
-  checkpoint_dir: /workspace/models/lerobot/pi05_base
+  checkpoint_dir: /workspace/models/lerobot/pi05_libero_base
   # Path to paligemma tokenizer
   tokenizer_path: /workspace/models/google/paligemma-3b-pt-224
   tokenizer_max_length: 200
-  action_steps: 50
+  gradient_checkpointing: true
+  freeze_vision_encoder: false
 
   optimizer:
     name: AdamW
@@ -37,14 +38,12 @@ model:
 
 data:
   # Path to the training data
-  data_path: /workspace/datasets/lerobot/aloha_mobile_cabinet
+  data_path: /workspace/datasets/IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot
   tolerance_s: 0.0001
   use_imagenet_stats: true
   # To match the input features naming from the dataset to the policy config
   # For example, for the aloha_mobile_cabinet dataset, the rename_map is:
   rename_map:
-    observation.images.cam_high: observation.images.base_0_rgb
-    observation.images.cam_left_wrist: observation.images.left_wrist_0_rgb
-    observation.images.cam_right_wrist: observation.images.right_wrist_0_rgb
+    "observation.images.wrist_image": "observation.images.image2"
   # By default, Pi0.5 uses quantiles for state and action normalization, if false, it uses mean and std instead
-  use_quantiles: false
+  use_quantiles: true

examples/qwen_gr00t/conf/serve/qwen_gr00t.yaml

@@ -3,5 +3,17 @@
     host: 0.0.0.0
     port: 5000
     model_variant: QwenGr00t
-    model: /workspace/models/qwen_gr00t_train/checkpoints/last
+    model: /share/project/fengyupu/github/FlagScale_2/outputs/260320_qwen_gr00t_train_libero_goal_old_dataset/checkpoints/last
     device: "cuda"
+    # Maps client-sent observation keys to the keys the model was trained with.
+    # Format: {key_from_client: key_expected_by_model}
+    rename_map:
+      "observation/image": "observation.images.image"
+      "observation/wrist_image": "observation.images.wrist_image"
+      "observation/state": "observation.state"
+      "prompt": "task"
+    serve_preprocessor:
+      steps:
+        - registry_name: image_resize_processor
+          config:
+            image_size: [224, 224]

examples/qwen_gr00t/conf/train/qwen_gr00t.yaml

@@ -1,47 +1,42 @@
 system:
-  batch_size: 16
+  batch_size: 8
   train_steps: 30000
   log_freq: 1
   grad_clip_norm: 1.0
   use_amp: true
   shuffle: true
   num_workers: 4
+  # Weight applied to VLM language modelling loss when co-training with vlm_data_path.
+  # Set to 0 or omit vlm_data_path to disable co-training.
+  vlm_loss_scale: 0.1
 
   checkpoint:
     output_directory: ${experiment.exp_dir}
     # Whether to save checkpoint
     save_checkpoint: true
     # Number of steps between checkpoints
-    save_freq: 1000
-  # TODO(yupu): Support resuming from checkpoint
+    save_freq: 10000
+    # Path to a checkpoint directory to resume training from (e.g. /path/to/checkpoints/005000)
+    # resume_from:
 
 model:
-  # TODO: (yupu) the config layout is still a mess
   model_name: qwen_gr00t
-  # Path to the checkpoint of the pretrained base VLM model, e.g. Qwen3-VL-4B-Instruct
-  checkpoint_dir: /workspace/models/Qwen/Qwen3-VL-4B-Instruct/
   vlm:
     type: qwen3-vl
-  qwenvl:
     base_vlm: /workspace/models/Qwen/Qwen3-VL-4B-Instruct/
     attn_implementation: flash_attention_2
-    vl_hidden_dim: 2048
-  dino:
-    dino_backbone: dinov2_vits14
   action_model:
     # Whether to condition the action model on proprioceptive state (observation.state)
     use_state: false
-    type: flow_matching
+    type: gr00t_action_head
     action_model_type: DiT-B
-    action_hidden_dim: 1024
     hidden_size: 1024
     add_pos_embed: True
     max_seq_len: 1024
     action_dim: 7
     state_dim: 7
     future_action_window_size: 7
     action_horizon: 8
-    past_action_window_size: 0
     repeated_diffusion_steps: 4
     noise_beta_alpha: 1.5
     noise_beta_beta: 1.0
@@ -58,7 +53,13 @@ model:
       num_layers: 16
       output_dim: 1024
       positional_embeddings: None
-  reduce_in_full_precision: True
+
+  prompt_template: "Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format."
+
+  normalization_mapping:
+    VISUAL: IDENTITY
+    STATE: MIN_MAX
+    ACTION: MIN_MAX
 
   optimizer:
     name: AdamW
@@ -117,24 +118,11 @@ data:
     vision_root: ""
     action_key: eepose
     state_key: eepose
-  # TODO: (yupu) Remove this once we have a proper dataset config
-  vla_data:
-    dataset_py: lerobot_datasets
-    data_root_dir: playground/Datasets/
-    data_mix: libero_goal_old
-    action_type: delta_qpos
-    CoT_prompt: Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format.
-    CoT_answer: bbox
-    default_image_resolution: [3, 224, 224]
-    load_all_data_for_training: True
-    obs: ["image_0"]
-    video_backend: torchvision_av
   # Path to the training data
   data_path: /workspace/datasets/IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot/
+  # Path to VLM co-training data (WDS/Energon format). Leave unset to disable co-training.
+  # vlm_data_path: /workspace/datasets/vlm_cotrain/
   tolerance_s: 0.0001
-  # TODO: (yupu) I think these indices should belong to the policy config, maybe put it in the model config?
-  observation_delta_indices: [0]
-  action_delta_indices: [0,1,2,3,4,5,6,7]
   preprocessor:
     name: policy_preprocessor
     steps:
@@ -151,18 +139,12 @@ data:
         config:
           eps: 1e-8
           features: {}
-          norm_map:
-            VISUAL: IDENTITY
-            STATE: MIN_MAX
-            ACTION: MIN_MAX
+          # norm_map is injected at runtime from model.normalization_mapping
   postprocessor:
     name: policy_postprocessor
     steps:
       - registry_name: unnormalizer_processor
         config:
           eps: 1e-8
           features: {}
-          norm_map:
-            VISUAL: IDENTITY
-            STATE: MIN_MAX
-            ACTION: MIN_MAX
+          # norm_map is injected at runtime from model.normalization_mapping