From e9332a283deb297dd4721d7c13e0e029a31e13a7 Mon Sep 17 00:00:00 2001 From: Lufang Chen <64068400+vincentccc@users.noreply.github.com> Date: Thu, 11 Sep 2025 13:18:33 +0800 Subject: [PATCH] [lint] Update lint (#16) * update lint * update readme * update ruff lint --- .pre-commit-config.yaml | 27 + CONTRIBUTING.md | 17 + README.md | 2 +- csrc/README.md | 4 +- csrc/ops.cu | 1 - csrc/rope.cu | 1 - pyproject.toml | 2 + requirements.txt | 2 +- scripts/draw_openloop_plot.py | 43 +- scripts/fake_inference.py | 27 +- scripts/merge_tokenizer.py | 7 +- train_qact.py | 60 +- wall_x/data/config.py | 87 +- wall_x/data/load_lerobot_dataset.py | 151 ++- wall_x/data/utils.py | 113 +- wall_x/fusions/backend.py | 209 ++-- wall_x/fusions/ops.py | 174 ++- wall_x/model/action_head.py | 185 ++-- wall_x/model/qwen2_5_based/__init__.py | 10 +- .../qwen2_5_based/configuration_qwen2_5_vl.py | 2 +- .../qwen2_5_based/modeling_qwen2_5_vl.py | 574 +++++++--- .../qwen2_5_based/modeling_qwen2_5_vl_act.py | 999 ++++++++++++------ wall_x/trainer/qwen_vl_act_trainer.py | 389 ++++--- wall_x/utils/constant.py | 285 ++++- wall_x/utils/timers.py | 93 +- workspace/README.md | 4 +- workspace/lerobot_example/config_qact.yml | 10 +- workspace/lerobot_example/run.sh | 2 +- 28 files changed, 2406 insertions(+), 1074 deletions(-) create mode 100644 .pre-commit-config.yaml create mode 100644 CONTRIBUTING.md create mode 100644 pyproject.toml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..8a9d4b3 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,27 @@ +# See https://pre-commit.com for more information +# See https://pre-commit.com/hooks.html for more hooks +exclude: ".git" + +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.2.2 + hooks: + - id: ruff + args: [ --fix, --exit-non-zero-on-fix ] + + - repo: https://github.com/psf/black + rev: 24.2.0 + hooks: + - id: black + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.5.0 + hooks: + - id: check-added-large-files + - id: check-ast + - id: check-case-conflict + - id: check-merge-conflict + - id: check-toml + - id: check-yaml + - id: end-of-file-fixer + - id: trailing-whitespace diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..26e430c --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,17 @@ +# Contributing to Wall-x + +## Submit a Pull Request + +Before opening a pull request, please make sure your code passes the lint checks. + +```bash +# Install pre-commit hooks (run once) +pre-commit install +``` + +Or + +```bash +# Manually run all checks +pre-commit run --all-files +``` diff --git a/README.md b/README.md index 74907ff..cb9f576 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ ## Building General-Purpose Robots Based on Embodied Foundation Model -We are building the embodied foundation model to capture and compress the world's most valuable data: the continuous, high-fidelity stream of physical interaction. +We are building the embodied foundation model to capture and compress the world's most valuable data: the continuous, high-fidelity stream of physical interaction. By creating a direct feedback loop between the model's decisions and the body's lived experience, we enable the emergence of a truly generalizable intelligence—one that understands not just how the world works, but how to act effectively within it. diff --git a/csrc/README.md b/csrc/README.md index d73f8d4..70d5e63 100644 --- a/csrc/README.md +++ b/csrc/README.md @@ -10,7 +10,7 @@ High-performance CUDA kernels for accelerating model training. ### Token Permutation - `permute`: Token permutation for MoE routing -- `unpermute`: Token recovery after expert computation +- `unpermute`: Token recovery after expert computation - `unpermute_bwd`: Backward pass for token recovery ### Multimodal RoPE @@ -20,4 +20,4 @@ High-performance CUDA kernels for accelerating model training. ## Acknowledgments -The `permute` and `unpermute` operators are adapted from [fanshiqing/grouped_gemm](https://github.com/fanshiqing/grouped_gemm). Thanks for their open-source contributions. \ No newline at end of file +The `permute` and `unpermute` operators are adapted from [fanshiqing/grouped_gemm](https://github.com/fanshiqing/grouped_gemm). Thanks for their open-source contributions. diff --git a/csrc/ops.cu b/csrc/ops.cu index 059c77e..d30bb5b 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -13,4 +13,3 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("rope", &launch_multimodal_rope_forward, "Multimodal RoPE forward kernel"); m.def("rope_bwd", &launch_multimodal_rope_backward, "Multimodal RoPE backward kernel"); } - diff --git a/csrc/rope.cu b/csrc/rope.cu index 09c1a11..6369f1c 100644 --- a/csrc/rope.cu +++ b/csrc/rope.cu @@ -549,4 +549,3 @@ void launch_multimodal_rope_backward( break; } } - diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..96581bb --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,2 @@ +[tool.ruff] +per-file-ignores = { "__init__.py" = ["F401", "E402"] } diff --git a/requirements.txt b/requirements.txt index 3f88666..65b6121 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,4 @@ accelerate==1.10.1 peft==0.17.1 scipy==1.15.3 torchdiffeq==0.2.5 -qwen_vl_utils==0.0.11 \ No newline at end of file +qwen_vl_utils==0.0.11 diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py index 2be5694..982f6c8 100644 --- a/scripts/draw_openloop_plot.py +++ b/scripts/draw_openloop_plot.py @@ -1,6 +1,7 @@ import os import yaml import torch +import matplotlib.pyplot as plt from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs @@ -8,20 +9,24 @@ from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs model_path = "path/to/model" action_tokenizer_path = "path/to/action_tokenizer" save_dir = "path/to/plot" -model = Qwen2_5_VLMoEForAction.from_pretrained(model_path, action_tokenizer_path=action_tokenizer_path) +model = Qwen2_5_VLMoEForAction.from_pretrained( + model_path, action_tokenizer_path=action_tokenizer_path +) model.eval() model = model.to("cuda") model = model.bfloat16() + def load_config(config_path): """Load configuration from YAML file.""" with open(config_path, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) - + config["data"]["model_type"] = config.get("model_type") - + return config + # get test dataloader path = "path/to/config" config = load_config(path) @@ -38,7 +43,7 @@ gt_traj = torch.zeros((total_frames, action_dim)) pred_traj = torch.zeros((total_frames, action_dim)) for idx, batch in enumerate(dataloader): - if idx % pred_horizon ==0 and idx + pred_horizon < total_frames: + if idx % pred_horizon == 0 and idx + pred_horizon < total_frames: batch = batch.to("cuda") with torch.no_grad(): outputs = model( @@ -46,36 +51,36 @@ for idx, batch in enumerate(dataloader): action_dim=action_dim, pred_horizon=pred_horizon, mode="predict", - predict_mode="fast" + predict_mode="fast", ) - pred_traj[idx : idx + pred_horizon] = outputs['predict_action'].detach().cpu() - + pred_traj[idx : idx + pred_horizon] = outputs["predict_action"].detach().cpu() + # Denormalize ground truth actions - gt_action_chunk = batch['action_chunk'][:, :, :action_dim] + gt_action_chunk = batch["action_chunk"][:, :, :action_dim] dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype) - denormalized_gt = model.action_preprocessor.normalizer_action.unnormalize_data(gt_action_chunk, ["x2_normal"], dof_mask) + denormalized_gt = model.action_preprocessor.normalizer_action.unnormalize_data( + gt_action_chunk, ["x2_normal"], dof_mask + ) gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu() - + gt_traj_np = gt_traj.numpy() pred_traj_np = pred_traj.numpy() timesteps = gt_traj.shape[0] -import matplotlib.pyplot as plt - fig, axs = plt.subplots(action_dim, 1, figsize=(15, 5 * action_dim), sharex=True) -fig.suptitle(f'Action Comparison for lerobot', fontsize=16) +fig.suptitle("Action Comparison for lerobot", fontsize=16) for i in range(action_dim): - axs[i].plot(range(timesteps), gt_traj_np[:, i], label='Ground Truth') - axs[i].plot(range(timesteps), pred_traj_np[:, i], label='Prediction') - axs[i].set_ylabel(f'Action Dim {i+1}') + axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth") + axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction") + axs[i].set_ylabel(f"Action Dim {i+1}") axs[i].legend() axs[i].grid(True) - -axs[-1].set_xlabel('Timestep') + +axs[-1].set_xlabel("Timestep") plt.tight_layout(rect=[0, 0.03, 1, 0.95]) os.makedirs(save_dir, exist_ok=True) -plt.savefig(os.path.join(save_dir, f"lerobot_comparison.png")) +plt.savefig(os.path.join(save_dir, "lerobot_comparison.png")) plt.close() diff --git a/scripts/fake_inference.py b/scripts/fake_inference.py index 8c2cdf6..428aec7 100644 --- a/scripts/fake_inference.py +++ b/scripts/fake_inference.py @@ -10,10 +10,14 @@ batch_size = 1 seq_length = 50 torch.manual_seed(0) -fake_input_ids = torch.randint(0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long) +fake_input_ids = torch.randint( + 0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long +) fake_attention_mask = torch.ones((batch_size, seq_length), dtype=torch.long) fake_moe_token_types = torch.zeros((batch_size, seq_length), dtype=torch.long) -fake_position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) +fake_position_ids = ( + torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) +) fake_proprioception = torch.randn((batch_size, 1, 20), dtype=torch.float32) fake_agent_pos_mask = torch.ones((batch_size, 1, 20), dtype=torch.float32) fake_dof_mask = torch.ones((batch_size, 32, 20), dtype=torch.float32) @@ -44,37 +48,38 @@ try: agent_pos_mask=fake_agent_pos_mask, dof_mask=fake_dof_mask, dataset_names=fake_dataset_names, - mode="validate" + mode="validate", ) - + print("✅ Fake inference test successful!") print(f"Output logits shape: {outputs.logits.shape}") print(f"Output logits dtype: {outputs.logits.dtype}") print(f"Output logits device: {outputs.logits.device}") - + # Check if output is reasonable if outputs.logits.shape == (batch_size, seq_length, model.config.vocab_size): print("✅ Output shape correct") else: print("❌ Output shape incorrect") - + if not torch.isnan(outputs.logits).any(): print("✅ Output contains no NaN values") else: print("❌ Output contains NaN values") - + if not torch.isinf(outputs.logits).any(): print("✅ Output contains no infinity values") else: print("❌ Output contains infinity values") - - print(f"Output logits statistics:") + + print("Output logits statistics:") print(f" Min value: {outputs.logits.min().item():.4f}") print(f" Max value: {outputs.logits.max().item():.4f}") print(f" Mean: {outputs.logits.mean().item():.4f}") print(f" Standard deviation: {outputs.logits.std().item():.4f}") - + except Exception as e: print(f"❌ Fake inference test failed: {e}") import traceback - traceback.print_exc() \ No newline at end of file + + traceback.print_exc() diff --git a/scripts/merge_tokenizer.py b/scripts/merge_tokenizer.py index 2ec453e..52ac1ae 100644 --- a/scripts/merge_tokenizer.py +++ b/scripts/merge_tokenizer.py @@ -8,13 +8,15 @@ use_fast_tokenizer = True processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) processor.tokenizer.padding_side = "left" -action_tokenizer = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True) +action_tokenizer = AutoProcessor.from_pretrained( + action_tokenizer_path, trust_remote_code=True +) new_tokens = ["<|propri|>", "<|action|>"] new_tokens += [f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)] num_added_tokens = processor.tokenizer.add_tokens(new_tokens) -begin_idx_token = f"<|action_token_0|>" +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 @@ -22,4 +24,3 @@ processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_tokenizer.vo new_tokenizer_dir = "/path/to/new_tokenizer" os.makedirs(new_tokenizer_dir, exist_ok=True) processor.save_pretrained(new_tokenizer_dir) - diff --git a/train_qact.py b/train_qact.py index e0281f7..2906c90 100755 --- a/train_qact.py +++ b/train_qact.py @@ -4,7 +4,11 @@ import time import yaml import wandb from argparse import ArgumentParser -from accelerate import Accelerator, DistributedDataParallelKwargs, DataLoaderConfiguration +from accelerate import ( + Accelerator, + DistributedDataParallelKwargs, + DataLoaderConfiguration, +) from wall_x.trainer.qwen_vl_act_trainer import QwenVlAct_Trainer @@ -18,29 +22,33 @@ def load_config(config_path): """Load configuration from YAML file.""" with open(config_path, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) - + # Set model_type in data config if not already set config["data"]["model_type"] = config.get("model_type") - + return config def setup_accelerator(config): """Initialize and configure the accelerator for distributed training.""" - print(f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Preparing accelerator") - + print( + f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Preparing accelerator" + ) + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator_dataloader_config = DataLoaderConfiguration(dispatch_batches=False) - + accelerator = Accelerator( kwargs_handlers=[ddp_kwargs], mixed_precision="bf16", dataloader_config=accelerator_dataloader_config, - gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1) + gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1), ) - - print(f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Accelerator initialization complete") - + + print( + f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Accelerator initialization complete" + ) + return accelerator @@ -48,18 +56,18 @@ def setup_logging(config, accelerator): """Set up logging with wandb for the main process.""" if not accelerator.is_main_process: return None - + # Create save directory if it doesn't exist save_path = config["save_path"] if not os.path.exists(save_path): print(f"Save path {save_path} does not exist, creating directory.") os.makedirs(save_path, exist_ok=True) - + print("Configuration:") print("=" * 50) print(json.dumps(config, indent=2, ensure_ascii=False)) print("=" * 50) - + # Initialize wandb logger logger = wandb.init( project=config["log_project"], @@ -67,23 +75,23 @@ def setup_logging(config, accelerator): save_code=False, force=False, ) - + return logger def main(args): """Main training function.""" setup_environment() - + # Load configuration config = load_config(args.config) - + # Set up accelerator accelerator = setup_accelerator(config) - + # Set up logging logger = setup_logging(config, accelerator) - + # Initialize trainer trainer = QwenVlAct_Trainer( config=config, @@ -92,15 +100,19 @@ def main(args): seed=args.seed, data_config_path=args.config, ) - + # Start training trainer.fit() -if __name__ == '__main__': +if __name__ == "__main__": parser = ArgumentParser(description="Training script for Wall-X model") - parser.add_argument("--config", type=str, required=True, help="Path to configuration YAML file") - parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility") - + parser.add_argument( + "--config", type=str, required=True, help="Path to configuration YAML file" + ) + parser.add_argument( + "--seed", type=int, default=42, help="Random seed for reproducibility" + ) + args = parser.parse_args() - main(args) \ No newline at end of file + main(args) diff --git a/wall_x/data/config.py b/wall_x/data/config.py index 240110c..a6964e8 100644 --- a/wall_x/data/config.py +++ b/wall_x/data/config.py @@ -5,57 +5,84 @@ from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR # Tactile sensor file mapping for data processing TACTILE_FILE_MAPPING = { - "tactile_data_left": "left_tactile", - "tactile_data_right": "right_tactile" + "tactile_data_left": "left_tactile", + "tactile_data_right": "right_tactile", } # Supported action datasets ACTION_DATASET_NAMES = [ - "x2_normal", "agibotworld_alpha", "droid", "fractal", "bridge_data_v2", - "DobbE", "RH20T", "UMI-biarm", "austin_buds", "austin_sailor", "austin_sirius", - "bc_z", "berkeley_autolab_ur5", "berkeley_cable_routing", "berkeley_fanuc_manipulation", - "dlr_edan_shared_control", "fmb", "furniture_bench", "jaco_play", "nyu_rot", - "stanford_hydra", "stanford_kuka_multimodal", "taco_play", "utaustin_mutex", "viola" + "x2_normal", + "agibotworld_alpha", + "droid", + "fractal", + "bridge_data_v2", + "DobbE", + "RH20T", + "UMI-biarm", + "austin_buds", + "austin_sailor", + "austin_sirius", + "bc_z", + "berkeley_autolab_ur5", + "berkeley_cable_routing", + "berkeley_fanuc_manipulation", + "dlr_edan_shared_control", + "fmb", + "furniture_bench", + "jaco_play", + "nyu_rot", + "stanford_hydra", + "stanford_kuka_multimodal", + "taco_play", + "utaustin_mutex", + "viola", ] # Supported multimodal datasets MULTIMODAL_DATASET_NAMES = [ - "x2_multimodal_from_action", "x2_multimodal", "x2_subtask_generation", - "multimodal_CapsFusion", "multimodal_Robo2VLM", "multimodal_RoboPoint", - "multimodal_EQA", "multimodal_Cambrian", "multimodal_pixmo", - "multimodal_VQAv2", "multimodal_COCO" + "x2_multimodal_from_action", + "x2_multimodal", + "x2_subtask_generation", + "multimodal_CapsFusion", + "multimodal_Robo2VLM", + "multimodal_RoboPoint", + "multimodal_EQA", + "multimodal_Cambrian", + "multimodal_pixmo", + "multimodal_VQAv2", + "multimodal_COCO", ] @dataclass class X2RDataProcessingConfig: """Configuration class for X2R data processing pipeline. - + This class contains all the necessary parameters for processing robotic data including camera mappings, tactile sensor configurations, action predictions, and various processing options. """ - + # Action prediction configuration predict_action_keys: List[str] = field(default_factory=list) obs_action_keys: List[str] = field(default_factory=list) - + # Image resolution settings for different views resolution: Dict[str, int] = field( default_factory=lambda: { - "face_view": -1, - "left_wrist_view": 128, - "right_wrist_view": 128 + "face_view": -1, + "left_wrist_view": 128, + "right_wrist_view": 128, } ) - + # Dataset splitting train_test_split: float = 0.9 split_seed: int = 42 - + # Instruction handling priority_order: Optional[Dict[str, float]] = None - + # Vision model parameters model_type: str = "qwen2_5" max_pixels: int = MAX_PIXELS @@ -63,27 +90,29 @@ class X2RDataProcessingConfig: image_factor: int = IMAGE_FACTOR generate_subtask_ratio: float = 0.0 - + def __post_init__(self): """Post-initialization validation and setup.""" # Validate train/test split if not 0 < self.train_test_split < 1: - raise ValueError(f"train_test_split must be between 0 and 1, got {self.train_test_split}") - + raise ValueError( + f"train_test_split must be between 0 and 1, got {self.train_test_split}" + ) + def as_dict(self) -> Dict: """Convert configuration to dictionary format. - + Returns: Dict: Configuration as dictionary """ return self.__dict__ - - def update(self, **kwargs) -> 'X2RDataProcessingConfig': + + def update(self, **kwargs) -> "X2RDataProcessingConfig": """Update configuration parameters. - + Args: **kwargs: Key-value pairs to update - + Returns: X2RDataProcessingConfig: Updated configuration instance """ @@ -92,4 +121,4 @@ class X2RDataProcessingConfig: setattr(self, key, value) else: raise ValueError(f"Unknown configuration parameter: {key}") - return self \ No newline at end of file + return self diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py index 373dd8a..cd26b50 100644 --- a/wall_x/data/load_lerobot_dataset.py +++ b/wall_x/data/load_lerobot_dataset.py @@ -8,7 +8,12 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from typing import Protocol, SupportsIndex, TypeVar from qwen_vl_utils.vision_process import smart_resize from wall_x.data.config import X2RDataProcessingConfig -from wall_x.data.utils import process_grounding_points, get_wallx_normal_text, replace_action_token, preprocesser_call +from wall_x.data.utils import ( + process_grounding_points, + get_wallx_normal_text, + replace_action_token, + preprocesser_call, +) from transformers import AutoProcessor @@ -67,7 +72,9 @@ class PreprocessedDataset(Dataset[T_co]): img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy()) orig_width, orig_height = img_pil.size # 2. Apply resolution constraints (if config is not -1) - target_size = self.data_config.resolution.get(self._cam_key_mapping[key], -1) + target_size = self.data_config.resolution.get( + self._cam_key_mapping[key], -1 + ) if target_size != -1: # Maintain aspect ratio logic if orig_width > orig_height: # Landscape image @@ -108,7 +115,9 @@ class PreprocessedDataset(Dataset[T_co]): self._cam_key_mapping, generate_subtask_ratio=generate_subtask_ratio, ) - text = process_grounding_points(complete_text, h, w, resize_h, resize_w, self.data_config.model_type) + text = process_grounding_points( + complete_text, h, w, resize_h, resize_w, self.data_config.model_type + ) result = { "image_inputs": image_inputs, "text": text, @@ -150,7 +159,9 @@ class PreprocessedDataset(Dataset[T_co]): batch_size=batch_size, sampler=sampler, # Use distributed sampler instead of shuffle=True num_workers=num_workers, - collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats), + collate_fn=DataCollator( + self.config, self.dataload_config, self._dataset.meta.stats + ), pin_memory=True, # Enable for GPU training persistent_workers=num_workers > 0, # Only if num_workers > 0 prefetch_factor=2, # Reduce memory usage @@ -164,7 +175,9 @@ class PreprocessedDataset(Dataset[T_co]): Get distributed evaluation dataloader (no shuffling for consistent evaluation) """ - batch_size = self.config.get("eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8)) + batch_size = self.config.get( + "eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8) + ) num_workers = self.config.get("num_workers", 4) # Create distributed sampler for evaluation (no shuffle) @@ -181,7 +194,9 @@ class PreprocessedDataset(Dataset[T_co]): batch_size=batch_size, sampler=sampler, num_workers=num_workers, - collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats), + collate_fn=DataCollator( + self.config, self.dataload_config, self._dataset.meta.stats + ), pin_memory=True, persistent_workers=num_workers > 0, prefetch_factor=2, @@ -212,19 +227,30 @@ class DataCollator: # 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) + 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: - self._action_tokenizer_cache[action_tokenizer_path] = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True) + if ( + self.use_fast_tokenizer + and action_tokenizer_path not in self._action_tokenizer_cache + ): + self._action_tokenizer_cache[action_tokenizer_path] = ( + AutoProcessor.from_pretrained( + action_tokenizer_path, trust_remote_code=True + ) + ) self.processor = self._processor_cache[processor_path] if not self.use_fast_tokenizer: self.train_action_tokenizer = None else: - self.train_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] + self.train_action_tokenizer = self._action_tokenizer_cache[ + action_tokenizer_path + ] if self.use_fast_tokenizer: self.action_mapper = {} @@ -254,9 +280,27 @@ class DataCollator: agent_pos.nan_to_num_(nan=0.0) agent_pos = self._normalize(agent_pos, self.min_stat, self.delta) if agent_pos.shape[-1] != 20: - agent_pos = torch.cat([agent_pos, torch.zeros(agent_pos.shape[0], agent_pos.shape[1], 20 - agent_pos.shape[-1])], dim=-1) + agent_pos = torch.cat( + [ + agent_pos, + torch.zeros( + agent_pos.shape[0], + agent_pos.shape[1], + 20 - agent_pos.shape[-1], + ), + ], + dim=-1, + ) agent_pos_mask = torch.cat( - [agent_pos_mask, torch.zeros(agent_pos_mask.shape[0], agent_pos_mask.shape[1], 20 - agent_pos_mask.shape[-1])], dim=-1 + [ + agent_pos_mask, + torch.zeros( + agent_pos_mask.shape[0], + agent_pos_mask.shape[1], + 20 - agent_pos_mask.shape[-1], + ), + ], + dim=-1, ) additional_inputs["proprioception"] = agent_pos additional_inputs["agent_pos_mask"] = agent_pos_mask @@ -268,18 +312,42 @@ class DataCollator: action.nan_to_num_(nan=0.0) action = self._normalize(action, self.min_stat, self.delta) if action.shape[-1] != 20: - action = torch.cat([action, torch.zeros(action.shape[0], action.shape[1], 20 - action.shape[-1])], dim=-1) - dof_mask = torch.cat([dof_mask, torch.zeros(dof_mask.shape[0], dof_mask.shape[1], 20 - dof_mask.shape[-1])], dim=-1) + action = torch.cat( + [ + action, + torch.zeros( + action.shape[0], action.shape[1], 20 - action.shape[-1] + ), + ], + dim=-1, + ) + dof_mask = torch.cat( + [ + dof_mask, + torch.zeros( + dof_mask.shape[0], + dof_mask.shape[1], + 20 - dof_mask.shape[-1], + ), + ], + dim=-1, + ) additional_inputs["action_chunk"] = action additional_inputs["dof_mask"] = dof_mask elif key == "image_inputs": - additional_inputs["image_inputs"] = [item["image_inputs"] for item in batch] + additional_inputs["image_inputs"] = [ + item["image_inputs"] for item in batch + ] elif key == "text": additional_inputs["text"] = [item["text"] for item in batch] elif key == "frame_index": - additional_inputs["frame_index"] = torch.stack([item["frame_index"] for item in batch]) + additional_inputs["frame_index"] = torch.stack( + [item["frame_index"] for item in batch] + ) else: - raise NotImplementedError(f"{key} input not implemented in preprocesser") + raise NotImplementedError( + f"{key} input not implemented in preprocesser" + ) additional_inputs["text"] = replace_action_token( additional_inputs["text"], @@ -342,20 +410,27 @@ def load_lerobot_data( delta_timestamps = { # action chunk - "action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)], + "action": [ + t / dataset_fps + for t in range(dataload_config.get("action_horizon", 32) - 1) + ], } batch_size = config.get("batch_size_per_gpu", 8) # repo_id = "lerobot/aloha_mobile_cabinet" repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet") - dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps, video_backend="pyav") + dataset = LeRobotDataset( + repo_id, delta_timestamps=delta_timestamps, video_backend="pyav" + ) if rank == 0: print(f"Selected episodes: {dataset.episodes}") print(f"Number of episodes selected: {dataset.num_episodes}") print(f"Number of frames selected: {dataset.num_frames}") - dataset = PreprocessedDataset(dataset, config, dataload_config, seed=seed, rank=rank, world_size=world_size) + dataset = PreprocessedDataset( + dataset, config, dataload_config, seed=seed, rank=rank, world_size=world_size + ) # Calculate samples per process if world_size > 1: @@ -385,7 +460,9 @@ def load_lerobot_data( return dataset, train_num -def get_distributed_dataloader(dataset, config, rank=0, world_size=1, seed=42, is_train=True): +def get_distributed_dataloader( + dataset, config, rank=0, world_size=1, seed=42, is_train=True +): """ Helper function to get distributed dataloader @@ -429,23 +506,29 @@ def get_data_configs(config): return data_config + class TestDataset(PreprocessedDataset): def __init__(self, dataset, config, dataload_config, seed=42): - super().__init__(dataset, config, dataload_config, seed=seed, rank=0, world_size=1) + super().__init__( + dataset, config, dataload_config, seed=seed, rank=0, world_size=1 + ) def get_dataloader(self): """ Get distributed evaluation dataloader (no shuffling for consistent evaluation) """ - + dataloader = torch.utils.data.DataLoader( self, batch_size=1, - collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats), + collate_fn=DataCollator( + self.config, self.dataload_config, self._dataset.meta.stats + ), ) return dataloader - + + def load_test_dataset( config, lerobot_config, @@ -471,16 +554,24 @@ def load_test_dataset( delta_timestamps = { # action chunk - "action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)], + "action": [ + t / dataset_fps + for t in range(dataload_config.get("action_horizon", 32) - 1) + ], } repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet") - dataset = LeRobotDataset(repo_id, episodes=[episode], delta_timestamps=delta_timestamps, video_backend="pyav") + dataset = LeRobotDataset( + repo_id, + episodes=[episode], + delta_timestamps=delta_timestamps, + video_backend="pyav", + ) print(f"Selected episodes: {dataset.episodes}") print(f"Number of episodes selected: {dataset.num_episodes}") print(f"Number of frames selected: {dataset.num_frames}") - + dataset = TestDataset(dataset, config, dataload_config, seed=seed) - - return dataset \ No newline at end of file + + return dataset diff --git a/wall_x/data/utils.py b/wall_x/data/utils.py index cc3517f..e6853a9 100644 --- a/wall_x/data/utils.py +++ b/wall_x/data/utils.py @@ -15,10 +15,10 @@ from transformers import BatchFeature CAMERA_NAME_MAPPING = { "face_view": "front view", - "left_wrist_view": "left wrist view", + "left_wrist_view": "left wrist view", "right_wrist_view": "right wrist view", "move1_view": "move view", - "move2_view": "move view", + "move2_view": "move view", "wall_view": "wall view", "top_view": "top view", } @@ -78,13 +78,13 @@ def preprocesser_call( return_tensors: str = "pt", ) -> BatchFeature: """Unified preprocessing function for Wall-X model handling text, image and video inputs. - + Processes inputs into format suitable for multimodal transformer models, including: - Text tokenization and special token handling - Image/video processing through image processor - Attention mask and label generation - Padding and truncation handling - + Args: processor: Multimodal processor containing tokenizer and image processor images: Input images (PIL, numpy arrays, or torch tensors) @@ -94,7 +94,7 @@ def preprocesser_call( truncation: Whether to truncate sequences longer than max_length max_length: Maximum length for truncation/padding return_tensors: Format for returned tensors ('pt', 'np', etc.) - + Returns: BatchFeature containing processed inputs with keys: - input_ids: Tokenized text @@ -131,15 +131,17 @@ def preprocesser_call( # Process image placeholder tokens in text if image_grid_thw is not None: - merge_length = processor.image_processor.merge_size ** 2 + merge_length = processor.image_processor.merge_size**2 index = 0 for i in range(len(text)): while "<|image_pad|>" in text[i]: # Add bounds checking to avoid index overflow if index >= len(image_grid_thw): - print(f"Warning: Number of image placeholders ({index + 1}) " - f"exceeds actual images ({len(image_grid_thw)}), " - f"skipping remaining placeholder processing") + print( + f"Warning: Number of image placeholders ({index + 1}) " + f"exceeds actual images ({len(image_grid_thw)}), " + f"skipping remaining placeholder processing" + ) break # Replace image placeholder with actual token count token_count = image_grid_thw[index].prod() // merge_length @@ -151,7 +153,7 @@ def preprocesser_call( # Process video placeholder tokens in text if video_grid_thw is not None: - merge_length = processor.image_processor.merge_size ** 2 + merge_length = processor.image_processor.merge_size**2 index = 0 for i in range(len(text)): while "<|video_pad|>" in text[i]: @@ -169,7 +171,7 @@ def preprocesser_call( return_tensors=return_tensors, padding=padding, truncation=truncation, - max_length=max_length + max_length=max_length, ) # Get pad token ID for label generation @@ -209,9 +211,9 @@ def preprocesser_call( # From second part onwards, each part starts with assistant response for k in range(current_pos + 1, len(text_inputs.input_ids[i])): if text_inputs.input_ids[i][k] == im_end_token_id: - assistant_regions.append(( - current_pos + len(assistant_tokens), k + 2 - )) + assistant_regions.append( + (current_pos + len(assistant_tokens), k + 2) + ) break current_pos += len(part_tokens) + 3 @@ -235,7 +237,14 @@ def preprocesser_call( return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) -def process_grounding_points(text: str, orig_height: int, orig_width: int, resized_height: int, resized_width: int, model_type: str) -> str: +def process_grounding_points( + text: str, + orig_height: int, + orig_width: int, + resized_height: int, + resized_width: int, + model_type: str, +) -> str: """Process grounding point coordinates in text based on image resizing. Adjusts coordinate values in tags to match resized image dimensions @@ -309,7 +318,9 @@ def process_grounding_points(text: str, orig_height: int, orig_width: int, resiz def get_frame_instruction( - instruction_info: Dict[str, Any], frame_idx: Optional[int] = None, truncate_keys: Optional[List[str]] = None + instruction_info: Dict[str, Any], + frame_idx: Optional[int] = None, + truncate_keys: Optional[List[str]] = None, ) -> Tuple[Dict[str, Any], Optional[int]]: """Extract frame-specific instruction from instruction dictionary. @@ -322,7 +333,12 @@ def get_frame_instruction( Tuple of (frame_instruction_dict, split_end_frame) """ if truncate_keys is None: - truncate_keys = ["subtask_generation", "distribute", "subtask_generation_zh", "distribute_zh"] + truncate_keys = [ + "subtask_generation", + "distribute", + "subtask_generation_zh", + "distribute_zh", + ] instruction_for_frame = {} split_end = None @@ -334,7 +350,11 @@ def get_frame_instruction( start_frame, end_frame = map(int, frame_range.split(" ")) if start_frame <= frame_idx < end_frame or (start_frame == frame_idx): instruction_for_frame[key] = frame_instruction - if truncate_keys is not None and split_end is None and key in truncate_keys: + if ( + truncate_keys is not None + and split_end is None + and key in truncate_keys + ): split_end = end_frame + 1 break else: @@ -343,7 +363,9 @@ def get_frame_instruction( return instruction_for_frame, split_end -def get_task_instruction(frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None) -> str: +def get_task_instruction( + frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None +) -> str: """Construct task instruction from available instruction fields using priority sampling. Args: @@ -428,7 +450,9 @@ def get_wallx_normal_text( action_fast_symbol = "<|action_fast|>" # System prologue - prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" + prologue = ( + f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" + ) # User request with observation user_request = f"{role_start_symbol}user\nObservation:" @@ -439,13 +463,18 @@ def get_wallx_normal_text( user_request += "\nInstruction:" # Get frame-specific instruction - frame_instruction_info, _ = get_frame_instruction(instruction_info, frame_idx=frame_idx) + frame_instruction_info, _ = get_frame_instruction( + instruction_info, frame_idx=frame_idx + ) generate_subtask = False priority_keys = ["subtask_generation", "distribute"] # Decide whether to generate subtask or actions - if bool(set(frame_instruction_info.keys()) & set(priority_keys)) and random.random() < generate_subtask_ratio: + if ( + bool(set(frame_instruction_info.keys()) & set(priority_keys)) + and random.random() < generate_subtask_ratio + ): # Generate subtask (equivalent to VQA task) instruction = frame_instruction_info.get("instruction", "") text_prompt = "\nPredict the next action in language.\n" @@ -457,11 +486,15 @@ def get_wallx_normal_text( output_instruction = frame_instruction_info[key] break - assistant_output = f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}" + assistant_output = ( + f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}" + ) generate_subtask = True else: # Generate actions - instruction = get_task_instruction(frame_instruction_info, priority_order=priority_order) + instruction = get_task_instruction( + frame_instruction_info, priority_order=priority_order + ) text_prompt = f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n" user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" assistant_output = f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n{action_symbol * action_chunk_size}" @@ -470,7 +503,9 @@ def get_wallx_normal_text( return complete_text, generate_subtask -def get_action_tokens(normalized_actions: Union[torch.Tensor, List], action_tokenizer) -> List[List[str]]: +def get_action_tokens( + normalized_actions: Union[torch.Tensor, List], action_tokenizer +) -> List[List[str]]: """Convert normalized actions to action token strings. Args: @@ -495,7 +530,9 @@ def get_action_tokens(normalized_actions: Union[torch.Tensor, List], action_toke return all_action_tokens -def pad_action_token_strs(actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>") -> List[str]: +def pad_action_token_strs( + actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>" +) -> List[str]: """Pad action token lists to same length and join as strings. Args: @@ -509,14 +546,20 @@ def pad_action_token_strs(actions_token_lists: List[List[str]], pad_token: str = padded_action_strs = [] for tokens in actions_token_lists: - padded_tokens = tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens)) + padded_tokens = ( + tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens)) + ) padded_action_strs.append("".join(padded_tokens)) return padded_action_strs def replace_action_token( - text: List[str], norm_action: Optional[torch.Tensor], action_tokenizer, dataset_names: List[str], dof_masks: Optional[torch.Tensor] = None + text: List[str], + norm_action: Optional[torch.Tensor], + action_tokenizer, + dataset_names: List[str], + dof_masks: Optional[torch.Tensor] = None, ) -> List[str]: """Replace action placeholders in text with actual action tokens. @@ -531,14 +574,19 @@ def replace_action_token( List of text strings with action tokens replaced """ # Filter out multimodal dataset names - dataset_names = [name for name in dataset_names if name not in MULTIMODAL_DATASET_NAMES] + dataset_names = [ + name for name in dataset_names if name not in MULTIMODAL_DATASET_NAMES + ] # Get required action chunk sizes required_chunk_sizes = [FREQUENCY_MAPPING.get(name, 32) for name in dataset_names] if action_tokenizer is not None and norm_action is not None: # Extract actions based on chunk sizes and DOF masks - norm_action = [action[: required_chunk_sizes[i], dof_masks[i, 0].bool()] for i, action in enumerate(norm_action)] + norm_action = [ + action[: required_chunk_sizes[i], dof_masks[i, 0].bool()] + for i, action in enumerate(norm_action) + ] # Convert to action tokens and pad actions_fast_tokens = get_action_tokens(norm_action, action_tokenizer) @@ -548,7 +596,10 @@ def replace_action_token( actions_fast_token_idx = 0 for i in range(len(text)): if "<|action_fast|>" in text[i]: - text[i] = text[i].replace("<|action_fast|><|im_end|>\n", actions_fast_token_strs[actions_fast_token_idx]) + text[i] = text[i].replace( + "<|action_fast|><|im_end|>\n", + actions_fast_token_strs[actions_fast_token_idx], + ) actions_fast_token_idx += 1 # Remove remaining action placeholders diff --git a/wall_x/fusions/backend.py b/wall_x/fusions/backend.py index b386be6..8bcd554 100644 --- a/wall_x/fusions/backend.py +++ b/wall_x/fusions/backend.py @@ -13,28 +13,29 @@ from typing import Tuple, Optional import wallx_csrc as backend - -def _allocate_asymmetric_dual_outputs(input_expert0: torch.Tensor, - input_expert1: torch.Tensor, - weight_expert0: torch.Tensor, - weight_expert1: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: +def _allocate_asymmetric_dual_outputs( + input_expert0: torch.Tensor, + input_expert1: torch.Tensor, + weight_expert0: torch.Tensor, + weight_expert1: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: """ Allocate output tensors for asymmetric dual expert GEMM operations. - + This function handles the case where two experts may have different output dimensions, which is common in heterogeneous MoE architectures. - + Args: input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k] input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k] weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0] weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1] - + Returns: Tuple[torch.Tensor, torch.Tensor]: Pre-allocated output tensors - output_expert0: Shape [m0, n0] - output_expert1: Shape [m1, n1] - + Raises: AssertionError: If tensor dimensions are incompatible """ @@ -43,42 +44,50 @@ def _allocate_asymmetric_dual_outputs(input_expert0: torch.Tensor, assert input_expert1.ndim == 2, "Expected 2D tensor for input_expert1" assert weight_expert0.ndim == 2, "Expected 2D tensor for weight_expert0" assert weight_expert1.ndim == 2, "Expected 2D tensor for weight_expert1" - + # Verify dimension compatibility for matrix multiplication - assert input_expert0.size(1) == weight_expert0.size(0), \ - f"Input expert0 K dimension {input_expert0.size(1)} != weight expert0 K dimension {weight_expert0.size(0)}" - assert input_expert1.size(1) == weight_expert1.size(0), \ - f"Input expert1 K dimension {input_expert1.size(1)} != weight expert1 K dimension {weight_expert1.size(0)}" - + assert input_expert0.size(1) == weight_expert0.size( + 0 + ), f"Input expert0 K dimension {input_expert0.size(1)} != weight expert0 K dimension {weight_expert0.size(0)}" + assert input_expert1.size(1) == weight_expert1.size( + 0 + ), f"Input expert1 K dimension {input_expert1.size(1)} != weight expert1 K dimension {weight_expert1.size(0)}" + # Calculate output shapes: [m, k] × [k, n] = [m, n] m0, n0 = input_expert0.size(0), weight_expert0.size(1) m1, n1 = input_expert1.size(0), weight_expert1.size(1) - + # Allocate output tensors with matching device and dtype - output_expert0 = torch.empty(m0, n0, device=input_expert0.device, dtype=input_expert0.dtype) - output_expert1 = torch.empty(m1, n1, device=input_expert1.device, dtype=input_expert1.dtype) - + output_expert0 = torch.empty( + m0, n0, device=input_expert0.device, dtype=input_expert0.dtype + ) + output_expert1 = torch.empty( + m1, n1, device=input_expert1.device, dtype=input_expert1.dtype + ) + return output_expert0, output_expert1 -def asym_dual_gmm_separated(input_expert0: torch.Tensor, - input_expert1: torch.Tensor, - weight_expert0: torch.Tensor, - weight_expert1: torch.Tensor, - output_expert0: Optional[torch.Tensor] = None, - output_expert1: Optional[torch.Tensor] = None, - trans_a: bool = False, - trans_b: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: +def asym_dual_gmm_separated( + input_expert0: torch.Tensor, + input_expert1: torch.Tensor, + weight_expert0: torch.Tensor, + weight_expert1: torch.Tensor, + output_expert0: Optional[torch.Tensor] = None, + output_expert1: Optional[torch.Tensor] = None, + trans_a: bool = False, + trans_b: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: """ Asymmetric dual expert grouped GEMM with separated inputs and outputs. - + This is the recommended interface for maximum flexibility and performance when dealing with two experts that may have different intermediate dimensions. The operation is equivalent to: output_expert0 = input_expert0 @ weight_expert0 output_expert1 = input_expert1 @ weight_expert1 But optimized as a single fused kernel call. - + Args: input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k] input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k] @@ -89,10 +98,10 @@ def asym_dual_gmm_separated(input_expert0: torch.Tensor, output_expert1 (torch.Tensor, optional): Pre-allocated output for expert 1 [m1, n1] trans_a (bool, optional): Whether to transpose input tensors. Defaults to False. trans_b (bool, optional): Whether to transpose weight tensors. Defaults to False. - + Returns: Tuple[torch.Tensor, torch.Tensor]: Output tensors (output_expert0, output_expert1) - + Example: >>> # Two experts with different output dimensions >>> input0 = torch.randn(512, 1024, device='cuda') # 512 tokens for expert 0 @@ -110,67 +119,77 @@ def asym_dual_gmm_separated(input_expert0: torch.Tensor, output_expert0 = alloc_out0 if output_expert1 is None: output_expert1 = alloc_out1 - + # Call optimized C++ backend kernel backend.asym_dual_gmm( - input_expert0, input_expert1, - weight_expert0, weight_expert1, - output_expert0, output_expert1, - trans_a, trans_b + input_expert0, + input_expert1, + weight_expert0, + weight_expert1, + output_expert0, + output_expert1, + trans_a, + trans_b, ) - + return output_expert0, output_expert1 -def permute(input: torch.Tensor, - indices: torch.Tensor, - num_out_tokens: int, - workspace: torch.Tensor, - max_expanded_token_num: int) -> torch.Tensor: +def permute( + input: torch.Tensor, + indices: torch.Tensor, + num_out_tokens: int, + workspace: torch.Tensor, + max_expanded_token_num: int, +) -> torch.Tensor: """ Permute input tokens according to expert assignment indices for MoE routing. - + This function reorders tokens based on their assigned experts to enable efficient grouped processing. Used in the forward pass of MoE layers. - + Args: input (torch.Tensor): Input tokens to permute indices (torch.Tensor): Expert assignment indices for each token num_out_tokens (int): Number of output tokens after expansion workspace (torch.Tensor): Temporary workspace tensor for intermediate computations max_expanded_token_num (int): Maximum number of tokens after top-k expansion - + Returns: torch.Tensor: Permuted tokens grouped by expert assignment - + Note: This is typically used with top-k expert selection where each token can be routed to multiple experts. """ - return backend.permute(input, indices, num_out_tokens, workspace, max_expanded_token_num) + return backend.permute( + input, indices, num_out_tokens, workspace, max_expanded_token_num + ) -def unpermute(input: torch.Tensor, - row_id_map: torch.Tensor, - prob: torch.Tensor, - max_tokens: int, - num_topK: int) -> torch.Tensor: +def unpermute( + input: torch.Tensor, + row_id_map: torch.Tensor, + prob: torch.Tensor, + max_tokens: int, + num_topK: int, +) -> torch.Tensor: """ Unpermute expert outputs back to original token order with probability weighting. - + This function reverses the permutation applied in the forward pass and combines outputs from multiple experts using their routing probabilities. - + Args: input (torch.Tensor): Permuted expert outputs to unpermute row_id_map (torch.Tensor): Mapping from permuted positions to original positions prob (torch.Tensor): Expert routing probabilities for weighted combination max_tokens (int): Maximum number of tokens in the sequence num_topK (int): Number of top experts selected per token - + Returns: torch.Tensor: Unpermuted tokens in original order with expert outputs combined - + Note: The output combines multiple expert predictions for each token using the routing probabilities as weights. @@ -178,57 +197,63 @@ def unpermute(input: torch.Tensor, return backend.unpermute(input, row_id_map, prob, max_tokens, num_topK) -def unpermute_bwd(input_bwd: torch.Tensor, - input_fwd: torch.Tensor, - row_id_map: torch.Tensor, - prob: Optional[torch.Tensor]) -> torch.Tensor: +def unpermute_bwd( + input_bwd: torch.Tensor, + input_fwd: torch.Tensor, + row_id_map: torch.Tensor, + prob: Optional[torch.Tensor], +) -> torch.Tensor: """ Backward pass for unpermute operation with gradient flow. - + This function handles the backward pass through the unpermute operation, ensuring proper gradient flow for training MoE models. - + Args: input_bwd (torch.Tensor): Backward gradients from the next layer input_fwd (torch.Tensor): Forward pass inputs (for gradient computation) row_id_map (torch.Tensor): Row mapping used in forward unpermute prob (torch.Tensor, optional): Expert probabilities. If None, uniform weights are used. - + Returns: torch.Tensor: Gradients with respect to the input of unpermute forward pass - + Note: If prob is None, uniform probabilities are assumed for gradient computation. """ # Handle case where probabilities are not provided if prob is None: - prob = torch.ones([input_bwd.size(0), 1], dtype=torch.float32, device=input_bwd.device) - + prob = torch.ones( + [input_bwd.size(0), 1], dtype=torch.float32, device=input_bwd.device + ) + return backend.unpermute_bwd(input_bwd, input_fwd, row_id_map, prob) -def rope(q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - q_out: torch.Tensor, - k_out: torch.Tensor, - mrope_section_doubled: bool) -> None: +def rope( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + mrope_section_doubled: bool, +) -> None: """ Apply RoPE (Rotary Position Embedding) to query and key tensors. - + Applies rotary position embeddings to query and key tensors using precomputed cosine and sine values. Supports both standard RoPE and multi-dimensional RoPE (mRoPE). - + Args: q (torch.Tensor): Query tensor to apply RoPE to - k (torch.Tensor): Key tensor to apply RoPE to + k (torch.Tensor): Key tensor to apply RoPE to cos (torch.Tensor): Precomputed cosine values for rotation sin (torch.Tensor): Precomputed sine values for rotation q_out (torch.Tensor): Output tensor for rotated queries (in-place operation supported) k_out (torch.Tensor): Output tensor for rotated keys (in-place operation supported) mrope_section_doubled (bool): Whether using multi-dimensional RoPE with doubled sections - + Note: This function performs in-place operations if q_out and k_out point to the same memory as q and k respectively. The rotation is applied using the standard @@ -237,21 +262,23 @@ def rope(q: torch.Tensor, return backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled) -def rope_bwd(grad_q_out: torch.Tensor, - grad_k_out: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - grad_q: torch.Tensor, - grad_k: torch.Tensor, - mrope_section_doubled: bool) -> None: +def rope_bwd( + grad_q_out: torch.Tensor, + grad_k_out: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + grad_q: torch.Tensor, + grad_k: torch.Tensor, + mrope_section_doubled: bool, +) -> None: """ Backward pass for RoPE operation with gradient computation. - + Computes gradients with respect to the input query and key tensors for the RoPE operation used in transformer attention mechanisms. - + Args: grad_q_out (torch.Tensor): Gradient with respect to output queries grad_k_out (torch.Tensor): Gradient with respect to output keys @@ -262,9 +289,11 @@ def rope_bwd(grad_q_out: torch.Tensor, grad_q (torch.Tensor): Output tensor for query gradients grad_k (torch.Tensor): Output tensor for key gradients mrope_section_doubled (bool): Whether using multi-dimensional RoPE configuration - + Note: This function computes the analytical gradient of the RoPE operation, which involves the inverse rotation compared to the forward pass. """ - return backend.rope_bwd(grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled) \ No newline at end of file + return backend.rope_bwd( + grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled + ) diff --git a/wall_x/fusions/ops.py b/wall_x/fusions/ops.py index defb39c..3beef5f 100644 --- a/wall_x/fusions/ops.py +++ b/wall_x/fusions/ops.py @@ -5,7 +5,9 @@ from wall_x.fusions import backend class AsymmetricDualExpertGemm(torch.autograd.Function): @staticmethod - def forward(ctx, input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False): + def forward( + ctx, input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False + ): """ Forward pass for asymmetric dual expert GEMM. @@ -27,14 +29,24 @@ class AsymmetricDualExpertGemm(torch.autograd.Function): # Dimension validation depends on trans_b if trans_b: - assert input_expert0.size(1) == weight_expert0.size(1), "Expert 0 dimension mismatch (trans_b=True)" - assert input_expert1.size(1) == weight_expert1.size(1), "Expert 1 dimension mismatch (trans_b=True)" + assert input_expert0.size(1) == weight_expert0.size( + 1 + ), "Expert 0 dimension mismatch (trans_b=True)" + assert input_expert1.size(1) == weight_expert1.size( + 1 + ), "Expert 1 dimension mismatch (trans_b=True)" else: - assert input_expert0.size(1) == weight_expert0.size(0), "Expert 0 dimension mismatch (trans_b=False)" - assert input_expert1.size(1) == weight_expert1.size(0), "Expert 1 dimension mismatch (trans_b=False)" + assert input_expert0.size(1) == weight_expert0.size( + 0 + ), "Expert 0 dimension mismatch (trans_b=False)" + assert input_expert1.size(1) == weight_expert1.size( + 0 + ), "Expert 1 dimension mismatch (trans_b=False)" # Save tensors and trans_b for backward pass - ctx.save_for_backward(input_expert0, input_expert1, weight_expert0, weight_expert1) + ctx.save_for_backward( + input_expert0, input_expert1, weight_expert0, weight_expert1 + ) ctx.trans_b = trans_b # Allocate output tensors @@ -43,11 +55,23 @@ class AsymmetricDualExpertGemm(torch.autograd.Function): n0 = weight_expert0.size(0) if trans_b else weight_expert0.size(1) n1 = weight_expert1.size(0) if trans_b else weight_expert1.size(1) - output_expert0 = torch.empty(m0, n0, device=input_expert0.device, dtype=input_expert0.dtype) - output_expert1 = torch.empty(m1, n1, device=input_expert1.device, dtype=input_expert1.dtype) + output_expert0 = torch.empty( + m0, n0, device=input_expert0.device, dtype=input_expert0.dtype + ) + output_expert1 = torch.empty( + m1, n1, device=input_expert1.device, dtype=input_expert1.dtype + ) # Call the backend C++ function - backend.asym_dual_gmm_separated(input_expert0, input_expert1, weight_expert0, weight_expert1, output_expert0, output_expert1, trans_b=trans_b) + backend.asym_dual_gmm_separated( + input_expert0, + input_expert1, + weight_expert0, + weight_expert1, + output_expert0, + output_expert1, + trans_b=trans_b, + ) return output_expert0, output_expert1 @@ -111,10 +135,18 @@ class AsymmetricDualExpertGemm(torch.autograd.Function): trans_b=False, ) - return grad_input_expert0, grad_input_expert1, grad_weight_expert0, grad_weight_expert1, None + return ( + grad_input_expert0, + grad_input_expert1, + grad_weight_expert0, + grad_weight_expert1, + None, + ) -def asym_dual_gmm(input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False): +def asym_dual_gmm( + input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False +): """ Convenience function for asymmetric dual expert GEMM. @@ -128,7 +160,9 @@ def asym_dual_gmm(input_expert0, input_expert1, weight_expert0, weight_expert1, Returns: Tuple of (output_expert0, output_expert1) """ - return AsymmetricDualExpertGemm.apply(input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b) + return AsymmetricDualExpertGemm.apply( + input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b + ) ################################################################################################ @@ -145,7 +179,13 @@ class PermuteMoE_topK(torch.autograd.Function): max_expanded_token_num = 0 @staticmethod - def forward(ctx, input_act: torch.Tensor, indices: torch.Tensor, num_out_tokens: int, max_token_num: int): + def forward( + ctx, + input_act: torch.Tensor, + indices: torch.Tensor, + num_out_tokens: int, + max_token_num: int, + ): """ indices: for topK=1, indices in a 1-d tensor of shape [num_tokens], otherwise, it's a 2-d tensor of shape [num_tokens, topK] @@ -160,18 +200,27 @@ class PermuteMoE_topK(torch.autograd.Function): # Device check if input_act.is_cpu: - raise RuntimeError("[Error] The input `input_act` of permute_topK op is on the device: CPU!") + raise RuntimeError( + "[Error] The input `input_act` of permute_topK op is on the device: CPU!" + ) if indices.is_cpu: - warnings.warn("The input `indices` of permute_topK op is on the device: CPU!") - expert_for_rows = expert_for_rows.cuda() + warnings.warn( + "The input `indices` of permute_topK op is on the device: CPU!" + ) # Shape check if input_act.size(0) != indices.size(0): - raise RuntimeError(f"[Error] permute_topK op input `indices` shape mismatch! " f"Expect {input_act.size(0)}, but got {indices.size(0)}.") + raise RuntimeError( + f"[Error] permute_topK op input `indices` shape mismatch! " + f"Expect {input_act.size(0)}, but got {indices.size(0)}." + ) # Data type check if indices.dtype != torch.int32: - warnings.warn(f"The data type of the input `indices` of permute_topK op is {indices.dtype}! " "The recommended type is torch.int32.") + warnings.warn( + f"The data type of the input `indices` of permute_topK op is {indices.dtype}! " + "The recommended type is torch.int32." + ) indices = indices.to(torch.int32) # Contiguous check @@ -194,7 +243,11 @@ class PermuteMoE_topK(torch.autograd.Function): PermuteMoE_topK.workspace_fw = [] permuted_act, row_id_map, PermuteMoE_topK.workspace_fw = backend.permute( - input_act, indices, num_out_tokens, PermuteMoE_topK.workspace_fw, PermuteMoE_topK.max_expanded_token_num + input_act, + indices, + num_out_tokens, + PermuteMoE_topK.workspace_fw, + PermuteMoE_topK.max_expanded_token_num, ) ctx.row_id_map = row_id_map @@ -215,7 +268,9 @@ class PermuteMoE_topK(torch.autograd.Function): num_tokens = ctx.num_tokens num_topK = ctx.num_topK - unpermuted_act_grad = backend.unpermute(permuted_act_grad, row_id_map, torch.tensor([]), num_tokens, num_topK) + unpermuted_act_grad = backend.unpermute( + permuted_act_grad, row_id_map, torch.tensor([]), num_tokens, num_topK + ) return unpermuted_act_grad, None, None, None @@ -229,7 +284,12 @@ class PermuteMoE_topK(torch.autograd.Function): class UnpermuteMoE_topK(torch.autograd.Function): @staticmethod - def forward(ctx, input_act: torch.Tensor, row_id_map: torch.Tensor, probs: torch.Tensor = None): + def forward( + ctx, + input_act: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor = None, + ): # Empty input check if not input_act.numel(): ctx.probs = probs @@ -237,36 +297,51 @@ class UnpermuteMoE_topK(torch.autograd.Function): # Device check if input_act.is_cpu: - raise RuntimeError("[Error] The input `input_act` of unpermute_topK op is on the device: CPU!") + raise RuntimeError( + "[Error] The input `input_act` of unpermute_topK op is on the device: CPU!" + ) if row_id_map.is_cpu: - warnings.warn("The input `row_id_map` of unpermute_topK op is on the device: CPU!") + warnings.warn( + "The input `row_id_map` of unpermute_topK op is on the device: CPU!" + ) row_id_map = row_id_map.cuda() if probs is not None and probs.is_cpu: - warnings.warn("The input `probs` of unpermute_topK op is on the device: CPU!") + warnings.warn( + "The input `probs` of unpermute_topK op is on the device: CPU!" + ) probs = probs.cuda() # Shape check if probs is not None and row_id_map.size(0) != probs.size(0) * probs.size(1): raise RuntimeError( - f"[Error] unpermute_topK op input `probs` shape mismatch! " f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}." + f"[Error] unpermute_topK op input `probs` shape mismatch! " + f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}." ) # Data type check if row_id_map.dtype != torch.int32: warnings.warn( - f"The data type of the input `row_id_map` of unpermute_topK op is {row_id_map.dtype}! " "The recommended type is torch.int32." + f"The data type of the input `row_id_map` of unpermute_topK op is {row_id_map.dtype}! " + "The recommended type is torch.int32." ) row_id_map = row_id_map.to(torch.int32) if probs is not None and probs.dtype != torch.float32: - warnings.warn(f"The data type of the input `probs` of unpermute_topK op is {probs.dtype}! " "The recommended type is torch.float32.") + warnings.warn( + f"The data type of the input `probs` of unpermute_topK op is {probs.dtype}! " + "The recommended type is torch.float32." + ) probs = probs.to(torch.float32) # Contiguous check if not input_act.is_contiguous(): - warnings.warn("The input `input_act` of unpermute_topK op is discontiguous!") + warnings.warn( + "The input `input_act` of unpermute_topK op is discontiguous!" + ) input_act = input_act.contiguous() if not row_id_map.is_contiguous(): - warnings.warn("The input `row_id_map` of unpermute_topK op is discontiguous!") + warnings.warn( + "The input `row_id_map` of unpermute_topK op is discontiguous!" + ) row_id_map = row_id_map.contiguous() if probs is not None and not probs.is_contiguous(): warnings.warn("The input `probs` of unpermute_topK op is discontiguous!") @@ -275,7 +350,13 @@ class UnpermuteMoE_topK(torch.autograd.Function): num_tokens = probs.size(0) if probs is not None else input_act.size(0) num_topK = probs.size(1) if probs is not None else 1 - unpermuted_output = backend.unpermute(input_act, row_id_map, probs if probs is not None else torch.tensor([]), num_tokens, num_topK) + unpermuted_output = backend.unpermute( + input_act, + row_id_map, + probs if probs is not None else torch.tensor([]), + num_tokens, + num_topK, + ) ctx.save_for_backward(input_act, row_id_map, probs) return unpermuted_output @@ -293,7 +374,9 @@ class UnpermuteMoE_topK(torch.autograd.Function): act_grad = None if ctx.needs_input_grad[0]: - act_grad, prob_grad = backend.unpermute_bwd(unpermuted_act_grad, input_act, row_id_map, probs) + act_grad, prob_grad = backend.unpermute_bwd( + unpermuted_act_grad, input_act, row_id_map, probs + ) if not ctx.needs_input_grad[2]: prob_grad = None @@ -319,19 +402,36 @@ def unpermute(input_act, row_id_map, probs=None): class MultimodalRoPE(torch.autograd.Function): @staticmethod - def forward(ctx, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, mrope_section: list): + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + mrope_section: list, + ): # Device check if q.is_cpu: - raise RuntimeError("[Error] The input `q` of multimodal_rope op is on the device: CPU!") + raise RuntimeError( + "[Error] The input `q` of multimodal_rope op is on the device: CPU!" + ) if k.is_cpu: - raise RuntimeError("[Error] The input `k` of multimodal_rope op is on the device: CPU!") + raise RuntimeError( + "[Error] The input `k` of multimodal_rope op is on the device: CPU!" + ) if cos.is_cpu: - raise RuntimeError("[Error] The input `cos` of multimodal_rope op is on the device: CPU!") + raise RuntimeError( + "[Error] The input `cos` of multimodal_rope op is on the device: CPU!" + ) if sin.is_cpu: - raise RuntimeError("[Error] The input `sin` of multimodal_rope op is on the device: CPU!") + raise RuntimeError( + "[Error] The input `sin` of multimodal_rope op is on the device: CPU!" + ) if len(mrope_section) != 3: - raise RuntimeError("[Error] The input `mrope_section` of multimodal_rope op must be a list of 3 integers!") + raise RuntimeError( + "[Error] The input `mrope_section` of multimodal_rope op must be a list of 3 integers!" + ) # Contiguous check if not q.is_contiguous(): diff --git a/wall_x/model/action_head.py b/wall_x/model/action_head.py index 5a29631..437b795 100755 --- a/wall_x/model/action_head.py +++ b/wall_x/model/action_head.py @@ -1,23 +1,23 @@ - import math import torch import torch.nn as nn from torch.distributions import Beta from wall_x.utils.constant import action_statistic_dof + class Normalizer(nn.Module): """ Action data normalizer for multi-robot systems. - + This module handles normalization and denormalization of action data for different robot configurations. It maintains per-robot statistics (min values and deltas) and applies normalization to map actions to the [-1, 1] range. """ - + def __init__(self, action_statistic_dof, dof_config): """ Initialize the normalizer with robot-specific action statistics. - + Args: action_statistic_dof (dict): Statistical data for each robot's degrees of freedom dof_config (dict): Configuration mapping for degrees of freedom per robot @@ -25,13 +25,13 @@ class Normalizer(nn.Module): super(Normalizer, self).__init__() action_statistic = {} - + # Process statistics for each robot for robot_name in action_statistic_dof.keys(): action_statistic[robot_name] = {} all_dof_min = [] all_dof_delta = [] - + # Collect min and delta values for all DOFs for k in dof_config: if k in action_statistic_dof[robot_name]: @@ -41,37 +41,41 @@ class Normalizer(nn.Module): # Use default values if statistics not available all_dof_min.extend([0.0] * dof_config[k]) all_dof_delta.extend([1.0] * dof_config[k]) - + all_dof_min = torch.tensor(all_dof_min) all_dof_delta = torch.tensor(all_dof_delta) action_statistic[robot_name]["min"] = all_dof_min action_statistic[robot_name]["delta"] = all_dof_delta # Register statistics as non-trainable parameters - self.min = nn.ParameterDict({ - k: nn.Parameter(action_statistic[k]["min"], requires_grad=False) - for k in action_statistic.keys() - }) - self.delta = nn.ParameterDict({ - k: nn.Parameter(action_statistic[k]["delta"], requires_grad=False) - for k in action_statistic.keys() - }) + self.min = nn.ParameterDict( + { + k: nn.Parameter(action_statistic[k]["min"], requires_grad=False) + for k in action_statistic.keys() + } + ) + self.delta = nn.ParameterDict( + { + k: nn.Parameter(action_statistic[k]["delta"], requires_grad=False) + for k in action_statistic.keys() + } + ) def normalize_data(self, xs, dataset_names): """ Normalize action data to [-1, 1] range using robot-specific statistics. - + Args: xs: Input action data tensors dataset_names: List of dataset/robot names corresponding to each tensor - + Returns: torch.Tensor: Normalized action data in [-1, 1] range """ new_xs = [] # Filter out multimodal dataset entries dataset_names = [name for name in dataset_names if name != "x2_multimodal"] - + for x, dataset_name in zip(xs, dataset_names): # Apply min-max normalization x = (x - self.min[dataset_name]) / (self.delta[dataset_name]) @@ -80,19 +84,19 @@ class Normalizer(nn.Module): # Clamp to ensure bounds x = torch.clamp(x, -1, 1) new_xs.append(x) - + new_xs = torch.stack(new_xs) return new_xs def unnormalize_data(self, xs, dataset_names, dof_mask=None): """ Convert normalized data back to original action space. - + Args: xs: Normalized action data in [-1, 1] range dataset_names: List of dataset/robot names dof_mask: Optional mask to select specific degrees of freedom - + Returns: torch.Tensor: Denormalized action data in original scale """ @@ -100,11 +104,11 @@ class Normalizer(nn.Module): # Filter out multimodal dataset entries dataset_names = [name for name in dataset_names if name != "x2_multimodal"] dof_mask = dof_mask if dof_mask is not None else [None] * len(xs) - + for x, dataset_name, mask in zip(xs, dataset_names, dof_mask): # Convert from [-1, 1] to [0, 1] range x = (x + 1) / 2 - + # Apply DOF mask if provided if mask is not None: mask = mask[0].bool() @@ -113,11 +117,11 @@ class Normalizer(nn.Module): else: action_space_delta = self.delta[dataset_name] action_space_min = self.min[dataset_name] - + # Scale back to original range x = x * action_space_delta + action_space_min new_xs.append(x) - + new_xs = torch.stack(new_xs) return new_xs @@ -125,15 +129,15 @@ class Normalizer(nn.Module): class SinusoidalPosEmb(nn.Module): """ Sinusoidal positional embedding for diffusion timesteps. - + Generates sinusoidal embeddings commonly used in diffusion models to encode timestep information with different frequencies. """ - + def __init__(self, dim): """ Initialize sinusoidal positional embedding. - + Args: dim (int): Embedding dimension (must be even) """ @@ -143,10 +147,10 @@ class SinusoidalPosEmb(nn.Module): def forward(self, x): """ Generate sinusoidal embeddings for input timesteps. - + Args: x (torch.Tensor): Input timesteps - + Returns: torch.Tensor: Sinusoidal embeddings of shape (..., dim) """ @@ -159,25 +163,24 @@ class SinusoidalPosEmb(nn.Module): return emb - class ActionProcessor(nn.Module): """ Action sequence processor for robotic control with flow matching. - + This module handles action sequence processing for robotic systems with the following capabilities: 1. Adds controlled noise to action sequences using Beta distribution scheduling 2. Generates temporal embeddings for timestep conditioning 3. Projects actions to model hidden space for transformer processing 4. Supports proprioceptive data integration and multi-robot configurations - + The Beta distribution provides more flexible noise injection strategies compared to traditional linear schedules, allowing better control over the noise scheduling process. """ - + def __init__(self, config): """ Initialize the action processor with multi-robot support. - + Args: config: Configuration object containing: - dof_config (dict): Degrees of freedom configuration per robot type @@ -186,7 +189,7 @@ class ActionProcessor(nn.Module): - noise_scheduler (dict): Noise scheduler configuration with Beta parameters """ super().__init__() - + # Calculate action and proprioception dimensions from configuration self.dof_config = config.dof_config self.agent_pos_config = config.agent_pos_config @@ -203,22 +206,28 @@ class ActionProcessor(nn.Module): print(" Agent position configuration:", flush=True) for key, value in self.agent_pos_config.items(): print(f" {key}: {value}", flush=True) - + self.hidden_size = config.hidden_size # Initialize data normalizers for actions and proprioception self.normalizer_action = Normalizer(action_statistic_dof, config.dof_config) - self.normalizer_propri = Normalizer(action_statistic_dof, config.agent_pos_config) + self.normalizer_propri = Normalizer( + action_statistic_dof, config.agent_pos_config + ) # Proprioception projection layer (includes history/current state) self.propri_proj = nn.Linear(self.propri_dim * 2, self.hidden_size, bias=False) # Beta distribution noise scheduler configuration noise_scheduler_config = config.noise_scheduler - self.beta_alpha = noise_scheduler_config.get('beta_alpha', 1.5) # Beta distribution α parameter - self.beta_beta = noise_scheduler_config.get('beta_beta', 1.0) # Beta distribution β parameter - self.s = noise_scheduler_config.get('s', 0.999) # Scaling factor - + self.beta_alpha = noise_scheduler_config.get( + "beta_alpha", 1.5 + ) # Beta distribution α parameter + self.beta_beta = noise_scheduler_config.get( + "beta_beta", 1.0 + ) # Beta distribution β parameter + self.s = noise_scheduler_config.get("s", 0.999) # Scaling factor + # Initialize Beta distribution for noise scheduling alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda") beta_tensor = torch.tensor(self.beta_beta, dtype=torch.float32).to("cuda") @@ -228,51 +237,59 @@ class ActionProcessor(nn.Module): self.time_embed = SinusoidalPosEmb(config.hidden_size) # Action embedding network: project to hidden space - self.w1 = nn.Linear(self.action_dim * 2, self.hidden_size, bias=False) # *2 for action + DOF mask - self.w2 = nn.Linear(self.hidden_size * 2, self.hidden_size, bias=False) # *2 for action + time embeddings + self.w1 = nn.Linear( + self.action_dim * 2, self.hidden_size, bias=False + ) # *2 for action + DOF mask + self.w2 = nn.Linear( + self.hidden_size * 2, self.hidden_size, bias=False + ) # *2 for action + time embeddings self.w3 = nn.Linear(self.hidden_size, self.hidden_size, bias=False) self.act_fn = nn.SiLU() - + # Project back to action space for flow matching loss self.action_proj_back = nn.Linear(self.hidden_size, self.action_dim, bias=False) - self.mse_loss = nn.MSELoss(reduction='none') + self.mse_loss = nn.MSELoss(reduction="none") def sample_time(self, batch_size, device, dtype): """ Sample timesteps using Beta distribution for noise scheduling. - + Generates random timesteps in [0,1] range using Beta distribution, then scales them. This provides more flexible control over the noise injection schedule compared to uniform sampling. - + Args: batch_size (int): Number of timesteps to sample device: Target device for tensors dtype: Target data type for tensors - + Returns: torch.Tensor: Sampled timesteps of shape [batch_size] """ sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype) time = (self.s - sample) / self.s return time - - def proprioception_proj(self, proprioception, dataset_names=None, dof_mask=None, use_history=False): + + def proprioception_proj( + self, proprioception, dataset_names=None, dof_mask=None, use_history=False + ): """ Project proprioceptive data (joint positions, orientations) to hidden space. - + Args: proprioception (torch.Tensor): Proprioceptive data of shape [batch_size, seq_len, propri_dim] dataset_names (list, optional): Dataset names for normalization. Defaults to None. dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, propri_dim]. Defaults to None. use_history (bool, optional): Whether to use historical proprioceptive data. Defaults to False. - + Returns: torch.Tensor: Projected proprioceptive features of shape [batch_size, seq_len, hidden_size] """ # Ensure proper device and dtype alignment - proprioception = proprioception.to(device=self.propri_proj.weight.device).to(dtype=self.propri_proj.weight.dtype) - + proprioception = proprioception.to(device=self.propri_proj.weight.device).to( + dtype=self.propri_proj.weight.dtype + ) + if dof_mask is not None: # Concatenate proprioception with DOF mask # TODO: Use variable-based dimension checking for better flexibility @@ -280,26 +297,28 @@ class ActionProcessor(nn.Module): proprioception = torch.cat([proprioception, dof_mask], dim=-1) else: proprioception = torch.cat([proprioception, dof_mask], dim=-1) - - proprioception = proprioception.to(device=self.propri_proj.weight.device).to(dtype=self.propri_proj.weight.dtype) + + proprioception = proprioception.to(device=self.propri_proj.weight.device).to( + dtype=self.propri_proj.weight.dtype + ) return self.propri_proj(proprioception) def forward(self, action_chunk, dataset_names, dof_mask=None): """ Process action sequences with noise injection and temporal embedding. - + This method implements the forward pass for flow matching training: 1. Adds Beta-distributed noise to action sequences 2. Generates sinusoidal timestep embeddings 3. Projects noisy actions to hidden space 4. Combines action and temporal features - + Args: action_chunk (torch.Tensor): Action sequences of shape [batch_size, seq_len, action_dim] dataset_names (list): Dataset names for normalization - dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, seq_len, action_dim]. + dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, seq_len, action_dim]. Defaults to None. - + Returns: tuple: (action_embeddings, flow_target) where: - action_embeddings: Processed action features of shape [batch_size, seq_len, hidden_size] @@ -324,48 +343,54 @@ class ActionProcessor(nn.Module): # 3. Project noisy actions with DOF mask to hidden space if dof_mask is not None: noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - + noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) action_embed = self.w1(noisy_action) - + # Repeat time embedding for each sequence position - time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1).to(dtype=self.w2.weight.dtype) - + time_embed = ( + time_embed.unsqueeze(1) + .repeat(1, action_embed.shape[1], 1) + .to(dtype=self.w2.weight.dtype) + ) + # Combine action and temporal embeddings concat_embed = torch.cat([action_embed, time_embed], dim=-1) concat_embed = self.w2(concat_embed) embed = self.w3(self.act_fn(concat_embed)) return embed, flow - + def step(self, timestep, noisy_action, dof_mask=None): """ Single denoising step for diffusion inference. - + Processes noisy actions at a specific timestep for iterative denoising during inference. - + Args: timestep (torch.Tensor): Current timesteps of shape [batch_size] noisy_action (torch.Tensor): Noisy actions of shape [batch_size, seq_len, action_dim] dof_mask (torch.Tensor, optional): DOF mask for action space. Defaults to None. - + Returns: torch.Tensor: Processed action embeddings of shape [batch_size, seq_len, hidden_size] """ # Concatenate noisy action with DOF mask if provided if dof_mask is not None: noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - + # Generate timestep embeddings time_embed = self.time_embed(timestep) # [batch_size, hidden_size] - + # Project noisy actions action_embed = self.w1(noisy_action) - + # Broadcast time embeddings to sequence length time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) - time_embed = time_embed.to(device=noisy_action.device).to(dtype=noisy_action.dtype) - + time_embed = time_embed.to(device=noisy_action.device).to( + dtype=noisy_action.dtype + ) + # Combine embeddings and process through MLP concat_embed = torch.cat([action_embed, time_embed], dim=-1) concat_embed = self.w2(concat_embed) @@ -376,25 +401,25 @@ class ActionProcessor(nn.Module): def flow_loss(self, action_hidden_states, flow, dof_mask=None): """ Compute flow matching loss between predicted and target actions. - + Args: action_hidden_states (torch.Tensor): Hidden states from transformer flow (torch.Tensor): Target flow (action - noise) for matching dof_mask (torch.Tensor, optional): DOF mask to weight loss per dimension. Defaults to None. - + Returns: torch.Tensor: Flow matching loss (no reduction for channel loss computation) """ # Project hidden states back to action space action_pred = self.action_proj_back(action_hidden_states) - + # Compute MSE loss between predicted and target flow loss = self.mse_loss(action_pred, flow) - + # Apply DOF mask if provided if dof_mask is not None: dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) loss = loss * dof_mask - + # Return loss without reduction for channel-wise loss computation - return loss \ No newline at end of file + return loss diff --git a/wall_x/model/qwen2_5_based/__init__.py b/wall_x/model/qwen2_5_based/__init__.py index dd5b053..d9ed878 100644 --- a/wall_x/model/qwen2_5_based/__init__.py +++ b/wall_x/model/qwen2_5_based/__init__.py @@ -1,2 +1,8 @@ -from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel,Qwen2_5_VLMoEForAction -from .configuration_qwen2_5_vl import Qwen2_5_VLConfig \ No newline at end of file +from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel, Qwen2_5_VLMoEForAction +from .configuration_qwen2_5_vl import Qwen2_5_VLConfig + +__all__ = [ + "Qwen2_5_VLMoEModel", + "Qwen2_5_VLMoEForAction", + "Qwen2_5_VLConfig", +] diff --git a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py b/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py index 8420b80..439a369 100644 --- a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py +++ b/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py @@ -190,7 +190,7 @@ class Qwen2_5_VLConfig(PretrainedConfig): experts=None, dof_config=None, noise_scheduler=None, - dim_inputs=(1536,1536), + dim_inputs=(1536, 1536), attention_moe=False, mlp_moe=False, **kwargs, diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py index e0d891f..92f8f95 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py +++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py @@ -1,4 +1,3 @@ - import math import torch import torch.nn as nn @@ -8,7 +7,12 @@ from torch.nn import CrossEntropyLoss from typing import Any, Dict, List, Optional, Tuple, Union from transformers.activations import ACT2FN -from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache +from transformers.cache_utils import ( + Cache, + DynamicCache, + SlidingWindowCache, + StaticCache, +) from transformers.generation import GenerationMixin from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput @@ -29,13 +33,15 @@ from wall_x.fusions import ops if is_flash_attn_2_available(): from flash_attn import flash_attn_varlen_func from flash_attn.layers.rotary import apply_rotary_emb + from flash_attn import flash_attn_func else: flash_attn_varlen_func = None apply_rotary_emb = None + flash_attn_func = None if is_flash_attn_2_available(): - from transformers.modeling_flash_attention_utils import _flash_attention_forward + pass else: flash_attn_varlen_func = None @@ -56,7 +62,9 @@ class Qwen2_5_VLMLP(nn.Module): self.act_fn = ACT2FN[config.hidden_act] def forward(self, hidden_state): - return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + return self.down_proj( + self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state) + ) class Qwen2_5_VisionPatchEmbed(nn.Module): @@ -74,14 +82,26 @@ class Qwen2_5_VisionPatchEmbed(nn.Module): self.embed_dim = embed_dim kernel_size = [temporal_patch_size, patch_size, patch_size] - self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False) + self.proj = nn.Conv3d( + in_channels, + embed_dim, + kernel_size=kernel_size, + stride=kernel_size, + bias=False, + ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: target_dtype = self.proj.weight.dtype hidden_states = hidden_states.view( - -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( + -1, self.embed_dim ) - hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) return hidden_states @@ -92,7 +112,9 @@ class Qwen2_5_VisionRotaryEmbedding(nn.Module): self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, seqlen: int) -> torch.Tensor: - seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + seq = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) freqs = torch.outer(seq, self.inv_freq) return freqs @@ -159,7 +181,12 @@ class Qwen2_5_VLVisionFlashAttention2(nn.Module): position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] - q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + q, k, v = ( + self.qkv(hidden_states) + .reshape(seq_length, 3, self.num_heads, -1) + .permute(1, 0, 2, 3) + .unbind(0) + ) if position_embeddings is None: logger.warning_once( "The attention layers in this model are transitioning from computing the RoPE embeddings internally " @@ -178,9 +205,9 @@ class Qwen2_5_VLVisionFlashAttention2(nn.Module): if max_seqlen is None: max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen).reshape( - seq_length, -1 - ) + attn_output = flash_attn_varlen_func( + q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen + ).reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -223,7 +250,12 @@ class Qwen2_5_VLVisionAttention(nn.Module): position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] - q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + q, k, v = ( + self.qkv(hidden_states) + .reshape(seq_length, 3, self.num_heads, -1) + .permute(1, 0, 2, 3) + .unbind(0) + ) if position_embeddings is None: logger.warning_once( "The attention layers in this model are transitioning from computing the RoPE embeddings internally " @@ -239,17 +271,26 @@ class Qwen2_5_VLVisionAttention(nn.Module): q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) attention_mask = torch.full( - [1, seq_length, seq_length], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype + [1, seq_length, seq_length], + torch.finfo(q.dtype).min, + device=q.device, + dtype=q.dtype, ) for i in range(1, len(cu_seqlens)): - attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0 + attention_mask[ + ..., + cu_seqlens[i - 1] : cu_seqlens[i], + cu_seqlens[i - 1] : cu_seqlens[i], + ] = 0 q = q.transpose(0, 1) k = k.transpose(0, 1) v = v.transpose(0, 1) attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim) attn_weights = attn_weights + attention_mask - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(q.dtype) attn_output = torch.matmul(attn_weights, v) attn_output = attn_output.transpose(0, 1) attn_output = attn_output.reshape(seq_length, -1) @@ -273,7 +314,12 @@ class Qwen2_5_VLVisionSdpaAttention(nn.Module): position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] - q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + q, k, v = ( + self.qkv(hidden_states) + .reshape(seq_length, 3, self.num_heads, -1) + .permute(1, 0, 2, 3) + .unbind(0) + ) if position_embeddings is None: logger.warning_once( "The attention layers in this model are transitioning from computing the RoPE embeddings internally " @@ -288,13 +334,21 @@ class Qwen2_5_VLVisionSdpaAttention(nn.Module): cos, sin = position_embeddings q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) - attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool) + attention_mask = torch.zeros( + [1, seq_length, seq_length], device=q.device, dtype=torch.bool + ) for i in range(1, len(cu_seqlens)): - attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = True + attention_mask[ + ..., + cu_seqlens[i - 1] : cu_seqlens[i], + cu_seqlens[i - 1] : cu_seqlens[i], + ] = True q = q.transpose(0, 1) k = k.transpose(0, 1) v = v.transpose(0, 1) - attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0) + attn_output = F.scaled_dot_product_attention( + q, k, v, attention_mask, dropout_p=0.0 + ) attn_output = attn_output.transpose(0, 1) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) @@ -404,7 +458,10 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) self.blocks = nn.ModuleList( - [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)] + [ + Qwen2_5_VLVisionBlock(config, config._attn_implementation) + for _ in range(config.depth) + ] ) self.merger = Qwen2_5_VLPatchMerger( dim=config.out_hidden_size, @@ -446,14 +503,18 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): window_index: list = [] cu_window_seqlens: list = [0] window_index_id = 0 - vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size + vit_merger_window_size = ( + self.window_size // self.spatial_merge_size // self.patch_size + ) for grid_t, grid_h, grid_w in grid_thw: llm_grid_h, llm_grid_w = ( grid_h // self.spatial_merge_size, grid_w // self.spatial_merge_size, ) - index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape( + grid_t, llm_grid_h, llm_grid_w + ) pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size @@ -476,14 +537,18 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): index_padded = index_padded.reshape(-1) index_new = index_padded[index_padded != -100] window_index.append(index_new + window_index_id) - cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + cu_seqlens_tmp = ( + seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + ) cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) window_index_id += (grid_t * llm_grid_h * llm_grid_w).item() window_index = torch.cat(window_index, dim=0) return window_index, cu_window_seqlens - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor + ) -> torch.Tensor: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -506,16 +571,22 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) seq_len, _ = hidden_states.size() - hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + hidden_states = hidden_states.reshape( + seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 + ) hidden_states = hidden_states[window_index, :, :] hidden_states = hidden_states.reshape(seq_len, -1) - rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb.reshape( + seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 + ) rotary_pos_emb = rotary_pos_emb[window_index, :, :] rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) position_embeddings = (emb.cos(), emb.sin()) - cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum( dim=0, # Select dtype based on the following factors: # - FA2 requires that cu_seqlens_q must have dtype int32 @@ -525,7 +596,9 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): ) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) max_seqlen_full = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - max_seqlen_window = (cu_window_seqlens[1:] - cu_window_seqlens[:-1]).max().item() + max_seqlen_window = ( + (cu_window_seqlens[1:] - cu_window_seqlens[:-1]).max().item() + ) for layer_num, blk in enumerate(self.blocks): if layer_num in self.fullatt_block_indexes: @@ -536,10 +609,19 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): max_seqlen_now = max_seqlen_window if self.gradient_checkpointing and self.training: hidden_states = self._gradient_checkpointing_func( - blk.__call__, hidden_states, cu_seqlens_now, None, position_embeddings + blk.__call__, + hidden_states, + cu_seqlens_now, + None, + position_embeddings, ) else: - hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens_now, max_seqlen=max_seqlen_now, position_embeddings=position_embeddings) + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens_now, + max_seqlen=max_seqlen_now, + position_embeddings=position_embeddings, + ) hidden_states = self.merger(hidden_states) reverse_indices = torch.argsort(window_index) @@ -553,7 +635,9 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): super().__init__() # BC: "rope_type" was originally "type" if hasattr(config, "rope_scaling") and config.rope_scaling is not None: - self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + self.rope_type = config.rope_scaling.get( + "rope_type", config.rope_scaling.get("type") + ) else: self.rope_type = "default" self.max_seq_len_cached = config.max_position_embeddings @@ -577,10 +661,15 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): inv_freq, self.attention_scaling = self.rope_init_fn( self.config, device, seq_len=seq_len, **self.rope_kwargs ) - self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.register_buffer( + "inv_freq", inv_freq, persistent=False + ) # TODO joao: may break with compilation self.max_seq_len_cached = seq_len - if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + if ( + seq_len < self.original_max_seq_len + and self.max_seq_len_cached > self.original_max_seq_len + ): # reset self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len @@ -591,13 +680,25 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): # Core RoPE block. In contrast to other models, Qwen2_5_VL has different position ids for thw grids # So we expand the inv_freq to shape (3, ...) - inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) - position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + inv_freq_expanded = ( + self.inv_freq[None, None, :, None] + .float() + .expand(3, position_ids.shape[1], -1, 1) + ) + position_ids_expanded = position_ids[ + :, :, None, : + ].float() # shape (3, bs, 1, positions) # Force float32 (see https://github.com/huggingface/transformers/pull/29285) device_type = x.device.type - device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + device_type = ( + device_type + if isinstance(device_type, str) and device_type != "mps" + else "cpu" + ) with torch.autocast(device_type=device_type, enabled=False): - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(2, 3) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() @@ -658,12 +759,12 @@ def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ mrope_section = mrope_section * 2 - cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim - ) - sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim - ) + cos = torch.cat( + [m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1 + ).unsqueeze(unsqueeze_dim) + sin = torch.cat( + [m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1 + ).unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) @@ -678,7 +779,9 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) @@ -713,10 +816,18 @@ class Qwen2_5_VLAttention(nn.Module): f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) - self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) - self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) - self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) - self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.q_proj = nn.Linear( + self.hidden_size, self.num_heads * self.head_dim, bias=True + ) + self.k_proj = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True + ) + self.v_proj = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, self.hidden_size, bias=False + ) self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) @@ -729,7 +840,9 @@ class Qwen2_5_VLAttention(nn.Module): output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -747,14 +860,22 @@ class Qwen2_5_VLAttention(nn.Module): ) if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) # repeat k/v heads if n_kv_heads < n_heads key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + attn_weights = torch.matmul( + query_states, key_states.transpose(2, 3) + ) / math.sqrt(self.head_dim) if attention_mask is not None: # no matter the length, we just slice it causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] @@ -763,11 +884,17 @@ class Qwen2_5_VLAttention(nn.Module): # Fix precision issues in Qwen2-VL float16 inference # Replace inf values with zeros in attention weights to prevent NaN propagation if query_states.dtype == torch.float16: - attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights) + attn_weights = torch.where( + torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights + ) # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(query_states.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.attention_dropout, training=self.training + ) attn_output = torch.matmul(attn_weights, value_states) if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): @@ -786,7 +913,6 @@ class Qwen2_5_VLAttention(nn.Module): return attn_output, attn_weights, past_key_value -from flash_attn import flash_attn_func class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): """ @@ -814,7 +940,9 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC ): bsz, q_len, _ = hidden_states.size() @@ -829,10 +957,18 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): # Because the input can be padded, the absolute sequence length depends on the max position id. cos, sin = position_embeddings - query_states, key_states = ops.multimodal_rope(query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]) + query_states, key_states = ops.multimodal_rope( + query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + ) if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) # repeat k/v heads if n_kv_heads < n_heads # key_states = repeat_kv(key_states, self.num_key_value_groups) @@ -868,7 +1004,12 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): value_states = value_states.transpose(1, 2) attn_output = flash_attn_func( - query_states, key_states, value_states, dropout_rate, softmax_scale=None, causal=self.is_causal + query_states, + key_states, + value_states, + dropout_rate, + softmax_scale=None, + causal=self.is_causal, ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() @@ -897,7 +1038,9 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if output_attentions: # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. @@ -932,8 +1075,14 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): ) if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) @@ -983,16 +1132,23 @@ class Qwen2_5_VLDecoderLayer(nn.Module): super().__init__() self.hidden_size = config.hidden_size - if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + if ( + config.use_sliding_window + and config._attn_implementation != "flash_attention_2" + ): logger.warning_once( f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " "unexpected results may be encountered." ) - self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation]( + config, layer_idx + ) self.mlp = Qwen2MLP(config) self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) def forward( self, @@ -1003,9 +1159,13 @@ class Qwen2_5_VLDecoderLayer(nn.Module): output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` @@ -1072,9 +1232,14 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) self.layers = nn.ModuleList( - [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + [ + Qwen2_5_VLDecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] ) self._attn_implementation = config._attn_implementation self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -1103,16 +1268,26 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) if self.gradient_checkpointing and self.training: if use_cache: @@ -1129,19 +1304,29 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): inputs_embeds = self.embed_tokens(input_ids) if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) cache_position = torch.arange( - past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, ) # the hard coded `3` is for temporal, height and width. if position_ids is None: - position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + position_ids = cache_position.view(1, 1, -1).expand( + 3, inputs_embeds.shape[0], -1 + ) elif position_ids.dim() == 2: position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, ) hidden_states = inputs_embeds @@ -1199,7 +1384,11 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): next_cache = next_decoder_cache if use_cache else None if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -1217,7 +1406,9 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): ): if self.config._attn_implementation == "flash_attention_2": if attention_mask is not None and past_key_values is not None: - is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] + is_padding_right = ( + attention_mask[:, -1].sum().item() != input_tensor.size()[0] + ) if is_padding_right: raise ValueError( "You are attempting to perform batched generation with padding_side='right'" @@ -1231,7 +1422,9 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) using_static_cache = isinstance(past_key_values, StaticCache) using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) @@ -1286,7 +1479,9 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + causal_mask = AttentionMaskConverter._unmask_unattended( + causal_mask, min_dtype + ) return causal_mask @@ -1332,31 +1527,41 @@ class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): else: min_dtype = torch.finfo(dtype).min causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + (sequence_length, target_length), + fill_value=min_dtype, + dtype=dtype, + device=device, ) - diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + diagonal_attend_mask = torch.arange( + target_length, device=device + ) > cache_position.reshape(-1, 1) if config.sliding_window is not None: # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also # the check is needed to verify is current checkpoint was trained with sliding window or not - if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: - sliding_attend_mask = torch.arange(target_length, device=device) <= ( - cache_position.reshape(-1, 1) - config.sliding_window - ) + if ( + not isinstance(past_key_values, SlidingWindowCache) + or sequence_length > target_length + ): + sliding_attend_mask = torch.arange( + target_length, device=device + ) <= (cache_position.reshape(-1, 1) - config.sliding_window) diagonal_attend_mask.bitwise_or_(sliding_attend_mask) causal_mask *= diagonal_attend_mask causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + causal_mask = ( + causal_mask.clone() + ) # copy to contiguous memory for in-place edit if attention_mask.shape[-1] > target_length: attention_mask = attention_mask[:, :target_length] mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[ + :, None, None, : + ].to(causal_mask.device) padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) + causal_mask[:, :, :, :mask_length] = causal_mask[ + :, :, :, :mask_length + ].masked_fill(padding_mask, min_dtype) return causal_mask @@ -1482,7 +1687,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi def __init__(self, config): super().__init__(config) - self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( + config.vision_config + ) self.model = Qwen2_5_VLModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) @@ -1575,7 +1782,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi video_token_id = self.config.video_token_id vision_start_token_id = self.config.vision_start_token_id mrope_position_deltas = [] - if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): total_input_ids = input_ids if attention_mask is None: attention_mask = torch.ones_like(total_input_ids) @@ -1591,7 +1800,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi for i, input_ids in enumerate(total_input_ids): input_ids = input_ids[attention_mask[i] == 1] image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id + ).squeeze(1) vision_tokens = input_ids[vision_start_indices + 1] image_nums = (vision_tokens == image_token_id).sum() video_nums = (vision_tokens == video_token_id).sum() @@ -1639,38 +1850,78 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi ) text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) range_tensor = torch.arange(llm_grid_t).view(-1, 1) expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) - time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second + time_tensor = ( + expanded_range + * second_per_grid_t + * self.config.vision_config.tokens_per_second + ) time_tensor_long = time_tensor.long() t_index = time_tensor_long.flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) st = ed + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) return position_ids, mrope_position_deltas else: if attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) - max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] else: position_ids = ( @@ -1687,7 +1938,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi return position_ids, mrope_position_deltas @add_start_docstrings_to_model_forward(QWEN2_5_VL_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Qwen2_5_VLCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + @replace_return_docstrings( + output_type=Qwen2_5_VLCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC + ) def forward( self, input_ids: torch.LongTensor = None, @@ -1748,11 +2001,19 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..." ```""" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) @@ -1771,7 +2032,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) image_mask = mask_expanded.to(inputs_embeds.device) - image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: @@ -1789,14 +2052,18 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) video_mask = mask_expanded.to(inputs_embeds.device) - video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme - if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): # calculate RoPE index once per generation in the pre-fill stage only if ( (cache_position is not None and cache_position[0] == 0) @@ -1898,12 +2165,13 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi if past_key_values is not None: if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] - elif ( - inputs_embeds is not None # Exception 1 - or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]) # Exception 3 - ): + elif inputs_embeds is not None or ( # Exception 1 + is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1] + ): # Exception 3 input_ids = input_ids[:, -cache_position.shape[0] :] - elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) + elif ( + input_ids.shape[1] != cache_position.shape[0] + ): # Default case (the "else", a no op, is Exception 2) input_ids = input_ids[:, cache_position] if cache_position[0] != 0: @@ -1924,16 +2192,18 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi batch_size, sequence_length = input_ids.shape device = input_ids.device - attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=past_key_values.get_max_cache_shape(), - dtype=self.lm_head.weight.dtype, - device=device, - cache_position=cache_position, - batch_size=batch_size, - config=self.config, - past_key_values=past_key_values, + attention_mask = ( + self.model._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.lm_head.weight.dtype, + device=device, + cache_position=cache_position, + batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, + ) ) model_inputs.update( @@ -1996,7 +2266,13 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi if expand_size == 1: return input_ids, model_kwargs - visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"] + visual_keys = [ + "pixel_values", + "image_grid_thw", + "pixel_values_videos", + "video_grid_thw", + "second_per_grid_ts", + ] def _expand_dict_for_generation_visual(dict_to_expand): image_grid_thw = model_kwargs.get("image_grid_thw", None) @@ -2006,7 +2282,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi def _repeat_interleave_samples(x, lengths, repeat_times): samples = torch.split(x, lengths) repeat_args = [repeat_times] + [1] * (x.dim() - 1) - result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0) + result = torch.cat( + [sample.repeat(*repeat_args) for sample in samples], dim=0 + ) return result for key in dict_to_expand: @@ -2042,7 +2320,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi ) tensor = torch.tensor(dict_to_expand[key]) lengths = list(video_nums) - tensor = _repeat_interleave_samples(tensor, lengths=lengths, repeat_times=expand_size) + tensor = _repeat_interleave_samples( + tensor, lengths=lengths, repeat_times=expand_size + ) dict_to_expand[key] = tensor.tolist() return dict_to_expand @@ -2054,7 +2334,9 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi and isinstance(dict_to_expand[key], torch.Tensor) and key not in visual_keys ): - dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) + dict_to_expand[key] = dict_to_expand[key].repeat_interleave( + expand_size, dim=0 + ) return dict_to_expand # input_ids is required for expanding visual inputs @@ -2069,10 +2351,18 @@ class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMi if is_encoder_decoder: if model_kwargs.get("encoder_outputs") is None: - raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") - model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + raise ValueError( + "If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined." + ) + model_kwargs["encoder_outputs"] = _expand_dict_for_generation( + model_kwargs["encoder_outputs"] + ) return input_ids, model_kwargs -__all__ = ["Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLModel", "Qwen2_5_VLPreTrainedModel"] +__all__ = [ + "Qwen2_5_VLForConditionalGeneration", + "Qwen2_5_VLModel", + "Qwen2_5_VLPreTrainedModel", +] diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py index cf0edb8..e5a4f85 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py +++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py @@ -13,7 +13,12 @@ from typing import Optional, List, Tuple, Any, Dict, Union from transformers import AutoConfig, AutoProcessor from transformers.activations import ACT2FN from transformers.utils import logging, is_torchdynamo_compiling -from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache +from transformers.cache_utils import ( + Cache, + DynamicCache, + SlidingWindowCache, + StaticCache, +) from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.models.qwen2_vl.modeling_qwen2_vl import ( Qwen2RMSNorm, @@ -32,7 +37,12 @@ from transformers.modeling_outputs import ( from wall_x.fusions import ops from wall_x.model.action_head import ActionProcessor from wall_x.model.qwen2_5_based.configuration_qwen2_5_vl import Qwen2_5_VLConfig -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import Qwen2_5_VisionTransformerPretrainedModel, Qwen2_5_VLAttention, Qwen2_5_VLFlashAttention2, Qwen2_5_VLSdpaAttention +from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( + Qwen2_5_VisionTransformerPretrainedModel, + Qwen2_5_VLAttention, + Qwen2_5_VLFlashAttention2, + Qwen2_5_VLSdpaAttention, +) from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES @@ -67,104 +77,120 @@ class BlockSparseMLP(nn.Module): self.act_fn = ACT2FN[self.hidden_act] def forward(self, hidden_state): - return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + return self.down_proj( + self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state) + ) + class SparseMoeBlock(nn.Module): """Sparse Mixture of Experts (MoE) Block optimized with Grouped GEMM. - + This module implements a sparse MoE layer where tokens are dynamically routed to different expert networks. Uses grouped GEMM operations for efficient computation across multiple experts. - + Args: config: Configuration object containing expert specifications num_experts: Number of expert networks in the MoE block """ - + def __init__(self, config, num_experts: int): super().__init__() self.num_experts = num_experts # Initialize expert networks based on configuration - self.experts = nn.ModuleList([ - BlockSparseMLP(config.experts[i]) for i in range(num_experts) - ]) - + self.experts = nn.ModuleList( + [BlockSparseMLP(config.experts[i]) for i in range(num_experts)] + ) + def forward( - self, - hidden_states: torch.Tensor, - experts_indices: torch.Tensor, - start_indices: torch.Tensor, - end_indices: torch.Tensor + self, + hidden_states: torch.Tensor, + experts_indices: torch.Tensor, + start_indices: torch.Tensor, + end_indices: torch.Tensor, ) -> torch.Tensor: """Forward pass through the Sparse MoE block. - + Routes different hidden states to their corresponding expert networks for processing. Uses efficient grouped operations to minimize overhead. - + Args: hidden_states: Input tensor of shape (batch_size, seq_length, hidden_dim) experts_indices: Expert assignment indices of shape (batch_size, seq_length) indicating which expert each token should be routed to start_indices: Starting indices for each expert's assigned tokens end_indices: Ending indices for each expert's assigned tokens - + Returns: output: Processed tensor of shape (batch_size, seq_length, hidden_dim) after expert processing and token reordering """ batch_size, seq_length, hidden_dim = hidden_states.size() - + # Flatten inputs for efficient grouped processing hidden_states = hidden_states.view(-1, hidden_dim) # [total_tokens, hidden_dim] experts_indices = experts_indices.view(-1) # [total_tokens] - + # Create uniform probabilities for all tokens (can be modified for weighted routing) probs = torch.ones_like(experts_indices, dtype=torch.float32).view(-1, 1) - + # Permute inputs to group tokens by expert assignment permuted_inputs, row_id_map = ops.permute(hidden_states, experts_indices) final_output = torch.zeros_like(permuted_inputs) - + # Process tokens through their assigned experts for expert_idx, expert in enumerate(self.experts): # Skip experts with no assigned tokens if start_indices[expert_idx] == end_indices[expert_idx]: continue - + # Process tokens assigned to this expert - expert_input = permuted_inputs[start_indices[expert_idx]:end_indices[expert_idx]] + expert_input = permuted_inputs[ + start_indices[expert_idx] : end_indices[expert_idx] + ] expert_output = expert(expert_input) - final_output[start_indices[expert_idx]:end_indices[expert_idx]] = expert_output - + final_output[start_indices[expert_idx] : end_indices[expert_idx]] = ( + expert_output + ) + # Restore original token ordering unpermuted_outputs = ops.unpermute(final_output, row_id_map, probs) - + # Reshape back to original dimensions output = unpermuted_outputs.view(batch_size, seq_length, hidden_dim) - + return output + QWEN2_5_VL_ATTENTION_CLASSES = { "eager": Qwen2_5_VLAttention, "flash_attention_2": Qwen2_5_VLFlashAttention2, "sdpa": Qwen2_5_VLSdpaAttention, } + class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, num_experts: int): super().__init__() self.hidden_size = config.hidden_size - if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + if ( + config.use_sliding_window + and config._attn_implementation != "flash_attention_2" + ): logger.warning_once( f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " "unexpected results may be encountered." ) - self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation]( + config, layer_idx + ) self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) if config.mlp_moe: self.moe = SparseMoeBlock(config, num_experts=num_experts) @@ -186,7 +212,9 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` @@ -227,8 +255,10 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) - if self.mlp is None: # using moe mlp - hidden_states = self.moe(hidden_states, token_types, start_indices, end_indices) + if self.mlp is None: # using moe mlp + hidden_states = self.moe( + hidden_states, token_types, start_indices, end_indices + ) else: hidden_states = self.mlp(hidden_states) @@ -242,51 +272,52 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): outputs += (present_key_value,) return outputs + class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): """Qwen2.5-VL model with Mixture of Experts (MoE) architecture. - + This model extends the base Qwen2.5-VL model by incorporating MoE layers for improved scalability and specialization across different token types. """ - + @classmethod def from_pretrained( - cls, - pretrained_model_name_or_path: str, - num_experts: Optional[int] = None, - *args, - **kwargs + cls, + pretrained_model_name_or_path: str, + num_experts: Optional[int] = None, + *args, + **kwargs, ): """Load a pretrained model with optional MoE configuration. - + Args: pretrained_model_name_or_path: Path or name of the pretrained model num_experts: Number of experts for MoE layers (if not in config) *args: Additional arguments passed to parent class **kwargs: Additional keyword arguments passed to parent class - + Returns: Initialized model instance with MoE configuration """ config = kwargs.get("config", None) if config is None: config = AutoConfig.from_pretrained(pretrained_model_name_or_path) - + # Override number of experts if specified if num_experts is not None: config.num_experts = num_experts - + kwargs["config"] = config return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) def __init__(self, config: Qwen2_5_VLConfig): """Initialize the Qwen2.5-VL MoE model. - + Args: config: Model configuration containing architecture parameters """ super().__init__(config) - + # Basic model parameters self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size @@ -295,25 +326,27 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) - + # Decoder layers with MoE support - self.layers = nn.ModuleList([ - Qwen2_5_VLDecoderLayer_with_MoE(config, layer_idx, config.num_experts) - for layer_idx in range(config.num_hidden_layers) - ]) - + self.layers = nn.ModuleList( + [ + Qwen2_5_VLDecoderLayer_with_MoE(config, layer_idx, config.num_experts) + for layer_idx in range(config.num_hidden_layers) + ] + ) + # Model configuration self._attn_implementation = config._attn_implementation self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) self.gradient_checkpointing = False - + # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self) -> nn.Embedding: """Get the input embedding layer. - + Returns: The token embedding layer """ @@ -321,7 +354,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): def set_input_embeddings(self, value: nn.Embedding) -> None: """Set the input embedding layer. - + Args: value: New embedding layer to use """ @@ -345,20 +378,30 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): **kwargs, ) -> Union[Tuple, BaseModelOutputWithPast]: # Set default output options - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) # Validate inputs if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + if moe_token_types is None: raise ValueError("moe_token_types must be provided for MoE routing") - + # Handle gradient checkpointing compatibility if self.gradient_checkpointing and self.training: if use_cache: @@ -377,22 +420,31 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): # Set up cache position if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) cache_position = torch.arange( - past_seen_tokens, - past_seen_tokens + inputs_embeds.shape[1], - device=inputs_embeds.device + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, ) # Set up position IDs (hardcoded 3 dimensions for temporal, height, width) if position_ids is None: - position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + position_ids = cache_position.view(1, 1, -1).expand( + 3, inputs_embeds.shape[0], -1 + ) elif position_ids.dim() == 2: position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) # Create causal attention mask causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions, moe_token_types + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + moe_token_types, ) hidden_states = inputs_embeds @@ -439,7 +491,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): cache_position=cache_position, position_embeddings=position_embeddings, ) - + hidden_states = layer_outputs[0] # Update cache if using it @@ -462,10 +514,11 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): # Return outputs in requested format if not return_dict: return tuple( - v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None ) - + return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -483,11 +536,11 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): moe_token_types: Optional[torch.LongTensor] = None, ): """Update causal attention mask with support for bidirectional attention for specific token types. - + This method creates and modifies attention masks to support different attention patterns: - Standard causal (unidirectional) attention for most tokens - Bidirectional attention for specific token types (e.g., MoE routing tokens) - + Args: attention_mask: Input attention mask to avoid attending to padding tokens input_tensor: Input embeddings tensor for shape and device information @@ -496,7 +549,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): output_attentions: Whether attention weights will be returned moe_token_types: Optional tensor indicating token types for MoE routing (type 1 tokens will use bidirectional attention) - + Returns: Updated causal attention mask, or None if using Flash Attention 2 """ @@ -505,7 +558,9 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): return None # Calculate sequence lengths for cache management - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) using_static_cache = isinstance(past_key_values, StaticCache) using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) @@ -531,7 +586,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): dtype, device = input_tensor.dtype, input_tensor.device min_dtype = torch.finfo(dtype).min sequence_length = input_tensor.shape[1] - + # Determine target length based on cache type if using_sliding_window_cache or using_static_cache: # Use maximum cache shape for sliding window or static caches @@ -556,25 +611,31 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): config=self.config, past_key_values=past_key_values, ) - + # Modify mask to support bidirectional attention for specific token types if moe_token_types is not None: # Identify positions of type 1 tokens (MoE routing tokens) - type1_tokens = (moe_token_types == 1).unsqueeze(1).unsqueeze(2) # Shape: [B, 1, 1, S] + type1_tokens = ( + (moe_token_types == 1).unsqueeze(1).unsqueeze(2) + ) # Shape: [B, 1, 1, S] # Create bidirectional attention region for type 1 tokens # This allows type 1 tokens to attend to each other bidirectionally type1_mask = torch.zeros_like(causal_mask) # Shape: [B, num_heads, S, S] - type1_region = type1_tokens & type1_tokens.transpose(-1, -2) # Shape: [B, 1, S, S] + type1_region = type1_tokens & type1_tokens.transpose( + -1, -2 + ) # Shape: [B, 1, S, S] type1_mask = type1_mask.masked_fill(type1_region, 1.0).to(torch.bool) - + # Apply bidirectional attention: zero out causal constraints in type 1 regions causal_mask = torch.where( type1_mask, # Where type 1 tokens interact with each other - torch.zeros_like(causal_mask), # Remove causal masking (allow bidirectional) - causal_mask # Keep original causal masking for other regions + torch.zeros_like( + causal_mask + ), # Remove causal masking (allow bidirectional) + causal_mask, # Keep original causal masking for other regions ) - + # Handle special case for SDPA with CUDA/XPU devices if ( self.config._attn_implementation == "sdpa" @@ -585,7 +646,9 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): # Ensure attention to all tokens in fully masked rows for memory-efficient attention # This is required for F.scaled_dot_product_attention's memory-efficient path # when using left padding. See: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + causal_mask = AttentionMaskConverter._unmask_unattended( + causal_mask, min_dtype + ) return causal_mask @@ -631,92 +694,119 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): else: min_dtype = torch.finfo(dtype).min causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + (sequence_length, target_length), + fill_value=min_dtype, + dtype=dtype, + device=device, ) - diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + diagonal_attend_mask = torch.arange( + target_length, device=device + ) > cache_position.reshape(-1, 1) if config.sliding_window is not None: # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also # the check is needed to verify is current checkpoint was trained with sliding window or not - if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: - sliding_attend_mask = torch.arange(target_length, device=device) <= ( - cache_position.reshape(-1, 1) - config.sliding_window - ) + if ( + not isinstance(past_key_values, SlidingWindowCache) + or sequence_length > target_length + ): + sliding_attend_mask = torch.arange( + target_length, device=device + ) <= (cache_position.reshape(-1, 1) - config.sliding_window) diagonal_attend_mask.bitwise_or_(sliding_attend_mask) causal_mask *= diagonal_attend_mask causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + causal_mask = ( + causal_mask.clone() + ) # copy to contiguous memory for in-place edit if attention_mask.shape[-1] > target_length: attention_mask = attention_mask[:, :target_length] mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[ + :, None, None, : + ].to(causal_mask.device) padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) + causal_mask[:, :, :, :mask_length] = causal_mask[ + :, :, :, :mask_length + ].masked_fill(padding_mask, min_dtype) return causal_mask + class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): """ Qwen2.5 Vision-Language Mixture of Experts model for action processing. - + This model extends the base Qwen2.5 VL model with action token processing capabilities and optional LoRA fine-tuning support. """ + _tied_weights_keys = ["lm_head.weight"] config_class = Qwen2_5_VLConfig _no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"] @classmethod - def from_pretrained(cls, pretrained_model_path, config_path=None, processor_path=None, action_tokenizer_path=None, **kwargs): + def from_pretrained( + cls, + pretrained_model_path, + config_path=None, + processor_path=None, + action_tokenizer_path=None, + **kwargs, + ): """ Load model from pretrained model path. - + Args: pretrained_model_path (str): Model directory path containing model.safetensors file config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path processor_path (str, optional): Processor path, if None will load from default config action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config **kwargs: Additional arguments - + Returns: Qwen2_5_VLMoEForAction: Loaded model instance """ - + # Load model components from pretrained path config_path = os.path.join(pretrained_model_path, "config.json") config = cls.config_class.from_pretrained(config_path) processor = AutoProcessor.from_pretrained(pretrained_model_path, use_fast=True) if action_tokenizer_path is not None: - processor.action_processor = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True) - + processor.action_processor = AutoProcessor.from_pretrained( + action_tokenizer_path, trust_remote_code=True + ) + # Initialize model with configuration and processor - model = cls( - config, - processor=processor, - **kwargs - ) - + model = cls(config, processor=processor, **kwargs) + # Resize token embeddings to match processor tokenizer vocabulary size model.resize_token_embeddings(len(processor.tokenizer)) - + # Load model state dict from safetensors file - safetensor_files = glob.glob(os.path.join(pretrained_model_path, "*.safetensors")) + safetensor_files = glob.glob( + os.path.join(pretrained_model_path, "*.safetensors") + ) state_dict = {} for file in safetensor_files: sd = load_file(file, device="cpu") state_dict.update(sd) model.load_state_dict(state_dict, strict=False) - + return model - def __init__(self, config, use_fast_tokenizer=False, processor=None, action_tokenizer=None, action_mapper=None, flow_loss_weight=1.0): + def __init__( + self, + config, + use_fast_tokenizer=False, + processor=None, + action_tokenizer=None, + action_mapper=None, + flow_loss_weight=1.0, + ): """ Initialize the Qwen2.5 VLMoE model for action processing. - + Args: config: Model configuration use_fast_tokenizer (bool): Whether to use fast tokenizer @@ -726,9 +816,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): flow_loss_weight (float): Weight for flow loss computation """ super().__init__(config) - + # Initialize vision transformer and language model components - self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( + config.vision_config + ) self.model = Qwen2_5_VLMoEModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) @@ -738,13 +830,13 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): self.flow_loss_weight = flow_loss_weight self.use_fast_tokenizer = use_fast_tokenizer self.processor = processor - + # Define action token IDs self.define_action_token_id() # Cache for rope deltas self.rope_deltas = None - + # Initialize action preprocessor self.action_preprocessor = ActionProcessor(config) @@ -754,28 +846,30 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): r=config.lora_r, lora_alpha=config.lora_alpha, target_modules=config.lora_target_modules, - lora_dropout=config.lora_dropout + lora_dropout=config.lora_dropout, ) - + # Initialize weights and apply final processing self.post_init() def define_action_token_id(self): """ Define action token IDs based on tokenizer configuration. - + Creates mappings for fast action tokens, proprioception tokens, and general action tokens. """ # Create list of fast action token IDs fast_action_token_list = [] for i in range(self.processor.tokenizer.init_kwargs["action_token_vocab_size"]): - action_token_id = self.processor.tokenizer.convert_tokens_to_ids(f"<|action_token_{i}|>") + action_token_id = self.processor.tokenizer.convert_tokens_to_ids( + f"<|action_token_{i}|>" + ) fast_action_token_list.append(action_token_id) # Get special action token IDs - action_token_id = self.processor.tokenizer.convert_tokens_to_ids(f"<|action|>") + action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>") - + # Store action token ID mappings self.action_token_id_set = { "fast_action_token_list": fast_action_token_list, @@ -783,10 +877,12 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): "action_token_id": action_token_id, } - def add_lora(self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1): + def add_lora( + self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 + ): """ Add LoRA (Low-Rank Adaptation) adapters to the model. - + Args: r (int): Rank of adaptation lora_alpha (int): LoRA scaling parameter @@ -799,10 +895,10 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): target_modules=target_modules, lora_dropout=lora_dropout, bias="none", - task_type="CAUSAL_LM" + task_type="CAUSAL_LM", ) self.model = get_peft_model(self.model, config) - + # Print information about trainable parameters self.model.print_trainable_parameters() @@ -847,7 +943,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): For vision tokens, 3D position embeddings are calculated based on: - Temporal dimension: Time patches in videos - - Height dimension: Vertical patches in images/video frames + - Height dimension: Vertical patches in images/video frames - Width dimension: Horizontal patches in images/video frames For text tokens, standard 1D position embeddings are used, continuing from the maximum @@ -856,12 +952,12 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): Args: input_ids (torch.LongTensor, optional): Input token IDs of shape (batch_size, sequence_length) image_grid_thw (torch.LongTensor, optional): Image grid dimensions (num_images, 3) for [temporal, height, width] - video_grid_thw (torch.LongTensor, optional): Video grid dimensions (num_videos, 3) for [temporal, height, width] + video_grid_thw (torch.LongTensor, optional): Video grid dimensions (num_videos, 3) for [temporal, height, width] second_per_grid_ts (torch.Tensor, optional): Time interval per temporal grid (num_videos,) attention_mask (torch.Tensor, optional): Attention mask (batch_size, sequence_length) Returns: - tuple: + tuple: - position_ids (torch.LongTensor): 3D position IDs of shape (3, batch_size, sequence_length) - mrope_position_deltas (torch.Tensor): Position deltas for mRoPE of shape (batch_size, 1) """ @@ -870,12 +966,14 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): video_token_id = self.config.video_token_id vision_start_token_id = self.config.vision_start_token_id mrope_position_deltas = [] - - if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): total_input_ids = input_ids if attention_mask is None: attention_mask = torch.ones_like(total_input_ids) - + # Initialize 3D position IDs tensor position_ids = torch.ones( 3, @@ -884,26 +982,28 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): dtype=input_ids.dtype, device=input_ids.device, ) - + image_index, video_index = 0, 0 attention_mask = attention_mask.to(total_input_ids.device) - + # Process each sequence in the batch for i, input_ids in enumerate(total_input_ids): input_ids = input_ids[attention_mask[i] == 1] image_nums, video_nums = 0, 0 - + # Find vision tokens and count images/videos - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id + ).squeeze(1) vision_tokens = input_ids[vision_start_indices + 1] image_nums = (vision_tokens == image_token_id).sum() video_nums = (vision_tokens == video_token_id).sum() - + input_tokens = input_ids.tolist() llm_pos_ids_list: list = [] st = 0 remain_images, remain_videos = image_nums, video_nums - + # Process each vision token (image or video) for _ in range(image_nums + video_nums): # Find next image or video token @@ -911,12 +1011,12 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ed_image = input_tokens.index(image_token_id, st) else: ed_image = len(input_tokens) + 1 - + if video_token_id in input_tokens and remain_videos > 0: ed_video = input_tokens.index(video_token_id, st) else: ed_video = len(input_tokens) + 1 - + # Determine if processing image or video token if ed_image < ed_video: # Process image token @@ -943,7 +1043,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): video_index += 1 remain_videos -= 1 ed = ed_video - + # Calculate grid dimensions after spatial merging llm_grid_t, llm_grid_h, llm_grid_w = ( t.item(), @@ -953,46 +1053,86 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): text_len = ed - st # Add position IDs for text tokens before vision token - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) # Calculate 3D position embeddings for vision tokens range_tensor = torch.arange(llm_grid_t).view(-1, 1) expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) # Calculate temporal position IDs with time scaling - time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second + time_tensor = ( + expanded_range + * second_per_grid_t + * self.config.vision_config.tokens_per_second + ) time_tensor_long = time_tensor.long() t_index = time_tensor_long.flatten() # Calculate spatial position IDs - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + # Add 3D position IDs for vision tokens - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) st = ed + llm_grid_t * llm_grid_h * llm_grid_w # Add position IDs for remaining text tokens if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) # Concatenate all position IDs for this sequence llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - - mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) return position_ids, mrope_position_deltas else: # Handle case without vision tokens - use standard 1D position embeddings if attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) - max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] else: position_ids = ( @@ -1015,7 +1155,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - moe_token_types: Optional[torch.LongTensor] = None, # MoE token type assignments + moe_token_types: Optional[ + torch.LongTensor + ] = None, # MoE token type assignments labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -1026,7 +1168,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, action_chunk: Optional[torch.FloatTensor] = None, # Action trajectory chunks - proprioception: Optional[torch.FloatTensor] = None, # Joint position/orientation data + proprioception: Optional[ + torch.FloatTensor + ] = None, # Joint position/orientation data rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, @@ -1037,11 +1181,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) -> Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: """ Forward pass for training with multi-modal inputs including vision, text, and action data. - + This method handles the complete forward pass during training, processing various input modalities including images, videos, text, proprioceptive data, and action sequences. It computes losses for both language modeling and action prediction using flow matching. - + Args: input_ids (torch.LongTensor, optional): Input token IDs attention_mask (torch.Tensor, optional): Attention mask for input tokens @@ -1067,23 +1211,33 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): dof_mask (torch.FloatTensor, optional): Degrees of freedom mask for action tokens agent_pos_mask (torch.FloatTensor, optional): Agent position mask for proprioceptive data **kwargs: Additional keyword arguments - + Returns: - Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: Model outputs including losses, logits, + Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: Model outputs including losses, logits, and auxiliary information, or tuple if return_dict=False """ batch_size, seq_length = input_ids.shape - + # Set output configuration from model config if not specified - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict # Calculate RoPE position IDs if not provided # Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation - if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): # Calculate RoPE index once per generation in the pre-fill stage only if ( (cache_position is not None and cache_position[0] == 0) @@ -1111,9 +1265,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - + # Calculate token distribution across MoE expert groups - group_size = torch.zeros(self.config.num_experts, dtype=torch.long, device="cpu") + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) for i in range(self.config.num_experts): group_size[i] = (moe_token_types == i).sum() @@ -1124,7 +1280,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): # Process input embeddings with multi-modal data if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) - + # Process image embeddings if pixel_values is not None: pixel_values = pixel_values.type(self.visual.dtype) @@ -1134,7 +1290,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) image_mask = mask_expanded.to(inputs_embeds.device) - image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # Process video embeddings @@ -1143,7 +1301,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) n_video_tokens = (input_ids == self.config.video_token_id).sum().item() n_video_features = video_embeds.shape[0] - + # Validate video token and feature count match if n_video_tokens != n_video_features: raise ValueError( @@ -1154,45 +1312,71 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) video_mask = mask_expanded.to(inputs_embeds.device) - video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) # Process proprioceptive data (joint positions, orientations, etc.) if proprioception is not None: - proprioception = proprioception.to(inputs_embeds.device).to(inputs_embeds.dtype) - agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) - proprioception = self.action_preprocessor.proprioception_proj( - proprioception, dataset_names, agent_pos_mask, use_history=proprioception.shape[1] > 1 + proprioception = proprioception.to(inputs_embeds.device).to( + inputs_embeds.dtype ) - mask = input_ids == self.action_token_id_set['propri_token_id'] + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + proprioception = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + mask = input_ids == self.action_token_id_set["propri_token_id"] mask_unsqueezed = mask.unsqueeze(-1) mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) proprioception_mask = mask_expanded.to(inputs_embeds.device) - - proprioception = proprioception.to(inputs_embeds.device, inputs_embeds.dtype) - inputs_embeds = inputs_embeds.masked_scatter(proprioception_mask, proprioception) + + proprioception = proprioception.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter( + proprioception_mask, proprioception + ) elif self.training: # Dummy forward pass to ensure gradient registration in DDP # This handles cases where one process has proprioception data while another doesn't # Without this, DDP would hang waiting for a gradient that will never be computed - dummy_output = sum(p.sum() for p in self.action_preprocessor.proprioception_proj.parameters()) - dummy_input = torch.randn(2, self.action_preprocessor.propri_dim*2, device=inputs_embeds.device) - dummy_forward = self.action_preprocessor.proprioception_proj(dummy_input) + dummy_input = torch.randn( + 2, + self.action_preprocessor.propri_dim * 2, + device=inputs_embeds.device, + ) + dummy_forward = self.action_preprocessor.proprioception_proj( + dummy_input + ) dummy_loss = sum(p.sum() for p in dummy_forward) inputs_embeds = inputs_embeds + 0 * dummy_loss - + # Process action chunk data if action_chunk is not None: - action_chunk = action_chunk.to(inputs_embeds.device).to(inputs_embeds.dtype) + action_chunk = action_chunk.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) - noisy_action_emb, flow = self.action_preprocessor(action_chunk, dataset_names, dof_mask) - mask = input_ids == self.action_token_id_set['action_token_id'] + noisy_action_emb, flow = self.action_preprocessor( + action_chunk, dataset_names, dof_mask + ) + mask = input_ids == self.action_token_id_set["action_token_id"] mask_unsqueezed = mask.unsqueeze(-1) mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) action_mask = mask_expanded.to(inputs_embeds.device) - - noisy_action_emb = noisy_action_emb.to(inputs_embeds.device, inputs_embeds.dtype) - inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb) + + noisy_action_emb = noisy_action_emb.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter( + action_mask, noisy_action_emb + ) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) @@ -1200,7 +1384,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): # Forward pass through the main model outputs = self.model( input_ids=None, - position_ids=position_ids, + position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, @@ -1221,45 +1405,55 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): cross_entropy_loss, flow_loss = None, None channel_loss_dict = None channel_loss_count_dict = None - + # Compute losses if labels are provided if labels is not None: loss = 0 action_accuracy = 0 unique_datasets_name = list(set(dataset_names)) - + # Initialize per-dataset loss tracking dictionaries channel_loss_dict = { - dataset_name: torch.tensor(0.0, device=logits.device) + dataset_name: torch.tensor(0.0, device=logits.device) for dataset_name in ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES } channel_loss_count_dict = { - dataset_name: torch.tensor(0, device=logits.device) + dataset_name: torch.tensor(0, device=logits.device) for dataset_name in ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES } - + # Compute standard cross-entropy loss for language modeling shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() + shift_labels = labels[..., 1:].contiguous() shift_logits = shift_logits.view(-1, self.config.vocab_size) shift_labels = shift_labels.view(-1) - + # Enable model parallelism by moving labels to correct device shift_labels = shift_labels.to(shift_logits.device) - non_ignored_mask = (shift_labels != -100) + non_ignored_mask = shift_labels != -100 _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels) - cross_entropy_loss = _cross_entropy_loss[non_ignored_mask].mean() if non_ignored_mask.any() else torch.tensor(0.0, device=shift_logits.device) + cross_entropy_loss = ( + _cross_entropy_loss[non_ignored_mask].mean() + if non_ignored_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) # Compute per-dataset channel losses - _cross_entropy_loss = _cross_entropy_loss.view(batch_size, seq_length-1) - non_ignored_mask = non_ignored_mask.view(batch_size, seq_length-1) + _cross_entropy_loss = _cross_entropy_loss.view(batch_size, seq_length - 1) + non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1) for dataset_name_i in unique_datasets_name: - dataset_mask = torch.tensor([name == dataset_name_i for name in dataset_names], - device=logits.device) + dataset_mask = torch.tensor( + [name == dataset_name_i for name in dataset_names], + device=logits.device, + ) combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask - channel_loss_dict[dataset_name_i] = _cross_entropy_loss[combined_mask].sum() if combined_mask.any() else torch.tensor(0.0, device=shift_logits.device) + channel_loss_dict[dataset_name_i] = ( + _cross_entropy_loss[combined_mask].sum() + if combined_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) channel_loss_count_dict[dataset_name_i] += combined_mask.sum() - + # Add cross-entropy loss to total loss if valid if not torch.isnan(cross_entropy_loss): loss += cross_entropy_loss @@ -1270,26 +1464,34 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): # Compute action token prediction accuracy shift_logits = logits[..., :-1, :].contiguous() action_preds = shift_logits.argmax(dim=-1) - shift_labels = labels[..., 1:].contiguous() + shift_labels = labels[..., 1:].contiguous() if self.use_fast_tokenizer: - action_mask = shift_labels > self.action_token_id_set['fast_action_token_list'][0] + action_mask = ( + shift_labels > self.action_token_id_set["fast_action_token_list"][0] + ) correct_preds = (action_preds == shift_labels) & action_mask - action_accuracy = correct_preds.sum().float() / action_mask.sum().float() + action_accuracy = ( + correct_preds.sum().float() / action_mask.sum().float() + ) channel_loss_dict["action_accuracy"] = action_accuracy if action_chunk is not None: - action_mask = input_ids == self.action_token_id_set['action_token_id'] + action_mask = input_ids == self.action_token_id_set["action_token_id"] if action_mask.any(): action_hidden_states = hidden_states[action_mask] flow = flow.reshape(-1, flow.shape[-1]) - _flow_loss = self.action_preprocessor.flow_loss(action_hidden_states, flow, dof_mask) + _flow_loss = self.action_preprocessor.flow_loss( + action_hidden_states, flow, dof_mask + ) if isinstance(_flow_loss, torch.Tensor): flow_loss = _flow_loss.mean() if loss is not None: loss += self.flow_loss_weight * flow_loss else: loss = self.flow_loss_weight * flow_loss - _flow_loss = _flow_loss.view(dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2]) + _flow_loss = _flow_loss.view( + dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2] + ) # Return outputs based on return_dict setting if not return_dict: @@ -1298,7 +1500,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): return Qwen2_5_VLACausalLMOutputWithPast( loss=loss, - cross_entropy_loss=cross_entropy_loss.clone() if cross_entropy_loss is not None else None, + cross_entropy_loss=( + cross_entropy_loss.clone() if cross_entropy_loss is not None else None + ), flow_loss=flow_loss, logits=logits, past_key_values=outputs.past_key_values, @@ -1312,22 +1516,19 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): def predict_action(self, predict_mode: str, **kwargs): """ Predict actions using specified prediction mode. - + Args: predict_mode (str): Prediction mode, either "fast" or "diffusion" **kwargs: Additional arguments passed to the predict method - + Returns: tuple: (predicted_action, ground_truth_action) where ground_truth_action may be None """ assert predict_mode in ["fast", "diffusion"] - - output = self.predict( - predict_mode=predict_mode, - **kwargs - ) - return output['predict_action'], output.get('gt_action', None) + output = self.predict(predict_mode=predict_mode, **kwargs) + + return output["predict_action"], output.get("gt_action", None) @torch.no_grad() def predict( @@ -1364,12 +1565,12 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ): """ Multi-modal prediction method supporting text generation, fast action prediction, and diffusion-based action prediction. - + This method handles three prediction modes: 1. "text": Pure text generation using autoregressive decoding 2. "fast": Fast action prediction using discrete action tokens 3. "diffusion": Continuous action prediction using diffusion/flow matching - + Args: predict_mode (str): Prediction mode ("text", "fast", or "diffusion") pred_horizon (int, optional): Prediction horizon for action sequences @@ -1400,7 +1601,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): agent_pos_mask (torch.FloatTensor, optional): Agent position mask re_generate (bool, optional): Whether to use sampling for regeneration **kwargs: Additional keyword arguments - + Returns: dict: Dictionary containing prediction results with keys like: - 'predict_action': Predicted action sequences @@ -1409,30 +1610,42 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): - 'predict_output_text': Generated text (for text/fast modes) - 'gt_output_text': Ground truth text (for text/fast modes) """ - batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] - + batch_size = ( + input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + ) + # Text and fast modes require batch size 1 for autoregressive generation if predict_mode in ["text", "fast"]: - assert batch_size == 1, "predict only support batch size 1 for ar generation" - + assert ( + batch_size == 1 + ), "predict only support batch size 1 for ar generation" + # Set output configuration from model config if not specified - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + # Process input embeddings with multi-modal data if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) - + # Process image embeddings if pixel_values is not None: pixel_values = pixel_values.type(self.visual.dtype) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) n_image_tokens = (input_ids == self.config.image_token_id).sum().item() n_image_features = image_embeds.shape[0] - + # Validate image token and feature count match if n_image_tokens != n_image_features: raise ValueError( @@ -1444,7 +1657,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) image_mask = mask_expanded.to(inputs_embeds.device) - image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # Process video embeddings @@ -1453,7 +1668,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) n_video_tokens = (input_ids == self.config.video_token_id).sum().item() n_video_features = video_embeds.shape[0] - + # Validate video token and feature count match if n_video_tokens != n_video_features: raise ValueError( @@ -1465,25 +1680,40 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) video_mask = mask_expanded.to(inputs_embeds.device) - video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) # Process proprioceptive data if proprioception is not None: - proprioception = proprioception.to(inputs_embeds.device).to(inputs_embeds.dtype) - agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) + proprioception = proprioception.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) proprio_embed = self.action_preprocessor.proprioception_proj( - proprioception, dataset_names, agent_pos_mask, use_history=proprioception.shape[1] > 1 + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + proprioception_mask = ( + input_ids == self.action_token_id_set["propri_token_id"] + ) + inputs_embeds[proprioception_mask] = proprio_embed.reshape( + -1, inputs_embeds.shape[-1] ) - proprioception_mask = input_ids == self.action_token_id_set['propri_token_id'] - inputs_embeds[proprioception_mask] = proprio_embed.reshape(-1, inputs_embeds.shape[-1]) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) # Calculate RoPE position IDs if not provided # Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation - if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): # Calculate RoPE index once per generation in the pre-fill stage only if ( (cache_position is not None and cache_position[0] == 0) @@ -1518,30 +1748,38 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): action_chunk = action_chunk.to(inputs_embeds.device).to(inputs_embeds.dtype) output = {} - + # Split input sequence for text and fast modes (not needed for diffusion) if predict_mode == "text" or predict_mode == "fast": # Look for generation prompt tokens: <|im_start|>assistant - generation_prompt_ids = torch.tensor([151644, 77091], device=input_ids.device, dtype=input_ids.dtype) - matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (input_ids[0, 1:] == generation_prompt_ids[1]) - + generation_prompt_ids = torch.tensor( + [151644, 77091], device=input_ids.device, dtype=input_ids.dtype + ) + matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & ( + input_ids[0, 1:] == generation_prompt_ids[1] + ) + if matches.any(): split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item() # Extract ground truth output tokens (including newline) - gt_output_ids = input_ids[:, split_pos + 3:] + gt_output_ids = input_ids[:, split_pos + 3 :] # Remove output part from input, keeping prompt - input_ids = input_ids[:, :split_pos + 3] - inputs_embeds = inputs_embeds[:, :split_pos + 3, :] + input_ids = input_ids[:, : split_pos + 3] + inputs_embeds = inputs_embeds[:, : split_pos + 3, :] if attention_mask is not None: - attention_mask = attention_mask[:, :split_pos + 3] + attention_mask = attention_mask[:, : split_pos + 3] if labels is not None: - labels = labels[:, split_pos + 3:] + labels = labels[:, split_pos + 3 :] else: - raise Warning("input_ids does not contain the generation prompt tokens <|im_start|>assistant") - + raise Warning( + "input_ids does not contain the generation prompt tokens <|im_start|>assistant" + ) + # Decode input text for output - input_text = self.processor.batch_decode(input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True) - output['input_text'] = input_text + input_text = self.processor.batch_decode( + input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True + ) + output["input_text"] = input_text # Handle text and fast prediction modes using autoregressive generation if predict_mode == "text" or predict_mode == "fast": @@ -1558,7 +1796,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): "proprioception": proprioception, "dataset_names": dataset_names, } - + # Generate output tokens predict_output_ids = self.generate( **batch, @@ -1566,80 +1804,109 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): eos_token_id=[self.processor.tokenizer.eos_token_id], use_cache=True, pad_token_id=self.processor.tokenizer.pad_token_id, - temperature=1.0 if not re_generate else 0.7, # Higher temperature for regeneration - do_sample=False if not re_generate else True, # Enable sampling for regeneration + temperature=( + 1.0 if not re_generate else 0.7 + ), # Higher temperature for regeneration + do_sample=( + False if not re_generate else True + ), # Enable sampling for regeneration ) - + # Decode generated and ground truth text - gt_output_text = self.processor.batch_decode(gt_output_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True) - predict_output_text = self.processor.batch_decode(predict_output_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True) - output['gt_output_text'] = gt_output_text - output['predict_output_text'] = predict_output_text + gt_output_text = self.processor.batch_decode( + gt_output_ids, + skip_special_tokens=False, + clean_up_tokenization_spaces=True, + ) + predict_output_text = self.processor.batch_decode( + predict_output_ids, + skip_special_tokens=False, + clean_up_tokenization_spaces=True, + ) + output["gt_output_text"] = gt_output_text + output["predict_output_text"] = predict_output_text # Convert tokens to actions for fast prediction mode if predict_mode == "fast": action_id = [] # Extract action tokens from generated sequence for token_id_i in predict_output_ids[0]: - if token_id_i.item() >= self.processor.tokenizer.init_kwargs["action_token_start_index"]: - action_id.append(token_id_i.item() - self.processor.tokenizer.init_kwargs["action_token_start_index"]) + if ( + token_id_i.item() + >= self.processor.tokenizer.init_kwargs["action_token_start_index"] + ): + action_id.append( + token_id_i.item() + - self.processor.tokenizer.init_kwargs[ + "action_token_start_index" + ] + ) - predict_action = self.processor.action_processor.decode([action_id], time_horizon=pred_horizon, action_dim=action_dim) + predict_action = self.processor.action_processor.decode( + [action_id], time_horizon=pred_horizon, action_dim=action_dim + ) # Handle action decoding errors if np.sum(predict_action) == 0: print("Error in decoding action, predict_action is None") - output['predict_action'] = None + output["predict_action"] = None else: # Convert discrete tokens to continuous actions predict_action = torch.tensor(predict_action, device=self.device) dof_mask = dof_mask.to(self.device).to(pixel_values.dtype) - predict_action = self.action_preprocessor.normalizer_action.unnormalize_data(predict_action, dataset_names, dof_mask) - output['predict_action'] = predict_action - + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, dof_mask + ) + ) + output["predict_action"] = predict_action + # Process ground truth actions if available if action_chunk is not None: # Apply DOF mask and unnormalize action chunk to get ground truth actions action_chunk = action_chunk[:, :, dof_mask[0, 0, :].bool()] - output['gt_action'] = self.action_preprocessor.normalizer_action.unnormalize_data(action_chunk, dataset_names, dof_mask) + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names, dof_mask + ) + ) else: - output['gt_action'] = None - + output["gt_action"] = None + # Handle diffusion-based action prediction if predict_mode == "diffusion": # Initialize with random noise noisy_action = torch.randn( size=(batch_size, pred_horizon, action_dim), - dtype=inputs_embeds.dtype, device=inputs_embeds.device + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, ) dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) def step(timestep, noisy_action): """ Single denoising step for diffusion process. - + Args: timestep: Current diffusion timestep noisy_action: Current noisy action estimate - + Returns: torch.Tensor: Predicted clean action """ - action_mask = input_ids == self.action_token_id_set['action_token_id'] + action_mask = input_ids == self.action_token_id_set["action_token_id"] assert action_mask.any(), "No action token found in input_ids" - + # Prepare timestep for batch processing timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) action_embed = self.action_preprocessor.step( - timestep=timestep, - noisy_action=noisy_action, - dof_mask=dof_mask + timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask ) action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]) - + # Create temporary copy of embeddings for thread safety temp_inputs_embeds = inputs_embeds.clone() temp_inputs_embeds[action_mask] = action_embed - + # Forward pass through transformer transformer_outputs = self.model( input_ids=None, @@ -1653,69 +1920,78 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): output_hidden_states=False, return_dict=True, ) - + # Extract action predictions from hidden states hidden_states = transformer_outputs.last_hidden_state - action_mask = input_ids == self.action_token_id_set['action_token_id'] + action_mask = input_ids == self.action_token_id_set["action_token_id"] action_hidden_states = hidden_states[action_mask] pred = self.action_preprocessor.action_proj_back(action_hidden_states) return pred.reshape(batch_size, pred_horizon, action_dim) # Perform ODE integration for diffusion sampling - times = torch.linspace(0, 1, num_inference_timesteps, - device=inputs_embeds.device, dtype=inputs_embeds.dtype) + times = torch.linspace( + 0, + 1, + num_inference_timesteps, + device=inputs_embeds.device, + dtype=inputs_embeds.dtype, + ) action_trajectory = odeint(step, noisy_action, times, method="euler") # Extract final predicted action and unnormalize predict_action = action_trajectory[-1] - predict_action = self.action_preprocessor.normalizer_action.unnormalize_data(predict_action, dataset_names) - output['predict_action'] = predict_action - + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names + ) + ) + output["predict_action"] = predict_action + # Process ground truth actions if available if action_chunk is not None: - output['gt_action'] = self.action_preprocessor.normalizer_action.unnormalize_data(action_chunk, dataset_names) + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names + ) + ) return output - def forward( - self, - mode: Optional[str] = None, - predict_mode: Optional[str] = "text", - **kwargs + self, mode: Optional[str] = None, predict_mode: Optional[str] = "text", **kwargs ): """ Main forward pass dispatcher for different execution modes. - + This method routes execution to appropriate forward functions based on the specified mode: - - No mode (None): Training step with gradient disabled + - No mode (None): Training step with gradient disabled - 'predict': Prediction/inference mode - 'train': Training mode with gradients enabled - 'validate': Validation mode with gradients disabled - + Args: mode (str, optional): Execution mode. If None, defaults to training step without gradients predict_mode (str, optional): Prediction mode for 'predict' mode ("text", "fast", or "diffusion") **kwargs: Additional arguments passed to the selected forward function - + Returns: Model outputs appropriate for the selected mode - + Todo: - Add support for distinguishing multi-modal data types in prediction mode """ if not mode: with torch.no_grad(): return self.train_step_forward(**kwargs) - elif mode == 'predict': + elif mode == "predict": return self.predict(predict_mode=predict_mode, **kwargs) - elif mode == 'train': + elif mode == "train": return self.train_step_forward(use_cache=False, **kwargs) - elif mode == 'validate': + elif mode == "validate": with torch.no_grad(): return self.train_step_forward(use_cache=False, **kwargs) else: - raise NotImplementedError('invalid key') + raise NotImplementedError("invalid key") def prepare_inputs_for_generation( self, @@ -1740,11 +2016,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ): """ Prepare inputs for autoregressive generation with multi-modal support. - + This method handles input preparation for generation, including proper slicing of inputs based on cache position, MoE token type management, and multi-modal data handling. Vision inputs are selectively forwarded only when needed during generation. - + Args: input_ids: Input token IDs past_key_values: Cached key-value pairs from previous generation steps @@ -1764,13 +2040,13 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): dof_mask: Degrees of freedom mask agent_pos_mask: Agent position mask **kwargs: Additional arguments - + Returns: dict: Prepared model inputs for generation step - + Todo: - Test this function thoroughly with various input configurations - + Note: This is an overridden method that handles specific cases for multi-modal generation: - Slices input_ids through cache_position to keep only unprocessed tokens @@ -1779,31 +2055,38 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): """ # Initialize MoE token types if not provided if moe_token_types is None: - moe_token_types = torch.zeros_like(input_ids) # FIXME: Handle case when input_embeds is used instead + moe_token_types = torch.zeros_like( + input_ids + ) # FIXME: Handle case when input_embeds is used instead else: # Ensure moe_token_types length matches input_ids if moe_token_types.shape[1] < input_ids.shape[1]: # Calculate required padding length pad_length = input_ids.shape[1] - moe_token_types.shape[1] # Create padding tensor with default token type (0) - pad_tensor = torch.zeros((moe_token_types.shape[0], pad_length), - dtype=moe_token_types.dtype, - device=moe_token_types.device) + pad_tensor = torch.zeros( + (moe_token_types.shape[0], pad_length), + dtype=moe_token_types.dtype, + device=moe_token_types.device, + ) # Concatenate padding to existing moe_token_types moe_token_types = torch.cat([moe_token_types, pad_tensor], dim=1) # Handle input slicing based on cache state and special cases if past_key_values is not None: - if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4: input_embeds case + if ( + inputs_embeds is not None and input_ids.shape[1] == 0 + ): # Exception 4: input_embeds case inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] moe_token_types = moe_token_types[:, -cache_position.shape[0] :] - elif ( - inputs_embeds is not None # Exception 1: input_embeds provided - or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]) # Exception 3: GPU sync edge case - ): + elif inputs_embeds is not None or ( # Exception 1: input_embeds provided + is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1] + ): # Exception 3: GPU sync edge case input_ids = input_ids[:, -cache_position.shape[0] :] moe_token_types = moe_token_types[:, -cache_position.shape[0] :] - elif input_ids.shape[1] != cache_position.shape[0]: # Default case (Exception 2 is no-op) + elif ( + input_ids.shape[1] != cache_position.shape[0] + ): # Default case (Exception 2 is no-op) cache_pos = cache_position.clone() input_ids = input_ids[:, cache_pos] moe_token_types = moe_token_types[:, cache_pos] @@ -1828,16 +2111,18 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): batch_size, sequence_length = input_ids.shape device = input_ids.device - attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=past_key_values.get_max_cache_shape(), - dtype=self.lm_head.weight.dtype, - device=device, - cache_position=cache_position, - batch_size=batch_size, - config=self.config, - past_key_values=past_key_values, + attention_mask = ( + self.model._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.lm_head.weight.dtype, + device=device, + cache_position=cache_position, + batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, + ) ) # Assemble all model inputs for generation @@ -1868,7 +2153,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) -> Tuple[torch.Tensor, torch.Tensor]: """ Get the number of images and videos for each sample to calculate tensor separation lengths. - + These parameters are computed directly from input_ids rather than being passed through the processor to avoid unpredictable impacts from interface modifications. @@ -1876,7 +2161,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): input_ids (torch.LongTensor): Input token IDs of shape (batch_size, sequence_length) Returns: - tuple: + tuple: - image_nums (torch.LongTensor): Number of images per sample - video_nums (torch.LongTensor): Number of videos per sample """ @@ -1889,7 +2174,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) image_mask = input_ids == image_token_id video_mask = input_ids == video_token_id - + # Count images and videos following vision start tokens image_nums = torch.sum(vision_first_mask & image_mask, dim=1) video_nums = torch.sum(vision_first_mask & video_mask, dim=1) @@ -1905,19 +2190,19 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) -> Tuple[torch.LongTensor, Dict[str, Any]]: """ Expand inputs for generation with support for multi-modal tensors. - + This is an overridden method that supports expanding tensors without a standard batch size dimension, specifically for vision-related tensors: - pixel_values.shape[0] = sum(sequence_lengths for all image samples) - image_grid_thw.shape[0] = sum(num_images for all samples) - Similar patterns for video tensors - + Args: expand_size (int): Factor by which to expand inputs (for beam search, etc.) is_encoder_decoder (bool): Whether using encoder-decoder architecture input_ids (torch.LongTensor, optional): Input token IDs **model_kwargs: Additional model arguments to expand - + Returns: tuple: (expanded_input_ids, expanded_model_kwargs) """ @@ -1925,7 +2210,13 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): return input_ids, model_kwargs # Define keys for vision-related tensors that need special handling - visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"] + visual_keys = [ + "pixel_values", + "image_grid_thw", + "pixel_values_videos", + "video_grid_thw", + "second_per_grid_ts", + ] def _expand_dict_for_generation_visual(dict_to_expand): """Expand vision-related tensors based on image/video counts per sample.""" @@ -1937,7 +2228,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): """Split tensor by lengths and repeat each sample.""" samples = torch.split(x, lengths) repeat_args = [repeat_times] + [1] * (x.dim() - 1) - result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0) + result = torch.cat( + [sample.repeat(*repeat_args) for sample in samples], dim=0 + ) return result for key in dict_to_expand: @@ -1975,7 +2268,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) tensor = torch.tensor(dict_to_expand[key]) lengths = list(video_nums) - tensor = _repeat_interleave_samples(tensor, lengths=lengths, repeat_times=expand_size) + tensor = _repeat_interleave_samples( + tensor, lengths=lengths, repeat_times=expand_size + ) dict_to_expand[key] = tensor.tolist() return dict_to_expand @@ -1988,7 +2283,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): and isinstance(dict_to_expand[key], torch.Tensor) and key not in visual_keys ): - dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) + dict_to_expand[key] = dict_to_expand[key].repeat_interleave( + expand_size, dim=0 + ) return dict_to_expand # Expand visual inputs only if input_ids is available for counting images/videos @@ -2006,7 +2303,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): # Handle encoder-decoder specific expansion if is_encoder_decoder: if model_kwargs.get("encoder_outputs") is None: - raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") - model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + raise ValueError( + "If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined." + ) + model_kwargs["encoder_outputs"] = _expand_dict_for_generation( + model_kwargs["encoder_outputs"] + ) - return input_ids, model_kwargs \ No newline at end of file + return input_ids, model_kwargs diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py index ed693ee..736b936 100644 --- a/wall_x/trainer/qwen_vl_act_trainer.py +++ b/wall_x/trainer/qwen_vl_act_trainer.py @@ -12,25 +12,29 @@ from datetime import datetime from torch.optim import AdamW from accelerate import Accelerator from safetensors.torch import load_file -from accelerate.utils import DistributedType from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup from wall_x.utils.timers import Timers from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES -from wall_x.data.load_lerobot_dataset import PreprocessedDataset, get_data_configs, load_lerobot_data +from wall_x.data.load_lerobot_dataset import ( + PreprocessedDataset, + get_data_configs, + load_lerobot_data, +) def timer(func): """ Decorator to measure function execution time. - + Args: func: Function to be timed - + Returns: Wrapped function with timing functionality """ + @wraps(func) def wrapper(*args, **kwargs): start_time = time.time() @@ -40,13 +44,14 @@ def timer(func): f"\033[92m[current time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Function {func.__name__} took {end_time - start_time:.2f} seconds to execute\033[0m" ) return result + return wrapper def print_rank_last(message): """ Print message only on the last rank in distributed training. - + Args: message (str): Message to print """ @@ -60,7 +65,7 @@ def print_rank_last(message): def seed_all(seed): """ Set random seeds for reproducible training. - + Args: seed (int): Random seed value """ @@ -72,11 +77,11 @@ def seed_all(seed): class QwenVlAct_Trainer: """ Vision-Language-Action trainer for Qwen-VL models with robotic action prediction. - + This trainer handles multi-modal learning combining vision, language, and action data for robotic control applications. It supports distributed training, mixed precision, gradient accumulation, and various optimization strategies including MoE (Mixture of Experts). - + Features: - Multi-modal data processing (vision + language + actions) - Distributed training with Accelerate @@ -87,10 +92,17 @@ class QwenVlAct_Trainer: """ @timer - def __init__(self, config, logger, accelerator: Accelerator = None, seed=42, data_config_path=None): + def __init__( + self, + config, + logger, + accelerator: Accelerator = None, + seed=42, + data_config_path=None, + ): """ Initialize the Vision-Language-Action trainer. - + Args: config (dict): Training configuration dictionary containing: - processor_path (str): Path to data preprocessing processor @@ -103,7 +115,7 @@ class QwenVlAct_Trainer: accelerator (Accelerator, optional): Hugging Face Accelerate instance for distributed training seed (int, optional): Random seed for reproducibility. Defaults to 42. data_config_path (str, optional): Path to data configuration file - + Raises: ValueError: If required configuration keys are missing """ @@ -117,41 +129,51 @@ class QwenVlAct_Trainer: self.logger = logger self.accelerator = accelerator self.seed = seed - + # Initialize random seeds for reproducibility seed_all(self.seed) - + # Training state variables self.start_epoch = 0 self.global_step = 0 self.num_epoch = self.config["num_epoch"] self.initial_step = 0 - + # Data and model configuration self.dataload_config = get_data_configs(self.config["data"]) self.data_config_path = data_config_path self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - + # Load model and initialize training components self.load_model() self.action_dim = sum(self.config["dof_config"].values()) - + # Distributed training setup self.rank = self.accelerator.process_index self.world_size = self.accelerator.num_processes - print(f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", flush=True) - + print( + f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", + flush=True, + ) + # Load training data self.load_qact_data() - print(f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", flush=True) + print( + f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", + flush=True, + ) # Resume from checkpoint if specified if "resume" in self.config: self.resume_from_checkpoint() # Initialize special token IDs - self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>") - self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") + self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids( + "<|propri|>" + ) + self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids( + "<|action|>" + ) # Initialize evaluation metrics self.base_l1_loss = None @@ -162,12 +184,14 @@ class QwenVlAct_Trainer: # Adjust global step if resuming from checkpoint if self.initial_step != 0: - self.global_step = self.initial_step // self.config.get("gradient_accumulation_steps", 1) + self.global_step = self.initial_step // self.config.get( + "gradient_accumulation_steps", 1 + ) def print_rank0(self, msg, flush=True): """ Print message only on rank 0 (main process). - + Args: msg: Message to print flush (bool): Whether to flush output buffer @@ -178,7 +202,7 @@ class QwenVlAct_Trainer: def fit(self): """ Main training loop executing multiple epochs with validation. - + Handles the complete training process including: - Training loop execution - Validation after each epoch @@ -186,9 +210,11 @@ 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): + if self.config.get("resume", None) is not None and self.config["resume"].get( + "validate_first", False + ): self.val_loop() self.accelerator.wait_for_everyone() @@ -196,24 +222,24 @@ class QwenVlAct_Trainer: for epoch in range(self.start_epoch, self.num_epoch): self.train_loop(epoch) self.accelerator.wait_for_everyone() - + if (epoch + 1) % self.config.get("epoch_save_interval", 10) == 0: self.save_checkpoint(epoch) - + # Validation after each epoch self.val_loop() self.accelerator.wait_for_everyone() - + # Memory cleanup gc.collect() def train_loop(self, epoch): """ Execute training for a single epoch. - + Args: epoch (int): Current epoch number - + Handles: - Data loading and batching - Forward/backward passes @@ -227,7 +253,9 @@ class QwenVlAct_Trainer: if getattr(self, "train_dataloader", None) is not None: self.train_sampler.set_epoch(epoch) else: - self.train_dataloader, self.train_sampler = self.dataset.get_train_dataloader() + self.train_dataloader, self.train_sampler = ( + self.dataset.get_train_dataloader() + ) self.train_sampler.set_epoch(epoch) else: self.train_dataloader = self.dataset.get_train_dataloader() @@ -236,16 +264,23 @@ class QwenVlAct_Trainer: grad_accum_steps = self.config.get("gradient_accumulation_steps", 1) total = len(self.train_dataloader) t0 = time.time() - enable_profiling = self.config['profile'] + enable_profiling = self.config["profile"] # Optional PyTorch profiler for performance analysis if enable_profiling: profiler = torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], - schedule=torch.profiler.schedule(wait=self.config['profile_wait_iters'], - warmup=self.config['profile_warmup_iters'], - active=self.config['profile_active_iters']), - on_trace_ready=torch.profiler.tensorboard_trace_handler(self.config['profile_save_path'], worker_name="worker0"), + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule( + wait=self.config["profile_wait_iters"], + warmup=self.config["profile_warmup_iters"], + active=self.config["profile_active_iters"], + ), + on_trace_ready=torch.profiler.tensorboard_trace_handler( + self.config["profile_save_path"], worker_name="worker0" + ), record_shapes=True, profile_memory=True, with_stack=True, @@ -253,7 +288,7 @@ class QwenVlAct_Trainer: profiler.__enter__() try: - + # Setup timers for First iteration self.timers("interval-time", log_level=0).start(barrier=False) self.timers("data-load", log_level=0).start(barrier=False) @@ -261,28 +296,38 @@ class QwenVlAct_Trainer: for i, batch in enumerate(self.train_dataloader, self.initial_step): # Move batch to device if isinstance(self.dataset, PreprocessedDataset): - batch = {k: v.to(self.accelerator.device, non_blocking=True) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} - + batch = { + k: ( + v.to(self.accelerator.device, non_blocking=True) + if isinstance(v, torch.Tensor) + else v + ) + for k, v in batch.items() + } + self.timers("data-load").stop() - + with self.accelerator.accumulate(self.model): # Forward pass self.timers("forward-compute", log_level=0).start(barrier=False) outputs = self.model(**batch, mode="train") self.timers("forward-compute").stop() - + loss = outputs.loss - + # Check for NaN loss if torch.isnan(loss): - print(f"Warning: NaN loss detected in epoch: {epoch}, step: {i}", flush=True) + print( + f"Warning: NaN loss detected in epoch: {epoch}, step: {i}", + flush=True, + ) continue - + # Backward pass self.timers("backward-compute", log_level=0).start(barrier=False) self.accelerator.backward(loss) self.timers("backward-compute").stop() - + # Gradient clipping total_norm = self.accelerator.clip_grad_norm_( self.model.parameters(), self.config.get("max_grad_norm", 1.0) @@ -299,33 +344,79 @@ class QwenVlAct_Trainer: self.lr_scheduler.step() self.global_step += 1 lr = self.lr_scheduler.get_last_lr()[0] - + # Gather loss across all processes for logging - train_loss = self.accelerator.gather(loss.detach()).mean().item() + train_loss = ( + self.accelerator.gather(loss.detach()).mean().item() + ) _log_dict = { "lr": lr, "train_loss": train_loss, } # Log component losses - if "cross_entropy_loss" in outputs and outputs.cross_entropy_loss is not None: - _log_dict["cross_entropy_loss"] = self.accelerator.gather(outputs.cross_entropy_loss.detach()).mean().item() - + if ( + "cross_entropy_loss" in outputs + and outputs.cross_entropy_loss is not None + ): + _log_dict["cross_entropy_loss"] = ( + self.accelerator.gather( + outputs.cross_entropy_loss.detach() + ) + .mean() + .item() + ) + if "flow_loss" in outputs and outputs.flow_loss is not None: - _log_dict["flow_loss"] = self.accelerator.gather(outputs.flow_loss.detach()).mean().item() - + _log_dict["flow_loss"] = ( + self.accelerator.gather(outputs.flow_loss.detach()) + .mean() + .item() + ) + # Log per-dataset channel losses - if "channel_loss_dict" in outputs and outputs.channel_loss_dict is not None: - for dataset_name_i in ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES: - count_sum = self.accelerator.gather(outputs.channel_loss_count_dict[dataset_name_i]).sum().item() + if ( + "channel_loss_dict" in outputs + and outputs.channel_loss_dict is not None + ): + for dataset_name_i in ( + ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES + ): + count_sum = ( + self.accelerator.gather( + outputs.channel_loss_count_dict[dataset_name_i] + ) + .sum() + .item() + ) if count_sum > 0: - channel_loss = self.accelerator.gather(outputs.channel_loss_dict[dataset_name_i].detach()).sum().item() / count_sum - _log_dict[f"channel_loss_{dataset_name_i}"] = channel_loss - + channel_loss = ( + self.accelerator.gather( + outputs.channel_loss_dict[ + dataset_name_i + ].detach() + ) + .sum() + .item() + / count_sum + ) + _log_dict[f"channel_loss_{dataset_name_i}"] = ( + channel_loss + ) + # Log action accuracy for fast tokenizer - if "action_accuracy" in outputs.channel_loss_dict and self.use_fast_tokenizer: + if ( + "action_accuracy" in outputs.channel_loss_dict + and self.use_fast_tokenizer + ): _log_dict["action_accuracy"] = ( - self.accelerator.gather(outputs.channel_loss_dict["action_accuracy"].detach()).mean().item() + self.accelerator.gather( + outputs.channel_loss_dict[ + "action_accuracy" + ].detach() + ) + .mean() + .item() ) # Log metrics @@ -334,23 +425,26 @@ class QwenVlAct_Trainer: # Log gradient norm if self.logger is not None and self.accelerator.sync_gradients: - self.logger.log({"total_norm": total_norm}, step=self.global_step) + self.logger.log( + {"total_norm": total_norm}, step=self.global_step + ) self.timers("interval-time").stop() - + # Setup timers for next iteration if i < len(self.train_dataloader) - 1: self.timers("interval-time", log_level=0).start(barrier=False) self.timers("data-load", log_level=0).start(barrier=False) - # Periodic logging t1 = time.time() if i % 1 == 0: lr = self.lr_scheduler.get_last_lr()[0] - self.training_log(epoch, self.num_epoch, i, total, loss, lr, t1 - t0) + self.training_log( + epoch, self.num_epoch, i, total, loss, lr, t1 - t0 + ) t0 = time.time() - + if enable_profiling: profiler.step() @@ -362,7 +456,7 @@ class QwenVlAct_Trainer: def val_loop(self): """ Execute validation loop with gradient computation disabled. - + Evaluates model performance on validation set and logs validation loss. """ # Initialize validation dataloader @@ -374,34 +468,45 @@ class QwenVlAct_Trainer: self.model.eval() self.val_loss = 0 - + # Validation loop for i, batch in enumerate( - tqdm(self.val_dataloader, desc="Validating", total=len(self.val_dataloader), - disable=not self.accelerator.is_main_process) + tqdm( + self.val_dataloader, + desc="Validating", + total=len(self.val_dataloader), + disable=not self.accelerator.is_main_process, + ) ): if isinstance(self.dataset, PreprocessedDataset): - batch = {k: v.to(self.accelerator.device, non_blocking=True) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} - + batch = { + k: ( + v.to(self.accelerator.device, non_blocking=True) + if isinstance(v, torch.Tensor) + else v + ) + for k, v in batch.items() + } + with torch.no_grad(): outputs = self.model(**batch, mode="train") loss = outputs.loss self.val_loss += self.accelerator.gather(loss.detach()).mean().item() - + # Calculate average validation loss self.val_loss /= len(self.val_dataloader) - + # Log validation metrics if self.logger is not None: self.logger.log({"val_loss": self.val_loss}, step=self.global_step) - + self.model.train() @timer def load_model(self): """ Load and configure the Vision-Language-Action model. - + Handles: - Model loading from pretrained weights - Processor initialization @@ -411,8 +516,8 @@ class QwenVlAct_Trainer: """ # Load pretrained model model = Qwen2_5_VLMoEForAction.from_pretrained( - self.config["pretrained_wallx_path"], - **{"use_fast_tokenizer": self.use_fast_tokenizer} + self.config["pretrained_wallx_path"], + **{"use_fast_tokenizer": self.use_fast_tokenizer}, ) self.processor = model.processor model = model.to(torch.bfloat16) @@ -428,7 +533,7 @@ class QwenVlAct_Trainer: moe_params.append(param) param_groups = [{"params": moe_params, "lr": self.config["learning_rate"]}] self.optimizer = AdamW(param_groups, weight_decay=0.1) - + elif "action_expert_learning_rate" in self.config: # Separate learning rates for VLM and action expert parameters moe_params = [] @@ -442,17 +547,26 @@ class QwenVlAct_Trainer: # Configure parameter groups if self.config.get("train_action_expert_only", False): self.print_rank0("Training action expert only", flush=True) - param_groups = [{"params": moe_params, "lr": self.config["action_expert_learning_rate"]}] + param_groups = [ + { + "params": moe_params, + "lr": self.config["action_expert_learning_rate"], + } + ] else: param_groups = [ {"params": vlm_params, "lr": self.config["learning_rate"]}, - {"params": moe_params, "lr": self.config["action_expert_learning_rate"]}, + { + "params": moe_params, + "lr": self.config["action_expert_learning_rate"], + }, ] self.optimizer = AdamW(param_groups, weight_decay=0.1) self.print_rank0( f"Setting MoE learning rate to {self.config['action_expert_learning_rate']}, " - f"VLM learning rate to {self.config['learning_rate']}", flush=True + f"VLM learning rate to {self.config['learning_rate']}", + flush=True, ) else: # Standard optimizer configuration @@ -479,9 +593,13 @@ class QwenVlAct_Trainer: if hasattr(model, "enable_input_require_grads"): self.model.enable_input_require_grads() else: + def make_inputs_require_grad(module, input, output): output.requires_grad_(True) - self.model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + self.model.get_input_embeddings().register_forward_hook( + make_inputs_require_grad + ) # Prepare model, optimizer, and scheduler for distributed training self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( @@ -492,7 +610,7 @@ class QwenVlAct_Trainer: def load_qact_data(self): """ Load and configure training data for Vision-Language-Action learning. - + Supports LeRobot dataset format and handles distributed data loading across multiple processes. """ @@ -511,18 +629,20 @@ class QwenVlAct_Trainer: def load_qwen_pretrain_weight(self, model, pretrain_weight_path): """ Load pretrained Qwen weights with MoE adaptation. - + Args: model: Model instance to load weights into pretrain_weight_path (str): Path to pretrained weight files - + Returns: Model with loaded pretrained weights - + Handles weight key renaming for MoE architecture compatibility. """ # Load all safetensors files - weight_files = sorted([f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")]) + weight_files = sorted( + [f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")] + ) merged_weights = {} # Merge weights from all files @@ -534,19 +654,31 @@ class QwenVlAct_Trainer: # Rename weights for MoE compatibility renamed_weights = {} for key, value in merged_weights.items(): - if key.startswith("model.layers") and "mlp." in key and model.config.mlp_moe: + if ( + key.startswith("model.layers") + and "mlp." in key + and model.config.mlp_moe + ): # Rename MLP weights for MoE structure layer_num = key.split(".layers.")[1].split(".mlp")[0] - new_key = key.replace(f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0.") + new_key = key.replace( + f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0." + ) renamed_weights[new_key] = value - elif key.startswith("model.layers") and "self_attn." in key and model.config.attention_moe: + elif ( + key.startswith("model.layers") + and "self_attn." in key + and model.config.attention_moe + ): # Rename attention weights for MoE structure layer_num = key.split(".layers.")[1].split(".self_attn")[0] proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"] for proj in proj_types: if proj in key: - new_key = key.replace(f"layers.{layer_num}.self_attn.{proj}", - f"layers.{layer_num}.self_attn.{proj}_experts.0") + new_key = key.replace( + f"layers.{layer_num}.self_attn.{proj}", + f"layers.{layer_num}.self_attn.{proj}_experts.0", + ) renamed_weights[new_key] = value break else: @@ -560,10 +692,19 @@ class QwenVlAct_Trainer: return model - def training_log(self, current_epoch, total_epoch, current_train_iter, total_train_iter, loss, lr, time_per_step): + def training_log( + self, + current_epoch, + total_epoch, + current_train_iter, + total_train_iter, + loss, + lr, + time_per_step, + ): """ Log training progress and performance metrics. - + Args: current_epoch (int): Current epoch number total_epoch (int): Total number of epochs @@ -573,26 +714,32 @@ class QwenVlAct_Trainer: lr (float): Current learning rate time_per_step (float): Time taken for current step """ - timers_to_log = ["interval-time", "data-load", "forward-compute", "backward-compute", "optimizer"] - + timers_to_log = [ + "interval-time", + "data-load", + "forward-compute", + "backward-compute", + "optimizer", + ] + log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]" log_string += " epoch {:3d}/{:3d} |".format(current_epoch, total_epoch) log_string += " iter {:6d}/{:6d} |".format(current_train_iter, total_train_iter) log_string += " loss {:.6f} |".format(loss) log_string += " lr {:.6f} |".format(lr) log_string += " time_per_step_avg {:.6f}s |".format(time_per_step) - + print_rank_last(log_string) self.timers.log(timers_to_log, normalizer=1) def save_checkpoint(self, epoch, step=0): """ Save training checkpoint. - + Args: epoch (int): Current epoch number step (int, optional): Current step number. Defaults to 0. - + Saves model state, optimizer state, and training progress information. """ save_path = self.config["save_path"] @@ -600,7 +747,7 @@ class QwenVlAct_Trainer: ckpt_path = f"{save_path}/{epoch}" else: ckpt_path = f"{save_path}/{epoch}_{step}" - + self.accelerator.save_state(ckpt_path) # Save current iteration steps for dataset resuming @@ -608,14 +755,16 @@ class QwenVlAct_Trainer: _rank = self.accelerator.process_index if isinstance(self.dataset, PreprocessedDataset): torch.save( - {"epoch": epoch, "step": step}, - os.path.join(ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth") + {"epoch": epoch, "step": step}, + os.path.join( + ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth" + ), ) def resume_from_checkpoint(self): """ Resume training from a saved checkpoint. - + Handles both full checkpoint loading and model-only loading based on configuration. """ checkpoint_path = self.config["resume"]["ckpt"] @@ -624,45 +773,47 @@ class QwenVlAct_Trainer: # Load only model weights ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors" state_dict = load_file(ckpt_path, device="cpu") - + # 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] - - err = self.model.load_state_dict(new_state_dict, strict=False) + + 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) - + self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}") def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask): """ Log detailed L1 loss metrics by degrees of freedom. - + Args: all_label (torch.Tensor): Ground truth action labels all_pred (torch.Tensor): Predicted actions all_task (list): Task identifiers all_dof_mask (torch.Tensor): Degrees of freedom mask - + Computes and logs L1 loss for each DOF component separately for detailed analysis. """ - all_task = all_task[:len(all_label)] + all_task = all_task[: len(all_label)] # Apply DOF mask all_label = all_label * all_dof_mask all_pred = all_pred * all_dof_mask - + # Compute baseline L1 loss (predict mean action) if self.base_l1_loss is None: mean_action = all_label.mean(dim=0) self.base_l1_loss = nn.functional.l1_loss(all_label, mean_action) - self.logger.log({"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step) + self.logger.log( + {"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step + ) # Log L1 loss for each DOF component start_idx = 0 @@ -672,8 +823,10 @@ class QwenVlAct_Trainer: dof_label = all_label[:, :, start_idx:end_idx] dof_pred = all_pred[:, :, start_idx:end_idx] dof_l1 = nn.functional.l1_loss(dof_pred, dof_label) - + self.print_rank0(f"DOF {dof}, L1 loss: {dof_l1.item()}", flush=True) - self.logger.log({f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step) - - start_idx = end_idx \ No newline at end of file + self.logger.log( + {f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step + ) + + start_idx = end_idx diff --git a/wall_x/utils/constant.py b/wall_x/utils/constant.py index 9da35fd..3074591 100755 --- a/wall_x/utils/constant.py +++ b/wall_x/utils/constant.py @@ -9,133 +9,292 @@ action_statistic_dof = { "min": [-3.6176], "delta": [8.5015], }, - "follow_left_ee_cartesian_pos": {"min": [-0.036, -0.3241, -0.1245], "delta": [0.4389, 0.557, 0.479]}, - "follow_left_ee_rotation": {"min": [-1.2373, -0.1929, -1.5182], "delta": [2.2009, 1.5669, 2.0936]}, + "follow_left_ee_cartesian_pos": { + "min": [-0.036, -0.3241, -0.1245], + "delta": [0.4389, 0.557, 0.479], + }, + "follow_left_ee_rotation": { + "min": [-1.2373, -0.1929, -1.5182], + "delta": [2.2009, 1.5669, 2.0936], + }, "follow_left_gripper": {"min": [-0.1196], "delta": [4.5226]}, - "follow_right_ee_cartesian_pos": {"min": [-0.0326, -0.2273, -0.1377], "delta": [0.4574, 0.5704, 0.4743]}, - "follow_right_ee_rotation": {"min": [-1.2201, -0.2611, -0.7427], "delta": [2.6623, 1.6622, 2.4186]}, + "follow_right_ee_cartesian_pos": { + "min": [-0.0326, -0.2273, -0.1377], + "delta": [0.4574, 0.5704, 0.4743], + }, + "follow_right_ee_rotation": { + "min": [-1.2201, -0.2611, -0.7427], + "delta": [2.6623, 1.6622, 2.4186], + }, "follow_right_gripper": {"min": [-0.1208], "delta": [4.5261]}, "height": {"min": [-0.0001], "delta": [0.5051]}, "head_actions": {"min": [-1.5000, -1.4167], "delta": [2.5000, 1.8879]}, - "base_velocity": {"min": [-0.0359, -0.084, -0.0162], "delta": [0.1539, 0.1848, 0.0322]}, + "base_velocity": { + "min": [-0.0359, -0.084, -0.0162], + "delta": [0.1539, 0.1848, 0.0322], + }, }, "DobbE": { - "follow_right_ee_cartesian_pos": {"min": [-0.6107, -0.3272, -0.4282], "delta": [1.2629, 1.5297, 0.8349]}, - "follow_right_ee_rotation": {"min": [-1.7378, -1.4597, -1.8712], "delta": [2.7031, 2.8182, 3.5921]}, + "follow_right_ee_cartesian_pos": { + "min": [-0.6107, -0.3272, -0.4282], + "delta": [1.2629, 1.5297, 0.8349], + }, + "follow_right_ee_rotation": { + "min": [-1.7378, -1.4597, -1.8712], + "delta": [2.7031, 2.8182, 3.5921], + }, "follow_right_gripper": {"min": [0.0], "delta": [0.9983]}, }, "RH20T": { - "follow_right_ee_cartesian_pos": {"min": [0.3646, -0.2722, 0.0066], "delta": [0.3813, 0.5973, 0.3277]}, - "follow_right_ee_rotation": {"min": [-1.8716, -0.4398, -3.1414], "delta": [3.4145, 1.0225, 6.2828]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3646, -0.2722, 0.0066], + "delta": [0.3813, 0.5973, 0.3277], + }, + "follow_right_ee_rotation": { + "min": [-1.8716, -0.4398, -3.1414], + "delta": [3.4145, 1.0225, 6.2828], + }, "follow_right_gripper": {"min": [0.0], "delta": [95.0]}, }, "agibotworld_alpha": { - "follow_left_ee_cartesian_pos": {"min": [0.4954, 0.0166, 0.1729], "delta": [0.3336, 0.5123, 0.9189]}, - "follow_left_ee_rotation": {"min": [-3.1064, -1.2629, -3.1238], "delta": [6.2127, 2.5923, 6.2496]}, + "follow_left_ee_cartesian_pos": { + "min": [0.4954, 0.0166, 0.1729], + "delta": [0.3336, 0.5123, 0.9189], + }, + "follow_left_ee_rotation": { + "min": [-3.1064, -1.2629, -3.1238], + "delta": [6.2127, 2.5923, 6.2496], + }, "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]}, - "follow_right_ee_cartesian_pos": {"min": [0.4615, -0.5975, 0.1638], "delta": [0.3823, 0.5577, 0.8873]}, - "follow_right_ee_rotation": {"min": [-3.0891, -1.0739, -2.5091], "delta": [6.1707, 2.3074, 3.8533]}, + "follow_right_ee_cartesian_pos": { + "min": [0.4615, -0.5975, 0.1638], + "delta": [0.3823, 0.5577, 0.8873], + }, + "follow_right_ee_rotation": { + "min": [-3.0891, -1.0739, -2.5091], + "delta": [6.1707, 2.3074, 3.8533], + }, "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]}, "height": {"min": [0.0], "delta": [0.4535]}, "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]}, }, "austin_buds": { - "follow_right_ee_cartesian_pos": {"min": [0.3496, -0.2855, 0.0105], "delta": [0.3748, 0.492, 0.3116]}, - "follow_right_ee_rotation": {"min": [-3.1405, -0.151, -0.0737], "delta": [6.2813, 0.3218, 0.1536]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3496, -0.2855, 0.0105], + "delta": [0.3748, 0.492, 0.3116], + }, + "follow_right_ee_rotation": { + "min": [-3.1405, -0.151, -0.0737], + "delta": [6.2813, 0.3218, 0.1536], + }, "follow_right_gripper": {"min": [0.0076], "delta": [0.0724]}, }, "austin_sailor": { - "follow_right_ee_cartesian_pos": {"min": [0.387, -0.3165, 0.0244], "delta": [0.2999, 0.5252, 0.2308]}, - "follow_right_ee_rotation": {"min": [-3.1402, -0.1618, -1.5918], "delta": [6.2804, 0.337, 2.9478]}, + "follow_right_ee_cartesian_pos": { + "min": [0.387, -0.3165, 0.0244], + "delta": [0.2999, 0.5252, 0.2308], + }, + "follow_right_ee_rotation": { + "min": [-3.1402, -0.1618, -1.5918], + "delta": [6.2804, 0.337, 2.9478], + }, "follow_right_gripper": {"min": [0.0005], "delta": [0.0773]}, }, "austin_sirius": { - "follow_right_ee_cartesian_pos": {"min": [0.0, -0.1182, 0.0], "delta": [0.5329, 0.3812, 0.2723]}, - "follow_right_ee_rotation": {"min": [-3.1407, -0.1243, -1.7434], "delta": [6.2823, 0.1975, 1.8073]}, + "follow_right_ee_cartesian_pos": { + "min": [0.0, -0.1182, 0.0], + "delta": [0.5329, 0.3812, 0.2723], + }, + "follow_right_ee_rotation": { + "min": [-3.1407, -0.1243, -1.7434], + "delta": [6.2823, 0.1975, 1.8073], + }, "follow_right_gripper": {"min": [0.0334], "delta": [0.046]}, }, "bc_z": { - "follow_right_ee_cartesian_pos": {"min": [-0.3883, -0.1116, 0.6113], "delta": [0.7199, 0.4288, 0.3709]}, - "follow_right_ee_rotation": {"min": [-1.056, -1.0587, -2.6295], "delta": [1.9142, 1.9455, 4.8064]}, + "follow_right_ee_cartesian_pos": { + "min": [-0.3883, -0.1116, 0.6113], + "delta": [0.7199, 0.4288, 0.3709], + }, + "follow_right_ee_rotation": { + "min": [-1.056, -1.0587, -2.6295], + "delta": [1.9142, 1.9455, 4.8064], + }, "follow_right_gripper": {"min": [0.2], "delta": [0.8]}, }, "berkeley_autolab_ur5": { - "follow_right_ee_cartesian_pos": {"min": [0.3018, -0.2129, -0.1888], "delta": [0.3121, 0.52, 0.3107]}, - "follow_right_ee_rotation": {"min": [-3.1396, -0.2278, 1.1413], "delta": [6.279, 0.454, 0.9841]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3018, -0.2129, -0.1888], + "delta": [0.3121, 0.52, 0.3107], + }, + "follow_right_ee_rotation": { + "min": [-3.1396, -0.2278, 1.1413], + "delta": [6.279, 0.454, 0.9841], + }, "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, }, "berkeley_cable_routing": { - "follow_right_ee_cartesian_pos": {"min": [0.4617, -0.28, 0.03], "delta": [0.1838, 0.5665, 0.1272]}, - "follow_right_ee_rotation": {"min": [-3.1413, -0.0299, -0.7665], "delta": [6.2826, 0.0692, 3.322]}, + "follow_right_ee_cartesian_pos": { + "min": [0.4617, -0.28, 0.03], + "delta": [0.1838, 0.5665, 0.1272], + }, + "follow_right_ee_rotation": { + "min": [-3.1413, -0.0299, -0.7665], + "delta": [6.2826, 0.0692, 3.322], + }, }, "berkeley_fanuc_manipulation": { - "follow_right_ee_cartesian_pos": {"min": [0.3718, -0.4072, 0.0184], "delta": [0.3483, 0.7201, 0.5229]}, - "follow_right_ee_rotation": {"min": [-3.1399, -1.0166, -1.6988], "delta": [6.2802, 1.4498, 3.2074]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3718, -0.4072, 0.0184], + "delta": [0.3483, 0.7201, 0.5229], + }, + "follow_right_ee_rotation": { + "min": [-3.1399, -1.0166, -1.6988], + "delta": [6.2802, 1.4498, 3.2074], + }, "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, }, "bridge_data_v2": { - "follow_right_ee_cartesian_pos": {"min": [0.1498, -0.2178, -0.0901], "delta": [0.3012, 0.469, 0.298]}, - "follow_right_ee_rotation": {"min": [-0.3279, -0.6105, -1.0578], "delta": [0.7378, 1.0353, 2.2552]}, + "follow_right_ee_cartesian_pos": { + "min": [0.1498, -0.2178, -0.0901], + "delta": [0.3012, 0.469, 0.298], + }, + "follow_right_ee_rotation": { + "min": [-0.3279, -0.6105, -1.0578], + "delta": [0.7378, 1.0353, 2.2552], + }, "follow_right_gripper": {"min": [0.0692], "delta": [0.9426]}, }, "dlr_edan_shared_control": { - "follow_right_ee_cartesian_pos": {"min": [-0.8387, 0.1473, -0.3934], "delta": [0.6579, 0.6025, 1.1566]}, - "follow_right_ee_rotation": {"min": [-3.1217, -1.5197, -2.2516], "delta": [6.2505, 1.5594, 4.2831]}, + "follow_right_ee_cartesian_pos": { + "min": [-0.8387, 0.1473, -0.3934], + "delta": [0.6579, 0.6025, 1.1566], + }, + "follow_right_ee_rotation": { + "min": [-3.1217, -1.5197, -2.2516], + "delta": [6.2505, 1.5594, 4.2831], + }, "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, }, "droid": { - "follow_right_ee_cartesian_pos": {"min": [0.2667, -0.4396, -0.0472], "delta": [0.5159, 0.8806, 0.8331]}, - "follow_right_ee_rotation": {"min": [-3.1374, -1.216, -2.1741], "delta": [6.2749, 2.1075, 4.2259]}, + "follow_right_ee_cartesian_pos": { + "min": [0.2667, -0.4396, -0.0472], + "delta": [0.5159, 0.8806, 0.8331], + }, + "follow_right_ee_rotation": { + "min": [-3.1374, -1.216, -2.1741], + "delta": [6.2749, 2.1075, 4.2259], + }, "follow_right_gripper": {"min": [0.0], "delta": [0.9912]}, }, "fmb": { - "follow_right_ee_cartesian_pos": {"min": [0.3554, -0.2844, 0.0354], "delta": [0.336, 0.4961, 0.2943]}, - "follow_right_ee_rotation": {"min": [-3.1404, -0.9302, -0.0599], "delta": [6.2807, 1.724, 1.8284]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3554, -0.2844, 0.0354], + "delta": [0.336, 0.4961, 0.2943], + }, + "follow_right_ee_rotation": { + "min": [-3.1404, -0.9302, -0.0599], + "delta": [6.2807, 1.724, 1.8284], + }, "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, }, "fractal": { - "follow_right_ee_cartesian_pos": {"min": [0.3242, -0.2836, 0.1405], "delta": [0.5518, 0.4963, 0.9328]}, - "follow_right_ee_rotation": {"min": [-3.1308, -0.2421, -2.9685], "delta": [6.2609, 1.7343, 5.819]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3242, -0.2836, 0.1405], + "delta": [0.5518, 0.4963, 0.9328], + }, + "follow_right_ee_rotation": { + "min": [-3.1308, -0.2421, -2.9685], + "delta": [6.2609, 1.7343, 5.819], + }, "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, }, "furniture_bench": { - "follow_right_ee_cartesian_pos": {"min": [0.3691, -0.181, 0.0058], "delta": [0.2962, 0.3582, 0.1775]}, - "follow_right_ee_rotation": {"min": [-3.1394, -0.6121, -1.9958], "delta": [6.2786, 1.6114, 3.7748]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3691, -0.181, 0.0058], + "delta": [0.2962, 0.3582, 0.1775], + }, + "follow_right_ee_rotation": { + "min": [-3.1394, -0.6121, -1.9958], + "delta": [6.2786, 1.6114, 3.7748], + }, "follow_right_gripper": {"min": [0.0035], "delta": [0.0762]}, }, "jaco_play": { - "follow_right_ee_cartesian_pos": {"min": [-0.3787, -0.6294, 0.1682], "delta": [0.5898, 0.3587, 0.2183]}, - "follow_right_ee_rotation": {"min": [0.9792, -0.0668, -0.0498], "delta": [0.0175, 0.1277, 0.0686]}, + "follow_right_ee_cartesian_pos": { + "min": [-0.3787, -0.6294, 0.1682], + "delta": [0.5898, 0.3587, 0.2183], + }, + "follow_right_ee_rotation": { + "min": [0.9792, -0.0668, -0.0498], + "delta": [0.0175, 0.1277, 0.0686], + }, "follow_right_gripper": {"min": [0.0791], "delta": [0.1033]}, }, "nyu_rot": { - "follow_right_ee_cartesian_pos": {"min": [0.25, -1.0, -0.2], "delta": [0.75, 2.0, 1.2]}, - "follow_right_ee_rotation": {"min": [-3.1416, -3.1416, 6.2831], "delta": [9.4248, 4.1416, 0.0]}, + "follow_right_ee_cartesian_pos": { + "min": [0.25, -1.0, -0.2], + "delta": [0.75, 2.0, 1.2], + }, + "follow_right_ee_rotation": { + "min": [-3.1416, -3.1416, 6.2831], + "delta": [9.4248, 4.1416, 0.0], + }, "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, }, "stanford_hydra": { - "follow_right_ee_cartesian_pos": {"min": [0.2068, -0.274, 0.1317], "delta": [0.4929, 0.4981, 0.4588]}, - "follow_right_ee_rotation": {"min": [-3.1321, -0.7496, -3.0269], "delta": [6.2658, 1.5176, 5.8261]}, + "follow_right_ee_cartesian_pos": { + "min": [0.2068, -0.274, 0.1317], + "delta": [0.4929, 0.4981, 0.4588], + }, + "follow_right_ee_rotation": { + "min": [-3.1321, -0.7496, -3.0269], + "delta": [6.2658, 1.5176, 5.8261], + }, "follow_right_gripper": {"min": [0.0], "delta": [0.0811]}, }, "stanford_kuka_multimodal": { - "follow_right_ee_cartesian_pos": {"min": [0.4781, -0.0659, 0.3424], "delta": [0.0868, 0.0864, 0.1863]}, - "follow_right_ee_rotation": {"min": [-3.136, -0.0521, -3.1413], "delta": [6.2727, 0.1109, 6.2825]}, + "follow_right_ee_cartesian_pos": { + "min": [0.4781, -0.0659, 0.3424], + "delta": [0.0868, 0.0864, 0.1863], + }, + "follow_right_ee_rotation": { + "min": [-3.136, -0.0521, -3.1413], + "delta": [6.2727, 0.1109, 6.2825], + }, "follow_right_gripper": {"min": [-0.4713], "delta": [0.9485]}, }, "taco_play": { - "follow_right_ee_cartesian_pos": {"min": [0.1375, -0.4291, 0.2052], "delta": [0.5327, 1.0237, 0.3913]}, - "follow_right_ee_rotation": {"min": [-3.1391, -0.6946, -1.2808], "delta": [6.2784, 0.8196, 3.0856]}, + "follow_right_ee_cartesian_pos": { + "min": [0.1375, -0.4291, 0.2052], + "delta": [0.5327, 1.0237, 0.3913], + }, + "follow_right_ee_rotation": { + "min": [-3.1391, -0.6946, -1.2808], + "delta": [6.2784, 0.8196, 3.0856], + }, "follow_right_gripper": {"min": [0.0001], "delta": [0.0806]}, }, "utaustin_mutex": { - "follow_right_ee_cartesian_pos": {"min": [0.3213, -0.4734, 0.0141], "delta": [0.2108, 0.8471, 0.5644]}, - "follow_right_ee_rotation": {"min": [-3.1404, -0.2202, -1.5489], "delta": [6.2805, 0.582, 1.9282]}, + "follow_right_ee_cartesian_pos": { + "min": [0.3213, -0.4734, 0.0141], + "delta": [0.2108, 0.8471, 0.5644], + }, + "follow_right_ee_rotation": { + "min": [-3.1404, -0.2202, -1.5489], + "delta": [6.2805, 0.582, 1.9282], + }, "follow_right_gripper": {"min": [0.0019], "delta": [0.0738]}, }, "viola": { - "follow_right_ee_cartesian_pos": {"min": [0.4011, -0.2521, 0.0103], "delta": [0.2444, 0.4305, 0.4355]}, - "follow_right_ee_rotation": {"min": [-3.1403, -0.2737, -1.8626], "delta": [6.2804, 0.4901, 2.0618]}, + "follow_right_ee_cartesian_pos": { + "min": [0.4011, -0.2521, 0.0103], + "delta": [0.2444, 0.4305, 0.4355], + }, + "follow_right_ee_rotation": { + "min": [-3.1403, -0.2737, -1.8626], + "delta": [6.2804, 0.4901, 2.0618], + }, "follow_right_gripper": {"min": [0.0002], "delta": [0.0773]}, }, "kuka": { @@ -179,11 +338,23 @@ action_statistic_dof = { }, }, "agibotworld_beta": { - "follow_left_ee_cartesian_pos": {"min": [0.4954, 0.0166, 0.1729], "delta": [0.3336, 0.5123, 0.9189]}, - "follow_left_ee_rotation": {"min": [-3.1064, -1.2629, -3.1238], "delta": [6.2127, 2.5923, 6.2496]}, + "follow_left_ee_cartesian_pos": { + "min": [0.4954, 0.0166, 0.1729], + "delta": [0.3336, 0.5123, 0.9189], + }, + "follow_left_ee_rotation": { + "min": [-3.1064, -1.2629, -3.1238], + "delta": [6.2127, 2.5923, 6.2496], + }, "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]}, - "follow_right_ee_cartesian_pos": {"min": [0.4615, -0.5975, 0.1638], "delta": [0.3823, 0.5577, 0.8873]}, - "follow_right_ee_rotation": {"min": [-3.0891, -1.0739, -2.5091], "delta": [6.1707, 2.3074, 3.8533]}, + "follow_right_ee_cartesian_pos": { + "min": [0.4615, -0.5975, 0.1638], + "delta": [0.3823, 0.5577, 0.8873], + }, + "follow_right_ee_rotation": { + "min": [-3.0891, -1.0739, -2.5091], + "delta": [6.1707, 2.3074, 3.8533], + }, "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]}, "height": {"min": [0.0], "delta": [0.4535]}, "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]}, diff --git a/wall_x/utils/timers.py b/wall_x/utils/timers.py index 0a67fb3..32b03b6 100644 --- a/wall_x/utils/timers.py +++ b/wall_x/utils/timers.py @@ -4,23 +4,28 @@ from torch.cuda import nvtx from abc import ABC, abstractmethod from typing import List + def _is_distributed(): return torch.distributed.is_available() and torch.distributed.is_initialized() + def _get_world_size(): if _is_distributed(): return torch.distributed.get_world_size() return 1 + def _get_rank(): if _is_distributed(): return torch.distributed.get_rank() return 0 + def _barrier(group=None): if _is_distributed(): torch.distributed.barrier(group=group) + if torch.distributed.is_available(): try: dist_all_gather_func = torch.distributed.all_gather_into_tensor @@ -29,6 +34,7 @@ if torch.distributed.is_available(): else: dist_all_gather_func = None + class TimerBase(ABC): """Timer base class.""" @@ -76,7 +82,7 @@ class DummyTimer(TimerBase): """Dummy Timer.""" def __init__(self): - super().__init__('dummy timer') + super().__init__("dummy timer") def start(self, barrier=False, nvtx_push=False): return @@ -89,8 +95,8 @@ class DummyTimer(TimerBase): def elapsed(self, reset=True, barrier=False): raise Exception( - 'dummy timer should not be used to calculate elapsed time, ' - 'check if timer\'s log_level <= self._log_level.' + "dummy timer should not be used to calculate elapsed time, " + "check if timer's log_level <= self._log_level." ) def active_time(self): @@ -98,8 +104,8 @@ class DummyTimer(TimerBase): Note: Not supported for DummyTimer. """ raise Exception( - 'active timer should not be used to calculate elapsed time, ' - 'check if timer\'s log_level <= self._log_level.' + "active timer should not be used to calculate elapsed time, " + "check if timer's log_level <= self._log_level." ) @@ -144,7 +150,7 @@ class Timer(TimerBase): Args: barrier (bool, optional): Synchronizes ranks before starting. Defaults to False. """ - assert not self._started, 'timer has already been started' + assert not self._started, "timer has already been started" if barrier: _barrier(group=self._barrier_group) if torch.cuda.is_available(): @@ -154,7 +160,6 @@ class Timer(TimerBase): if nvtx_push: nvtx.range_push("{}".format(self.name)) self.nvtx = True - def stop(self, barrier=False, sync=False): """Stop the timer. @@ -164,7 +169,7 @@ class Timer(TimerBase): """ if self.nvtx: nvtx.range_pop() - assert self._started, 'timer is not started' + assert self._started, "timer is not started" if barrier: _barrier(group=self._barrier_group) if torch.cuda.is_available() and sync: @@ -221,10 +226,10 @@ class Timers: Allowed: ['max', 'minmax', 'all']. """ self._log_level = log_level - allowed_log_options = set(['max', 'minmax', 'all']) + allowed_log_options = set(["max", "minmax", "all"]) assert ( log_option in allowed_log_options - ), 'input log option {} is invalid. It must be one of {}'.format( + ), "input log option {} is invalid. It must be one of {}".format( log_option, allowed_log_options ) self._log_option = log_option @@ -240,8 +245,10 @@ class Timers: if name in self._timers: if log_level is not None: assert log_level == self._log_levels[name], ( - 'input log level {} does not match already existing ' - 'log level {} for {} timer'.format(log_level, self._log_levels[name], name) + "input log level {} does not match already existing " + "log level {} for {} timer".format( + log_level, self._log_levels[name], name + ) ) return self._timers[name] # If timer does not exist and no log level is provided, @@ -250,7 +257,7 @@ class Timers: log_level = self._max_log_level assert ( log_level <= self._max_log_level - ), 'log level {} is larger than max supported log level {}'.format( + ), "log level {} is larger than max supported log level {}".format( log_level, self._max_log_level ) # Now if the input log level is larger than the one set for @@ -284,7 +291,7 @@ class Timers: if torch.cuda.is_available(): device = torch.cuda.current_device() else: - device = torch.device('cpu') + device = torch.device("cpu") rank_name_to_time = torch.zeros( (world_size, len(names)), dtype=torch.float, device=device @@ -296,7 +303,9 @@ class Timers: if world_size > 1 and _is_distributed() and dist_all_gather_func is not None: try: - dist_all_gather_func(rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1)) + dist_all_gather_func( + rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1) + ) except Exception as e: print(f"Warning: all_gather failed: {e}. Using single rank timing.") @@ -319,30 +328,38 @@ class Timers: ) return name_to_min_max_time - def _get_global_min_max_time_string(self, names, reset, barrier, normalizer, max_only): + def _get_global_min_max_time_string( + self, names, reset, barrier, normalizer, max_only + ): """Report strings for max/minmax times across all ranks.""" - name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer) + name_to_min_max_time = self._get_global_min_max_time( + names, reset, barrier, normalizer + ) if not name_to_min_max_time: return None - + world_size = _get_world_size() if world_size == 1: - output_string = 'time (ms):' + output_string = "time (ms):" for name in name_to_min_max_time: _, max_time = name_to_min_max_time[name] - output_string += '\n {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time) + output_string += "\n {}: {:.2f}".format( + (name + " ").ljust(48, "."), max_time + ) else: if max_only: - output_string = 'max time across ranks (ms):' + output_string = "max time across ranks (ms):" else: - output_string = '(min, max) time across ranks (ms):' + output_string = "(min, max) time across ranks (ms):" for name in name_to_min_max_time: min_time, max_time = name_to_min_max_time[name] if max_only: - output_string += '\n {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time) + output_string += "\n {}: {:.2f}".format( + (name + " ").ljust(48, "."), max_time + ) else: - output_string += '\n {}: ({:.2f}, {:.2f})'.format( - (name + ' ').ljust(48, '.'), min_time, max_time + output_string += "\n {}: ({:.2f}, {:.2f})".format( + (name + " ").ljust(48, "."), min_time, max_time ) return output_string @@ -351,7 +368,7 @@ class Timers: rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier) world_size = _get_world_size() - output_string = 'times across ranks (ms):' + output_string = "times across ranks (ms):" no_reported_timing = True for i, name in enumerate(names): not_yet_found = True @@ -360,13 +377,13 @@ class Timers: no_reported_timing = False if not_yet_found: not_yet_found = False - output_string += '\n {}:'.format(name) + output_string += "\n {}:".format(name) if world_size == 1: - output_string += '\n {:.2f}'.format( + output_string += "\n {:.2f}".format( rank_name_to_time[rank, i] / normalizer ) else: - output_string += '\n rank {:2d}: {:.2f}'.format( + output_string += "\n rank {:2d}: {:.2f}".format( rank, rank_name_to_time[rank, i] / normalizer ) if no_reported_timing: @@ -398,23 +415,23 @@ class Timers: str: Formatted string with the timer values. """ - if names == None: # get all registered timers + if names is None: # get all registered timers names = list(self._timers.keys()) assert normalizer > 0.0 - if self._log_option in ['max', 'minmax']: + if self._log_option in ["max", "minmax"]: max_only = False - if self._log_option == 'max': + if self._log_option == "max": max_only = True output_string = self._get_global_min_max_time_string( names, reset, barrier, normalizer / 1000.0, max_only ) - elif self._log_option == 'all': + elif self._log_option == "all": output_string = self._get_all_ranks_time_string( names, reset, barrier, normalizer / 1000.0 ) else: - raise Exception('unknown timing log option {}'.format(self._log_option)) + raise Exception("unknown timing log option {}".format(self._log_option)) return output_string def log( @@ -444,7 +461,7 @@ class Timers: # If no input rank is provided, log on last rank. world_size = _get_world_size() current_rank = _get_rank() - + if rank is None: rank = world_size - 1 if rank == current_rank and output_string is not None: @@ -476,8 +493,10 @@ class Timers: # torch.utils.add_scalars makes each timer its own run, which # polutes the runs list, so we just add each as a scalar assert normalizer > 0.0 - name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer) + name_to_min_max_time = self._get_global_min_max_time( + names, reset, barrier, normalizer + ) if writer is not None: for name in name_to_min_max_time: _, max_time = name_to_min_max_time[name] - writer.add_scalar(name + '-time', max_time, iteration) \ No newline at end of file + writer.add_scalar(name + "-time", max_time, iteration) diff --git a/workspace/README.md b/workspace/README.md index a74fdf4..b9dbc14 100644 --- a/workspace/README.md +++ b/workspace/README.md @@ -4,7 +4,7 @@ This document explains the key configuration parameters that can be modified for ## Quick Start Checklist 1. **Update run.sh**: Set `code_dir` and `config_path` to your actual paths -2. **Configure GPUs**: Set `CUDA_VISIBLE_DEVICES` for your available GPUs +2. **Configure GPUs**: Set `CUDA_VISIBLE_DEVICES` for your available GPUs 3. **Update config paths**: Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths 4. **Configure robot**: Set `dof_config` and `agent_pos_config` for your robot 5. **Set dataset**: Choose appropriate `repo_id` for your dataset @@ -69,4 +69,4 @@ Keep `agent_pos_config` consistent with `dof_config`. ## Performance Settings (Optional) - `profile`: Enable PyTorch profiling (true/false) -- `padding_side`: Token padding side (left/right) \ No newline at end of file +- `padding_side`: Token padding side (left/right) diff --git a/workspace/lerobot_example/config_qact.yml b/workspace/lerobot_example/config_qact.yml index 60c9d77..7f72500 100644 --- a/workspace/lerobot_example/config_qact.yml +++ b/workspace/lerobot_example/config_qact.yml @@ -60,7 +60,7 @@ agent_pos_config: # Data configuration data: use_lerobot: true - + # LeRobot dataset configuration lerobot_config: repo_id: "lerobot/aloha_mobile_cabinet" @@ -73,10 +73,10 @@ data: 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 @@ -88,7 +88,7 @@ data: - head_actions - height - car_pose - + predict_action_keys: - follow_left_ee_cartesian_pos - follow_left_ee_rotation @@ -99,7 +99,7 @@ data: - head_actions - height - car_pose - + # Image resolution configuration for different camera views resolution: face_view: 256 diff --git a/workspace/lerobot_example/run.sh b/workspace/lerobot_example/run.sh index 8baeb43..380b65c 100644 --- a/workspace/lerobot_example/run.sh +++ b/workspace/lerobot_example/run.sh @@ -21,4 +21,4 @@ export SCRIPT_ARGS="--config ${config_path}/config_qact.yml --seed $MASTER_PORT" echo "Running command: $LAUNCHER $SCRIPT $SCRIPT_ARGS" -$LAUNCHER $SCRIPT $SCRIPT_ARGS \ No newline at end of file +$LAUNCHER $SCRIPT $SCRIPT_ARGS