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

[lint] Update lint (#16)

Browse Source 
* update lint

* update readme

* update ruff lint
This commit is contained in:
Lufang Chen
2025-09-11 13:18:33 +08:00
committed by GitHub
parent a89dce95aa
commit e9332a283d
28 changed files with 2406 additions and 1074 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/data/config.py
+58 -29
View File
@@ -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
return self
wall_x/data/load_lerobot_dataset.py
+121 -30
View File
@@ -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
return dataset
wall_x/data/utils.py
+82 -31
View File
@@ -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 <point> 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
wall_x/fusions/backend.py
+119 -90
View File
@@ -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)
return backend.rope_bwd(
grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled
)
wall_x/fusions/ops.py
+137 -37
View File
@@ -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():
wall_x/model/action_head.py
+105 -80
View File
@@ -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
return loss
wall_x/model/qwen2_5_based/__init__.py
+8 -2
View File
@@ -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
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",
]
wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py
+1 -1
View File
@@ -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,
wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py
+432 -142
View File
File diff suppressed because it is too large Load Diff
wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
+650 -349
View File
File diff suppressed because it is too large Load Diff
wall_x/trainer/qwen_vl_act_trainer.py
+271 -118
View File
@@ -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
self.logger.log(
{f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step
)
start_idx = end_idx
wall_x/utils/constant.py
+228 -57
View File
@@ -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]},
wall_x/utils/timers.py
+56 -37
View File
@@ -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)
writer.add_scalar(name + "-time", max_time, iteration)