huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix normalizer (#57)

Browse Source 
* fix normalizer

* fix val

* update compute stats

* delete norm

* update readme

* minor fix

* fix action normalizer

* fix

* fix

* update

* update

* update

* update

* update

* lint

* lint

* lint

* lint
This commit is contained in:
Lufang Chen
2025-10-24 17:29:12 +08:00
committed by GitHub
parent 35399d187a
commit d821b0cb26
14 changed files with 837 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/trainer/qwen_vl_act_trainer.py
+26 -7
View File
@@ -263,7 +263,6 @@ class QwenVlAct_Trainer:
self.train_dataloader = self.dataset.get_train_dataloader()
self.model.train()
grad_accum_steps = self.config.get("gradient_accumulation_steps", 1)
total = len(self.train_dataloader)
t0 = time.time()
enable_profiling = self.config["profile"]
@@ -341,7 +340,7 @@ class QwenVlAct_Trainer:
self.timers("optimizer").stop()
# Update global step and learning rate after gradient accumulation
if (i + 1) % grad_accum_steps == 0:
if self.accelerator.sync_gradients:
self.lr_scheduler.step()
self.global_step += 1
lr = self.lr_scheduler.get_last_lr()[0]
@@ -522,7 +521,12 @@ class QwenVlAct_Trainer:
if model_type == "wall-oss":
model = Qwen2_5_VLMoEForAction.from_pretrained(
self.config["pretrained_wallx_path"],
**{"use_fast_tokenizer": self.use_fast_tokenizer},
train_config=self.config,
action_tokenizer_path=(
self.config["action_tokenizer_path"]
if self.use_fast_tokenizer
else None
),
)
self.processor = model.processor
model = model.to(torch.bfloat16)
@@ -535,14 +539,15 @@ class QwenVlAct_Trainer:
self.processor = AutoProcessor.from_pretrained(
self.config["pretrained_wallx_path"], use_fast=True
)
new_tokens = ["<|propri|>", "<|action|>"]
self.processor.tokenizer.add_tokens(new_tokens)
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 += [
new_tokens = [
f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
]
self.processor.tokenizer.add_tokens(new_tokens)
@@ -557,6 +562,19 @@ class QwenVlAct_Trainer:
action_tokenizer.vocab_size
)
self.processor.action_processor = action_tokenizer
# Set the customized robot configuration to ensure consistency between cross-embodiment
# representations and the Wall-X action dimensionality.
Qwen2_5_VLMoEForAction._set_customized_config(self.config)
customized_dof_config = self.config["customized_robot_config"][
"customized_dof_config"
]
customized_agent_pos_config = self.config["customized_robot_config"][
"customized_agent_pos_config"
]
setattr(config, "customized_dof_config", customized_dof_config)
setattr(config, "customized_agent_pos_config", customized_agent_pos_config)
model = Qwen2_5_VLMoEForAction(
config,
self.use_fast_tokenizer,
@@ -803,7 +821,8 @@ class QwenVlAct_Trainer:
# 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"))
if self.accelerator.is_main_process:
self.processor.save_pretrained(os.path.join(ckpt_path, "processor"))
# Save current iteration steps for dataset resuming
if step != 0:
@@ -845,7 +864,7 @@ class QwenVlAct_Trainer:
# Load full checkpoint including optimizer and scheduler states
self.accelerator.load_state(checkpoint_path)
self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}")
self.print_rank0(f"\033[32mResumed from checkpoint: {checkpoint_path}\033[0m")
def _load_fsdp_state_dict_with_distribute_tensor(self):