From 35399d187ae9594d46655fabffb944950e190e4b Mon Sep 17 00:00:00 2001 From: Lufang Chen <64068400+vincentccc@users.noreply.github.com> Date: Thu, 16 Oct 2025 10:53:51 +0800 Subject: [PATCH] Update train from QwenVL (#50) * update from vlm * update * update * update --- .gitattributes | 1 + wall_x/data/load_lerobot_dataset.py | 39 +++--- wall_x/trainer/qwen_vl_act_trainer.py | 101 +++++++++++---- workspace/README.md | 19 ++- workspace/lerobot_example/config_qact.yml | 2 +- .../lerobot_example/config_qact_from_vlm.yml | 117 ++++++++++++++++++ .../evaluation/lerobot_openloop.png | 3 + workspace/lerobot_example/qwen25_config.json | 106 ++++++++++++++++ 8 files changed, 344 insertions(+), 44 deletions(-) create mode 100644 .gitattributes create mode 100644 workspace/lerobot_example/config_qact_from_vlm.yml create mode 100644 workspace/lerobot_example/evaluation/lerobot_openloop.png create mode 100644 workspace/lerobot_example/qwen25_config.json diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..f8f1794 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +workspace/lerobot_example/evaluation/lerobot_openloop.png filter=lfs diff=lfs merge=lfs -text diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py index f5ce2fb..646831b 100644 --- a/wall_x/data/load_lerobot_dataset.py +++ b/wall_x/data/load_lerobot_dataset.py @@ -255,14 +255,6 @@ class DataCollator: processor_path = self.config["pretrained_wallx_path"] action_tokenizer_path = self.config["action_tokenizer_path"] - # Use cached processors if available - if processor_path not in self._processor_cache: - self._processor_cache[processor_path] = AutoProcessor.from_pretrained( - processor_path, use_fast=True - ) - if self.config.get("padding_side", "left") == "left": - self._processor_cache[processor_path].tokenizer.padding_side = "left" - if ( self.use_fast_tokenizer and action_tokenizer_path not in self._action_tokenizer_cache @@ -273,6 +265,28 @@ class DataCollator: ) ) + # Use cached processors if available + if processor_path not in self._processor_cache: + processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) + if self.config.get("padding_side", "left") == "left": + processor.tokenizer.padding_side = "left" + + if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5": + action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] + new_tokens = ["<|propri|>", "<|action|>"] + new_tokens += [ + f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) + ] + processor.tokenizer.add_tokens(new_tokens) + begin_idx_token = "<|action_token_0|>" + token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token) + processor.tokenizer.init_kwargs["action_token_start_index"] = token_id + processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( + action_tokenizer.vocab_size + ) + + self._processor_cache[processor_path] = processor + self.processor = self._processor_cache[processor_path] if not self.use_fast_tokenizer: @@ -282,15 +296,6 @@ class DataCollator: action_tokenizer_path ] - if self.use_fast_tokenizer: - self.action_mapper = {} - for i in range(self.train_action_tokenizer.vocab_size): - token = f"<|action_token_{i}|>" - token_id = self.processor.tokenizer.convert_tokens_to_ids(token) - self.action_mapper[token_id] = i - else: - self.action_mapper = None - @classmethod def _normalize(cls, action, min_stat, delta): """ diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py index 18dd181..01dd2c0 100644 --- a/wall_x/trainer/qwen_vl_act_trainer.py +++ b/wall_x/trainer/qwen_vl_act_trainer.py @@ -15,9 +15,9 @@ from torch.distributed.tensor import distribute_tensor from accelerate import Accelerator from safetensors.torch import load_file from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup - +from transformers import AutoProcessor from wall_x.utils.timers import Timers -from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction +from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction, Qwen2_5_VLConfig from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES from wall_x.data.load_lerobot_dataset import ( PreprocessedDataset, @@ -212,7 +212,6 @@ class QwenVlAct_Trainer: - Memory cleanup """ self.accelerator.wait_for_everyone() - # Optional validation before training starts if self.config.get("resume", None) is not None and self.config["resume"].get( "validate_first", False @@ -225,7 +224,7 @@ class QwenVlAct_Trainer: self.train_loop(epoch) self.accelerator.wait_for_everyone() - if (epoch + 1) % self.config.get("epoch_save_interval", 10) == 0: + if (epoch + 1) % self.config.get("epoch_save_interval", 1) == 0: self.save_checkpoint(epoch) # Validation after each epoch @@ -317,7 +316,6 @@ class QwenVlAct_Trainer: self.timers("forward-compute").stop() loss = outputs.loss - # Check for NaN loss if torch.isnan(loss): print( @@ -519,12 +517,61 @@ class QwenVlAct_Trainer: - Model preparation for distributed training """ # Load pretrained model - model = Qwen2_5_VLMoEForAction.from_pretrained( - self.config["pretrained_wallx_path"], - **{"use_fast_tokenizer": self.use_fast_tokenizer}, - ) - self.processor = model.processor - model = model.to(torch.bfloat16) + model_type = self.config.get("model_type", "qwen2_5") + assert model_type in ["wall-oss", "qwen2_5"] + if model_type == "wall-oss": + model = Qwen2_5_VLMoEForAction.from_pretrained( + self.config["pretrained_wallx_path"], + **{"use_fast_tokenizer": self.use_fast_tokenizer}, + ) + self.processor = model.processor + model = model.to(torch.bfloat16) + elif model_type == "qwen2_5": + + config = Qwen2_5_VLConfig.from_pretrained( + self.config["qwen_vl_act_config_path"] + ) + flow_loss_weight = self.config.get("flow_loss_weight", 1.0) + self.processor = AutoProcessor.from_pretrained( + self.config["pretrained_wallx_path"], use_fast=True + ) + if self.config.get("use_fast_tokenizer", False): + action_tokenizer_path = self.config["action_tokenizer_path"] + action_tokenizer = AutoProcessor.from_pretrained( + action_tokenizer_path, trust_remote_code=True + ) + # process for use fast + new_tokens = ["<|propri|>", "<|action|>"] + new_tokens += [ + f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) + ] + self.processor.tokenizer.add_tokens(new_tokens) + begin_idx_token = "<|action_token_0|>" + token_id = self.processor.tokenizer.convert_tokens_to_ids( + begin_idx_token + ) + self.processor.tokenizer.init_kwargs["action_token_start_index"] = ( + token_id + ) + self.processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( + action_tokenizer.vocab_size + ) + self.processor.action_processor = action_tokenizer + model = Qwen2_5_VLMoEForAction( + config, + self.use_fast_tokenizer, + self.processor, + flow_loss_weight=flow_loss_weight, + ) + + model = model.to(torch.bfloat16) + model = self.load_qwen_pretrain_weight( + model, self.config["pretrained_wallx_path"] + ) + model.resize_token_embeddings(len(self.processor.tokenizer)) + model = model.to(torch.bfloat16) + else: + raise NotImplementedError(f"Invalid model type: {model_type}") # Configure optimizer based on training strategy if "freeze_vlm" in self.config and self.config["freeze_vlm"]: @@ -752,8 +799,12 @@ class QwenVlAct_Trainer: else: ckpt_path = f"{save_path}/{epoch}_{step}" + # If FSDP SHARDED_STATE_DICT is used, please refer to the wall-x/workspace/README.md + # merge checkpoint section to merge the weights into a single safetensors if needed. self.accelerator.save_state(ckpt_path) + self.processor.save_pretrained(os.path.join(ckpt_path, "processor")) + # Save current iteration steps for dataset resuming if step != 0: _rank = self.accelerator.process_index @@ -773,21 +824,23 @@ class QwenVlAct_Trainer: """ checkpoint_path = self.config["resume"]["ckpt"] - if self.config.get("FSDP2", False): - self._load_fsdp_state_dict_with_distribute_tensor() - elif self.config.get("resume", {}).get("load_ckpt_only", False): - # Load only model weights - ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors" - state_dict = load_file(ckpt_path, device="cpu") + if self.config.get("resume", {}).get("load_ckpt_only", False): + if self.config.get("FSDP2", False): + self._load_fsdp_state_dict_with_distribute_tensor() - # Add module prefix if needed for distributed training - new_state_dict = {} - for key in state_dict: - if not key.startswith("module."): - new_key = "module." + key - new_state_dict[new_key] = state_dict[key] + else: + # Load only model weights + ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors" + state_dict = load_file(ckpt_path, device="cpu") - self.model.load_state_dict(new_state_dict, strict=False) + # Add module prefix if needed for distributed training + new_state_dict = {} + for key in state_dict: + if not key.startswith("module."): + new_key = "module." + key + new_state_dict[new_key] = state_dict[key] + + self.model.load_state_dict(new_state_dict, strict=False) else: # Load full checkpoint including optimizer and scheduler states self.accelerator.load_state(checkpoint_path) diff --git a/workspace/README.md b/workspace/README.md index a1acfe5..add2803 100644 --- a/workspace/README.md +++ b/workspace/README.md @@ -4,10 +4,12 @@ This document explains the key configuration parameters and memory requirements ## Quick Start Checklist -### 🚀 **Step 1: Download Pre-trained Model** -Choose one of the available models: +### 🚀 **Step 1: Prepare Model** +Choose one of our pretrained models: - **WALL-OSS-FLOW**: https://huggingface.co/x-square-robot/wall-oss-flow - **WALL-OSS-FAST**: https://huggingface.co/x-square-robot/wall-oss-fast +Or from Qwen-2.5-VL +- Download https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct, settings refer to `config_qact_from_vlm.yml` ### ⚙️ **Step 2: Configure Environment** - Update `run.sh`: Set `code_dir` and `config_path` to your actual paths @@ -89,6 +91,13 @@ Keep `agent_pos_config` consistent with `dof_config`. - `resume.ckpt`: Path to checkpoint for resuming training - `resume.load_ckpt_only`: Only load model weights, not optimizer state +## Merge checkpoint +- If FSDP SHARDED_STATE_DICT is used, please run command below to merge checkpoint into a single safetensors +```bash + # refer to accelerate/commands/merge.py + accelerate merge-weights /path/to/sharded_tensors /path/to/model.safetensors +``` + ## Memory Usage Below are the memory consumption benchmarks for different training configurations using the `lerobot/aloha_mobile_cabinet` dataset: @@ -106,3 +115,9 @@ Below are the memory consumption benchmarks for different training configuration - For single GPU training: Ensure at least 48GB VRAM (e.g., RTX 6000 Ada, A6000) - For multi-GPU training: Enable FSDP2 for optimal memory distribution + +## Reproduce + +Openloop plot `wall-x/workspace/lerobot_example/evaluation/lerobot_openloop.png` + +To reproduce the results, use the config file wall-x/workspace/lerobot_example/config_qact_from_vlm.yml with a global batch size of 128, adjusted via `gradient_accumulation_steps` and numbers of gpu. diff --git a/workspace/lerobot_example/config_qact.yml b/workspace/lerobot_example/config_qact.yml index 6f167ba..dd076e4 100644 --- a/workspace/lerobot_example/config_qact.yml +++ b/workspace/lerobot_example/config_qact.yml @@ -4,7 +4,7 @@ # Model and paths configuration log_name: "robotic_training" log_project: "vla_training" -model_type: qwen2_5 +model_type: wall-oss pretrained_wallx_path: "/path/to/wallx_model/" # Must set save_path: "/path/to/workspace/" # Must set use_fast_tokenizer: False # True: train FAST, False: train Flow diff --git a/workspace/lerobot_example/config_qact_from_vlm.yml b/workspace/lerobot_example/config_qact_from_vlm.yml new file mode 100644 index 0000000..a9fb992 --- /dev/null +++ b/workspace/lerobot_example/config_qact_from_vlm.yml @@ -0,0 +1,117 @@ +# Train from Qwen-2.5-VL + +# Model and paths configuration +log_name: "robotic_training" +log_project: "vla_training" +model_type: qwen2_5 +pretrained_wallx_path: "/path/to/wallx_model/" # Must set +save_path: "/path/to/workspace/" # Must set +use_fast_tokenizer: True # True: train FAST, False: train Flow +action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is true +qwen_vl_act_config_path: "wall-x/workspace/lerobot_example/qwen25_config.json" + + +# Torch Profile +profile: False +profile_save_path: /path/to/profile/ +profile_wait_iters: 10 +profile_warmup_iters: 5 +profile_active_iters: 2 + +# Training hyperparameters +num_warmup_steps: 100 +num_training_steps: 64000000 +learning_rate: 0.00009 +min_lr: 0.00005 +num_epoch: 100 +gradient_accumulation_steps: 1 +batch_size_per_gpu: 8 +padding_side: left +epoch_save_interval: 10 + +# Training optimization settings +FSDP2: True +torch_compile: False + +# Robot configuration - Define degrees of freedom for each component +dof_config: + follow_left_ee_cartesian_pos: 3 # Left end-effector Cartesian position + follow_left_ee_rotation: 3 # Left end-effector rotation + follow_left_gripper: 1 # Left gripper control + follow_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position + follow_right_ee_rotation: 3 # Right end-effector rotation + follow_right_gripper: 1 # Right gripper control + head_actions: 2 # Head/camera movement + height: 1 # Mobile base height control + car_pose: 3 # Mobile base pose (x, y, theta) + +# Agent proprioception configuration (typically matches DOF config) +agent_pos_config: + follow_left_ee_cartesian_pos: 3 + follow_left_ee_rotation: 3 + follow_left_gripper: 1 + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation: 3 + follow_right_gripper: 1 + head_actions: 2 + height: 1 + car_pose: 3 + +# # Checkpoint resuming configuration +# resume: +# ckpt: "/path/to/resume_model/" +# load_ckpt_only: true + +# Data configuration +data: + use_lerobot: true + + # LeRobot dataset configuration + lerobot_config: + repo_id: "lerobot/aloha_mobile_cabinet" + root: null + episodes: null + image_transforms: null + delta_timestamps: null + tolerance_s: 1e-4 + revision: null + force_cache_sync: false + download_videos: true + video_backend: null + + action_horizon: 32 + train_test_split: 0.95 + + # Action keys for observation and prediction + obs_action_keys: + - follow_left_ee_cartesian_pos + - follow_left_ee_rotation + - follow_left_gripper + - follow_right_ee_cartesian_pos + - follow_right_ee_rotation + - follow_right_gripper + - head_actions + - height + - car_pose + + predict_action_keys: + - follow_left_ee_cartesian_pos + - follow_left_ee_rotation + - follow_left_gripper + - follow_right_ee_cartesian_pos + - follow_right_ee_rotation + - follow_right_gripper + - head_actions + - height + - car_pose + + # Image resolution configuration for different camera views + resolution: + face_view: 256 + left_wrist_view: 256 + right_wrist_view: 256 + move1_view: 256 + move2_view: 256 + top_view: 256 + wall_view: 256 + multi_modal: 256 diff --git a/workspace/lerobot_example/evaluation/lerobot_openloop.png b/workspace/lerobot_example/evaluation/lerobot_openloop.png new file mode 100644 index 0000000..dfc72e7 --- /dev/null +++ b/workspace/lerobot_example/evaluation/lerobot_openloop.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aae8646e566b79b64a669956d6d2c779d72809010c802f75b2623ee371444b47 +size 985979 diff --git a/workspace/lerobot_example/qwen25_config.json b/workspace/lerobot_example/qwen25_config.json new file mode 100644 index 0000000..c92cc20 --- /dev/null +++ b/workspace/lerobot_example/qwen25_config.json @@ -0,0 +1,106 @@ +{ + "architectures": [ + "Qwen2_5_VLForConditionalGeneration" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "vision_start_token_id": 151652, + "vision_end_token_id": 151653, + "vision_token_id": 151654, + "image_token_id": 151655, + "video_token_id": 151656, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 128000, + "max_window_layers": 70, + "model_type": "qwen2_5_vl", + "num_attention_heads": 16, + "num_hidden_layers": 36, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.41.2", + "_attn_implementation": "flash_attention_2", + "use_cache": true, + "use_sliding_window": false, + "vision_config": { + "depth": 32, + "hidden_act": "silu", + "hidden_size": 1280, + "intermediate_size": 3420, + "num_heads": 16, + "in_chans": 3, + "out_hidden_size": 2048, + "patch_size": 14, + "spatial_merge_size": 2, + "spatial_patch_size": 14, + "window_size": 112, + "fullatt_block_indexes": [ + 7, + 15, + 23, + 31 + ], + "tokens_per_second": 2, + "temporal_patch_size": 2 + }, + "rope_scaling": { + "type": "mrope", + "mrope_section": [ + 16, + 24, + 24 + ] + }, + "vocab_size": 151936, + "num_experts": 2, + "experts":[ + { + "hidden_size": 2048, + "intermediate_size": 11008, + "hidden_act": "silu" + }, + { + "hidden_size": 2048, + "intermediate_size": 2048, + "hidden_act": "silu" + } + ], + "dof_config": { + "follow_left_ee_cartesian_pos": 3, + "follow_left_ee_rotation": 3, + "follow_left_gripper": 1, + "follow_right_ee_cartesian_pos": 3, + "follow_right_ee_rotation": 3, + "follow_right_gripper": 1, + "head_actions": 2, + "height": 1, + "car_pose": 3 + }, + "agent_pos_config": { + "follow_left_ee_cartesian_pos": 3, + "follow_left_ee_rotation": 3, + "follow_left_gripper": 1, + "follow_right_ee_cartesian_pos": 3, + "follow_right_ee_rotation": 3, + "follow_right_gripper": 1, + "head_actions": 2, + "height": 1, + "car_pose": 3 + }, + "noise_scheduler": { + "beta_alpha": 1.5, + "beta_beta": 1.0, + "s": 0.999, + "num_inference_timesteps": 5 + }, + "dim_inputs": [2048,2048], + "attention_moe": false, + "mlp_moe": true + }