From f9641696bc0c5e274db27029cdac8958d6c20699 Mon Sep 17 00:00:00 2001 From: HappyDog060713 Date: Tue, 28 Jul 2026 11:52:50 +0800 Subject: [PATCH 01/15] feat: lingbot-vla-v2 init --- examples/lingbot_vla_v2/README.md | 47 + .../lingbot_vla_v2_inference.py | 119 + telefuser/models/lingbot_vla_v2.py | 3168 +++++++++++++++++ telefuser/models/lingbot_vla_v2_loader.py | 1787 ++++++++++ telefuser/models/lingbot_vla_v2_moe.py | 829 +++++ telefuser/models/lingbot_vla_v2_qwen.py | 746 ++++ .../pipelines/lingbot_vla_v2/__init__.py | 18 + telefuser/pipelines/lingbot_vla_v2/data.py | 155 + .../pipelines/lingbot_vla_v2/pipeline.py | 67 + telefuser/pipelines/lingbot_vla_v2/policy.py | 73 + .../pipelines/lingbot_vla_v2/robot_profile.py | 170 + .../unit/models/test_lingbot_vla_v2_loader.py | 58 + .../unit/pipelines/lingbot_vla_v2/__init__.py | 1 + .../pipelines/lingbot_vla_v2/test_data.py | 101 + .../pipelines/lingbot_vla_v2/test_pipeline.py | 87 + .../lingbot_vla_v2/test_robot_profile.py | 70 + 16 files changed, 7496 insertions(+) create mode 100644 examples/lingbot_vla_v2/README.md create mode 100644 examples/lingbot_vla_v2/lingbot_vla_v2_inference.py create mode 100644 telefuser/models/lingbot_vla_v2.py create mode 100644 telefuser/models/lingbot_vla_v2_loader.py create mode 100644 telefuser/models/lingbot_vla_v2_moe.py create mode 100644 telefuser/models/lingbot_vla_v2_qwen.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/__init__.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/data.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/pipeline.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/policy.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/robot_profile.py create mode 100644 tests/unit/models/test_lingbot_vla_v2_loader.py create mode 100644 tests/unit/pipelines/lingbot_vla_v2/__init__.py create mode 100644 tests/unit/pipelines/lingbot_vla_v2/test_data.py create mode 100644 tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py create mode 100644 tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md new file mode 100644 index 00000000..582bedc5 --- /dev/null +++ b/examples/lingbot_vla_v2/README.md @@ -0,0 +1,47 @@ +# LingBot-VLA v2 RobotWin SDK + +This example loads the official LingBot-VLA v2 6B base checkpoint through TeleFuser and returns a structured +RobotWin action chunk. The current integration verifies the SDK contract without claiming policy quality: every +result is marked `policy_verified=False` and `verification_status="unverified_official_6b_base"`. + +## Inputs + +- Three RGB cameras in the upstream RobotWin order: high, left wrist, right wrist. +- A raw 14-dimensional RobotWin state. +- A non-empty task string. + +The SDK applies the bundled upstream RobotWin `bounds_99_woclip` statistics and maps the observation into +LingBot's 55-dimensional canonical state. + +## Output + +The pipeline returns `LingBotVlaV2ActionChunk` with: + +- `fields["action.arm.position"]`: `[H, 12]`. +- `fields["action.effector.position"]`: `[H, 2]`. +- `raw_actions` and `fields["action"]`: reconstructed `[H, 14]` RobotWin actions. +- `action_mask`: the 55-dimensional canonical RobotWin action mask. +- `horizon`: action chunk length, normally 50 for the official base config. +- `canonical_normalized_actions`: optional `[H, 55]` debugging output. + +## Checkpoints + +The VLA directory must contain `model.safetensors.index.json` and every referenced shard. The Qwen3-VL directory +supplies the visual-language backbone configuration and processor. + +## Example + +```bash +python examples/lingbot_vla_v2/lingbot_vla_v2_inference.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --camera-high /data/cam_high.png \ + --camera-left-wrist /data/cam_left_wrist.png \ + --camera-right-wrist /data/cam_right_wrist.png \ + --task "pick up the red block" \ + --state-json '[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]' \ + --output action_chunk.npz +``` + +The example saves named arrays and verification metadata in an `.npz` file. Do not send the output to a robot until +the official 6B GPU smoke test and policy-level parity validation are complete. diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py new file mode 100644 index 00000000..0d3dcecb --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py @@ -0,0 +1,119 @@ +"""Run LingBot-VLA v2 with a RobotWin observation.""" + +from __future__ import annotations + +import json + +import click +import numpy as np +import torch +from transformers import AutoProcessor + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2 +from telefuser.pipelines.lingbot_vla_v2 import ( + LingBotVlaV2Observation, + LingBotVlaV2Pipeline, + LingBotVlaV2PipelineConfig, + ROBOTWIN_CAMERA_KEYS, +) + + +def get_pipeline( + model_root: str, + qwen3vl_root: str, + device: str = "cuda", + include_canonical_actions: bool = False, +) -> LingBotVlaV2Pipeline: + """Load the official 6B checkpoint and Qwen3-VL processor.""" + dtype = torch.bfloat16 if torch.device(device).type == "cuda" else torch.float32 + processor = AutoProcessor.from_pretrained(qwen3vl_root, local_files_only=True, padding_side="right") + manager = ModuleManager(torch_dtype=dtype, device="cpu") + manager.add_module(processor, "lingbot_vla_v2_processor", path=qwen3vl_root) + load_lingbot_vla_v2(manager, model_root, qwen3vl_root, torch_dtype=dtype) + pipeline = LingBotVlaV2Pipeline(device=device, torch_dtype=dtype) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig(device_type=torch.device(device).type, torch_dtype=dtype), + include_canonical_actions=include_canonical_actions, + ), + ) + return pipeline + + +@click.command() +@click.option("--model-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--qwen3vl-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--camera-high", required=True, type=click.Path(exists=True, dir_okay=False)) +@click.option("--camera-left-wrist", required=True, type=click.Path(exists=True, dir_okay=False)) +@click.option("--camera-right-wrist", required=True, type=click.Path(exists=True, dir_okay=False)) +@click.option("--task", required=True) +@click.option("--state-json", required=True, help="Raw 14-D RobotWin state as a JSON list") +@click.option("--output", default="action_chunk.npz", type=click.Path(dir_okay=False)) +@click.option("--include-canonical-actions", is_flag=True) +@click.option("--seed", default=None, type=int) +@click.option("--device", default="cuda") +def main( + model_root: str, + qwen3vl_root: str, + camera_high: str, + camera_left_wrist: str, + camera_right_wrist: str, + task: str, + state_json: str, + output: str, + include_canonical_actions: bool, + seed: int | None, + device: str, +) -> None: + """Predict and save a structured RobotWin action chunk.""" + try: + state = json.loads(state_json) + except json.JSONDecodeError as error: + raise click.BadParameter("state-json must be valid JSON") from error + if not isinstance(state, list) or len(state) != 14: + raise click.BadParameter("state-json must decode to a 14-element JSON list") + observation = LingBotVlaV2Observation( + task=task, + state=state, + images=dict( + zip( + ROBOTWIN_CAMERA_KEYS, + (camera_high, camera_left_wrist, camera_right_wrist), + strict=True, + ) + ), + ) + pipeline = get_pipeline( + model_root, + qwen3vl_root, + device=device, + include_canonical_actions=include_canonical_actions, + ) + try: + chunk = pipeline(observation, seed=seed) + arrays = { + "action": chunk.raw_actions.numpy(), + "action_arm_position": chunk.fields["action.arm.position"].numpy(), + "action_effector_position": chunk.fields["action.effector.position"].numpy(), + "action_mask": chunk.action_mask.numpy(), + "horizon": np.asarray(chunk.horizon), + "robot_profile": np.asarray(chunk.robot_profile), + "policy_verified": np.asarray(chunk.policy_verified), + "verification_status": np.asarray(chunk.verification_status), + } + if chunk.canonical_normalized_actions is not None: + arrays["canonical_normalized_actions"] = chunk.canonical_normalized_actions.numpy() + np.savez(output, **arrays) + click.echo( + f"Saved {chunk.horizon}-step RobotWin action chunk to {output}; " + f"policy status: {chunk.verification_status}" + ) + finally: + pipeline.close() + + +if __name__ == "__main__": + main() diff --git a/telefuser/models/lingbot_vla_v2.py b/telefuser/models/lingbot_vla_v2.py new file mode 100644 index 00000000..d2837f64 --- /dev/null +++ b/telefuser/models/lingbot_vla_v2.py @@ -0,0 +1,3168 @@ +"""Native LingBot-VLA v2 policy and flow-matching implementation. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + +# Copyright 2026 Robbyant Team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from copy import deepcopy +from typing import Any, Dict, Literal, Optional + +from transformers import AutoConfig, PretrainedConfig + +class LingbotVLAConfig(PretrainedConfig): + """Configuration class for Lingbot-VLA. + This is the configuration class to store the configuration of a [`Lingbot-VLA`]. + """ + + model_type = "lingbotvla" + is_composition = True + + def __init__( + self, + vlm_repo_id: Optional[str] = None, + expert_vision_path: Optional[str] = None, + tokenizer_path: Optional[str] = None, + + post_training: bool = False, + adanorm_time: bool = False, + split_gate_liner: bool = False, + nosplit_gate_liner: bool = False, + separate_time_proj: bool = False, + final_norm_adanorm: bool = False, + + enable_expert_vision: bool = False, + expert_vision_type: Optional[str] = None, + freeze_vision_encoder: bool = False, + + incremental_training: bool = False, + depth_incremental_training: bool = False, + reinit_mismatched_weights: bool = False, + + action_dim: int = 14, + max_action_dim: int = 14, + max_state_dim: int = 14, + chunk_size: int = 50, + vlm_causal: bool = False, + tokenizer_max_length: int = 48, + loss_type: str = "fm", + norm_qkv: bool = False, + align_params: Optional[Dict[str, Any]] = None, + use_compile: bool = False, + + use_moe: bool = False, + token_moe_layers: Optional[list] = None, + token_num_experts: int = 32, + token_top_k: int = 1, + token_moe_intermediate_size: int = 256, + token_shared_intermediate_size: int = 256, + bias_update_speed: float = 0.001, + sequence_wise_loss_coeff: float = 0.001, + sequence_wise_mode: str = "per_sequence", + router_z_loss_coeff: float = 0.0, + router_activation: str = "softmax", + routed_scaling_factor: float = 1.0, + use_shared_expert_gate: bool = True, + moe_implementation: Optional[Literal[None, "eager", "fused"]] = None, + split_fused_experts_from_decoder_fsdp: bool = False, + expert_hidden_size: int = 768, + expert_intermediate_size: int = 2752, + action_num_attention_heads: int = 16, + action_num_key_value_heads: int = 2, + action_head_dim: int = 128, + action_fp32: bool = False, + use_qwen3_chat_template: bool = False, + return_image_grid_thw: bool = False, + qwen3vl_use_vision_boundaries: bool = False, + precompute_grid_thw: bool = False, + use_qwen3_fixed_grid_cache: bool = False, + + use_lm_head: bool = False, + vocab_size: int = 0, + vit_attn_implementation: str = "flash_attention_2", + attention_implementation: str = "flex", + + train_expert_only: bool = False, + train_state_proj: bool = True, + + **kwargs + ): + super().__init__() + if moe_implementation is None: + moe_implementation = kwargs.pop("_moe_implementation", None) + self.architectures = ["LingbotVlaPolicy"] + self.train_state_proj = train_state_proj + self.train_expert_only = train_expert_only + self.use_cache = False + self.attention_implementation = attention_implementation + self.num_steps = 10 + self.n_obs_steps = 1 + + assert not (split_gate_liner and nosplit_gate_liner), \ + "split_gate_liner and nosplit_gate_liner can not be both True." + + self.vlm_repo_id = vlm_repo_id + self.expert_vision_path = expert_vision_path + self.tokenizer_path = tokenizer_path + self.post_training = post_training + self.adanorm_time = adanorm_time + self.split_gate_liner = split_gate_liner + self.nosplit_gate_liner = nosplit_gate_liner + self.enable_expert_vision = enable_expert_vision + self.expert_vision_type = expert_vision_type + self.incremental_training = incremental_training + self.depth_incremental_training = depth_incremental_training + self.reinit_mismatched_weights = reinit_mismatched_weights + self.norm_qkv = norm_qkv + self.use_compile = use_compile + self.loss_type = loss_type + self.separate_time_proj = separate_time_proj + self.final_norm_adanorm = final_norm_adanorm + self.freeze_vision_encoder = freeze_vision_encoder + self.tokenizer_max_length = tokenizer_max_length + self.action_dim = action_dim + self.max_action_dim = max_action_dim + self.max_state_dim = max_state_dim + self.chunk_size = chunk_size + self.n_action_steps = chunk_size + self.vlm_causal = vlm_causal + self.align_params = align_params + self.use_moe = use_moe + if self.use_moe: + self.token_moe_layers = token_moe_layers + self.token_num_experts = token_num_experts + self.token_top_k = token_top_k + self.token_moe_intermediate_size = token_moe_intermediate_size + self.token_shared_intermediate_size = token_shared_intermediate_size + self.bias_update_speed = bias_update_speed + self.sequence_wise_loss_coeff = sequence_wise_loss_coeff + self.sequence_wise_mode = sequence_wise_mode + self.router_z_loss_coeff = router_z_loss_coeff + self.router_activation = router_activation + self.routed_scaling_factor = routed_scaling_factor + self.use_shared_expert_gate = use_shared_expert_gate + self.moe_implementation = moe_implementation + if moe_implementation is not None: + if moe_implementation not in ("eager", "fused"): + raise ValueError(f"Invalid moe_implementation: {moe_implementation}") + self._moe_implementation = moe_implementation + self.split_fused_experts_from_decoder_fsdp = split_fused_experts_from_decoder_fsdp + self.expert_hidden_size = expert_hidden_size + self.expert_intermediate_size = expert_intermediate_size + self.action_num_attention_heads = action_num_attention_heads + self.action_num_key_value_heads = action_num_key_value_heads + self.action_head_dim = action_head_dim + self.action_fp32 = action_fp32 + self.use_qwen3_chat_template = use_qwen3_chat_template + self.return_image_grid_thw = return_image_grid_thw + self.qwen3vl_use_vision_boundaries = qwen3vl_use_vision_boundaries + self.precompute_grid_thw = precompute_grid_thw + self.use_qwen3_fixed_grid_cache = use_qwen3_fixed_grid_cache + self.use_lm_head = use_lm_head + if vocab_size == 0: + if vlm_repo_id and 'paligemma' in vlm_repo_id.lower(): + self.vocab_size = 257216 + elif vlm_repo_id and 'qwen' in vlm_repo_id.lower(): + self.vocab_size = 151936 + else: + self.vocab_size = 257152 + else: + self.vocab_size = vocab_size + self.vit_attn_implementation = vit_attn_implementation + +class LingbotVLAV2Config(LingbotVLAConfig): + def __init__(self, **kwargs): + kwargs.setdefault("attention_implementation", "flex_cached") + kwargs.setdefault("vit_attn_implementation", "flash_attention_2") + kwargs.setdefault("action_num_attention_heads", 32) + kwargs.setdefault("action_num_key_value_heads", 8) + kwargs.setdefault("action_head_dim", 128) + kwargs.setdefault("expert_hidden_size", 768) + kwargs.setdefault("use_qwen3_chat_template", True) + kwargs.setdefault("return_image_grid_thw", True) + kwargs.setdefault("qwen3vl_use_vision_boundaries", True) + kwargs.setdefault("use_qwen3_fixed_grid_cache", True) + super().__init__(**kwargs) + self.architectures = ["LingbotVlaV2Policy"] + self.vlm_family = "qwen3_vl" + + +ConfigClass = [LingbotVLAConfig, LingbotVLAV2Config] +__all__ = ["LingbotVLAConfig", "LingbotVLAV2Config"] + + + +# Shared V1 flow-matching base retained by the V2 architecture. +from logging import raiseExceptions +import einops +import numpy as np +import torch +from torch import nn +import torch.nn.functional as F +from torch import Tensor, nn +from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Union +from functools import partial +import math +from transformers import ( + AutoConfig, + PretrainedConfig, + PreTrainedModel, +) +from transformers.models.auto import CONFIG_MAPPING +from transformers import AutoTokenizer +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from transformers.utils import ( + is_torchdynamo_compiling, + logging, +) + + +class LossKwargs(TypedDict, total=False): + labels: Optional[torch.LongTensor] +from transformers.utils.deprecation import deprecate_kwarg +from transformers.activations import ACT2FN +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, is_flash_attn_available +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.processing_utils import Unpack +from telefuser.models.lingbot_vla_v2_qwen import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLTextModel, Qwen2_5_VLPreTrainedModel + +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2RMSNorm, +) + +try: + from dinov3.hub.backbones import ( + dinov3_vits16, + dinov3_vits16plus, + dinov3_vitb16, + ) +except: pass +from telefuser.models.lingbot_vla_v2_loader import ( + create_sinusoidal_pos_embedding, + make_att_2d_masks, + resize_with_pad, + sample_beta, +) +from telefuser.models.lingbot_vla_v2_loader import apply_rope, our_eager_attention_forward +from telefuser.models.lingbot_vla_v2_loader import flex_attention_forward +from telefuser.models.lingbot_vla_v2_loader import build_block_mask, flex_attention_with_block_mask +import time + +from telefuser.models.lingbot_vla_v2_loader import LingBotVLAWeightLoader, TaskTokenDepthHead +from telefuser.models.lingbot_vla_v2_moe import ( + Qwen2ForCausalLM, + Qwen2FusedExperts, + Qwen2TokenMoeBlock, + FixQwen2RMSNorm, +) + +logger = logging.get_logger(__name__) + +class QwenvlWithExpertConfig(PretrainedConfig): + model_type = "QwenvlWithExpertModel" + sub_configs = {"qwenvl_config": AutoConfig, "qwen_expert_config": AutoConfig} + + def __init__( + self, + qwenvl_config: dict | None = None, + qwen_expert_config: dict | None = None, + freeze_vision_encoder: bool = False, + train_expert_only: bool = False, + vocab_size: int = 257152, + use_lm_head: bool = False, + attention_implementation: str = "eager", + tokenizer_path: str | None = None, + enable_expert_vision: bool = False, + expert_vision_type: str | None = None, + use_cache: bool = False, + expert_hidden_size: int = 768, + expert_intermediate_size: int = 2752, + **kwargs, + ): + self.freeze_vision_encoder = freeze_vision_encoder + self.train_expert_only = train_expert_only + self.attention_implementation = attention_implementation + self.tokenizer_path = tokenizer_path + self.enable_expert_vision = enable_expert_vision + self.expert_vision_type = expert_vision_type + self.vocab_size = vocab_size + self.use_lm_head = use_lm_head + if qwenvl_config is None: + self.qwenvl_config = CONFIG_MAPPING["qwen2_5_vl"]( + attention_dropout=0.0, + bos_token_id=151643, + eos_token_id=151645, + vision_start_token_id=151652, + vision_end_token_id=151653, + vision_token_id=151654, + image_token_id=151655, + video_token_id=151656, + hidden_act="silu", + hidden_size=2048, + initializer_range=0.02, + intermediate_size=11008, + max_position_embeddings=128000, + max_window_layers=70, + model_type="qwen2_5_vl", + num_attention_heads=16, + num_hidden_layers=36, + num_key_value_heads=2, + rms_norm_eps=1e-06, + rope_theta=1000000.0, + sliding_window=32768, + tie_word_embeddings=True, + torch_dtype="bfloat16", + transformers_version="4.41.2", + use_cache=True, + use_sliding_window=False, + vision_config={ + "depth": 32, + "hidden_act": "silu", + "hidden_size": 1280, + "intermediate_size": 3420, + "num_heads": 16, + "in_chans": 3, + "out_hidden_size": 2048, + "patch_size": 14, + "spatial_merge_size": 2, + "spatial_patch_size": 14, + "window_size": 112, + "fullatt_block_indexes": [ + 7, + 15, + 23, + 31 + ], + "tokens_per_second": 2, + "temporal_patch_size": 2 + }, + rope_scaling={ + "type": "mrope", + "mrope_section": [ + 16, + 24, + 24 + ] + }, + vocab_size=151936, + ) + elif isinstance(self.qwenvl_config, dict): + if "model_type" not in qwen_expert_config: + qwenvl_config["model_type"] = "qwen2_5_vl" + + cfg_cls = CONFIG_MAPPING[qwenvl_config["model_type"]] + self.qwenvl_config = cfg_cls(**qwenvl_config) + + if qwen_expert_config is None: + self.qwen_expert_config = CONFIG_MAPPING["qwen2"]( + attention_dropout=0.0, + bos_token_id=151643, + eos_token_id=151645, + hidden_act="silu", + hidden_size=expert_hidden_size, + head_dim=128, + initializer_range=0.02, + intermediate_size=expert_intermediate_size, + max_position_embeddings=32768, + max_window_layers=21, + model_type="qwen2", + num_attention_heads=16, + num_hidden_layers=36, + num_key_value_heads=2, + rms_norm_eps=1e-06, + rope_theta=1000000.0, + sliding_window=32768, + tie_word_embeddings=True, + torch_dtype="bfloat16", + transformers_version="4.43.1", + use_cache=use_cache, + use_sliding_window=False, + vocab_size=151936, + ) + elif isinstance(self.qwen_expert_config, dict): + if "model_type" not in qwen_expert_config: + qwen_expert_config["model_type"] = "qwen2" + + cfg_cls = CONFIG_MAPPING[qwenvl_config["model_type"]] + self.qwen_expert_config = cfg_cls(**qwen_expert_config) + + super().__init__(**kwargs) + + def __post_init__(self): + super().__post_init__() + if self.train_expert_only and not self.freeze_vision_encoder: + raise ValueError( + "You set `freeze_vision_encoder=False` and `train_expert_only=True` which are not compatible." + ) + + if self.attention_implementation not in ["eager", "fa2", "flex"]: + raise ValueError( + f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). Expected 'eager', 'fa2' or 'flex'." + ) + +class AdaRMSNorm(nn.Module): + def __init__(self, hidden_size, cond_dim, eps=1e-6): + """ + AdaRMSNorm: RMSNorm + FiLM + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.gamma = nn.Linear(cond_dim, hidden_size) + self.beta = nn.Linear(cond_dim, hidden_size) + + # DiT style init: gamma.weight=0, gamma.bias=1; beta.weight=0, beta.bias=0 + nn.init.zeros_(self.gamma.weight) + nn.init.zeros_(self.gamma.bias) + nn.init.zeros_(self.beta.weight) + nn.init.zeros_(self.beta.bias) + + def forward(self, hidden_states, cond): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + hidden_states = self.weight * hidden_states + # cond = cond.to(torch.float32) + gamma = self.gamma(cond).unsqueeze(1) # [B, 1, H] + beta = self.beta(cond).unsqueeze(1) # [B, 1, H] + hidden_states = (1 + gamma.to(torch.float32)) * hidden_states + beta.to(torch.float32) + return hidden_states.to(input_dtype) + +class FixAdaRMSNorm(nn.Module): + def __init__(self, hidden_size, cond_dim, eps=1e-6): + """ + AdaRMSNorm: RMSNorm + FiLM + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.gamma = nn.Linear(cond_dim, hidden_size) + self.beta = nn.Linear(cond_dim, hidden_size) + + # DiT style init: gamma.weight=0, gamma.bias=1; beta.weight=0, beta.bias=0 + nn.init.zeros_(self.gamma.weight) + nn.init.zeros_(self.gamma.bias) + nn.init.zeros_(self.beta.weight) + nn.init.zeros_(self.beta.bias) + + def forward(self, hidden_states, cond): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + hidden_states = self.weight * hidden_states + cond = cond.to(torch.float32) + gamma = self.gamma(cond).unsqueeze(1) # [B, 1, H] + beta = self.beta(cond).unsqueeze(1) # [B, 1, H] + hidden_states = (1 + gamma.to(torch.float32)) * hidden_states + beta.to(torch.float32) + return hidden_states.to(input_dtype) + +# HACK: show directly use this norm during initialization +# TODO: clear the logics +def replace_lnorm_with_adanorm(module, hidden_size, cond_dim, final_norm_adanorm): + for name, child in module.named_children(): + if final_norm_adanorm: + if isinstance(child, Qwen2RMSNorm): + if 'q_layernorm' not in name and 'k_layernorm' not in name: + setattr(module, name, AdaRMSNorm(hidden_size, cond_dim)) + elif isinstance(child, FixQwen2RMSNorm): + if 'q_layernorm' not in name and 'k_layernorm' not in name: + setattr(module, name, FixAdaRMSNorm(hidden_size, cond_dim)) + else: + replace_lnorm_with_adanorm(child, hidden_size, cond_dim, final_norm_adanorm) + else: + if isinstance(child, Qwen2RMSNorm): + if 'q_layernorm' not in name and 'k_layernorm' not in name: + setattr(module, name, AdaRMSNorm(hidden_size, cond_dim)) + else: + replace_lnorm_with_adanorm(child, hidden_size, cond_dim, final_norm_adanorm) + +class QwenvlWithExpertModel(PreTrainedModel): + config_class = QwenvlWithExpertConfig + + def __init__(self, config: QwenvlWithExpertConfig, eval=False): + super().__init__(config=config) + self.config = config + vlm_config = AutoConfig.from_pretrained(self.config.tokenizer_path, local_files_only=True) + vlm_config.vision_config.initializer_range = 0.02 + print(f'=====Vocab_size in Config is {self.config.vocab_size}=====') + if self.config.vocab_size != 0 and self.config.vocab_size != 257152 and vlm_config.vocab_size != self.config.vocab_size: + vlm_config.vocab_size = self.config.vocab_size + print(f'====Vocabulary Size is {vlm_config.vocab_size}====') + vlm_config._attn_implementation = "flash_attention_2" + vlm_config.vision_config._attn_implementation = self.config.vit_attn_implementation + self.qwenvl = Qwen2_5_VLForConditionalGeneration._from_config(vlm_config) + if self.config.use_lm_head: + self.qwenvl.tie_weights() + self.config.qwen_expert_config._attn_implementation = "flash_attention_2" + self.qwen_expert = Qwen2ForCausalLM._from_config(self.config.qwen_expert_config, eval=eval) + + if getattr(self.config, 'adanorm_time', False): + replace_lnorm_with_adanorm(self.qwen_expert, self.config.qwen_expert_config.hidden_size, self.config.qwen_expert_config.hidden_size, config.final_norm_adanorm) + if getattr(self.config, 'use_moe', False): + bias_update_speed = getattr(self.config, 'bias_update_speed', 0.001) + hidden_size = self.config.qwen_expert_config.hidden_size # 768 + + token_moe_layers = getattr(self.config, 'token_moe_layers', None) or [] + + if token_moe_layers: + token_config = CONFIG_MAPPING['qwen2_moe']( + num_experts=getattr(self.config, 'token_num_experts', 32), + num_experts_per_tok=getattr(self.config, 'token_top_k', 1), + norm_topk_prob=True, + hidden_size=hidden_size, + moe_intermediate_size=getattr(self.config, 'token_moe_intermediate_size', 256), + shared_expert_intermediate_size=getattr(self.config, 'token_shared_intermediate_size', 256), + output_router_logits=False, + ) + token_config.bias_update_speed = bias_update_speed + token_config._moe_implementation = getattr(self.config, '_moe_implementation', None) + token_config.router_activation = getattr(self.config, 'router_activation', 'softmax') + token_config.routed_scaling_factor = getattr(self.config, 'routed_scaling_factor', 1.0) + token_config.use_shared_expert_gate = getattr(self.config, 'use_shared_expert_gate', True) + for idx in token_moe_layers: + self.qwen_expert.model.layers[idx].mlp = Qwen2TokenMoeBlock(token_config) + # Precomputed grid_thw cache (populated on first call when precompute_grid_thw=True) + self.rotary_pos_emb = None + self.window_index = None + self.cu_window_seqlens = None + self.cu_seqlens = None + + # Remove unused embed_tokens + del self.qwen_expert.model.embed_tokens + if self.config.enable_expert_vision: + if 'dinov3_vitb16' in self.config.expert_vision_type: + self.expert_visual = dinov3_vitb16(pretrained=False) + self.expert_visual_mlp = nn.Sequential( + nn.Linear(self.expert_visual.embed_dim, self.expert_visual.embed_dim * 2), + nn.GELU(), + nn.Linear(self.expert_visual.embed_dim * 2, self.config.qwen_expert_config.hidden_size), + ) + self.attention_interface = self.get_attention_interface() + + # self.to_bfloat16_like_physical_intelligence() + self.set_requires_grad() + + def set_requires_grad(self): + """sets the requires_grad attribute of the model parameters based on the configuration. + If `freeze_vision_encoder` is True, the vision tower parameters are frozen. + If `train_expert_only` is True, the entire Qwenvl model is frozen. + """ + if self.config.freeze_vision_encoder: + self.qwenvl.visual.eval() + for params in self.qwenvl.visual.parameters(): + params.requires_grad = False + + if self.config.train_expert_only: + self.qwenvl.eval() + for params in self.qwenvl.parameters(): + params.requires_grad = False + + def train(self, mode: bool = True): + super().train(mode) + if self.config.freeze_vision_encoder: + self.qwenvl.visual.eval() + if self.config.train_expert_only: + self.qwenvl.eval() + + def to_bfloat16_like_physical_intelligence(self): + """casts the model to bfloat16. + + Modules not casted to bfloat16: + - .qwenvl.model.embed_tokens.weight + - .qwenvl.model.norm.weight + - qwen_expert.model.norm.weight + - qwen_expert.lm_head.weight + """ + self.qwenvl = self.qwenvl.to(dtype=torch.bfloat16) + + params_to_change_dtype = [ + ".qwenvl.model.layers", + "qwen_expert.model.layers", + "visual", + "multi_modal", + ] + for name, param in self.named_parameters(): + if any(selector in name for selector in params_to_change_dtype): + param.data = param.data.to(dtype=torch.bfloat16) + + def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None, precompute_grid_thw: bool = False): + """ + Encodes images into continuous embeddings that can be forwarded to the language model. + + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + precompute_grid_thw (`bool`): If True, compute and cache rotary_pos_emb/window_index/cu_seqlens on first call. + """ + if precompute_grid_thw and self.rotary_pos_emb is None: + ( + self.rotary_pos_emb, + self.window_index, + self.cu_window_seqlens, + self.cu_seqlens + ) = self.qwenvl.visual.preprcess_grid_thw(grid_thw=image_grid_thw) + image_embeds = self.qwenvl.visual( + pixel_values, + grid_thw=image_grid_thw, + rotary_pos_emb=self.rotary_pos_emb, + window_index=self.window_index, + cu_window_seqlens=self.cu_window_seqlens, + cu_seqlens=self.cu_seqlens, + ) + split_sizes = (image_grid_thw.prod(-1) // self.qwenvl.visual.spatial_merge_size**2).tolist() + image_embeds = torch.split(image_embeds, split_sizes) + image_embeds = torch.stack(image_embeds, dim=0) + return image_embeds + + def embed_image(self, image: torch.Tensor, patch_size=14, temporal_patch_size=2, precompute_grid_thw=False): + h = w = int(image.shape[1] ** 0.5) + image_grid_thw = torch.tensor([[1, h, w]]*image.shape[0], device=image.device) + image_embeds = self.get_image_features(image, image_grid_thw=image_grid_thw, precompute_grid_thw=precompute_grid_thw) + return image_embeds + # return torch.randn(72, 64, 2048).to(device=image.device, dtype=torch.bfloat16) + + def embed_language_tokens(self, tokens: torch.Tensor): + return self.qwenvl.model.embed_tokens(tokens) + + def handle_kv_cache( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + use_cache: Optional[bool] = None, + fill_kv_cache: Optional[bool] = None, + ): + if use_cache: + if past_key_values is None: + past_key_values = {} + + if fill_kv_cache: + past_key_values[layer_idx] = { + "key_states": key_states, + "value_states": value_states, + } + else: + key_states = torch.cat( + [past_key_values[layer_idx]["key_states"], key_states], dim=1 + ) + value_states = torch.cat( + [past_key_values[layer_idx]["value_states"], value_states], + dim=1, + ) + return key_states, value_states, past_key_values + + def forward( + self, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + vlm_position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + inputs_embeds: List[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + fill_kv_cache: Optional[bool] = None, + ada_cond: List[torch.FloatTensor] = None, + ): + """ + Args: + attention_mask (Optional[torch.Tensor], optional): + Attention mask with shape (b, seq_len, seq_len). Defaults to None. + position_ids (Optional[torch.LongTensor], optional): + Position indices for applying RoPE. Defaults to None. + past_key_values (Optional[Union[List[torch.FloatTensor], Cache]], optional): + Optional kv cache. Defaults to None. + inputs_embeds (List[torch.FloatTensor], optional): + Input embeddings. Defaults to None. + use_cache (Optional[bool], optional): + Whether to use kv cache. Defaults to None. + fill_kv_cache (Optional[bool], optional): + Whether to return kv tensors in this forward pass as cache. Defaults to None. + + Returns: + outputs_embeds (torch.Tensor): Output embeddings. + past_key_values (Optional[Union[List[torch.FloatTensor], Cache]]): + Optional kv cache. + """ + models = [self.qwenvl.model, self.qwen_expert.model] # Qwen2_5_VLTextModel, Qwen2Model (We have re-writeen their forward as follows:) + + # RMSNorm + num_layers = self.qwenvl.config.num_hidden_layers # 36 + action_num_layers = self.config.qwen_expert_config.num_hidden_layers # 36 + assert action_num_layers == num_layers, ( + "Action expert and VLM must have the same number of layers " + f"(got action={action_num_layers}, vlm={num_layers})." + ) + + router_logits_list = [] + for layer_idx in range(num_layers): + query_states = [] + key_states = [] + value_states = [] + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is None: + continue + if i == 1: # For action expert + query_state, key_state, value_state = models[i].layers[layer_idx](hidden_states, compute_kqv=True, ada_cond = ada_cond) + else: # For VLM + query_state, key_state, value_state = models[i].layers[layer_idx](hidden_states, compute_kqv=True) + + if query_state.dtype != torch.float32: + query_state, key_state, value_state = query_state.to(torch.float32), key_state.to(torch.float32), value_state.to(torch.float32) + query_states.append(query_state) + key_states.append(key_state) + value_states.append(value_state) + + # B,L,H,D with L sequence length (img, lang, state, action), H number of heads, D head dim + # concatenate on the number of embeddings/tokens + query_states = torch.cat(query_states, dim=1) + key_states = torch.cat(key_states, dim=1) + value_states = torch.cat(value_states, dim=1) + + query_states = apply_rope(query_states, position_ids) + key_states = apply_rope(key_states, position_ids) + + key_states, value_states, past_key_values = self.handle_kv_cache( + key_states, + value_states, + layer_idx, + past_key_values=past_key_values, + use_cache=use_cache, + fill_kv_cache=fill_kv_cache, + ) + if self.config.attention_implementation == "flex_cached": + if layer_idx == 0: + _full_len = query_states.shape[1] + _full_block_mask = build_block_mask(attention_mask, self.qwenvl.config.num_attention_heads, _full_len, _full_len) + att_output = flex_attention_with_block_mask(query_states, key_states, value_states, _full_block_mask, query_states.shape[1]) + else: + att_output = self.attention_interface(query_states, key_states, value_states, attention_mask) + + # first part of att_output is prefix (up to sequence length, [:, 0:prefix_seq_len]) + outputs_embeds = [] + start = 0 + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is not None: + end = start + hidden_states.shape[1] + if i == 1: + out_emb, _router_logits = models[i].layers[layer_idx](hidden_states, att_output, start, end, output_atten=True, ada_cond = ada_cond) + if _router_logits is not None: + router_logits_list.append(_router_logits) + else: + out_emb = models[i].layers[layer_idx](hidden_states, att_output, start, end, output_atten=True) + outputs_embeds.append(out_emb) + start = end + else: + outputs_embeds.append(None) + + inputs_embeds = outputs_embeds + + # final norm + outputs_embeds = [] + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is not None: + if self.config.final_norm_adanorm: + if i == 1: + out_emb, _ = models[i].norm(hidden_states, ada_cond) + else: + out_emb = models[i].norm(hidden_states) + else: + out_emb = models[i].norm(hidden_states) + outputs_embeds.append(out_emb) + else: + outputs_embeds.append(None) + + return outputs_embeds, past_key_values, router_logits_list + + def get_attention_interface(self): + if self.config.attention_implementation == "fa2": + raise NotImplementedError("FA2 is not implemented (yet)") + elif self.config.attention_implementation == "flex": + print('=====Using Flex Attn=====') + attention_interface = flex_attention_forward + elif self.config.attention_implementation == "eager": + print('=====Using Eager Attn=====') + attention_interface = our_eager_attention_forward + elif self.config.attention_implementation == "flex_cached": + print('=====Using Flex Cached (prebuilt BlockMask) Attn=====') + attention_interface = flex_attention_forward # fallback + elif self.config.attention_implementation == "xformer": + # attention_interface = xformer_attention_forward + raise NotImplementedError("Xformer attention is not implemented (yet)") + else: + raise ValueError( + f"Invalid attention implementation: {self.config.attention_implementation}. " + "Expected one of ['fa2', 'flex', 'flex_cached', 'eager', 'xformer']." + ) + return attention_interface + +class LingbotVlaPolicy(PreTrainedModel): + config_class = LingbotVLAConfig + name = "torch_lingbot_vla" + supports_gradient_checkpointing = True + + _no_split_modules = ["Qwen2DecoderLayer", "FixQwen2RMSNorm", "FixAdaRMSNorm"] # NOTE: if moudule in Qwen2DecoderLayer, it doesn't need to specify in _no_split_modules + + def get_parallel_plan(self): + from telefuser.models.lingbot_vla_v2_loader import NativeParallelPlan as ParallelPlan + from torch.distributed._tensor import Shard + ep_plan = { + "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.gate_proj": Shard(0), + "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.up_proj": Shard(0), + "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.down_proj": Shard(0), + } + return ParallelPlan(ep_plan=ep_plan) + + @classmethod + def get_weight_loader(cls): + return LingBotVLAWeightLoader() + + def __init__( + self, + config: LingbotVLAConfig, + eval: bool=False, + ): + """ + Args: + config: Policy configuration class instance or None, in which case the default instantiation of + the configuration class is used. + """ + + super().__init__(config) + self.config = config + self.language_tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, local_files_only=True) + self.model = FlowMatching(config, eval) + + if not getattr(self.config,"use_lm_head", False): + del self.model.qwenvl_with_expert.qwenvl.lm_head + del self.model.qwenvl_with_expert.qwen_expert.lm_head + + self.reset() + torch.set_float32_matmul_precision("high") + + def reset(self): + return None + + def get_optim_params(self) -> dict: + return self.parameters() + + def forward( + self, images, img_masks, state, lang_tokens, lang_masks, actions, joint_mask=None, action_is_pad=None, noise=None, time=None, vlm_causal=False, depth_targets=None, precompute_grid_thw=False, future_depth_targets=None, **kwargs + ) -> tuple[Tensor, dict[str, Tensor]]: + loss_dict = {} + # Keep state and actions in fp32 for action learning precision + if getattr(self.config, 'action_fp32', False): + state = state.float() + actions = actions.float() + losses, loss_depth, depth_preds, seq_wise_loss, moe_metrics = self.model.forward( + images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, vlm_causal, self.config.loss_type, depth_targets, precompute_grid_thw=precompute_grid_thw, future_depth_targets=future_depth_targets, + ) + + if joint_mask is not None: + if 'repeat' in self.config.loss_type: + joint_mask = joint_mask.repeat(2,1) + mask_expanded = joint_mask.unsqueeze(1).expand(-1, losses.size(1), -1) # (B, T, D) + masked_losses = losses * mask_expanded + + valid_counts = mask_expanded.sum(dim=(1, 2)).clamp(min=1) + batch_mean_losses = masked_losses.sum(dim=(1, 2)) / valid_counts + loss_vla = masked_losses.sum() / mask_expanded.sum().clamp(min=1) + else: + losses = losses[:, :, :self.config.action_dim] + batch_mean_losses = losses.mean(dim=(1, 2)) + loss_vla = losses.mean() + + loss_dict["batch_mean_losses"] = batch_mean_losses.detach() + + total_loss = ( + loss_vla + + loss_depth + + seq_wise_loss + ) + + # Attach MoE monitoring metrics to loss_dict + if moe_metrics: + loss_dict.update(moe_metrics) + + return total_loss, loss_vla, loss_depth, seq_wise_loss, loss_dict, depth_preds + +class FlowMatching(nn.Module): + def __init__(self, config, eval): + super().__init__() + self.config = config + + # qwenvl with action expert + qwenvl_with_export_config = QwenvlWithExpertConfig( + freeze_vision_encoder=self.config.freeze_vision_encoder, + train_expert_only=self.config.train_expert_only, + vocab_size=getattr(self.config,"vocab_size", 0), + use_lm_head=getattr(self.config,"use_lm_head", False), + attention_implementation=self.config.attention_implementation, + tokenizer_path=self.config.tokenizer_path, + enable_expert_vision=self.config.enable_expert_vision, + expert_vision_type=self.config.expert_vision_type, + use_cache=getattr(self.config,"use_cache", True), + expert_hidden_size=getattr(self.config, 'expert_hidden_size', 768), + expert_intermediate_size=getattr(self.config, 'expert_intermediate_size', 2752), + ) + qwenvl_with_export_config.adanorm_time = getattr(config, "adanorm_time", False) + qwenvl_with_export_config.final_norm_adanorm = getattr(config, "final_norm_adanorm", False) + qwenvl_with_export_config.vit_attn_implementation = getattr(config, "vit_attn_implementation", "flash_attention_2") + if getattr(config, "use_moe", False): + qwenvl_with_export_config.use_moe = config.use_moe + qwenvl_with_export_config.bias_update_speed = getattr(config, "bias_update_speed", 0.001) + qwenvl_with_export_config.token_moe_layers = getattr(config, "token_moe_layers", None) + qwenvl_with_export_config.token_num_experts = getattr(config, "token_num_experts", 32) + qwenvl_with_export_config.token_top_k = getattr(config, "token_top_k", 1) + qwenvl_with_export_config.token_moe_intermediate_size = getattr(config, "token_moe_intermediate_size", 256) + qwenvl_with_export_config.token_shared_intermediate_size = getattr(config, "token_shared_intermediate_size", 256) + # Pass _moe_implementation through for EP/fused support + qwenvl_with_export_config._moe_implementation = getattr(config, '_moe_implementation', None) + self.qwenvl_with_expert = QwenvlWithExpertModel( + qwenvl_with_export_config, eval + ) + self.config.proj_width = qwenvl_with_export_config.qwen_expert_config.hidden_size + self.config.initializer_range = getattr(qwenvl_with_export_config.qwen_expert_config, "initializer_range", None) + # projection layers + self.state_proj = nn.Linear(self.config.max_state_dim, self.config.proj_width) + self.action_in_proj = nn.Linear( + self.config.max_action_dim, self.config.proj_width + ) + self.action_out_proj = nn.Linear( + self.config.proj_width, self.config.max_action_dim + ) + self.action_time_mlp_in = nn.Linear( + self.config.proj_width * 2, self.config.proj_width + ) + self.action_time_mlp_out = nn.Linear( + self.config.proj_width, self.config.proj_width + ) + self.config.align_params = getattr(self.config, 'align_params', {}) + if self.config.align_params != {}: + self.steps=0 + self.use_depth_align = True + self.init_depth_heads(self.config.align_params) + self.use_future_video = self.config.align_params.get('use_future_video', False) + if self.use_future_video: + self.init_video_heads(self.config.align_params) + else: + self.use_depth_align = False + self.use_future_video = False + self.use_future_video_patch = False + self.use_current_video_patch = False + self.use_current_shared_task_proj = False + self.use_future_video_cls = False + self.use_shared_future_task_proj = False + self.future_video_share_future_depth_query = False + + self.set_requires_grad() + + def init_depth_heads(self, config): + self.llm_image_token_size = config['llm']['image_token_size'] + self.llm_image_input_size = config['llm']['image_input_size'] + self.depth_token_size = config['depth']['token_size'] + self.depth_input_size = config['depth']['input_size'] + self.align_type = config.get('mode', None) + self.model_type = config['depth']['model_type'] + if self.align_type != "query": + raise ValueError(f"Only query depth alignment is supported, got {self.align_type!r}.") + if self.model_type != "MoRGBD": + raise ValueError(f"Only MoRGBD depth distillation is supported, got {self.model_type!r}.") + self.use_future_depth = (config.get('depth') or {}).get('use_future_depth', False) + self.block_future_depth_to_action = (config.get('depth') or {}).get('block_future_depth_to_action', False) + self.detach_future_depth_image_feats = bool( + (config.get('depth') or {}).get('detach_future_image_feats', False) + ) + self.use_future_video = bool(config.get('use_future_video', False)) + self.use_future_video_patch = False + self.use_current_video_patch = False + self.use_current_shared_task_proj = False + self.use_future_video_cls = False + self.use_shared_future_task_proj = False + self.future_video_share_future_depth_query = False + self.num_task_tokens = config['num_task_tokens'] + assert config['depth']['num_backbone_tokens'] % self.num_task_tokens == 0 + self.depth_align_embs = nn.Parameter( + torch.randn( + config['depth']['num_backbone_tokens'], config['llm']['dim_out'] + ) + ) + self.depth_align_embs.requires_grad = True + + self.depth_align_head = TaskTokenDepthHead(config['depth'], llm_hidden_size=config['llm']['dim_out']).to(dtype=torch.bfloat16) + + for p in self.depth_align_head.parameters(): + p.requires_grad = True + + if self.use_future_depth: + self.future_depth_align_embs = nn.Parameter( + torch.randn( + config['depth']['num_backbone_tokens'], config['llm']['dim_out'] + ) + ) + self.future_depth_align_embs.requires_grad = True + + self.future_depth_align_head = TaskTokenDepthHead( + config['depth'], llm_hidden_size=config['llm']['dim_out'] + ).to(dtype=torch.bfloat16) + + for p in self.future_depth_align_head.parameters(): + p.requires_grad = True + + def init_video_heads(self, config): + if self.align_type != "query": + raise ValueError("future-video alignment is only supported for query align mode.") + + video_config = dict(config.get('depth', {})) + video_config.update(config.get('video', {})) + required_keys = ("num_backbone_tokens", "dim_out", "num_layers", "num_heads", "dim_head", "ff_mult") + missing = [key for key in required_keys if key not in video_config] + if missing: + raise ValueError(f"video align config missing required keys: {missing}") + self.use_future_video_patch = bool(video_config.get("use_patch_loss", True)) + self.use_current_video_patch = bool(video_config.get("use_current_patch_loss", False)) + if self.use_current_video_patch and not self.use_future_video_patch: + raise ValueError( + "align_params.video.use_current_patch_loss=True requires " + "align_params.video.use_patch_loss=True." + ) + self.use_current_shared_task_proj = bool( + video_config.get("use_current_shared_task_proj", self.use_current_video_patch) + ) + if self.use_current_shared_task_proj and not self.use_current_video_patch: + raise ValueError( + "align_params.video.use_current_shared_task_proj=True requires " + "align_params.video.use_current_patch_loss=True." + ) + self.use_future_video_cls = bool(video_config.get("use_cls_loss", False)) + self.future_video_share_future_depth_query = bool( + video_config.get("share_future_depth_query", False) + ) + self.use_shared_future_task_proj = bool( + video_config.get("use_shared_future_task_proj", False) + ) + if self.use_shared_future_task_proj and not self.use_future_video_patch: + raise ValueError( + "align_params.video.use_shared_future_task_proj=True requires " + "align_params.video.use_patch_loss=True." + ) + if self.use_shared_future_task_proj and not self.future_video_share_future_depth_query: + raise ValueError( + "align_params.video.use_shared_future_task_proj=True requires " + "align_params.video.share_future_depth_query=True." + ) + if self.future_video_share_future_depth_query: + if not self.use_future_depth: + raise ValueError( + "align_params.video.share_future_depth_query=True requires " + "align_params.depth.use_future_depth=True." + ) + if int(video_config["num_backbone_tokens"]) != int(config["depth"]["num_backbone_tokens"]): + raise ValueError( + "future-video shared query requires video.num_backbone_tokens " + "to match depth.num_backbone_tokens." + ) + + self.block_suffix_to_future_video = bool(video_config.get("block_suffix_to_future_video", False)) + self.future_video_context_mode = str(video_config.get("context_mode", "img_query")).lower() + if self.future_video_context_mode not in ("img_query", "query_only"): + raise ValueError( + "future-video context_mode must be 'img_query' or 'query_only', " + f"got {self.future_video_context_mode!r}." + ) + if self.use_future_video_patch: + if self.use_current_video_patch: + self.current_video_align_embs = nn.Parameter( + torch.randn( + video_config['num_backbone_tokens'], config['llm']['dim_out'] + ) + ) + self.current_video_align_embs.requires_grad = True + if self.use_current_shared_task_proj: + self.current_shared_task_proj = nn.Linear( + config['llm']['dim_out'] * 2, + config['llm']['dim_out'], + ) + for p in self.current_shared_task_proj.parameters(): + p.requires_grad = True + self.current_video_align_head = TaskTokenDepthHead( + video_config, llm_hidden_size=config['llm']['dim_out'] + ).to(dtype=torch.bfloat16) + for p in self.current_video_align_head.parameters(): + p.requires_grad = True + + if ( + not self.future_video_share_future_depth_query + or self.use_shared_future_task_proj + ): + self.future_video_align_embs = nn.Parameter( + torch.randn( + video_config['num_backbone_tokens'], config['llm']['dim_out'] + ) + ) + self.future_video_align_embs.requires_grad = True + if self.use_shared_future_task_proj: + self.future_shared_task_proj = nn.Linear( + config['llm']['dim_out'] * 2, + config['llm']['dim_out'], + ) + for p in self.future_shared_task_proj.parameters(): + p.requires_grad = True + self.future_video_align_head = TaskTokenDepthHead( + video_config, llm_hidden_size=config['llm']['dim_out'] + ).to(dtype=torch.bfloat16) + for p in self.future_video_align_head.parameters(): + p.requires_grad = True + + if self.use_future_video_cls: + self.future_video_cls_align_emb = nn.Embedding(1, config['llm']['dim_out']) + self.future_video_cls_head = nn.Sequential( + nn.LayerNorm(config['llm']['dim_out']), + nn.Linear(config['llm']['dim_out'], video_config['dim_out']), + ).to(dtype=torch.bfloat16) + for p in self.future_video_cls_head.parameters(): + p.requires_grad = True + + def _future_depth_token_count(self): + return self.num_task_tokens if getattr(self, "use_future_depth", False) else 0 + + def _future_video_own_token_count(self): + if not getattr(self, "use_future_video", False): + return 0 + count = 1 if getattr(self, "use_future_video_cls", False) else 0 + if ( + getattr(self, "use_future_video_patch", True) + and not getattr(self, "future_video_share_future_depth_query", False) + ): + count += self.num_task_tokens + return count + + def _future_video_own_span(self, hidden_states): + own_count = self._future_video_own_token_count() + future_depth_count = self._future_depth_token_count() + end = hidden_states.shape[1] - future_depth_count + start = end - own_count + return start, end + + def _future_depth_task_tokens(self, hidden_states): + if not getattr(self, "use_future_depth", False): + raise ValueError("future-depth query tokens are not enabled.") + return hidden_states[:, -self.num_task_tokens:, :] + + def _future_video_cls_task_tokens(self, hidden_states): + if not getattr(self, "use_future_video_cls", False): + return None + start, _ = self._future_video_own_span(hidden_states) + return hidden_states[:, start : start + 1, :] + + def _future_video_patch_task_tokens(self, hidden_states): + if getattr(self, "future_video_share_future_depth_query", False): + return self._future_depth_task_tokens(hidden_states) + start, end = self._future_video_own_span(hidden_states) + if getattr(self, "use_future_video_cls", False): + start += 1 + return hidden_states[:, start:end, :] + + def _current_depth_task_tokens(self, hidden_states, num_images=3): + chunk_size = self.llm_image_token_size * self.llm_image_token_size + image_token_len = chunk_size + (2 if getattr(self.config, "qwen3vl_use_vision_boundaries", False) else 0) + if getattr(self, "use_future_depth", False): + start = num_images * image_token_len + return hidden_states[:, start : start + self.num_task_tokens, :] + end = hidden_states.shape[1] - self._future_video_own_token_count() + start = end - self.num_task_tokens + return hidden_states[:, start:end, :] + + def _future_video_query_span(self, prefix_len): + if not getattr(self, "use_future_video", False): + return prefix_len, prefix_len + future_depth_count = self._future_depth_token_count() + own_count = self._future_video_own_token_count() + end = prefix_len - future_depth_count + return end - own_count, end + + def _block_suffix_to_future_video_(self, att_2d_masks, suffix_row_start, prefix_len): + start, end = self._future_video_query_span(prefix_len) + if end <= start: + return att_2d_masks + att_2d_masks[:, suffix_row_start:, start:end] = False + return att_2d_masks + + def _block_suffix_to_future_video_if_enabled_( + self, + att_2d_masks, + suffix_row_start, + prefix_len, + ): + if not getattr(self, "block_suffix_to_future_video", False): + return att_2d_masks + return self._block_suffix_to_future_video_( + att_2d_masks, + suffix_row_start=suffix_row_start, + prefix_len=prefix_len, + ) + + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, (nn.Linear, nn.Conv3d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.LayerNorm): + if module.weight is not None: + module.weight.data.fill_(1.0) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Qwen2FusedExperts): + module.initializer_range = std + module.reset_parameters() + reset_post_init = getattr(module, "_reset_post_init_parameters", None) + if reset_post_init is not None: + reset_post_init() + + def set_requires_grad(self): + for params in self.state_proj.parameters(): + params.requires_grad = self.config.train_state_proj + + @staticmethod + def _fp32_linear(module, x): + """Compute linear layer in fp32 regardless of module's current parameter dtype.""" + return F.linear( + x.float(), + module.weight.float(), + module.bias.float() if module.bias is not None else None + ) + + def sample_time(self, bsize, device): + time_beta = sample_beta(1.5, 1.0, bsize, device) + time = time_beta * 0.999 + 0.001 + return time.to(dtype=torch.float32, device=device) + + def embed_prefix( + self, images, img_masks, lang_tokens, lang_masks, vlm_causal, precompute_grid_thw=False + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bsize = images.shape[0] + device = images.device + dtype = images.dtype + + # embed image + if images.ndim == 5: + images = einops.rearrange(images, "b n c h w -> (b n) c h w") + elif images.ndim == 4: + images = einops.rearrange(images, "b n l d -> (b n) l d") + elif images.ndim == 3: # For inference bs=1 + bsize = 1 + img_emb = self.qwenvl_with_expert.embed_image(images, precompute_grid_thw=precompute_grid_thw) + num_patch = img_emb.shape[1] + img_emb = einops.rearrange(img_emb, "(b n) l d -> b (n l) d", b=bsize) # bsize = 24 + num_img_embs = img_emb.shape[1] + if img_masks.ndim ==1: # For inference bs=1 + img_masks = img_masks.unsqueeze(0) + if self.use_depth_align and self.align_type == "query": + align_masks = einops.repeat(img_masks, "b n -> b (n l)", l=self.num_task_tokens) + img_masks = einops.repeat(img_masks, "b n -> b (n l)", l=num_patch) + + # embed language + lang_emb = self.qwenvl_with_expert.embed_language_tokens(lang_tokens) + num_lang_embs = lang_emb.shape[1] + + if self.use_depth_align and self.align_type == "query": + def _get_align_tokens(tokens): + tk_weights = tokens.view(self.num_task_tokens, tokens.shape[0] // self.num_task_tokens, tokens.shape[1]) + tk_weights = tk_weights.mean(dim=1) + return tk_weights + + align_embs = _get_align_tokens(self.depth_align_embs).repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) + # align_masks = einops.rearrange(img_masks, "b (n l) -> b n l", n=3) + # align_masks = align_masks[:, :, 0] + # align_masks = einops.repeat(align_masks, "b n -> b (n l)", l=self.num_task_tokens) + embs = torch.cat([img_emb, align_embs, align_embs, align_embs, lang_emb], dim=1) + pad_masks = torch.cat([img_masks, align_masks, lang_masks], dim=1) + else: + # assemble embeddings + embs = torch.cat([img_emb, lang_emb], dim=1) + pad_masks = torch.cat([img_masks, lang_masks], dim=1) + + # (see `make_att_2d_masks` to understand why zeros means bidirection) + if not vlm_causal: + if self.use_depth_align and self.align_type == "query": + att_masks = torch.zeros( + (img_emb.size(0), num_img_embs + 3 * self.num_task_tokens + num_lang_embs), device=device, dtype=torch.bool + ) # 1, bs_img*(768+48) + else: + att_masks = torch.zeros( + (img_emb.size(0), num_img_embs + num_lang_embs), device=device, dtype=torch.bool + ) # 1, bs_img*(768+48) + else: + if self.use_depth_align and self.align_type == "query": + att_masks = torch.ones( + (img_emb.size(0), num_img_embs + 3 * self.num_task_tokens + num_lang_embs), device=device, dtype=torch.bool + ) # 1, bs_img*(768+48) + else: + att_masks = torch.ones( + (img_emb.size(0), num_img_embs + num_lang_embs), device=device, dtype=torch.bool + ) # 1, bs_img*(768+48) + return embs, pad_masks, att_masks + + def embed_suffix(self, state, noisy_actions, timestep): # (torch.Size([state_bs, 32]), torch.Size([1, state_bs*50, 32]), torch.Size([1])) + bsize = state.shape[0] # state_bs = img_bs + device = state.device + dtype = state.dtype + _fp32 = getattr(self.config, 'action_fp32', False) + # embed state + state_emb = self._fp32_linear(self.state_proj, state) if _fp32 else self.state_proj(state) + + # embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1] + time_emb = create_sinusoidal_pos_embedding( # 1, 1024 + timestep, # torch.Size([1])) + self.config.proj_width, # 1024 + min_period=4e-3, + max_period=4.0, + device=device, + ) + time_emb = time_emb.type(dtype=dtype) + + time_emb_ori = time_emb + + # Fuse timestep + action information using an MLP + action_emb = self._fp32_linear(self.action_in_proj, noisy_actions) if _fp32 else self.action_in_proj(noisy_actions) # torch.Size([1, state_bs*50, 1024]) + time_emb = einops.repeat(time_emb, "b d -> b n d", n=action_emb.shape[1]) # [1, 1024] -> [1, state_bs*50, 1024] + action_time_emb = torch.cat([action_emb, time_emb], dim=-1) # [1, state_bs*50, 2048] + + action_time_emb = self._fp32_linear(self.action_time_mlp_in, action_time_emb) if _fp32 else self.action_time_mlp_in(action_time_emb) + action_time_emb = F.silu(action_time_emb) # swish == silu + action_time_emb = self._fp32_linear(self.action_time_mlp_out, action_time_emb) if _fp32 else self.action_time_mlp_out(action_time_emb) # [1, state_bs*50, 1024] + action_time_dim = action_time_emb.shape[1] + + embs = torch.cat([state_emb[:, None], action_time_emb], dim=1) + pad_masks = torch.ones( + (bsize, action_time_dim + 1), device=device, dtype=torch.bool + ) + + # Set attention masks for suffix tokens so that prefix tokens cannot attend to suffix tokens. + # And state token cannot attend action tokens. + # Action tokens use a bidirectional attention. + att_masks = torch.zeros( + (bsize, action_time_dim + 1), device=device, dtype=torch.bool + ) + att_masks[:, :2] = True + + return time_emb_ori, embs, pad_masks, att_masks + + def forward( + self, + images, + img_masks, + lang_tokens, + lang_masks, + state, + actions, + noise=None, + time=None, + vlm_causal=False, + loss_type='fm', + depth_targets=None, + precompute_grid_thw=False, + future_depth_targets=None, + ) -> Tensor: + dtype = state.dtype + device = state.device + if noise is None: + noise = torch.randn(actions.shape, device=device, dtype=dtype) + + if time is None: + time = self.sample_time(actions.size(0), device).to(dtype) + + time_expanded = time[:, None, None] + x_t = time_expanded * noise + (1 - time_expanded) * actions + u_t = noise - actions + + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( + images, img_masks, lang_tokens, lang_masks, vlm_causal, precompute_grid_thw=precompute_grid_thw + ) # 1,bs_img*(768+48),2048 1,bs_img*(768+48) 1,bs_img*(768+48) + time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( + state, x_t, time + ) # [1, state_bs*(50+1), 1024], [1, state_bs*(50+1)], [1, state_bs*(50+1)] state_bs=bs_img + + pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1) # 1,state_bs*(768+48+50+1) + att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)# 1,state_bs*(768+48+50+1) + + # pad_masks = pad_masks.reshape(state.size(0), -1) + # att_masks = att_masks.reshape(state.size(0), -1) + att_2d_masks = make_att_2d_masks(pad_masks, att_masks) # torch.Size([state_bs, 768+48+50+1, 768+48+50+1]) + position_ids = torch.cumsum(pad_masks, dim=1) - 1 # torch.Size([state_bs, 768+48+50+1]) + vlm_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + + # prefix_embs = prefix_embs.reshape(state.size(0), -1, prefix_embs.size(-1)) + # suffix_embs = suffix_embs.reshape(state.size(0), -1, suffix_embs.size(-1)) + (outputs_embeds, suffix_out), _, router_logits_list = self.qwenvl_with_expert.forward( + attention_mask=att_2d_masks, + position_ids=position_ids, + vlm_position_ids=vlm_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, suffix_embs], # bs_img,(768+48),2048 [state_bs, (50+1), 1024] + use_cache=self.config.use_cache, + fill_kv_cache=True, + ada_cond = time_embs if getattr(self.config, 'adanorm_time', False) else None, + ) + if self.config.align_params != {}: + loss_depth, depth_preds = self.depth_emb_forward(outputs_embeds, depth_targets, img_masks) + loss_depth = loss_depth * self.config.align_params['depth_loss_weight'] + self.steps+=1 + else: + loss_depth = 0 + depth_preds = None + suffix_out = suffix_out[:, -self.config.n_action_steps :] + if getattr(self.config, 'action_fp32', False): + v_t = self._fp32_linear(self.action_out_proj, suffix_out) + else: + if suffix_out.dtype != self.action_out_proj.weight.dtype: + suffix_out = suffix_out.to(self.action_out_proj.weight.dtype) + v_t = self.action_out_proj(suffix_out) + # u_t = u_t.reshape(images.size(0), -1, u_t.size(-1)) + if loss_type == 'fm': + losses = F.mse_loss(u_t, v_t, reduction="none") + # losses = torch.mean((v_t - u_t)**2, dim=-1) + elif loss_type == 'L1_fm': + losses = F.l1_loss(u_t, v_t, reduction="none") + + # Sequence-wise balance loss (DeepSeek-V3 style, for token-MoE only) + seq_wise_loss_coeff = getattr(self.config, 'sequence_wise_loss_coeff', 0) + seq_wise_loss = 0 + + if seq_wise_loss_coeff > 0 and router_logits_list: + from telefuser.models.lingbot_vla_v2_loader import triton_sequence_wise_balance_loss + + token_moe_layers_set = set(getattr(self.config, 'token_moe_layers', None) or []) + token_moe_layers_list = sorted(token_moe_layers_set) + token_router_logits = tuple( + logits for i, logits in enumerate(router_logits_list) + if not token_moe_layers_list or (token_moe_layers_list[i] if i < len(token_moe_layers_list) else i) in token_moe_layers_set + ) + + if token_router_logits: + token_top_k = getattr(self.config, 'token_top_k', 4) + + # Batch-wise balance loss: treat all B脳T tokens as one group. + # seq_lengths=None makes the function use all tokens at once, + # giving stable f_i statistics (B脳T脳K assignments / E experts). + layer_losses = triton_sequence_wise_balance_loss( + router_logits_list=token_router_logits, + top_k=token_top_k, + seq_lengths=None, + padding_len=0, + ) + if layer_losses: + seq_wise_loss = seq_wise_loss_coeff * torch.stack(layer_losses).mean() + + # MoE monitoring metrics for token-MoE. + moe_metrics = {} + if router_logits_list: + all_moe_indices = sorted(getattr(self.config, 'token_moe_layers', None) or []) + token_expert_counts = [] + + with torch.no_grad(): + for i, logits in enumerate(router_logits_list): + layer_id = all_moe_indices[i] if i < len(all_moe_indices) else i + num_experts = logits.shape[-1] + routing_probs = F.softmax(logits, dim=1, dtype=torch.float) + + moe_block = self.qwenvl_with_expert.qwen_expert.model.layers[layer_id].mlp + if hasattr(moe_block, 'last_tokens_per_expert'): + counts = moe_block.last_tokens_per_expert.clone() + else: + _, selected = torch.topk(routing_probs, 1, dim=-1) + expert_indices = selected.squeeze(-1) + counts = F.one_hot(expert_indices, num_classes=num_experts).float().sum(dim=0) + + token_expert_counts.append((layer_id, counts)) + + # MaxVio: (max_load - avg_load) / avg_load (paper 2408.15664) + avg_load = counts.mean() + maxvio = (counts.max() - avg_load) / avg_load.clamp(min=1e-9) + moe_metrics[f"token_moe/layer{layer_id}_maxvio"] = maxvio + + per_sample_entropy = -(routing_probs * routing_probs.clamp(min=1e-9).log()).sum(dim=-1) + moe_metrics[f"token_moe/layer{layer_id}_entropy"] = per_sample_entropy.mean() + + # Compute average MaxVio across token-MoE layers + token_maxvio_values = [ + moe_metrics[k] for k in moe_metrics + if k.startswith("token_moe/") and k.endswith("_maxvio") + ] + if token_maxvio_values: + moe_metrics["token_moe/avg_maxvio"] = torch.stack(token_maxvio_values).mean() + + # Avg top-K sigmoid score (before norm) across token-MoE layers + token_moe_layers_list = sorted(getattr(self.config, 'token_moe_layers', None) or []) + if token_moe_layers_list: + sigmoid_scores = [] + for lid in token_moe_layers_list: + moe_block = self.qwenvl_with_expert.qwen_expert.model.layers[lid].mlp + if hasattr(moe_block, 'avg_topk_sigmoid_score'): + sigmoid_scores.append(moe_block.avg_topk_sigmoid_score.detach().to(losses.device)) + if sigmoid_scores: + moe_metrics["token_moe/avg_topk_sigmoid"] = torch.stack(sigmoid_scores).mean() + + if token_expert_counts: + moe_metrics["_token_moe_expert_counts"] = token_expert_counts + + return losses, loss_depth, depth_preds, seq_wise_loss, moe_metrics + + def sample_actions( + self, images, img_masks, lang_tokens, lang_masks, state, vlm_causal=False, noise=None + ) -> Tensor: + """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)""" + bsize = state.shape[0] + device = state.device + dtype = state.dtype + + if noise is None: + actions_shape = ( + bsize, + self.config.n_action_steps, + self.config.max_action_dim, + ) + noise = torch.randn(actions_shape, device=device, dtype=dtype) + + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( + images, img_masks, lang_tokens, lang_masks, vlm_causal + ) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) # bs, prefix_len, prefix_len + prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + + # Compute image and language key value cache + _, past_key_values, _ = self.qwenvl_with_expert.forward( + attention_mask=prefix_att_2d_masks, + position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=self.config.use_cache, + fill_kv_cache=True, + ) + + dt = torch.tensor(-1.0 / self.config.num_steps, dtype=dtype, device=device) + x_t = noise + time = torch.tensor(1.0, dtype=dtype, device=device) + count = 0 + while time >= -dt / 2: + count += 1 + expanded_time = time.expand(bsize) + + v_t = self.predict_velocity( + state, prefix_pad_masks, past_key_values, x_t, expanded_time + ) + + # Euler step + x_t += dt * v_t + time += dt + print(f'Denoise {count} steps') + return x_t + + def predict_velocity(self, state, prefix_pad_masks, past_key_values, x_t, timestep): + """predict velocity at time t using the suffix model.""" + time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( + state, x_t, timestep + ) + + suffix_len = suffix_pad_masks.shape[1] + batch_size = prefix_pad_masks.shape[0] + prefix_len = prefix_pad_masks.shape[1] + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand( + batch_size, suffix_len, prefix_len + ) + + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) # bs, suffix_len, prefix_len+suffix_len + + prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] + position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 + + outputs_embeds, _, _ = self.qwenvl_with_expert.forward( + attention_mask=full_att_2d_masks, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=self.config.use_cache, + fill_kv_cache=False, + ada_cond = time_embs if getattr(self.config, 'adanorm_time', False) else None, + ) + suffix_out = outputs_embeds[1] + suffix_out = suffix_out[:, -self.config.n_action_steps :] + if getattr(self.config, 'action_fp32', False): + v_t = self._fp32_linear(self.action_out_proj, suffix_out) + else: + v_t = self.action_out_proj(suffix_out) + return v_t + + def depth_emb_forward(self, hidden_states, depth_targets=None, img_masks=None, future_depth_targets=None): + chunk_size = self.llm_image_token_size * self.llm_image_token_size + num_images = img_masks.shape[1] if img_masks is not None and img_masks.ndim == 2 else 3 + if img_masks is not None: + img_masks = einops.rearrange(img_masks, 'b n -> (b n)') + image_embs = hidden_states[:, chunk_size * 0 + 1 : chunk_size * 1 + 1, :] + align_embs = self._current_depth_task_tokens(hidden_states, num_images=num_images) + align_embs = torch.cat([image_embs, align_embs], dim=1) + depth_preds = self.depth_align_embs.repeat(align_embs.shape[0], 1, 1).to(dtype=align_embs.dtype, device=align_embs.device) + depth_preds = self.depth_align_head(align_embs, depth_preds).contiguous().float() + current_loss = self._emb_loss(depth_preds, depth_targets) + + if self.use_future_depth: + future_align_embs = self._future_depth_task_tokens(hidden_states) + future_image_embs = ( + image_embs.detach() + if getattr(self, "detach_future_depth_image_feats", False) + else image_embs + ) + future_align_embs = torch.cat([future_image_embs, future_align_embs], dim=1) + future_depth_preds = self.future_depth_align_embs.repeat(future_align_embs.shape[0], 1, 1).to(dtype=future_align_embs.dtype, device=future_align_embs.device) + future_depth_preds = self.future_depth_align_head(future_align_embs, future_depth_preds).contiguous().float() + future_loss = self._emb_loss(future_depth_preds, future_depth_targets) + return current_loss, future_loss, depth_preds, future_depth_preds + + return current_loss, 0, depth_preds, None + + def video_emb_forward( + self, + hidden_states, + future_video_targets=None, + future_video_cls_targets=None, + future_video_current_patch=None, + ): + if self.align_type != 'query': + raise ValueError("future-video alignment is only supported for query align mode.") + + use_patch = getattr(self, "use_future_video_patch", True) + use_cls = getattr(self, "use_future_video_cls", False) + if not use_patch and not use_cls: + raise ValueError("future-video alignment requires use_patch_loss or use_cls_loss to be enabled.") + if use_patch and future_video_targets is None: + raise ValueError("future_video_targets is required when use_patch_loss=True.") + + align_params = getattr(getattr(self, "config", None), "align_params", {}) or {} + video_cfg = align_params.get("video", {}) if hasattr(align_params, "get") else {} + chunk_size = self.llm_image_token_size * self.llm_image_token_size + image_embs = hidden_states[:, chunk_size * 0 + 1 : chunk_size * 1 + 1, :] + image_embs_for_video = image_embs.detach() if bool(video_cfg.get("detach_image_feats", False)) else image_embs + + cls_preds = None + if use_cls: + if future_video_cls_targets is None: + raise ValueError("future_video_cls_targets is required when use_cls_loss=True.") + cls_task_embs = self._future_video_cls_task_tokens(hidden_states) + cls_delta = self.future_video_cls_head(cls_task_embs.squeeze(1)) + cls_preds = cls_delta.contiguous().float() + + loss = None + metrics = {} + video_preds = None + if use_patch: + video_task_embs = self._future_video_patch_task_tokens(hidden_states) + context_mode = str( + video_cfg.get( + "context_mode", + getattr(self, "future_video_context_mode", "img_query"), + ) + ).lower() + if context_mode == "query_only": + video_align_embs = video_task_embs + else: + video_align_embs = torch.cat([image_embs_for_video, video_task_embs], dim=1) + if ( + getattr(self, "future_video_share_future_depth_query", False) + and not getattr(self, "use_shared_future_task_proj", False) + ): + query_embs = self.future_depth_align_embs + else: + query_embs = self.future_video_align_embs + video_preds = query_embs.repeat(video_align_embs.shape[0], 1, 1).to( + dtype=video_align_embs.dtype, device=video_align_embs.device + ) + video_preds = self.future_video_align_head(video_align_embs, video_preds).contiguous().float() + loss, metrics = self._video_emb_loss(video_preds, future_video_targets) + if use_cls: + cls_loss, cls_metrics = self._video_cls_loss(cls_preds, future_video_cls_targets) + loss = cls_loss if loss is None else loss + cls_loss + metrics.update(cls_metrics) + return loss, video_preds, metrics + + def current_video_emb_forward( + self, + hidden_states, + current_video_targets=None, + ): + if self.align_type != 'query': + raise ValueError("current-video alignment is only supported for query align mode.") + if not getattr(self, "use_current_video_patch", False): + raise ValueError("current-video alignment requires use_current_patch_loss=True.") + if current_video_targets is None: + raise ValueError("current_video_targets is required for current-video alignment.") + + chunk_size = self.llm_image_token_size * self.llm_image_token_size + image_embs = hidden_states[:, chunk_size * 0 + 1 : chunk_size * 1 + 1, :] + current_task_embs = self._current_depth_task_tokens(hidden_states) + align_embs = torch.cat([image_embs, current_task_embs], dim=1) + queries = self.current_video_align_embs.repeat(align_embs.shape[0], 1, 1).to( + dtype=align_embs.dtype, + device=align_embs.device, + ) + preds = self.current_video_align_head(align_embs, queries).contiguous().float() + loss, metrics = self._video_emb_loss( + preds, + current_video_targets, + metric_prefix="current_video", + ) + return loss, preds, metrics + + def _video_emb_loss(self, video_preds, future_video_targets, metric_prefix="future_video"): + align_params = getattr(getattr(self, "config", None), "align_params", {}) or {} + video_cfg = align_params.get("video", {}) if hasattr(align_params, "get") else {} + use_smooth_l1 = bool(video_cfg.get("use_smooth_l1_loss", True)) + use_mse = bool(video_cfg.get("use_mse_loss", False)) + use_cosine = bool(video_cfg.get("use_cosine_loss", False)) + if not use_smooth_l1 and not use_mse and not use_cosine: + raise ValueError(f"{metric_prefix} loss requires smooth-L1, MSE, and/or cosine loss.") + + metrics = {} + loss = None + if use_smooth_l1: + smooth_l1_loss = self._emb_loss(video_preds, future_video_targets) + metrics[f"align/{metric_prefix}_smooth_l1_loss"] = smooth_l1_loss.detach() + loss = smooth_l1_loss + if use_mse: + target = future_video_targets.to(dtype=video_preds.dtype, device=video_preds.device) + mse_loss = F.mse_loss(video_preds.float(), target.float().detach()) + mse_weight = float(video_cfg.get("mse_loss_weight", 1.0)) + metrics[f"align/{metric_prefix}_mse_loss"] = mse_loss.detach() + metrics[f"align/{metric_prefix}_mse_loss_weighted"] = (mse_loss * mse_weight).detach() + weighted_mse_loss = mse_loss * mse_weight + loss = weighted_mse_loss if loss is None else loss + weighted_mse_loss + if use_cosine: + target = future_video_targets.to(dtype=video_preds.dtype, device=video_preds.device) + pred_norm = F.normalize(video_preds.float(), dim=-1, eps=1e-6) + target_norm = F.normalize(target.float().detach(), dim=-1, eps=1e-6) + cosine_loss = 1.0 - F.cosine_similarity(pred_norm, target_norm, dim=-1, eps=1e-6).mean() + cosine_weight = float(video_cfg.get("cosine_loss_weight", 1.0)) + metrics[f"align/{metric_prefix}_cosine_loss"] = cosine_loss.detach() + metrics[f"align/{metric_prefix}_cosine_loss_weighted"] = (cosine_loss * cosine_weight).detach() + weighted_cosine_loss = cosine_loss * cosine_weight + loss = weighted_cosine_loss if loss is None else loss + weighted_cosine_loss + return loss, metrics + + def _video_cls_loss(self, cls_preds, future_video_cls_targets): + align_params = getattr(getattr(self, "config", None), "align_params", {}) or {} + video_cfg = align_params.get("video", {}) if hasattr(align_params, "get") else {} + cls_loss_type = str(video_cfg.get("cls_loss_type", "cosine")).lower() + cls_weight = float(video_cfg.get("cls_loss_weight", 1.0)) + target = future_video_cls_targets.to(dtype=cls_preds.dtype, device=cls_preds.device) + if target.ndim == 3 and target.shape[1] == 1: + target = target.squeeze(1) + + metrics = {} + loss = None + if cls_loss_type in ("smooth_l1", "smoothl1", "huber"): + smooth_l1_loss = F.smooth_l1_loss(cls_preds.float(), target.float().detach()) + metrics["align/future_video_cls_smooth_l1_loss"] = smooth_l1_loss.detach() + loss = smooth_l1_loss + if cls_loss_type in ("mse", "mse_cosine", "cosine_mse"): + mse_loss = F.mse_loss(cls_preds.float(), target.float().detach()) + metrics["align/future_video_cls_mse_loss"] = mse_loss.detach() + loss = mse_loss + if cls_loss_type in ("cosine", "mse_cosine", "cosine_mse"): + pred_norm = F.normalize(cls_preds.float(), dim=-1, eps=1e-6) + target_norm = F.normalize(target.float().detach(), dim=-1, eps=1e-6) + cosine_loss = 1.0 - F.cosine_similarity(pred_norm, target_norm, dim=-1, eps=1e-6).mean() + metrics["align/future_video_cls_cosine_loss"] = cosine_loss.detach() + loss = cosine_loss if loss is None else loss + cosine_loss + if loss is None: + raise ValueError(f"Unsupported future-video CLS loss type: {cls_loss_type}") + weighted_loss = loss * cls_weight + metrics["align/future_video_cls_loss"] = loss.detach() + metrics["align/future_video_cls_loss_weighted"] = weighted_loss.detach() + return weighted_loss, metrics + + def _emb_loss(self, emb_preds, emb_targets): + l1_loss = F.smooth_l1_loss(emb_preds.float(), emb_targets.float().detach(), reduction="none") + return l1_loss.mean() + +ModelClass = LingbotVlaPolicy + +__all__ = ["LingbotVlaPolicy", "Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLTextModel", "Qwen2ForCausalLM", "Qwen2_5_VLPreTrainedModel"] +# __V1_END__ + +# Qwen3-VL LingBot-VLA v2 policy. +FlowMatchingV1 = FlowMatching +import einops +import torch +from torch import Tensor, nn +import torch.nn.functional as F +from typing import List, Optional, Tuple, Union + +from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedModel +from transformers.models.auto import CONFIG_MAPPING +from transformers.cache_utils import Cache +from transformers.utils import logging + +from telefuser.models.lingbot_vla_v2_qwen import ( + Qwen3VLForConditionalGeneration, + Qwen3VLTextModel, + Qwen3VLPreTrainedModel, + apply_rotary_pos_emb, +) +from telefuser.models.lingbot_vla_v2_loader import ( + block_suffix_to_fv_, + create_sinusoidal_pos_embedding, + make_att_2d_masks, + our_eager_attention_forward, + prefix_query_segments, + prefix_query_token_spans, + sample_beta, +) +from telefuser.models.lingbot_vla_v2_loader import build_block_mask, flex_attention_forward, flex_attention_with_block_mask +from telefuser.models.lingbot_vla_v2_loader import LingBotVLAWeightLoader +from telefuser.models.lingbot_vla_v2_loader import triton_sequence_wise_balance_loss +from telefuser.models.lingbot_vla_v2_moe import ( + Qwen2ForCausalLM, + Qwen2TokenMoeBlock, + Qwen2FusedExperts, + FixQwen2RMSNorm, +) + +try: + from dinov3.hub.backbones import dinov3_vitb16 +except Exception: + dinov3_vitb16 = None + + +logger = logging.get_logger(__name__) + + +class QwenvlWithExpertV2Config(PretrainedConfig): + model_type = "QwenvlWithExpertV2Model" + + def __init__( + self, + freeze_vision_encoder: bool = False, + train_expert_only: bool = False, + vocab_size: int = 0, + use_lm_head: bool = False, + attention_implementation: str = "flex_cached", + tokenizer_path: str | None = None, + enable_expert_vision: bool = False, + expert_vision_type: str | None = None, + use_cache: bool = False, + expert_hidden_size: int = 768, + expert_intermediate_size: int = 2752, + action_num_attention_heads: int = 32, + action_num_key_value_heads: int = 8, + action_head_dim: int = 128, + **kwargs, + ): + self.freeze_vision_encoder = freeze_vision_encoder + self.train_expert_only = train_expert_only + self.attention_implementation = attention_implementation + self.tokenizer_path = tokenizer_path + self.enable_expert_vision = enable_expert_vision + self.expert_vision_type = expert_vision_type + self.vocab_size = vocab_size + self.use_lm_head = use_lm_head + self.action_num_attention_heads = action_num_attention_heads + self.action_num_key_value_heads = action_num_key_value_heads + self.action_head_dim = action_head_dim + num_layers = 36 + + self.qwen_expert_config = CONFIG_MAPPING["qwen2"]( + attention_dropout=0.0, + bos_token_id=151643, + eos_token_id=151645, + hidden_act="silu", + hidden_size=expert_hidden_size, + head_dim=action_head_dim, + initializer_range=0.02, + intermediate_size=expert_intermediate_size, + max_position_embeddings=32768, + max_window_layers=21, + model_type="qwen2", + num_attention_heads=action_num_attention_heads, + num_hidden_layers=num_layers, + num_key_value_heads=action_num_key_value_heads, + rms_norm_eps=1e-06, + rope_theta=1000000.0, + sliding_window=32768, + tie_word_embeddings=True, + torch_dtype="bfloat16", + transformers_version="4.57.3", + use_cache=use_cache, + use_sliding_window=False, + vocab_size=151936, + ) + print( + "=====Action Expert V2 init " + f"{num_layers} Layers, hidden={expert_hidden_size}, " + f"q_heads={action_num_attention_heads}, kv_heads={action_num_key_value_heads}, " + f"head_dim={action_head_dim}.=====" + ) + super().__init__(**kwargs) + + +class QwenvlWithExpertV2Model(PreTrainedModel): + config_class = QwenvlWithExpertV2Config + + def __init__(self, config: QwenvlWithExpertV2Config, eval=False): + super().__init__(config=config) + self.config = config + vlm_config = AutoConfig.from_pretrained(self.config.tokenizer_path, local_files_only=True) + if self.config.vocab_size not in (0, 257152): + vlm_config.text_config.vocab_size = self.config.vocab_size + base_attn_implementation = "flash_attention_2" if is_flash_attn_available() else "eager" + vision_attn_implementation = self.config.vit_attn_implementation + if vision_attn_implementation == "flash_attention_2" and not is_flash_attn_available(): + logger.warning_once("flash-attn is unavailable; using eager attention for Qwen3-VL vision") + vision_attn_implementation = "eager" + vlm_config._attn_implementation = base_attn_implementation + vlm_config.text_config._attn_implementation = base_attn_implementation + vlm_config.vision_config._attn_implementation = vision_attn_implementation + self.qwenvl = Qwen3VLForConditionalGeneration._from_config(vlm_config) + if self.config.use_lm_head: + self.qwenvl.tie_weights() + + self.config.qwen_expert_config._attn_implementation = "flash_attention_2" + self.qwen_expert = Qwen2ForCausalLM._from_config(self.config.qwen_expert_config, eval=eval) + + if getattr(self.config, "adanorm_time", False): + replace_lnorm_with_adanorm( + self.qwen_expert, + self.config.qwen_expert_config.hidden_size, + self.config.qwen_expert_config.hidden_size, + config.final_norm_adanorm, + ) + + self._install_moe_blocks() + self.pos_embeds = None + self.position_embeddings = None + self.cu_seqlens = None + self.visual_split_sizes = None + self.visual_max_seqlen = None + + del self.qwen_expert.model.embed_tokens + if self.config.enable_expert_vision: + if dinov3_vitb16 is None: + raise ImportError("dinov3 is required when enable_expert_vision=True") + if "dinov3_vitb16" in self.config.expert_vision_type: + self.expert_visual = dinov3_vitb16(pretrained=False) + self.expert_visual_mlp = nn.Sequential( + nn.Linear(self.expert_visual.embed_dim, self.expert_visual.embed_dim * 2), + nn.GELU(), + nn.Linear(self.expert_visual.embed_dim * 2, self.config.qwen_expert_config.hidden_size), + ) + + self.attention_interface = self.get_attention_interface() + self.set_requires_grad() + + def _install_moe_blocks(self): + if not getattr(self.config, "use_moe", False): + return + bias_update_speed = getattr(self.config, "bias_update_speed", 0.001) + hidden_size = self.config.qwen_expert_config.hidden_size + token_moe_layers = getattr(self.config, "token_moe_layers", None) or [] + + _moe_impl = getattr(self.config, "_moe_implementation", None) + + if token_moe_layers: + token_config = CONFIG_MAPPING["qwen2_moe"]( + num_experts=getattr(self.config, "token_num_experts", 32), + num_experts_per_tok=getattr(self.config, "token_top_k", 1), + norm_topk_prob=True, + hidden_size=hidden_size, + moe_intermediate_size=getattr(self.config, "token_moe_intermediate_size", 256), + shared_expert_intermediate_size=getattr(self.config, "token_shared_intermediate_size", 256), + output_router_logits=False, + ) + token_config.bias_update_speed = bias_update_speed + token_config._moe_implementation = _moe_impl + token_config.router_activation = getattr(self.config, "router_activation", "softmax") + token_config.routed_scaling_factor = getattr(self.config, "routed_scaling_factor", 1.0) + token_config.use_shared_expert_gate = getattr(self.config, "use_shared_expert_gate", True) + for idx in token_moe_layers: + self.qwen_expert.model.layers[idx].mlp = Qwen2TokenMoeBlock(token_config) + + def set_requires_grad(self): + if self.config.freeze_vision_encoder: + self.qwenvl.visual.eval() + for params in self.qwenvl.visual.parameters(): + params.requires_grad = False + if self.config.train_expert_only: + self.qwenvl.eval() + for params in self.qwenvl.parameters(): + params.requires_grad = False + + def train(self, mode: bool = True): + super().train(mode) + if self.config.freeze_vision_encoder: + self.qwenvl.visual.eval() + if self.config.train_expert_only: + self.qwenvl.eval() + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor, + ): + precompute_grid_thw = getattr(self.config, "precompute_grid_thw", False) + if precompute_grid_thw and self.position_embeddings is None: + ( + self.pos_embeds, + self.position_embeddings, + self.cu_seqlens, + self.visual_split_sizes, + self.visual_max_seqlen, + ) = self.qwenvl.visual.preprcess_grid_thw(grid_thw=image_grid_thw) + image_embeds, deepstack_image_embeds = self.qwenvl.visual( + pixel_values, + grid_thw=image_grid_thw, + pos_embeds=self.pos_embeds, + position_embeddings=self.position_embeddings, + cu_seqlens=self.cu_seqlens, + max_seqlen=self.visual_max_seqlen, + ) + split_sizes = self.visual_split_sizes + if split_sizes is None: + split_sizes = (image_grid_thw.prod(-1) // self.qwenvl.visual.spatial_merge_size**2).tolist() + image_chunks = list(torch.split(image_embeds, split_sizes)) + deepstack_chunks = [ + list(torch.split(deepstack_embeds, split_sizes)) + for deepstack_embeds in deepstack_image_embeds + ] + image_embeds = torch.stack(image_chunks, dim=0) + deepstack_image_embeds = [ + torch.stack(chunks, dim=0) + for chunks in deepstack_chunks + ] + return image_embeds, deepstack_image_embeds + + def embed_image(self, image: torch.Tensor, image_grid_thw: torch.LongTensor): + return self.get_image_features( + image, + image_grid_thw=image_grid_thw, + ) + + def embed_language_tokens(self, tokens: torch.Tensor): + return self.qwenvl.model.language_model.embed_tokens(tokens) + + def embed_special_token(self, token_id: int, batch: int, count: int, device, dtype): + token = torch.tensor([token_id], device=device, dtype=torch.long) + emb = self.embed_language_tokens(token).to(dtype=dtype) + return emb.view(1, 1, 1, -1).expand(batch, count, 1, -1) + + def build_prefix_position_ids(self, input_ids, attention_mask, + image_grid_thw=None, video_grid_thw=None): + position_ids, _ = self.qwenvl.model.get_rope_index( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + ) + return position_ids + + def apply_mrope(self, query_states, key_states, position_ids): + position_embeddings = self.qwenvl.model.language_model.rotary_emb(query_states, position_ids) + return apply_rotary_pos_emb(query_states, key_states, *position_embeddings, unsqueeze_dim=2) + + def handle_kv_cache( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + use_cache: Optional[bool] = None, + fill_kv_cache: Optional[bool] = None, + ): + if use_cache: + if past_key_values is None: + past_key_values = {} + if fill_kv_cache: + past_key_values[layer_idx] = {"key_states": key_states, "value_states": value_states} + else: + key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1) + value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1) + return key_states, value_states, past_key_values + + def _apply_deepstack(self, hidden_states, layer_idx, visual_pos_masks, deepstack_visual_embeds): + if ( + deepstack_visual_embeds is not None + and visual_pos_masks is not None + and layer_idx < len(deepstack_visual_embeds) + ): + hidden_states = self.qwenvl.model.language_model._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + return hidden_states + + def forward( + self, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + vlm_position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + inputs_embeds: List[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + fill_kv_cache: Optional[bool] = None, + ada_cond: List[torch.FloatTensor] = None, + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, + ): + models = [self.qwenvl.model.language_model, self.qwen_expert.model] + num_layers = self.qwenvl.config.text_config.num_hidden_layers + action_num_layers = self.config.qwen_expert_config.num_hidden_layers + router_logits_list = [] + + assert action_num_layers == num_layers, ( + "Action expert and VLM must have the same number of layers " + f"(got action={action_num_layers}, vlm={num_layers})." + ) + + for layer_idx in range(num_layers): + query_states = [] + key_states = [] + value_states = [] + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is None: + continue + if i == 1: + q, k, v = models[i].layers[layer_idx]( + hidden_states, compute_kqv=True, ada_cond=ada_cond + ) + else: + q, k, v = models[i].layers[layer_idx](hidden_states, compute_kqv=True) + query_states.append(q.float()) + key_states.append(k.float()) + value_states.append(v.float()) + + query_states = torch.cat(query_states, dim=1) + key_states = torch.cat(key_states, dim=1) + value_states = torch.cat(value_states, dim=1) + query_states, key_states = self.apply_mrope(query_states, key_states, position_ids) + key_states, value_states, past_key_values = self.handle_kv_cache( + key_states, + value_states, + layer_idx, + past_key_values=past_key_values, + use_cache=use_cache, + fill_kv_cache=fill_kv_cache, + ) + if self.config.attention_implementation == "flex_cached": + if layer_idx == 0: + _full_len = query_states.shape[1] + _full_block_mask = build_block_mask( + attention_mask, + self.qwenvl.config.text_config.num_attention_heads, + _full_len, + _full_len, + ) + att_output = flex_attention_with_block_mask( + query_states, key_states, value_states, _full_block_mask, query_states.shape[1] + ) + else: + att_output = self.attention_interface(query_states, key_states, value_states, attention_mask) + + outputs_embeds = [] + start = 0 + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is None: + outputs_embeds.append(None) + continue + end = start + hidden_states.shape[1] + if i == 1: + out_emb, router_logits = models[i].layers[layer_idx]( + hidden_states, + att_output, + start, + end, + output_atten=True, + ada_cond=ada_cond, + ) + if router_logits is not None: + router_logits_list.append(router_logits) + else: + out_emb = models[i].layers[layer_idx]( + hidden_states, att_output, start, end, output_atten=True + ) + out_emb = self._apply_deepstack(out_emb, layer_idx, visual_pos_masks, deepstack_visual_embeds) + outputs_embeds.append(out_emb) + start = end + inputs_embeds = outputs_embeds + + outputs_embeds = [] + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is None: + outputs_embeds.append(None) + elif self.config.final_norm_adanorm and i == 1: + out_emb, _ = models[i].norm(hidden_states, ada_cond) + outputs_embeds.append(out_emb) + else: + outputs_embeds.append(models[i].norm(hidden_states)) + return outputs_embeds, past_key_values, router_logits_list + + def get_attention_interface(self): + if self.config.attention_implementation == "flex": + print("=====Using Flex Attn=====") + return flex_attention_forward + if self.config.attention_implementation == "flex_cached": + print("=====Using Flex Cached (prebuilt BlockMask) Attn=====") + return flex_attention_forward + if self.config.attention_implementation == "eager": + print("=====Using Eager Attn=====") + return our_eager_attention_forward + raise ValueError(f"Invalid attention implementation: {self.config.attention_implementation}") + + +class FlowMatchingV2(FlowMatchingV1): + def __init__(self, config, eval): + nn.Module.__init__(self) + self.config = config + qwenvl_with_export_config = QwenvlWithExpertV2Config( + freeze_vision_encoder=self.config.freeze_vision_encoder, + train_expert_only=self.config.train_expert_only, + vocab_size=getattr(self.config, "vocab_size", 0), + use_lm_head=getattr(self.config, "use_lm_head", False), + attention_implementation=self.config.attention_implementation, + tokenizer_path=self.config.tokenizer_path, + enable_expert_vision=self.config.enable_expert_vision, + expert_vision_type=self.config.expert_vision_type, + use_cache=getattr(self.config, "use_cache", True), + expert_hidden_size=getattr(self.config, "expert_hidden_size", 768), + expert_intermediate_size=getattr(self.config, "expert_intermediate_size", 2752), + action_num_attention_heads=getattr(self.config, "action_num_attention_heads", 32), + action_num_key_value_heads=getattr(self.config, "action_num_key_value_heads", 8), + action_head_dim=getattr(self.config, "action_head_dim", 128), + ) + for name in [ + "adanorm_time", + "final_norm_adanorm", + "precompute_grid_thw", + "vit_attn_implementation", + "use_moe", + "bias_update_speed", + "token_moe_layers", + "token_num_experts", + "token_top_k", + "token_moe_intermediate_size", + "token_shared_intermediate_size", + "router_activation", + "routed_scaling_factor", + "use_shared_expert_gate", + "_moe_implementation", + ]: + if hasattr(config, name): + setattr(qwenvl_with_export_config, name, getattr(config, name)) + self.qwenvl_with_expert = QwenvlWithExpertV2Model(qwenvl_with_export_config, eval) + self.config.proj_width = qwenvl_with_export_config.qwen_expert_config.hidden_size + self.config.initializer_range = getattr(qwenvl_with_export_config.qwen_expert_config, "initializer_range", None) + + self.state_proj = nn.Linear(self.config.max_state_dim, self.config.proj_width) + self.action_in_proj = nn.Linear(self.config.max_action_dim, self.config.proj_width) + self.action_out_proj = nn.Linear(self.config.proj_width, self.config.max_action_dim) + self.action_time_mlp_in = nn.Linear(self.config.proj_width * 2, self.config.proj_width) + self.action_time_mlp_out = nn.Linear(self.config.proj_width, self.config.proj_width) + + self.config.align_params = getattr(self.config, "align_params", None) or {} + if self.config.align_params != {}: + self.steps = 0 + self.use_depth_align = True + self.init_depth_heads(self.config.align_params) + self.use_future_video = self.config.align_params.get("use_future_video", False) + if self.use_future_video: + self.init_video_heads(self.config.align_params) + else: + self.use_depth_align = False + self.use_future_video = False + self.use_future_video_patch = False + self.use_current_video_patch = False + self.use_current_shared_task_proj = False + self.use_future_video_cls = False + self.use_shared_future_task_proj = False + self.future_video_share_future_depth_query = False + self.block_future_depth_to_action = False + + self.set_requires_grad() + + def embed_prefix( + self, + images, + img_masks, + lang_tokens, + lang_masks, + image_grid_thw=None, + ): + if image_grid_thw is None: + raise ValueError("LingbotVlaV2Policy requires image_grid_thw from the Qwen3-VL image processor.") + bsize = images.shape[0] + device = images.device + dtype = images.dtype + if images.ndim == 3: + bsize = 1 + num_images = images.shape[0] + else: + num_images = images.shape[1] if images.ndim >= 4 else 1 + if images.ndim == 4: + images = einops.rearrange(images, "b n l d -> (b n) l d") + elif images.ndim == 5: + images = einops.rearrange(images, "b n c h w -> (b n) c h w") + if image_grid_thw.ndim == 3: + flat_grid_thw = einops.rearrange(image_grid_thw, "b n d -> (b n) d") + else: + flat_grid_thw = image_grid_thw + + img_emb, deepstack_embs = self.qwenvl_with_expert.embed_image( + images, + flat_grid_thw, + ) + embed_dtype = img_emb.dtype + num_patch = img_emb.shape[1] + img_emb = einops.rearrange(img_emb, "(b n) l d -> b n l d", b=bsize, n=num_images) + deepstack_embs = [ + einops.rearrange(x, "(b n) l d -> b n l d", b=bsize, n=num_images) + for x in deepstack_embs + ] + if img_masks.ndim == 1: + img_masks = img_masks.unsqueeze(0) + + cfg = self.qwenvl_with_expert.qwenvl.config + visual_token_id = cfg.image_token_id + + if getattr(self.config, "qwen3vl_use_vision_boundaries", True): + start_emb = self.qwenvl_with_expert.embed_special_token( + cfg.vision_start_token_id, bsize, num_images, device, embed_dtype + ) + end_emb = self.qwenvl_with_expert.embed_special_token( + cfg.vision_end_token_id, bsize, num_images, device, embed_dtype + ) + img_chunks = torch.cat([start_emb, img_emb, end_emb], dim=2) + image_token_len = num_patch + 2 + image_pad_masks = einops.repeat(img_masks, "b n -> b n l", l=image_token_len) + image_visual_masks = torch.zeros_like(image_pad_masks) + image_visual_masks[:, :, 1 : 1 + num_patch] = einops.repeat(img_masks, "b n -> b n l", l=num_patch) + fake_image_ids = torch.full( + (bsize, num_images, image_token_len), + visual_token_id, + dtype=torch.long, + device=device, + ) + fake_image_ids[:, :, 0] = cfg.vision_start_token_id + fake_image_ids[:, :, -1] = cfg.vision_end_token_id + else: + img_chunks = img_emb + image_token_len = num_patch + image_pad_masks = einops.repeat(img_masks, "b n -> b n l", l=image_token_len) + image_visual_masks = image_pad_masks + fake_image_ids = torch.full( + (bsize, num_images, image_token_len), + visual_token_id, + dtype=torch.long, + device=device, + ) + + img_emb = einops.rearrange(img_chunks, "b n l d -> b (n l) d") + image_pad_masks = einops.rearrange(image_pad_masks, "b n l -> b (n l)") + visual_pos_masks = einops.rearrange(image_visual_masks, "b n l -> b (n l)") + fake_image_ids = einops.rearrange(fake_image_ids, "b n l -> b (n l)") + + lang_emb = self.qwenvl_with_expert.embed_language_tokens(lang_tokens).to(dtype=embed_dtype) + + if self.use_depth_align and self.align_type == "query": + def _get_align_tokens(tokens): + tk_weights = tokens.view(self.num_task_tokens, tokens.shape[0] // self.num_task_tokens, tokens.shape[1]) + tk_weights = tk_weights.mean(dim=1) + return tk_weights + + align_pad_masks = torch.ones( + bsize, + self.num_task_tokens, + device=device, + dtype=lang_masks.dtype + ) + fake_align_ids = torch.full( + (bsize, self.num_task_tokens), + cfg.text_config.eos_token_id, + dtype=torch.long, + device=device + ) + + current_task = _get_align_tokens(self.depth_align_embs) + if ( + getattr(self, "use_future_video", False) + and getattr(self, "use_current_video_patch", False) + and getattr(self, "use_current_shared_task_proj", False) + ): + current_video_task = _get_align_tokens(self.current_video_align_embs) + current_task = self.current_shared_task_proj( + torch.cat([current_task, current_video_task], dim=-1) + ) + align_embs = current_task.repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) + parts = [img_emb] + masks = [image_pad_masks] + input_ids = [fake_image_ids] + visual_masks = [visual_pos_masks] + + def _append( + tokens, + token_masks, + token_ids, + token_visual_masks=None, + ): + parts.append(tokens) + masks.append(token_masks) + input_ids.append(token_ids) + if token_visual_masks is None: + token_visual_masks = torch.zeros_like(token_masks) + visual_masks.append(token_visual_masks) + + future_align_embs = None + if self.use_future_depth: + future_task = _get_align_tokens(self.future_depth_align_embs) + if ( + getattr(self, "use_future_video", False) + and getattr(self, "use_future_video_patch", True) + and getattr(self, "future_video_share_future_depth_query", False) + and getattr(self, "use_shared_future_task_proj", False) + ): + future_video_task = _get_align_tokens(self.future_video_align_embs) + future_task = self.future_shared_task_proj( + torch.cat([future_task, future_video_task], dim=-1) + ) + future_align_embs = future_task.repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) + + if ( + not self.use_future_depth + and getattr(self, "use_future_video", False) + and getattr(self, "future_video_share_future_depth_query", False) + ): + raise ValueError( + "share_future_depth_query=True requires depth.use_future_depth=True." + ) + + for segment_name in prefix_query_segments( + use_depth_align=True, + use_future_depth=self.use_future_depth, + use_future_video=getattr(self, "use_future_video", False), + use_future_video_cls=getattr(self, "use_future_video_cls", False), + use_future_video_patch=getattr(self, "use_future_video_patch", True), + future_video_share_future_depth_query=getattr( + self, + "future_video_share_future_depth_query", + False, + ), + ): + if segment_name == "language": + _append( + lang_emb, + lang_masks, + lang_tokens.to(device), + ) + elif segment_name == "current_depth": + _append(align_embs, align_pad_masks, fake_align_ids) + elif segment_name == "future_video_cls": + future_video_cls_align_emb = self.future_video_cls_align_emb.weight.repeat( + img_emb.size(0), 1, 1 + ).to(img_emb.device, img_emb.dtype) + cls_align_pad_masks = torch.ones( + bsize, + 1, + device=device, + dtype=lang_masks.dtype, + ) + fake_cls_align_ids = torch.full( + (bsize, 1), + cfg.text_config.eos_token_id, + dtype=torch.long, + device=device, + ) + _append(future_video_cls_align_emb, cls_align_pad_masks, fake_cls_align_ids) + elif segment_name == "future_video": + future_video_align_embs = _get_align_tokens(self.future_video_align_embs).repeat( + img_emb.size(0), 1, 1 + ).to(img_emb.device, img_emb.dtype) + _append(future_video_align_embs, align_pad_masks, fake_align_ids) + elif segment_name == "future_depth": + _append(future_align_embs, align_pad_masks, fake_align_ids) + else: + raise ValueError(f"Unsupported prefix query segment: {segment_name}") + + embs = torch.cat(parts, dim=1) + pad_masks = torch.cat(masks, dim=1) + prefix_input_ids = torch.cat(input_ids, dim=1) + full_visual_pos_masks = torch.cat(visual_masks, dim=1) + else: + embs = torch.cat([img_emb, lang_emb], dim=1) + pad_masks = torch.cat([image_pad_masks, lang_masks], dim=1) + prefix_input_ids = torch.cat([fake_image_ids, lang_tokens.to(device)], dim=1) + full_visual_pos_masks = torch.cat([visual_pos_masks, torch.zeros_like(lang_masks)], dim=1) + + if getattr(self.config, "vlm_causal", False): + att_masks = torch.ones((bsize, embs.shape[1]), device=device, dtype=torch.bool) + else: + att_masks = torch.zeros((bsize, embs.shape[1]), device=device, dtype=torch.bool) + + flat_img_masks = einops.rearrange(img_masks, "b n -> (b n)") + rope_grid_thw = flat_grid_thw[flat_img_masks] + if rope_grid_thw.numel() == 0: + rope_grid_thw = flat_grid_thw[:1] + prefix_position_ids = self.qwenvl_with_expert.build_prefix_position_ids( + prefix_input_ids, + pad_masks.long(), + image_grid_thw=rope_grid_thw, + video_grid_thw=None, + ) + filtered_deepstack = [] + img_visual_only = einops.repeat(img_masks, "b n -> b n l", l=num_patch) + for deepstack in deepstack_embs: + filtered_deepstack.append(deepstack[img_visual_only]) + + result = ( + embs, + pad_masks, + att_masks, + prefix_position_ids, + full_visual_pos_masks, + filtered_deepstack, + ) + return result + + def _build_full_position_ids(self, prefix_position_ids, prefix_pad_masks, suffix_pad_masks): + valid_prefix_pos = prefix_position_ids.masked_fill(~prefix_pad_masks.unsqueeze(0), 0) + prefix_offsets = valid_prefix_pos.amax(dim=(0, 2)) + 1 + suffix_1d = prefix_offsets[:, None] + torch.cumsum(suffix_pad_masks.long(), dim=1) - 1 + suffix_1d = suffix_1d.masked_fill(~suffix_pad_masks, 1) + suffix_position_ids = suffix_1d.unsqueeze(0).expand(3, -1, -1) + return torch.cat([prefix_position_ids, suffix_position_ids], dim=-1) + + def _current_depth_task_tokens(self, hidden_states, num_images=3): + query_spans = prefix_query_token_spans( + prefix_len=hidden_states.shape[1], + num_task_tokens=self.num_task_tokens, + use_depth_align=True, + use_future_depth=getattr(self, "use_future_depth", False), + use_future_video=getattr(self, "use_future_video", False), + use_future_video_cls=getattr(self, "use_future_video_cls", False), + use_future_video_patch=getattr(self, "use_future_video_patch", True), + future_video_share_future_depth_query=getattr( + self, + "future_video_share_future_depth_query", + False, + ), + ) + start, end = query_spans["current_depth"] + return hidden_states[:, start:end, :] + + def forward( + self, + images, + img_masks, + lang_tokens, + lang_masks, + state, + actions, + noise=None, + time=None, + loss_type="fm", + depth_targets=None, + image_grid_thw=None, + future_depth_targets=None, + future_video_targets=None, + future_video_cls_targets=None, + future_video_current_patch=None, + ) -> Tensor: + dtype = state.dtype + device = state.device + if noise is None: + noise = torch.randn(actions.shape, device=device, dtype=dtype) + if time is None: + time = self.sample_time(actions.size(0), device).to(dtype) + + time_expanded = time[:, None, None] + x_t = time_expanded * noise + (1 - time_expanded) * actions + u_t = noise - actions + + ( + prefix_embs, + prefix_pad_masks, + prefix_att_masks, + prefix_position_ids, + visual_pos_masks, + deepstack_visual_embeds, + ) = self.embed_prefix( + images, + img_masks, + lang_tokens, + lang_masks, + image_grid_thw=image_grid_thw, + ) + time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( + state, x_t, time + ) + + pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1) + att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1) + att_2d_masks = make_att_2d_masks(pad_masks, att_masks) + prefix_len = prefix_pad_masks.shape[1] + if self.block_future_depth_to_action: + att_2d_masks = block_suffix_to_fv_( + att_2d_masks, + suffix_row_start=prefix_len, + prefix_len=prefix_len, + num_task_tokens=self.num_task_tokens, + ) + + att_2d_masks = self._block_suffix_to_future_video_if_enabled_( + att_2d_masks, + suffix_row_start=prefix_len, + prefix_len=prefix_len, + ) + position_ids = self._build_full_position_ids(prefix_position_ids, prefix_pad_masks, suffix_pad_masks) + + (outputs_embeds, suffix_out), _, router_logits_list = self.qwenvl_with_expert.forward( + attention_mask=att_2d_masks, + position_ids=position_ids, + vlm_position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, suffix_embs], + use_cache=self.config.use_cache, + fill_kv_cache=True, + ada_cond=time_embs if getattr(self.config, "adanorm_time", False) else None, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + align_metrics = {} + if self.config.align_params != {}: + loss_depth, loss_future_depth, depth_preds, future_depth_preds = self.depth_emb_forward(outputs_embeds, depth_targets, img_masks,future_depth_targets,) + loss_depth = loss_depth * self.config.align_params["depth_loss_weight"] + loss_future_depth = loss_future_depth * self.config.align_params.get("future_depth_loss_weight", 1.0) + loss_future_video = 0 + future_video_preds = None + current_video_preds = None + if getattr(self, "use_future_video", False): + loss_video, future_video_preds, video_metrics = self.video_emb_forward( + outputs_embeds, + future_video_targets, + future_video_cls_targets=future_video_cls_targets, + future_video_current_patch=future_video_current_patch, + ) + video_total_loss = loss_video + if ( + getattr(self, "use_current_video_patch", False) + and future_video_current_patch is not None + ): + current_video_loss, current_video_preds, current_video_metrics = self.current_video_emb_forward( + outputs_embeds, + future_video_current_patch, + ) + video_total_loss = video_total_loss + current_video_loss + video_metrics.update(current_video_metrics) + video_metrics["align/current_video_loss"] = current_video_loss.detach() + video_cfg = self.config.align_params.get("video", {}) + video_weight = video_cfg.get( + "future_video_loss_weight", + self.config.align_params.get( + "future_video_loss_weight", + self.config.align_params["depth_loss_weight"], + ), + ) + loss_future_video = video_total_loss * video_weight + align_metrics.update(video_metrics) + if "align/current_video_loss" in align_metrics: + align_metrics["align/current_video_loss_weighted"] = ( + align_metrics["align/current_video_loss"] * video_weight + ) + align_metrics["align/future_video_loss"] = loss_video.detach() + align_metrics["align/future_video_loss_weighted"] = (loss_video * video_weight).detach() + align_metrics["align/video_loss"] = video_total_loss.detach() + align_metrics["align/video_loss_weighted"] = loss_future_video.detach() + self.steps += 1 + else: + loss_depth = 0 + loss_future_depth = 0 + loss_future_video = 0 + depth_preds = None + future_depth_preds = None + future_video_preds = None + current_video_preds = None + + suffix_out = suffix_out[:, -self.config.n_action_steps :] + if getattr(self.config, "action_fp32", False): + v_t = self._fp32_linear(self.action_out_proj, suffix_out) + else: + if suffix_out.dtype != self.action_out_proj.weight.dtype: + suffix_out = suffix_out.to(self.action_out_proj.weight.dtype) + v_t = self.action_out_proj(suffix_out) + + if loss_type == "fm": + losses = F.mse_loss(u_t, v_t, reduction="none") + elif loss_type == "L1_fm": + losses = F.l1_loss(u_t, v_t, reduction="none") + + seq_wise_loss, router_z_loss, moe_metrics = self._moe_losses_and_metrics( + router_logits_list, losses + ) + if align_metrics: + moe_metrics.update(align_metrics) + return losses, loss_depth, loss_future_depth, loss_future_video, depth_preds, seq_wise_loss, router_z_loss, moe_metrics, future_depth_preds, future_video_preds, current_video_preds + + def sample_actions( + self, + images, + img_masks, + lang_tokens, + lang_masks, + state, + noise=None, + image_grid_thw=None, + ) -> Tensor: + """Do a full Qwen3-VL inference forward and compute the action.""" + bsize = state.shape[0] + device = state.device + dtype = state.dtype + + if noise is None: + actions_shape = ( + bsize, + self.config.n_action_steps, + self.config.max_action_dim, + ) + noise = torch.randn(actions_shape, device=device, dtype=dtype) + + ( + prefix_embs, + prefix_pad_masks, + prefix_att_masks, + prefix_position_ids, + visual_pos_masks, + deepstack_visual_embeds, + ) = self.embed_prefix( + images, + img_masks, + lang_tokens, + lang_masks, + image_grid_thw=image_grid_thw, + ) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + + _, past_key_values, _ = self.qwenvl_with_expert.forward( + attention_mask=prefix_att_2d_masks, + position_ids=prefix_position_ids, + vlm_position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=self.config.use_cache, + fill_kv_cache=True, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + + dt = torch.tensor(-1.0 / self.config.num_steps, dtype=dtype, device=device) + x_t = noise + time = torch.tensor(1.0, dtype=dtype, device=device) + count = 0 + predict_velocity_fn = self.predict_velocity + if getattr(self, "_use_compile_predict_velocity", False): + predict_velocity_fn = getattr(self, "_compiled_predict_velocity", None) + if predict_velocity_fn is None: + predict_velocity_fn = torch.compile( + self.predict_velocity, + fullgraph=False, + dynamic=False, + options={"triton.cudagraphs": False}, + ) + self._compiled_predict_velocity = predict_velocity_fn + + while time >= -dt / 2: + count += 1 + expanded_time = time.expand(bsize) + v_t = predict_velocity_fn( + state, + prefix_pad_masks, + past_key_values, + x_t, + expanded_time, + prefix_position_ids=prefix_position_ids, + ) + + x_t += dt * v_t + time += dt + print(f"Denoise {count} steps") + return x_t + + def predict_velocity( + self, + state, + prefix_pad_masks, + past_key_values, + x_t, + timestep, + prefix_position_ids=None, + ): + """Predict velocity at time t using cached Qwen3-VL prefix states.""" + if prefix_position_ids is None: + raise ValueError("FlowMatchingV2.predict_velocity requires Qwen3-VL prefix_position_ids.") + + time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( + state, + x_t, + timestep, + ) + + suffix_len = suffix_pad_masks.shape[1] + batch_size = prefix_pad_masks.shape[0] + prefix_len = prefix_pad_masks.shape[1] + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand( + batch_size, + suffix_len, + prefix_len, + ) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + if self.block_future_depth_to_action: + # Query rows here are all suffix (state/action), so row start is 0. + full_att_2d_masks = block_suffix_to_fv_( + full_att_2d_masks, + suffix_row_start=0, + prefix_len=prefix_len, + num_task_tokens=self.num_task_tokens, + ) + full_att_2d_masks = self._block_suffix_to_future_video_if_enabled_( + full_att_2d_masks, + suffix_row_start=0, + prefix_len=prefix_len, + ) + + full_position_ids = self._build_full_position_ids( + prefix_position_ids, + prefix_pad_masks, + suffix_pad_masks, + ) + position_ids = full_position_ids[:, :, -suffix_len:] + + outputs_embeds, _, _ = self.qwenvl_with_expert.forward( + attention_mask=full_att_2d_masks, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=self.config.use_cache, + fill_kv_cache=False, + ada_cond=time_embs if getattr(self.config, "adanorm_time", False) else None, + ) + suffix_out = outputs_embeds[1] + suffix_out = suffix_out[:, -self.config.n_action_steps :] + if getattr(self.config, "action_fp32", False): + v_t = self._fp32_linear(self.action_out_proj, suffix_out) + else: + if suffix_out.dtype != self.action_out_proj.weight.dtype: + suffix_out = suffix_out.to(self.action_out_proj.weight.dtype) + v_t = self.action_out_proj(suffix_out) + return v_t + + def _moe_losses_and_metrics(self, router_logits_list, losses): + router_z_loss_coeff = getattr(self.config, "router_z_loss_coeff", 0) + router_z_loss = losses.new_zeros(()) + router_z_layer_losses = None # per-layer raw z-loss (pre-coeff), for monitoring + if router_z_loss_coeff > 0 and router_logits_list: + router_z_layer_losses = [ + torch.logsumexp(logits.float(), dim=-1).pow(2).mean() + for logits in router_logits_list + ] + router_z_loss = router_z_loss_coeff * torch.stack(router_z_layer_losses).mean() + + seq_wise_loss_coeff = getattr(self.config, "sequence_wise_loss_coeff", 0) + seq_wise_loss = 0 + seqwise_layer_losses = None # per-layer raw seq-wise balance loss (pre-coeff), for monitoring + if seq_wise_loss_coeff > 0 and router_logits_list: + # router_logits are [B*T, E] (action-expert tokens, fixed length T per sample). + # per_sequence -> balance experts within each sample's T tokens (DeepSeek-V3 intent); + # global -> treat the whole B*T batch as one sequence. + mode = getattr(self.config, "sequence_wise_mode", "per_sequence") + score_func = getattr(self.config, "router_activation", "softmax") + if mode == "global": + seq_lengths = None + else: + B = losses.shape[0] + N = router_logits_list[0].shape[0] + seq_lengths = [N // B] * B + seqwise_layer_losses = triton_sequence_wise_balance_loss( + router_logits_list=tuple(router_logits_list), + top_k=getattr(self.config, "token_top_k", 4), + seq_lengths=seq_lengths, + padding_len=0, + score_func=score_func, + ) + if seqwise_layer_losses: + seq_wise_loss = seq_wise_loss_coeff * torch.stack(seqwise_layer_losses).mean() + + moe_metrics = {} + if router_logits_list: + token_moe_layers_list = sorted(getattr(self.config, "token_moe_layers", None) or []) + all_moe_indices = token_moe_layers_list + token_expert_counts = [] + # Per-layer token-MoE stats, collected for moe_summary/* cross-layer aggregates. + tok_maxvio, tok_minvio, tok_minload, tok_entropy, tok_sigmoid = [], [], [], [], [] + tok_bias = [] # per-layer max(|e_score_correction_bias|) (loss-free); >1 -> bias dominates sigmoid score + any_dead = None # OR-accumulated bool: any token-MoE layer with a 0-count expert + with torch.no_grad(): + for i, logits in enumerate(router_logits_list): + layer_id = all_moe_indices[i] if i < len(all_moe_indices) else i + num_experts = logits.shape[-1] + routing_probs = F.softmax(logits, dim=1, dtype=torch.float) + moe_block = self.qwenvl_with_expert.qwen_expert.model.layers[layer_id].mlp + if hasattr(moe_block, "last_tokens_per_expert"): + # Global (all-reduced), biased, true top-k load from the load-balance hook. + counts = moe_block.last_tokens_per_expert.clone() + if counts.sum() == 0: + # Buffer not yet populated by the load-balance hook (first step + # after run start / resume) -> skip this layer to avoid a spurious + # has_dead_expert / min_load_ratio spike on the very first viz. + continue + else: + _, selected = torch.topk(routing_probs, 1, dim=-1) + counts = F.one_hot(selected.squeeze(-1), num_classes=num_experts).float().sum(dim=0) + avg_load = counts.mean() + denom = avg_load.clamp(min=1e-9) + maxvio = (counts.max() - avg_load) / denom # peak overload (>=0, larger=worse) + minvio = (avg_load - counts.min()) / denom # valley underload (=1 -> dead expert) + min_load_ratio = counts.min() / denom # =0 -> dead expert + # entropy is rank-local (this rank's routing_probs, last micro-batch). + per_sample_entropy = -(routing_probs * routing_probs.clamp(min=1e-9).log()).sum(dim=-1) + entropy = per_sample_entropy.mean() + ll = f"{layer_id:02d}" + token_expert_counts.append((layer_id, counts)) + moe_metrics[f"moe_maxvio/layer{ll}"] = maxvio + moe_metrics[f"moe_minvio/layer{ll}"] = minvio + moe_metrics[f"moe_minload/layer{ll}"] = min_load_ratio + moe_metrics[f"moe_entropy_rank0/layer{ll}"] = entropy + tok_maxvio.append(maxvio) + tok_minvio.append(minvio) + tok_minload.append(min_load_ratio) + tok_entropy.append(entropy) + dead = counts.min() == 0 + any_dead = dead if any_dead is None else (any_dead | dead) + if hasattr(moe_block, "avg_topk_sigmoid_score"): + sig = moe_block.avg_topk_sigmoid_score.detach().reshape(()).to(denom) + moe_metrics[f"moe_topksigmoid_rank0/layer{ll}"] = sig + tok_sigmoid.append(sig) + if hasattr(moe_block, "e_score_correction_bias"): + bias_absmax = moe_block.e_score_correction_bias.detach().abs().max().to(denom) + moe_metrics[f"moe_bias/layer{ll}"] = bias_absmax + tok_bias.append(bias_absmax) + # ---- moe_summary/* : cross-layer aggregates over token-MoE layers (written every step) ---- + if tok_maxvio: + moe_metrics["moe_summary/maxvio_avg"] = torch.stack(tok_maxvio).mean() + moe_metrics["moe_summary/maxvio_max"] = torch.stack(tok_maxvio).max() + moe_metrics["moe_summary/minvio_avg"] = torch.stack(tok_minvio).mean() + moe_metrics["moe_summary/minvio_max"] = torch.stack(tok_minvio).max() + moe_metrics["moe_summary/min_load_ratio"] = torch.stack(tok_minload).min() + moe_metrics["moe_summary/has_dead_expert"] = any_dead.float() + moe_metrics["moe_summary/entropy_avg_rank0"] = torch.stack(tok_entropy).mean() + if tok_sigmoid: + moe_metrics["moe_summary/topk_sigmoid_avg_rank0"] = torch.stack(tok_sigmoid).mean() + if tok_bias: + moe_metrics["moe_summary/bias_absmax"] = torch.stack(tok_bias).max() + # ---- moe_seqwise/* : per-layer raw sequence-wise balance loss (pre-coeff) + average ---- + if seqwise_layer_losses and len(seqwise_layer_losses) == len(all_moe_indices): + sw_vals = [] + for lid, sw in zip(all_moe_indices, seqwise_layer_losses): + v = sw.detach() + moe_metrics[f"moe_seqwise/layer{lid:02d}"] = v + sw_vals.append(v) + moe_metrics["moe_seqwise/avg"] = torch.stack(sw_vals).mean() + # ---- moe_zloss/* : per-layer raw router z-loss (pre-coeff) + average/weighted loss ---- + if router_z_layer_losses and len(router_z_layer_losses) == len(all_moe_indices): + zl_vals = [] + for lid, zl in zip(all_moe_indices, router_z_layer_losses): + v = zl.detach() + moe_metrics[f"moe_zloss/layer{lid:02d}"] = v + zl_vals.append(v) + moe_metrics["moe_zloss/avg_raw"] = torch.stack(zl_vals).mean() + moe_metrics["moe_zloss/weighted"] = router_z_loss.detach() + if token_expert_counts: + moe_metrics["_token_moe_expert_counts"] = token_expert_counts + return seq_wise_loss, router_z_loss, moe_metrics + + +class LingbotVlaV2Policy(PreTrainedModel): + config_class = LingbotVLAV2Config + name = "torch_lingbot_vla_v2" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2DecoderLayer", "FixQwen2RMSNorm", "FixAdaRMSNorm"] + + def get_parallel_plan(self): + from telefuser.models.lingbot_vla_v2_loader import NativeParallelPlan as ParallelPlan + from torch.distributed._tensor import Shard + + ep_plan = { + "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.gate_proj": Shard(0), + "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.up_proj": Shard(0), + "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.down_proj": Shard(0), + } + return ParallelPlan(ep_plan=ep_plan) + + @classmethod + def get_weight_loader(cls): + return LingBotVLAWeightLoader() + + def __init__(self, config: LingbotVLAV2Config, eval: bool = False): + super().__init__(config) + self.config = config + self.language_tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, local_files_only=True) + self.model = FlowMatchingV2(config, eval) + if not getattr(self.config, "use_lm_head", False): + del self.model.qwenvl_with_expert.qwenvl.lm_head + del self.model.qwenvl_with_expert.qwen_expert.lm_head + self.reset() + torch.set_float32_matmul_precision("high") + + def reset(self): + return None + + def get_optim_params(self) -> dict: + return self.parameters() + + def forward( + self, + images, + img_masks, + state, + lang_tokens, + lang_masks, + actions, + joint_mask=None, + action_is_pad=None, + noise=None, + time=None, + depth_targets=None, + image_grid_thw=None, + future_depth_targets=None, + future_video_targets=None, + future_video_cls_targets=None, + future_video_current_patch=None, + **kwargs + ) -> tuple[Tensor, dict[str, Tensor]]: + loss_dict = {} + if getattr(self.config, "action_fp32", False): + state = state.float() + actions = actions.float() + ( + losses, + loss_depth, + loss_future_depth, + loss_future_video, + depth_preds, + seq_wise_loss, + router_z_loss, + moe_metrics, + future_depth_preds, + future_video_preds, + current_video_preds, + ) = self.model.forward( + images, + img_masks, + lang_tokens, + lang_masks, + state, + actions, + noise, + time, + loss_type=self.config.loss_type, + depth_targets=depth_targets, + image_grid_thw=image_grid_thw, + future_depth_targets=future_depth_targets, + future_video_targets=future_video_targets, + future_video_cls_targets=future_video_cls_targets, + future_video_current_patch=future_video_current_patch, + ) + + if joint_mask is not None: + if "repeat" in self.config.loss_type: + joint_mask = joint_mask.repeat(2, 1, 1) + assert len(joint_mask.shape) == 3 + + masked_losses = losses * joint_mask + valid_counts = joint_mask.sum(dim=(1, 2)).clamp(min=1) + batch_mean_losses = masked_losses.sum(dim=(1, 2)) / valid_counts + loss_vla = masked_losses.sum() / joint_mask.sum().clamp(min=1) + else: + losses = losses[:, :, : self.config.action_dim] + batch_mean_losses = losses.mean(dim=(1, 2)) + loss_vla = losses.mean() + + loss_dict["batch_mean_losses"] = batch_mean_losses.detach() + total_loss = ( + loss_vla + + loss_depth + + loss_future_depth + + loss_future_video + + seq_wise_loss + + router_z_loss + ) + loss_dict["router_z_loss"] = router_z_loss.detach() if torch.is_tensor(router_z_loss) else router_z_loss + if moe_metrics: + loss_dict.update(moe_metrics) + return total_loss, loss_vla, loss_depth, loss_future_depth, loss_future_video, seq_wise_loss, loss_dict, depth_preds, future_depth_preds, future_video_preds, current_video_preds + + def sample_actions(self, *args, **kwargs) -> Tensor: + return self.model.sample_actions(*args, **kwargs) + + +ModelClass = LingbotVlaV2Policy + +__all__ = [ + "LingbotVlaV2Policy", + "Qwen3VLForConditionalGeneration", + "Qwen3VLTextModel", + "Qwen3VLPreTrainedModel", + "Qwen2ForCausalLM", +] +# __V2_END__ + +from telefuser.models.lingbot_vla_v2_loader import LingBotVlaV2StateDictConverter +from telefuser.models.lingbot_vla_v2_qwen import apply_lingbot_qwen3_vl_patch + + +class LingBotVlaV2Model(LingbotVlaV2Policy): + """TeleFuser-native entry point preserving official checkpoint key names.""" + + name = "lingbot_vla_v2" + + def __init__(self, config, eval=True): + apply_lingbot_qwen3_vl_patch() + super().__init__(config=config, eval=eval) + + @staticmethod + def state_dict_converter(**kwargs): + return LingBotVlaV2StateDictConverter(**kwargs) + +# __WRAPPER_END__ diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py new file mode 100644 index 00000000..38caa903 --- /dev/null +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -0,0 +1,1787 @@ +"""Native utility, alignment, and checkpoint support for LingBot-VLA v2. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + +import math + +import einops +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from packaging.version import Version +# from xformers.ops import memory_efficient_attention + + +def find_next_divisible_by_8_numpy(n: np.ndarray) -> np.ndarray: + """ + Finds the smallest integers greater than each element in a NumPy array 'n' + that are divisible by 8. Assumes non-negative integers. + + Args: + n: A NumPy array of integers. + + Returns: + A NumPy array containing the smallest integers greater than each input element + that are divisible by 8. + """ + remainder = n % 8 + # Calculate the amount to add: 0 if already divisible, otherwise 8 - remainder + # np.where is efficient for conditional operations on arrays + amount_to_add = np.where(remainder == 0, 8, 8 - remainder) + return n + amount_to_add + + +def create_sinusoidal_pos_embedding( + time: torch.tensor, + dimension: int, + min_period: float, + max_period: float, + device="cpu", +) -> Tensor: + """Computes sine-cosine positional embedding vectors for scalar positions.""" + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + + if time.ndim != 1: + raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") + + fraction = torch.linspace( + 0.0, 1.0, dimension // 2, dtype=torch.float32, device=device + ) + period = min_period * (max_period / min_period) ** fraction + + # Compute the outer product + scaling_factor = 1.0 / period * 2 * math.pi + sin_input = scaling_factor[None, :] * time[:, None] + pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + return pos_emb + + +def sample_beta(alpha, beta, bsize, device): + gamma1 = torch.rand((bsize,), device=device).pow(1 / alpha) + gamma2 = torch.rand((bsize,), device=device).pow(1 / beta) + return gamma1 / (gamma1 + gamma2) + + +def make_att_2d_masks(pad_masks, att_masks): + """Copied from big_vision. + + Tokens can attend to valid inputs tokens which have a cumulative mask_ar + smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to + setup several types of attention, for example: + + [[1 1 1 1 1 1]]: pure causal attention. + + [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between + themselves and the last 3 tokens have a causal attention. The first + entry could also be a 1 without changing behaviour. + + [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a + block can attend all previous blocks and all tokens on the same block. + + Args: + input_mask: bool[B, N] true if its part of the input, false if padding. + mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on + it and 0 where it shares the same attention mask as the previous token. + """ + if att_masks.ndim != 2: + raise ValueError(att_masks.ndim) + if pad_masks.ndim != 2: + raise ValueError(pad_masks.ndim) + + cumsum = torch.cumsum(att_masks, dim=1) + att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] + att_2d_masks = att_2d_masks & pad_2d_masks + return att_2d_masks + + +def prefix_query_segments( + use_depth_align, + use_future_depth, + use_future_video=False, + use_future_video_cls=False, + use_future_video_patch=True, + future_video_share_future_depth_query=False, +): + """Return prefix segment order after the image block. + + Task-specific query tokens are always placed after language tokens. Current + task queries precede future task queries; future-depth remains the last + query segment so the existing suffix-to-future-depth blocking can keep using + the tail span. + """ + segments = ["language"] + if not use_depth_align: + return tuple(segments) + + segments.append("current_depth") + if use_future_video: + if use_future_video_cls: + segments.append("future_video_cls") + if use_future_video_patch and not future_video_share_future_depth_query: + segments.append("future_video") + if use_future_depth: + segments.append("future_depth") + return tuple(segments) + + +def prefix_query_token_spans( + prefix_len, + num_task_tokens, + use_depth_align, + use_future_depth, + use_future_video=False, + use_future_video_cls=False, + use_future_video_patch=True, + future_video_share_future_depth_query=False, +): + """Return [start, end) spans for non-language task query segments.""" + counts = { + "current_depth": num_task_tokens, + "future_video_cls": 1, + "future_video": num_task_tokens, + "future_depth": num_task_tokens, + } + ordered = prefix_query_segments( + use_depth_align=use_depth_align, + use_future_depth=use_future_depth, + use_future_video=use_future_video, + use_future_video_cls=use_future_video_cls, + use_future_video_patch=use_future_video_patch, + future_video_share_future_depth_query=future_video_share_future_depth_query, + ) + query_segments = [name for name in ordered if name != "language"] + cursor = prefix_len - sum(counts[name] for name in query_segments) + spans = {} + for name in query_segments: + count = counts[name] + spans[name] = (cursor, cursor + count) + cursor += count + return spans + + +def fv_col_span(prefix_len, num_task_tokens, use_cls, use_patch): + """Return [start, end) of a tail query block inside the prefix. + + This legacy helper is still used for future-depth tail blocking in V2. + New prefix layout code should prefer prefix_query_token_spans(), which also + handles current-depth and separate future-video spans. + """ + fv_len = (1 if use_cls else 0) + (num_task_tokens if use_patch else 0) + return prefix_len - fv_len, prefix_len + +def block_suffix_to_fv_(att_2d_masks, suffix_row_start, prefix_len, + num_task_tokens, use_cls=False, use_patch=True, drop_mask=None): + """In-place mask out the suffix-to-future-video attention edge. + + `make_att_2d_masks`' cumsum scheme cannot express "a query cannot see a + segment that precedes it", so we zero the rectangular [suffix rows, FV cols] + block on the already-built 2D mask instead of touching mask_ar. + + att_2d_masks: bool[B, Q, K], True == visible. `suffix_row_start` is the first + query row belonging to the suffix: prefix_len in the square training mask, + 0 in the suffix-only inference mask. Leaves FV -> img/lang rows untouched so + the distillation query still reads the current observation. + + `drop_mask`: optional bool[B], True where this sample's suffix must NOT see + FV. None == block every sample (hard mask). Used for per-sample stochastic + masking (FV-attention dropout): keep = visible iff not dropped, applied via + broadcast multiply so it stays a static graph under torch.compile. + """ + fv_start, fv_end = fv_col_span(prefix_len, num_task_tokens, use_cls, use_patch) + if fv_end <= fv_start: + return att_2d_masks + if drop_mask is None: + att_2d_masks[:, suffix_row_start:, fv_start:fv_end] = False + else: + # keep[b] = True where the sample is NOT dropped -> AND keeps those rows + # visible and zeros the dropped ones, with no data-dependent indexing. + keep = (~drop_mask).view(-1, 1, 1) + block = att_2d_masks[:, suffix_row_start:, fv_start:fv_end] + att_2d_masks[:, suffix_row_start:, fv_start:fv_end] = block & keep + return att_2d_masks + +def resize_with_pad(img, width, height, pad_value=-1): + # assume no-op when width height fits already + if img.ndim != 4: + raise ValueError(f"(b,c,h,w) expected, but {img.shape}") + + cur_height, cur_width = img.shape[2:] + + ratio = max(cur_width / width, cur_height / height) + resized_height = int(cur_height / ratio) + resized_width = int(cur_width / ratio) + resized_img = F.interpolate( + img, size=(resized_height, resized_width), mode="bilinear", align_corners=False + ) + + pad_height = max(0, int(height - resized_height)) + pad_width = max(0, int(width - resized_width)) + + # pad on left and top of image + padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value) + return padded_img + + +def our_eager_attention_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: torch.Tensor, +): + """ + Performs eager attention, optimized with torch.einsum. + + Args: + query_states: Query tensor of shape [batch_size, seq_len, num_attention_heads, head_dim]. + key_states: Key tensor of shape [batch_size, seq_len, num_key_value_heads, head_dim]. + value_states: Value tensor of shape [batch_size, seq_len, num_key_value_heads, head_dim]. + attention_mask: Attention mask tensor, typically [batch_size, 1, seq_len, seq_len] or [batch_size, seq_len, seq_len]. + + Returns: + Output tensor of shape [batch_size, seq_len, num_attention_heads * head_dim]. + """ + bsize, seq_len, num_att_heads, head_dim = query_states.shape + num_key_value_heads = key_states.shape[2] + num_key_value_groups = num_att_heads // num_key_value_heads + + key_states = einops.repeat( + key_states, "b l h d -> b l (h g) d", g=num_key_value_groups + ) + value_states = einops.repeat( + value_states, "b l h d -> b l (h g) d", g=num_key_value_groups + ) + + query_states_permuted = torch.einsum("blhd->bhld", query_states) + key_states_permuted = torch.einsum("blhd->bhld", key_states) + + att_weights = torch.einsum( + "bhqd,bhkd->bhqk", query_states_permuted, key_states_permuted + ) + att_weights *= head_dim**-0.5 + + big_neg = -2.3819763e38 + masked_att_weights = torch.where( + attention_mask[:, None, :, :], att_weights, big_neg + ) + + probs = nn.functional.softmax(masked_att_weights, dim=-1) + probs = probs.to(dtype=value_states.dtype) + + value_states_permuted = torch.einsum("blhd->bhld", value_states) # [B, H, L_v, D] + att_output = torch.einsum( + "bhqk,bhkv->bhqv", probs, value_states_permuted + ) # [B, H, L_q, D] + att_output = torch.einsum("bhld->blhd", att_output) # [B, L, H, D] + att_output = att_output.reshape(bsize, seq_len, num_att_heads * head_dim) + + return att_output + + +# @torch.jit.script +def apply_rope( + x: torch.Tensor, + positions: torch.Tensor, + max_wavelength: float = 10_000.0, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Applies RoPE positions [B, L] to x [B, L, H, D].""" + original_dtype = x.dtype # bf16 + d = x.shape[-1] + d_half = d // 2 + device = x.device + + # Cast input to compute_dtype for all internal operations + x_casted = x.to(dtype) + positions_casted = positions.to(dtype) + + freq_exponents = (2.0 / d) * torch.arange(d_half, dtype=dtype, device=device) + timescale = max_wavelength**freq_exponents + radians = torch.einsum("bl,h->blh", positions_casted, 1.0 / timescale) # fp32 -> bf16 + + radians = radians[..., None, :] # [B, L, 1, D_half] + + sin = torch.sin(radians) # bf16 + cos = torch.cos(radians) # bf16 + + x1, x2 = x_casted.split(d_half, dim=-1) # fp32 + + res = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) # fp32 + + return res.to(original_dtype) # bf16 + + + +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn.functional as F # noqa: N812 +from packaging.version import Version +import einops + +FLEX_SPARSE_BLOCK_SIZE = 128 +FLEX_KERNEL_OPTIONS = {"BLOCK_M": 32, "BLOCK_N": 64, "num_warps": 4, "num_stages": 2} + +if Version(torch.__version__) > Version("2.5.0"): + # Ffex attention is only available from torch 2.5 onwards + from torch.nn.attention.flex_attention import ( + _mask_mod_signature, + _round_up_to_multiple, + create_block_mask, + create_mask, + flex_attention, + ) + +# @torch.compile(dynamic=False) +def flex_attention_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: torch.Tensor, + scaling=None, +): + """ + This is defined out of classes to make compile happy. + """ + batch_size, seq_len, num_att_heads, head_dim = query_states.shape + original_dtype = query_states.dtype + num_key_value_heads = key_states.shape[2] + # num_key_value_groups = num_att_heads // num_key_value_heads # 16 // 2 = 8 + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + query_states = query_states.to(torch.float32) + key_states = key_states.to(torch.float32) + value_states = value_states.to(torch.float32) + + causal_mask = attention_mask + if causal_mask is not None: + causal_mask = causal_mask[:, None, :, : key_states.shape[2]] + + if causal_mask.shape[1] == 1 and query_states.shape[1] > 1: + causal_mask = causal_mask.expand(-1, query_states.shape[1], -1, -1) + + def precomputed_mask_factory(precomputed_mask: torch.Tensor) -> _mask_mod_signature: + def mask_mod(b, h, q_idx, kv_idx): + # Danger zone: if b,h,q_idx,kv_idx exceed the shape, device-side assert occurs. + return precomputed_mask[b][h][q_idx][kv_idx] + + return mask_mod + + b_mask, h_mask, q_len, kv_len = causal_mask.shape # The shape of your mask + # ipdb.set_trace() + block_size = FLEX_SPARSE_BLOCK_SIZE + q_len_rounded = _round_up_to_multiple(q_len, block_size) + kv_len_rounded = _round_up_to_multiple(kv_len, block_size) + + # *CRITICAL* we do need to expand here, else we get a CUDA index error + + pad_q = q_len_rounded - q_len + pad_k = kv_len_rounded - kv_len + + if pad_q > 0: + query_states = F.pad(query_states, (0, 0, 0, pad_q), value=0.0) # [B, H, q_len_rounded, D] + if pad_k > 0: + key_states = F.pad(key_states, (0, 0, 0, pad_k), value=0.0) + value_states = F.pad(value_states, (0, 0, 0, pad_k), value=0.0) + padded_causal_mask = F.pad(causal_mask, (0, pad_k, 0, pad_q), value=0.0) + mask_mod_fn_orig = precomputed_mask_factory(padded_causal_mask) + + mask_4d = create_mask( + mod_fn=mask_mod_fn_orig, + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, + device=causal_mask.device, + ) + + mask_mod_fn_padded = precomputed_mask_factory(mask_4d) + block_mask = create_block_mask( + mask_mod=mask_mod_fn_padded, + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, + BLOCK_SIZE=block_size, + device=causal_mask.device, + _compile=False, + ) + + # mask is applied inside the kernel, ideally more efficiently than score_mod. + attn_output, attention_weights = flex_attention( + query_states, + key_states, + value_states, + block_mask=block_mask, + enable_gqa=True, # because we shaped query/key states for GQA + scale=head_dim**-0.5 if scaling is None else scaling, + return_lse=True, + kernel_options=FLEX_KERNEL_OPTIONS, + ) + attn_output = attn_output[:, :, :seq_len, :].to(dtype=original_dtype) + attn_output = attn_output.transpose(1, 2).contiguous() # [B, Q_LEN, H, head_dim] + attn_output = attn_output.reshape( + batch_size, + -1, + attn_output.shape[2] * attn_output.shape[3], # merges [H, head_dim] + ) + return attn_output + + +@torch.compiler.disable +def build_block_mask( + attention_mask_3d: torch.Tensor, + num_heads: int, + q_len: int, + kv_len: int, + block_size: int = FLEX_SPARSE_BLOCK_SIZE, +): + """ + Build a reusable BlockMask from a 3D attention mask [B, Q, KV]. + This allocates the dense 4D mask once; the returned BlockMask can be reused across layers. + """ + from torch.nn.attention.flex_attention import ( + _mask_mod_signature, + _round_up_to_multiple, + create_block_mask, + create_mask, + ) + + causal_mask = attention_mask_3d[:, None, :, :].expand(-1, num_heads, -1, -1).contiguous() + b_mask, h_mask = causal_mask.shape[0], causal_mask.shape[1] + + q_len_rounded = _round_up_to_multiple(q_len, block_size) + kv_len_rounded = _round_up_to_multiple(kv_len, block_size) + + pad_q = q_len_rounded - q_len + pad_k = kv_len_rounded - kv_len + padded_mask = F.pad(causal_mask, (0, pad_k, 0, pad_q), value=0.0) + + def precomputed_mask_factory(precomputed_mask: torch.Tensor): + def mask_mod(b, h, q_idx, kv_idx): + return precomputed_mask[b][h][q_idx][kv_idx] + return mask_mod + + mask_4d = create_mask( + mod_fn=precomputed_mask_factory(padded_mask), + B=b_mask, H=h_mask, + Q_LEN=q_len_rounded, KV_LEN=kv_len_rounded, + device=causal_mask.device, + ) + + block_mask = create_block_mask( + mask_mod=precomputed_mask_factory(mask_4d), + B=b_mask, H=h_mask, + Q_LEN=q_len_rounded, KV_LEN=kv_len_rounded, + BLOCK_SIZE=block_size, + device=causal_mask.device, + _compile=False, + ) + return block_mask + + +def flex_attention_with_block_mask( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + block_mask, + seq_len: int, + scaling=None, +): + """ + Run flex_attention with a pre-built BlockMask (no create_mask allocation per call). + """ + batch_size = query_states.shape[0] + num_att_heads = query_states.shape[2] + head_dim = query_states.shape[3] + original_dtype = query_states.dtype + + query_states = query_states.transpose(1, 2).to(torch.float32) + key_states = key_states.transpose(1, 2).to(torch.float32) + value_states = value_states.transpose(1, 2).to(torch.float32) + + q_len_rounded = block_mask.shape[-2] if hasattr(block_mask, 'shape') else query_states.shape[2] + kv_len_rounded = block_mask.shape[-1] if hasattr(block_mask, 'shape') else key_states.shape[2] + + pad_q = q_len_rounded - query_states.shape[2] + pad_k = kv_len_rounded - key_states.shape[2] + + if pad_q > 0: + query_states = F.pad(query_states, (0, 0, 0, pad_q), value=0.0) + if pad_k > 0: + key_states = F.pad(key_states, (0, 0, 0, pad_k), value=0.0) + value_states = F.pad(value_states, (0, 0, 0, pad_k), value=0.0) + + attn_output, _ = flex_attention( + query_states, + key_states, + value_states, + block_mask=block_mask, + enable_gqa=True, + scale=head_dim**-0.5 if scaling is None else scaling, + return_lse=True, + kernel_options=FLEX_KERNEL_OPTIONS, + ) + attn_output = attn_output[:, :, :seq_len, :].to(dtype=original_dtype) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(batch_size, -1, attn_output.shape[2] * attn_output.shape[3]) + return attn_output + + + +# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + #(bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + + +class AttentionPool2d(nn.Module): + + def __init__(self, seq_len: int, embed_dim: int, num_heads: int, output_dim: int = None): + super().__init__() + self.positional_embedding = nn.Parameter(torch.randn(seq_len + 1, embed_dim) / embed_dim**0.5) + self.k_proj = nn.Linear(embed_dim, embed_dim) + self.q_proj = nn.Linear(embed_dim, embed_dim) + self.v_proj = nn.Linear(embed_dim, embed_dim) + self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) + self.num_heads = num_heads + + def forward(self, x, return_all_tokens=False): + # x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1) # NCHW -> (HW)NC + x = x.permute(1, 0, 2) # (N(HW)C) => (HW)NC + x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC + x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC + x, _ = F.multi_head_attention_forward(query=x, + key=x, + value=x, + embed_dim_to_check=x.shape[-1], + num_heads=self.num_heads, + q_proj_weight=self.q_proj.weight, + k_proj_weight=self.k_proj.weight, + v_proj_weight=self.v_proj.weight, + in_proj_weight=None, + in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), + bias_k=None, + bias_v=None, + add_zero_attn=False, + dropout_p=0, + out_proj_weight=self.c_proj.weight, + out_proj_bias=self.c_proj.bias, + use_separate_proj_weight=True, + training=self.training, + need_weights=False) + if return_all_tokens: + return x + else: + return x[0] + + +class Resampler(nn.Module): + + def __init__( + self, + dim_in=768, + dim_mid=1024, + dim_head=64, + dim_out=1024, + num_layers=8, + num_queries=8, + num_heads=16, + ff_mult=4, + ): + super().__init__() + + self.queries = nn.Parameter(torch.randn(1, num_queries, dim_in) / dim_mid ** 0.5) + + self.proj_in = nn.Linear(dim_in, dim_mid) + self.proj_out = nn.Linear(dim_mid, dim_out) + self.norm_out = nn.LayerNorm(dim_out) + + self.layers = nn.ModuleList([]) + for _ in range(num_layers): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim_mid, dim_head=dim_head, heads=num_heads), + FeedForward(dim=dim_mid, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + queries = self.queries.repeat(x.size(0), 1, 1) + x = self.proj_in(x) + + for attn, ff in self.layers: + queries = attn(x, queries) + queries + queries = ff(queries) + queries + + queries = self.proj_out(queries) + queries = self.norm_out(queries) + return queries + +class TaskTokenResampler(nn.Module): + + def __init__( + self, + dim_in=768, + dim_mid=1024, + dim_head=64, + dim_out=1024, + num_layers=8, + num_queries=8, + num_heads=16, + ff_mult=4, + ): + super().__init__() + + self.num_queries = num_queries + self.proj_in1 = nn.Linear(dim_in, dim_mid) + self.proj_in2 = nn.Linear(dim_in, dim_mid) + self.proj_out = nn.Linear(dim_mid, dim_out) + self.norm_out = nn.LayerNorm(dim_out) + + self.layers = nn.ModuleList([]) + for _ in range(num_layers): + self.layers.append( + nn.ModuleList([ + PerceiverAttention(dim=dim_mid, dim_head=dim_head, heads=num_heads), + FeedForward(dim=dim_mid, mult=ff_mult), + ])) + + def forward(self, x, queries): + queries = self.proj_in1(queries) + x = self.proj_in2(x) + + for attn, ff in self.layers: + queries = attn(x, queries) + queries + queries = ff(queries) + queries + + queries = self.proj_out(queries) + queries = self.norm_out(queries) + return queries + + +class ResamplerXL(nn.Module): + + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output1_dim=768, + output2_dim=1280, + ff_mult=4, + ): + super().__init__() + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.proj_in = nn.Linear(embedding_dim, dim) + + # self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(dim) + + self.in_dim = dim + self.out_dim = output1_dim + output2_dim + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList([ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ])) + + self.unet_proj_1 = nn.Linear(self.in_dim, output1_dim) + self.unet_proj_2 = nn.Linear(self.in_dim, output2_dim) + self.unet_attnpool = AttentionPool2d(num_queries, self.in_dim, heads, output2_dim) + + def forward(self, x): + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + hidden_embeds = self.norm_out(latents) + + encoder_hidden_1 = self.unet_proj_1(hidden_embeds) # [bs, 256, 768] + encoder_hidden_2 = self.unet_proj_2(hidden_embeds) # [bs, 256, 1280] + prompt_embeds = torch.cat([encoder_hidden_1, encoder_hidden_2], dim=-1) # [bs, 256, 2048] + pooled_prompt_embeds = self.unet_attnpool(hidden_embeds) # [bs, 1280] + + return prompt_embeds, pooled_prompt_embeds + + +class ResamplerXLV2(nn.Module): + + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output1_dim=768, + output2_dim=1280, + ff_mult=4, + normalize=True + ): + super().__init__() + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.normalize = normalize + self.proj_in = nn.Linear(embedding_dim, dim) + + # self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(dim) + + self.in_dim = dim + self.out_dim = output1_dim + output2_dim + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList([ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ])) + + self.unet_proj_1 = nn.Linear(self.in_dim, output1_dim) + self.unet_proj_2 = nn.Linear(self.in_dim, output2_dim) + self.unet_attnpool = AttentionPool2d(num_queries, self.in_dim, heads, output2_dim) + + def forward(self, x,pooled_text_embeds=None): + + latents = self.latents.repeat(x.size(0), 1, 1) + + if self.normalize: + x = F.normalize(x) + + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + hidden_embeds = self.norm_out(latents) + + encoder_hidden_1 = self.unet_proj_1(hidden_embeds) # [bs, 256, 768] + encoder_hidden_2 = self.unet_proj_2(hidden_embeds) # [bs, 256, 1280] + prompt_embeds = torch.cat([encoder_hidden_1, encoder_hidden_2], dim=-1) # [bs, 256, 2048] + pooled_prompt_embeds = self.unet_attnpool(hidden_embeds) # [bs, 1280] + + return prompt_embeds, pooled_prompt_embeds + +class ResamplerXLIdentity(nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x, pooled_text_embeds=None): + return x, pooled_text_embeds + + +if __name__ == '__main__': + image_proj_model = Resampler(dim=1024, + depth=4, + dim_head=64, + heads=12, + num_queries=1024, + embedding_dim=1024, + output_dim=1024, + ff_mult=4) + numel = 0 + for name, param in image_proj_model.named_parameters(): + numel += param.numel() + + print(f'Total params: {numel}') + + + +import torch.nn as nn + +def build_mlp(in_hidden_size, hidden_size): + modules = [nn.Linear(in_hidden_size, hidden_size)] + modules.append(nn.ReLU()) + modules.append(nn.Linear(hidden_size, hidden_size)) + return nn.Sequential(*modules) + +def build_expand_mlp(in_hidden_size, hidden_size, out_size): + modules = [nn.Linear(in_hidden_size, hidden_size)] + modules.append(nn.ReLU()) + modules.append(nn.Linear(hidden_size, hidden_size)) + modules.append(nn.ReLU()) + modules.append(nn.Linear(hidden_size, out_size)) + return nn.Sequential(*modules) + +class DepthHead(nn.Module): + def __init__( + self, + proj_config=None, + llm_hidden_size=4096, + use_intermediate_depth=False, + ): + super(DepthHead, self).__init__() + + self.projector = Resampler( + dim_in=llm_hidden_size, + dim_mid=llm_hidden_size, + dim_head=proj_config["dim_head"], + dim_out=proj_config["dim_out"], + num_layers=proj_config["num_layers"], + num_heads=proj_config["num_heads"], + num_queries=proj_config["num_backbone_tokens"], + ff_mult=proj_config["ff_mult"], + ) + + def forward(self, llm_feats): + queries = self.projector(llm_feats) + return queries + +class TaskTokenDepthHead(nn.Module): + def __init__( + self, + proj_config=None, + llm_hidden_size=4096, + use_intermediate_depth=False, + ): + super(TaskTokenDepthHead, self).__init__() + + self.projector = TaskTokenResampler( + dim_in=llm_hidden_size, + dim_mid=llm_hidden_size, + dim_head=proj_config["dim_head"], + dim_out=proj_config["dim_out"], + num_layers=proj_config["num_layers"], + num_heads=proj_config["num_heads"], + num_queries=proj_config["num_backbone_tokens"], + ff_mult=proj_config["ff_mult"], + ) + + def forward(self, llm_feats, queries): + queries = self.projector(llm_feats, queries) + return queries + + + +# Copyright 2025 Ant Group Co., Ltd. All Rights Reserved. +# Developer: xiancun +# Project锛?Lumos VIdeo Generation Foundation Model +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Triton-optimized MoE auxiliary loss functions. + +Provides numerically equivalent replacements for the functions in loss.py, +with two key optimizations: + 1. Eliminate Python for-loops via vectorized segment-wise operations. + 2. Fuse topK + counting into Triton kernels to avoid huge intermediate tensors. + +Usage: + from telefuser.models.lingbot_vla_v2_loader import ( + triton_load_balancing_loss_func, + triton_sequence_wise_balance_loss, + ) + # Drop-in replacement 鈥?same signature and return type as loss.py +""" + +import torch +import torch.nn.functional as F +from typing import List, Optional, Tuple, Union + + +def _next_power_of_2(n: int) -> int: + """Return the smallest power of 2 >= n.""" + if n <= 0: + return 1 + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + return n + 1 + + +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? +# Section 1: Triton availability check + kernel definitions +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? + +_HAS_TRITON = False +try: + import triton + import triton.language as tl + + _HAS_TRITON = True + + # 鈹€鈹€ Kernel 1: per-segment topK counting (for sequence_wise_balance_loss) 鈹€鈹€ + @triton.jit + def _topk_segment_count_kernel( + logits_ptr, # [N_total, E] + seg_starts_ptr, # [S_total] + seg_lengths_ptr, # [S_total] + f_out_ptr, # [S_total, E] + stride_logits_n, # stride of logits along token dim + E: tl.constexpr, # actual num_experts + K: tl.constexpr, # top_k + BLOCK_E: tl.constexpr, # next power of 2 >= E + ): + """Each program computes f_i (expert counts) for one segment. + + Iterates over tokens in the segment, performs K rounds of argmax to + find top-K experts, and accumulates per-expert hit counts. + No gradient needed 鈥?f_i is always detached in the loss. + """ + seg_id = tl.program_id(0) + seg_start = tl.load(seg_starts_ptr + seg_id) + seg_len = tl.load(seg_lengths_ptr + seg_id) + + expert_offs = tl.arange(0, BLOCK_E) # [BLOCK_E] + mask_e = expert_offs < E + f_acc = tl.zeros((BLOCK_E,), dtype=tl.float32) + + for t in range(0, seg_len): + row_ptr = logits_ptr + (seg_start + t) * stride_logits_n + logits_row = tl.load(row_ptr + expert_offs, mask=mask_e, other=float('-inf')) + + # K rounds of argmax to find top-K indices + row_copy = logits_row + for _k in range(K): + max_val = tl.max(row_copy, axis=0) + is_max = (row_copy == max_val) + # Distribute count evenly among ties (rare with float32) + n_ties = tl.sum(is_max.to(tl.float32), axis=0) + f_acc += tl.where(is_max, 1.0 / n_ties, 0.0) + row_copy = tl.where(is_max, float('-inf'), row_copy) + + # Write f_count (unnormalized) 鈥?caller normalizes by (E / K) / seg_len + out_ptr = f_out_ptr + seg_id * E + tl.store(out_ptr + expert_offs, f_acc, mask=mask_e) + + # 鈹€鈹€ Kernel 2: blocked topK counting (for load_balancing_loss_func) 鈹€鈹€ + @triton.jit + def _topk_count_with_mask_kernel( + routing_weights_ptr, # [N, E] 鈥?softmax probabilities + mask_ptr, # [N] 鈥?1.0 for valid, 0.0 for padding + partial_f_ptr, # [num_blocks, E] 鈥?partial expert counts + partial_p_ptr, # [num_blocks, E] 鈥?partial masked prob sums + N, + stride_rw_n, # stride along token dim + has_mask: tl.constexpr, + E: tl.constexpr, + K: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + """Each program accumulates topK counts + masked probs for a token block. + + Two-phase reduction: writes partial [E] results per block; + caller sums across blocks in PyTorch. + """ + pid = tl.program_id(0) + n_start = pid * BLOCK_N + + expert_offs = tl.arange(0, BLOCK_E) + mask_e = expert_offs < E + f_local = tl.zeros((BLOCK_E,), dtype=tl.float32) + p_local = tl.zeros((BLOCK_E,), dtype=tl.float32) + + for t_offset in range(BLOCK_N): + t = n_start + t_offset + # Guard: skip if t >= N (handles last block) + if t < N: + row_ptr = routing_weights_ptr + t * stride_rw_n + rw_row = tl.load(row_ptr + expert_offs, mask=mask_e, other=0.0) + + if has_mask: + m = tl.load(mask_ptr + t) + else: + m = 1.0 + + # Accumulate masked probs + p_local += rw_row * m + + # TopK counting + row_copy = rw_row + for _k in range(K): + max_val = tl.max(row_copy, axis=0) + is_max = (row_copy == max_val) + n_ties = tl.sum(is_max.to(tl.float32), axis=0) + f_local += tl.where(is_max, m / n_ties, 0.0) + row_copy = tl.where(is_max, float('-inf'), row_copy) + + # Write partial results (only E valid elements) + f_ptr = partial_f_ptr + pid * E + p_ptr = partial_p_ptr + pid * E + tl.store(f_ptr + expert_offs, f_local, mask=mask_e) + tl.store(p_ptr + expert_offs, p_local, mask=mask_e) + +except ImportError: + pass + + +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? +# Section 2: Vectorized PyTorch fallback (no Triton needed) +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? + +def _build_segment_info( + seq_lengths_per_layer: List[int], + num_layers: int, + N_per_layer: int, + device: torch.device, +): + """Build segment IDs and metadata for all layers combined. + + Returns: + segment_ids: [num_layers * N_valid_per_layer] int64 + seg_starts: [total_segments] int64 + seg_lengths: [total_segments] int64 + total_segments: int + """ + S = len(seq_lengths_per_layer) + total_segments = num_layers * S + seg_lengths_t = torch.tensor(seq_lengths_per_layer, dtype=torch.int64, device=device) + + # Repeat for all layers + all_seg_lengths = seg_lengths_t.repeat(num_layers) # [total_segments] + + # Per-layer segment starts: cumsum of seq_lengths + per_layer_starts = torch.cumsum(seg_lengths_t, 0) - seg_lengths_t # [S] + # Layer offsets in the concatenated tensor + layer_offsets = torch.arange(num_layers, device=device, dtype=torch.int64) * N_per_layer + # [L, S] 鈫?[L*S] + all_seg_starts = (per_layer_starts.unsqueeze(0) + layer_offsets.unsqueeze(1)).reshape(-1) + + # Segment IDs: [num_layers * N_valid] + segment_ids = torch.repeat_interleave( + torch.arange(total_segments, device=device), all_seg_lengths, + ) + + return segment_ids, all_seg_starts, all_seg_lengths, total_segments + + +def _vectorized_segment_f_i( + logits: torch.Tensor, # [N_total, E] + seg_starts: torch.Tensor, # [S_total] + seg_lengths: torch.Tensor, # [S_total] + top_k: int, +) -> torch.Tensor: + """Compute f_i per segment without Triton, using vectorized PyTorch ops. + + Returns: f_i [S_total, E] + """ + N, E = logits.shape + S_total = seg_starts.shape[0] + + # TopK over all tokens at once + _, topk_idx = torch.topk(logits, k=top_k, dim=-1) # [N, K] + + # Build one-hot mask efficiently: scatter into [N, E] + mask = torch.zeros(N, E, device=logits.device, dtype=torch.float32) + mask.scatter_(1, topk_idx, 1.0) + + # Segment-wise sum using scatter_add + segment_ids = torch.repeat_interleave( + torch.arange(S_total, device=logits.device), seg_lengths, + ) + seg_ids_exp = segment_ids.unsqueeze(1).expand(-1, E) # [N, E] + + f_sum = torch.zeros(S_total, E, device=logits.device, dtype=torch.float32) + f_sum.scatter_add_(0, seg_ids_exp, mask) + + # Normalize: f_i = (E / K) * f_sum / T_s + inv_lens = (float(E) / top_k) / seg_lengths.unsqueeze(1).float().clamp(min=1) + f_i = f_sum * inv_lens + + return f_i + + +def _vectorized_topk_count( + routing_weights: torch.Tensor, # [N, E] + top_k: int, + flat_mask: Optional[torch.Tensor] = None, # [N] float +) -> torch.Tensor: + """Count per-expert topK selections, optionally masked. Returns [E].""" + N, E = routing_weights.shape + _, topk_idx = torch.topk(routing_weights, k=top_k, dim=-1) # [N, K] + + tokens_per_expert = torch.zeros(E, device=routing_weights.device, dtype=torch.float32) + weight = flat_mask if flat_mask is not None else torch.ones(N, device=routing_weights.device, dtype=torch.float32) + + # K rounds of scatter_add 鈥?K is small (typically 2-8), no Python overhead concern + for k in range(top_k): + tokens_per_expert.scatter_add_(0, topk_idx[:, k], weight) + + return tokens_per_expert + + +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? +# Section 3: Triton-accelerated wrappers +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? + +def _triton_segment_f_i( + logits: torch.Tensor, # [N_total, E] + seg_starts: torch.Tensor, # [S_total] + seg_lengths: torch.Tensor, # [S_total] + top_k: int, +) -> torch.Tensor: + """Compute f_i per segment using the Triton kernel. Returns [S_total, E].""" + N, E = logits.shape + S_total = seg_starts.shape[0] + BLOCK_E = _next_power_of_2(E) + + f_counts = torch.zeros(S_total, E, device=logits.device, dtype=torch.float32) + + _topk_segment_count_kernel[(S_total,)]( + logits, + seg_starts, + seg_lengths, + f_counts, + logits.stride(0), + E=E, + K=top_k, + BLOCK_E=BLOCK_E, + ) + + # Normalize: f_i = (E / K) * counts / T_s + inv_lens = (float(E) / top_k) / seg_lengths.unsqueeze(1).float().clamp(min=1) + return f_counts * inv_lens + + +def _triton_topk_count( + routing_weights: torch.Tensor, # [N, E] + top_k: int, + flat_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Compute tokens_per_expert [E] using the Triton kernel.""" + N, E = routing_weights.shape + BLOCK_E = _next_power_of_2(E) + BLOCK_N = 256 + num_blocks = (N + BLOCK_N - 1) // BLOCK_N + + partial_f = torch.zeros(num_blocks, E, device=routing_weights.device, dtype=torch.float32) + partial_p = torch.zeros(num_blocks, E, device=routing_weights.device, dtype=torch.float32) + + has_mask = flat_mask is not None + _topk_count_with_mask_kernel[(num_blocks,)]( + routing_weights, + flat_mask if has_mask else routing_weights, # dummy ptr when no mask + partial_f, + partial_p, + N, + routing_weights.stride(0), + has_mask=has_mask, + E=E, + K=top_k, + BLOCK_E=BLOCK_E, + BLOCK_N=BLOCK_N, + ) + + # Phase 2: reduce across blocks + tokens_per_expert = partial_f.sum(dim=0) # [E] + + if flat_mask is not None: + n_valid = flat_mask.sum().clamp(min=1) + else: + n_valid = float(N) + tokens_per_expert = tokens_per_expert / n_valid + + return tokens_per_expert + + +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? +# Section 4: Main API 鈥?drop-in replacements for loss.py +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? + +def triton_sequence_wise_balance_loss( + router_logits_list: tuple, + top_k: int, + seq_lengths: Optional[List[int]] = None, + padding_len: int = 0, + score_func: str = "softmax", +) -> List[torch.Tensor]: + """Triton-optimized DeepSeek-V3 sequence-wise balance loss. + + Numerically equivalent to sequence_wise_balance_loss() in loss.py, + but eliminates all Python for-loops by: + - Processing all layers simultaneously via concatenation + - Using segment-wise parallel reduction (scatter_add) instead of per-sequence loops + - Fusing topK + counting in a Triton kernel (with PyTorch vectorized fallback) + + Args / Returns: same as sequence_wise_balance_loss in loss.py. + """ + if router_logits_list is None or not isinstance(router_logits_list, (tuple, list)): + return [] + + valid_logits = [rl for rl in router_logits_list if rl is not None] + if len(valid_logits) == 0: + return [] + + num_layers = len(valid_logits) + device = valid_logits[0].device + E = valid_logits[0].shape[1] + + # 鈹€鈹€ Step 1: Concatenate all layers, remove padding 鈹€鈹€ + all_logits_list = [] + N_per_layer = None + for logits in valid_logits: + logits_f32 = logits.to(dtype=torch.float32) + N = logits_f32.shape[0] + if padding_len > 0: + logits_f32 = logits_f32[:N - padding_len] + all_logits_list.append(logits_f32) + if N_per_layer is None: + N_per_layer = logits_f32.shape[0] + + # Check if all layers have the same valid length (common case) + same_length = all(l.shape[0] == N_per_layer for l in all_logits_list) + + if not same_length: + # Rare: different MoE layers have different token counts + return _fallback_per_layer(valid_logits, top_k, seq_lengths, padding_len, score_func) + + if seq_lengths is None or len(seq_lengths) == 0: + seq_lengths_effective = [N_per_layer] + else: + seq_lengths_effective = seq_lengths + + S = len(seq_lengths_effective) + all_logits = torch.cat(all_logits_list, dim=0) # [L * N_valid, E] + + # 鈹€鈹€ Step 2: Build segment metadata 鈹€鈹€ + segment_ids, seg_starts, seg_lengths_t, total_segments = _build_segment_info( + seq_lengths_effective, num_layers, N_per_layer, device + ) + + # 鈹€鈹€ Step 3: P_i via PyTorch (gradient path) 鈹€鈹€ + if score_func == "sigmoid": + all_scores = all_logits.sigmoid() + all_probs = all_scores / all_scores.sum(dim=-1, keepdim=True) + else: + all_probs = F.softmax(all_logits, dim=-1) # [L * N_valid, E] + seg_ids_exp = segment_ids.unsqueeze(1).expand(-1, E) # [L * N_valid, E] + + P_sum = torch.zeros(total_segments, E, device=device, dtype=torch.float32) + P_sum.scatter_add_(0, seg_ids_exp, all_probs) + P_i = P_sum / seg_lengths_t.unsqueeze(1).float().clamp(min=1) # [total_segments, E] + + # 鈹€鈹€ Step 4: f_i (no gradient needed) 鈹€鈹€ + with torch.no_grad(): + if _HAS_TRITON and all_logits.is_cuda: + f_i = _triton_segment_f_i(all_logits, seg_starts, seg_lengths_t, top_k) + else: + f_i = _vectorized_segment_f_i(all_logits, seg_starts, seg_lengths_t, top_k) + + # 鈹€鈹€ Step 5: Per-segment loss 鈫?per-layer mean 鈹€鈹€ + loss_per_seg = (f_i * P_i).sum(dim=-1) # [total_segments] + loss_per_seg = loss_per_seg.reshape(num_layers, S) + layer_losses = loss_per_seg.mean(dim=1) # [L] + + return list(layer_losses.unbind(0)) + + +def triton_load_balancing_loss_func( + gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None], + num_experts: Optional[int] = None, + top_k: int = 2, + attention_mask: Optional[torch.Tensor] = None, +) -> Union[torch.Tensor, int]: + """Triton-optimized Switch Transformer load balancing loss. + + Numerically equivalent to load_balancing_loss_func() in loss.py, + but avoids the huge [N, K, E] one_hot intermediate tensor by + directly counting expert assignments via Triton or scatter_add. + + Memory reduction: O(N*K*E) 鈫?O(N*E + num_blocks*E) + + Args / Returns: same as load_balancing_loss_func in loss.py. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + gate_logits = tuple(g for g in gate_logits if g is not None) + if len(gate_logits) == 0: + return 0 + + compute_device = gate_logits[0].device + concatenated = torch.cat( + [g.to(device=compute_device, dtype=torch.float32) for g in gate_logits], + dim=0, + ) # [L*N, E] + + # 鈹€鈹€ Step 1: softmax (gradient path) 鈹€鈹€ + routing_weights = F.softmax(concatenated, dim=-1) # [L*N, E] + + # 鈹€鈹€ Step 2: Build flat mask 鈹€鈹€ + N_total = routing_weights.shape[0] + if attention_mask is not None: + batch_size, seq_len = attention_mask.shape + num_layers = N_total // (batch_size * seq_len) + flat_mask = ( + attention_mask + .unsqueeze(0) + .expand(num_layers, -1, -1) + .reshape(-1) + .to(device=compute_device, dtype=torch.float32) + ) + else: + flat_mask = None + + # 鈹€鈹€ Step 3: tokens_per_expert (no gradient) 鈹€鈹€ + with torch.no_grad(): + if _HAS_TRITON and routing_weights.is_cuda: + tokens_per_expert = _triton_topk_count(routing_weights, top_k, flat_mask) + else: + tokens_per_expert = _vectorized_topk_count(routing_weights, top_k, flat_mask) + if flat_mask is not None: + tokens_per_expert = tokens_per_expert / flat_mask.sum().clamp(min=1) + else: + tokens_per_expert = tokens_per_expert / float(N_total) + + # 鈹€鈹€ Step 4: router_prob_per_expert (gradient path) 鈹€鈹€ + if flat_mask is not None: + n_valid = flat_mask.sum().clamp(min=1) + router_prob_per_expert = (routing_weights * flat_mask.unsqueeze(1)).sum(0) / n_valid + else: + router_prob_per_expert = routing_weights.mean(dim=0) + + # 鈹€鈹€ Step 5: loss 鈹€鈹€ + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert) + return overall_loss * num_experts + + +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? +# Section 5: Fallback for edge cases +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? + +def _fallback_per_layer( + valid_logits: List[torch.Tensor], + top_k: int, + seq_lengths: Optional[List[int]], + padding_len: int, + score_func: str = "softmax", +) -> List[torch.Tensor]: + """Fallback when layers have different valid token counts. + + Still vectorized within each layer (no per-sequence for-loop). + """ + layer_loss_list = [] + for logits in valid_logits: + logits = logits.to(dtype=torch.float32) + N, E = logits.shape + if padding_len > 0: + logits = logits[:N - padding_len] + if logits.shape[0] == 0: + continue + + if seq_lengths is not None and len(seq_lengths) > 0: + S = len(seq_lengths) + device = logits.device + seg_lengths_t = torch.tensor(seq_lengths, dtype=torch.int64, device=device) + seg_starts = torch.cumsum(seg_lengths_t, 0) - seg_lengths_t + + # P_i (gradient path) + if score_func == "sigmoid": + scores = logits.sigmoid() + probs = scores / scores.sum(dim=-1, keepdim=True) + else: + probs = F.softmax(logits, dim=-1) + segment_ids = torch.repeat_interleave(torch.arange(S, device=device), seg_lengths_t) + seg_ids_exp = segment_ids.unsqueeze(1).expand(-1, E) + P_sum = torch.zeros(S, E, device=device, dtype=torch.float32) + P_sum.scatter_add_(0, seg_ids_exp, probs) + P_i = P_sum / seg_lengths_t.unsqueeze(1).float().clamp(min=1) + + # f_i (no gradient) + with torch.no_grad(): + if _HAS_TRITON and logits.is_cuda: + f_i = _triton_segment_f_i(logits, seg_starts, seg_lengths_t, top_k) + else: + f_i = _vectorized_segment_f_i(logits, seg_starts, seg_lengths_t, top_k) + + loss_per_seq = (f_i * P_i).sum(dim=-1) + layer_loss_list.append(loss_per_seq.mean()) + else: + if score_func == "sigmoid": + scores = logits.sigmoid() + probs = scores / scores.sum(dim=-1, keepdim=True) + else: + probs = F.softmax(logits, dim=-1) + P_i = probs.mean(dim=0) + + with torch.no_grad(): + _, topk_idx = torch.topk(logits, k=top_k, dim=-1) + mask = torch.zeros_like(logits) + mask.scatter_(1, topk_idx, 1.0) + f_i = (E / top_k) * mask.mean(dim=0) + + layer_loss_list.append(torch.sum(f_i * P_i)) + + return layer_loss_list + + +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? +# Section 6: Numerical alignment test +# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? + + + + + +import json +from copy import deepcopy +from pathlib import Path +from typing import Any + +from transformers import AutoConfig + + +class NativeParallelPlan: + """Compatibility container for the upstream training-only parallel plan.""" + + def __init__(self, ep_plan=None): + self.ep_plan = ep_plan or {} + + +class LingBotVLAWeightLoader: + """Minimal native weight-name mapper retained for model compatibility.""" + + def get_vlm_submodule(self, model): + return model.model.qwenvl_with_expert.qwenvl + + def get_expert_vision_submodule(self, model): + return getattr(model.model.qwenvl_with_expert, "expert_visual", None) + + def map_ckpt_key(self, key, load_vlm_only=False, post_training=False): + if key.startswith("expert_visual.") and not post_training: + return "model.qwenvl_with_expert." + key + if load_vlm_only: + return "model.qwenvl_with_expert.qwenvl." + key + return key + + +OFFICIAL_6B_MODEL_CONFIG: dict[str, Any] = { + "post_training": True, + "adanorm_time": True, + "moe_implementation": "fused", + "attention_implementation": "eager", + "precompute_grid_thw": True, + "vlm_causal": True, + "use_moe": True, + "token_moe_layers": list(range(36)), + "token_num_experts": 32, + "token_top_k": 4, + "token_moe_intermediate_size": 512, + "token_shared_intermediate_size": 704, + "bias_update_speed": 0.0, + "sequence_wise_mode": "per_sequence", + "sequence_wise_loss_coeff": 1e-3, + "router_z_loss_coeff": 1e-4, + "router_activation": "sigmoid", + "routed_scaling_factor": 4.0, + "use_shared_expert_gate": False, + "freeze_vision_encoder": False, + "tokenizer_max_length": 72, + "loss_type": "L1_fm", + "action_dim": 55, + "max_action_dim": 55, + "max_state_dim": 55, + "align_params": { + "mode": "query", + "num_task_tokens": 8, + "depth_loss_weight": 0.004, + "future_depth_loss_weight": 0.004, + "use_future_video": True, + "llm": { + "dim_out": 2560, + "image_token_size": 8, + "image_input_size": 224, + }, + "depth": { + "model_type": "MoRGBD", + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "token_size": 16, + "dim_out": 1024, + "input_size": 224, + "use_future_depth": True, + "block_future_depth_to_action": True, + "future_depth_head_type": "resampler", + "detach_future_image_feats": True, + }, + "video": { + "attention_mode": "flex_block_causal", + "input_size": 256, + "block_suffix_to_future_video": True, + "share_future_depth_query": True, + "use_shared_future_task_proj": True, + "use_current_shared_task_proj": True, + "num_future_frames": 1, + "use_warmup_frame": True, + "effective_fps": 1.0, + "n_blocks": 1, + "cls_pool": "last", + "detach_image_feats": True, + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "dim_out": 1024, + "future_video_loss_weight": 0.004, + "use_smooth_l1_loss": False, + "use_mse_loss": True, + "mse_loss_weight": 1.0, + "use_patch_loss": True, + "use_current_patch_loss": True, + "use_cosine_loss": False, + "cosine_loss_weight": 0.2, + "use_cls_loss": False, + "cls_loss_type": "mse", + "cls_loss_weight": 0.2, + }, + }, +} + + +def resolve_lingbot_vla_v2_checkpoint(model_path: str | Path) -> Path: + path = Path(model_path).expanduser().resolve() + if path.is_file(): + if path.name != "model.safetensors.index.json": + raise ValueError(f"Expected model.safetensors.index.json, got: {path}") + return path + if not path.is_dir(): + raise FileNotFoundError(f"LingBot-VLA v2 model path does not exist: {path}") + index_path = path / "model.safetensors.index.json" + if not index_path.is_file(): + raise FileNotFoundError(f"Missing sharded checkpoint index: {index_path}") + return index_path + + +def resolve_lingbot_vla_v2_shards(model_path: str | Path) -> list[str]: + index_path = resolve_lingbot_vla_v2_checkpoint(model_path) + with index_path.open("r", encoding="utf-8") as handle: + index = json.load(handle) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Invalid safetensors index without weight_map: {index_path}") + + shard_paths = [index_path.parent / name for name in sorted(set(weight_map.values()))] + missing = [str(path) for path in shard_paths if not path.is_file()] + if missing: + raise FileNotFoundError(f"Missing LingBot-VLA v2 checkpoint shards: {missing}") + return [str(path) for path in shard_paths] + + +def build_official_6b_config(qwen3vl_path: str | Path): + from telefuser.models.lingbot_vla_v2 import LingbotVLAV2Config + + qwen_path = Path(qwen3vl_path).expanduser().resolve() + qwen_config = AutoConfig.from_pretrained(str(qwen_path), local_files_only=True) + if not hasattr(qwen_config, "text_config") or not hasattr(qwen_config, "vision_config"): + raise ValueError( + "LingBot-VLA v2 requires the local Qwen3-VL-4B-Instruct architecture/tokenizer " + f"directory; this is not a complete Qwen3-VL directory: {qwen_path}" + ) + + text_config = qwen_config.text_config + expected_architecture = {"hidden_size": 2560, "num_hidden_layers": 36} + mismatches = { + key: (expected, getattr(text_config, key, None)) + for key, expected in expected_architecture.items() + if getattr(text_config, key, None) != expected + } + if mismatches: + raise ValueError( + "LingBot-VLA v2 6B was trained with Qwen3-VL-4B-Instruct; " + f"the supplied architecture is incompatible: {mismatches}" + ) + + values = deepcopy(OFFICIAL_6B_MODEL_CONFIG) + values["tokenizer_path"] = str(qwen_path) + config = LingbotVLAV2Config(**values) + for key in ( + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "rms_norm_eps", + "rope_theta", + "vocab_size", + "max_position_embeddings", + "hidden_act", + "tie_word_embeddings", + ): + if hasattr(text_config, key): + setattr(config, key, getattr(text_config, key)) + config.vision_config = qwen_config.vision_config + config.tokenizer_path = str(qwen_path) + config.use_cache = True + config.attention_implementation = "eager" + return config + + +def validate_official_6b_checkpoint(state_dict): + gate = "model.qwenvl_with_expert.qwen_expert.model.layers.0.mlp.experts.gate_proj" + last_gate = "model.qwenvl_with_expert.qwen_expert.model.layers.35.mlp.experts.gate_proj" + expected = (32, 512, 768) + for key in (gate, last_gate): + if key not in state_dict: + raise ValueError(f"Missing official LingBot-VLA v2 weight: {key}") + if tuple(state_dict[key].shape) != expected: + raise ValueError( + f"Unexpected shape for {key}: expected {expected}, got {tuple(state_dict[key].shape)}" + ) + + +class LingBotVlaV2StateDictConverter: + def __init__(self, qwen3vl_path: str | Path): + self.qwen3vl_path = Path(qwen3vl_path) + + def from_official(self, state_dict): + validate_official_6b_checkpoint(state_dict) + config = build_official_6b_config(self.qwen3vl_path) + return state_dict, {"config": config, "eval": True} + + def from_diffusers(self, state_dict): + del state_dict + raise ValueError("LingBot-VLA v2 does not provide a Diffusers checkpoint") + + +def load_lingbot_vla_v2( + module_manager, + model_path: str | Path, + qwen3vl_path: str | Path, + *, + torch_dtype=torch.bfloat16, + device=None, +): + from telefuser.models.lingbot_vla_v2 import LingBotVlaV2Model + + shard_paths = resolve_lingbot_vla_v2_shards(model_path) + module_manager.load_model( + shard_paths, + device=device, + torch_dtype=torch_dtype, + low_cpu_mem_usage=True, + name="lingbot_vla_v2", + model_class=LingBotVlaV2Model, + model_resource="official", + converter_kwargs={"qwen3vl_path": str(qwen3vl_path)}, + strict=True, + ) + return module_manager.fetch_module("lingbot_vla_v2") diff --git a/telefuser/models/lingbot_vla_v2_moe.py b/telefuser/models/lingbot_vla_v2_moe.py new file mode 100644 index 00000000..82ea43af --- /dev/null +++ b/telefuser/models/lingbot_vla_v2_moe.py @@ -0,0 +1,829 @@ +"""Native fused-MoE action expert used by LingBot-VLA v2. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _zero_i32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.int32), mask=offs < N) + + +@triton.jit +def _zero_fp32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.float32), mask=offs < N) + + +@triton.jit +def _moe_pack_selected_kernel( + selected_ptr, + route_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + T: tl.constexpr, + TOPK: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_K: tl.constexpr, +): + row = tl.program_id(0) + slots = tl.arange(0, BLOCK_K) + mask = slots < TOPK + experts = tl.load(selected_ptr + row * TOPK + slots, mask=mask, other=0).to(tl.int32) + routes = tl.load(route_ptr + row * TOPK + slots, mask=mask, other=0.0).to(tl.float32) + pos = tl.atomic_add(counts_ptr + experts, 1, sem="relaxed", mask=mask) + store_mask = mask & (pos < MAX_ROUTES) + tl.store(rows_ptr + experts * MAX_ROUTES + pos, row, mask=store_mask) + tl.store(slots_ptr + experts * MAX_ROUTES + pos, slots, mask=store_mask) + + +@triton.jit +def _moe_gate_up_grouped_kernel( + x_ptr, + gate_ptr, + up_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + route_ptr, + inter_ptr, + T: tl.constexpr, + D: tl.constexpr, + E: tl.constexpr, + TOPK: tl.constexpr, + I: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_I: tl.constexpr, + BLOCK_D: tl.constexpr, +): + expert = tl.program_id(0) + bid_m = tl.program_id(1) + bid_i = tl.program_id(2) + count = tl.load(counts_ptr + expert).to(tl.int32) + start_m = bid_m * BLOCK_M + if start_m >= count: + return + route_idx = start_m + tl.arange(0, BLOCK_M) + offs_i = bid_i * BLOCK_I + tl.arange(0, BLOCK_I) + offs_d = tl.arange(0, BLOCK_D) + valid_m = route_idx < count + rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + acc_g = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) + acc_u = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) + for d0 in range(0, D, BLOCK_D): + ds = d0 + offs_d + x = tl.load( + x_ptr + rows[:, None] * D + ds[None, :], + mask=valid_m[:, None] & (ds[None, :] < D), + other=0.0, + ) + gw = tl.load( + gate_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], + mask=(offs_i[None, :] < I) & (ds[:, None] < D), + other=0.0, + ) + uw = tl.load( + up_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], + mask=(offs_i[None, :] < I) & (ds[:, None] < D), + other=0.0, + ) + acc_g += tl.dot(x, gw) + acc_u += tl.dot(x, uw) + route = tl.load(route_ptr + rows * TOPK + slots, mask=valid_m, other=0.0).to(tl.float32) + silu = acc_g * (1.0 / (1.0 + tl.exp(-acc_g))) + val = silu * acc_u * route[:, None] + tl.store( + inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + offs_i[None, :]), + val.to(inter_ptr.dtype.element_ty), + mask=valid_m[:, None] & (offs_i[None, :] < I), + ) + + +@triton.jit +def _moe_down_grouped_kernel( + inter_ptr, + down_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + out_ptr, + T: tl.constexpr, + D: tl.constexpr, + E: tl.constexpr, + TOPK: tl.constexpr, + I: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_I: tl.constexpr, +): + expert = tl.program_id(0) + bid_m = tl.program_id(1) + bid_d = tl.program_id(2) + count = tl.load(counts_ptr + expert).to(tl.int32) + start_m = bid_m * BLOCK_M + if start_m >= count: + return + route_idx = start_m + tl.arange(0, BLOCK_M) + offs_d = bid_d * BLOCK_D + tl.arange(0, BLOCK_D) + offs_i = tl.arange(0, BLOCK_I) + valid_m = route_idx < count + rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + acc = tl.zeros((BLOCK_M, BLOCK_D), dtype=tl.float32) + for i0 in range(0, I, BLOCK_I): + is_ = i0 + offs_i + x = tl.load( + inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + is_[None, :]), + mask=valid_m[:, None] & (is_[None, :] < I), + other=0.0, + ) + w = tl.load( + down_ptr + (expert * D + offs_d[None, :]) * I + is_[:, None], + mask=(offs_d[None, :] < D) & (is_[:, None] < I), + other=0.0, + ) + acc += tl.dot(x, w) + tl.atomic_add( + out_ptr + rows[:, None] * D + offs_d[None, :], + acc, + sem="relaxed", + mask=valid_m[:, None] & (offs_d[None, :] < D), + ) + + +def robby_moe_forward( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + workspace: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Inference-only grouped MoE path migrated from robbyvla_infer _moe.""" + if hidden_states.ndim != 2: + raise ValueError(f"hidden_states must be 2D, got {tuple(hidden_states.shape)}") + if selected_experts.ndim != 2 or routing_weights.ndim != 2: + raise ValueError("selected_experts and routing_weights must be 2D") + if not hidden_states.is_cuda: + raise ValueError("robby_moe_forward requires CUDA tensors") + + T, D = hidden_states.shape + E, I, weight_d = gate_weight.shape + top_k = selected_experts.shape[1] + if weight_d != D or up_weight.shape != gate_weight.shape or down_weight.shape != (E, D, I): + raise ValueError( + "Unexpected MoE weight shapes: " + f"hidden={tuple(hidden_states.shape)} gate={tuple(gate_weight.shape)} " + f"up={tuple(up_weight.shape)} down={tuple(down_weight.shape)}" + ) + + max_routes = T * top_k + if workspace is None: + counts = torch.empty((E,), device=hidden_states.device, dtype=torch.int32) + rows = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) + slots = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) + inter = torch.empty((T, top_k, I), device=hidden_states.device, dtype=hidden_states.dtype) + out = torch.empty((T, D), device=hidden_states.device, dtype=torch.float32) + else: + counts = workspace["counts"] + rows = workspace["rows"] + slots = workspace["slots"] + inter = workspace["inter"] + out = workspace["out"] + + selected_i32 = selected_experts.to(torch.int32).contiguous() + route = routing_weights.contiguous() + + _zero_i32_kernel[(1,)](counts, E, BLOCK=triton.next_power_of_2(E), num_warps=1) + _moe_pack_selected_kernel[(T,)]( + selected_i32, + route, + counts, + rows, + slots, + T, + top_k, + max_routes, + BLOCK_K=triton.next_power_of_2(top_k), + num_warps=1, + ) + _moe_gate_up_grouped_kernel[ + (E, triton.cdiv(max_routes, 16), triton.cdiv(I, 32)) + ]( + hidden_states, + gate_weight, + up_weight, + counts, + rows, + slots, + route, + inter, + T, + D, + E, + top_k, + I, + max_routes, + BLOCK_M=16, + BLOCK_I=32, + BLOCK_D=64, + num_warps=4, + ) + _zero_fp32_kernel[(triton.cdiv(out.numel(), 1024),)]( + out, + out.numel(), + BLOCK=1024, + num_warps=4, + ) + _moe_down_grouped_kernel[ + (E, triton.cdiv(max_routes, 16), triton.cdiv(D, 64)) + ]( + inter, + down_weight, + counts, + rows, + slots, + out, + T, + D, + E, + top_k, + I, + max_routes, + BLOCK_M=16, + BLOCK_D=64, + BLOCK_I=64, + num_warps=4, + ) + return out.reshape_as(hidden_states) + + + +def fused_moe_forward( + module, + num_experts, + routing_weights, + selected_experts, + hidden_states, + fc1_1_weight, + fc1_2_weight, + fc2_weight, +): + """Single-device PyTorch fallback for the fused 3D expert layout.""" + del module + output = torch.zeros_like(hidden_states) + for expert_id in range(num_experts): + routes = (selected_experts == expert_id).nonzero(as_tuple=False) + if routes.numel() == 0: + continue + token_ids = routes[:, 0] + route_ids = routes[:, 1] + expert_input = hidden_states.index_select(0, token_ids) + gate = torch.nn.functional.linear(expert_input, fc1_1_weight[expert_id]) + up = torch.nn.functional.linear(expert_input, fc1_2_weight[expert_id]) + intermediate = torch.nn.functional.silu(gate) * up + expert_output = torch.nn.functional.linear(intermediate, fc2_weight[expert_id]) + weights = routing_weights[token_ids, route_ids].unsqueeze(-1) + output.index_add_(0, token_ids, expert_output * weights) + return output + + + +from logging import raiseExceptions +import einops +import numpy as np +import torch +from torch import nn +import torch.nn.functional as F +from torch import Tensor, nn +from typing import List, Optional, Tuple +from transformers import AutoTokenizer +from dataclasses import dataclass +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.cache_utils import Cache, SlidingWindowCache, StaticCache, DynamicCache +from transformers.generation import GenerationMixin +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from transformers.utils import ( + ModelOutput, + is_torchdynamo_compiling, + logging, + can_return_tuple, + auto_docstring +) +from transformers.utils.deprecation import deprecate_kwarg +from transformers.activations import ACT2FN +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, is_flash_attn_available +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.processing_utils import Unpack + +try: + from dinov3.hub.backbones import ( + dinov3_vits16, + dinov3_vits16plus, + dinov3_vitb16, + ) +except: pass +def _update_moe_runtime_stats(block, routing_weights, selected_experts): + """Update MoE runtime buffers outside torch.compile graphs.""" + with torch.no_grad(): + if routing_weights is not None and hasattr(block, 'avg_topk_sigmoid_score'): + avg_score = routing_weights.detach().float().mean() + block.avg_topk_sigmoid_score.copy_( + avg_score.reshape_as(block.avg_topk_sigmoid_score).to( + device=block.avg_topk_sigmoid_score.device, + dtype=block.avg_topk_sigmoid_score.dtype, + ) + ) + + if hasattr(block, 'tokens_per_expert'): + counts = F.one_hot( + selected_experts.detach().reshape(-1), + num_classes=block.num_experts, + ).sum(dim=0) + block.tokens_per_expert.add_( + counts.to( + device=block.tokens_per_expert.device, + dtype=block.tokens_per_expert.dtype, + ) + ) + + +import transformers.models.qwen2.modeling_qwen2 as hf_qwen2 +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2MLP, + rotate_half, + apply_rotary_pos_emb, + repeat_kv, + eager_attention_forward, + Qwen2Attention, + Qwen2RMSNorm, + Qwen2RotaryEmbedding, + PreTrainedModel, +) + +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2Model as _Qwen2Model, + Qwen2ForCausalLM as _Qwen2ForCausalLM, +) +logger = logging.get_logger(__name__) +# from transformers.models.mistral.modeling_mistral import MistralMLP + +# Modified from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2Moe +class Qwen2MoeRoutedExpertMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class Qwen2MoeSharedExpertMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class Qwen2FusedExperts(nn.Module): + """Fused expert module: stores E experts' weights as 3D tensors for group_gemm. + + Shape convention matches nn.Linear(in, out).weight = [out, in]: + gate_proj: [E, intermediate_size, hidden_size] + up_proj: [E, intermediate_size, hidden_size] + down_proj: [E, hidden_size, intermediate_size] + + The forward() method runs the full fused_moe computation. This is critical + for FSDP2: calling self.experts(...) triggers FSDP2's forward pre-hook to + unshard the expert params on ep_fsdp_mesh BEFORE they are used by kernels. + """ + def __init__(self, num_experts, hidden_size, intermediate_size, initializer_range=0.02): + super().__init__() + self.num_experts = num_experts + self.intermediate_size = intermediate_size + self.initializer_range = initializer_range + self.gate_proj = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.up_proj = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.down_proj = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) + self.register_buffer("_gate_up_proj_cache", None, persistent=False) + self._gate_up_proj_cache_key = None + self._robby_moe_workspace = None + self._robby_moe_workspace_key = None + self.reset_parameters() + + def reset_parameters(self): + nn.init.normal_(self.gate_proj, mean=0.0, std=self.initializer_range) + nn.init.normal_(self.up_proj, mean=0.0, std=self.initializer_range) + nn.init.normal_(self.down_proj, mean=0.0, std=self.initializer_range) + self.clear_inference_cache() + + def clear_inference_cache(self): + self._gate_up_proj_cache = None + self._gate_up_proj_cache_key = None + self._robby_moe_workspace = None + self._robby_moe_workspace_key = None + + def _get_robby_moe_workspace(self, hidden_states, top_k): + if self.training or torch.is_grad_enabled() or not hidden_states.is_cuda: + return None + num_tokens, hidden_size = hidden_states.shape + key = ( + num_tokens, + int(top_k), + self.num_experts, + hidden_size, + self.intermediate_size, + hidden_states.dtype, + hidden_states.device, + ) + if self._robby_moe_workspace is None or self._robby_moe_workspace_key != key: + max_routes = num_tokens * int(top_k) + self._robby_moe_workspace = { + "counts": torch.empty((self.num_experts,), device=hidden_states.device, dtype=torch.int32), + "rows": torch.empty((self.num_experts, max_routes), device=hidden_states.device, dtype=torch.int32), + "slots": torch.empty((self.num_experts, max_routes), device=hidden_states.device, dtype=torch.int32), + "inter": torch.empty( + (num_tokens, int(top_k), self.intermediate_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ), + "out": torch.empty((num_tokens, hidden_size), device=hidden_states.device, dtype=torch.float32), + } + self._robby_moe_workspace_key = key + return self._robby_moe_workspace + + def forward(self, module, num_experts, routing_weights, selected_experts, hidden_states): + """Run fused_moe_forward with FSDP2-managed weights. + + Must be called via self.experts(...) so FSDP2 unshards params first. + """ + return fused_moe_forward( + module=module, + num_experts=num_experts, + routing_weights=routing_weights, + selected_experts=selected_experts, + hidden_states=hidden_states, + fc1_1_weight=self.gate_proj, + fc1_2_weight=self.up_proj, + fc2_weight=self.down_proj, + ) + + +class FixQwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + FixQwen2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + # print(f'self.weight dtype is {self.weight.dtype}') + input_dtype = hidden_states.dtype + # print(f'input_dtype is {input_dtype}') + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + # print(f'hidden_states dtype is {hidden_states.dtype}') + # print(f'output dtype is {(self.weight * hidden_states.to(input_dtype)).dtype}') + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen2TokenMoeBlock(nn.Module): + """Token-level routing MoE block with all-to-all computation for torch.compile compatibility.""" + def __init__(self, config): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.norm_topk_prob = config.norm_topk_prob + + # Loss-free balancing support. With zero correction bias this is + # equivalent to unbiased top-k selection; the optimizer pre-hook updates + # the bias when bias_update_speed > 0. + self.register_buffer( + "e_score_correction_bias", torch.zeros(config.num_experts), + persistent=True, + ) + self.register_buffer( + "tokens_per_expert", torch.zeros(config.num_experts, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "last_tokens_per_expert", torch.zeros(config.num_experts, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "avg_topk_sigmoid_score", torch.zeros(1, dtype=torch.float32), + persistent=False, + ) + + # gating (per-token) + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + + # EP/fused support: choose expert storage based on moe_implementation + self._moe_implementation = getattr(config, '_moe_implementation', None) or 'eager' + if self._moe_implementation == 'fused': + self.experts = Qwen2FusedExperts( + self.num_experts, + config.hidden_size, + config.moe_intermediate_size, + initializer_range=getattr(config, "initializer_range", 0.02), + ) + else: + self.experts = nn.ModuleList( + [Qwen2MoeRoutedExpertMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)] + ) + + self.shared_expert = Qwen2MoeSharedExpertMLP(config, intermediate_size=config.shared_expert_intermediate_size) + self._router_activation = getattr(config, 'router_activation', 'softmax') + self.routed_scaling_factor = getattr(config, 'routed_scaling_factor', 1.0) + self._use_shared_expert_gate = getattr(config, 'use_shared_expert_gate', True) + if self._use_shared_expert_gate: + self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Token-level routing with all-to-all computation for torch.compile compatibility.""" + batch_size, sequence_length, hidden_dim = hidden_states.shape + num_tokens = batch_size * sequence_length + + # Token-level routing: each token individually + hidden_flat = hidden_states.reshape(-1, hidden_dim) # (B*T, D) + # Gate in true fp32 (autocast disabled): bf16 gate logits can flip top-k + # selection on near-equal scores -> routing jitter / rotating dead experts. + # cf. VideoPretrain lumos/moe/router.py TokenChoiceTopKRouter. + with torch.amp.autocast(hidden_flat.device.type, enabled=False): + router_logits = F.linear(hidden_flat.float(), self.gate.weight.float()) # (B*T, num_experts) + + if self._router_activation == 'sigmoid': + routing_scores = router_logits.sigmoid() + else: + routing_scores = F.softmax(router_logits, dim=1, dtype=torch.float) + + scores_for_choice = routing_scores + self.e_score_correction_bias.unsqueeze(0) + _, selected_experts = torch.topk(scores_for_choice, self.top_k, dim=-1) + routing_weights = routing_scores.gather(1, selected_experts) + if self.training: + _update_moe_runtime_stats(self, routing_weights, selected_experts) + if self.norm_topk_prob: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-20) + if self.routed_scaling_factor != 1.0: + routing_weights = routing_weights * self.routed_scaling_factor + routing_weights = routing_weights.to(hidden_states.dtype) + + # Expert computation: fused (group_gemm) or eager (per-expert loop) + if self._moe_implementation == 'fused': + use_robby_moe = ( + robby_moe_forward is not None + and hidden_flat.is_cuda + and not self.training + and not torch.is_grad_enabled() + ) + if use_robby_moe: + try: + final_hidden_states = robby_moe_forward( + hidden_flat, + routing_weights, + selected_experts, + self.experts.gate_proj, + self.experts.up_proj, + self.experts.down_proj, + workspace=self.experts._get_robby_moe_workspace( + hidden_flat, + selected_experts.shape[1], + ), + ) + except Exception as exc: + logger.warning_once(f"robby_moe_forward failed, falling back to fused_moe_forward: {exc}") + final_hidden_states = self.experts( + module=self, + num_experts=self.num_experts, + routing_weights=routing_weights, + selected_experts=selected_experts, + hidden_states=hidden_flat, + ) + else: + final_hidden_states = self.experts( + module=self, + num_experts=self.num_experts, + routing_weights=routing_weights, + selected_experts=selected_experts, + hidden_states=hidden_flat, + ) + else: + # Original eager path: every expert processes all tokens + expert_outputs = torch.stack( + [expert(hidden_flat) for expert in self.experts], dim=0 + ) # (num_experts, B*T, D) + expert_mask = F.one_hot( + selected_experts, num_classes=self.num_experts + ).float() # (B*T, top_k, num_experts) + weights = (expert_mask * routing_weights.unsqueeze(-1).float()).sum(dim=1).to(hidden_states.dtype) # (B*T, num_experts) + final_hidden_states = torch.einsum('ebd,be->bd', expert_outputs, weights) # (B*T, D) + + # Shared expert: applied to all tokens (fixed shape) + if final_hidden_states.dtype != hidden_flat.dtype: + final_hidden_states = final_hidden_states.to(hidden_flat.dtype) + shared_expert_output = self.shared_expert(hidden_flat) + if self._use_shared_expert_gate: + shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_flat)) * shared_expert_output + final_hidden_states = final_hidden_states + shared_expert_output + + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits + + +class Qwen2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen2Attention(config=config, layer_idx=layer_idx) + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + + def forward( + self, + hidden_states: torch.Tensor, + att_output: Optional[torch.Tensor] = None, + start: Optional[int] = 0, + end: Optional[int] = 0, + compute_kqv: bool = False, + output_atten: bool = False, + ada_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + # Ensure input dtypes match weight dtype (needed for gradient checkpointing + # recomputation where autocast context is lost) + param_dtype = self.self_attn.q_proj.weight.dtype + hidden_states = hidden_states.to(param_dtype) + if att_output is not None: + att_output = att_output.to(param_dtype) + if ada_cond is not None: + ada_cond = ada_cond.to(param_dtype) + + if compute_kqv: + if ada_cond is not None: + hidden_states = self.input_layernorm(hidden_states, ada_cond) + else: + hidden_states = self.input_layernorm(hidden_states) + hidden_shape = (*hidden_states.shape[:-1], -1, self.self_attn.head_dim) + + query_state = self.self_attn.q_proj(hidden_states).view(hidden_shape) + key_state = self.self_attn.k_proj(hidden_states).view(hidden_shape) + value_state = self.self_attn.v_proj(hidden_states).view(hidden_shape) + + return query_state, key_state, value_state + + elif output_atten: + if att_output.dtype != self.self_attn.o_proj.weight.dtype: + att_output = att_output.to(self.self_attn.o_proj.weight.dtype) + out_emb = self.self_attn.o_proj(att_output[:, start:end]) + + # first residual + out_emb += hidden_states + after_first_residual = out_emb.clone() + if ada_cond is not None: + out_emb = self.post_attention_layernorm(out_emb, ada_cond) + else: + out_emb = self.post_attention_layernorm(out_emb) + out_emb = self.mlp(out_emb) + # Handle MoE block returning (hidden_states, router_logits) + router_logits = None + if isinstance(out_emb, tuple): + out_emb, router_logits = out_emb + # second residual + out_emb += after_first_residual + + return out_emb, router_logits + + else: + raise ValueError(f"Invaild Operation compute_kqv={compute_kqv} and output_atten={output_atten} with Qwen2DecoderLayer in LingBot-VLA") + + +@auto_docstring +class Qwen2PreTrainedModel(PreTrainedModel): + config: Qwen2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Qwen2DecoderLayer, + "attentions": Qwen2Attention, + } + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Qwen2FusedExperts): + module.initializer_range = std + module.reset_parameters() + + +class Qwen2Model(Qwen2PreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`] + + Args: + config: Qwen2Config + """ + get_input_embeddings = _Qwen2Model.get_input_embeddings + set_input_embeddings = _Qwen2Model.set_input_embeddings + forward = _Qwen2Model.forward + + def __init__(self, config: Qwen2Config, eval=False): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = FixQwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen2RotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + if eval: + self._init_weights = lambda module: None + self.post_init() + + +class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + get_input_embeddings = _Qwen2ForCausalLM.get_input_embeddings + set_input_embeddings = _Qwen2ForCausalLM.set_input_embeddings + get_output_embeddings = _Qwen2ForCausalLM.get_output_embeddings + set_output_embeddings = _Qwen2ForCausalLM.set_output_embeddings + forward = _Qwen2ForCausalLM.forward + set_decoder = _Qwen2ForCausalLM.set_decoder + get_decoder = _Qwen2ForCausalLM.get_decoder + def __init__(self, config, eval): + super().__init__(config) + self.model = Qwen2Model(config, eval) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + +def apply_lingbot_qwen2_patch(): + hf_qwen2.Qwen2DecoderLayer = Qwen2DecoderLayer + hf_qwen2.Qwen2PreTrainedModel = Qwen2PreTrainedModel + hf_qwen2.Qwen2Model = Qwen2Model + hf_qwen2.Qwen2ForCausalLM = Qwen2ForCausalLM + + diff --git a/telefuser/models/lingbot_vla_v2_qwen.py b/telefuser/models/lingbot_vla_v2_qwen.py new file mode 100644 index 00000000..189194d2 --- /dev/null +++ b/telefuser/models/lingbot_vla_v2_qwen.py @@ -0,0 +1,746 @@ +"""Native Qwen vision-language layers used by LingBot-VLA v2. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + +# Qwen2.5-VL compatibility is retained because the shared flow-matching base +# defines the legacy V1 policy alongside the V2 policy. + +import torch +from torch import nn +import torch.nn.functional as F +from torch.nn import CrossEntropyLoss +from torch import Tensor, nn +from typing import List, Optional, Tuple, Union, Callable, Dict, Any +import math +from transformers import ( + PreTrainedModel, +) +from dataclasses import dataclass +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig +from transformers.cache_utils import Cache, SlidingWindowCache, StaticCache, DynamicCache +from transformers.generation import GenerationMixin +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.utils import ( + ModelOutput, + logging, +) +from transformers.activations import ACT2FN +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, flash_attn_supports_top_left_mask, is_flash_attn_available +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.processing_utils import Unpack +import torch.distributed._tensor as dt + +if is_flash_attn_available(): + from flash_attn.layers.rotary import apply_rotary_emb + from flash_attn.flash_attn_interface import flash_attn_varlen_func + from transformers.modeling_flash_attention_utils import _flash_attention_forward +import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as hf_qwen25vl +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2RMSNorm, + Qwen2_5_VLMLP, + Qwen2_5_VLAttention, + Qwen2MLP, + Qwen2_5_VisionTransformerPretrainedModel, + Qwen2_5_VLRotaryEmbedding, + apply_rotary_pos_emb_vision, + eager_attention_forward +) + +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLTextModel as _Qwen2_5_VLTextModel, + Qwen2_5_VLForConditionalGeneration as _Qwen2_5_VLForConditionalGeneration +) +logger = logging.get_logger(__name__) + + +class Qwen2_5_VLVisionAttention(nn.Module): + def __init__(self, config: Qwen2_5_VLVisionConfig) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 # needed for eager attention + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + self.config = config + self.attention_dropout = 0.0 + self.is_causal = False + # print(f"ViT Attention Type is {self.config._attn_implementation}") + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + if self.config._attn_implementation == "flash_attention_2": + # Flash Attention 2: Use cu_seqlens for variable length attention + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + out_fp32_atten = False + if key_states.dtype == torch.float32: + out_fp32_atten = True + query_states, key_states, value_states = query_states.to(torch.bfloat16), key_states.to(torch.bfloat16), value_states.to(torch.bfloat16) + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + **kwargs, + ) + if out_fp32_atten: + attn_output = attn_output.to(torch.float32) + else: + # Other implementations: Process each chunk separately + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) + ] + + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + is_causal=False, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class Qwen2_5_VLVisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "flash_attention_2") -> None: + super().__init__() + self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) + self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen2_5_VLVisionAttention(config=config) + self.mlp = Qwen2_5_VLMLP(config, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class Qwen2_5_VLPreTrainedModel(PreTrainedModel): + config_class = Qwen2_5_VLConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + _supports_static_cache = False # TODO (joao): fix. torch.compile failing probably due to `cache_positions` + + # def _init_weights(self, module): + # std = self.config.initializer_range + # if isinstance(module, (nn.Linear, nn.Conv3d)): + # module.weight.data.normal_(mean=0.0, std=std) + # if module.bias is not None: + # module.bias.data.zero_() + # elif isinstance(module, nn.Embedding): + # module.weight.data.normal_(mean=0.0, std=std) + # if module.padding_idx is not None: + # module.weight.data[module.padding_idx].zero_() + + +class Qwen2_5_VLDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + self.self_attn = Qwen2_5_VLAttention(config, layer_idx) + + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + if config.norm_qkv: + self.q_layernorm = Qwen2RMSNorm(self.self_attn.head_dim, eps=config.rms_norm_eps) + self.k_layernorm = Qwen2RMSNorm(self.self_attn.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + att_output: Optional[torch.Tensor] = None, + start: Optional[int] = 0, + end: Optional[int] = 0, + compute_kqv: bool = False, + norm_qkv: bool = False, + output_atten: bool = False, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + if compute_kqv: + hidden_states = self.input_layernorm(hidden_states) + hidden_shape = (*hidden_states.shape[:-1], -1, self.self_attn.head_dim) + + query_state = self.self_attn.q_proj(hidden_states).view(hidden_shape) + key_state = self.self_attn.k_proj(hidden_states).view(hidden_shape) + value_state = self.self_attn.v_proj(hidden_states).view(hidden_shape) + + if norm_qkv: + query_state = self.q_layernorm(query_state) + key_state = self.k_layernorm(key_state) + + return query_state, key_state, value_state + + elif output_atten: + if att_output.dtype != self.self_attn.o_proj.weight.dtype: + att_output = att_output.to(self.self_attn.o_proj.weight.dtype) + out_emb = self.self_attn.o_proj(att_output[:, start:end]) + + # first residual + out_emb += hidden_states + after_first_residual = out_emb.clone() + + out_emb = self.post_attention_layernorm(out_emb) + out_emb = self.mlp(out_emb) + + # second residual + out_emb += after_first_residual + + return out_emb + + else: + raise ValueError(f"Invaild Operation compute_kqv={compute_kqv} and output_atten={output_atten} with Qwen2_5_VLDecoderLayer in LingBot-VLA") + + +class Qwen2_5_VLTextModel(Qwen2_5_VLPreTrainedModel): + get_input_embeddings = _Qwen2_5_VLTextModel.get_input_embeddings + set_input_embeddings = _Qwen2_5_VLTextModel.set_input_embeddings + forward = _Qwen2_5_VLTextModel.forward + + def __init__(self, config: Qwen2_5_VLConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self._init_weights = lambda module: None + self.post_init() + + +class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + config_class = Qwen2_5_VLConfig + _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + get_input_embeddings = _Qwen2_5_VLForConditionalGeneration.get_input_embeddings + set_input_embeddings = _Qwen2_5_VLForConditionalGeneration.set_input_embeddings + get_output_embeddings = _Qwen2_5_VLForConditionalGeneration.get_output_embeddings + set_output_embeddings = _Qwen2_5_VLForConditionalGeneration.set_output_embeddings + get_decoder = _Qwen2_5_VLForConditionalGeneration.get_decoder + set_decoder = _Qwen2_5_VLForConditionalGeneration.set_decoder + forward = _Qwen2_5_VLForConditionalGeneration.forward + prepare_inputs_for_generation = _Qwen2_5_VLForConditionalGeneration.prepare_inputs_for_generation + def __init__(self, config): + super().__init__(config) + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) + self.model = Qwen2_5_VLTextModel._from_config(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.rope_deltas = None # cache rope_deltas here + + # Initialize weights and apply final processing + self.post_init() + + + +def qwen25_preprcess_grid_thw(self, grid_thw: torch.Tensor): + rotary_pos_emb = self.rot_pos_emb(grid_thw) + window_index, cu_window_seqlens = self.get_window_index(grid_thw) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, + device=grid_thw.device, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + return rotary_pos_emb, window_index, cu_window_seqlens, cu_seqlens + + +def qwen25_forward_without_grid_thw( + self, + hidden_states: torch.Tensor, + grid_thw: torch.Tensor = None, + rotary_pos_emb = None, + window_index = None, + cu_window_seqlens = None, + cu_seqlens = None, + **kwargs +) -> torch.Tensor: + hidden_states = self.patch_embed(hidden_states) + + if rotary_pos_emb is None or window_index is None or cu_window_seqlens is None or cu_seqlens is None: + rotary_pos_emb, window_index, cu_window_seqlens, cu_seqlens = self.preprcess_grid_thw(grid_thw) + + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + hidden_states = hidden_states[window_index, :, :] + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb[window_index, :, :] + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + for layer_num, blk in enumerate(self.blocks): + if layer_num in self.fullatt_block_indexes: + cu_seqlens_now = cu_seqlens + else: + cu_seqlens_now = cu_window_seqlens + + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens_now, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.merger(hidden_states) + reverse_indices = torch.argsort(window_index) + hidden_states = hidden_states[reverse_indices, :] + + return hidden_states + + +def apply_lingbot_qwen25_vl_patch(): + logger.info_rank0("apply patch") + hf_qwen25vl.Qwen2_5_VLPreTrainedModel = Qwen2_5_VLPreTrainedModel + hf_qwen25vl.Qwen2_5_VLDecoderLayer = Qwen2_5_VLDecoderLayer + hf_qwen25vl.Qwen2_5_VLTextModel = Qwen2_5_VLTextModel + hf_qwen25vl.Qwen2_5_VLForConditionalGeneration = Qwen2_5_VLForConditionalGeneration + hf_qwen25vl.Qwen2_5_VLVisionAttention = Qwen2_5_VLVisionAttention + hf_qwen25vl.Qwen2_5_VLVisionBlock = Qwen2_5_VLVisionBlock + hf_qwen25vl.Qwen2_5_VisionTransformerPretrainedModel.forward = qwen25_forward_without_grid_thw + hf_qwen25vl.Qwen2_5_VisionTransformerPretrainedModel.preprcess_grid_thw = qwen25_preprcess_grid_thw + + +# Qwen3-VL implementation used by LingBot-VLA v2. + +import torch +from torch import nn +import torch.nn.functional as F +from typing import Callable, Optional, Tuple + +from transformers.generation import GenerationMixin +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.utils import logging +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig +import transformers.models.qwen3_vl.modeling_qwen3_vl as hf_qwen3vl +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLForConditionalGeneration as _Qwen3VLForConditionalGeneration, + Qwen3VLModel as _Qwen3VLModel, + Qwen3VLTextModel as _Qwen3VLTextModel, + Qwen3VLPreTrainedModel as _Qwen3VLPreTrainedModel, + Qwen3VLTextAttention, + Qwen3VLTextMLP, + Qwen3VLTextRMSNorm, + Qwen3VLTextRotaryEmbedding, + Qwen3VLVisionModel, + Qwen3VLVisionMLP, + apply_rotary_pos_emb, + apply_rotary_pos_emb_vision, + eager_attention_forward, +) + + +logger = logging.get_logger(__name__) + + +def _qwen3vl_no_init_weights(self, module): + return + +_Qwen3VLPreTrainedModel._init_weights = _qwen3vl_no_init_weights +Qwen3VLPreTrainedModel = _Qwen3VLPreTrainedModel + + +class Qwen3VLVisionAttention(nn.Module): + def __init__(self, config: Qwen3VLVisionConfig) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + self.config = config + self.attention_dropout = 0.0 + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + max_seqlen: Optional[int] = None, + **kwargs, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + if self.config._attn_implementation == "flash_attention_2": + if max_seqlen is None: + max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item()) + out_fp32_atten = False + if key_states.dtype == torch.float32: + out_fp32_atten = True + query_states = query_states.to(torch.bfloat16) + key_states = key_states.to(torch.bfloat16) + value_states = value_states.to(torch.bfloat16) + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + **kwargs, + ) + if out_fp32_atten: + attn_output = attn_output.to(torch.float32) + else: + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) + ] + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + is_causal=False, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class Qwen3VLVisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen3VLVisionAttention(config=config) + self.mlp = Qwen3VLVisionMLP(config=config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class Qwen3VLTextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3VLTextMLP(config) + self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + att_output: Optional[torch.Tensor] = None, + start: Optional[int] = 0, + end: Optional[int] = 0, + compute_kqv: bool = False, + output_atten: bool = False, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + param_dtype = self.self_attn.q_proj.weight.dtype + hidden_states = hidden_states.to(param_dtype) + if att_output is not None: + att_output = att_output.to(param_dtype) + + if compute_kqv: + hidden_states = self.input_layernorm(hidden_states) + hidden_shape = (*hidden_states.shape[:-1], -1, self.self_attn.head_dim) + query_state = self.self_attn.q_norm(self.self_attn.q_proj(hidden_states).view(hidden_shape)) + key_state = self.self_attn.k_norm(self.self_attn.k_proj(hidden_states).view(hidden_shape)) + value_state = self.self_attn.v_proj(hidden_states).view(hidden_shape) + return query_state, key_state, value_state + + if output_atten: + if att_output.dtype != self.self_attn.o_proj.weight.dtype: + att_output = att_output.to(self.self_attn.o_proj.weight.dtype) + out_emb = self.self_attn.o_proj(att_output[:, start:end]) + out_emb += hidden_states + after_first_residual = out_emb.clone() + out_emb = self.post_attention_layernorm(out_emb) + out_emb = self.mlp(out_emb) + out_emb += after_first_residual + return out_emb + + position_embeddings = kwargs.pop("position_embeddings", None) + attention_mask = kwargs.pop("attention_mask", None) + if position_embeddings is not None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + raise ValueError( + f"Invalid operation compute_kqv={compute_kqv} and output_atten={output_atten} " + "with Qwen3VLTextDecoderLayer in LingBot-VLA" + ) + + +class Qwen3VLTextModel(_Qwen3VLTextModel): + def __init__(self, config: Qwen3VLTextConfig): + Qwen3VLPreTrainedModel.__init__(self, config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config) + self.gradient_checkpointing = False + self.post_init() + + +class Qwen3VLModel(_Qwen3VLModel): + def __init__(self, config: Qwen3VLConfig): + Qwen3VLPreTrainedModel.__init__(self, config) + self.visual = Qwen3VLVisionModel._from_config(config.vision_config) + self.language_model = Qwen3VLTextModel._from_config(config.text_config) + self.rope_deltas = None + self.post_init() + + +class Qwen3VLForConditionalGeneration(_Qwen3VLForConditionalGeneration, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + config_class = Qwen3VLConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"] + + def __init__(self, config): + Qwen3VLPreTrainedModel.__init__(self, config) + self.model = Qwen3VLModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + +@torch.compiler.disable +def preprcess_grid_thw(self, grid_thw: torch.Tensor): + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + seq_len = int(torch.prod(grid_thw, dim=1).sum().item()) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + split_sizes = (grid_thw.prod(-1) // self.spatial_merge_size**2).tolist() + max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item()) + return None, position_embeddings, cu_seqlens, split_sizes, max_seqlen + + +def forward_without_grid_thw( + self, + hidden_states: torch.Tensor, + grid_thw: torch.Tensor = None, + pos_embeds: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + **kwargs, +) -> torch.Tensor: + hidden_states = self.patch_embed(hidden_states) + + if pos_embeds is None or position_embeddings is None or cu_seqlens is None or max_seqlen is None: + pos_embeds, position_embeddings, cu_seqlens, _, max_seqlen = self.preprcess_grid_thw(grid_thw) + if pos_embeds is None: + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + + hidden_states = hidden_states + pos_embeds + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len, -1) + + deepstack_feature_lists = [] + for layer_num, blk in enumerate(self.blocks): + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + max_seqlen=max_seqlen, + **kwargs, + ) + if layer_num in self.deepstack_visual_indexes: + deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)]( + hidden_states + ) + deepstack_feature_lists.append(deepstack_feature) + + hidden_states = self.merger(hidden_states) + return hidden_states, deepstack_feature_lists + + +def apply_lingbot_qwen3_vl_patch(): + logger.info_rank0("apply Qwen3-VL Lingbot patch") + hf_qwen3vl.Qwen3VLPreTrainedModel = Qwen3VLPreTrainedModel + hf_qwen3vl.Qwen3VLTextDecoderLayer = Qwen3VLTextDecoderLayer + hf_qwen3vl.Qwen3VLTextModel = Qwen3VLTextModel + hf_qwen3vl.Qwen3VLModel = Qwen3VLModel + hf_qwen3vl.Qwen3VLForConditionalGeneration = Qwen3VLForConditionalGeneration + hf_qwen3vl.Qwen3VLVisionAttention = Qwen3VLVisionAttention + hf_qwen3vl.Qwen3VLVisionBlock = Qwen3VLVisionBlock + hf_qwen3vl.Qwen3VLVisionModel.forward = forward_without_grid_thw + hf_qwen3vl.Qwen3VLVisionModel.preprcess_grid_thw = preprcess_grid_thw + diff --git a/telefuser/pipelines/lingbot_vla_v2/__init__.py b/telefuser/pipelines/lingbot_vla_v2/__init__.py new file mode 100644 index 00000000..3ecd6c76 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/__init__.py @@ -0,0 +1,18 @@ +"""TeleFuser pipeline components for LingBot-VLA v2 action inference.""" + +from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation +from .pipeline import LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig +from .policy import LingBotVlaV2PolicyStage +from .robot_profile import ROBOTWIN_CAMERA_KEYS, LingBotVlaV2ActionChunk, RobotWinProfile + +__all__ = [ + "LingBotVlaV2ActionChunk", + "LingBotVlaV2InputProcessor", + "LingBotVlaV2Inputs", + "LingBotVlaV2Observation", + "LingBotVlaV2Pipeline", + "LingBotVlaV2PipelineConfig", + "LingBotVlaV2PolicyStage", + "ROBOTWIN_CAMERA_KEYS", + "RobotWinProfile", +] diff --git a/telefuser/pipelines/lingbot_vla_v2/data.py b/telefuser/pipelines/lingbot_vla_v2/data.py new file mode 100644 index 00000000..a4055003 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/data.py @@ -0,0 +1,155 @@ +"""RobotWin input preparation for LingBot-VLA v2.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np +import torch +from PIL import Image + +from .robot_profile import ROBOTWIN_CAMERA_KEYS, RobotWinProfile + + +ImageInput = Image.Image | np.ndarray | torch.Tensor | str | Path + + +@dataclass(frozen=True) +class LingBotVlaV2Observation: + """One RobotWin observation accepted by the public SDK.""" + + task: str + state: torch.Tensor | Sequence[float] + images: Mapping[str, ImageInput] + + +@dataclass(frozen=True) +class LingBotVlaV2Inputs: + """Tensor contract consumed by ``LingBotVlaV2PolicyStage``.""" + + images: torch.Tensor + img_masks: torch.Tensor + lang_tokens: torch.Tensor + lang_masks: torch.Tensor + state: torch.Tensor + image_grid_thw: torch.Tensor + + +def _image_to_chw_uint8(image: ImageInput) -> torch.Tensor: + """Convert one RGB image to the format used by the upstream processor.""" + if isinstance(image, (str, Path)): + with Image.open(image) as opened: + image = np.asarray(opened.convert("RGB")) + elif isinstance(image, Image.Image): + image = np.asarray(image.convert("RGB")) + if isinstance(image, np.ndarray): + image = torch.from_numpy(np.asarray(image).copy()) + if not isinstance(image, torch.Tensor): + raise TypeError(f"unsupported image type: {type(image)!r}") + image = image.detach().to(device="cpu") + if image.ndim != 3: + raise ValueError(f"each image must have three dimensions, got {tuple(image.shape)}") + if image.shape[0] == 3: + chw = image + elif image.shape[-1] == 3: + chw = image.permute(2, 0, 1) + else: + raise ValueError(f"each image must have three RGB channels, got {tuple(image.shape)}") + + if chw.dtype == torch.uint8: + return chw.contiguous() + chw = chw.to(dtype=torch.float32) + if not torch.isfinite(chw).all(): + raise ValueError("images must contain only finite values") + if chw.numel() and float(chw.max()) <= 2.0 and float(chw.min()) >= 0.0: + chw = chw * 255.0 + return chw.round().clamp_(0, 255).to(dtype=torch.uint8).contiguous() + + +class LingBotVlaV2InputProcessor: + """Prepare RobotWin images, task text, and canonical state tensors.""" + + def __init__(self, processor: Any, model_config: Any, robot_profile: RobotWinProfile) -> None: + if processor is None or not hasattr(processor, "image_processor") or not hasattr(processor, "tokenizer"): + raise TypeError("LingBot-VLA v2 requires a Qwen3-VL AutoProcessor") + self.processor = processor + self.robot_profile = robot_profile + self.max_state_dim = int(getattr(model_config, "max_state_dim", 55)) + self.tokenizer_max_length = int(getattr(model_config, "tokenizer_max_length", 72)) + if self.max_state_dim != robot_profile.canonical_dim: + raise ValueError( + f"model max_state_dim is {self.max_state_dim}, RobotWin requires {robot_profile.canonical_dim}" + ) + + def _process_images( + self, images: Mapping[str, ImageInput] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + missing = [key for key in ROBOTWIN_CAMERA_KEYS if key not in images] + if missing: + raise ValueError(f"RobotWin observation is missing camera keys: {missing}") + + processed_images: list[torch.Tensor] = [] + grids: list[torch.Tensor] = [] + for key in self.robot_profile.camera_keys: + output = self.processor.image_processor(_image_to_chw_uint8(images[key])) + pixels = output["pixel_values"] if isinstance(output, dict) else output.pixel_values + grid = output.get("image_grid_thw") if isinstance(output, dict) else getattr(output, "image_grid_thw", None) + pixels = torch.as_tensor(pixels) + grid = None if grid is None else torch.as_tensor(grid) + if pixels.ndim == 3 and pixels.shape[0] == 1: + pixels = pixels.squeeze(0) + if pixels.ndim != 2: + raise ValueError( + f"Qwen3-VL image processor must return [patches, features], got {tuple(pixels.shape)}" + ) + if grid is None or grid.numel() < 3: + raise ValueError("Qwen3-VL image processor must return image_grid_thw") + processed_images.append(pixels) + grids.append(grid.reshape(-1, 3)[0].to(dtype=torch.long)) + + first_shape = processed_images[0].shape + if any(image.shape != first_shape for image in processed_images[1:]): + shapes = [tuple(image.shape) for image in processed_images] + raise ValueError(f"all RobotWin cameras must produce equal patch shapes, got {shapes}") + return ( + torch.stack(processed_images, dim=0).unsqueeze(0), + torch.ones(1, len(processed_images), dtype=torch.bool), + torch.stack(grids, dim=0).unsqueeze(0), + ) + + def _process_language(self, task: str) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance(task, str) or not task.strip(): + raise ValueError("task must be a non-empty string") + tokenizer = self.processor.tokenizer + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": task}], + tokenize=False, + add_generation_prompt=False, + ) + tokens = tokenizer( + [rendered], + padding="max_length", + padding_side="right", + truncation=True, + max_length=self.tokenizer_max_length, + return_tensors="pt", + ) + return tokens["input_ids"], tokens["attention_mask"].to(dtype=torch.bool) + + def prepare(self, observation: LingBotVlaV2Observation) -> LingBotVlaV2Inputs: + """Prepare one RobotWin observation for model inference.""" + if not isinstance(observation, LingBotVlaV2Observation): + raise TypeError("observation must be a LingBotVlaV2Observation") + image_tensors, image_masks, image_grid_thw = self._process_images(observation.images) + state = self.robot_profile.normalize_state(observation.state).unsqueeze(0) + lang_tokens, lang_masks = self._process_language(observation.task) + return LingBotVlaV2Inputs( + images=image_tensors, + img_masks=image_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + state=state, + image_grid_thw=image_grid_thw, + ) diff --git a/telefuser/pipelines/lingbot_vla_v2/pipeline.py b/telefuser/pipelines/lingbot_vla_v2/pipeline.py new file mode 100644 index 00000000..7d184e17 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/pipeline.py @@ -0,0 +1,67 @@ +"""BasePipeline integration for LingBot-VLA v2 RobotWin inference.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager + +from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation +from .policy import LingBotVlaV2PolicyStage +from .robot_profile import LingBotVlaV2ActionChunk, RobotWinProfile + + +@dataclass +class LingBotVlaV2PipelineConfig: + """Runtime configuration for one LingBot-VLA v2 pipeline replica.""" + + policy_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + robot_profile: RobotWinProfile = field(default_factory=RobotWinProfile.default) + include_canonical_actions: bool = False + enable_metrics: bool = False + + +class LingBotVlaV2Pipeline(BasePipeline): + """Single-replica LingBot-VLA v2 structured action SDK.""" + + def _get_stages(self) -> list: + return [self.policy_stage] + + def init(self, module_manager: ModuleManager, config: LingBotVlaV2PipelineConfig) -> None: + self._model_info = module_manager.get_model_info() + self.config = config + policy = module_manager.fetch_module("lingbot_vla_v2") + processor = module_manager.fetch_module("lingbot_vla_v2_processor") + if policy is None or processor is None: + raise RuntimeError("LingBot-VLA v2 requires policy and lingbot_vla_v2_processor modules") + self.input_processor = LingBotVlaV2InputProcessor(processor, policy.config, config.robot_profile) + self.policy_stage = LingBotVlaV2PolicyStage("policy", module_manager, config.policy_config) + if config.enable_metrics: + self.enable_metrics() + + @torch.inference_mode() + def predict(self, inputs: LingBotVlaV2Inputs, seed: int | None = None) -> LingBotVlaV2ActionChunk: + """Run prepared tensors and convert the canonical result to RobotWin fields.""" + actions = self.policy_stage.process(inputs, seed=seed) + return self.config.robot_profile.structure_actions( + actions, + include_canonical=self.config.include_canonical_actions, + ) + + @torch.inference_mode() + def __call__( + self, + observation: LingBotVlaV2Observation, + seed: int | None = None, + ) -> LingBotVlaV2ActionChunk: + """Predict one structured RobotWin action chunk.""" + return self.predict(self.input_processor.prepare(observation), seed=seed) + + def close(self) -> None: + """Release policy device memory.""" + if hasattr(self, "policy_stage"): + self.policy_stage.offload_models() diff --git a/telefuser/pipelines/lingbot_vla_v2/policy.py b/telefuser/pipelines/lingbot_vla_v2/policy.py new file mode 100644 index 00000000..51a5b09e --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/policy.py @@ -0,0 +1,73 @@ +"""TeleFuser stage for LingBot-VLA v2 flow-matching action inference.""" + +from __future__ import annotations + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.metrics import with_metrics + +from .data import LingBotVlaV2Inputs + + +class LingBotVlaV2PolicyStage(BaseStage): + """Run Qwen3-VL prefix encoding and the complete 10-step action sampler.""" + + def __init__(self, name: str, module_manager: ModuleManager, runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, runtime_config) + self.policy = module_manager.fetch_module("lingbot_vla_v2") + if self.policy is None: + raise RuntimeError("ModuleManager does not contain 'lingbot_vla_v2'") + self.model_names = ["policy"] + self._validate_parallelism() + + def _validate_parallelism(self) -> None: + parallel_config = self.model_runtime_config.parallel_config + if getattr(parallel_config, "world_size", 1) != 1: + raise ValueError("LingBot-VLA v2 currently supports one GPU per pipeline replica") + + @with_model_offload(["policy"]) + @torch.inference_mode() + @with_metrics + def process(self, inputs: LingBotVlaV2Inputs, seed: int | None = None) -> torch.Tensor: + """Return a CPU float32 normalized action chunk with shape ``[1, H, 55]``.""" + device = self.device + dtype = self.torch_dtype + tensors = { + "images": inputs.images.to(device=device, dtype=dtype), + "img_masks": inputs.img_masks.to(device=device), + "lang_tokens": inputs.lang_tokens.to(device=device), + "lang_masks": inputs.lang_masks.to(device=device), + "state": inputs.state.to(device=device, dtype=dtype), + "image_grid_thw": inputs.image_grid_thw.to(device=device, dtype=torch.long), + } + noise = None + if seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + config = self.policy.config + noise = torch.randn( + tensors["state"].shape[0], + int(config.n_action_steps), + int(config.max_action_dim), + device=device, + dtype=dtype, + generator=generator, + ) + actions = self.policy.sample_actions(**tensors, noise=noise) + if not isinstance(actions, torch.Tensor) or actions.ndim != 3: + raise RuntimeError(f"LingBot-VLA v2 policy returned an invalid action tensor: {type(actions)!r}") + config = self.policy.config + expected_shape = ( + tensors["state"].shape[0], + int(config.n_action_steps), + int(config.max_action_dim), + ) + if tuple(actions.shape) != expected_shape: + raise RuntimeError( + f"LingBot-VLA v2 policy returned shape {tuple(actions.shape)}, expected {expected_shape}" + ) + if not torch.isfinite(actions).all(): + raise RuntimeError("LingBot-VLA v2 policy returned non-finite actions") + return actions.detach().to(device="cpu", dtype=torch.float32) diff --git a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py new file mode 100644 index 00000000..44f98100 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py @@ -0,0 +1,170 @@ +"""RobotWin feature mapping for LingBot-VLA v2 inference.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Mapping, Sequence + +import torch + + +ROBOTWIN_CAMERA_KEYS = ( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", +) +ROBOTWIN_STATE_DIM = 14 +CANONICAL_DIM = 55 +ARM_SLICE = slice(0, 12) +EFFECTOR_SLICE = slice(28, 30) + + +@dataclass(frozen=True) +class LingBotVlaV2ActionChunk: + """Structured RobotWin action chunk returned by the SDK.""" + + fields: Mapping[str, torch.Tensor] + raw_actions: torch.Tensor + action_mask: torch.Tensor + horizon: int + robot_profile: str = "robotwin" + policy_verified: bool = False + verification_status: str = "unverified_official_6b_base" + canonical_normalized_actions: torch.Tensor | None = None + + +class RobotWinProfile: + """Map RobotWin observations and actions to LingBot's canonical space.""" + + name = "robotwin" + camera_keys = ROBOTWIN_CAMERA_KEYS + canonical_dim = CANONICAL_DIM + raw_state_dim = ROBOTWIN_STATE_DIM + _REQUIRED_STATS = ( + "observation.state.arm.position", + "observation.state.effector.position", + "action.arm.position", + "action.effector.position", + ) + + def __init__(self, norm_stats: Mapping[str, Mapping[str, object]]) -> None: + self._stats = { + key: { + stat_name: torch.as_tensor(stat_value, dtype=torch.float32) + for stat_name, stat_value in values.items() + } + for key, values in norm_stats.items() + } + self._validate_stats() + + @classmethod + def from_json(cls, path: str | Path) -> "RobotWinProfile": + """Load RobotWin normalization statistics from an upstream-format JSON file.""" + payload = json.loads(Path(path).read_text(encoding="utf-8")) + norm_stats = payload.get("norm_stats") + if not isinstance(norm_stats, dict): + raise ValueError("RobotWin normalization file must contain a norm_stats object") + return cls(norm_stats) + + @classmethod + def default(cls) -> "RobotWinProfile": + """Load the RobotWin statistics bundled with TeleFuser.""" + path = Path(__file__).with_name("assets") / "robotwin_norm_stats.json" + return cls.from_json(path) + + @property + def action_mask(self) -> torch.Tensor: + """Return the canonical dimensions used by RobotWin actions.""" + mask = torch.zeros(self.canonical_dim, dtype=torch.bool) + mask[ARM_SLICE] = True + mask[EFFECTOR_SLICE] = True + return mask + + def normalize_state(self, raw_state: torch.Tensor | Sequence[float]) -> torch.Tensor: + """Convert one raw 14-D RobotWin state to normalized canonical 55-D space.""" + state = torch.as_tensor(raw_state, dtype=torch.float32, device="cpu") + if state.shape != (self.raw_state_dim,): + raise ValueError(f"RobotWin state must have shape ({self.raw_state_dim},), got {tuple(state.shape)}") + if not torch.isfinite(state).all(): + raise ValueError("RobotWin state must contain only finite values") + + arm = torch.cat((state[0:6], state[7:13])) + effector = state[[6, 13]] + canonical = torch.zeros(self.canonical_dim, dtype=torch.float32) + canonical[ARM_SLICE] = self._normalize("observation.state.arm.position", arm) + canonical[EFFECTOR_SLICE] = self._normalize("observation.state.effector.position", effector) + return canonical + + def structure_actions( + self, + canonical_normalized_actions: torch.Tensor, + *, + include_canonical: bool = False, + ) -> LingBotVlaV2ActionChunk: + """Convert a normalized canonical action chunk to RobotWin action fields.""" + actions = torch.as_tensor(canonical_normalized_actions, dtype=torch.float32, device="cpu") + if actions.ndim == 3: + if actions.shape[0] != 1: + raise ValueError("RobotWin structured output currently supports a single observation") + actions = actions[0] + if actions.ndim != 2 or actions.shape[-1] != self.canonical_dim: + raise ValueError( + f"canonical actions must have shape [H,{self.canonical_dim}] or [1,H,{self.canonical_dim}], " + f"got {tuple(actions.shape)}" + ) + if not torch.isfinite(actions).all(): + raise ValueError("canonical actions must contain only finite values") + + arm = self._unnormalize("action.arm.position", actions[:, ARM_SLICE]) + effector = self._unnormalize("action.effector.position", actions[:, EFFECTOR_SLICE]) + raw = torch.empty(actions.shape[0], self.raw_state_dim, dtype=torch.float32) + raw[:, 0:6] = arm[:, 0:6] + raw[:, 6] = effector[:, 0] + raw[:, 7:13] = arm[:, 6:12] + raw[:, 13] = effector[:, 1] + fields = MappingProxyType( + { + "action.arm.position": arm, + "action.effector.position": effector, + "action": raw, + } + ) + return LingBotVlaV2ActionChunk( + fields=fields, + raw_actions=raw, + action_mask=self.action_mask, + horizon=int(actions.shape[0]), + canonical_normalized_actions=actions.clone() if include_canonical else None, + ) + + def _validate_stats(self) -> None: + expected_dims = { + "observation.state.arm.position": 12, + "observation.state.effector.position": 2, + "action.arm.position": 12, + "action.effector.position": 2, + } + missing = [key for key in self._REQUIRED_STATS if key not in self._stats] + if missing: + raise ValueError(f"RobotWin normalization statistics are missing keys: {missing}") + for key, expected_dim in expected_dims.items(): + values = self._stats[key] + for stat_name in ("q01", "q99"): + if stat_name not in values or values[stat_name].shape != (expected_dim,): + shape = None if stat_name not in values else tuple(values[stat_name].shape) + raise ValueError( + f"RobotWin statistic {key}.{stat_name} must have shape ({expected_dim},), got {shape}" + ) + + def _normalize(self, key: str, value: torch.Tensor) -> torch.Tensor: + low = self._stats[key]["q01"] + high = self._stats[key]["q99"] + return (value - low) / (high - low + 1e-6) * 2.0 - 1.0 + + def _unnormalize(self, key: str, value: torch.Tensor) -> torch.Tensor: + low = self._stats[key]["q01"] + high = self._stats[key]["q99"] + return (value + 1.0) / 2.0 * (high - low + 1e-6) + low diff --git a/tests/unit/models/test_lingbot_vla_v2_loader.py b/tests/unit/models/test_lingbot_vla_v2_loader.py new file mode 100644 index 00000000..48f42d76 --- /dev/null +++ b/tests/unit/models/test_lingbot_vla_v2_loader.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from telefuser.models.lingbot_vla_v2_loader import ( + resolve_lingbot_vla_v2_shards, + validate_official_6b_checkpoint, +) + + +def test_resolve_lingbot_vla_v2_shards_uses_index_manifest(tmp_path) -> None: + shard_names = ["model-00002-of-00002.safetensors", "model-00001-of-00002.safetensors"] + for name in shard_names: + (tmp_path / name).write_bytes(b"") + index = { + "weight_map": { + "layer.0": shard_names[0], + "layer.1": shard_names[1], + "layer.2": shard_names[0], + } + } + (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index), encoding="utf-8") + + resolved = resolve_lingbot_vla_v2_shards(tmp_path) + + assert resolved == [str(tmp_path / name) for name in sorted(shard_names)] + + +def test_resolve_lingbot_vla_v2_shards_rejects_missing_files(tmp_path) -> None: + index = {"weight_map": {"layer.0": "missing.safetensors"}} + (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index), encoding="utf-8") + + with pytest.raises(FileNotFoundError, match="checkpoint shards"): + resolve_lingbot_vla_v2_shards(tmp_path) + + +def test_validate_official_6b_checkpoint_accepts_expected_gate_shapes() -> None: + prefix = "model.qwenvl_with_expert.qwen_expert.model.layers" + state_dict = { + f"{prefix}.0.mlp.experts.gate_proj": SimpleNamespace(shape=(32, 512, 768)), + f"{prefix}.35.mlp.experts.gate_proj": SimpleNamespace(shape=(32, 512, 768)), + } + + validate_official_6b_checkpoint(state_dict) + + +def test_validate_official_6b_checkpoint_rejects_wrong_shape() -> None: + prefix = "model.qwenvl_with_expert.qwen_expert.model.layers" + state_dict = { + f"{prefix}.0.mlp.experts.gate_proj": SimpleNamespace(shape=(1, 2, 3)), + f"{prefix}.35.mlp.experts.gate_proj": SimpleNamespace(shape=(32, 512, 768)), + } + + with pytest.raises(ValueError, match="Unexpected shape"): + validate_official_6b_checkpoint(state_dict) diff --git a/tests/unit/pipelines/lingbot_vla_v2/__init__.py b/tests/unit/pipelines/lingbot_vla_v2/__init__.py new file mode 100644 index 00000000..c7d9de08 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/__init__.py @@ -0,0 +1 @@ +"""LingBot-VLA v2 pipeline unit tests.""" diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_data.py b/tests/unit/pipelines/lingbot_vla_v2/test_data.py new file mode 100644 index 00000000..ce8dc7a2 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_data.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from telefuser.pipelines.lingbot_vla_v2.data import LingBotVlaV2InputProcessor, LingBotVlaV2Observation +from telefuser.pipelines.lingbot_vla_v2.robot_profile import ROBOTWIN_CAMERA_KEYS, RobotWinProfile + + +class _ImageProcessor: + def __init__(self) -> None: + self.values: list[int] = [] + + def __call__(self, image: torch.Tensor) -> dict[str, torch.Tensor]: + assert image.dtype == torch.uint8 + assert image.shape == (3, 8, 8) + value = int(image[0, 0, 0]) + self.values.append(value) + return { + "pixel_values": torch.full((4, 6), float(value)), + "image_grid_thw": torch.tensor([[1, 4, 4]]), + } + + +class _Tokenizer: + def __init__(self) -> None: + self.rendered_task: str | None = None + self.padding_side: str | None = None + + def apply_chat_template(self, messages, *, tokenize: bool, add_generation_prompt: bool) -> str: + assert tokenize is False + assert add_generation_prompt is False + self.rendered_task = messages[0]["content"] + return f"chat:{self.rendered_task}" + + def __call__(self, prompts, **kwargs) -> dict[str, torch.Tensor]: + assert prompts == [f"chat:{self.rendered_task}"] + self.padding_side = kwargs["padding_side"] + length = kwargs["max_length"] + return { + "input_ids": torch.arange(length).unsqueeze(0), + "attention_mask": torch.ones(1, length), + } + + +def _processor() -> tuple[LingBotVlaV2InputProcessor, _ImageProcessor, _Tokenizer]: + image_processor = _ImageProcessor() + tokenizer = _Tokenizer() + processor = SimpleNamespace(image_processor=image_processor, tokenizer=tokenizer) + config = SimpleNamespace(max_state_dim=55, tokenizer_max_length=6) + return LingBotVlaV2InputProcessor(processor, config, RobotWinProfile.default()), image_processor, tokenizer + + +def _observation() -> LingBotVlaV2Observation: + images = { + ROBOTWIN_CAMERA_KEYS[0]: np.full((8, 8, 3), 10, dtype=np.uint8), + ROBOTWIN_CAMERA_KEYS[1]: np.full((8, 8, 3), 20, dtype=np.uint8), + ROBOTWIN_CAMERA_KEYS[2]: np.full((8, 8, 3), 30, dtype=np.uint8), + } + return LingBotVlaV2Observation(task="pick up the block", state=[0.0] * 14, images=images) + + +def test_prepare_preserves_robotwin_camera_order_and_tensor_contract() -> None: + processor, image_processor, tokenizer = _processor() + observation = _observation() + + inputs = processor.prepare(observation) + + assert image_processor.values == [10, 20, 30] + assert tokenizer.rendered_task == observation.task + assert tokenizer.padding_side == "right" + assert inputs.images.shape == (1, 3, 4, 6) + assert inputs.img_masks.tolist() == [[True, True, True]] + assert inputs.image_grid_thw.shape == (1, 3, 3) + assert inputs.lang_tokens.shape == (1, 6) + assert inputs.lang_masks.dtype == torch.bool + assert torch.equal(inputs.state, processor.robot_profile.normalize_state(observation.state).unsqueeze(0)) + + +def test_prepare_rejects_a_missing_robotwin_camera() -> None: + processor, _, _ = _processor() + observation = _observation() + images = dict(observation.images) + del images[ROBOTWIN_CAMERA_KEYS[1]] + + with pytest.raises(ValueError, match="missing camera keys"): + processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) + + +def test_prepare_scales_unit_float_images_to_uint8() -> None: + processor, image_processor, _ = _processor() + observation = _observation() + images = dict(observation.images) + images[ROBOTWIN_CAMERA_KEYS[0]] = np.full((3, 8, 8), 0.5, dtype=np.float32) + + processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) + + assert image_processor.values[0] == 128 diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py new file mode 100644 index 00000000..436be6bb --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import torch +from torch import nn + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_vla_v2 import ( + LingBotVlaV2Observation, + LingBotVlaV2Pipeline, + LingBotVlaV2PipelineConfig, +) +from telefuser.pipelines.lingbot_vla_v2.robot_profile import ROBOTWIN_CAMERA_KEYS + + +class _ImageProcessor: + def __call__(self, image: torch.Tensor) -> dict[str, torch.Tensor]: + assert image.shape == (3, 8, 8) + return { + "pixel_values": torch.zeros(4, 6), + "image_grid_thw": torch.tensor([[1, 4, 4]]), + } + + +class _Tokenizer: + def apply_chat_template(self, messages, **kwargs) -> str: + return messages[0]["content"] + + def __call__(self, prompts, **kwargs) -> dict[str, torch.Tensor]: + length = kwargs["max_length"] + return { + "input_ids": torch.zeros(1, length, dtype=torch.long), + "attention_mask": torch.ones(1, length, dtype=torch.long), + } + + +class _Policy(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + max_state_dim=55, + max_action_dim=55, + n_action_steps=4, + tokenizer_max_length=6, + ) + + def sample_actions(self, **inputs) -> torch.Tensor: + assert inputs["state"].shape == (1, 55) + assert inputs["images"].shape == (1, 3, 4, 6) + return torch.zeros(1, self.config.n_action_steps, self.config.max_action_dim, device=self.anchor.device) + + +def test_pipeline_returns_structured_robotwin_action_chunk() -> None: + policy = _Policy() + processor = SimpleNamespace(image_processor=_ImageProcessor(), tokenizer=_Tokenizer()) + manager = ModuleManager(torch_dtype=torch.float32, device="cpu") + manager.add_module(policy, "lingbot_vla_v2") + manager.add_module(processor, "lingbot_vla_v2_processor") + pipeline = LingBotVlaV2Pipeline(device="cpu", torch_dtype=torch.float32) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + include_canonical_actions=True, + ), + ) + observation = LingBotVlaV2Observation( + task="pick up the block", + state=[0.0] * 14, + images={key: np.zeros((8, 8, 3), dtype=np.uint8) for key in ROBOTWIN_CAMERA_KEYS}, + ) + + try: + chunk = pipeline(observation, seed=7) + finally: + pipeline.close() + + assert chunk.horizon == 4 + assert chunk.raw_actions.shape == (4, 14) + assert chunk.fields["action.arm.position"].shape == (4, 12) + assert chunk.fields["action.effector.position"].shape == (4, 2) + assert chunk.canonical_normalized_actions is not None + assert chunk.policy_verified is False diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py new file mode 100644 index 00000000..83d00c2f --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import pytest +import torch + +from telefuser.pipelines.lingbot_vla_v2.robot_profile import RobotWinProfile + + +def _stats() -> dict[str, dict[str, list[float]]]: + return { + "observation.state.arm.position": {"q01": [0.0] * 12, "q99": [2.0] * 12}, + "observation.state.effector.position": {"q01": [-1.0] * 2, "q99": [1.0] * 2}, + "action.arm.position": {"q01": [0.0] * 12, "q99": [2.0] * 12}, + "action.effector.position": {"q01": [-1.0] * 2, "q99": [1.0] * 2}, + } + + +def test_normalize_state_uses_robotwin_joint_order() -> None: + profile = RobotWinProfile(_stats()) + state = torch.arange(14, dtype=torch.float32) / 10.0 + + canonical = profile.normalize_state(state) + + arm = torch.cat((state[0:6], state[7:13])) + effector = state[[6, 13]] + assert canonical.shape == (55,) + assert torch.allclose(canonical[0:12], arm / (2.0 + 1e-6) * 2.0 - 1.0) + assert torch.allclose(canonical[28:30], (effector + 1.0) / (2.0 + 1e-6) * 2.0 - 1.0) + assert torch.count_nonzero(canonical[12:28]) == 0 + assert torch.count_nonzero(canonical[30:]) == 0 + + +def test_structure_actions_reconstructs_raw_robotwin_layout() -> None: + profile = RobotWinProfile(_stats()) + canonical = torch.zeros(1, 3, 55) + + chunk = profile.structure_actions(canonical, include_canonical=True) + + arm = chunk.fields["action.arm.position"] + effector = chunk.fields["action.effector.position"] + assert arm.shape == (3, 12) + assert effector.shape == (3, 2) + assert torch.allclose(arm, torch.full_like(arm, 1.0000005)) + assert torch.allclose(effector, torch.zeros_like(effector), atol=1e-6) + assert torch.equal(chunk.raw_actions[:, 0:6], arm[:, 0:6]) + assert torch.equal(chunk.raw_actions[:, 6], effector[:, 0]) + assert torch.equal(chunk.raw_actions[:, 7:13], arm[:, 6:12]) + assert torch.equal(chunk.raw_actions[:, 13], effector[:, 1]) + assert chunk.horizon == 3 + assert chunk.canonical_normalized_actions is not None + + +def test_action_chunk_is_marked_unverified() -> None: + chunk = RobotWinProfile(_stats()).structure_actions(torch.zeros(2, 55)) + + assert chunk.policy_verified is False + assert chunk.verification_status == "unverified_official_6b_base" + assert chunk.robot_profile == "robotwin" + assert chunk.action_mask.shape == (55,) + assert chunk.action_mask.nonzero().flatten().tolist() == list(range(12)) + [28, 29] + assert chunk.canonical_normalized_actions is None + + +def test_profile_rejects_invalid_state_and_action_shapes() -> None: + profile = RobotWinProfile(_stats()) + + with pytest.raises(ValueError, match="state must have shape"): + profile.normalize_state(torch.zeros(13)) + with pytest.raises(ValueError, match="canonical actions must have shape"): + profile.structure_actions(torch.zeros(2, 54)) From 87960aca0e1acd3173a9c17bc78eedfc63c58505 Mon Sep 17 00:00:00 2001 From: HappyDog060713 Date: Tue, 28 Jul 2026 17:37:35 +0800 Subject: [PATCH 02/15] feat: add LingBot VLA v2 base inference support --- examples/lingbot_vla_v2/README.md | 25 +- .../lingbot_vla_v2_inference.py | 34 +- telefuser/models/lingbot_vla_v2.py | 28 +- telefuser/models/lingbot_vla_v2_loader.py | 42 +- telefuser/models/lingbot_vla_v2_moe.py | 5 +- telefuser/models/lingbot_vla_v2_qwen.py | 5 +- .../pipelines/lingbot_vla_v2/__init__.py | 3 +- telefuser/pipelines/lingbot_vla_v2/data.py | 18 +- .../pipelines/lingbot_vla_v2/pipeline.py | 55 ++- tests/unit/models/test_lingbot_vla_v2.py | 53 +++ .../unit/models/test_lingbot_vla_v2_loader.py | 6 + .../pipelines/lingbot_vla_v2/test_data.py | 19 +- .../pipelines/lingbot_vla_v2/test_pipeline.py | 15 +- .../telefuser/core/base_pipeline.py | 423 ++++++++++++++++++ 14 files changed, 662 insertions(+), 69 deletions(-) create mode 100644 tests/unit/models/test_lingbot_vla_v2.py create mode 100644 tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 582bedc5..bf19e7d3 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -1,8 +1,8 @@ -# LingBot-VLA v2 RobotWin SDK +# LingBot-VLA v2 Base Model SDK -This example loads the official LingBot-VLA v2 6B base checkpoint through TeleFuser and returns a structured -RobotWin action chunk. The current integration verifies the SDK contract without claiming policy quality: every -result is marked `policy_verified=False` and `verification_status="unverified_official_6b_base"`. +This example loads the official LingBot-VLA v2 6B base checkpoint through TeleFuser and returns its normalized +55-dimensional canonical action chunk. The RobotWin profile is used only to prepare the example observation; the +result is not converted to physical RobotWin actions. ## Inputs @@ -15,14 +15,13 @@ LingBot's 55-dimensional canonical state. ## Output -The pipeline returns `LingBotVlaV2ActionChunk` with: +The pipeline returns `LingBotVlaV2CanonicalActionChunk` with: -- `fields["action.arm.position"]`: `[H, 12]`. -- `fields["action.effector.position"]`: `[H, 2]`. -- `raw_actions` and `fields["action"]`: reconstructed `[H, 14]` RobotWin actions. -- `action_mask`: the 55-dimensional canonical RobotWin action mask. +- `canonical_normalized_actions`: `[H, 55]` base-model output. - `horizon`: action chunk length, normally 50 for the official base config. -- `canonical_normalized_actions`: optional `[H, 55]` debugging output. +- `action_dim`: canonical action dimension, normally 55. +- `checkpoint_variant`: `base`. +- `policy_verified=False` and `verification_status="unverified_official_6b_base"`. ## Checkpoints @@ -40,8 +39,8 @@ python examples/lingbot_vla_v2/lingbot_vla_v2_inference.py \ --camera-right-wrist /data/cam_right_wrist.png \ --task "pick up the red block" \ --state-json '[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]' \ - --output action_chunk.npz + --output canonical_action_chunk.npz ``` -The example saves named arrays and verification metadata in an `.npz` file. Do not send the output to a robot until -the official 6B GPU smoke test and policy-level parity validation are complete. +The example saves canonical actions and checkpoint metadata in an `.npz` file. The base output must not be sent to +a robot without an embodiment-specific post-training checkpoint, action mapping, and policy validation. diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py index 0d3dcecb..3484c0d1 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py @@ -1,4 +1,4 @@ -"""Run LingBot-VLA v2 with a RobotWin observation.""" +"""Run the LingBot-VLA v2 base checkpoint with a RobotWin observation adapter.""" from __future__ import annotations @@ -13,10 +13,10 @@ from telefuser.core.module_manager import ModuleManager from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2 from telefuser.pipelines.lingbot_vla_v2 import ( + ROBOTWIN_CAMERA_KEYS, LingBotVlaV2Observation, LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig, - ROBOTWIN_CAMERA_KEYS, ) @@ -24,10 +24,10 @@ def get_pipeline( model_root: str, qwen3vl_root: str, device: str = "cuda", - include_canonical_actions: bool = False, ) -> LingBotVlaV2Pipeline: """Load the official 6B checkpoint and Qwen3-VL processor.""" - dtype = torch.bfloat16 if torch.device(device).type == "cuda" else torch.float32 + target_device = torch.device(device) + dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 processor = AutoProcessor.from_pretrained(qwen3vl_root, local_files_only=True, padding_side="right") manager = ModuleManager(torch_dtype=dtype, device="cpu") manager.add_module(processor, "lingbot_vla_v2_processor", path=qwen3vl_root) @@ -36,8 +36,11 @@ def get_pipeline( pipeline.init( manager, LingBotVlaV2PipelineConfig( - policy_config=ModelRuntimeConfig(device_type=torch.device(device).type, torch_dtype=dtype), - include_canonical_actions=include_canonical_actions, + policy_config=ModelRuntimeConfig( + device_type=target_device.type, + device_id=target_device.index or 0, + torch_dtype=dtype, + ), ), ) return pipeline @@ -51,8 +54,7 @@ def get_pipeline( @click.option("--camera-right-wrist", required=True, type=click.Path(exists=True, dir_okay=False)) @click.option("--task", required=True) @click.option("--state-json", required=True, help="Raw 14-D RobotWin state as a JSON list") -@click.option("--output", default="action_chunk.npz", type=click.Path(dir_okay=False)) -@click.option("--include-canonical-actions", is_flag=True) +@click.option("--output", default="canonical_action_chunk.npz", type=click.Path(dir_okay=False)) @click.option("--seed", default=None, type=int) @click.option("--device", default="cuda") def main( @@ -64,11 +66,10 @@ def main( task: str, state_json: str, output: str, - include_canonical_actions: bool, seed: int | None, device: str, ) -> None: - """Predict and save a structured RobotWin action chunk.""" + """Predict and save a normalized canonical action chunk.""" try: state = json.loads(state_json) except json.JSONDecodeError as error: @@ -90,25 +91,20 @@ def main( model_root, qwen3vl_root, device=device, - include_canonical_actions=include_canonical_actions, ) try: chunk = pipeline(observation, seed=seed) arrays = { - "action": chunk.raw_actions.numpy(), - "action_arm_position": chunk.fields["action.arm.position"].numpy(), - "action_effector_position": chunk.fields["action.effector.position"].numpy(), - "action_mask": chunk.action_mask.numpy(), + "canonical_normalized_actions": chunk.canonical_normalized_actions.numpy(), "horizon": np.asarray(chunk.horizon), - "robot_profile": np.asarray(chunk.robot_profile), + "action_dim": np.asarray(chunk.action_dim), + "checkpoint_variant": np.asarray(chunk.checkpoint_variant), "policy_verified": np.asarray(chunk.policy_verified), "verification_status": np.asarray(chunk.verification_status), } - if chunk.canonical_normalized_actions is not None: - arrays["canonical_normalized_actions"] = chunk.canonical_normalized_actions.numpy() np.savez(output, **arrays) click.echo( - f"Saved {chunk.horizon}-step RobotWin action chunk to {output}; " + f"Saved {chunk.horizon}-step normalized canonical action chunk to {output}; " f"policy status: {chunk.verification_status}" ) finally: diff --git a/telefuser/models/lingbot_vla_v2.py b/telefuser/models/lingbot_vla_v2.py index d2837f64..30687951 100644 --- a/telefuser/models/lingbot_vla_v2.py +++ b/telefuser/models/lingbot_vla_v2.py @@ -77,6 +77,7 @@ def __init__( routed_scaling_factor: float = 1.0, use_shared_expert_gate: bool = True, moe_implementation: Optional[Literal[None, "eager", "fused"]] = None, + use_robby_moe_kernel: bool = False, split_fused_experts_from_decoder_fsdp: bool = False, expert_hidden_size: int = 768, expert_intermediate_size: int = 2752, @@ -155,6 +156,7 @@ def __init__( self.routed_scaling_factor = routed_scaling_factor self.use_shared_expert_gate = use_shared_expert_gate self.moe_implementation = moe_implementation + self.use_robby_moe_kernel = use_robby_moe_kernel if moe_implementation is not None: if moe_implementation not in ("eager", "fused"): raise ValueError(f"Invalid moe_implementation: {moe_implementation}") @@ -1952,7 +1954,7 @@ def __init__(self, config: QwenvlWithExpertV2Config, eval=False): if self.config.use_lm_head: self.qwenvl.tie_weights() - self.config.qwen_expert_config._attn_implementation = "flash_attention_2" + self.config.qwen_expert_config._attn_implementation = base_attn_implementation self.qwen_expert = Qwen2ForCausalLM._from_config(self.config.qwen_expert_config, eval=eval) if getattr(self.config, "adanorm_time", False): @@ -1969,6 +1971,7 @@ def __init__(self, config: QwenvlWithExpertV2Config, eval=False): self.cu_seqlens = None self.visual_split_sizes = None self.visual_max_seqlen = None + self._cached_image_grid_signature = None del self.qwen_expert.model.embed_tokens if self.config.enable_expert_vision: @@ -1985,6 +1988,20 @@ def __init__(self, config: QwenvlWithExpertV2Config, eval=False): self.attention_interface = self.get_attention_interface() self.set_requires_grad() + def _apply(self, fn): + super()._apply(fn) + for name in ("pos_embeds", "position_embeddings", "cu_seqlens"): + value = getattr(self, name, None) + if isinstance(value, torch.Tensor): + setattr(self, name, fn(value)) + elif isinstance(value, tuple): + setattr( + self, + name, + tuple(fn(item) if isinstance(item, torch.Tensor) else item for item in value), + ) + return self + def _install_moe_blocks(self): if not getattr(self.config, "use_moe", False): return @@ -2009,6 +2026,7 @@ def _install_moe_blocks(self): token_config.router_activation = getattr(self.config, "router_activation", "softmax") token_config.routed_scaling_factor = getattr(self.config, "routed_scaling_factor", 1.0) token_config.use_shared_expert_gate = getattr(self.config, "use_shared_expert_gate", True) + token_config.use_robby_moe_kernel = getattr(self.config, "use_robby_moe_kernel", False) for idx in token_moe_layers: self.qwen_expert.model.layers[idx].mlp = Qwen2TokenMoeBlock(token_config) @@ -2035,7 +2053,9 @@ def get_image_features( image_grid_thw: torch.LongTensor, ): precompute_grid_thw = getattr(self.config, "precompute_grid_thw", False) - if precompute_grid_thw and self.position_embeddings is None: + grid_signature = tuple(image_grid_thw.detach().to(device="cpu").reshape(-1).tolist()) + cache_miss = self.position_embeddings is None or self._cached_image_grid_signature != grid_signature + if precompute_grid_thw and cache_miss: ( self.pos_embeds, self.position_embeddings, @@ -2043,6 +2063,7 @@ def get_image_features( self.visual_split_sizes, self.visual_max_seqlen, ) = self.qwenvl.visual.preprcess_grid_thw(grid_thw=image_grid_thw) + self._cached_image_grid_signature = grid_signature image_embeds, deepstack_image_embeds = self.qwenvl.visual( pixel_values, grid_thw=image_grid_thw, @@ -2279,6 +2300,7 @@ def __init__(self, config, eval): "router_activation", "routed_scaling_factor", "use_shared_expert_gate", + "use_robby_moe_kernel", "_moe_implementation", ]: if hasattr(config, name): @@ -2816,7 +2838,7 @@ def sample_actions( x_t += dt * v_t time += dt - print(f"Denoise {count} steps") + logger.debug("Denoised actions in %d steps", count) return x_t def predict_velocity( diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py index 38caa903..b06c5e0f 100644 --- a/telefuser/models/lingbot_vla_v2_loader.py +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -1571,9 +1571,10 @@ def map_ckpt_key(self, key, load_vlm_only=False, post_training=False): OFFICIAL_6B_MODEL_CONFIG: dict[str, Any] = { - "post_training": True, + "post_training": False, "adanorm_time": True, "moe_implementation": "fused", + "use_robby_moe_kernel": False, "attention_implementation": "eager", "precompute_grid_thw": True, "vlm_causal": True, @@ -1686,9 +1687,17 @@ def resolve_lingbot_vla_v2_shards(model_path: str | Path) -> list[str]: return [str(path) for path in shard_paths] -def build_official_6b_config(qwen3vl_path: str | Path): +def build_official_6b_config( + qwen3vl_path: str | Path, + *, + checkpoint_variant: str = "base", + checkpoint_path: str | Path | None = None, +): from telefuser.models.lingbot_vla_v2 import LingbotVLAV2Config + if checkpoint_variant != "base": + raise ValueError(f"Unsupported LingBot-VLA v2 checkpoint variant: {checkpoint_variant!r}") + qwen_path = Path(qwen3vl_path).expanduser().resolve() qwen_config = AutoConfig.from_pretrained(str(qwen_path), local_files_only=True) if not hasattr(qwen_config, "text_config") or not hasattr(qwen_config, "vision_config"): @@ -1732,6 +1741,10 @@ def build_official_6b_config(qwen3vl_path: str | Path): config.tokenizer_path = str(qwen_path) config.use_cache = True config.attention_implementation = "eager" + config.checkpoint_variant = checkpoint_variant + config.checkpoint_path = None if checkpoint_path is None else str(Path(checkpoint_path).expanduser().resolve()) + config.policy_verified = False + config.verification_status = "unverified_official_6b_base" return config @@ -1749,12 +1762,23 @@ def validate_official_6b_checkpoint(state_dict): class LingBotVlaV2StateDictConverter: - def __init__(self, qwen3vl_path: str | Path): + def __init__( + self, + qwen3vl_path: str | Path, + checkpoint_variant: str = "base", + checkpoint_path: str | Path | None = None, + ): self.qwen3vl_path = Path(qwen3vl_path) + self.checkpoint_variant = checkpoint_variant + self.checkpoint_path = checkpoint_path def from_official(self, state_dict): validate_official_6b_checkpoint(state_dict) - config = build_official_6b_config(self.qwen3vl_path) + config = build_official_6b_config( + self.qwen3vl_path, + checkpoint_variant=self.checkpoint_variant, + checkpoint_path=self.checkpoint_path, + ) return state_dict, {"config": config, "eval": True} def from_diffusers(self, state_dict): @@ -1769,10 +1793,12 @@ def load_lingbot_vla_v2( *, torch_dtype=torch.bfloat16, device=None, + checkpoint_variant: str = "base", ): from telefuser.models.lingbot_vla_v2 import LingBotVlaV2Model - shard_paths = resolve_lingbot_vla_v2_shards(model_path) + checkpoint_path = resolve_lingbot_vla_v2_checkpoint(model_path).parent + shard_paths = resolve_lingbot_vla_v2_shards(checkpoint_path) module_manager.load_model( shard_paths, device=device, @@ -1781,7 +1807,11 @@ def load_lingbot_vla_v2( name="lingbot_vla_v2", model_class=LingBotVlaV2Model, model_resource="official", - converter_kwargs={"qwen3vl_path": str(qwen3vl_path)}, + converter_kwargs={ + "qwen3vl_path": str(qwen3vl_path), + "checkpoint_variant": checkpoint_variant, + "checkpoint_path": str(checkpoint_path), + }, strict=True, ) return module_manager.fetch_module("lingbot_vla_v2") diff --git a/telefuser/models/lingbot_vla_v2_moe.py b/telefuser/models/lingbot_vla_v2_moe.py index 82ea43af..9c4e790d 100644 --- a/telefuser/models/lingbot_vla_v2_moe.py +++ b/telefuser/models/lingbot_vla_v2_moe.py @@ -553,6 +553,7 @@ def __init__(self, config): # EP/fused support: choose expert storage based on moe_implementation self._moe_implementation = getattr(config, '_moe_implementation', None) or 'eager' + self._use_robby_moe_kernel = bool(getattr(config, "use_robby_moe_kernel", False)) if self._moe_implementation == 'fused': self.experts = Qwen2FusedExperts( self.num_experts, @@ -604,7 +605,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Expert computation: fused (group_gemm) or eager (per-expert loop) if self._moe_implementation == 'fused': use_robby_moe = ( - robby_moe_forward is not None + self._use_robby_moe_kernel + and robby_moe_forward is not None and hidden_flat.is_cuda and not self.training and not torch.is_grad_enabled() @@ -826,4 +828,3 @@ def apply_lingbot_qwen2_patch(): hf_qwen2.Qwen2Model = Qwen2Model hf_qwen2.Qwen2ForCausalLM = Qwen2ForCausalLM - diff --git a/telefuser/models/lingbot_vla_v2_qwen.py b/telefuser/models/lingbot_vla_v2_qwen.py index 189194d2..bb67141e 100644 --- a/telefuser/models/lingbot_vla_v2_qwen.py +++ b/telefuser/models/lingbot_vla_v2_qwen.py @@ -401,7 +401,7 @@ def qwen25_forward_without_grid_thw( def apply_lingbot_qwen25_vl_patch(): - logger.info_rank0("apply patch") + logger.info("apply Qwen2.5-VL LingBot patch") hf_qwen25vl.Qwen2_5_VLPreTrainedModel = Qwen2_5_VLPreTrainedModel hf_qwen25vl.Qwen2_5_VLDecoderLayer = Qwen2_5_VLDecoderLayer hf_qwen25vl.Qwen2_5_VLTextModel = Qwen2_5_VLTextModel @@ -733,7 +733,7 @@ def forward_without_grid_thw( def apply_lingbot_qwen3_vl_patch(): - logger.info_rank0("apply Qwen3-VL Lingbot patch") + logger.info("apply Qwen3-VL LingBot patch") hf_qwen3vl.Qwen3VLPreTrainedModel = Qwen3VLPreTrainedModel hf_qwen3vl.Qwen3VLTextDecoderLayer = Qwen3VLTextDecoderLayer hf_qwen3vl.Qwen3VLTextModel = Qwen3VLTextModel @@ -743,4 +743,3 @@ def apply_lingbot_qwen3_vl_patch(): hf_qwen3vl.Qwen3VLVisionBlock = Qwen3VLVisionBlock hf_qwen3vl.Qwen3VLVisionModel.forward = forward_without_grid_thw hf_qwen3vl.Qwen3VLVisionModel.preprcess_grid_thw = preprcess_grid_thw - diff --git a/telefuser/pipelines/lingbot_vla_v2/__init__.py b/telefuser/pipelines/lingbot_vla_v2/__init__.py index 3ecd6c76..de38aed1 100644 --- a/telefuser/pipelines/lingbot_vla_v2/__init__.py +++ b/telefuser/pipelines/lingbot_vla_v2/__init__.py @@ -1,12 +1,13 @@ """TeleFuser pipeline components for LingBot-VLA v2 action inference.""" from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation -from .pipeline import LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig +from .pipeline import LingBotVlaV2CanonicalActionChunk, LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig from .policy import LingBotVlaV2PolicyStage from .robot_profile import ROBOTWIN_CAMERA_KEYS, LingBotVlaV2ActionChunk, RobotWinProfile __all__ = [ "LingBotVlaV2ActionChunk", + "LingBotVlaV2CanonicalActionChunk", "LingBotVlaV2InputProcessor", "LingBotVlaV2Inputs", "LingBotVlaV2Observation", diff --git a/telefuser/pipelines/lingbot_vla_v2/data.py b/telefuser/pipelines/lingbot_vla_v2/data.py index a4055003..374e6489 100644 --- a/telefuser/pipelines/lingbot_vla_v2/data.py +++ b/telefuser/pipelines/lingbot_vla_v2/data.py @@ -9,10 +9,10 @@ import numpy as np import torch from PIL import Image +from torchvision.transforms.v2 import Resize from .robot_profile import ROBOTWIN_CAMERA_KEYS, RobotWinProfile - ImageInput = Image.Image | np.ndarray | torch.Tensor | str | Path @@ -71,11 +71,22 @@ def _image_to_chw_uint8(image: ImageInput) -> torch.Tensor: class LingBotVlaV2InputProcessor: """Prepare RobotWin images, task text, and canonical state tensors.""" - def __init__(self, processor: Any, model_config: Any, robot_profile: RobotWinProfile) -> None: + def __init__( + self, + processor: Any, + model_config: Any, + robot_profile: RobotWinProfile, + *, + image_size: int = 256, + ) -> None: if processor is None or not hasattr(processor, "image_processor") or not hasattr(processor, "tokenizer"): raise TypeError("LingBot-VLA v2 requires a Qwen3-VL AutoProcessor") self.processor = processor self.robot_profile = robot_profile + if image_size <= 0: + raise ValueError(f"image_size must be positive, got {image_size}") + self.image_size = int(image_size) + self.image_resize = Resize((self.image_size, self.image_size), antialias=True) self.max_state_dim = int(getattr(model_config, "max_state_dim", 55)) self.tokenizer_max_length = int(getattr(model_config, "tokenizer_max_length", 72)) if self.max_state_dim != robot_profile.canonical_dim: @@ -93,7 +104,8 @@ def _process_images( processed_images: list[torch.Tensor] = [] grids: list[torch.Tensor] = [] for key in self.robot_profile.camera_keys: - output = self.processor.image_processor(_image_to_chw_uint8(images[key])) + image = self.image_resize(_image_to_chw_uint8(images[key])) + output = self.processor.image_processor(image) pixels = output["pixel_values"] if isinstance(output, dict) else output.pixel_values grid = output.get("image_grid_thw") if isinstance(output, dict) else getattr(output, "image_grid_thw", None) pixels = torch.as_tensor(pixels) diff --git a/telefuser/pipelines/lingbot_vla_v2/pipeline.py b/telefuser/pipelines/lingbot_vla_v2/pipeline.py index 7d184e17..4f941154 100644 --- a/telefuser/pipelines/lingbot_vla_v2/pipeline.py +++ b/telefuser/pipelines/lingbot_vla_v2/pipeline.py @@ -1,4 +1,4 @@ -"""BasePipeline integration for LingBot-VLA v2 RobotWin inference.""" +"""BasePipeline integration for LingBot-VLA v2 base-model inference.""" from __future__ import annotations @@ -12,7 +12,7 @@ from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation from .policy import LingBotVlaV2PolicyStage -from .robot_profile import LingBotVlaV2ActionChunk, RobotWinProfile +from .robot_profile import RobotWinProfile @dataclass @@ -21,12 +21,24 @@ class LingBotVlaV2PipelineConfig: policy_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) robot_profile: RobotWinProfile = field(default_factory=RobotWinProfile.default) - include_canonical_actions: bool = False + image_size: int = 256 enable_metrics: bool = False +@dataclass(frozen=True) +class LingBotVlaV2CanonicalActionChunk: + """Normalized canonical actions produced by the base checkpoint.""" + + canonical_normalized_actions: torch.Tensor + horizon: int + action_dim: int + checkpoint_variant: str = "base" + policy_verified: bool = False + verification_status: str = "unverified_official_6b_base" + + class LingBotVlaV2Pipeline(BasePipeline): - """Single-replica LingBot-VLA v2 structured action SDK.""" + """Single-replica LingBot-VLA v2 canonical action SDK.""" def _get_stages(self) -> list: return [self.policy_stage] @@ -38,18 +50,37 @@ def init(self, module_manager: ModuleManager, config: LingBotVlaV2PipelineConfig processor = module_manager.fetch_module("lingbot_vla_v2_processor") if policy is None or processor is None: raise RuntimeError("LingBot-VLA v2 requires policy and lingbot_vla_v2_processor modules") - self.input_processor = LingBotVlaV2InputProcessor(processor, policy.config, config.robot_profile) + self.input_processor = LingBotVlaV2InputProcessor( + processor, + policy.config, + config.robot_profile, + image_size=config.image_size, + ) self.policy_stage = LingBotVlaV2PolicyStage("policy", module_manager, config.policy_config) if config.enable_metrics: self.enable_metrics() @torch.inference_mode() - def predict(self, inputs: LingBotVlaV2Inputs, seed: int | None = None) -> LingBotVlaV2ActionChunk: - """Run prepared tensors and convert the canonical result to RobotWin fields.""" + def predict( + self, + inputs: LingBotVlaV2Inputs, + seed: int | None = None, + ) -> LingBotVlaV2CanonicalActionChunk: + """Run prepared tensors and return normalized canonical actions.""" actions = self.policy_stage.process(inputs, seed=seed) - return self.config.robot_profile.structure_actions( - actions, - include_canonical=self.config.include_canonical_actions, + if actions.shape[0] != 1: + raise RuntimeError(f"LingBot-VLA v2 pipeline expects batch size 1, got {actions.shape[0]}") + canonical_actions = actions[0] + policy_config = self.policy_stage.policy.config + return LingBotVlaV2CanonicalActionChunk( + canonical_normalized_actions=canonical_actions, + horizon=int(canonical_actions.shape[0]), + action_dim=int(canonical_actions.shape[1]), + checkpoint_variant=str(getattr(policy_config, "checkpoint_variant", "base")), + policy_verified=bool(getattr(policy_config, "policy_verified", False)), + verification_status=str( + getattr(policy_config, "verification_status", "unverified_official_6b_base") + ), ) @torch.inference_mode() @@ -57,8 +88,8 @@ def __call__( self, observation: LingBotVlaV2Observation, seed: int | None = None, - ) -> LingBotVlaV2ActionChunk: - """Predict one structured RobotWin action chunk.""" + ) -> LingBotVlaV2CanonicalActionChunk: + """Predict one normalized canonical action chunk.""" return self.predict(self.input_processor.prepare(observation), seed=seed) def close(self) -> None: diff --git a/tests/unit/models/test_lingbot_vla_v2.py b/tests/unit/models/test_lingbot_vla_v2.py new file mode 100644 index 00000000..b80e0924 --- /dev/null +++ b/tests/unit/models/test_lingbot_vla_v2.py @@ -0,0 +1,53 @@ +from types import SimpleNamespace + +import torch + +from telefuser.models.lingbot_vla_v2 import QwenvlWithExpertV2Model + + +class _Visual: + spatial_merge_size = 1 + + def __init__(self) -> None: + self.preprocess_calls = 0 + + def preprcess_grid_thw(self, grid_thw: torch.Tensor): + self.preprocess_calls += 1 + token_count = int(grid_thw.prod(dim=-1).sum()) + position_embeddings = (torch.zeros(token_count, 2), torch.ones(token_count, 2)) + cu_seqlens = torch.tensor([0, token_count], dtype=torch.int32) + split_sizes = grid_thw.prod(dim=-1).tolist() + return None, position_embeddings, cu_seqlens, split_sizes, token_count + + def __call__(self, pixel_values: torch.Tensor, **kwargs): + del kwargs + embeddings = torch.zeros(pixel_values.shape[0], 3) + return embeddings, [embeddings.clone()] + + +def _model(visual: _Visual) -> SimpleNamespace: + return SimpleNamespace( + config=SimpleNamespace(precompute_grid_thw=True), + qwenvl=SimpleNamespace(visual=visual), + pos_embeds=None, + position_embeddings=None, + cu_seqlens=None, + visual_split_sizes=None, + visual_max_seqlen=None, + _cached_image_grid_signature=None, + ) + + +def test_image_grid_cache_is_reused_and_invalidated_by_grid_shape() -> None: + visual = _Visual() + model = _model(visual) + first_grid = torch.tensor([[1, 2, 2], [1, 2, 2]]) + second_grid = torch.tensor([[1, 1, 2], [1, 1, 2]]) + + first = QwenvlWithExpertV2Model.get_image_features(model, torch.zeros(8, 6), first_grid) + repeated = QwenvlWithExpertV2Model.get_image_features(model, torch.zeros(8, 6), first_grid.clone()) + changed = QwenvlWithExpertV2Model.get_image_features(model, torch.zeros(4, 6), second_grid) + + assert visual.preprocess_calls == 2 + assert first[0].shape == repeated[0].shape == (2, 4, 3) + assert changed[0].shape == (2, 2, 3) diff --git a/tests/unit/models/test_lingbot_vla_v2_loader.py b/tests/unit/models/test_lingbot_vla_v2_loader.py index 48f42d76..b8cb2771 100644 --- a/tests/unit/models/test_lingbot_vla_v2_loader.py +++ b/tests/unit/models/test_lingbot_vla_v2_loader.py @@ -6,6 +6,7 @@ import pytest from telefuser.models.lingbot_vla_v2_loader import ( + build_official_6b_config, resolve_lingbot_vla_v2_shards, validate_official_6b_checkpoint, ) @@ -56,3 +57,8 @@ def test_validate_official_6b_checkpoint_rejects_wrong_shape() -> None: with pytest.raises(ValueError, match="Unexpected shape"): validate_official_6b_checkpoint(state_dict) + + +def test_build_official_6b_config_rejects_non_base_variant(tmp_path) -> None: + with pytest.raises(ValueError, match="Unsupported LingBot-VLA v2 checkpoint variant"): + build_official_6b_config(tmp_path, checkpoint_variant="robotwin") diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_data.py b/tests/unit/pipelines/lingbot_vla_v2/test_data.py index ce8dc7a2..f73bf740 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_data.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_data.py @@ -51,7 +51,11 @@ def _processor() -> tuple[LingBotVlaV2InputProcessor, _ImageProcessor, _Tokenize tokenizer = _Tokenizer() processor = SimpleNamespace(image_processor=image_processor, tokenizer=tokenizer) config = SimpleNamespace(max_state_dim=55, tokenizer_max_length=6) - return LingBotVlaV2InputProcessor(processor, config, RobotWinProfile.default()), image_processor, tokenizer + return ( + LingBotVlaV2InputProcessor(processor, config, RobotWinProfile.default(), image_size=8), + image_processor, + tokenizer, + ) def _observation() -> LingBotVlaV2Observation: @@ -99,3 +103,16 @@ def test_prepare_scales_unit_float_images_to_uint8() -> None: processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) assert image_processor.values[0] == 128 + + +def test_prepare_resizes_each_camera_before_qwen_processing() -> None: + processor, image_processor, _ = _processor() + observation = _observation() + images = { + key: np.full((12, 16, 3), value, dtype=np.uint8) + for key, value in zip(ROBOTWIN_CAMERA_KEYS, (10, 20, 30), strict=True) + } + + processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) + + assert image_processor.values == [10, 20, 30] diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py index 436be6bb..e95e28f8 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py @@ -46,6 +46,9 @@ def __init__(self) -> None: max_action_dim=55, n_action_steps=4, tokenizer_max_length=6, + checkpoint_variant="base", + policy_verified=False, + verification_status="unverified_official_6b_base", ) def sample_actions(self, **inputs) -> torch.Tensor: @@ -54,7 +57,7 @@ def sample_actions(self, **inputs) -> torch.Tensor: return torch.zeros(1, self.config.n_action_steps, self.config.max_action_dim, device=self.anchor.device) -def test_pipeline_returns_structured_robotwin_action_chunk() -> None: +def test_pipeline_returns_normalized_canonical_action_chunk() -> None: policy = _Policy() processor = SimpleNamespace(image_processor=_ImageProcessor(), tokenizer=_Tokenizer()) manager = ModuleManager(torch_dtype=torch.float32, device="cpu") @@ -65,7 +68,7 @@ def test_pipeline_returns_structured_robotwin_action_chunk() -> None: manager, LingBotVlaV2PipelineConfig( policy_config=ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), - include_canonical_actions=True, + image_size=8, ), ) observation = LingBotVlaV2Observation( @@ -80,8 +83,8 @@ def test_pipeline_returns_structured_robotwin_action_chunk() -> None: pipeline.close() assert chunk.horizon == 4 - assert chunk.raw_actions.shape == (4, 14) - assert chunk.fields["action.arm.position"].shape == (4, 12) - assert chunk.fields["action.effector.position"].shape == (4, 2) - assert chunk.canonical_normalized_actions is not None + assert chunk.action_dim == 55 + assert chunk.canonical_normalized_actions.shape == (4, 55) + assert chunk.checkpoint_variant == "base" assert chunk.policy_verified is False + assert chunk.verification_status == "unverified_official_6b_base" diff --git a/tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py b/tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py new file mode 100644 index 00000000..6f72a596 --- /dev/null +++ b/tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py @@ -0,0 +1,423 @@ +"""Base pipeline for multimodal generation.""" + +from __future__ import annotations + +import json +from abc import ABC +from datetime import datetime +from functools import wraps +from pathlib import Path +from typing import TYPE_CHECKING, Any, Sequence + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from einops import rearrange + +from telefuser.utils.logging import logger + +if TYPE_CHECKING: + from telefuser.metrics import StageMetricsManager + + +class BasePipeline(ABC): + """Base pipeline for generation tasks.""" + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + + # Wrap init method to print config after initialization + if "init" in cls.__dict__: + original_init = cls.__dict__["init"] + + @wraps(original_init) + def wrapped_init(self, *args, **kwargs): + result = original_init(self, *args, **kwargs) + if hasattr(self, "config"): + self._print_config_banner() + return result + + cls.init = wrapped_init + + # Wrap __call__ method to print parameters and reset timing registry + if "__call__" in cls.__dict__: + original_call = cls.__dict__["__call__"] + + @wraps(original_call) + def wrapped_call(self, *args, **kwargs): + from telefuser.utils.profiler import reset_timing_registry + + reset_timing_registry() + self._print_call_banner(args, kwargs) + try: + return original_call(self, *args, **kwargs) + finally: + if getattr(self, "clear_memory_after_call", True): + # Clear GPU memory after pipeline execution. + import gc + + from telefuser.platforms import current_platform + + gc.collect() + current_platform.empty_cache() + + cls.__call__ = wrapped_call + + # ANSI color codes for banner formatting + _ANSI_CYAN = "\033[36m" + _ANSI_GREEN = "\033[32m" + _ANSI_YELLOW = "\033[33m" + _ANSI_BLUE = "\033[34m" + _ANSI_DIM = "\033[2m" + _ANSI_BOLD = "\033[1m" + _ANSI_RESET = "\033[0m" + + def _get_config_defaults(self, config: Any) -> dict[str, Any]: + """Get default values for config fields. + + Args: + config: Config object (typically a dataclass) + + Returns: + Dict mapping field names to their default values + """ + defaults = {} + if hasattr(config, "__dataclass_fields__"): + from dataclasses import MISSING, fields + + for field in fields(config): + if field.default is not MISSING: + defaults[field.name] = field.default + elif field.default_factory is not MISSING: + defaults[field.name] = field.default_factory() + else: + # No default value - treat as always changed + defaults[field.name] = None + return defaults + + def _print_config_banner(self) -> None: + """Print pipeline config initialization banner with formatted output.""" + SEP = f"{self._ANSI_DIM}─{'─' * 50}─{self._ANSI_RESET}" + + lines = [ + SEP, + f"{self._ANSI_BOLD}{self._ANSI_CYAN}Pipeline Config{self._ANSI_RESET} " + f"{self._ANSI_DIM}{self.__class__.__name__}{self._ANSI_RESET}", + ] + + # Format config fields, only showing values that differ from defaults + config = self.config + defaults = self._get_config_defaults(config) + + if hasattr(config, "__dataclass_fields__"): + # Dataclass config - format each field + from dataclasses import asdict + + config_dict = asdict(config) + changed_count = 0 + for key, value in config_dict.items(): + default_value = defaults.get(key) + is_changed = default_value is None or value != default_value + if is_changed: + changed_count += 1 + formatted_value = self._format_config_value(value) + lines.append(f" {self._ANSI_DIM}{key}:{self._ANSI_RESET} {formatted_value}") + + if changed_count == 0: + lines.append(f" {self._ANSI_DIM}(all defaults){self._ANSI_RESET}") + else: + # Non-dataclass config - try to format as dict + try: + config_dict = dict(config) if hasattr(config, "items") else {} + for key, value in config_dict.items(): + formatted_value = self._format_config_value(value) + lines.append(f" {self._ANSI_DIM}{key}:{self._ANSI_RESET} {formatted_value}") + except Exception: + lines.append(f" {config}") + + lines.append(SEP) + + # Print to stderr for visibility (like logging.py) + import sys + + print("\n".join(lines), file=sys.stderr) + + def _format_config_value(self, value: Any) -> str: + """Format a config value for display. + + Args: + value: Config value to format + + Returns: + Formatted string representation + """ + if isinstance(value, str): + return f"{self._ANSI_GREEN}{value}{self._ANSI_RESET}" + elif isinstance(value, bool): + color = self._ANSI_GREEN if value else self._ANSI_YELLOW + return f"{color}{value}{self._ANSI_RESET}" + elif isinstance(value, (int, float)): + return f"{self._ANSI_BLUE}{value}{self._ANSI_RESET}" + elif isinstance(value, dict): + # Nested dict - show key count + return f"{self._ANSI_DIM}dict({len(value)} items){self._ANSI_RESET}" + elif isinstance(value, list): + return f"{self._ANSI_DIM}list({len(value)} items){self._ANSI_RESET}" + elif hasattr(value, "__class__"): + return f"{self._ANSI_DIM}{value.__class__.__name__}{self._ANSI_RESET}" + return str(value) + + def _print_call_banner(self, args: tuple, kwargs: dict) -> None: + """Print pipeline __call__ parameters banner with formatted output.""" + SEP = f"{self._ANSI_DIM}─{'─' * 50}─{self._ANSI_RESET}" + + lines = [ + SEP, + f"{self._ANSI_BOLD}{self._ANSI_YELLOW}Pipeline Call{self._ANSI_RESET} " + f"{self._ANSI_DIM}{self.__class__.__name__}{self._ANSI_RESET}", + ] + + # Format kwargs (primary parameters for pipeline calls) + if kwargs: + for key, value in kwargs.items(): + formatted_value = self._format_call_value(value) + lines.append(f" {self._ANSI_DIM}{key}:{self._ANSI_RESET} {formatted_value}") + + # Format positional args (if any) + if args: + display_args = args[1:] if args and getattr(args[0], "__class__", None) else args + if display_args: + lines.append(f" {self._ANSI_DIM}args:{self._ANSI_RESET} {self._format_call_value(display_args)}") + + if not kwargs and not args: + lines.append(f" {self._ANSI_DIM}(no parameters){self._ANSI_RESET}") + + lines.append(SEP) + + import sys + + print("\n".join(lines), file=sys.stderr) + + def _format_call_value(self, value: Any) -> str: + """Format a __call__ parameter value for display. + + Args: + value: Parameter value to format + + Returns: + Formatted string representation + """ + if isinstance(value, str): + # Truncate long strings + if len(value) > 80: + return ( + f"{self._ANSI_GREEN}{value[:80]}...{self._ANSI_RESET} " + f"{self._ANSI_DIM}({len(value)} chars){self._ANSI_RESET}" + ) + return f"{self._ANSI_GREEN}{value}{self._ANSI_RESET}" + elif isinstance(value, bool): + color = self._ANSI_GREEN if value else self._ANSI_YELLOW + return f"{color}{value}{self._ANSI_RESET}" + elif isinstance(value, (int, float)): + return f"{self._ANSI_BLUE}{value}{self._ANSI_RESET}" + elif isinstance(value, torch.Tensor): + return ( + f"{self._ANSI_CYAN}Tensor{self._ANSI_RESET}(" + f"{self._ANSI_DIM}shape={list(value.shape)}, dtype={value.dtype}{self._ANSI_RESET})" + ) + elif isinstance(value, Image.Image): + return f"{self._ANSI_CYAN}PIL.Image{self._ANSI_RESET}({self._ANSI_DIM}size={value.size}{self._ANSI_RESET})" + elif isinstance(value, list): + if len(value) > 5: + return f"{self._ANSI_DIM}list({len(value)} items){self._ANSI_RESET}" + formatted_items = [self._format_call_value(v) for v in value] + return f"[{', '.join(formatted_items)}]" + elif isinstance(value, tuple): + if len(value) > 5: + return f"{self._ANSI_DIM}tuple({len(value)} items){self._ANSI_RESET}" + formatted_items = [self._format_call_value(v) for v in value] + return f"({', '.join(formatted_items)})" + elif isinstance(value, dict): + if len(value) > 8: + return f"{self._ANSI_DIM}dict({len(value)} items){self._ANSI_RESET}" + # Recurse so inner Tensor/Image show as shape summary, not raw dump. + formatted_items = [f"{k}={self._format_call_value(v)}" for k, v in value.items()] + return "{" + ", ".join(formatted_items) + "}" + elif value is None: + return f"{self._ANSI_DIM}None{self._ANSI_RESET}" + return str(value) + + def __init__(self, device: str | torch.device, torch_dtype: torch.dtype) -> None: + self.device = device + self.torch_dtype = torch_dtype + # VAE typically requires dimensions divisible by 16 + self.height_division_factor = 16 + self.width_division_factor = 16 + self._metrics_manager: StageMetricsManager | None = None + + def _get_stages(self) -> list: + """Get list of pipeline stages for metrics collection. + + Subclasses should override this method to return their stages. + + Returns: + List of stage instances that support metrics collection. + """ + return [] + + def enable_metrics(self) -> None: + """Enable metrics collection for all pipeline stages. + + This method enables metrics tracking for stages returned by _get_stages(). + + Example: + >>> pipeline = MyPipeline.from_pretrained(...) + >>> pipeline.enable_metrics() + >>> # Run inference with metrics enabled + >>> result = pipeline(...) + >>> # Get metrics + >>> print(pipeline.get_prometheus_metrics()) + """ + if self._metrics_manager is not None: + logger.warning("Metrics already enabled, skipping") + return + + from telefuser.metrics import StageMetricsManager + + self._metrics_manager = StageMetricsManager() + stages = self._get_stages() + + if stages: + self._metrics_manager.enable_all_stages(stages) + logger.info(f"Enabled metrics for stages: {self._metrics_manager.enabled_stages}") + else: + logger.warning("No stages available for metrics collection") + + def disable_metrics(self) -> None: + """Disable metrics collection for all pipeline stages. + + Example: + >>> pipeline.disable_metrics() + """ + if self._metrics_manager is None: + return + + self._metrics_manager.disable_all_stages() + self._metrics_manager = None + logger.info("Disabled metrics for all stages") + + def get_prometheus_metrics(self) -> str: + """Get metrics in Prometheus text exposition format. + + Returns: + String containing all metrics in Prometheus format. + + Example: + >>> metrics = pipeline.get_prometheus_metrics() + >>> print(metrics) + # HELP stage_vae_duration_seconds Execution duration + # TYPE stage_vae_duration_seconds histogram + ... + + Raises: + RuntimeError: If metrics are not enabled. + """ + if self._metrics_manager is None: + raise RuntimeError("Metrics not enabled. Call enable_metrics() first.") + return self._metrics_manager.registry.get_prometheus_format() + + @property + def metrics_enabled(self) -> bool: + """Check if metrics collection is enabled.""" + return self._metrics_manager is not None + + def check_resize_height_width(self, height: int, width: int) -> tuple[int, int]: + """Ensure dimensions are divisible by division factors.""" + if height % self.height_division_factor != 0: + factor = self.height_division_factor + height = ((height + factor - 1) // factor) * factor + print(f"Height rounded up to {height}") + if width % self.width_division_factor != 0: + factor = self.width_division_factor + width = ((width + factor - 1) // factor) * factor + print(f"Width rounded up to {width}") + return height, width + + def preprocess_image(self, image: Image.Image, height: int | None = None, width: int | None = None) -> torch.Tensor: + """Preprocess PIL image to tensor.""" + if height is not None and width is not None: + if height != image.size[1] or width != image.size[0]: + image = image.resize((width, height), Image.LANCZOS) + return torch.Tensor(np.array(image, dtype=np.float32) * (2 / 255) - 1).permute(2, 0, 1).unsqueeze(0) + + def preprocess_images( + self, images: Sequence[Image.Image], height: int | None = None, width: int | None = None + ) -> list[torch.Tensor]: + """Preprocess multiple images.""" + return [self.preprocess_image(image, height, width) for image in images] + + def tensor2video( + self, frames: torch.Tensor, height: int | None = None, width: int | None = None + ) -> list[Image.Image]: + """Convert tensor to list of PIL images.""" + if height is not None and width is not None: + if height != frames.shape[2] or width != frames.shape[3]: + logger.info(f"Resizing video to {width}x{height}") + # Bicubic antialias resize does not support bf16 inputs in PyTorch. + resize_dtype = frames.dtype + resize_frames = frames.float() if frames.dtype == torch.bfloat16 else frames + frames = F.interpolate( + resize_frames, size=(height, width), mode="bicubic", align_corners=False, antialias=True + ) + frames = frames.to(resize_dtype) + frames = rearrange(frames, "C T H W -> T H W C") + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + return [Image.fromarray(frame) for frame in frames] + + def generate_noise( + self, shape: Sequence[int], seed: int | None = None, device: str = "cpu", dtype: torch.dtype = torch.float16 + ) -> torch.Tensor: + """Generate random noise.""" + generator = None if seed is None else torch.Generator(device).manual_seed(seed) + return torch.randn(shape, generator=generator, device=device, dtype=dtype) + + def dump_config(self, path: str | Path | None = None) -> dict[str, Any]: + """Dump pipeline configuration to dict or file. + + Captures model definition (Layer 1) and inference algorithm + parameters (Layer 2) for reproducibility and debugging. + + Args: + path: Optional file path to save config (JSON format) + + Returns: + Configuration dictionary with model and inference settings + + Example: + >>> pipeline = MyPipeline.from_pretrained(...) + >>> config = pipeline.dump_config() + >>> pipeline.dump_config("output/config.json") + """ + from .config_serializer import serialize_config + + config: dict[str, Any] = { + "version": "1.0", + "timestamp": datetime.now().isoformat(), + "pipeline_type": self.__class__.__name__, + "device": str(self.device), + "torch_dtype": str(self.torch_dtype).replace("torch.", ""), + "layer1_model_definition": { + "models": self._model_info if hasattr(self, "_model_info") else [], + }, + "layer2_inference_config": serialize_config(self.config) if hasattr(self, "config") else {}, + } + + if path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=2, ensure_ascii=False) + logger.info(f"Config dumped to {path}") + + return config From 6ee8998ec4854189733db88005b9f1176727a563 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 3 Aug 2026 07:20:45 +0000 Subject: [PATCH 03/15] fix(vla): harden LingBot VLA v2 integration Fix the public official-checkpoint loader and bundle pinned RobotWin normalization statistics with public entry and profile coverage. Localize the Qwen3-VL adaptations, move the optional Triton MoE implementation behind the ops/kernel boundary, and keep PyTorch fallbacks for unsupported execution. Pin the upstream LingBot-VLA v2 revision in a preprocessing, velocity, and final-action parity comparator, and remove the tracked temporary source copy. Verification: - .venv-vla/bin/python -m pytest -q tests/unit/models/test_lingbot_vla_v2.py tests/unit/models/test_lingbot_vla_v2_loader.py tests/unit/pipelines/lingbot_vla_v2 (17 passed) - ruff check on new VLA ops, kernel, validation, and tests - ruff check --select F821 on modified VLA model modules - .venv-vla/bin/python -m py_compile on modified VLA modules - git diff --check --- .gitignore | 1 + telefuser/kernel/triton/lingbot_vla_v2_moe.py | 265 +++++++++++ telefuser/models/lingbot_vla_v2.py | 2 - telefuser/models/lingbot_vla_v2_loader.py | 181 +------- telefuser/models/lingbot_vla_v2_moe.py | 264 +---------- telefuser/models/lingbot_vla_v2_qwen.py | 26 +- telefuser/ops/lingbot_vla_v2_moe.py | 33 ++ .../assets/robotwin_norm_stats.json | 25 ++ .../unit/models/test_lingbot_vla_v2_loader.py | 56 +++ .../lingbot_vla_v2/test_robot_profile.py | 12 + .../telefuser/core/base_pipeline.py | 423 ------------------ tools/validation/run_lingbot_vla_v2_parity.py | 149 ++++++ 12 files changed, 558 insertions(+), 879 deletions(-) create mode 100644 telefuser/kernel/triton/lingbot_vla_v2_moe.py create mode 100644 telefuser/ops/lingbot_vla_v2_moe.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json delete mode 100644 tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py create mode 100644 tools/validation/run_lingbot_vla_v2_parity.py diff --git a/.gitignore b/.gitignore index 99e68f0e..5d23efb8 100755 --- a/.gitignore +++ b/.gitignore @@ -158,6 +158,7 @@ INSTALL_HYS.md AGENTS.md _version.py.mcp.json telefuser/_version.py +!telefuser/pipelines/lingbot_vla_v2/assets/*.json # LingBot regression example assets !examples/data/lingbot_world_fast/image.jpg !examples/data/lingbot_world_fast/poses.npy diff --git a/telefuser/kernel/triton/lingbot_vla_v2_moe.py b/telefuser/kernel/triton/lingbot_vla_v2_moe.py new file mode 100644 index 00000000..eb851d58 --- /dev/null +++ b/telefuser/kernel/triton/lingbot_vla_v2_moe.py @@ -0,0 +1,265 @@ +# ruff: noqa: E741 +"""Triton grouped-MoE kernels for LingBot-VLA v2. + +Internal implementation; callers should use ``telefuser.ops.lingbot_vla_v2_moe``. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _zero_i32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.int32), mask=offs < N) + + +@triton.jit +def _zero_fp32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.float32), mask=offs < N) + + +@triton.jit +def _moe_pack_selected_kernel( + selected_ptr, + route_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + T: tl.constexpr, + TOPK: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_K: tl.constexpr, +): + row = tl.program_id(0) + slots = tl.arange(0, BLOCK_K) + mask = slots < TOPK + experts = tl.load(selected_ptr + row * TOPK + slots, mask=mask, other=0).to(tl.int32) + pos = tl.atomic_add(counts_ptr + experts, 1, sem="relaxed", mask=mask) + store_mask = mask & (pos < MAX_ROUTES) + tl.store(rows_ptr + experts * MAX_ROUTES + pos, row, mask=store_mask) + tl.store(slots_ptr + experts * MAX_ROUTES + pos, slots, mask=store_mask) + + +@triton.jit +def _moe_gate_up_grouped_kernel( + x_ptr, + gate_ptr, + up_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + route_ptr, + inter_ptr, + T: tl.constexpr, + D: tl.constexpr, + E: tl.constexpr, + TOPK: tl.constexpr, + I: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_I: tl.constexpr, + BLOCK_D: tl.constexpr, +): + expert = tl.program_id(0) + bid_m = tl.program_id(1) + bid_i = tl.program_id(2) + count = tl.load(counts_ptr + expert).to(tl.int32) + start_m = bid_m * BLOCK_M + if start_m >= count: + return + route_idx = start_m + tl.arange(0, BLOCK_M) + offs_i = bid_i * BLOCK_I + tl.arange(0, BLOCK_I) + offs_d = tl.arange(0, BLOCK_D) + valid_m = route_idx < count + rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + acc_g = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) + acc_u = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) + for d0 in range(0, D, BLOCK_D): + ds = d0 + offs_d + x = tl.load( + x_ptr + rows[:, None] * D + ds[None, :], + mask=valid_m[:, None] & (ds[None, :] < D), + other=0.0, + ) + gw = tl.load( + gate_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], + mask=(offs_i[None, :] < I) & (ds[:, None] < D), + other=0.0, + ) + uw = tl.load( + up_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], + mask=(offs_i[None, :] < I) & (ds[:, None] < D), + other=0.0, + ) + acc_g += tl.dot(x, gw) + acc_u += tl.dot(x, uw) + route = tl.load(route_ptr + rows * TOPK + slots, mask=valid_m, other=0.0).to(tl.float32) + silu = acc_g * (1.0 / (1.0 + tl.exp(-acc_g))) + val = silu * acc_u * route[:, None] + tl.store( + inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + offs_i[None, :]), + val.to(inter_ptr.dtype.element_ty), + mask=valid_m[:, None] & (offs_i[None, :] < I), + ) + + +@triton.jit +def _moe_down_grouped_kernel( + inter_ptr, + down_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + out_ptr, + T: tl.constexpr, + D: tl.constexpr, + E: tl.constexpr, + TOPK: tl.constexpr, + I: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_I: tl.constexpr, +): + expert = tl.program_id(0) + bid_m = tl.program_id(1) + bid_d = tl.program_id(2) + count = tl.load(counts_ptr + expert).to(tl.int32) + start_m = bid_m * BLOCK_M + if start_m >= count: + return + route_idx = start_m + tl.arange(0, BLOCK_M) + offs_d = bid_d * BLOCK_D + tl.arange(0, BLOCK_D) + offs_i = tl.arange(0, BLOCK_I) + valid_m = route_idx < count + rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + acc = tl.zeros((BLOCK_M, BLOCK_D), dtype=tl.float32) + for i0 in range(0, I, BLOCK_I): + is_ = i0 + offs_i + x = tl.load( + inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + is_[None, :]), + mask=valid_m[:, None] & (is_[None, :] < I), + other=0.0, + ) + w = tl.load( + down_ptr + (expert * D + offs_d[None, :]) * I + is_[:, None], + mask=(offs_d[None, :] < D) & (is_[:, None] < I), + other=0.0, + ) + acc += tl.dot(x, w) + tl.atomic_add( + out_ptr + rows[:, None] * D + offs_d[None, :], + acc, + sem="relaxed", + mask=valid_m[:, None] & (offs_d[None, :] < D), + ) + + +def robby_moe_forward( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + workspace: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Inference-only grouped MoE path migrated from robbyvla_infer _moe.""" + if hidden_states.ndim != 2: + raise ValueError(f"hidden_states must be 2D, got {tuple(hidden_states.shape)}") + if selected_experts.ndim != 2 or routing_weights.ndim != 2: + raise ValueError("selected_experts and routing_weights must be 2D") + if not hidden_states.is_cuda: + raise ValueError("robby_moe_forward requires CUDA tensors") + + T, D = hidden_states.shape + E, I, weight_d = gate_weight.shape + top_k = selected_experts.shape[1] + if weight_d != D or up_weight.shape != gate_weight.shape or down_weight.shape != (E, D, I): + raise ValueError( + "Unexpected MoE weight shapes: " + f"hidden={tuple(hidden_states.shape)} gate={tuple(gate_weight.shape)} " + f"up={tuple(up_weight.shape)} down={tuple(down_weight.shape)}" + ) + + max_routes = T * top_k + if workspace is None: + counts = torch.empty((E,), device=hidden_states.device, dtype=torch.int32) + rows = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) + slots = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) + inter = torch.empty((T, top_k, I), device=hidden_states.device, dtype=hidden_states.dtype) + out = torch.empty((T, D), device=hidden_states.device, dtype=torch.float32) + else: + counts = workspace["counts"] + rows = workspace["rows"] + slots = workspace["slots"] + inter = workspace["inter"] + out = workspace["out"] + + selected_i32 = selected_experts.to(torch.int32).contiguous() + route = routing_weights.contiguous() + + _zero_i32_kernel[(1,)](counts, E, BLOCK=triton.next_power_of_2(E), num_warps=1) + _moe_pack_selected_kernel[(T,)]( + selected_i32, + route, + counts, + rows, + slots, + T, + top_k, + max_routes, + BLOCK_K=triton.next_power_of_2(top_k), + num_warps=1, + ) + _moe_gate_up_grouped_kernel[(E, triton.cdiv(max_routes, 16), triton.cdiv(I, 32))]( + hidden_states, + gate_weight, + up_weight, + counts, + rows, + slots, + route, + inter, + T, + D, + E, + top_k, + I, + max_routes, + BLOCK_M=16, + BLOCK_I=32, + BLOCK_D=64, + num_warps=4, + ) + _zero_fp32_kernel[(triton.cdiv(out.numel(), 1024),)]( + out, + out.numel(), + BLOCK=1024, + num_warps=4, + ) + _moe_down_grouped_kernel[(E, triton.cdiv(max_routes, 16), triton.cdiv(D, 64))]( + inter, + down_weight, + counts, + rows, + slots, + out, + T, + D, + E, + top_k, + I, + max_routes, + BLOCK_M=16, + BLOCK_D=64, + BLOCK_I=64, + num_warps=4, + ) + return out.reshape_as(hidden_states) diff --git a/telefuser/models/lingbot_vla_v2.py b/telefuser/models/lingbot_vla_v2.py index 30687951..2eec3a13 100644 --- a/telefuser/models/lingbot_vla_v2.py +++ b/telefuser/models/lingbot_vla_v2.py @@ -3171,7 +3171,6 @@ def sample_actions(self, *args, **kwargs) -> Tensor: # __V2_END__ from telefuser.models.lingbot_vla_v2_loader import LingBotVlaV2StateDictConverter -from telefuser.models.lingbot_vla_v2_qwen import apply_lingbot_qwen3_vl_patch class LingBotVlaV2Model(LingbotVlaV2Policy): @@ -3180,7 +3179,6 @@ class LingBotVlaV2Model(LingbotVlaV2Policy): name = "lingbot_vla_v2" def __init__(self, config, eval=True): - apply_lingbot_qwen3_vl_patch() super().__init__(config=config, eval=eval) @staticmethod diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py index b06c5e0f..b6aad310 100644 --- a/telefuser/models/lingbot_vla_v2_loader.py +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -1018,116 +1018,9 @@ def _next_power_of_2(n: int) -> int: # 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? _HAS_TRITON = False -try: - import triton - import triton.language as tl - - _HAS_TRITON = True - - # 鈹€鈹€ Kernel 1: per-segment topK counting (for sequence_wise_balance_loss) 鈹€鈹€ - @triton.jit - def _topk_segment_count_kernel( - logits_ptr, # [N_total, E] - seg_starts_ptr, # [S_total] - seg_lengths_ptr, # [S_total] - f_out_ptr, # [S_total, E] - stride_logits_n, # stride of logits along token dim - E: tl.constexpr, # actual num_experts - K: tl.constexpr, # top_k - BLOCK_E: tl.constexpr, # next power of 2 >= E - ): - """Each program computes f_i (expert counts) for one segment. - - Iterates over tokens in the segment, performs K rounds of argmax to - find top-K experts, and accumulates per-expert hit counts. - No gradient needed 鈥?f_i is always detached in the loss. - """ - seg_id = tl.program_id(0) - seg_start = tl.load(seg_starts_ptr + seg_id) - seg_len = tl.load(seg_lengths_ptr + seg_id) - - expert_offs = tl.arange(0, BLOCK_E) # [BLOCK_E] - mask_e = expert_offs < E - f_acc = tl.zeros((BLOCK_E,), dtype=tl.float32) - - for t in range(0, seg_len): - row_ptr = logits_ptr + (seg_start + t) * stride_logits_n - logits_row = tl.load(row_ptr + expert_offs, mask=mask_e, other=float('-inf')) - - # K rounds of argmax to find top-K indices - row_copy = logits_row - for _k in range(K): - max_val = tl.max(row_copy, axis=0) - is_max = (row_copy == max_val) - # Distribute count evenly among ties (rare with float32) - n_ties = tl.sum(is_max.to(tl.float32), axis=0) - f_acc += tl.where(is_max, 1.0 / n_ties, 0.0) - row_copy = tl.where(is_max, float('-inf'), row_copy) - - # Write f_count (unnormalized) 鈥?caller normalizes by (E / K) / seg_len - out_ptr = f_out_ptr + seg_id * E - tl.store(out_ptr + expert_offs, f_acc, mask=mask_e) - - # 鈹€鈹€ Kernel 2: blocked topK counting (for load_balancing_loss_func) 鈹€鈹€ - @triton.jit - def _topk_count_with_mask_kernel( - routing_weights_ptr, # [N, E] 鈥?softmax probabilities - mask_ptr, # [N] 鈥?1.0 for valid, 0.0 for padding - partial_f_ptr, # [num_blocks, E] 鈥?partial expert counts - partial_p_ptr, # [num_blocks, E] 鈥?partial masked prob sums - N, - stride_rw_n, # stride along token dim - has_mask: tl.constexpr, - E: tl.constexpr, - K: tl.constexpr, - BLOCK_E: tl.constexpr, - BLOCK_N: tl.constexpr, - ): - """Each program accumulates topK counts + masked probs for a token block. - - Two-phase reduction: writes partial [E] results per block; - caller sums across blocks in PyTorch. - """ - pid = tl.program_id(0) - n_start = pid * BLOCK_N - - expert_offs = tl.arange(0, BLOCK_E) - mask_e = expert_offs < E - f_local = tl.zeros((BLOCK_E,), dtype=tl.float32) - p_local = tl.zeros((BLOCK_E,), dtype=tl.float32) - - for t_offset in range(BLOCK_N): - t = n_start + t_offset - # Guard: skip if t >= N (handles last block) - if t < N: - row_ptr = routing_weights_ptr + t * stride_rw_n - rw_row = tl.load(row_ptr + expert_offs, mask=mask_e, other=0.0) - - if has_mask: - m = tl.load(mask_ptr + t) - else: - m = 1.0 - - # Accumulate masked probs - p_local += rw_row * m +# Training-only MoE auxiliary losses keep the PyTorch vectorized path in the +# model loader. Triton kernels live only under telefuser.kernel.triton. - # TopK counting - row_copy = rw_row - for _k in range(K): - max_val = tl.max(row_copy, axis=0) - is_max = (row_copy == max_val) - n_ties = tl.sum(is_max.to(tl.float32), axis=0) - f_local += tl.where(is_max, m / n_ties, 0.0) - row_copy = tl.where(is_max, float('-inf'), row_copy) - - # Write partial results (only E valid elements) - f_ptr = partial_f_ptr + pid * E - p_ptr = partial_p_ptr + pid * E - tl.store(f_ptr + expert_offs, f_local, mask=mask_e) - tl.store(p_ptr + expert_offs, p_local, mask=mask_e) - -except ImportError: - pass # 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? @@ -1229,74 +1122,13 @@ def _vectorized_topk_count( # Section 3: Triton-accelerated wrappers # 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? -def _triton_segment_f_i( - logits: torch.Tensor, # [N_total, E] - seg_starts: torch.Tensor, # [S_total] - seg_lengths: torch.Tensor, # [S_total] - top_k: int, -) -> torch.Tensor: - """Compute f_i per segment using the Triton kernel. Returns [S_total, E].""" - N, E = logits.shape - S_total = seg_starts.shape[0] - BLOCK_E = _next_power_of_2(E) - - f_counts = torch.zeros(S_total, E, device=logits.device, dtype=torch.float32) - - _topk_segment_count_kernel[(S_total,)]( - logits, - seg_starts, - seg_lengths, - f_counts, - logits.stride(0), - E=E, - K=top_k, - BLOCK_E=BLOCK_E, - ) +def _triton_segment_f_i(*args, **kwargs): + raise RuntimeError("LingBot-VLA v2 sequence-wise Triton loss was removed from models; use the PyTorch fallback") - # Normalize: f_i = (E / K) * counts / T_s - inv_lens = (float(E) / top_k) / seg_lengths.unsqueeze(1).float().clamp(min=1) - return f_counts * inv_lens +def _triton_topk_count(*args, **kwargs): + raise RuntimeError("LingBot-VLA v2 load-balancing Triton loss was removed from models; use the PyTorch fallback") -def _triton_topk_count( - routing_weights: torch.Tensor, # [N, E] - top_k: int, - flat_mask: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Compute tokens_per_expert [E] using the Triton kernel.""" - N, E = routing_weights.shape - BLOCK_E = _next_power_of_2(E) - BLOCK_N = 256 - num_blocks = (N + BLOCK_N - 1) // BLOCK_N - - partial_f = torch.zeros(num_blocks, E, device=routing_weights.device, dtype=torch.float32) - partial_p = torch.zeros(num_blocks, E, device=routing_weights.device, dtype=torch.float32) - - has_mask = flat_mask is not None - _topk_count_with_mask_kernel[(num_blocks,)]( - routing_weights, - flat_mask if has_mask else routing_weights, # dummy ptr when no mask - partial_f, - partial_p, - N, - routing_weights.stride(0), - has_mask=has_mask, - E=E, - K=top_k, - BLOCK_E=BLOCK_E, - BLOCK_N=BLOCK_N, - ) - - # Phase 2: reduce across blocks - tokens_per_expert = partial_f.sum(dim=0) # [E] - - if flat_mask is not None: - n_valid = flat_mask.sum().clamp(min=1) - else: - n_valid = float(N) - tokens_per_expert = tokens_per_expert / n_valid - - return tokens_per_expert # 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? @@ -1812,6 +1644,5 @@ def load_lingbot_vla_v2( "checkpoint_variant": checkpoint_variant, "checkpoint_path": str(checkpoint_path), }, - strict=True, ) return module_manager.fetch_module("lingbot_vla_v2") diff --git a/telefuser/models/lingbot_vla_v2_moe.py b/telefuser/models/lingbot_vla_v2_moe.py index 9c4e790d..a11e4024 100644 --- a/telefuser/models/lingbot_vla_v2_moe.py +++ b/telefuser/models/lingbot_vla_v2_moe.py @@ -4,270 +4,8 @@ """ import torch -import triton -import triton.language as tl - - -@triton.jit -def _zero_i32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): - offs = tl.arange(0, BLOCK) - tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.int32), mask=offs < N) - - -@triton.jit -def _zero_fp32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.float32), mask=offs < N) - - -@triton.jit -def _moe_pack_selected_kernel( - selected_ptr, - route_ptr, - counts_ptr, - rows_ptr, - slots_ptr, - T: tl.constexpr, - TOPK: tl.constexpr, - MAX_ROUTES: tl.constexpr, - BLOCK_K: tl.constexpr, -): - row = tl.program_id(0) - slots = tl.arange(0, BLOCK_K) - mask = slots < TOPK - experts = tl.load(selected_ptr + row * TOPK + slots, mask=mask, other=0).to(tl.int32) - routes = tl.load(route_ptr + row * TOPK + slots, mask=mask, other=0.0).to(tl.float32) - pos = tl.atomic_add(counts_ptr + experts, 1, sem="relaxed", mask=mask) - store_mask = mask & (pos < MAX_ROUTES) - tl.store(rows_ptr + experts * MAX_ROUTES + pos, row, mask=store_mask) - tl.store(slots_ptr + experts * MAX_ROUTES + pos, slots, mask=store_mask) - - -@triton.jit -def _moe_gate_up_grouped_kernel( - x_ptr, - gate_ptr, - up_ptr, - counts_ptr, - rows_ptr, - slots_ptr, - route_ptr, - inter_ptr, - T: tl.constexpr, - D: tl.constexpr, - E: tl.constexpr, - TOPK: tl.constexpr, - I: tl.constexpr, - MAX_ROUTES: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_I: tl.constexpr, - BLOCK_D: tl.constexpr, -): - expert = tl.program_id(0) - bid_m = tl.program_id(1) - bid_i = tl.program_id(2) - count = tl.load(counts_ptr + expert).to(tl.int32) - start_m = bid_m * BLOCK_M - if start_m >= count: - return - route_idx = start_m + tl.arange(0, BLOCK_M) - offs_i = bid_i * BLOCK_I + tl.arange(0, BLOCK_I) - offs_d = tl.arange(0, BLOCK_D) - valid_m = route_idx < count - rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) - slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) - acc_g = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) - acc_u = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) - for d0 in range(0, D, BLOCK_D): - ds = d0 + offs_d - x = tl.load( - x_ptr + rows[:, None] * D + ds[None, :], - mask=valid_m[:, None] & (ds[None, :] < D), - other=0.0, - ) - gw = tl.load( - gate_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], - mask=(offs_i[None, :] < I) & (ds[:, None] < D), - other=0.0, - ) - uw = tl.load( - up_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], - mask=(offs_i[None, :] < I) & (ds[:, None] < D), - other=0.0, - ) - acc_g += tl.dot(x, gw) - acc_u += tl.dot(x, uw) - route = tl.load(route_ptr + rows * TOPK + slots, mask=valid_m, other=0.0).to(tl.float32) - silu = acc_g * (1.0 / (1.0 + tl.exp(-acc_g))) - val = silu * acc_u * route[:, None] - tl.store( - inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + offs_i[None, :]), - val.to(inter_ptr.dtype.element_ty), - mask=valid_m[:, None] & (offs_i[None, :] < I), - ) - - -@triton.jit -def _moe_down_grouped_kernel( - inter_ptr, - down_ptr, - counts_ptr, - rows_ptr, - slots_ptr, - out_ptr, - T: tl.constexpr, - D: tl.constexpr, - E: tl.constexpr, - TOPK: tl.constexpr, - I: tl.constexpr, - MAX_ROUTES: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_I: tl.constexpr, -): - expert = tl.program_id(0) - bid_m = tl.program_id(1) - bid_d = tl.program_id(2) - count = tl.load(counts_ptr + expert).to(tl.int32) - start_m = bid_m * BLOCK_M - if start_m >= count: - return - route_idx = start_m + tl.arange(0, BLOCK_M) - offs_d = bid_d * BLOCK_D + tl.arange(0, BLOCK_D) - offs_i = tl.arange(0, BLOCK_I) - valid_m = route_idx < count - rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) - slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) - acc = tl.zeros((BLOCK_M, BLOCK_D), dtype=tl.float32) - for i0 in range(0, I, BLOCK_I): - is_ = i0 + offs_i - x = tl.load( - inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + is_[None, :]), - mask=valid_m[:, None] & (is_[None, :] < I), - other=0.0, - ) - w = tl.load( - down_ptr + (expert * D + offs_d[None, :]) * I + is_[:, None], - mask=(offs_d[None, :] < D) & (is_[:, None] < I), - other=0.0, - ) - acc += tl.dot(x, w) - tl.atomic_add( - out_ptr + rows[:, None] * D + offs_d[None, :], - acc, - sem="relaxed", - mask=valid_m[:, None] & (offs_d[None, :] < D), - ) - - -def robby_moe_forward( - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - gate_weight: torch.Tensor, - up_weight: torch.Tensor, - down_weight: torch.Tensor, - workspace: dict[str, torch.Tensor] | None = None, -) -> torch.Tensor: - """Inference-only grouped MoE path migrated from robbyvla_infer _moe.""" - if hidden_states.ndim != 2: - raise ValueError(f"hidden_states must be 2D, got {tuple(hidden_states.shape)}") - if selected_experts.ndim != 2 or routing_weights.ndim != 2: - raise ValueError("selected_experts and routing_weights must be 2D") - if not hidden_states.is_cuda: - raise ValueError("robby_moe_forward requires CUDA tensors") - - T, D = hidden_states.shape - E, I, weight_d = gate_weight.shape - top_k = selected_experts.shape[1] - if weight_d != D or up_weight.shape != gate_weight.shape or down_weight.shape != (E, D, I): - raise ValueError( - "Unexpected MoE weight shapes: " - f"hidden={tuple(hidden_states.shape)} gate={tuple(gate_weight.shape)} " - f"up={tuple(up_weight.shape)} down={tuple(down_weight.shape)}" - ) - - max_routes = T * top_k - if workspace is None: - counts = torch.empty((E,), device=hidden_states.device, dtype=torch.int32) - rows = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) - slots = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) - inter = torch.empty((T, top_k, I), device=hidden_states.device, dtype=hidden_states.dtype) - out = torch.empty((T, D), device=hidden_states.device, dtype=torch.float32) - else: - counts = workspace["counts"] - rows = workspace["rows"] - slots = workspace["slots"] - inter = workspace["inter"] - out = workspace["out"] - - selected_i32 = selected_experts.to(torch.int32).contiguous() - route = routing_weights.contiguous() - - _zero_i32_kernel[(1,)](counts, E, BLOCK=triton.next_power_of_2(E), num_warps=1) - _moe_pack_selected_kernel[(T,)]( - selected_i32, - route, - counts, - rows, - slots, - T, - top_k, - max_routes, - BLOCK_K=triton.next_power_of_2(top_k), - num_warps=1, - ) - _moe_gate_up_grouped_kernel[ - (E, triton.cdiv(max_routes, 16), triton.cdiv(I, 32)) - ]( - hidden_states, - gate_weight, - up_weight, - counts, - rows, - slots, - route, - inter, - T, - D, - E, - top_k, - I, - max_routes, - BLOCK_M=16, - BLOCK_I=32, - BLOCK_D=64, - num_warps=4, - ) - _zero_fp32_kernel[(triton.cdiv(out.numel(), 1024),)]( - out, - out.numel(), - BLOCK=1024, - num_warps=4, - ) - _moe_down_grouped_kernel[ - (E, triton.cdiv(max_routes, 16), triton.cdiv(D, 64)) - ]( - inter, - down_weight, - counts, - rows, - slots, - out, - T, - D, - E, - top_k, - I, - max_routes, - BLOCK_M=16, - BLOCK_D=64, - BLOCK_I=64, - num_warps=4, - ) - return out.reshape_as(hidden_states) +from telefuser.ops.lingbot_vla_v2_moe import robby_moe_forward def fused_moe_forward( diff --git a/telefuser/models/lingbot_vla_v2_qwen.py b/telefuser/models/lingbot_vla_v2_qwen.py index bb67141e..5ba74ccc 100644 --- a/telefuser/models/lingbot_vla_v2_qwen.py +++ b/telefuser/models/lingbot_vla_v2_qwen.py @@ -417,6 +417,7 @@ def apply_lingbot_qwen25_vl_patch(): import torch from torch import nn import torch.nn.functional as F +from types import MethodType from typing import Callable, Optional, Tuple from transformers.generation import GenerationMixin @@ -426,7 +427,6 @@ def apply_lingbot_qwen25_vl_patch(): from transformers.utils import logging from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig -import transformers.models.qwen3_vl.modeling_qwen3_vl as hf_qwen3vl from transformers.models.qwen3_vl.modeling_qwen3_vl import ( Qwen3VLForConditionalGeneration as _Qwen3VLForConditionalGeneration, Qwen3VLModel as _Qwen3VLModel, @@ -447,11 +447,9 @@ def apply_lingbot_qwen25_vl_patch(): logger = logging.get_logger(__name__) -def _qwen3vl_no_init_weights(self, module): - return - -_Qwen3VLPreTrainedModel._init_weights = _qwen3vl_no_init_weights -Qwen3VLPreTrainedModel = _Qwen3VLPreTrainedModel +class Qwen3VLPreTrainedModel(_Qwen3VLPreTrainedModel): + def _init_weights(self, module): + return class Qwen3VLVisionAttention(nn.Module): @@ -656,6 +654,9 @@ class Qwen3VLModel(_Qwen3VLModel): def __init__(self, config: Qwen3VLConfig): Qwen3VLPreTrainedModel.__init__(self, config) self.visual = Qwen3VLVisionModel._from_config(config.vision_config) + self.visual.blocks = nn.ModuleList([Qwen3VLVisionBlock(config.vision_config) for _ in self.visual.blocks]) + self.visual.forward = MethodType(forward_without_grid_thw, self.visual) + self.visual.preprcess_grid_thw = MethodType(preprcess_grid_thw, self.visual) self.language_model = Qwen3VLTextModel._from_config(config.text_config) self.rope_deltas = None self.post_init() @@ -733,13 +734,6 @@ def forward_without_grid_thw( def apply_lingbot_qwen3_vl_patch(): - logger.info("apply Qwen3-VL LingBot patch") - hf_qwen3vl.Qwen3VLPreTrainedModel = Qwen3VLPreTrainedModel - hf_qwen3vl.Qwen3VLTextDecoderLayer = Qwen3VLTextDecoderLayer - hf_qwen3vl.Qwen3VLTextModel = Qwen3VLTextModel - hf_qwen3vl.Qwen3VLModel = Qwen3VLModel - hf_qwen3vl.Qwen3VLForConditionalGeneration = Qwen3VLForConditionalGeneration - hf_qwen3vl.Qwen3VLVisionAttention = Qwen3VLVisionAttention - hf_qwen3vl.Qwen3VLVisionBlock = Qwen3VLVisionBlock - hf_qwen3vl.Qwen3VLVisionModel.forward = forward_without_grid_thw - hf_qwen3vl.Qwen3VLVisionModel.preprcess_grid_thw = preprcess_grid_thw + logger.warning_once( + "apply_lingbot_qwen3_vl_patch is deprecated; LingBot-VLA v2 now installs Qwen3-VL changes per instance." + ) diff --git a/telefuser/ops/lingbot_vla_v2_moe.py b/telefuser/ops/lingbot_vla_v2_moe.py new file mode 100644 index 00000000..f734ebdc --- /dev/null +++ b/telefuser/ops/lingbot_vla_v2_moe.py @@ -0,0 +1,33 @@ +"""Compile-aware LingBot-VLA v2 MoE operation dispatch.""" + +from __future__ import annotations + +import torch + + +def robby_moe_forward( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + workspace: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Run the optional Triton grouped-MoE path for VLA eager inference.""" + if torch.compiler.is_compiling(): + raise RuntimeError("LingBot-VLA v2 Triton MoE is disabled during torch.compile") + if hidden_states.device.type != "cuda": + raise RuntimeError("LingBot-VLA v2 Triton MoE requires CUDA tensors") + + from telefuser.kernel.triton.lingbot_vla_v2_moe import robby_moe_forward as _triton_robby_moe_forward + + return _triton_robby_moe_forward( + hidden_states, + routing_weights, + selected_experts, + gate_weight, + up_weight, + down_weight, + workspace=workspace, + ) diff --git a/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json b/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json new file mode 100644 index 00000000..1ab5c1da --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json @@ -0,0 +1,25 @@ +{ + "norm_stats": { + "observation.state.arm.position": { + "q01": [-1.3382688760757446, -0.40607330203056335, -1.4083482027053833, -3.058554172515869, -1.423754096031189, -3.192993402481079, -1.591109275817871, -0.7457540035247803, -1.4451789855957031, -3.0523548126220703, -1.4595792293548584, -3.1854426860809326], + "q99": [2.061160087585449, 1.0003128051757812, 1.2696261405944824, 2.941908836364746, 1.4975149631500244, 3.0741331577301025, 1.3934801816940308, 0.3905077278614044, 1.4333486557006836, 3.020704507827759, 1.444725751876831, 3.1354587078094482] + }, + "observation.state.effector.position": { + "q01": [0.3143864572048187, 0.0005160411237739027], + "q99": [1.0, 1.0] + }, + "action.arm.position": { + "q01": [-1.3382868766784668, -0.40629321336746216, -1.407132625579834, -3.0591986179351807, -1.4246528148651123, -3.192993402481079, -1.5909051895141602, -0.7457385063171387, -1.44444739818573, -3.0523548126220703, -1.4598064422607422, -3.1850173473358154], + "q99": [2.060295343399048, 1.0005663633346558, 1.2670165300369263, 2.941908836364746, 1.4981306791305542, 3.0859344005584717, 1.3933433294296265, 0.3887772858142853, 1.4337434768676758, 3.019955635070801, 1.4448832273483276, 3.133007049560547] + }, + "action.effector.position": { + "q01": [0.3143986165523529, 0.0004720990259665996], + "q99": [1.0, 1.0] + } + }, + "source": { + "repository": "https://github.com/Robbyant/lingbot-vla-v2", + "path": "assets/norm_stats/robotwin.json", + "commit": "be27333c9b5f2663b0ec33f069dd7dfd67fa32b5" + } +} diff --git a/tests/unit/models/test_lingbot_vla_v2_loader.py b/tests/unit/models/test_lingbot_vla_v2_loader.py index b8cb2771..b0e01136 100644 --- a/tests/unit/models/test_lingbot_vla_v2_loader.py +++ b/tests/unit/models/test_lingbot_vla_v2_loader.py @@ -1,12 +1,15 @@ from __future__ import annotations import json +import sys from types import SimpleNamespace import pytest +import torch from telefuser.models.lingbot_vla_v2_loader import ( build_official_6b_config, + load_lingbot_vla_v2, resolve_lingbot_vla_v2_shards, validate_official_6b_checkpoint, ) @@ -62,3 +65,56 @@ def test_validate_official_6b_checkpoint_rejects_wrong_shape() -> None: def test_build_official_6b_config_rejects_non_base_variant(tmp_path) -> None: with pytest.raises(ValueError, match="Unsupported LingBot-VLA v2 checkpoint variant"): build_official_6b_config(tmp_path, checkpoint_variant="robotwin") + + +def test_public_loader_routes_official_shards_through_module_manager(tmp_path, monkeypatch) -> None: + shard_names = ["model-00002-of-00002.safetensors", "model-00001-of-00002.safetensors"] + for name in shard_names: + (tmp_path / name).write_bytes(b"") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"layer.0": shard_names[0], "layer.1": shard_names[1]}}), + encoding="utf-8", + ) + + fake_model_class = type("FakeLingBotVlaV2Model", (), {}) + monkeypatch.setitem( + sys.modules, + "telefuser.models.lingbot_vla_v2", + SimpleNamespace(LingBotVlaV2Model=fake_model_class), + ) + + class _RecordingManager: + def __init__(self) -> None: + self.load_kwargs = None + + def load_model(self, file_path, **kwargs) -> None: + self.load_kwargs = {"file_path": file_path, **kwargs} + + def fetch_module(self, name: str): + return SimpleNamespace(name=name) + + manager = _RecordingManager() + + loaded = load_lingbot_vla_v2( + manager, + tmp_path, + tmp_path / "qwen3vl", + torch_dtype=torch.bfloat16, + device="cpu", + ) + + assert loaded.name == "lingbot_vla_v2" + assert manager.load_kwargs == { + "file_path": [str(tmp_path / name) for name in sorted(shard_names)], + "device": "cpu", + "torch_dtype": torch.bfloat16, + "low_cpu_mem_usage": True, + "name": "lingbot_vla_v2", + "model_class": fake_model_class, + "model_resource": "official", + "converter_kwargs": { + "qwen3vl_path": str(tmp_path / "qwen3vl"), + "checkpoint_variant": "base", + "checkpoint_path": str(tmp_path), + }, + } diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py index 83d00c2f..264f7181 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py @@ -61,6 +61,18 @@ def test_action_chunk_is_marked_unverified() -> None: assert chunk.canonical_normalized_actions is None +def test_default_profile_loads_bundled_upstream_stats() -> None: + profile = RobotWinProfile.default() + + canonical = profile.normalize_state(torch.zeros(14)) + chunk = profile.structure_actions(torch.zeros(1, 55)) + + assert canonical.shape == (55,) + assert torch.isfinite(canonical).all() + assert chunk.raw_actions.shape == (1, 14) + assert torch.isfinite(chunk.raw_actions).all() + + def test_profile_rejects_invalid_state_and_action_shapes() -> None: profile = RobotWinProfile(_stats()) diff --git a/tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py b/tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py deleted file mode 100644 index 6f72a596..00000000 --- a/tmp/pycharm_project_0306494b/telefuser/core/base_pipeline.py +++ /dev/null @@ -1,423 +0,0 @@ -"""Base pipeline for multimodal generation.""" - -from __future__ import annotations - -import json -from abc import ABC -from datetime import datetime -from functools import wraps -from pathlib import Path -from typing import TYPE_CHECKING, Any, Sequence - -import numpy as np -import torch -import torch.nn.functional as F -from PIL import Image -from einops import rearrange - -from telefuser.utils.logging import logger - -if TYPE_CHECKING: - from telefuser.metrics import StageMetricsManager - - -class BasePipeline(ABC): - """Base pipeline for generation tasks.""" - - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - - # Wrap init method to print config after initialization - if "init" in cls.__dict__: - original_init = cls.__dict__["init"] - - @wraps(original_init) - def wrapped_init(self, *args, **kwargs): - result = original_init(self, *args, **kwargs) - if hasattr(self, "config"): - self._print_config_banner() - return result - - cls.init = wrapped_init - - # Wrap __call__ method to print parameters and reset timing registry - if "__call__" in cls.__dict__: - original_call = cls.__dict__["__call__"] - - @wraps(original_call) - def wrapped_call(self, *args, **kwargs): - from telefuser.utils.profiler import reset_timing_registry - - reset_timing_registry() - self._print_call_banner(args, kwargs) - try: - return original_call(self, *args, **kwargs) - finally: - if getattr(self, "clear_memory_after_call", True): - # Clear GPU memory after pipeline execution. - import gc - - from telefuser.platforms import current_platform - - gc.collect() - current_platform.empty_cache() - - cls.__call__ = wrapped_call - - # ANSI color codes for banner formatting - _ANSI_CYAN = "\033[36m" - _ANSI_GREEN = "\033[32m" - _ANSI_YELLOW = "\033[33m" - _ANSI_BLUE = "\033[34m" - _ANSI_DIM = "\033[2m" - _ANSI_BOLD = "\033[1m" - _ANSI_RESET = "\033[0m" - - def _get_config_defaults(self, config: Any) -> dict[str, Any]: - """Get default values for config fields. - - Args: - config: Config object (typically a dataclass) - - Returns: - Dict mapping field names to their default values - """ - defaults = {} - if hasattr(config, "__dataclass_fields__"): - from dataclasses import MISSING, fields - - for field in fields(config): - if field.default is not MISSING: - defaults[field.name] = field.default - elif field.default_factory is not MISSING: - defaults[field.name] = field.default_factory() - else: - # No default value - treat as always changed - defaults[field.name] = None - return defaults - - def _print_config_banner(self) -> None: - """Print pipeline config initialization banner with formatted output.""" - SEP = f"{self._ANSI_DIM}─{'─' * 50}─{self._ANSI_RESET}" - - lines = [ - SEP, - f"{self._ANSI_BOLD}{self._ANSI_CYAN}Pipeline Config{self._ANSI_RESET} " - f"{self._ANSI_DIM}{self.__class__.__name__}{self._ANSI_RESET}", - ] - - # Format config fields, only showing values that differ from defaults - config = self.config - defaults = self._get_config_defaults(config) - - if hasattr(config, "__dataclass_fields__"): - # Dataclass config - format each field - from dataclasses import asdict - - config_dict = asdict(config) - changed_count = 0 - for key, value in config_dict.items(): - default_value = defaults.get(key) - is_changed = default_value is None or value != default_value - if is_changed: - changed_count += 1 - formatted_value = self._format_config_value(value) - lines.append(f" {self._ANSI_DIM}{key}:{self._ANSI_RESET} {formatted_value}") - - if changed_count == 0: - lines.append(f" {self._ANSI_DIM}(all defaults){self._ANSI_RESET}") - else: - # Non-dataclass config - try to format as dict - try: - config_dict = dict(config) if hasattr(config, "items") else {} - for key, value in config_dict.items(): - formatted_value = self._format_config_value(value) - lines.append(f" {self._ANSI_DIM}{key}:{self._ANSI_RESET} {formatted_value}") - except Exception: - lines.append(f" {config}") - - lines.append(SEP) - - # Print to stderr for visibility (like logging.py) - import sys - - print("\n".join(lines), file=sys.stderr) - - def _format_config_value(self, value: Any) -> str: - """Format a config value for display. - - Args: - value: Config value to format - - Returns: - Formatted string representation - """ - if isinstance(value, str): - return f"{self._ANSI_GREEN}{value}{self._ANSI_RESET}" - elif isinstance(value, bool): - color = self._ANSI_GREEN if value else self._ANSI_YELLOW - return f"{color}{value}{self._ANSI_RESET}" - elif isinstance(value, (int, float)): - return f"{self._ANSI_BLUE}{value}{self._ANSI_RESET}" - elif isinstance(value, dict): - # Nested dict - show key count - return f"{self._ANSI_DIM}dict({len(value)} items){self._ANSI_RESET}" - elif isinstance(value, list): - return f"{self._ANSI_DIM}list({len(value)} items){self._ANSI_RESET}" - elif hasattr(value, "__class__"): - return f"{self._ANSI_DIM}{value.__class__.__name__}{self._ANSI_RESET}" - return str(value) - - def _print_call_banner(self, args: tuple, kwargs: dict) -> None: - """Print pipeline __call__ parameters banner with formatted output.""" - SEP = f"{self._ANSI_DIM}─{'─' * 50}─{self._ANSI_RESET}" - - lines = [ - SEP, - f"{self._ANSI_BOLD}{self._ANSI_YELLOW}Pipeline Call{self._ANSI_RESET} " - f"{self._ANSI_DIM}{self.__class__.__name__}{self._ANSI_RESET}", - ] - - # Format kwargs (primary parameters for pipeline calls) - if kwargs: - for key, value in kwargs.items(): - formatted_value = self._format_call_value(value) - lines.append(f" {self._ANSI_DIM}{key}:{self._ANSI_RESET} {formatted_value}") - - # Format positional args (if any) - if args: - display_args = args[1:] if args and getattr(args[0], "__class__", None) else args - if display_args: - lines.append(f" {self._ANSI_DIM}args:{self._ANSI_RESET} {self._format_call_value(display_args)}") - - if not kwargs and not args: - lines.append(f" {self._ANSI_DIM}(no parameters){self._ANSI_RESET}") - - lines.append(SEP) - - import sys - - print("\n".join(lines), file=sys.stderr) - - def _format_call_value(self, value: Any) -> str: - """Format a __call__ parameter value for display. - - Args: - value: Parameter value to format - - Returns: - Formatted string representation - """ - if isinstance(value, str): - # Truncate long strings - if len(value) > 80: - return ( - f"{self._ANSI_GREEN}{value[:80]}...{self._ANSI_RESET} " - f"{self._ANSI_DIM}({len(value)} chars){self._ANSI_RESET}" - ) - return f"{self._ANSI_GREEN}{value}{self._ANSI_RESET}" - elif isinstance(value, bool): - color = self._ANSI_GREEN if value else self._ANSI_YELLOW - return f"{color}{value}{self._ANSI_RESET}" - elif isinstance(value, (int, float)): - return f"{self._ANSI_BLUE}{value}{self._ANSI_RESET}" - elif isinstance(value, torch.Tensor): - return ( - f"{self._ANSI_CYAN}Tensor{self._ANSI_RESET}(" - f"{self._ANSI_DIM}shape={list(value.shape)}, dtype={value.dtype}{self._ANSI_RESET})" - ) - elif isinstance(value, Image.Image): - return f"{self._ANSI_CYAN}PIL.Image{self._ANSI_RESET}({self._ANSI_DIM}size={value.size}{self._ANSI_RESET})" - elif isinstance(value, list): - if len(value) > 5: - return f"{self._ANSI_DIM}list({len(value)} items){self._ANSI_RESET}" - formatted_items = [self._format_call_value(v) for v in value] - return f"[{', '.join(formatted_items)}]" - elif isinstance(value, tuple): - if len(value) > 5: - return f"{self._ANSI_DIM}tuple({len(value)} items){self._ANSI_RESET}" - formatted_items = [self._format_call_value(v) for v in value] - return f"({', '.join(formatted_items)})" - elif isinstance(value, dict): - if len(value) > 8: - return f"{self._ANSI_DIM}dict({len(value)} items){self._ANSI_RESET}" - # Recurse so inner Tensor/Image show as shape summary, not raw dump. - formatted_items = [f"{k}={self._format_call_value(v)}" for k, v in value.items()] - return "{" + ", ".join(formatted_items) + "}" - elif value is None: - return f"{self._ANSI_DIM}None{self._ANSI_RESET}" - return str(value) - - def __init__(self, device: str | torch.device, torch_dtype: torch.dtype) -> None: - self.device = device - self.torch_dtype = torch_dtype - # VAE typically requires dimensions divisible by 16 - self.height_division_factor = 16 - self.width_division_factor = 16 - self._metrics_manager: StageMetricsManager | None = None - - def _get_stages(self) -> list: - """Get list of pipeline stages for metrics collection. - - Subclasses should override this method to return their stages. - - Returns: - List of stage instances that support metrics collection. - """ - return [] - - def enable_metrics(self) -> None: - """Enable metrics collection for all pipeline stages. - - This method enables metrics tracking for stages returned by _get_stages(). - - Example: - >>> pipeline = MyPipeline.from_pretrained(...) - >>> pipeline.enable_metrics() - >>> # Run inference with metrics enabled - >>> result = pipeline(...) - >>> # Get metrics - >>> print(pipeline.get_prometheus_metrics()) - """ - if self._metrics_manager is not None: - logger.warning("Metrics already enabled, skipping") - return - - from telefuser.metrics import StageMetricsManager - - self._metrics_manager = StageMetricsManager() - stages = self._get_stages() - - if stages: - self._metrics_manager.enable_all_stages(stages) - logger.info(f"Enabled metrics for stages: {self._metrics_manager.enabled_stages}") - else: - logger.warning("No stages available for metrics collection") - - def disable_metrics(self) -> None: - """Disable metrics collection for all pipeline stages. - - Example: - >>> pipeline.disable_metrics() - """ - if self._metrics_manager is None: - return - - self._metrics_manager.disable_all_stages() - self._metrics_manager = None - logger.info("Disabled metrics for all stages") - - def get_prometheus_metrics(self) -> str: - """Get metrics in Prometheus text exposition format. - - Returns: - String containing all metrics in Prometheus format. - - Example: - >>> metrics = pipeline.get_prometheus_metrics() - >>> print(metrics) - # HELP stage_vae_duration_seconds Execution duration - # TYPE stage_vae_duration_seconds histogram - ... - - Raises: - RuntimeError: If metrics are not enabled. - """ - if self._metrics_manager is None: - raise RuntimeError("Metrics not enabled. Call enable_metrics() first.") - return self._metrics_manager.registry.get_prometheus_format() - - @property - def metrics_enabled(self) -> bool: - """Check if metrics collection is enabled.""" - return self._metrics_manager is not None - - def check_resize_height_width(self, height: int, width: int) -> tuple[int, int]: - """Ensure dimensions are divisible by division factors.""" - if height % self.height_division_factor != 0: - factor = self.height_division_factor - height = ((height + factor - 1) // factor) * factor - print(f"Height rounded up to {height}") - if width % self.width_division_factor != 0: - factor = self.width_division_factor - width = ((width + factor - 1) // factor) * factor - print(f"Width rounded up to {width}") - return height, width - - def preprocess_image(self, image: Image.Image, height: int | None = None, width: int | None = None) -> torch.Tensor: - """Preprocess PIL image to tensor.""" - if height is not None and width is not None: - if height != image.size[1] or width != image.size[0]: - image = image.resize((width, height), Image.LANCZOS) - return torch.Tensor(np.array(image, dtype=np.float32) * (2 / 255) - 1).permute(2, 0, 1).unsqueeze(0) - - def preprocess_images( - self, images: Sequence[Image.Image], height: int | None = None, width: int | None = None - ) -> list[torch.Tensor]: - """Preprocess multiple images.""" - return [self.preprocess_image(image, height, width) for image in images] - - def tensor2video( - self, frames: torch.Tensor, height: int | None = None, width: int | None = None - ) -> list[Image.Image]: - """Convert tensor to list of PIL images.""" - if height is not None and width is not None: - if height != frames.shape[2] or width != frames.shape[3]: - logger.info(f"Resizing video to {width}x{height}") - # Bicubic antialias resize does not support bf16 inputs in PyTorch. - resize_dtype = frames.dtype - resize_frames = frames.float() if frames.dtype == torch.bfloat16 else frames - frames = F.interpolate( - resize_frames, size=(height, width), mode="bicubic", align_corners=False, antialias=True - ) - frames = frames.to(resize_dtype) - frames = rearrange(frames, "C T H W -> T H W C") - frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) - return [Image.fromarray(frame) for frame in frames] - - def generate_noise( - self, shape: Sequence[int], seed: int | None = None, device: str = "cpu", dtype: torch.dtype = torch.float16 - ) -> torch.Tensor: - """Generate random noise.""" - generator = None if seed is None else torch.Generator(device).manual_seed(seed) - return torch.randn(shape, generator=generator, device=device, dtype=dtype) - - def dump_config(self, path: str | Path | None = None) -> dict[str, Any]: - """Dump pipeline configuration to dict or file. - - Captures model definition (Layer 1) and inference algorithm - parameters (Layer 2) for reproducibility and debugging. - - Args: - path: Optional file path to save config (JSON format) - - Returns: - Configuration dictionary with model and inference settings - - Example: - >>> pipeline = MyPipeline.from_pretrained(...) - >>> config = pipeline.dump_config() - >>> pipeline.dump_config("output/config.json") - """ - from .config_serializer import serialize_config - - config: dict[str, Any] = { - "version": "1.0", - "timestamp": datetime.now().isoformat(), - "pipeline_type": self.__class__.__name__, - "device": str(self.device), - "torch_dtype": str(self.torch_dtype).replace("torch.", ""), - "layer1_model_definition": { - "models": self._model_info if hasattr(self, "_model_info") else [], - }, - "layer2_inference_config": serialize_config(self.config) if hasattr(self, "config") else {}, - } - - if path: - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - json.dump(config, f, indent=2, ensure_ascii=False) - logger.info(f"Config dumped to {path}") - - return config diff --git a/tools/validation/run_lingbot_vla_v2_parity.py b/tools/validation/run_lingbot_vla_v2_parity.py new file mode 100644 index 00000000..38b8c0be --- /dev/null +++ b/tools/validation/run_lingbot_vla_v2_parity.py @@ -0,0 +1,149 @@ +"""Compare LingBot-VLA v2 upstream and TeleFuser parity artifacts. + +The artifact contract is intentionally file-based so the upstream checkout and +TeleFuser model do not need to live in one Python process. Capture both sides at +the fixed upstream commit and save ``.npz`` files with matching array keys. +""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Iterable + +import numpy as np + +UPSTREAM_REPOSITORY = "https://github.com/Robbyant/lingbot-vla-v2" +UPSTREAM_COMMIT = "be27333c9b5f2663b0ec33f069dd7dfd67fa32b5" +PREPROCESSING_KEYS = ( + "images", + "img_masks", + "image_grid_thw", + "lang_tokens", + "lang_masks", + "state", +) +FINAL_ACTION_KEYS = ("actions", "canonical_normalized_actions") + + +@dataclass(frozen=True) +class ArrayParity: + layer: str + key: str + shape: tuple[int, ...] + max_abs: float + mean_abs: float + rtol: float + atol: float + passed: bool + + +def _load_npz(path: Path) -> dict[str, np.ndarray]: + if not path.is_file(): + raise FileNotFoundError(path) + with np.load(path, allow_pickle=False) as payload: + return {key: payload[key] for key in payload.files} + + +def _velocity_keys(left: dict[str, np.ndarray], right: dict[str, np.ndarray]) -> list[str]: + prefixes = ("velocity_step_", "v_t_step_") + keys = sorted(key for key in left if key in right and key.startswith(prefixes)) + if not keys and "velocity" in left and "velocity" in right: + keys = ["velocity"] + return keys + + +def _first_present(keys: Iterable[str], left: dict[str, np.ndarray], right: dict[str, np.ndarray]) -> str | None: + for key in keys: + if key in left and key in right: + return key + return None + + +def _compare_array( + layer: str, + key: str, + expected: np.ndarray, + actual: np.ndarray, + *, + rtol: float, + atol: float, +) -> ArrayParity: + if expected.shape != actual.shape: + return ArrayParity(layer, key, tuple(actual.shape), float("inf"), float("inf"), rtol, atol, False) + diff = np.abs(expected.astype(np.float64) - actual.astype(np.float64)) + max_abs = float(diff.max()) if diff.size else 0.0 + mean_abs = float(diff.mean()) if diff.size else 0.0 + passed = bool(np.allclose(expected, actual, rtol=rtol, atol=atol)) + return ArrayParity(layer, key, tuple(actual.shape), max_abs, mean_abs, rtol, atol, passed) + + +def compare_artifacts(reference: Path, candidate: Path, *, rtol: float, atol: float) -> dict[str, object]: + expected = _load_npz(reference) + actual = _load_npz(candidate) + results: list[ArrayParity] = [] + + missing_reference = sorted(set(actual) - set(expected)) + missing_candidate = sorted(set(expected) - set(actual)) + + for key in PREPROCESSING_KEYS: + if key in expected and key in actual: + results.append(_compare_array("preprocessing", key, expected[key], actual[key], rtol=0.0, atol=0.0)) + + for key in _velocity_keys(expected, actual): + results.append(_compare_array("velocity", key, expected[key], actual[key], rtol=rtol, atol=atol)) + + action_key = _first_present(FINAL_ACTION_KEYS, expected, actual) + if action_key is not None: + results.append( + _compare_array("action", action_key, expected[action_key], actual[action_key], rtol=rtol, atol=atol) + ) + + if not any(item.layer == "preprocessing" for item in results): + raise ValueError("No shared preprocessing keys were found in the parity artifacts") + if not any(item.layer == "velocity" for item in results): + raise ValueError("No shared velocity keys were found; capture intermediate flow-matching velocity tensors") + if not any(item.layer == "action" for item in results): + raise ValueError("No shared final action key was found; expected actions or canonical_normalized_actions") + + return { + "upstream_repository": UPSTREAM_REPOSITORY, + "upstream_commit": UPSTREAM_COMMIT, + "reference": str(reference), + "candidate": str(candidate), + "passed": all(item.passed for item in results) and not missing_candidate, + "missing_reference_keys": missing_reference, + "missing_candidate_keys": missing_candidate, + "results": [asdict(item) for item in results], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--reference", + type=Path, + required=True, + help="Upstream .npz artifact captured at the pinned commit", + ) + parser.add_argument("--candidate", type=Path, required=True, help="TeleFuser .npz artifact from the same inputs") + parser.add_argument("--output", type=Path, default=None, help="Optional JSON report path") + parser.add_argument("--rtol", type=float, default=1e-3) + parser.add_argument("--atol", type=float, default=1e-3) + args = parser.parse_args() + + report = compare_artifacts(args.reference, args.candidate, rtol=args.rtol, atol=args.atol) + payload = json.dumps(report, indent=2, sort_keys=True) + if args.output is None: + print(payload) + else: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload + "\n", encoding="utf-8") + if not report["passed"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() From 024e782d6960ef33607999a7fe64db214ea91517 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Tue, 4 Aug 2026 03:22:27 +0000 Subject: [PATCH 04/15] test(vla): add layered inference regression baseline Capture LingBot-VLA v2 preprocessing, fixed noise, per-step flow state and velocity, and final canonical actions through the public TeleFuser pipeline. Record checkpoint, processor, input, runtime, shape, and dtype identities alongside each artifact. Harden the parity comparator to require complete artifact contracts, contiguous denoising steps, finite values, matching metadata, and strict or portable tolerances with first-failure reporting. Add CPU regression tests and document the local baseline workflow. Verification: 25 focused VLA tests passed in .venv-vla; ruff check and format checks passed; two real 6B H100 captures matched across 38 comparisons with max_abs=0.0. --- examples/lingbot_vla_v2/README.md | 35 ++ .../test_lingbot_vla_v2_artifacts.py | 173 +++++++++ .../capture_lingbot_vla_v2_telefuser.py | 294 +++++++++++++++ tools/validation/run_lingbot_vla_v2_parity.py | 335 +++++++++++++++--- 4 files changed, 781 insertions(+), 56 deletions(-) create mode 100644 tests/unit/validation/test_lingbot_vla_v2_artifacts.py create mode 100644 tools/validation/capture_lingbot_vla_v2_telefuser.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index bf19e7d3..3c408eac 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -44,3 +44,38 @@ python examples/lingbot_vla_v2/lingbot_vla_v2_inference.py \ The example saves canonical actions and checkpoint metadata in an `.npz` file. The base output must not be sent to a robot without an embodiment-specific post-training checkpoint, action mapping, and policy validation. + +## TeleFuser Regression Baseline + +The validation capture runs through the public loader and pipeline, then records preprocessing tensors, fixed initial +noise, every flow-matching `x_t` and velocity step, and the final canonical action. Run it twice before changing VLA +model code to establish and verify a strict local baseline: + +```bash +.venv-vla/bin/python tools/validation/capture_lingbot_vla_v2_telefuser.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --camera-high /data/cam_high.png \ + --camera-left-wrist /data/cam_left_wrist.png \ + --camera-right-wrist /data/cam_right_wrist.png \ + --task "pick up the red block" \ + --state-json '[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]' \ + --seed 7 \ + --output work_dirs/vla_regression/baseline_seed7.npz + +# Repeat the same command with: +# --output work_dirs/vla_regression/replay_seed7.npz + +.venv-vla/bin/python tools/validation/run_lingbot_vla_v2_parity.py \ + --reference work_dirs/vla_regression/baseline_seed7.npz \ + --candidate work_dirs/vla_regression/replay_seed7.npz \ + --profile strict \ + --output work_dirs/vla_regression/strict_report.json +``` + +Each `.npz` has a same-name `.json` sidecar containing the checkpoint, processor, input, runtime, and tensor contract +metadata. The default checkpoint identity is a fast filename-and-size manifest. Add `--full-checkpoint-hash` when a +content hash of every checkpoint shard is required. Keep generated artifacts under `work_dirs`; do not commit them. + +This is a TeleFuser regression check, not upstream parity. It detects changes to the current implementation but does +not establish equivalence with the official repository. diff --git a/tests/unit/validation/test_lingbot_vla_v2_artifacts.py b/tests/unit/validation/test_lingbot_vla_v2_artifacts.py new file mode 100644 index 00000000..fc735301 --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_artifacts.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tools.validation.capture_lingbot_vla_v2_telefuser import TensorCapture, trace_predict_velocity +from tools.validation.run_lingbot_vla_v2_parity import compare_artifacts + + +def _arrays() -> dict[str, np.ndarray]: + arrays = { + "images": np.zeros((1, 3, 2, 2), dtype=np.float32), + "img_masks": np.ones((1, 3), dtype=np.bool_), + "image_grid_thw": np.ones((1, 3, 3), dtype=np.int64), + "lang_tokens": np.arange(4, dtype=np.int64).reshape(1, 4), + "lang_masks": np.ones((1, 4), dtype=np.bool_), + "state": np.zeros((1, 55), dtype=np.float32), + "initial_noise": np.ones((1, 2, 55), dtype=np.float32), + "canonical_normalized_actions": np.zeros((2, 55), dtype=np.float32), + } + for step in range(2): + suffix = f"{step:02d}" + arrays[f"timestep_step_{suffix}"] = np.asarray([1.0 - 0.5 * step], dtype=np.float32) + arrays[f"x_t_step_{suffix}"] = np.full((1, 2, 55), step, dtype=np.float32) + arrays[f"velocity_step_{suffix}"] = np.full((1, 2, 55), step + 0.25, dtype=np.float32) + return arrays + + +def _metadata() -> dict[str, object]: + return { + "schema_version": 1, + "artifact_kind": "telefuser_regression", + "checkpoint_manifest_sha256": "checkpoint", + "processor_manifest_sha256": "processor", + "input_sha256": "input", + "seed": 7, + "num_steps": 2, + "torch_dtype": "bfloat16", + "attention_backend": "eager", + } + + +def _write_artifact( + root: Path, + name: str, + arrays: dict[str, np.ndarray], + metadata: dict[str, object] | None = None, +) -> Path: + path = root / f"{name}.npz" + np.savez(path, **arrays) + payload = dict(_metadata() if metadata is None else metadata) + payload["arrays"] = { + key: { + "shape": list(array.shape), + "original_dtype": str(array.dtype), + "stored_dtype": str(array.dtype), + } + for key, array in arrays.items() + } + path.with_suffix(".json").write_text( + json.dumps(payload), + encoding="utf-8", + ) + return path + + +def test_compare_artifacts_accepts_a_complete_strict_replay(tmp_path: Path) -> None: + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", _arrays()) + + report = compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + assert report["passed"] is True + assert report["first_failed_step"] is None + assert len(report["results"]) == 14 + + +def test_compare_artifacts_requires_every_preprocessing_array(tmp_path: Path) -> None: + arrays = _arrays() + del arrays["image_grid_thw"] + reference = _write_artifact(tmp_path, "reference", arrays) + candidate = _write_artifact(tmp_path, "candidate", _arrays()) + + with pytest.raises(ValueError, match="image_grid_thw"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + +def test_compare_artifacts_requires_contiguous_sampling_steps(tmp_path: Path) -> None: + arrays = _arrays() + del arrays["velocity_step_01"] + reference = _write_artifact(tmp_path, "reference", arrays) + candidate = _write_artifact(tmp_path, "candidate", _arrays()) + + with pytest.raises(ValueError, match="velocity steps"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + +def test_compare_artifacts_reports_the_first_failed_step(tmp_path: Path) -> None: + candidate_arrays = _arrays() + candidate_arrays["velocity_step_01"][0, 0, 0] += 0.5 + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", candidate_arrays) + + report = compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + assert report["passed"] is False + assert report["first_failed_step"] == 1 + failed = [item for item in report["results"] if not item["passed"]] + assert [item["key"] for item in failed] == ["velocity_step_01"] + assert failed[0]["mismatch_count"] == 1 + + +def test_compare_artifacts_rejects_non_finite_values(tmp_path: Path) -> None: + candidate_arrays = _arrays() + candidate_arrays["x_t_step_00"][0, 0, 0] = np.nan + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", candidate_arrays) + + report = compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + assert report["passed"] is False + assert report["first_failed_step"] == 0 + + +def test_compare_artifacts_rejects_different_artifact_identity(tmp_path: Path) -> None: + candidate_metadata = _metadata() + candidate_metadata["input_sha256"] = "different" + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", _arrays(), candidate_metadata) + + with pytest.raises(ValueError, match="input_sha256"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + +def test_velocity_trace_snapshots_inputs_and_restores_the_model_instance() -> None: + class _FlowModel: + _use_compile_predict_velocity = True + + def predict_velocity(self, state, prefix_masks, cache, x_t, timestep, **kwargs): + del state, prefix_masks, cache, timestep, kwargs + return torch.ones_like(x_t) + + flow_model = _FlowModel() + capture = TensorCapture() + x_t = torch.zeros(1, 2, 3) + + with trace_predict_velocity(flow_model, capture) as trace: + velocity = flow_model.predict_velocity(None, None, None, x_t, torch.ones(1)) + x_t.add_(velocity) + + assert trace.step == 1 + assert np.array_equal(capture.arrays["initial_noise"], np.zeros((1, 2, 3), dtype=np.float32)) + assert np.array_equal(capture.arrays["x_t_step_00"], np.zeros((1, 2, 3), dtype=np.float32)) + assert np.array_equal(capture.arrays["velocity_step_00"], np.ones((1, 2, 3), dtype=np.float32)) + assert "predict_velocity" not in vars(flow_model) + assert flow_model._use_compile_predict_velocity is True + + +def test_tensor_capture_records_original_bfloat16_dtype() -> None: + capture = TensorCapture() + + capture.add("value", torch.ones(2, dtype=torch.bfloat16)) + + assert capture.arrays["value"].dtype == np.float32 + assert capture.array_metadata["value"] == { + "shape": [2], + "original_dtype": "bfloat16", + "stored_dtype": "float32", + } diff --git a/tools/validation/capture_lingbot_vla_v2_telefuser.py b/tools/validation/capture_lingbot_vla_v2_telefuser.py new file mode 100644 index 00000000..2ff3b1f6 --- /dev/null +++ b/tools/validation/capture_lingbot_vla_v2_telefuser.py @@ -0,0 +1,294 @@ +"""Capture a layered TeleFuser LingBot-VLA v2 regression artifact.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, Sequence + +import numpy as np +import torch +import transformers +from transformers import AutoProcessor + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2, resolve_lingbot_vla_v2_shards +from telefuser.pipelines.lingbot_vla_v2 import ( + ROBOTWIN_CAMERA_KEYS, + LingBotVlaV2Observation, + LingBotVlaV2Pipeline, + LingBotVlaV2PipelineConfig, +) + +ARTIFACT_SCHEMA_VERSION = 1 + + +def _sha256_file(path: Path, digest: Any | None = None) -> str: + result = hashlib.sha256() if digest is None else digest + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + result.update(block) + return result.hexdigest() + + +def _manifest_sha256(paths: Sequence[Path], *, include_contents: bool) -> str: + digest = hashlib.sha256() + for path in sorted((item.resolve() for item in paths), key=lambda item: item.name): + stat = path.stat() + digest.update(path.name.encode("utf-8")) + digest.update(str(stat.st_size).encode("ascii")) + if include_contents: + _sha256_file(path, digest) + return digest.hexdigest() + + +def _processor_files(root: Path) -> list[Path]: + return sorted(path for path in root.iterdir() if path.is_file() and path.suffix != ".safetensors") + + +def _input_sha256(task: str, state: Sequence[float], image_paths: Sequence[Path]) -> str: + digest = hashlib.sha256() + canonical = json.dumps({"task": task, "state": list(state)}, sort_keys=True, separators=(",", ":")) + digest.update(canonical.encode("utf-8")) + for path in image_paths: + digest.update(path.name.encode("utf-8")) + _sha256_file(path, digest) + return digest.hexdigest() + + +def _git_commit() -> str: + repository_root = Path(__file__).resolve().parents[2] + completed = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repository_root, + check=True, + capture_output=True, + text=True, + ) + return completed.stdout.strip() + + +class TensorCapture: + """Own CPU snapshots and their original tensor contracts.""" + + def __init__(self) -> None: + self.arrays: dict[str, np.ndarray] = {} + self.array_metadata: dict[str, dict[str, object]] = {} + + def add(self, key: str, tensor: torch.Tensor) -> None: + if key in self.arrays: + raise ValueError(f"Duplicate capture key: {key}") + snapshot = tensor.detach().cpu().clone() + stored = snapshot.float() if snapshot.is_floating_point() else snapshot + array = stored.numpy() + self.arrays[key] = array + self.array_metadata[key] = { + "shape": list(snapshot.shape), + "original_dtype": str(snapshot.dtype).removeprefix("torch."), + "stored_dtype": str(array.dtype), + } + + +class VelocityTrace: + """Capture the state around each call to the flow-matching velocity model.""" + + def __init__(self, capture: TensorCapture) -> None: + self.capture = capture + self.step = 0 + + def record(self, original: Any, *args: Any, **kwargs: Any) -> torch.Tensor: + if len(args) < 5: + raise RuntimeError("LingBot-VLA v2 predict_velocity trace received an unexpected call signature") + x_t = args[3] + timestep = args[4] + suffix = f"{self.step:02d}" + if self.step == 0: + self.capture.add("initial_noise", x_t) + self.capture.add(f"timestep_step_{suffix}", timestep) + self.capture.add(f"x_t_step_{suffix}", x_t) + velocity = original(*args, **kwargs) + self.capture.add(f"velocity_step_{suffix}", velocity) + self.step += 1 + return velocity + + +@contextmanager +def trace_predict_velocity(flow_model: Any, capture: TensorCapture) -> Iterator[VelocityTrace]: + """Temporarily trace one model instance without changing global classes.""" + original = flow_model.predict_velocity + had_instance_override = "predict_velocity" in vars(flow_model) + previous_override = vars(flow_model).get("predict_velocity") + compile_enabled = bool(getattr(flow_model, "_use_compile_predict_velocity", False)) + trace = VelocityTrace(capture) + + flow_model._use_compile_predict_velocity = False + flow_model.predict_velocity = lambda *args, **kwargs: trace.record(original, *args, **kwargs) + try: + yield trace + finally: + if had_instance_override: + flow_model.predict_velocity = previous_override + else: + del flow_model.predict_velocity + flow_model._use_compile_predict_velocity = compile_enabled + + +def _build_pipeline(model_root: Path, qwen3vl_root: Path, device: str) -> LingBotVlaV2Pipeline: + target_device = torch.device(device) + dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 + processor = AutoProcessor.from_pretrained(str(qwen3vl_root), local_files_only=True, padding_side="right") + manager = ModuleManager(torch_dtype=dtype, device="cpu") + manager.add_module(processor, "lingbot_vla_v2_processor", path=str(qwen3vl_root)) + load_lingbot_vla_v2(manager, model_root, qwen3vl_root, torch_dtype=dtype) + pipeline = LingBotVlaV2Pipeline(device=device, torch_dtype=dtype) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig( + device_type=target_device.type, + device_id=target_device.index or 0, + torch_dtype=dtype, + ) + ), + ) + return pipeline + + +def capture_artifact( + *, + model_root: Path, + qwen3vl_root: Path, + image_paths: Sequence[Path], + task: str, + state: Sequence[float], + seed: int, + output: Path, + device: str, + full_checkpoint_hash: bool, +) -> tuple[Path, Path]: + if len(image_paths) != len(ROBOTWIN_CAMERA_KEYS): + raise ValueError(f"expected {len(ROBOTWIN_CAMERA_KEYS)} camera paths, got {len(image_paths)}") + output = output.with_suffix(".npz") + metadata_path = output.with_suffix(".json") + output.parent.mkdir(parents=True, exist_ok=True) + + pipeline = _build_pipeline(model_root, qwen3vl_root, device) + capture = TensorCapture() + try: + observation = LingBotVlaV2Observation( + task=task, + state=state, + images=dict(zip(ROBOTWIN_CAMERA_KEYS, image_paths, strict=True)), + ) + inputs = pipeline.input_processor.prepare(observation) + for key in ("images", "img_masks", "image_grid_thw", "lang_tokens", "lang_masks", "state"): + capture.add(key, getattr(inputs, key)) + + flow_model = pipeline.policy_stage.policy.model + with trace_predict_velocity(flow_model, capture) as trace: + chunk = pipeline.predict(inputs, seed=seed) + capture.add("canonical_normalized_actions", chunk.canonical_normalized_actions) + + expected_steps = int(flow_model.config.num_steps) + if trace.step != expected_steps: + raise RuntimeError(f"captured {trace.step} denoising steps, expected {expected_steps}") + + target_device = torch.device(device) + checkpoint_paths = [Path(path) for path in resolve_lingbot_vla_v2_shards(model_root)] + metadata = { + "schema_version": ARTIFACT_SCHEMA_VERSION, + "artifact_kind": "telefuser_regression", + "telefuser_commit": _git_commit(), + "checkpoint_manifest_sha256": _manifest_sha256( + checkpoint_paths, + include_contents=full_checkpoint_hash, + ), + "checkpoint_hash_mode": "full_sha256" if full_checkpoint_hash else "filename_and_size", + "processor_manifest_sha256": _manifest_sha256( + _processor_files(qwen3vl_root), + include_contents=True, + ), + "input_sha256": _input_sha256(task, state, image_paths), + "seed": seed, + "num_steps": trace.step, + "torch_dtype": str(pipeline.torch_dtype).removeprefix("torch."), + "attention_backend": str(flow_model.config.attention_implementation), + "device": str(target_device), + "device_name": torch.cuda.get_device_name(target_device) if target_device.type == "cuda" else "cpu", + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + "arrays": capture.array_metadata, + } + np.savez(output, **capture.arrays) + metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + finally: + pipeline.close() + + return output, metadata_path + + +def _parse_state(value: str) -> list[float]: + try: + state = json.loads(value) + except json.JSONDecodeError as error: + raise argparse.ArgumentTypeError("state-json must be valid JSON") from error + if not isinstance(state, list) or len(state) != 14 or any(isinstance(item, bool) for item in state): + raise argparse.ArgumentTypeError("state-json must be a 14-element numeric JSON list") + try: + return [float(item) for item in state] + except (TypeError, ValueError) as error: + raise argparse.ArgumentTypeError("state-json must contain only numeric values") from error + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", required=True, type=Path) + parser.add_argument("--qwen3vl-root", required=True, type=Path) + parser.add_argument("--camera-high", required=True, type=Path) + parser.add_argument("--camera-left-wrist", required=True, type=Path) + parser.add_argument("--camera-right-wrist", required=True, type=Path) + parser.add_argument("--task", required=True) + parser.add_argument("--state-json", required=True, type=_parse_state) + parser.add_argument("--seed", required=True, type=int) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--full-checkpoint-hash", + action="store_true", + help="Hash all checkpoint bytes instead of the faster filename-and-size manifest", + ) + args = parser.parse_args() + + paths = ( + args.model_root, + args.qwen3vl_root, + args.camera_high, + args.camera_left_wrist, + args.camera_right_wrist, + ) + missing = [str(path) for path in paths if not path.exists()] + if missing: + parser.error(f"input paths do not exist: {missing}") + + artifact, metadata = capture_artifact( + model_root=args.model_root, + qwen3vl_root=args.qwen3vl_root, + image_paths=(args.camera_high, args.camera_left_wrist, args.camera_right_wrist), + task=args.task, + state=args.state_json, + seed=args.seed, + output=args.output, + device=args.device, + full_checkpoint_hash=args.full_checkpoint_hash, + ) + print(f"Saved LingBot-VLA v2 capture: {artifact}") + print(f"Saved LingBot-VLA v2 metadata: {metadata}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/run_lingbot_vla_v2_parity.py b/tools/validation/run_lingbot_vla_v2_parity.py index 38b8c0be..97c1c514 100644 --- a/tools/validation/run_lingbot_vla_v2_parity.py +++ b/tools/validation/run_lingbot_vla_v2_parity.py @@ -1,22 +1,24 @@ -"""Compare LingBot-VLA v2 upstream and TeleFuser parity artifacts. +"""Compare layered LingBot-VLA v2 capture artifacts. -The artifact contract is intentionally file-based so the upstream checkout and -TeleFuser model do not need to live in one Python process. Capture both sides at -the fixed upstream commit and save ``.npz`` files with matching array keys. +The same comparator is used for local TeleFuser regression artifacts and for +future upstream parity artifacts. Captures stay file based so implementations +with incompatible Python dependencies never need to share a process. """ from __future__ import annotations import argparse import json +import re from dataclasses import asdict, dataclass from pathlib import Path -from typing import Iterable +from typing import Any import numpy as np UPSTREAM_REPOSITORY = "https://github.com/Robbyant/lingbot-vla-v2" UPSTREAM_COMMIT = "be27333c9b5f2663b0ec33f069dd7dfd67fa32b5" +ARTIFACT_SCHEMA_VERSION = 1 PREPROCESSING_KEYS = ( "images", "img_masks", @@ -25,7 +27,18 @@ "lang_masks", "state", ) -FINAL_ACTION_KEYS = ("actions", "canonical_normalized_actions") +STEP_LAYERS = ("timestep", "x_t", "velocity") +FINAL_ACTION_KEYS = ("canonical_normalized_actions", "actions") +IDENTITY_METADATA_KEYS = ( + "checkpoint_manifest_sha256", + "processor_manifest_sha256", + "input_sha256", + "seed", + "num_steps", + "torch_dtype", + "attention_backend", +) +_STEP_KEY = re.compile(r"^(timestep|x_t|velocity)_step_([0-9]+)$") @dataclass(frozen=True) @@ -33,8 +46,11 @@ class ArrayParity: layer: str key: str shape: tuple[int, ...] + expected_dtype: str + actual_dtype: str max_abs: float mean_abs: float + mismatch_count: int rtol: float atol: float passed: bool @@ -47,19 +63,81 @@ def _load_npz(path: Path) -> dict[str, np.ndarray]: return {key: payload[key] for key in payload.files} -def _velocity_keys(left: dict[str, np.ndarray], right: dict[str, np.ndarray]) -> list[str]: - prefixes = ("velocity_step_", "v_t_step_") - keys = sorted(key for key in left if key in right and key.startswith(prefixes)) - if not keys and "velocity" in left and "velocity" in right: - keys = ["velocity"] - return keys +def _metadata_path(artifact_path: Path, explicit_path: Path | None) -> Path: + return explicit_path if explicit_path is not None else artifact_path.with_suffix(".json") -def _first_present(keys: Iterable[str], left: dict[str, np.ndarray], right: dict[str, np.ndarray]) -> str | None: - for key in keys: - if key in left and key in right: - return key - return None +def _load_metadata(artifact_path: Path, explicit_path: Path | None = None) -> dict[str, Any]: + path = _metadata_path(artifact_path, explicit_path) + if not path.is_file(): + raise FileNotFoundError(f"Missing LingBot-VLA v2 artifact metadata: {path}") + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise ValueError(f"Artifact metadata must be a JSON object: {path}") + if payload.get("schema_version") != ARTIFACT_SCHEMA_VERSION: + raise ValueError( + f"Unsupported artifact schema in {path}: expected {ARTIFACT_SCHEMA_VERSION}, " + f"got {payload.get('schema_version')!r}" + ) + missing = [key for key in IDENTITY_METADATA_KEYS if key not in payload] + if missing: + raise ValueError(f"Artifact metadata {path} is missing identity fields: {missing}") + return payload + + +def _step_keys(arrays: dict[str, np.ndarray]) -> dict[str, dict[int, str]]: + result: dict[str, dict[int, str]] = {layer: {} for layer in STEP_LAYERS} + for key in arrays: + match = _STEP_KEY.fullmatch(key) + if match is not None: + layer, step_text = match.groups() + step = int(step_text) + if step in result[layer]: + raise ValueError(f"Duplicate {layer} capture for step {step}") + result[layer][step] = key + return result + + +def _validate_contract(arrays: dict[str, np.ndarray], metadata: dict[str, Any], *, side: str) -> None: + required = (*PREPROCESSING_KEYS, "initial_noise") + missing = [key for key in required if key not in arrays] + if missing: + raise ValueError(f"{side} artifact is missing required arrays: {missing}") + + if not any(key in arrays for key in FINAL_ACTION_KEYS): + raise ValueError(f"{side} artifact is missing a final action array: {FINAL_ACTION_KEYS}") + + array_metadata = metadata.get("arrays") + if not isinstance(array_metadata, dict): + raise ValueError(f"{side} metadata must contain an arrays contract") + missing_contracts = sorted(set(arrays) - set(array_metadata)) + if missing_contracts: + raise ValueError(f"{side} metadata is missing array contracts: {missing_contracts}") + for key, array in arrays.items(): + contract = array_metadata[key] + if not isinstance(contract, dict): + raise ValueError(f"{side} metadata contract for {key} must be an object") + expected_contract = { + "shape": list(array.shape), + "stored_dtype": str(array.dtype), + } + mismatches = { + field: {"metadata": contract.get(field), "artifact": value} + for field, value in expected_contract.items() + if contract.get(field) != value + } + if "original_dtype" not in contract: + mismatches["original_dtype"] = {"metadata": None, "artifact": "required"} + if mismatches: + raise ValueError(f"{side} metadata contract for {key} does not match the artifact: {mismatches}") + + num_steps = metadata["num_steps"] + if not isinstance(num_steps, int) or isinstance(num_steps, bool) or num_steps <= 0: + raise ValueError(f"{side} metadata num_steps must be a positive integer, got {num_steps!r}") + expected_steps = set(range(num_steps)) + for layer, keys in _step_keys(arrays).items(): + if set(keys) != expected_steps: + raise ValueError(f"{side} artifact {layer} steps must be {sorted(expected_steps)}, got {sorted(keys)}") def _compare_array( @@ -70,71 +148,216 @@ def _compare_array( *, rtol: float, atol: float, + expected_original_dtype: str, + actual_original_dtype: str, ) -> ArrayParity: - if expected.shape != actual.shape: - return ArrayParity(layer, key, tuple(actual.shape), float("inf"), float("inf"), rtol, atol, False) - diff = np.abs(expected.astype(np.float64) - actual.astype(np.float64)) - max_abs = float(diff.max()) if diff.size else 0.0 - mean_abs = float(diff.mean()) if diff.size else 0.0 - passed = bool(np.allclose(expected, actual, rtol=rtol, atol=atol)) - return ArrayParity(layer, key, tuple(actual.shape), max_abs, mean_abs, rtol, atol, passed) + expected_dtype = expected_original_dtype + actual_dtype = actual_original_dtype + shape_matches = expected.shape == actual.shape + dtype_matches = expected.dtype == actual.dtype and expected_original_dtype == actual_original_dtype + finite = True + if np.issubdtype(expected.dtype, np.number) and not np.isfinite(expected).all(): + finite = False + if np.issubdtype(actual.dtype, np.number) and not np.isfinite(actual).all(): + finite = False + + if not shape_matches or not finite: + max_abs = float("inf") + mean_abs = float("inf") + mismatch_count = max(expected.size, actual.size) + values_match = False + elif expected.dtype == np.bool_ or np.issubdtype(expected.dtype, np.integer): + difference = np.abs(expected.astype(np.int64) - actual.astype(np.int64)) + mismatch_count = int(np.count_nonzero(difference)) + max_abs = float(difference.max()) if difference.size else 0.0 + mean_abs = float(difference.mean()) if difference.size else 0.0 + values_match = mismatch_count == 0 + else: + difference = np.abs(expected.astype(np.float64) - actual.astype(np.float64)) + close = np.isclose(expected, actual, rtol=rtol, atol=atol, equal_nan=False) + mismatch_count = int(np.count_nonzero(~close)) + max_abs = float(difference.max()) if difference.size else 0.0 + mean_abs = float(difference.mean()) if difference.size else 0.0 + values_match = mismatch_count == 0 + + return ArrayParity( + layer=layer, + key=key, + shape=tuple(actual.shape), + expected_dtype=expected_dtype, + actual_dtype=actual_dtype, + max_abs=max_abs, + mean_abs=mean_abs, + mismatch_count=mismatch_count, + rtol=rtol, + atol=atol, + passed=shape_matches and dtype_matches and finite and values_match, + ) + + +def _action_key(arrays: dict[str, np.ndarray]) -> str: + for key in FINAL_ACTION_KEYS: + if key in arrays: + return key + raise ValueError(f"No final action key found: {FINAL_ACTION_KEYS}") + +def _original_dtype(metadata: dict[str, Any], key: str) -> str: + return str(metadata["arrays"][key]["original_dtype"]) -def compare_artifacts(reference: Path, candidate: Path, *, rtol: float, atol: float) -> dict[str, object]: + +def compare_artifacts( + reference: Path, + candidate: Path, + *, + rtol: float, + atol: float, + reference_metadata: Path | None = None, + candidate_metadata: Path | None = None, +) -> dict[str, object]: expected = _load_npz(reference) actual = _load_npz(candidate) - results: list[ArrayParity] = [] + expected_metadata = _load_metadata(reference, reference_metadata) + actual_metadata = _load_metadata(candidate, candidate_metadata) + _validate_contract(expected, expected_metadata, side="reference") + _validate_contract(actual, actual_metadata, side="candidate") - missing_reference = sorted(set(actual) - set(expected)) - missing_candidate = sorted(set(expected) - set(actual)) + metadata_mismatches = { + key: {"reference": expected_metadata[key], "candidate": actual_metadata[key]} + for key in IDENTITY_METADATA_KEYS + if expected_metadata[key] != actual_metadata[key] + } + if metadata_mismatches: + raise ValueError(f"Artifact identity metadata does not match: {metadata_mismatches}") + results: list[ArrayParity] = [] for key in PREPROCESSING_KEYS: - if key in expected and key in actual: - results.append(_compare_array("preprocessing", key, expected[key], actual[key], rtol=0.0, atol=0.0)) + results.append( + _compare_array( + "preprocessing", + key, + expected[key], + actual[key], + rtol=0.0, + atol=0.0, + expected_original_dtype=_original_dtype(expected_metadata, key), + actual_original_dtype=_original_dtype(actual_metadata, key), + ) + ) + results.append( + _compare_array( + "noise", + "initial_noise", + expected["initial_noise"], + actual["initial_noise"], + rtol=0.0, + atol=0.0, + expected_original_dtype=_original_dtype(expected_metadata, "initial_noise"), + actual_original_dtype=_original_dtype(actual_metadata, "initial_noise"), + ) + ) - for key in _velocity_keys(expected, actual): - results.append(_compare_array("velocity", key, expected[key], actual[key], rtol=rtol, atol=atol)) + expected_steps = _step_keys(expected) + actual_steps = _step_keys(actual) + for step in range(expected_metadata["num_steps"]): + for layer in STEP_LAYERS: + expected_key = expected_steps[layer][step] + actual_key = actual_steps[layer][step] + layer_rtol = 0.0 if layer == "timestep" else rtol + layer_atol = 0.0 if layer == "timestep" else atol + results.append( + _compare_array( + layer, + expected_key, + expected[expected_key], + actual[actual_key], + rtol=layer_rtol, + atol=layer_atol, + expected_original_dtype=_original_dtype(expected_metadata, expected_key), + actual_original_dtype=_original_dtype(actual_metadata, actual_key), + ) + ) - action_key = _first_present(FINAL_ACTION_KEYS, expected, actual) - if action_key is not None: - results.append( - _compare_array("action", action_key, expected[action_key], actual[action_key], rtol=rtol, atol=atol) + expected_action_key = _action_key(expected) + actual_action_key = _action_key(actual) + results.append( + _compare_array( + "action", + expected_action_key, + expected[expected_action_key], + actual[actual_action_key], + rtol=rtol, + atol=atol, + expected_original_dtype=_original_dtype(expected_metadata, expected_action_key), + actual_original_dtype=_original_dtype(actual_metadata, actual_action_key), ) + ) - if not any(item.layer == "preprocessing" for item in results): - raise ValueError("No shared preprocessing keys were found in the parity artifacts") - if not any(item.layer == "velocity" for item in results): - raise ValueError("No shared velocity keys were found; capture intermediate flow-matching velocity tensors") - if not any(item.layer == "action" for item in results): - raise ValueError("No shared final action key was found; expected actions or canonical_normalized_actions") + expected_compared_keys = { + *PREPROCESSING_KEYS, + "initial_noise", + expected_action_key, + *(key for layer in expected_steps.values() for key in layer.values()), + } + actual_compared_keys = { + *PREPROCESSING_KEYS, + "initial_noise", + actual_action_key, + *(key for layer in actual_steps.values() for key in layer.values()), + } + unexpected_reference = sorted(set(expected) - expected_compared_keys) + unexpected_candidate = sorted(set(actual) - actual_compared_keys) + first_failed_step = next( + ( + int(match.group(2)) + for item in results + if not item.passed and (match := _STEP_KEY.fullmatch(item.key)) is not None + ), + None, + ) return { + "schema_version": ARTIFACT_SCHEMA_VERSION, "upstream_repository": UPSTREAM_REPOSITORY, "upstream_commit": UPSTREAM_COMMIT, "reference": str(reference), "candidate": str(candidate), - "passed": all(item.passed for item in results) and not missing_candidate, - "missing_reference_keys": missing_reference, - "missing_candidate_keys": missing_candidate, + "reference_kind": expected_metadata.get("artifact_kind"), + "candidate_kind": actual_metadata.get("artifact_kind"), + "passed": all(item.passed for item in results), + "first_failed_step": first_failed_step, + "unexpected_reference_keys": unexpected_reference, + "unexpected_candidate_keys": unexpected_candidate, "results": [asdict(item) for item in results], } def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--reference", - type=Path, - required=True, - help="Upstream .npz artifact captured at the pinned commit", - ) - parser.add_argument("--candidate", type=Path, required=True, help="TeleFuser .npz artifact from the same inputs") - parser.add_argument("--output", type=Path, default=None, help="Optional JSON report path") - parser.add_argument("--rtol", type=float, default=1e-3) - parser.add_argument("--atol", type=float, default=1e-3) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--reference-metadata", type=Path, default=None) + parser.add_argument("--candidate-metadata", type=Path, default=None) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument("--profile", choices=("strict", "portable"), default="strict") + parser.add_argument("--rtol", type=float, default=None) + parser.add_argument("--atol", type=float, default=None) args = parser.parse_args() - report = compare_artifacts(args.reference, args.candidate, rtol=args.rtol, atol=args.atol) + default_tolerance = 0.0 if args.profile == "strict" else 1e-3 + rtol = default_tolerance if args.rtol is None else args.rtol + atol = default_tolerance if args.atol is None else args.atol + if rtol < 0 or atol < 0: + parser.error("rtol and atol must be non-negative") + + report = compare_artifacts( + args.reference, + args.candidate, + rtol=rtol, + atol=atol, + reference_metadata=args.reference_metadata, + candidate_metadata=args.candidate_metadata, + ) payload = json.dumps(report, indent=2, sort_keys=True) if args.output is None: print(payload) From 8d03aa0a1614defd94a5723b03f8e8d07c47f983 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Tue, 4 Aug 2026 06:41:35 +0000 Subject: [PATCH 05/15] test(vla): establish official upstream parity Pin the LingBot-VLA v2 reference checkout and add an isolated upstream capture runtime for preprocessing, velocity, and final-action comparison. Commit the official RobotWin normalization statistics, align public image preprocessing and normalization precision, and record attention, MoE, and norm-stat identities in parity artifacts. Keep the upstream Triton MoE path enabled for production while providing a deterministic reference backend for strict cross-process comparison. Verification: .venv-vla/bin/python -m pytest tests/unit/models/test_lingbot_vla_v2_loader.py tests/unit/pipelines/lingbot_vla_v2/test_data.py tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py tests/unit/validation/test_lingbot_vla_v2_artifacts.py -q (24 passed); ruff check and ruff format --check passed for the changed VLA pipeline, tests, and validation scripts; git diff --cached --check passed. --- examples/lingbot_vla_v2/README.md | 31 + telefuser/models/lingbot_vla_v2_loader.py | 2 +- .../assets/robotwin_norm_stats.json | 240 +++++++- telefuser/pipelines/lingbot_vla_v2/data.py | 10 +- .../pipelines/lingbot_vla_v2/robot_profile.py | 10 +- .../pipelines/lingbot_vla_v2/test_data.py | 10 +- .../lingbot_vla_v2/test_robot_profile.py | 6 +- .../test_lingbot_vla_v2_artifacts.py | 12 + .../capture_lingbot_vla_v2_telefuser.py | 16 + .../capture_lingbot_vla_v2_upstream.py | 541 ++++++++++++++++++ .../requirements-lingbot-vla-v2-upstream.txt | 17 + tools/validation/run_lingbot_vla_v2_parity.py | 2 + 12 files changed, 859 insertions(+), 38 deletions(-) create mode 100644 tools/validation/capture_lingbot_vla_v2_upstream.py create mode 100644 tools/validation/requirements-lingbot-vla-v2-upstream.txt diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 3c408eac..b3ff9000 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -61,6 +61,7 @@ model code to establish and verify a strict local baseline: --task "pick up the red block" \ --state-json '[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]' \ --seed 7 \ + --deterministic-moe \ --output work_dirs/vla_regression/baseline_seed7.npz # Repeat the same command with: @@ -79,3 +80,33 @@ content hash of every checkpoint shard is required. Keep generated artifacts und This is a TeleFuser regression check, not upstream parity. It detects changes to the current implementation but does not establish equivalence with the official repository. + +## Official Upstream Parity + +The strict upstream baseline pins `Robbyant/lingbot-vla-v2` at commit +`be27333c9b5f2663b0ec33f069dd7dfd67fa32b5`. Keep the checkout, uv environment, cache, and artifacts under +`work_dirs`; Git ignores them. Create the isolated runtime with: + +```bash +mkdir -p work_dirs/.uv-cache-upstream work_dirs/.uv-tmp-upstream +UV_CACHE_DIR="$PWD/work_dirs/.uv-cache-upstream" TMPDIR="$PWD/work_dirs/.uv-tmp-upstream" uv venv work_dirs/.venv-lingbot-upstream --python .venv-vla/bin/python +UV_CACHE_DIR="$PWD/work_dirs/.uv-cache-upstream" TMPDIR="$PWD/work_dirs/.uv-tmp-upstream" uv pip install --python work_dirs/.venv-lingbot-upstream/bin/python -r tools/validation/requirements-lingbot-vla-v2-upstream.txt +UV_CACHE_DIR="$PWD/work_dirs/.uv-cache-upstream" TMPDIR="$PWD/work_dirs/.uv-tmp-upstream" uv pip install --python work_dirs/.venv-lingbot-upstream/bin/python --no-deps "lerobot @ https://github.com/huggingface/lerobot/archive/refs/tags/v0.4.2.tar.gz" +git clone https://github.com/Robbyant/lingbot-vla-v2 work_dirs/lingbot-vla-v2-upstream +git -C work_dirs/lingbot-vla-v2-upstream checkout be27333c9b5f2663b0ec33f069dd7dfd67fa32b5 +``` + +Generate the reference with `capture_lingbot_vla_v2_upstream.py` in the upstream uv environment and the candidate +with `capture_lingbot_vla_v2_telefuser.py` in `.venv-vla`. Pass identical model, processor, camera, task, state, seed, +and device arguments to both commands, add `--deterministic-moe`, and pass `--upstream-root` to the upstream command. +Then compare them with the strict comparator shown above. Generated artifacts belong in `work_dirs/vla_upstream_parity`. + +This is a minimal inference-parity runtime, not a LeRobot training environment. The upstream setup itself combines +LeRobot 0.4.2 metadata constraints with versions outside those constraints, so LeRobot is installed with `--no-deps`; +the capture import and end-to-end run are the runtime checks. + +The official code hard-codes FlashAttention during construction. The upstream capture replaces that selection only +inside its validation process so both sides use eager attention on the Python 3.10.12 / PyTorch 2.11 stack. Production +inference keeps the upstream Triton MoE path through `telefuser.ops`; strict capture uses `--deterministic-moe` because +the upstream kernel uses atomic accumulation and is not bitwise repeatable across separate processes. Artifact metadata +records both `attention_backend` and `moe_backend`, and the comparator rejects mixed-backend artifacts. diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py index b6aad310..5073a1b0 100644 --- a/telefuser/models/lingbot_vla_v2_loader.py +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -1406,7 +1406,7 @@ def map_ckpt_key(self, key, load_vlm_only=False, post_training=False): "post_training": False, "adanorm_time": True, "moe_implementation": "fused", - "use_robby_moe_kernel": False, + "use_robby_moe_kernel": True, "attention_implementation": "eager", "precompute_grid_thw": True, "vlm_causal": True, diff --git a/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json b/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json index 1ab5c1da..71b222fb 100644 --- a/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json +++ b/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json @@ -1,25 +1,229 @@ { "norm_stats": { - "observation.state.arm.position": { - "q01": [-1.3382688760757446, -0.40607330203056335, -1.4083482027053833, -3.058554172515869, -1.423754096031189, -3.192993402481079, -1.591109275817871, -0.7457540035247803, -1.4451789855957031, -3.0523548126220703, -1.4595792293548584, -3.1854426860809326], - "q99": [2.061160087585449, 1.0003128051757812, 1.2696261405944824, 2.941908836364746, 1.4975149631500244, 3.0741331577301025, 1.3934801816940308, 0.3905077278614044, 1.4333486557006836, 3.020704507827759, 1.444725751876831, 3.1354587078094482] - }, - "observation.state.effector.position": { - "q01": [0.3143864572048187, 0.0005160411237739027], - "q99": [1.0, 1.0] - }, "action.arm.position": { - "q01": [-1.3382868766784668, -0.40629321336746216, -1.407132625579834, -3.0591986179351807, -1.4246528148651123, -3.192993402481079, -1.5909051895141602, -0.7457385063171387, -1.44444739818573, -3.0523548126220703, -1.4598064422607422, -3.1850173473358154], - "q99": [2.060295343399048, 1.0005663633346558, 1.2670165300369263, 2.941908836364746, 1.4981306791305542, 3.0859344005584717, 1.3933433294296265, 0.3887772858142853, 1.4337434768676758, 3.019955635070801, 1.4448832273483276, 3.133007049560547] + "mean": [ + -0.2395261526107788, + 1.1349077224731445, + 0.8105292320251465, + -0.31164029240608215, + 0.055165957659482956, + -0.03837483748793602, + 0.21789361536502838, + 1.0841481685638428, + 0.7913455963134766, + -0.32266682386398315, + -0.017583254724740982, + 0.03598335385322571 + ], + "std": [ + 0.41590750217437744, + 1.003921389579773, + 0.7768228650093079, + 0.6747837662696838, + 0.2715570628643036, + 0.6217551827430725, + 0.3499184548854828, + 1.024870753288269, + 0.7947020530700684, + 0.6865031123161316, + 0.24511374533176422, + 0.6167412400245667 + ], + "q01": [ + -1.0185421916961674, + -0.0010411963462829688, + -0.004309860050678266, + -1.5662350454330445, + -0.6512136519670486, + -2.2326875198364258, + -0.1715596118927003, + -0.003369329285621614, + -0.0018556645691394785, + -1.6451744033813476, + -1.0230259281158447, + -1.6478794967651362 + ], + "q99": [ + 0.17221172838211007, + 2.601926616668701, + 2.450952765509486, + 1.3516903750896456, + 1.2373228998184205, + 1.6001025575637815, + 0.9952441711425788, + 2.6186830965638164, + 2.453483357307315, + 1.2904379455566408, + 0.875431350231171, + 2.2640067550659193 + ], + "q02": [ + -0.948497843456269, + -0.0010411963462829688, + -0.0020969149172306023, + -1.474721703195572, + -0.39062818100452423, + -1.6791497331619256, + -0.0868722405433644, + -0.003369329285621614, + -0.0007056228727102265, + -1.5527206142425536, + -0.8707946321487425, + -1.5263084461212157 + ], + "q98": [ + 0.14699576301574702, + 2.5201579942703245, + 2.2860883530676364, + 1.221171345996857, + 1.0305539935112003, + 1.4689324659347545, + 0.9262396463394165, + 2.5411415680885314, + 2.298227728289366, + 1.154620874786377, + 0.577619640159607, + 1.8101414993286138 + ] }, "action.effector.position": { - "q01": [0.3143986165523529, 0.0004720990259665996], - "q99": [1.0, 1.0] + "mean": [ + 0.664304256439209, + 0.6785873174667358 + ], + "std": [ + 0.45511099696159363, + 0.45013949275016785 + ], + "q01": [ + -1e-10, + -1e-10 + ], + "q99": [ + 0.99980000009996, + 0.99980000009996 + ], + "q02": [ + -1e-10, + -1e-10 + ], + "q98": [ + 0.99980000009996, + 0.99980000009996 + ] + }, + "observation.state.arm.position": { + "mean": [ + -0.2384551614522934, + 1.1301639080047607, + 0.8070681095123291, + -0.31032508611679077, + 0.05487748235464096, + -0.037838324904441833, + 0.21649977564811707, + 1.0786004066467285, + 0.78729248046875, + -0.3211615979671478, + -0.017434170469641685, + 0.03533728048205376 + ], + "std": [ + 0.41534423828125, + 1.0045241117477417, + 0.7767844200134277, + 0.6732552647590637, + 0.2708289623260498, + 0.619656503200531, + 0.34926003217697144, + 1.024977207183838, + 0.7943728566169739, + 0.6845870018005371, + 0.24417006969451904, + 0.6140097379684448 + ], + "q01": [ + -1.0185421916961674, + -0.0010411963462829688, + -0.004309860050678266, + -1.56473482670784, + -0.6483812011957168, + -2.224817314338684, + -0.1715596118927003, + -0.003369329285621614, + -0.0018556645691394785, + -1.6435380531311035, + -1.022286941242218, + -1.645177917861938 + ], + "q99": [ + 0.17221172838211007, + 2.601926616668701, + 2.4498462929427625, + 1.350190156364441, + 1.234490449047089, + 1.6001025575637815, + 0.9952441711425788, + 2.616768490922451, + 2.4511832739144563, + 1.2879834201812748, + 0.8739533764839176, + 2.26130517616272 + ], + "q02": [ + -0.948497843456269, + -0.0010411963462829688, + -0.0020969149172306023, + -1.4724713751077652, + -0.38921195561885824, + -1.6712795276641845, + -0.0837356712341304, + -0.003369329285621614, + -0.0007056228727102265, + -1.5510842639923097, + -0.8685776715278626, + -1.520905288314819 + ], + "q98": [ + 0.14419398908615033, + 2.5201579942703245, + 2.2838754079341888, + 1.2204212366342548, + 1.0291377681255343, + 1.4663090641021732, + 0.9231030770301825, + 2.5392269624471666, + 2.2959276448965076, + 1.152166349411011, + 0.575402679538727, + 1.7993351837158205 + ] + }, + "observation.state.effector.position": { + "mean": [ + 0.6655541062355042, + 0.6796996593475342 + ], + "std": [ + 0.45465800166130066, + 0.4496912956237793 + ], + "q01": [ + -1e-10, + -1e-10 + ], + "q99": [ + 0.99980000009996, + 0.99980000009996 + ], + "q02": [ + -1e-10, + -1e-10 + ], + "q98": [ + 0.99980000009996, + 0.99980000009996 + ] } }, - "source": { - "repository": "https://github.com/Robbyant/lingbot-vla-v2", - "path": "assets/norm_stats/robotwin.json", - "commit": "be27333c9b5f2663b0ec33f069dd7dfd67fa32b5" - } -} + "count": 6062592 +} \ No newline at end of file diff --git a/telefuser/pipelines/lingbot_vla_v2/data.py b/telefuser/pipelines/lingbot_vla_v2/data.py index 374e6489..5aedda3f 100644 --- a/telefuser/pipelines/lingbot_vla_v2/data.py +++ b/telefuser/pipelines/lingbot_vla_v2/data.py @@ -94,9 +94,7 @@ def __init__( f"model max_state_dim is {self.max_state_dim}, RobotWin requires {robot_profile.canonical_dim}" ) - def _process_images( - self, images: Mapping[str, ImageInput] - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def _process_images(self, images: Mapping[str, ImageInput]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: missing = [key for key in ROBOTWIN_CAMERA_KEYS if key not in images] if missing: raise ValueError(f"RobotWin observation is missing camera keys: {missing}") @@ -104,7 +102,7 @@ def _process_images( processed_images: list[torch.Tensor] = [] grids: list[torch.Tensor] = [] for key in self.robot_profile.camera_keys: - image = self.image_resize(_image_to_chw_uint8(images[key])) + image = self.image_resize(_image_to_chw_uint8(images[key]).to(dtype=torch.float32)) output = self.processor.image_processor(image) pixels = output["pixel_values"] if isinstance(output, dict) else output.pixel_values grid = output.get("image_grid_thw") if isinstance(output, dict) else getattr(output, "image_grid_thw", None) @@ -113,9 +111,7 @@ def _process_images( if pixels.ndim == 3 and pixels.shape[0] == 1: pixels = pixels.squeeze(0) if pixels.ndim != 2: - raise ValueError( - f"Qwen3-VL image processor must return [patches, features], got {tuple(pixels.shape)}" - ) + raise ValueError(f"Qwen3-VL image processor must return [patches, features], got {tuple(pixels.shape)}") if grid is None or grid.numel() < 3: raise ValueError("Qwen3-VL image processor must return image_grid_thw") processed_images.append(pixels) diff --git a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py index 44f98100..e460b88d 100644 --- a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py +++ b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py @@ -10,7 +10,6 @@ import torch - ROBOTWIN_CAMERA_KEYS = ( "observation.images.cam_high", "observation.images.cam_left_wrist", @@ -53,8 +52,7 @@ class RobotWinProfile: def __init__(self, norm_stats: Mapping[str, Mapping[str, object]]) -> None: self._stats = { key: { - stat_name: torch.as_tensor(stat_value, dtype=torch.float32) - for stat_name, stat_value in values.items() + stat_name: torch.as_tensor(stat_value, dtype=torch.float64) for stat_name, stat_value in values.items() } for key, values in norm_stats.items() } @@ -162,9 +160,11 @@ def _validate_stats(self) -> None: def _normalize(self, key: str, value: torch.Tensor) -> torch.Tensor: low = self._stats[key]["q01"] high = self._stats[key]["q99"] - return (value - low) / (high - low + 1e-6) * 2.0 - 1.0 + normalized = (value.to(dtype=torch.float64) - low) / (high - low + 1e-6) * 2.0 - 1.0 + return normalized.to(dtype=value.dtype) def _unnormalize(self, key: str, value: torch.Tensor) -> torch.Tensor: low = self._stats[key]["q01"] high = self._stats[key]["q99"] - return (value + 1.0) / 2.0 * (high - low + 1e-6) + low + unnormalized = (value.to(dtype=torch.float64) + 1.0) / 2.0 * (high - low + 1e-6) + low + return unnormalized.to(dtype=value.dtype) diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_data.py b/tests/unit/pipelines/lingbot_vla_v2/test_data.py index f73bf740..3459331a 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_data.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_data.py @@ -12,12 +12,12 @@ class _ImageProcessor: def __init__(self) -> None: - self.values: list[int] = [] + self.values: list[float] = [] def __call__(self, image: torch.Tensor) -> dict[str, torch.Tensor]: - assert image.dtype == torch.uint8 + assert image.dtype == torch.float32 assert image.shape == (3, 8, 8) - value = int(image[0, 0, 0]) + value = float(image[0, 0, 0]) self.values.append(value) return { "pixel_values": torch.full((4, 6), float(value)), @@ -73,7 +73,7 @@ def test_prepare_preserves_robotwin_camera_order_and_tensor_contract() -> None: inputs = processor.prepare(observation) - assert image_processor.values == [10, 20, 30] + assert image_processor.values == pytest.approx([10.0, 20.0, 30.0], abs=1e-5) assert tokenizer.rendered_task == observation.task assert tokenizer.padding_side == "right" assert inputs.images.shape == (1, 3, 4, 6) @@ -115,4 +115,4 @@ def test_prepare_resizes_each_camera_before_qwen_processing() -> None: processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) - assert image_processor.values == [10, 20, 30] + assert image_processor.values == pytest.approx([10.0, 20.0, 30.0], abs=1e-5) diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py index 264f7181..f2879e53 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py @@ -24,8 +24,10 @@ def test_normalize_state_uses_robotwin_joint_order() -> None: arm = torch.cat((state[0:6], state[7:13])) effector = state[[6, 13]] assert canonical.shape == (55,) - assert torch.allclose(canonical[0:12], arm / (2.0 + 1e-6) * 2.0 - 1.0) - assert torch.allclose(canonical[28:30], (effector + 1.0) / (2.0 + 1e-6) * 2.0 - 1.0) + expected_arm = (arm.to(torch.float64) / (2.0 + 1e-6) * 2.0 - 1.0).to(torch.float32) + expected_effector = ((effector.to(torch.float64) + 1.0) / (2.0 + 1e-6) * 2.0 - 1.0).to(torch.float32) + assert torch.equal(canonical[0:12], expected_arm) + assert torch.equal(canonical[28:30], expected_effector) assert torch.count_nonzero(canonical[12:28]) == 0 assert torch.count_nonzero(canonical[30:]) == 0 diff --git a/tests/unit/validation/test_lingbot_vla_v2_artifacts.py b/tests/unit/validation/test_lingbot_vla_v2_artifacts.py index fc735301..09865983 100644 --- a/tests/unit/validation/test_lingbot_vla_v2_artifacts.py +++ b/tests/unit/validation/test_lingbot_vla_v2_artifacts.py @@ -36,11 +36,13 @@ def _metadata() -> dict[str, object]: "artifact_kind": "telefuser_regression", "checkpoint_manifest_sha256": "checkpoint", "processor_manifest_sha256": "processor", + "norm_stats_sha256": "norm-stats", "input_sha256": "input", "seed": 7, "num_steps": 2, "torch_dtype": "bfloat16", "attention_backend": "eager", + "moe_backend": "deterministic_torch_reference", } @@ -136,6 +138,16 @@ def test_compare_artifacts_rejects_different_artifact_identity(tmp_path: Path) - compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) +def test_compare_artifacts_rejects_different_moe_backends(tmp_path: Path) -> None: + candidate_metadata = _metadata() + candidate_metadata["moe_backend"] = "upstream_triton" + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", _arrays(), candidate_metadata) + + with pytest.raises(ValueError, match="moe_backend"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + def test_velocity_trace_snapshots_inputs_and_restores_the_model_instance() -> None: class _FlowModel: _use_compile_predict_velocity = True diff --git a/tools/validation/capture_lingbot_vla_v2_telefuser.py b/tools/validation/capture_lingbot_vla_v2_telefuser.py index 2ff3b1f6..f440510c 100644 --- a/tools/validation/capture_lingbot_vla_v2_telefuser.py +++ b/tools/validation/capture_lingbot_vla_v2_telefuser.py @@ -170,6 +170,7 @@ def capture_artifact( output: Path, device: str, full_checkpoint_hash: bool, + deterministic_moe: bool, ) -> tuple[Path, Path]: if len(image_paths) != len(ROBOTWIN_CAMERA_KEYS): raise ValueError(f"expected {len(ROBOTWIN_CAMERA_KEYS)} camera paths, got {len(image_paths)}") @@ -178,6 +179,10 @@ def capture_artifact( output.parent.mkdir(parents=True, exist_ok=True) pipeline = _build_pipeline(model_root, qwen3vl_root, device) + if deterministic_moe: + for module in pipeline.policy_stage.policy.modules(): + if hasattr(module, "_use_robby_moe_kernel"): + module._use_robby_moe_kernel = False capture = TensorCapture() try: observation = LingBotVlaV2Observation( @@ -209,6 +214,10 @@ def capture_artifact( include_contents=full_checkpoint_hash, ), "checkpoint_hash_mode": "full_sha256" if full_checkpoint_hash else "filename_and_size", + "norm_stats_sha256": _sha256_file( + Path(__file__).resolve().parents[2] + / "telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json" + ), "processor_manifest_sha256": _manifest_sha256( _processor_files(qwen3vl_root), include_contents=True, @@ -218,6 +227,7 @@ def capture_artifact( "num_steps": trace.step, "torch_dtype": str(pipeline.torch_dtype).removeprefix("torch."), "attention_backend": str(flow_model.config.attention_implementation), + "moe_backend": "deterministic_torch_reference" if deterministic_moe else "upstream_triton", "device": str(target_device), "device_name": torch.cuda.get_device_name(target_device) if target_device.type == "cuda" else "cpu", "torch_version": torch.__version__, @@ -262,6 +272,11 @@ def main() -> None: action="store_true", help="Hash all checkpoint bytes instead of the faster filename-and-size manifest", ) + parser.add_argument( + "--deterministic-moe", + action="store_true", + help="Disable the upstream atomic Triton MoE kernel for bitwise cross-process parity", + ) args = parser.parse_args() paths = ( @@ -285,6 +300,7 @@ def main() -> None: output=args.output, device=args.device, full_checkpoint_hash=args.full_checkpoint_hash, + deterministic_moe=args.deterministic_moe, ) print(f"Saved LingBot-VLA v2 capture: {artifact}") print(f"Saved LingBot-VLA v2 metadata: {metadata}") diff --git a/tools/validation/capture_lingbot_vla_v2_upstream.py b/tools/validation/capture_lingbot_vla_v2_upstream.py new file mode 100644 index 00000000..cc724e9e --- /dev/null +++ b/tools/validation/capture_lingbot_vla_v2_upstream.py @@ -0,0 +1,541 @@ +"""Capture a layered artifact from the fixed official LingBot-VLA v2 checkout. + +Run this script with the dedicated upstream uv environment and with the fixed +upstream checkout as ``--upstream-root``. The official code forces +FlashAttention during construction. For reproducible comparison on the local +PyTorch 2.11 stack, this runner intercepts model construction inside this +process only and selects the eager attention implementation used by TeleFuser. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import sys +from contextlib import contextmanager +from pathlib import Path +from types import MethodType, SimpleNamespace +from typing import Any, Iterator, Sequence + +import numpy as np +import torch +import transformers +from PIL import Image +from accelerate import init_empty_weights +from safetensors.torch import load_file +from torchvision.transforms.v2 import Resize +from transformers import AutoConfig, AutoProcessor, PreTrainedModel + +ARTIFACT_SCHEMA_VERSION = 1 +UPSTREAM_COMMIT = "be27333c9b5f2663b0ec33f069dd7dfd67fa32b5" +CAMERA_KEYS = ( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", +) + + +def _sha256_file(path: Path, digest: Any | None = None) -> str: + result = hashlib.sha256() if digest is None else digest + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + result.update(block) + return result.hexdigest() + + +def _manifest_sha256(paths: Sequence[Path], *, include_contents: bool) -> str: + digest = hashlib.sha256() + for path in sorted((item.resolve() for item in paths), key=lambda item: item.name): + stat = path.stat() + digest.update(path.name.encode("utf-8")) + digest.update(str(stat.st_size).encode("ascii")) + if include_contents: + _sha256_file(path, digest) + return digest.hexdigest() + + +def _processor_files(root: Path) -> list[Path]: + return sorted(path for path in root.iterdir() if path.is_file() and path.suffix != ".safetensors") + + +def _input_sha256(task: str, state: Sequence[float], image_paths: Sequence[Path]) -> str: + digest = hashlib.sha256() + canonical = json.dumps({"task": task, "state": list(state)}, sort_keys=True, separators=(",", ":")) + digest.update(canonical.encode("utf-8")) + for path in image_paths: + digest.update(path.name.encode("utf-8")) + _sha256_file(path, digest) + return digest.hexdigest() + + +def _git_commit(repository: Path) -> str: + completed = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + status = subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + if status.stdout: + raise RuntimeError(f"Upstream checkout must be clean, got:\n{status.stdout}") + return completed.stdout.strip() + + +def _checkpoint_shards(model_root: Path) -> list[Path]: + index_path = model_root / "model.safetensors.index.json" + payload = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = payload.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Invalid checkpoint index: {index_path}") + shards = [model_root / name for name in sorted(set(weight_map.values()))] + missing = [str(path) for path in shards if not path.is_file()] + if missing: + raise FileNotFoundError(f"Missing checkpoint shards: {missing}") + return shards + + +class TensorCapture: + def __init__(self) -> None: + self.arrays: dict[str, np.ndarray] = {} + self.array_metadata: dict[str, dict[str, object]] = {} + + def add(self, key: str, tensor: torch.Tensor) -> None: + if key in self.arrays: + raise ValueError(f"Duplicate capture key: {key}") + snapshot = tensor.detach().cpu().clone() + stored = snapshot.float() if snapshot.is_floating_point() else snapshot + array = stored.numpy() + self.arrays[key] = array + self.array_metadata[key] = { + "shape": list(snapshot.shape), + "original_dtype": str(snapshot.dtype).removeprefix("torch."), + "stored_dtype": str(array.dtype), + } + + +class VelocityTrace: + def __init__(self, capture: TensorCapture) -> None: + self.capture = capture + self.step = 0 + + def record(self, original: Any, *args: Any, **kwargs: Any) -> torch.Tensor: + if len(args) < 5: + raise RuntimeError("Official predict_velocity trace received an unexpected call signature") + suffix = f"{self.step:02d}" + x_t, timestep = args[3], args[4] + if self.step == 0: + self.capture.add("initial_noise", x_t) + self.capture.add(f"timestep_step_{suffix}", timestep) + self.capture.add(f"x_t_step_{suffix}", x_t) + velocity = original(*args, **kwargs) + self.capture.add(f"velocity_step_{suffix}", velocity) + self.step += 1 + return velocity + + +@contextmanager +def _trace_predict_velocity(flow_model: Any, capture: TensorCapture) -> Iterator[VelocityTrace]: + original = flow_model.predict_velocity + trace = VelocityTrace(capture) + flow_model.predict_velocity = lambda *args, **kwargs: trace.record(original, *args, **kwargs) + try: + yield trace + finally: + del flow_model.predict_velocity + + +def _force_eager(config: Any) -> None: + for current in (config, getattr(config, "text_config", None), getattr(config, "vision_config", None)): + if current is not None: + current._attn_implementation = "eager" + + +@contextmanager +def _eager_construction() -> Iterator[None]: + """Override the upstream hard-coded FA2 selection in this process only.""" + original = PreTrainedModel._from_config.__func__ + + def from_config(cls: type[PreTrainedModel], config: Any, **kwargs: Any) -> PreTrainedModel: + _force_eager(config) + return original(cls, config, **kwargs) + + PreTrainedModel._from_config = classmethod(from_config) + try: + yield + finally: + PreTrainedModel._from_config = classmethod(original) + + +def _official_model_values(qwen3vl_root: Path) -> dict[str, Any]: + # Base-6B values are documented in the fixed upstream Training_Config.md. + return { + "post_training": False, + "adanorm_time": True, + "moe_implementation": "fused", + "use_robby_moe_kernel": False, + "attention_implementation": "eager", + "vit_attn_implementation": "eager", + "precompute_grid_thw": True, + "vlm_causal": True, + "use_moe": True, + "token_moe_layers": list(range(36)), + "token_num_experts": 32, + "token_top_k": 4, + "token_moe_intermediate_size": 512, + "token_shared_intermediate_size": 704, + "bias_update_speed": 0.0, + "sequence_wise_mode": "per_sequence", + "sequence_wise_loss_coeff": 1e-3, + "router_z_loss_coeff": 1e-4, + "router_activation": "sigmoid", + "routed_scaling_factor": 4.0, + "use_shared_expert_gate": False, + "freeze_vision_encoder": False, + "tokenizer_max_length": 72, + "loss_type": "L1_fm", + "action_dim": 55, + "max_action_dim": 55, + "max_state_dim": 55, + "tokenizer_path": str(qwen3vl_root), + "align_params": { + "mode": "query", + "num_task_tokens": 8, + "depth_loss_weight": 0.004, + "future_depth_loss_weight": 0.004, + "use_future_video": True, + "llm": {"dim_out": 2560, "image_token_size": 8, "image_input_size": 224}, + "depth": { + "model_type": "MoRGBD", + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "token_size": 16, + "dim_out": 1024, + "input_size": 224, + "use_future_depth": True, + "block_future_depth_to_action": True, + "future_depth_head_type": "resampler", + "detach_future_image_feats": True, + }, + "video": { + "attention_mode": "flex_block_causal", + "input_size": 256, + "block_suffix_to_future_video": True, + "share_future_depth_query": True, + "use_shared_future_task_proj": True, + "use_current_shared_task_proj": True, + "num_future_frames": 1, + "use_warmup_frame": True, + "effective_fps": 1.0, + "n_blocks": 1, + "cls_pool": "last", + "detach_image_feats": True, + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "dim_out": 1024, + "future_video_loss_weight": 0.004, + "use_smooth_l1_loss": False, + "use_mse_loss": True, + "mse_loss_weight": 1.0, + "use_patch_loss": True, + "use_current_patch_loss": True, + "use_cosine_loss": False, + "cosine_loss_weight": 0.2, + "use_cls_loss": False, + "cls_loss_type": "mse", + "cls_loss_weight": 0.2, + }, + }, + } + + +def _build_config(qwen3vl_root: Path) -> Any: + from lingbotvla.models.vla.lingbot_vla.configuration_lingbot_vla import LingbotVLAV2Config + + config = LingbotVLAV2Config(**_official_model_values(qwen3vl_root)) + qwen_config = AutoConfig.from_pretrained(str(qwen3vl_root), local_files_only=True) + for key in ( + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "rms_norm_eps", + "rope_theta", + "vocab_size", + "max_position_embeddings", + "hidden_act", + "tie_word_embeddings", + ): + if hasattr(qwen_config.text_config, key): + setattr(config, key, getattr(qwen_config.text_config, key)) + config.vision_config = qwen_config.vision_config + config.use_cache = True + return config + + +def _load_official_model(model_root: Path, config: Any, device: torch.device) -> Any: + from lingbotvla.models.vla.lingbot_vla.modeling_lingbot_vla_v2 import LingbotVlaV2Policy + from lingbotvla.models.vla.lingbot_vla.qwen3vl_in_vla import apply_lingbot_qwen3_vl_patch + + apply_lingbot_qwen3_vl_patch() + with _eager_construction(), init_empty_weights(): + model = LingbotVlaV2Policy(config, eval=True) + + index = json.loads((model_root / "model.safetensors.index.json").read_text(encoding="utf-8")) + checkpoint_keys = set(index["weight_map"]) + model_keys = set(model.state_dict()) + if checkpoint_keys != model_keys: + raise RuntimeError( + "Official model/checkpoint key mismatch: " + f"missing={sorted(model_keys - checkpoint_keys)[:10]}, " + f"unexpected={sorted(checkpoint_keys - model_keys)[:10]}" + ) + + for shard in _checkpoint_shards(model_root): + model.load_state_dict(load_file(shard, device="cpu"), strict=False, assign=True) + unmaterialized = [name for name, tensor in model.state_dict().items() if tensor.is_meta] + if unmaterialized: + raise RuntimeError(f"Official checkpoint left meta tensors: {unmaterialized[:10]}") + return model.to(device=device, dtype=torch.bfloat16).eval() + + +def _load_rgb(path: Path) -> torch.Tensor: + with Image.open(path) as image: + array = np.asarray(image.convert("RGB")).copy() + return torch.from_numpy(array).permute(2, 0, 1).contiguous() + + +def _prepare_inputs( + upstream_root: Path, + qwen3vl_root: Path, + config: Any, + image_paths: Sequence[Path], + task: str, + state: Sequence[float], +) -> dict[str, torch.Tensor]: + from lingbotvla.data.vla_data.utils import FeatureTransform + + processor = AutoProcessor.from_pretrained(str(qwen3vl_root), local_files_only=True, padding_side="right") + data_config = SimpleNamespace( + joints=["{'arm.position': 14}", "{'end.position': 14}", "{'effector.position': 2}"], + cameras=["camera_top", "camera_wrist_left", "camera_wrist_right"], + norm_type=[ + "{'arm.position': 'bounds_99_woclip'}", + "{'end.position': 'bounds_99_woclip'}", + "{'effector.position': 'bounds_99_woclip'}", + ], + ) + transform = FeatureTransform( + upstream_root / "configs/robot_configs/robotwin.yaml", + data_config, + config, + processor, + chunk_size=config.chunk_size, + norm_stats_path=upstream_root / "assets/norm_stats/robotwin.json", + ) + resize = Resize((256, 256), antialias=True) + item: dict[str, Any] = {"observation.state": torch.tensor(state, dtype=torch.float32), "task": task} + for key, path in zip(CAMERA_KEYS, image_paths, strict=True): + item[key] = resize(_load_rgb(path).to(dtype=torch.float32)) + prepared = transform.apply(item, policy_eval=True) + return { + "images": prepared["images"].unsqueeze(0), + "img_masks": prepared["img_masks"].unsqueeze(0), + "image_grid_thw": prepared["image_grid_thw"].unsqueeze(0), + "lang_tokens": prepared["lang_tokens"].unsqueeze(0), + "lang_masks": prepared["lang_masks"].unsqueeze(0), + "state": prepared["state"].unsqueeze(0), + } + + +def capture_artifact( + *, + upstream_root: Path, + model_root: Path, + qwen3vl_root: Path, + image_paths: Sequence[Path], + task: str, + state: Sequence[float], + seed: int, + output: Path, + device: str, + full_checkpoint_hash: bool, + deterministic_moe: bool, +) -> tuple[Path, Path]: + commit = _git_commit(upstream_root) + if commit != UPSTREAM_COMMIT: + raise RuntimeError(f"Expected upstream commit {UPSTREAM_COMMIT}, got {commit}") + if len(image_paths) != len(CAMERA_KEYS): + raise ValueError(f"expected {len(CAMERA_KEYS)} camera paths, got {len(image_paths)}") + + sys.path.insert(0, str(upstream_root)) + output = output.with_suffix(".npz") + metadata_path = output.with_suffix(".json") + output.parent.mkdir(parents=True, exist_ok=True) + target_device = torch.device(device) + if target_device.type != "cuda": + raise ValueError("Official 6B parity capture currently requires CUDA") + + config = _build_config(qwen3vl_root) + inputs = _prepare_inputs(upstream_root, qwen3vl_root, config, image_paths, task, state) + capture = TensorCapture() + for key, tensor in inputs.items(): + capture.add(key, tensor) + + model = _load_official_model(model_root, config, target_device) + if deterministic_moe: + import lingbotvla.models.vla.lingbot_vla.qwen2_action_expert as qwen2_action_expert + + qwen2_action_expert.robby_moe_forward = None + + def deterministic_forward(experts, module, num_experts, routing_weights, selected_experts, hidden_states): + del module + output = torch.zeros_like(hidden_states) + for expert_id in range(num_experts): + routes = (selected_experts == expert_id).nonzero(as_tuple=False) + if routes.numel() == 0: + continue + token_ids, route_ids = routes[:, 0], routes[:, 1] + expert_input = hidden_states.index_select(0, token_ids) + gate = torch.nn.functional.linear(expert_input, experts.gate_proj[expert_id]) + up = torch.nn.functional.linear(expert_input, experts.up_proj[expert_id]) + intermediate = torch.nn.functional.silu(gate) * up + expert_output = torch.nn.functional.linear(intermediate, experts.down_proj[expert_id]) + weights = routing_weights[token_ids, route_ids].unsqueeze(-1) + output.index_add_(0, token_ids, expert_output * weights) + return output + + for module in model.modules(): + if module.__class__.__name__ == "Qwen2FusedExperts": + module.forward = MethodType(deterministic_forward, module) + tensors = { + "images": inputs["images"].to(device=target_device, dtype=torch.bfloat16), + "img_masks": inputs["img_masks"].to(device=target_device), + "lang_tokens": inputs["lang_tokens"].to(device=target_device), + "lang_masks": inputs["lang_masks"].to(device=target_device), + "state": inputs["state"].to(device=target_device, dtype=torch.bfloat16), + "image_grid_thw": inputs["image_grid_thw"].to(device=target_device, dtype=torch.long), + } + generator = torch.Generator(device=target_device).manual_seed(seed) + noise = torch.randn( + 1, + int(config.n_action_steps), + int(config.max_action_dim), + device=target_device, + dtype=torch.bfloat16, + generator=generator, + ) + with torch.inference_mode(), _trace_predict_velocity(model.model, capture) as trace: + actions = model.sample_actions(**tensors, noise=noise) + capture.add("canonical_normalized_actions", actions.squeeze(0).to(device="cpu", dtype=torch.float32)) + if trace.step != int(config.num_steps): + raise RuntimeError(f"captured {trace.step} denoising steps, expected {config.num_steps}") + + checkpoint_paths = _checkpoint_shards(model_root) + metadata = { + "schema_version": ARTIFACT_SCHEMA_VERSION, + "artifact_kind": "official_upstream_common_eager", + "upstream_commit": commit, + "upstream_attention_override": "process_local_pretrained_model_from_config_intercept", + "checkpoint_manifest_sha256": _manifest_sha256(checkpoint_paths, include_contents=full_checkpoint_hash), + "checkpoint_hash_mode": "full_sha256" if full_checkpoint_hash else "filename_and_size", + "norm_stats_sha256": _sha256_file(upstream_root / "assets/norm_stats/robotwin.json"), + "processor_manifest_sha256": _manifest_sha256(_processor_files(qwen3vl_root), include_contents=True), + "input_sha256": _input_sha256(task, state, image_paths), + "seed": seed, + "num_steps": trace.step, + "torch_dtype": "bfloat16", + "attention_backend": "eager", + "moe_backend": "deterministic_torch_reference" if deterministic_moe else "upstream_triton", + "device": str(target_device), + "device_name": torch.cuda.get_device_name(target_device), + "python_version": sys.version.split()[0], + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + "arrays": capture.array_metadata, + } + np.savez(output, **capture.arrays) + metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return output, metadata_path + + +def _parse_state(value: str) -> list[float]: + try: + state = json.loads(value) + except json.JSONDecodeError as error: + raise argparse.ArgumentTypeError("state-json must be valid JSON") from error + if not isinstance(state, list) or len(state) != 14 or any(isinstance(item, bool) for item in state): + raise argparse.ArgumentTypeError("state-json must be a 14-element numeric JSON list") + try: + return [float(item) for item in state] + except (TypeError, ValueError) as error: + raise argparse.ArgumentTypeError("state-json must contain only numeric values") from error + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream-root", required=True, type=Path) + parser.add_argument("--model-root", required=True, type=Path) + parser.add_argument("--qwen3vl-root", required=True, type=Path) + parser.add_argument("--camera-high", required=True, type=Path) + parser.add_argument("--camera-left-wrist", required=True, type=Path) + parser.add_argument("--camera-right-wrist", required=True, type=Path) + parser.add_argument("--task", required=True) + parser.add_argument("--state-json", required=True, type=_parse_state) + parser.add_argument("--seed", required=True, type=int) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--full-checkpoint-hash", action="store_true") + parser.add_argument( + "--deterministic-moe", + action="store_true", + help="Disable the upstream atomic Triton MoE kernel for bitwise cross-process parity", + ) + args = parser.parse_args() + + paths = ( + args.upstream_root, + args.model_root, + args.qwen3vl_root, + args.camera_high, + args.camera_left_wrist, + args.camera_right_wrist, + ) + missing = [str(path) for path in paths if not path.exists()] + if missing: + parser.error(f"input paths do not exist: {missing}") + + artifact, metadata = capture_artifact( + upstream_root=args.upstream_root.resolve(), + model_root=args.model_root.resolve(), + qwen3vl_root=args.qwen3vl_root.resolve(), + image_paths=(args.camera_high, args.camera_left_wrist, args.camera_right_wrist), + task=args.task, + state=args.state_json, + seed=args.seed, + output=args.output, + device=args.device, + full_checkpoint_hash=args.full_checkpoint_hash, + deterministic_moe=args.deterministic_moe, + ) + print(f"Saved official LingBot-VLA v2 capture: {artifact}") + print(f"Saved official LingBot-VLA v2 metadata: {metadata}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/requirements-lingbot-vla-v2-upstream.txt b/tools/validation/requirements-lingbot-vla-v2-upstream.txt new file mode 100644 index 00000000..fa927ca4 --- /dev/null +++ b/tools/validation/requirements-lingbot-vla-v2-upstream.txt @@ -0,0 +1,17 @@ +# Isolated runtime for the fixed LingBot-VLA v2 upstream parity capture. +# Install this file into a dedicated uv venv; do not install it into TeleFuser's venv. +accelerate==1.7.0 +av==15.0.0 +datasets==3.6.0 +einops==0.8.1 +huggingface-hub==0.34.3 +numpy==2.2.6 +pillow==12.3.0 +psutil==7.0.0 +pydantic==2.13.4 +pyyaml==6.0.3 +safetensors==0.6.2 +torch==2.11.0 +torchdata==0.11.0 +torchvision==0.26.0 +transformers==4.57.3 diff --git a/tools/validation/run_lingbot_vla_v2_parity.py b/tools/validation/run_lingbot_vla_v2_parity.py index 97c1c514..901378d9 100644 --- a/tools/validation/run_lingbot_vla_v2_parity.py +++ b/tools/validation/run_lingbot_vla_v2_parity.py @@ -32,11 +32,13 @@ IDENTITY_METADATA_KEYS = ( "checkpoint_manifest_sha256", "processor_manifest_sha256", + "norm_stats_sha256", "input_sha256", "seed", "num_steps", "torch_dtype", "attention_backend", + "moe_backend", ) _STEP_KEY = re.compile(r"^(timestep|x_t|velocity)_step_([0-9]+)$") From 3c996d6b7c1b38587deb40188e1630f76e9f54f7 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Tue, 4 Aug 2026 07:37:59 +0000 Subject: [PATCH 06/15] refactor(vla): trim training-only runtime code Make the LingBot VLA v2 integration inference-only by rejecting training forwards, freezing policy parameters, and removing optimizer and training-parallel APIs. Remove action, depth, video, MoE balance, and router loss helpers while preserving checkpoint topology and inference prefix computation. Add regression coverage for the inference-only boundary and removed loader loss helpers. Verification:\n- 28 focused VLA tests passed\n- real LingBot VLA v2 6B checkpoint loaded and inferred on H100\n- strict upstream parity passed 38/38 comparisons with max_abs=0.0\n- targeted Ruff checks and git diff --check passed --- telefuser/models/lingbot_vla_v2.py | 909 +--------------------- telefuser/models/lingbot_vla_v2_loader.py | 420 ---------- tests/unit/models/test_lingbot_vla_v2.py | 19 +- 3 files changed, 36 insertions(+), 1312 deletions(-) diff --git a/telefuser/models/lingbot_vla_v2.py b/telefuser/models/lingbot_vla_v2.py index 2eec3a13..cf78cde1 100644 --- a/telefuser/models/lingbot_vla_v2.py +++ b/telefuser/models/lingbot_vla_v2.py @@ -215,7 +215,7 @@ def __init__(self, **kwargs): from torch import nn import torch.nn.functional as F from torch import Tensor, nn -from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union from functools import partial import math from transformers import ( @@ -234,8 +234,6 @@ def __init__(self, **kwargs): ) -class LossKwargs(TypedDict, total=False): - labels: Optional[torch.LongTensor] from transformers.utils.deprecation import deprecate_kwarg from transformers.activations import ACT2FN from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, is_flash_attn_available @@ -258,12 +256,10 @@ class LossKwargs(TypedDict, total=False): create_sinusoidal_pos_embedding, make_att_2d_masks, resize_with_pad, - sample_beta, ) from telefuser.models.lingbot_vla_v2_loader import apply_rope, our_eager_attention_forward from telefuser.models.lingbot_vla_v2_loader import flex_attention_forward from telefuser.models.lingbot_vla_v2_loader import build_block_mask, flex_attention_with_block_mask -import time from telefuser.models.lingbot_vla_v2_loader import LingBotVLAWeightLoader, TaskTokenDepthHead from telefuser.models.lingbot_vla_v2_moe import ( @@ -561,29 +557,7 @@ def __init__(self, config: QwenvlWithExpertConfig, eval=False): self.attention_interface = self.get_attention_interface() # self.to_bfloat16_like_physical_intelligence() - self.set_requires_grad() - def set_requires_grad(self): - """sets the requires_grad attribute of the model parameters based on the configuration. - If `freeze_vision_encoder` is True, the vision tower parameters are frozen. - If `train_expert_only` is True, the entire Qwenvl model is frozen. - """ - if self.config.freeze_vision_encoder: - self.qwenvl.visual.eval() - for params in self.qwenvl.visual.parameters(): - params.requires_grad = False - - if self.config.train_expert_only: - self.qwenvl.eval() - for params in self.qwenvl.parameters(): - params.requires_grad = False - - def train(self, mode: bool = True): - super().train(mode) - if self.config.freeze_vision_encoder: - self.qwenvl.visual.eval() - if self.config.train_expert_only: - self.qwenvl.eval() def to_bfloat16_like_physical_intelligence(self): """casts the model to bfloat16. @@ -818,96 +792,6 @@ def get_attention_interface(self): ) return attention_interface -class LingbotVlaPolicy(PreTrainedModel): - config_class = LingbotVLAConfig - name = "torch_lingbot_vla" - supports_gradient_checkpointing = True - - _no_split_modules = ["Qwen2DecoderLayer", "FixQwen2RMSNorm", "FixAdaRMSNorm"] # NOTE: if moudule in Qwen2DecoderLayer, it doesn't need to specify in _no_split_modules - - def get_parallel_plan(self): - from telefuser.models.lingbot_vla_v2_loader import NativeParallelPlan as ParallelPlan - from torch.distributed._tensor import Shard - ep_plan = { - "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.gate_proj": Shard(0), - "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.up_proj": Shard(0), - "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.down_proj": Shard(0), - } - return ParallelPlan(ep_plan=ep_plan) - - @classmethod - def get_weight_loader(cls): - return LingBotVLAWeightLoader() - - def __init__( - self, - config: LingbotVLAConfig, - eval: bool=False, - ): - """ - Args: - config: Policy configuration class instance or None, in which case the default instantiation of - the configuration class is used. - """ - - super().__init__(config) - self.config = config - self.language_tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, local_files_only=True) - self.model = FlowMatching(config, eval) - - if not getattr(self.config,"use_lm_head", False): - del self.model.qwenvl_with_expert.qwenvl.lm_head - del self.model.qwenvl_with_expert.qwen_expert.lm_head - - self.reset() - torch.set_float32_matmul_precision("high") - - def reset(self): - return None - - def get_optim_params(self) -> dict: - return self.parameters() - - def forward( - self, images, img_masks, state, lang_tokens, lang_masks, actions, joint_mask=None, action_is_pad=None, noise=None, time=None, vlm_causal=False, depth_targets=None, precompute_grid_thw=False, future_depth_targets=None, **kwargs - ) -> tuple[Tensor, dict[str, Tensor]]: - loss_dict = {} - # Keep state and actions in fp32 for action learning precision - if getattr(self.config, 'action_fp32', False): - state = state.float() - actions = actions.float() - losses, loss_depth, depth_preds, seq_wise_loss, moe_metrics = self.model.forward( - images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, vlm_causal, self.config.loss_type, depth_targets, precompute_grid_thw=precompute_grid_thw, future_depth_targets=future_depth_targets, - ) - - if joint_mask is not None: - if 'repeat' in self.config.loss_type: - joint_mask = joint_mask.repeat(2,1) - mask_expanded = joint_mask.unsqueeze(1).expand(-1, losses.size(1), -1) # (B, T, D) - masked_losses = losses * mask_expanded - - valid_counts = mask_expanded.sum(dim=(1, 2)).clamp(min=1) - batch_mean_losses = masked_losses.sum(dim=(1, 2)) / valid_counts - loss_vla = masked_losses.sum() / mask_expanded.sum().clamp(min=1) - else: - losses = losses[:, :, :self.config.action_dim] - batch_mean_losses = losses.mean(dim=(1, 2)) - loss_vla = losses.mean() - - loss_dict["batch_mean_losses"] = batch_mean_losses.detach() - - total_loss = ( - loss_vla - + loss_depth - + seq_wise_loss - ) - - # Attach MoE monitoring metrics to loss_dict - if moe_metrics: - loss_dict.update(moe_metrics) - - return total_loss, loss_vla, loss_depth, seq_wise_loss, loss_dict, depth_preds - class FlowMatching(nn.Module): def __init__(self, config, eval): super().__init__() @@ -977,7 +861,6 @@ def __init__(self, config, eval): self.use_shared_future_task_proj = False self.future_video_share_future_depth_query = False - self.set_requires_grad() def init_depth_heads(self, config): self.llm_image_token_size = config['llm']['image_token_size'] @@ -1009,12 +892,9 @@ def init_depth_heads(self, config): config['depth']['num_backbone_tokens'], config['llm']['dim_out'] ) ) - self.depth_align_embs.requires_grad = True self.depth_align_head = TaskTokenDepthHead(config['depth'], llm_hidden_size=config['llm']['dim_out']).to(dtype=torch.bfloat16) - for p in self.depth_align_head.parameters(): - p.requires_grad = True if self.use_future_depth: self.future_depth_align_embs = nn.Parameter( @@ -1022,14 +902,11 @@ def init_depth_heads(self, config): config['depth']['num_backbone_tokens'], config['llm']['dim_out'] ) ) - self.future_depth_align_embs.requires_grad = True self.future_depth_align_head = TaskTokenDepthHead( config['depth'], llm_hidden_size=config['llm']['dim_out'] ).to(dtype=torch.bfloat16) - for p in self.future_depth_align_head.parameters(): - p.requires_grad = True def init_video_heads(self, config): if self.align_type != "query": @@ -1099,19 +976,14 @@ def init_video_heads(self, config): video_config['num_backbone_tokens'], config['llm']['dim_out'] ) ) - self.current_video_align_embs.requires_grad = True if self.use_current_shared_task_proj: self.current_shared_task_proj = nn.Linear( config['llm']['dim_out'] * 2, config['llm']['dim_out'], ) - for p in self.current_shared_task_proj.parameters(): - p.requires_grad = True self.current_video_align_head = TaskTokenDepthHead( video_config, llm_hidden_size=config['llm']['dim_out'] ).to(dtype=torch.bfloat16) - for p in self.current_video_align_head.parameters(): - p.requires_grad = True if ( not self.future_video_share_future_depth_query @@ -1122,19 +994,14 @@ def init_video_heads(self, config): video_config['num_backbone_tokens'], config['llm']['dim_out'] ) ) - self.future_video_align_embs.requires_grad = True if self.use_shared_future_task_proj: self.future_shared_task_proj = nn.Linear( config['llm']['dim_out'] * 2, config['llm']['dim_out'], ) - for p in self.future_shared_task_proj.parameters(): - p.requires_grad = True self.future_video_align_head = TaskTokenDepthHead( video_config, llm_hidden_size=config['llm']['dim_out'] ).to(dtype=torch.bfloat16) - for p in self.future_video_align_head.parameters(): - p.requires_grad = True if self.use_future_video_cls: self.future_video_cls_align_emb = nn.Embedding(1, config['llm']['dim_out']) @@ -1142,8 +1009,6 @@ def init_video_heads(self, config): nn.LayerNorm(config['llm']['dim_out']), nn.Linear(config['llm']['dim_out'], video_config['dim_out']), ).to(dtype=torch.bfloat16) - for p in self.future_video_cls_head.parameters(): - p.requires_grad = True def _future_depth_token_count(self): return self.num_task_tokens if getattr(self, "use_future_depth", False) else 0 @@ -1247,10 +1112,6 @@ def _init_weights(self, module): if reset_post_init is not None: reset_post_init() - def set_requires_grad(self): - for params in self.state_proj.parameters(): - params.requires_grad = self.config.train_state_proj - @staticmethod def _fp32_linear(module, x): """Compute linear layer in fp32 regardless of module's current parameter dtype.""" @@ -1260,10 +1121,6 @@ def _fp32_linear(module, x): module.bias.float() if module.bias is not None else None ) - def sample_time(self, bsize, device): - time_beta = sample_beta(1.5, 1.0, bsize, device) - time = time_beta * 0.999 + 0.001 - return time.to(dtype=torch.float32, device=device) def embed_prefix( self, images, img_masks, lang_tokens, lang_masks, vlm_causal, precompute_grid_thw=False @@ -1376,165 +1233,10 @@ def embed_suffix(self, state, noisy_actions, timestep): # (torch.Size([state_bs, return time_emb_ori, embs, pad_masks, att_masks - def forward( - self, - images, - img_masks, - lang_tokens, - lang_masks, - state, - actions, - noise=None, - time=None, - vlm_causal=False, - loss_type='fm', - depth_targets=None, - precompute_grid_thw=False, - future_depth_targets=None, - ) -> Tensor: - dtype = state.dtype - device = state.device - if noise is None: - noise = torch.randn(actions.shape, device=device, dtype=dtype) - - if time is None: - time = self.sample_time(actions.size(0), device).to(dtype) - - time_expanded = time[:, None, None] - x_t = time_expanded * noise + (1 - time_expanded) * actions - u_t = noise - actions - - prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( - images, img_masks, lang_tokens, lang_masks, vlm_causal, precompute_grid_thw=precompute_grid_thw - ) # 1,bs_img*(768+48),2048 1,bs_img*(768+48) 1,bs_img*(768+48) - time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( - state, x_t, time - ) # [1, state_bs*(50+1), 1024], [1, state_bs*(50+1)], [1, state_bs*(50+1)] state_bs=bs_img - - pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1) # 1,state_bs*(768+48+50+1) - att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)# 1,state_bs*(768+48+50+1) - - # pad_masks = pad_masks.reshape(state.size(0), -1) - # att_masks = att_masks.reshape(state.size(0), -1) - att_2d_masks = make_att_2d_masks(pad_masks, att_masks) # torch.Size([state_bs, 768+48+50+1, 768+48+50+1]) - position_ids = torch.cumsum(pad_masks, dim=1) - 1 # torch.Size([state_bs, 768+48+50+1]) - vlm_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - - # prefix_embs = prefix_embs.reshape(state.size(0), -1, prefix_embs.size(-1)) - # suffix_embs = suffix_embs.reshape(state.size(0), -1, suffix_embs.size(-1)) - (outputs_embeds, suffix_out), _, router_logits_list = self.qwenvl_with_expert.forward( - attention_mask=att_2d_masks, - position_ids=position_ids, - vlm_position_ids=vlm_position_ids, - past_key_values=None, - inputs_embeds=[prefix_embs, suffix_embs], # bs_img,(768+48),2048 [state_bs, (50+1), 1024] - use_cache=self.config.use_cache, - fill_kv_cache=True, - ada_cond = time_embs if getattr(self.config, 'adanorm_time', False) else None, - ) - if self.config.align_params != {}: - loss_depth, depth_preds = self.depth_emb_forward(outputs_embeds, depth_targets, img_masks) - loss_depth = loss_depth * self.config.align_params['depth_loss_weight'] - self.steps+=1 - else: - loss_depth = 0 - depth_preds = None - suffix_out = suffix_out[:, -self.config.n_action_steps :] - if getattr(self.config, 'action_fp32', False): - v_t = self._fp32_linear(self.action_out_proj, suffix_out) - else: - if suffix_out.dtype != self.action_out_proj.weight.dtype: - suffix_out = suffix_out.to(self.action_out_proj.weight.dtype) - v_t = self.action_out_proj(suffix_out) - # u_t = u_t.reshape(images.size(0), -1, u_t.size(-1)) - if loss_type == 'fm': - losses = F.mse_loss(u_t, v_t, reduction="none") - # losses = torch.mean((v_t - u_t)**2, dim=-1) - elif loss_type == 'L1_fm': - losses = F.l1_loss(u_t, v_t, reduction="none") - - # Sequence-wise balance loss (DeepSeek-V3 style, for token-MoE only) - seq_wise_loss_coeff = getattr(self.config, 'sequence_wise_loss_coeff', 0) - seq_wise_loss = 0 - - if seq_wise_loss_coeff > 0 and router_logits_list: - from telefuser.models.lingbot_vla_v2_loader import triton_sequence_wise_balance_loss - - token_moe_layers_set = set(getattr(self.config, 'token_moe_layers', None) or []) - token_moe_layers_list = sorted(token_moe_layers_set) - token_router_logits = tuple( - logits for i, logits in enumerate(router_logits_list) - if not token_moe_layers_list or (token_moe_layers_list[i] if i < len(token_moe_layers_list) else i) in token_moe_layers_set - ) - - if token_router_logits: - token_top_k = getattr(self.config, 'token_top_k', 4) - - # Batch-wise balance loss: treat all B脳T tokens as one group. - # seq_lengths=None makes the function use all tokens at once, - # giving stable f_i statistics (B脳T脳K assignments / E experts). - layer_losses = triton_sequence_wise_balance_loss( - router_logits_list=token_router_logits, - top_k=token_top_k, - seq_lengths=None, - padding_len=0, - ) - if layer_losses: - seq_wise_loss = seq_wise_loss_coeff * torch.stack(layer_losses).mean() - - # MoE monitoring metrics for token-MoE. - moe_metrics = {} - if router_logits_list: - all_moe_indices = sorted(getattr(self.config, 'token_moe_layers', None) or []) - token_expert_counts = [] - - with torch.no_grad(): - for i, logits in enumerate(router_logits_list): - layer_id = all_moe_indices[i] if i < len(all_moe_indices) else i - num_experts = logits.shape[-1] - routing_probs = F.softmax(logits, dim=1, dtype=torch.float) - - moe_block = self.qwenvl_with_expert.qwen_expert.model.layers[layer_id].mlp - if hasattr(moe_block, 'last_tokens_per_expert'): - counts = moe_block.last_tokens_per_expert.clone() - else: - _, selected = torch.topk(routing_probs, 1, dim=-1) - expert_indices = selected.squeeze(-1) - counts = F.one_hot(expert_indices, num_classes=num_experts).float().sum(dim=0) - - token_expert_counts.append((layer_id, counts)) - - # MaxVio: (max_load - avg_load) / avg_load (paper 2408.15664) - avg_load = counts.mean() - maxvio = (counts.max() - avg_load) / avg_load.clamp(min=1e-9) - moe_metrics[f"token_moe/layer{layer_id}_maxvio"] = maxvio - - per_sample_entropy = -(routing_probs * routing_probs.clamp(min=1e-9).log()).sum(dim=-1) - moe_metrics[f"token_moe/layer{layer_id}_entropy"] = per_sample_entropy.mean() - - # Compute average MaxVio across token-MoE layers - token_maxvio_values = [ - moe_metrics[k] for k in moe_metrics - if k.startswith("token_moe/") and k.endswith("_maxvio") - ] - if token_maxvio_values: - moe_metrics["token_moe/avg_maxvio"] = torch.stack(token_maxvio_values).mean() - - # Avg top-K sigmoid score (before norm) across token-MoE layers - token_moe_layers_list = sorted(getattr(self.config, 'token_moe_layers', None) or []) - if token_moe_layers_list: - sigmoid_scores = [] - for lid in token_moe_layers_list: - moe_block = self.qwenvl_with_expert.qwen_expert.model.layers[lid].mlp - if hasattr(moe_block, 'avg_topk_sigmoid_score'): - sigmoid_scores.append(moe_block.avg_topk_sigmoid_score.detach().to(losses.device)) - if sigmoid_scores: - moe_metrics["token_moe/avg_topk_sigmoid"] = torch.stack(sigmoid_scores).mean() - - if token_expert_counts: - moe_metrics["_token_moe_expert_counts"] = token_expert_counts - - return losses, loss_depth, depth_preds, seq_wise_loss, moe_metrics + def forward(self, *args, **kwargs): + """Reject the upstream training API in the inference-only model.""" + del args, kwargs + raise RuntimeError("LingBot-VLA v2 is inference-only; use sample_actions()") def sample_actions( self, images, img_masks, lang_tokens, lang_masks, state, vlm_causal=False, noise=None @@ -1623,200 +1325,6 @@ def predict_velocity(self, state, prefix_pad_masks, past_key_values, x_t, timest v_t = self.action_out_proj(suffix_out) return v_t - def depth_emb_forward(self, hidden_states, depth_targets=None, img_masks=None, future_depth_targets=None): - chunk_size = self.llm_image_token_size * self.llm_image_token_size - num_images = img_masks.shape[1] if img_masks is not None and img_masks.ndim == 2 else 3 - if img_masks is not None: - img_masks = einops.rearrange(img_masks, 'b n -> (b n)') - image_embs = hidden_states[:, chunk_size * 0 + 1 : chunk_size * 1 + 1, :] - align_embs = self._current_depth_task_tokens(hidden_states, num_images=num_images) - align_embs = torch.cat([image_embs, align_embs], dim=1) - depth_preds = self.depth_align_embs.repeat(align_embs.shape[0], 1, 1).to(dtype=align_embs.dtype, device=align_embs.device) - depth_preds = self.depth_align_head(align_embs, depth_preds).contiguous().float() - current_loss = self._emb_loss(depth_preds, depth_targets) - - if self.use_future_depth: - future_align_embs = self._future_depth_task_tokens(hidden_states) - future_image_embs = ( - image_embs.detach() - if getattr(self, "detach_future_depth_image_feats", False) - else image_embs - ) - future_align_embs = torch.cat([future_image_embs, future_align_embs], dim=1) - future_depth_preds = self.future_depth_align_embs.repeat(future_align_embs.shape[0], 1, 1).to(dtype=future_align_embs.dtype, device=future_align_embs.device) - future_depth_preds = self.future_depth_align_head(future_align_embs, future_depth_preds).contiguous().float() - future_loss = self._emb_loss(future_depth_preds, future_depth_targets) - return current_loss, future_loss, depth_preds, future_depth_preds - - return current_loss, 0, depth_preds, None - - def video_emb_forward( - self, - hidden_states, - future_video_targets=None, - future_video_cls_targets=None, - future_video_current_patch=None, - ): - if self.align_type != 'query': - raise ValueError("future-video alignment is only supported for query align mode.") - - use_patch = getattr(self, "use_future_video_patch", True) - use_cls = getattr(self, "use_future_video_cls", False) - if not use_patch and not use_cls: - raise ValueError("future-video alignment requires use_patch_loss or use_cls_loss to be enabled.") - if use_patch and future_video_targets is None: - raise ValueError("future_video_targets is required when use_patch_loss=True.") - - align_params = getattr(getattr(self, "config", None), "align_params", {}) or {} - video_cfg = align_params.get("video", {}) if hasattr(align_params, "get") else {} - chunk_size = self.llm_image_token_size * self.llm_image_token_size - image_embs = hidden_states[:, chunk_size * 0 + 1 : chunk_size * 1 + 1, :] - image_embs_for_video = image_embs.detach() if bool(video_cfg.get("detach_image_feats", False)) else image_embs - - cls_preds = None - if use_cls: - if future_video_cls_targets is None: - raise ValueError("future_video_cls_targets is required when use_cls_loss=True.") - cls_task_embs = self._future_video_cls_task_tokens(hidden_states) - cls_delta = self.future_video_cls_head(cls_task_embs.squeeze(1)) - cls_preds = cls_delta.contiguous().float() - - loss = None - metrics = {} - video_preds = None - if use_patch: - video_task_embs = self._future_video_patch_task_tokens(hidden_states) - context_mode = str( - video_cfg.get( - "context_mode", - getattr(self, "future_video_context_mode", "img_query"), - ) - ).lower() - if context_mode == "query_only": - video_align_embs = video_task_embs - else: - video_align_embs = torch.cat([image_embs_for_video, video_task_embs], dim=1) - if ( - getattr(self, "future_video_share_future_depth_query", False) - and not getattr(self, "use_shared_future_task_proj", False) - ): - query_embs = self.future_depth_align_embs - else: - query_embs = self.future_video_align_embs - video_preds = query_embs.repeat(video_align_embs.shape[0], 1, 1).to( - dtype=video_align_embs.dtype, device=video_align_embs.device - ) - video_preds = self.future_video_align_head(video_align_embs, video_preds).contiguous().float() - loss, metrics = self._video_emb_loss(video_preds, future_video_targets) - if use_cls: - cls_loss, cls_metrics = self._video_cls_loss(cls_preds, future_video_cls_targets) - loss = cls_loss if loss is None else loss + cls_loss - metrics.update(cls_metrics) - return loss, video_preds, metrics - - def current_video_emb_forward( - self, - hidden_states, - current_video_targets=None, - ): - if self.align_type != 'query': - raise ValueError("current-video alignment is only supported for query align mode.") - if not getattr(self, "use_current_video_patch", False): - raise ValueError("current-video alignment requires use_current_patch_loss=True.") - if current_video_targets is None: - raise ValueError("current_video_targets is required for current-video alignment.") - - chunk_size = self.llm_image_token_size * self.llm_image_token_size - image_embs = hidden_states[:, chunk_size * 0 + 1 : chunk_size * 1 + 1, :] - current_task_embs = self._current_depth_task_tokens(hidden_states) - align_embs = torch.cat([image_embs, current_task_embs], dim=1) - queries = self.current_video_align_embs.repeat(align_embs.shape[0], 1, 1).to( - dtype=align_embs.dtype, - device=align_embs.device, - ) - preds = self.current_video_align_head(align_embs, queries).contiguous().float() - loss, metrics = self._video_emb_loss( - preds, - current_video_targets, - metric_prefix="current_video", - ) - return loss, preds, metrics - - def _video_emb_loss(self, video_preds, future_video_targets, metric_prefix="future_video"): - align_params = getattr(getattr(self, "config", None), "align_params", {}) or {} - video_cfg = align_params.get("video", {}) if hasattr(align_params, "get") else {} - use_smooth_l1 = bool(video_cfg.get("use_smooth_l1_loss", True)) - use_mse = bool(video_cfg.get("use_mse_loss", False)) - use_cosine = bool(video_cfg.get("use_cosine_loss", False)) - if not use_smooth_l1 and not use_mse and not use_cosine: - raise ValueError(f"{metric_prefix} loss requires smooth-L1, MSE, and/or cosine loss.") - - metrics = {} - loss = None - if use_smooth_l1: - smooth_l1_loss = self._emb_loss(video_preds, future_video_targets) - metrics[f"align/{metric_prefix}_smooth_l1_loss"] = smooth_l1_loss.detach() - loss = smooth_l1_loss - if use_mse: - target = future_video_targets.to(dtype=video_preds.dtype, device=video_preds.device) - mse_loss = F.mse_loss(video_preds.float(), target.float().detach()) - mse_weight = float(video_cfg.get("mse_loss_weight", 1.0)) - metrics[f"align/{metric_prefix}_mse_loss"] = mse_loss.detach() - metrics[f"align/{metric_prefix}_mse_loss_weighted"] = (mse_loss * mse_weight).detach() - weighted_mse_loss = mse_loss * mse_weight - loss = weighted_mse_loss if loss is None else loss + weighted_mse_loss - if use_cosine: - target = future_video_targets.to(dtype=video_preds.dtype, device=video_preds.device) - pred_norm = F.normalize(video_preds.float(), dim=-1, eps=1e-6) - target_norm = F.normalize(target.float().detach(), dim=-1, eps=1e-6) - cosine_loss = 1.0 - F.cosine_similarity(pred_norm, target_norm, dim=-1, eps=1e-6).mean() - cosine_weight = float(video_cfg.get("cosine_loss_weight", 1.0)) - metrics[f"align/{metric_prefix}_cosine_loss"] = cosine_loss.detach() - metrics[f"align/{metric_prefix}_cosine_loss_weighted"] = (cosine_loss * cosine_weight).detach() - weighted_cosine_loss = cosine_loss * cosine_weight - loss = weighted_cosine_loss if loss is None else loss + weighted_cosine_loss - return loss, metrics - - def _video_cls_loss(self, cls_preds, future_video_cls_targets): - align_params = getattr(getattr(self, "config", None), "align_params", {}) or {} - video_cfg = align_params.get("video", {}) if hasattr(align_params, "get") else {} - cls_loss_type = str(video_cfg.get("cls_loss_type", "cosine")).lower() - cls_weight = float(video_cfg.get("cls_loss_weight", 1.0)) - target = future_video_cls_targets.to(dtype=cls_preds.dtype, device=cls_preds.device) - if target.ndim == 3 and target.shape[1] == 1: - target = target.squeeze(1) - - metrics = {} - loss = None - if cls_loss_type in ("smooth_l1", "smoothl1", "huber"): - smooth_l1_loss = F.smooth_l1_loss(cls_preds.float(), target.float().detach()) - metrics["align/future_video_cls_smooth_l1_loss"] = smooth_l1_loss.detach() - loss = smooth_l1_loss - if cls_loss_type in ("mse", "mse_cosine", "cosine_mse"): - mse_loss = F.mse_loss(cls_preds.float(), target.float().detach()) - metrics["align/future_video_cls_mse_loss"] = mse_loss.detach() - loss = mse_loss - if cls_loss_type in ("cosine", "mse_cosine", "cosine_mse"): - pred_norm = F.normalize(cls_preds.float(), dim=-1, eps=1e-6) - target_norm = F.normalize(target.float().detach(), dim=-1, eps=1e-6) - cosine_loss = 1.0 - F.cosine_similarity(pred_norm, target_norm, dim=-1, eps=1e-6).mean() - metrics["align/future_video_cls_cosine_loss"] = cosine_loss.detach() - loss = cosine_loss if loss is None else loss + cosine_loss - if loss is None: - raise ValueError(f"Unsupported future-video CLS loss type: {cls_loss_type}") - weighted_loss = loss * cls_weight - metrics["align/future_video_cls_loss"] = loss.detach() - metrics["align/future_video_cls_loss_weighted"] = weighted_loss.detach() - return weighted_loss, metrics - - def _emb_loss(self, emb_preds, emb_targets): - l1_loss = F.smooth_l1_loss(emb_preds.float(), emb_targets.float().detach(), reduction="none") - return l1_loss.mean() - -ModelClass = LingbotVlaPolicy - -__all__ = ["LingbotVlaPolicy", "Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLTextModel", "Qwen2ForCausalLM", "Qwen2_5_VLPreTrainedModel"] -# __V1_END__ # Qwen3-VL LingBot-VLA v2 policy. FlowMatchingV1 = FlowMatching @@ -1844,11 +1352,9 @@ def _emb_loss(self, emb_preds, emb_targets): our_eager_attention_forward, prefix_query_segments, prefix_query_token_spans, - sample_beta, ) from telefuser.models.lingbot_vla_v2_loader import build_block_mask, flex_attention_forward, flex_attention_with_block_mask from telefuser.models.lingbot_vla_v2_loader import LingBotVLAWeightLoader -from telefuser.models.lingbot_vla_v2_loader import triton_sequence_wise_balance_loss from telefuser.models.lingbot_vla_v2_moe import ( Qwen2ForCausalLM, Qwen2TokenMoeBlock, @@ -1986,7 +1492,6 @@ def __init__(self, config: QwenvlWithExpertV2Config, eval=False): ) self.attention_interface = self.get_attention_interface() - self.set_requires_grad() def _apply(self, fn): super()._apply(fn) @@ -2030,23 +1535,6 @@ def _install_moe_blocks(self): for idx in token_moe_layers: self.qwen_expert.model.layers[idx].mlp = Qwen2TokenMoeBlock(token_config) - def set_requires_grad(self): - if self.config.freeze_vision_encoder: - self.qwenvl.visual.eval() - for params in self.qwenvl.visual.parameters(): - params.requires_grad = False - if self.config.train_expert_only: - self.qwenvl.eval() - for params in self.qwenvl.parameters(): - params.requires_grad = False - - def train(self, mode: bool = True): - super().train(mode) - if self.config.freeze_vision_encoder: - self.qwenvl.visual.eval() - if self.config.train_expert_only: - self.qwenvl.eval() - def get_image_features( self, pixel_values: torch.FloatTensor, @@ -2334,7 +1822,6 @@ def __init__(self, config, eval): self.future_video_share_future_depth_query = False self.block_future_depth_to_action = False - self.set_requires_grad() def embed_prefix( self, @@ -2604,158 +2091,10 @@ def _current_depth_task_tokens(self, hidden_states, num_images=3): start, end = query_spans["current_depth"] return hidden_states[:, start:end, :] - def forward( - self, - images, - img_masks, - lang_tokens, - lang_masks, - state, - actions, - noise=None, - time=None, - loss_type="fm", - depth_targets=None, - image_grid_thw=None, - future_depth_targets=None, - future_video_targets=None, - future_video_cls_targets=None, - future_video_current_patch=None, - ) -> Tensor: - dtype = state.dtype - device = state.device - if noise is None: - noise = torch.randn(actions.shape, device=device, dtype=dtype) - if time is None: - time = self.sample_time(actions.size(0), device).to(dtype) - - time_expanded = time[:, None, None] - x_t = time_expanded * noise + (1 - time_expanded) * actions - u_t = noise - actions - - ( - prefix_embs, - prefix_pad_masks, - prefix_att_masks, - prefix_position_ids, - visual_pos_masks, - deepstack_visual_embeds, - ) = self.embed_prefix( - images, - img_masks, - lang_tokens, - lang_masks, - image_grid_thw=image_grid_thw, - ) - time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( - state, x_t, time - ) - - pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1) - att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1) - att_2d_masks = make_att_2d_masks(pad_masks, att_masks) - prefix_len = prefix_pad_masks.shape[1] - if self.block_future_depth_to_action: - att_2d_masks = block_suffix_to_fv_( - att_2d_masks, - suffix_row_start=prefix_len, - prefix_len=prefix_len, - num_task_tokens=self.num_task_tokens, - ) - - att_2d_masks = self._block_suffix_to_future_video_if_enabled_( - att_2d_masks, - suffix_row_start=prefix_len, - prefix_len=prefix_len, - ) - position_ids = self._build_full_position_ids(prefix_position_ids, prefix_pad_masks, suffix_pad_masks) - - (outputs_embeds, suffix_out), _, router_logits_list = self.qwenvl_with_expert.forward( - attention_mask=att_2d_masks, - position_ids=position_ids, - vlm_position_ids=prefix_position_ids, - past_key_values=None, - inputs_embeds=[prefix_embs, suffix_embs], - use_cache=self.config.use_cache, - fill_kv_cache=True, - ada_cond=time_embs if getattr(self.config, "adanorm_time", False) else None, - visual_pos_masks=visual_pos_masks, - deepstack_visual_embeds=deepstack_visual_embeds, - ) - align_metrics = {} - if self.config.align_params != {}: - loss_depth, loss_future_depth, depth_preds, future_depth_preds = self.depth_emb_forward(outputs_embeds, depth_targets, img_masks,future_depth_targets,) - loss_depth = loss_depth * self.config.align_params["depth_loss_weight"] - loss_future_depth = loss_future_depth * self.config.align_params.get("future_depth_loss_weight", 1.0) - loss_future_video = 0 - future_video_preds = None - current_video_preds = None - if getattr(self, "use_future_video", False): - loss_video, future_video_preds, video_metrics = self.video_emb_forward( - outputs_embeds, - future_video_targets, - future_video_cls_targets=future_video_cls_targets, - future_video_current_patch=future_video_current_patch, - ) - video_total_loss = loss_video - if ( - getattr(self, "use_current_video_patch", False) - and future_video_current_patch is not None - ): - current_video_loss, current_video_preds, current_video_metrics = self.current_video_emb_forward( - outputs_embeds, - future_video_current_patch, - ) - video_total_loss = video_total_loss + current_video_loss - video_metrics.update(current_video_metrics) - video_metrics["align/current_video_loss"] = current_video_loss.detach() - video_cfg = self.config.align_params.get("video", {}) - video_weight = video_cfg.get( - "future_video_loss_weight", - self.config.align_params.get( - "future_video_loss_weight", - self.config.align_params["depth_loss_weight"], - ), - ) - loss_future_video = video_total_loss * video_weight - align_metrics.update(video_metrics) - if "align/current_video_loss" in align_metrics: - align_metrics["align/current_video_loss_weighted"] = ( - align_metrics["align/current_video_loss"] * video_weight - ) - align_metrics["align/future_video_loss"] = loss_video.detach() - align_metrics["align/future_video_loss_weighted"] = (loss_video * video_weight).detach() - align_metrics["align/video_loss"] = video_total_loss.detach() - align_metrics["align/video_loss_weighted"] = loss_future_video.detach() - self.steps += 1 - else: - loss_depth = 0 - loss_future_depth = 0 - loss_future_video = 0 - depth_preds = None - future_depth_preds = None - future_video_preds = None - current_video_preds = None - - suffix_out = suffix_out[:, -self.config.n_action_steps :] - if getattr(self.config, "action_fp32", False): - v_t = self._fp32_linear(self.action_out_proj, suffix_out) - else: - if suffix_out.dtype != self.action_out_proj.weight.dtype: - suffix_out = suffix_out.to(self.action_out_proj.weight.dtype) - v_t = self.action_out_proj(suffix_out) - - if loss_type == "fm": - losses = F.mse_loss(u_t, v_t, reduction="none") - elif loss_type == "L1_fm": - losses = F.l1_loss(u_t, v_t, reduction="none") - - seq_wise_loss, router_z_loss, moe_metrics = self._moe_losses_and_metrics( - router_logits_list, losses - ) - if align_metrics: - moe_metrics.update(align_metrics) - return losses, loss_depth, loss_future_depth, loss_future_video, depth_preds, seq_wise_loss, router_z_loss, moe_metrics, future_depth_preds, future_video_preds, current_video_preds + def forward(self, *args, **kwargs): + """Reject the upstream training API in the inference-only model.""" + del args, kwargs + raise RuntimeError("LingBot-VLA v2 is inference-only; use sample_actions()") def sample_actions( self, @@ -2910,153 +2249,19 @@ def predict_velocity( v_t = self.action_out_proj(suffix_out) return v_t - def _moe_losses_and_metrics(self, router_logits_list, losses): - router_z_loss_coeff = getattr(self.config, "router_z_loss_coeff", 0) - router_z_loss = losses.new_zeros(()) - router_z_layer_losses = None # per-layer raw z-loss (pre-coeff), for monitoring - if router_z_loss_coeff > 0 and router_logits_list: - router_z_layer_losses = [ - torch.logsumexp(logits.float(), dim=-1).pow(2).mean() - for logits in router_logits_list - ] - router_z_loss = router_z_loss_coeff * torch.stack(router_z_layer_losses).mean() - - seq_wise_loss_coeff = getattr(self.config, "sequence_wise_loss_coeff", 0) - seq_wise_loss = 0 - seqwise_layer_losses = None # per-layer raw seq-wise balance loss (pre-coeff), for monitoring - if seq_wise_loss_coeff > 0 and router_logits_list: - # router_logits are [B*T, E] (action-expert tokens, fixed length T per sample). - # per_sequence -> balance experts within each sample's T tokens (DeepSeek-V3 intent); - # global -> treat the whole B*T batch as one sequence. - mode = getattr(self.config, "sequence_wise_mode", "per_sequence") - score_func = getattr(self.config, "router_activation", "softmax") - if mode == "global": - seq_lengths = None - else: - B = losses.shape[0] - N = router_logits_list[0].shape[0] - seq_lengths = [N // B] * B - seqwise_layer_losses = triton_sequence_wise_balance_loss( - router_logits_list=tuple(router_logits_list), - top_k=getattr(self.config, "token_top_k", 4), - seq_lengths=seq_lengths, - padding_len=0, - score_func=score_func, - ) - if seqwise_layer_losses: - seq_wise_loss = seq_wise_loss_coeff * torch.stack(seqwise_layer_losses).mean() - - moe_metrics = {} - if router_logits_list: - token_moe_layers_list = sorted(getattr(self.config, "token_moe_layers", None) or []) - all_moe_indices = token_moe_layers_list - token_expert_counts = [] - # Per-layer token-MoE stats, collected for moe_summary/* cross-layer aggregates. - tok_maxvio, tok_minvio, tok_minload, tok_entropy, tok_sigmoid = [], [], [], [], [] - tok_bias = [] # per-layer max(|e_score_correction_bias|) (loss-free); >1 -> bias dominates sigmoid score - any_dead = None # OR-accumulated bool: any token-MoE layer with a 0-count expert - with torch.no_grad(): - for i, logits in enumerate(router_logits_list): - layer_id = all_moe_indices[i] if i < len(all_moe_indices) else i - num_experts = logits.shape[-1] - routing_probs = F.softmax(logits, dim=1, dtype=torch.float) - moe_block = self.qwenvl_with_expert.qwen_expert.model.layers[layer_id].mlp - if hasattr(moe_block, "last_tokens_per_expert"): - # Global (all-reduced), biased, true top-k load from the load-balance hook. - counts = moe_block.last_tokens_per_expert.clone() - if counts.sum() == 0: - # Buffer not yet populated by the load-balance hook (first step - # after run start / resume) -> skip this layer to avoid a spurious - # has_dead_expert / min_load_ratio spike on the very first viz. - continue - else: - _, selected = torch.topk(routing_probs, 1, dim=-1) - counts = F.one_hot(selected.squeeze(-1), num_classes=num_experts).float().sum(dim=0) - avg_load = counts.mean() - denom = avg_load.clamp(min=1e-9) - maxvio = (counts.max() - avg_load) / denom # peak overload (>=0, larger=worse) - minvio = (avg_load - counts.min()) / denom # valley underload (=1 -> dead expert) - min_load_ratio = counts.min() / denom # =0 -> dead expert - # entropy is rank-local (this rank's routing_probs, last micro-batch). - per_sample_entropy = -(routing_probs * routing_probs.clamp(min=1e-9).log()).sum(dim=-1) - entropy = per_sample_entropy.mean() - ll = f"{layer_id:02d}" - token_expert_counts.append((layer_id, counts)) - moe_metrics[f"moe_maxvio/layer{ll}"] = maxvio - moe_metrics[f"moe_minvio/layer{ll}"] = minvio - moe_metrics[f"moe_minload/layer{ll}"] = min_load_ratio - moe_metrics[f"moe_entropy_rank0/layer{ll}"] = entropy - tok_maxvio.append(maxvio) - tok_minvio.append(minvio) - tok_minload.append(min_load_ratio) - tok_entropy.append(entropy) - dead = counts.min() == 0 - any_dead = dead if any_dead is None else (any_dead | dead) - if hasattr(moe_block, "avg_topk_sigmoid_score"): - sig = moe_block.avg_topk_sigmoid_score.detach().reshape(()).to(denom) - moe_metrics[f"moe_topksigmoid_rank0/layer{ll}"] = sig - tok_sigmoid.append(sig) - if hasattr(moe_block, "e_score_correction_bias"): - bias_absmax = moe_block.e_score_correction_bias.detach().abs().max().to(denom) - moe_metrics[f"moe_bias/layer{ll}"] = bias_absmax - tok_bias.append(bias_absmax) - # ---- moe_summary/* : cross-layer aggregates over token-MoE layers (written every step) ---- - if tok_maxvio: - moe_metrics["moe_summary/maxvio_avg"] = torch.stack(tok_maxvio).mean() - moe_metrics["moe_summary/maxvio_max"] = torch.stack(tok_maxvio).max() - moe_metrics["moe_summary/minvio_avg"] = torch.stack(tok_minvio).mean() - moe_metrics["moe_summary/minvio_max"] = torch.stack(tok_minvio).max() - moe_metrics["moe_summary/min_load_ratio"] = torch.stack(tok_minload).min() - moe_metrics["moe_summary/has_dead_expert"] = any_dead.float() - moe_metrics["moe_summary/entropy_avg_rank0"] = torch.stack(tok_entropy).mean() - if tok_sigmoid: - moe_metrics["moe_summary/topk_sigmoid_avg_rank0"] = torch.stack(tok_sigmoid).mean() - if tok_bias: - moe_metrics["moe_summary/bias_absmax"] = torch.stack(tok_bias).max() - # ---- moe_seqwise/* : per-layer raw sequence-wise balance loss (pre-coeff) + average ---- - if seqwise_layer_losses and len(seqwise_layer_losses) == len(all_moe_indices): - sw_vals = [] - for lid, sw in zip(all_moe_indices, seqwise_layer_losses): - v = sw.detach() - moe_metrics[f"moe_seqwise/layer{lid:02d}"] = v - sw_vals.append(v) - moe_metrics["moe_seqwise/avg"] = torch.stack(sw_vals).mean() - # ---- moe_zloss/* : per-layer raw router z-loss (pre-coeff) + average/weighted loss ---- - if router_z_layer_losses and len(router_z_layer_losses) == len(all_moe_indices): - zl_vals = [] - for lid, zl in zip(all_moe_indices, router_z_layer_losses): - v = zl.detach() - moe_metrics[f"moe_zloss/layer{lid:02d}"] = v - zl_vals.append(v) - moe_metrics["moe_zloss/avg_raw"] = torch.stack(zl_vals).mean() - moe_metrics["moe_zloss/weighted"] = router_z_loss.detach() - if token_expert_counts: - moe_metrics["_token_moe_expert_counts"] = token_expert_counts - return seq_wise_loss, router_z_loss, moe_metrics - class LingbotVlaV2Policy(PreTrainedModel): config_class = LingbotVLAV2Config name = "torch_lingbot_vla_v2" - supports_gradient_checkpointing = True _no_split_modules = ["Qwen2DecoderLayer", "FixQwen2RMSNorm", "FixAdaRMSNorm"] - def get_parallel_plan(self): - from telefuser.models.lingbot_vla_v2_loader import NativeParallelPlan as ParallelPlan - from torch.distributed._tensor import Shard - - ep_plan = { - "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.gate_proj": Shard(0), - "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.up_proj": Shard(0), - "model.qwenvl_with_expert.qwen_expert.model.layers.*.mlp.experts.down_proj": Shard(0), - } - return ParallelPlan(ep_plan=ep_plan) - @classmethod def get_weight_loader(cls): return LingBotVLAWeightLoader() - def __init__(self, config: LingbotVLAV2Config, eval: bool = False): + def __init__(self, config: LingbotVLAV2Config, eval: bool = True): + if not eval: + raise ValueError("LingBot-VLA v2 only supports inference mode") super().__init__(config) self.config = config self.language_tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, local_files_only=True) @@ -3064,96 +2269,18 @@ def __init__(self, config: LingbotVLAV2Config, eval: bool = False): if not getattr(self.config, "use_lm_head", False): del self.model.qwenvl_with_expert.qwenvl.lm_head del self.model.qwenvl_with_expert.qwen_expert.lm_head + self.requires_grad_(False) + self.eval() self.reset() torch.set_float32_matmul_precision("high") def reset(self): return None - def get_optim_params(self) -> dict: - return self.parameters() - - def forward( - self, - images, - img_masks, - state, - lang_tokens, - lang_masks, - actions, - joint_mask=None, - action_is_pad=None, - noise=None, - time=None, - depth_targets=None, - image_grid_thw=None, - future_depth_targets=None, - future_video_targets=None, - future_video_cls_targets=None, - future_video_current_patch=None, - **kwargs - ) -> tuple[Tensor, dict[str, Tensor]]: - loss_dict = {} - if getattr(self.config, "action_fp32", False): - state = state.float() - actions = actions.float() - ( - losses, - loss_depth, - loss_future_depth, - loss_future_video, - depth_preds, - seq_wise_loss, - router_z_loss, - moe_metrics, - future_depth_preds, - future_video_preds, - current_video_preds, - ) = self.model.forward( - images, - img_masks, - lang_tokens, - lang_masks, - state, - actions, - noise, - time, - loss_type=self.config.loss_type, - depth_targets=depth_targets, - image_grid_thw=image_grid_thw, - future_depth_targets=future_depth_targets, - future_video_targets=future_video_targets, - future_video_cls_targets=future_video_cls_targets, - future_video_current_patch=future_video_current_patch, - ) - - if joint_mask is not None: - if "repeat" in self.config.loss_type: - joint_mask = joint_mask.repeat(2, 1, 1) - assert len(joint_mask.shape) == 3 - - masked_losses = losses * joint_mask - valid_counts = joint_mask.sum(dim=(1, 2)).clamp(min=1) - batch_mean_losses = masked_losses.sum(dim=(1, 2)) / valid_counts - loss_vla = masked_losses.sum() / joint_mask.sum().clamp(min=1) - else: - losses = losses[:, :, : self.config.action_dim] - batch_mean_losses = losses.mean(dim=(1, 2)) - loss_vla = losses.mean() - - loss_dict["batch_mean_losses"] = batch_mean_losses.detach() - total_loss = ( - loss_vla - + loss_depth - + loss_future_depth - + loss_future_video - + seq_wise_loss - + router_z_loss - ) - loss_dict["router_z_loss"] = router_z_loss.detach() if torch.is_tensor(router_z_loss) else router_z_loss - if moe_metrics: - loss_dict.update(moe_metrics) - return total_loss, loss_vla, loss_depth, loss_future_depth, loss_future_video, seq_wise_loss, loss_dict, depth_preds, future_depth_preds, future_video_preds, current_video_preds + def forward(self, *args, **kwargs): + """Reject the upstream training API in the inference-only model.""" + del args, kwargs + raise RuntimeError("LingBot-VLA v2 is inference-only; use sample_actions()") def sample_actions(self, *args, **kwargs) -> Tensor: return self.model.sample_actions(*args, **kwargs) diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py index 5073a1b0..95b0c160 100644 --- a/telefuser/models/lingbot_vla_v2_loader.py +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -60,10 +60,6 @@ def create_sinusoidal_pos_embedding( return pos_emb -def sample_beta(alpha, beta, bsize, device): - gamma1 = torch.rand((bsize,), device=device).pow(1 / alpha) - gamma2 = torch.rand((bsize,), device=device).pow(1 / beta) - return gamma1 / (gamma1 + gamma2) def make_att_2d_masks(pad_masks, att_masks): @@ -960,416 +956,6 @@ def __init__( def forward(self, llm_feats, queries): queries = self.projector(llm_feats, queries) return queries - - - -# Copyright 2025 Ant Group Co., Ltd. All Rights Reserved. -# Developer: xiancun -# Project锛?Lumos VIdeo Generation Foundation Model -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Triton-optimized MoE auxiliary loss functions. - -Provides numerically equivalent replacements for the functions in loss.py, -with two key optimizations: - 1. Eliminate Python for-loops via vectorized segment-wise operations. - 2. Fuse topK + counting into Triton kernels to avoid huge intermediate tensors. - -Usage: - from telefuser.models.lingbot_vla_v2_loader import ( - triton_load_balancing_loss_func, - triton_sequence_wise_balance_loss, - ) - # Drop-in replacement 鈥?same signature and return type as loss.py -""" - -import torch -import torch.nn.functional as F -from typing import List, Optional, Tuple, Union - - -def _next_power_of_2(n: int) -> int: - """Return the smallest power of 2 >= n.""" - if n <= 0: - return 1 - n -= 1 - n |= n >> 1 - n |= n >> 2 - n |= n >> 4 - n |= n >> 8 - n |= n >> 16 - return n + 1 - - -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? -# Section 1: Triton availability check + kernel definitions -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? - -_HAS_TRITON = False -# Training-only MoE auxiliary losses keep the PyTorch vectorized path in the -# model loader. Triton kernels live only under telefuser.kernel.triton. - - - -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? -# Section 2: Vectorized PyTorch fallback (no Triton needed) -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? - -def _build_segment_info( - seq_lengths_per_layer: List[int], - num_layers: int, - N_per_layer: int, - device: torch.device, -): - """Build segment IDs and metadata for all layers combined. - - Returns: - segment_ids: [num_layers * N_valid_per_layer] int64 - seg_starts: [total_segments] int64 - seg_lengths: [total_segments] int64 - total_segments: int - """ - S = len(seq_lengths_per_layer) - total_segments = num_layers * S - seg_lengths_t = torch.tensor(seq_lengths_per_layer, dtype=torch.int64, device=device) - - # Repeat for all layers - all_seg_lengths = seg_lengths_t.repeat(num_layers) # [total_segments] - - # Per-layer segment starts: cumsum of seq_lengths - per_layer_starts = torch.cumsum(seg_lengths_t, 0) - seg_lengths_t # [S] - # Layer offsets in the concatenated tensor - layer_offsets = torch.arange(num_layers, device=device, dtype=torch.int64) * N_per_layer - # [L, S] 鈫?[L*S] - all_seg_starts = (per_layer_starts.unsqueeze(0) + layer_offsets.unsqueeze(1)).reshape(-1) - - # Segment IDs: [num_layers * N_valid] - segment_ids = torch.repeat_interleave( - torch.arange(total_segments, device=device), all_seg_lengths, - ) - - return segment_ids, all_seg_starts, all_seg_lengths, total_segments - - -def _vectorized_segment_f_i( - logits: torch.Tensor, # [N_total, E] - seg_starts: torch.Tensor, # [S_total] - seg_lengths: torch.Tensor, # [S_total] - top_k: int, -) -> torch.Tensor: - """Compute f_i per segment without Triton, using vectorized PyTorch ops. - - Returns: f_i [S_total, E] - """ - N, E = logits.shape - S_total = seg_starts.shape[0] - - # TopK over all tokens at once - _, topk_idx = torch.topk(logits, k=top_k, dim=-1) # [N, K] - - # Build one-hot mask efficiently: scatter into [N, E] - mask = torch.zeros(N, E, device=logits.device, dtype=torch.float32) - mask.scatter_(1, topk_idx, 1.0) - - # Segment-wise sum using scatter_add - segment_ids = torch.repeat_interleave( - torch.arange(S_total, device=logits.device), seg_lengths, - ) - seg_ids_exp = segment_ids.unsqueeze(1).expand(-1, E) # [N, E] - - f_sum = torch.zeros(S_total, E, device=logits.device, dtype=torch.float32) - f_sum.scatter_add_(0, seg_ids_exp, mask) - - # Normalize: f_i = (E / K) * f_sum / T_s - inv_lens = (float(E) / top_k) / seg_lengths.unsqueeze(1).float().clamp(min=1) - f_i = f_sum * inv_lens - - return f_i - - -def _vectorized_topk_count( - routing_weights: torch.Tensor, # [N, E] - top_k: int, - flat_mask: Optional[torch.Tensor] = None, # [N] float -) -> torch.Tensor: - """Count per-expert topK selections, optionally masked. Returns [E].""" - N, E = routing_weights.shape - _, topk_idx = torch.topk(routing_weights, k=top_k, dim=-1) # [N, K] - - tokens_per_expert = torch.zeros(E, device=routing_weights.device, dtype=torch.float32) - weight = flat_mask if flat_mask is not None else torch.ones(N, device=routing_weights.device, dtype=torch.float32) - - # K rounds of scatter_add 鈥?K is small (typically 2-8), no Python overhead concern - for k in range(top_k): - tokens_per_expert.scatter_add_(0, topk_idx[:, k], weight) - - return tokens_per_expert - - -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? -# Section 3: Triton-accelerated wrappers -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? - -def _triton_segment_f_i(*args, **kwargs): - raise RuntimeError("LingBot-VLA v2 sequence-wise Triton loss was removed from models; use the PyTorch fallback") - - -def _triton_topk_count(*args, **kwargs): - raise RuntimeError("LingBot-VLA v2 load-balancing Triton loss was removed from models; use the PyTorch fallback") - - - -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? -# Section 4: Main API 鈥?drop-in replacements for loss.py -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? - -def triton_sequence_wise_balance_loss( - router_logits_list: tuple, - top_k: int, - seq_lengths: Optional[List[int]] = None, - padding_len: int = 0, - score_func: str = "softmax", -) -> List[torch.Tensor]: - """Triton-optimized DeepSeek-V3 sequence-wise balance loss. - - Numerically equivalent to sequence_wise_balance_loss() in loss.py, - but eliminates all Python for-loops by: - - Processing all layers simultaneously via concatenation - - Using segment-wise parallel reduction (scatter_add) instead of per-sequence loops - - Fusing topK + counting in a Triton kernel (with PyTorch vectorized fallback) - - Args / Returns: same as sequence_wise_balance_loss in loss.py. - """ - if router_logits_list is None or not isinstance(router_logits_list, (tuple, list)): - return [] - - valid_logits = [rl for rl in router_logits_list if rl is not None] - if len(valid_logits) == 0: - return [] - - num_layers = len(valid_logits) - device = valid_logits[0].device - E = valid_logits[0].shape[1] - - # 鈹€鈹€ Step 1: Concatenate all layers, remove padding 鈹€鈹€ - all_logits_list = [] - N_per_layer = None - for logits in valid_logits: - logits_f32 = logits.to(dtype=torch.float32) - N = logits_f32.shape[0] - if padding_len > 0: - logits_f32 = logits_f32[:N - padding_len] - all_logits_list.append(logits_f32) - if N_per_layer is None: - N_per_layer = logits_f32.shape[0] - - # Check if all layers have the same valid length (common case) - same_length = all(l.shape[0] == N_per_layer for l in all_logits_list) - - if not same_length: - # Rare: different MoE layers have different token counts - return _fallback_per_layer(valid_logits, top_k, seq_lengths, padding_len, score_func) - - if seq_lengths is None or len(seq_lengths) == 0: - seq_lengths_effective = [N_per_layer] - else: - seq_lengths_effective = seq_lengths - - S = len(seq_lengths_effective) - all_logits = torch.cat(all_logits_list, dim=0) # [L * N_valid, E] - - # 鈹€鈹€ Step 2: Build segment metadata 鈹€鈹€ - segment_ids, seg_starts, seg_lengths_t, total_segments = _build_segment_info( - seq_lengths_effective, num_layers, N_per_layer, device - ) - - # 鈹€鈹€ Step 3: P_i via PyTorch (gradient path) 鈹€鈹€ - if score_func == "sigmoid": - all_scores = all_logits.sigmoid() - all_probs = all_scores / all_scores.sum(dim=-1, keepdim=True) - else: - all_probs = F.softmax(all_logits, dim=-1) # [L * N_valid, E] - seg_ids_exp = segment_ids.unsqueeze(1).expand(-1, E) # [L * N_valid, E] - - P_sum = torch.zeros(total_segments, E, device=device, dtype=torch.float32) - P_sum.scatter_add_(0, seg_ids_exp, all_probs) - P_i = P_sum / seg_lengths_t.unsqueeze(1).float().clamp(min=1) # [total_segments, E] - - # 鈹€鈹€ Step 4: f_i (no gradient needed) 鈹€鈹€ - with torch.no_grad(): - if _HAS_TRITON and all_logits.is_cuda: - f_i = _triton_segment_f_i(all_logits, seg_starts, seg_lengths_t, top_k) - else: - f_i = _vectorized_segment_f_i(all_logits, seg_starts, seg_lengths_t, top_k) - - # 鈹€鈹€ Step 5: Per-segment loss 鈫?per-layer mean 鈹€鈹€ - loss_per_seg = (f_i * P_i).sum(dim=-1) # [total_segments] - loss_per_seg = loss_per_seg.reshape(num_layers, S) - layer_losses = loss_per_seg.mean(dim=1) # [L] - - return list(layer_losses.unbind(0)) - - -def triton_load_balancing_loss_func( - gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None], - num_experts: Optional[int] = None, - top_k: int = 2, - attention_mask: Optional[torch.Tensor] = None, -) -> Union[torch.Tensor, int]: - """Triton-optimized Switch Transformer load balancing loss. - - Numerically equivalent to load_balancing_loss_func() in loss.py, - but avoids the huge [N, K, E] one_hot intermediate tensor by - directly counting expert assignments via Triton or scatter_add. - - Memory reduction: O(N*K*E) 鈫?O(N*E + num_blocks*E) - - Args / Returns: same as load_balancing_loss_func in loss.py. - """ - if gate_logits is None or not isinstance(gate_logits, tuple): - return 0 - - gate_logits = tuple(g for g in gate_logits if g is not None) - if len(gate_logits) == 0: - return 0 - - compute_device = gate_logits[0].device - concatenated = torch.cat( - [g.to(device=compute_device, dtype=torch.float32) for g in gate_logits], - dim=0, - ) # [L*N, E] - - # 鈹€鈹€ Step 1: softmax (gradient path) 鈹€鈹€ - routing_weights = F.softmax(concatenated, dim=-1) # [L*N, E] - - # 鈹€鈹€ Step 2: Build flat mask 鈹€鈹€ - N_total = routing_weights.shape[0] - if attention_mask is not None: - batch_size, seq_len = attention_mask.shape - num_layers = N_total // (batch_size * seq_len) - flat_mask = ( - attention_mask - .unsqueeze(0) - .expand(num_layers, -1, -1) - .reshape(-1) - .to(device=compute_device, dtype=torch.float32) - ) - else: - flat_mask = None - - # 鈹€鈹€ Step 3: tokens_per_expert (no gradient) 鈹€鈹€ - with torch.no_grad(): - if _HAS_TRITON and routing_weights.is_cuda: - tokens_per_expert = _triton_topk_count(routing_weights, top_k, flat_mask) - else: - tokens_per_expert = _vectorized_topk_count(routing_weights, top_k, flat_mask) - if flat_mask is not None: - tokens_per_expert = tokens_per_expert / flat_mask.sum().clamp(min=1) - else: - tokens_per_expert = tokens_per_expert / float(N_total) - - # 鈹€鈹€ Step 4: router_prob_per_expert (gradient path) 鈹€鈹€ - if flat_mask is not None: - n_valid = flat_mask.sum().clamp(min=1) - router_prob_per_expert = (routing_weights * flat_mask.unsqueeze(1)).sum(0) / n_valid - else: - router_prob_per_expert = routing_weights.mean(dim=0) - - # 鈹€鈹€ Step 5: loss 鈹€鈹€ - overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert) - return overall_loss * num_experts - - -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? -# Section 5: Fallback for edge cases -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? - -def _fallback_per_layer( - valid_logits: List[torch.Tensor], - top_k: int, - seq_lengths: Optional[List[int]], - padding_len: int, - score_func: str = "softmax", -) -> List[torch.Tensor]: - """Fallback when layers have different valid token counts. - - Still vectorized within each layer (no per-sequence for-loop). - """ - layer_loss_list = [] - for logits in valid_logits: - logits = logits.to(dtype=torch.float32) - N, E = logits.shape - if padding_len > 0: - logits = logits[:N - padding_len] - if logits.shape[0] == 0: - continue - - if seq_lengths is not None and len(seq_lengths) > 0: - S = len(seq_lengths) - device = logits.device - seg_lengths_t = torch.tensor(seq_lengths, dtype=torch.int64, device=device) - seg_starts = torch.cumsum(seg_lengths_t, 0) - seg_lengths_t - - # P_i (gradient path) - if score_func == "sigmoid": - scores = logits.sigmoid() - probs = scores / scores.sum(dim=-1, keepdim=True) - else: - probs = F.softmax(logits, dim=-1) - segment_ids = torch.repeat_interleave(torch.arange(S, device=device), seg_lengths_t) - seg_ids_exp = segment_ids.unsqueeze(1).expand(-1, E) - P_sum = torch.zeros(S, E, device=device, dtype=torch.float32) - P_sum.scatter_add_(0, seg_ids_exp, probs) - P_i = P_sum / seg_lengths_t.unsqueeze(1).float().clamp(min=1) - - # f_i (no gradient) - with torch.no_grad(): - if _HAS_TRITON and logits.is_cuda: - f_i = _triton_segment_f_i(logits, seg_starts, seg_lengths_t, top_k) - else: - f_i = _vectorized_segment_f_i(logits, seg_starts, seg_lengths_t, top_k) - - loss_per_seq = (f_i * P_i).sum(dim=-1) - layer_loss_list.append(loss_per_seq.mean()) - else: - if score_func == "sigmoid": - scores = logits.sigmoid() - probs = scores / scores.sum(dim=-1, keepdim=True) - else: - probs = F.softmax(logits, dim=-1) - P_i = probs.mean(dim=0) - - with torch.no_grad(): - _, topk_idx = torch.topk(logits, k=top_k, dim=-1) - mask = torch.zeros_like(logits) - mask.scatter_(1, topk_idx, 1.0) - f_i = (E / top_k) * mask.mean(dim=0) - - layer_loss_list.append(torch.sum(f_i * P_i)) - - return layer_loss_list - - -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? -# Section 6: Numerical alignment test -# 鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺愨晲鈺? - - - - - import json from copy import deepcopy from pathlib import Path @@ -1378,12 +964,6 @@ def _fallback_per_layer( from transformers import AutoConfig -class NativeParallelPlan: - """Compatibility container for the upstream training-only parallel plan.""" - - def __init__(self, ep_plan=None): - self.ep_plan = ep_plan or {} - class LingBotVLAWeightLoader: """Minimal native weight-name mapper retained for model compatibility.""" diff --git a/tests/unit/models/test_lingbot_vla_v2.py b/tests/unit/models/test_lingbot_vla_v2.py index b80e0924..daa829da 100644 --- a/tests/unit/models/test_lingbot_vla_v2.py +++ b/tests/unit/models/test_lingbot_vla_v2.py @@ -1,8 +1,10 @@ from types import SimpleNamespace +import pytest import torch -from telefuser.models.lingbot_vla_v2 import QwenvlWithExpertV2Model +from telefuser.models import lingbot_vla_v2_loader +from telefuser.models.lingbot_vla_v2 import LingbotVlaV2Policy, QwenvlWithExpertV2Model class _Visual: @@ -51,3 +53,18 @@ def test_image_grid_cache_is_reused_and_invalidated_by_grid_shape() -> None: assert visual.preprocess_calls == 2 assert first[0].shape == repeated[0].shape == (2, 4, 3) assert changed[0].shape == (2, 2, 3) + + +def test_policy_rejects_training_entrypoints() -> None: + with pytest.raises(RuntimeError, match="inference-only"): + LingbotVlaV2Policy.forward(None) + with pytest.raises(ValueError, match="only supports inference mode"): + LingbotVlaV2Policy.__init__(None, SimpleNamespace(), eval=False) + + assert "get_optim_params" not in LingbotVlaV2Policy.__dict__ + assert "get_parallel_plan" not in LingbotVlaV2Policy.__dict__ + + +def test_loader_does_not_expose_training_loss_helpers() -> None: + assert not hasattr(lingbot_vla_v2_loader, "triton_sequence_wise_balance_loss") + assert not hasattr(lingbot_vla_v2_loader, "triton_load_balancing_loss_func") From 3b49c71c3dbbf98748739c30361dffa53770cc07 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Wed, 5 Aug 2026 02:35:51 +0000 Subject: [PATCH 07/15] feat(vla): add minimal single-gpu action service Add a LingBot VLA v2-specific FastAPI application with bounded Base64 image validation, one process-local policy replica, serialized inference, readiness reporting, and normalized canonical action responses. Factor the official 6B runtime construction into a reusable VLA helper, share it with the offline CLI, document local .venv-vla startup and request usage, and cover the HTTP contract, validation, lifecycle cleanup, and request serialization. Verification: 33 focused LingBot VLA v2 tests passed; Ruff lint and format checks passed; git diff --check passed; a real 6B single-GPU HTTP request returned a finite 50x55 action chunk. --- examples/lingbot_vla_v2/README.md | 57 +++++ .../lingbot_vla_v2_inference.py | 26 +-- .../lingbot_vla_v2/lingbot_vla_v2_server.py | 38 ++++ telefuser/pipelines/lingbot_vla_v2/runtime.py | 47 +++++ telefuser/pipelines/lingbot_vla_v2/service.py | 196 ++++++++++++++++++ .../pipelines/lingbot_vla_v2/test_service.py | 161 ++++++++++++++ 6 files changed, 501 insertions(+), 24 deletions(-) create mode 100644 examples/lingbot_vla_v2/lingbot_vla_v2_server.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/runtime.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/service.py create mode 100644 tests/unit/pipelines/lingbot_vla_v2/test_service.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index b3ff9000..b501da66 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -45,6 +45,63 @@ python examples/lingbot_vla_v2/lingbot_vla_v2_inference.py \ The example saves canonical actions and checkpoint metadata in an `.npz` file. The base output must not be sent to a robot without an embodiment-specific post-training checkpoint, action mapping, and policy validation. +## Minimal Single-GPU HTTP Service + +The VLA-specific server loads one policy replica and serializes all inference calls on the selected GPU. It does not +use the shared media service, Ray, multi-GPU execution, dynamic batching, or robot control. Start it from the repository +with the isolated VLA environment: + +```bash +.venv-vla/bin/python examples/lingbot_vla_v2/lingbot_vla_v2_server.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --device cuda:0 \ + --host 127.0.0.1 \ + --port 8000 +``` + +The process reports ready only after both the processor and policy have loaded: + +```bash +curl http://127.0.0.1:8000/health +``` + +`POST /v1/vla/actions` accepts raw Base64 or a Base64 data URL for each camera. The state must contain exactly 14 +finite values. For example: + +```bash +.venv-vla/bin/python - <<'PY' +import base64 +from pathlib import Path + +import httpx + + +def encode(path: str) -> str: + return base64.b64encode(Path(path).read_bytes()).decode("ascii") + + +response = httpx.post( + "http://127.0.0.1:8000/v1/vla/actions", + json={ + "task": "pick up the red block", + "state": [0.0] * 14, + "camera_high": encode("/data/cam_high.png"), + "camera_left_wrist": encode("/data/cam_left_wrist.png"), + "camera_right_wrist": encode("/data/cam_right_wrist.png"), + "seed": 7, + }, + timeout=300.0, +) +response.raise_for_status() +print(response.json()) +PY +``` + +The response contains `canonical_normalized_actions`, `horizon`, `action_dim`, `checkpoint_variant`, +`policy_verified`, and `verification_status`. A successful HTTP response confirms service and model execution only; +the normalized base-model output is not a physical robot command. + ## TeleFuser Regression Baseline The validation capture runs through the public loader and pipeline, then records preprocessing tensors, fixed initial diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py index 3484c0d1..608e9bff 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py @@ -6,18 +6,13 @@ import click import numpy as np -import torch -from transformers import AutoProcessor -from telefuser.core.config import ModelRuntimeConfig -from telefuser.core.module_manager import ModuleManager -from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2 from telefuser.pipelines.lingbot_vla_v2 import ( ROBOTWIN_CAMERA_KEYS, LingBotVlaV2Observation, LingBotVlaV2Pipeline, - LingBotVlaV2PipelineConfig, ) +from telefuser.pipelines.lingbot_vla_v2.runtime import get_lingbot_vla_v2_pipeline def get_pipeline( @@ -26,24 +21,7 @@ def get_pipeline( device: str = "cuda", ) -> LingBotVlaV2Pipeline: """Load the official 6B checkpoint and Qwen3-VL processor.""" - target_device = torch.device(device) - dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 - processor = AutoProcessor.from_pretrained(qwen3vl_root, local_files_only=True, padding_side="right") - manager = ModuleManager(torch_dtype=dtype, device="cpu") - manager.add_module(processor, "lingbot_vla_v2_processor", path=qwen3vl_root) - load_lingbot_vla_v2(manager, model_root, qwen3vl_root, torch_dtype=dtype) - pipeline = LingBotVlaV2Pipeline(device=device, torch_dtype=dtype) - pipeline.init( - manager, - LingBotVlaV2PipelineConfig( - policy_config=ModelRuntimeConfig( - device_type=target_device.type, - device_id=target_device.index or 0, - torch_dtype=dtype, - ), - ), - ) - return pipeline + return get_lingbot_vla_v2_pipeline(model_root, qwen3vl_root, device=device) @click.command() diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_server.py new file mode 100644 index 00000000..22dfecbc --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_server.py @@ -0,0 +1,38 @@ +"""Start a minimal single-GPU LingBot-VLA v2 HTTP service.""" + +from __future__ import annotations + +import click +import uvicorn + +from telefuser.pipelines.lingbot_vla_v2.service import LingBotVlaV2ServiceConfig, create_lingbot_vla_v2_app + + +@click.command() +@click.option("--model-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--qwen3vl-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--device", default="cuda:0", show_default=True) +@click.option("--host", default="127.0.0.1", show_default=True) +@click.option("--port", default=8000, show_default=True, type=click.IntRange(1, 65535)) +@click.option("--max-image-mb", default=10, show_default=True, type=click.IntRange(1, 100)) +def main( + model_root: str, + qwen3vl_root: str, + device: str, + host: str, + port: int, + max_image_mb: int, +) -> None: + """Load one policy replica and serve normalized canonical actions.""" + config = LingBotVlaV2ServiceConfig( + model_root=model_root, + qwen3vl_root=qwen3vl_root, + device=device, + max_image_bytes=max_image_mb * 1024 * 1024, + ) + app = create_lingbot_vla_v2_app(config) + uvicorn.run(app, host=host, port=port, workers=1) + + +if __name__ == "__main__": + main() diff --git a/telefuser/pipelines/lingbot_vla_v2/runtime.py b/telefuser/pipelines/lingbot_vla_v2/runtime.py new file mode 100644 index 00000000..059eed65 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/runtime.py @@ -0,0 +1,47 @@ +"""Runtime construction for single-replica LingBot-VLA v2 inference.""" + +from __future__ import annotations + +import torch +from transformers import AutoProcessor + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2 + +from .pipeline import LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig + + +def get_lingbot_vla_v2_pipeline( + model_root: str, + qwen3vl_root: str, + device: str = "cuda:0", +) -> LingBotVlaV2Pipeline: + """Load one official 6B base checkpoint replica for inference.""" + target_device = torch.device(device) + if target_device.type == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError(f"CUDA device {device!r} was requested, but CUDA is unavailable") + device_index = target_device.index or 0 + if device_index >= torch.cuda.device_count(): + raise ValueError( + f"CUDA device index {device_index} is unavailable; visible device count is {torch.cuda.device_count()}" + ) + target_device = torch.device("cuda", device_index) + dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 + processor = AutoProcessor.from_pretrained(qwen3vl_root, local_files_only=True, padding_side="right") + manager = ModuleManager(torch_dtype=dtype, device="cpu") + manager.add_module(processor, "lingbot_vla_v2_processor", path=qwen3vl_root) + load_lingbot_vla_v2(manager, model_root, qwen3vl_root, torch_dtype=dtype) + pipeline = LingBotVlaV2Pipeline(device=str(target_device), torch_dtype=dtype) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig( + device_type=target_device.type, + device_id=target_device.index or 0, + torch_dtype=dtype, + ), + ), + ) + return pipeline diff --git a/telefuser/pipelines/lingbot_vla_v2/service.py b/telefuser/pipelines/lingbot_vla_v2/service.py new file mode 100644 index 00000000..30c71e2e --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/service.py @@ -0,0 +1,196 @@ +"""Minimal single-GPU HTTP service for LingBot-VLA v2 action inference.""" + +from __future__ import annotations + +import base64 +import binascii +import io +import math +import threading +from collections.abc import Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Protocol + +from PIL import Image, UnidentifiedImageError +from fastapi import FastAPI, HTTPException +from fastapi.concurrency import run_in_threadpool +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from .data import LingBotVlaV2Observation +from .pipeline import LingBotVlaV2CanonicalActionChunk +from .robot_profile import ROBOTWIN_CAMERA_KEYS +from .runtime import get_lingbot_vla_v2_pipeline + + +@dataclass(frozen=True) +class LingBotVlaV2ServiceConfig: + """Configuration for one process-local LingBot-VLA v2 replica.""" + + model_root: str + qwen3vl_root: str + device: str = "cuda:0" + max_image_bytes: int = 10 * 1024 * 1024 + + def __post_init__(self) -> None: + if self.max_image_bytes <= 0: + raise ValueError("max_image_bytes must be positive") + + +class LingBotVlaV2ActionRequest(BaseModel): + """One RobotWin observation encoded for the HTTP boundary.""" + + model_config = ConfigDict(extra="forbid") + + task: str = Field(min_length=1) + state: list[float] = Field(min_length=14, max_length=14) + camera_high: str = Field(min_length=1) + camera_left_wrist: str = Field(min_length=1) + camera_right_wrist: str = Field(min_length=1) + seed: int | None = None + + @field_validator("task") + @classmethod + def validate_task(cls, value: str) -> str: + """Reject whitespace-only instructions.""" + value = value.strip() + if not value: + raise ValueError("task must be a non-empty string") + return value + + @field_validator("state") + @classmethod + def validate_state(cls, value: list[float]) -> list[float]: + """Reject non-finite robot state values.""" + if not all(math.isfinite(item) for item in value): + raise ValueError("state must contain only finite values") + return value + + +class LingBotVlaV2ActionResponse(BaseModel): + """Normalized canonical action chunk returned by the base checkpoint.""" + + canonical_normalized_actions: list[list[float]] + horizon: int + action_dim: int + checkpoint_variant: str + policy_verified: bool + verification_status: str + + +class LingBotVlaV2HealthResponse(BaseModel): + """Readiness state for the process-local model replica.""" + + status: str + model: str + device: str + policy_verified: bool + + +class _Pipeline(Protocol): + def __call__( + self, + observation: LingBotVlaV2Observation, + seed: int | None = None, + ) -> LingBotVlaV2CanonicalActionChunk: ... + + def close(self) -> None: ... + + +PipelineFactory = Callable[[LingBotVlaV2ServiceConfig], _Pipeline] + + +def _default_pipeline_factory(config: LingBotVlaV2ServiceConfig) -> _Pipeline: + return get_lingbot_vla_v2_pipeline(config.model_root, config.qwen3vl_root, device=config.device) + + +def _decode_image(value: str, *, max_image_bytes: int) -> Image.Image: + payload = value.strip() + if payload.startswith("data:"): + header, separator, payload = payload.partition(",") + if not separator or ";base64" not in header.lower(): + raise ValueError("image data URLs must use base64 encoding") + max_encoded_length = 4 * ((max_image_bytes + 2) // 3) + if len(payload) > max_encoded_length: + raise ValueError(f"decoded image must not exceed {max_image_bytes} bytes") + try: + decoded = base64.b64decode(payload, validate=True) + except (binascii.Error, ValueError) as error: + raise ValueError("image must be valid base64") from error + if not decoded or len(decoded) > max_image_bytes: + raise ValueError(f"decoded image must contain 1 to {max_image_bytes} bytes") + try: + with Image.open(io.BytesIO(decoded)) as image: + return image.convert("RGB").copy() + except (UnidentifiedImageError, OSError) as error: + raise ValueError("decoded payload must be a supported image") from error + + +class LingBotVlaV2Service: + """Serialize requests through one loaded policy replica.""" + + def __init__(self, pipeline: _Pipeline, config: LingBotVlaV2ServiceConfig) -> None: + self.pipeline = pipeline + self.config = config + self._inference_lock = threading.Lock() + + def predict(self, request: LingBotVlaV2ActionRequest) -> LingBotVlaV2ActionResponse: + """Decode one request and run it on the process-local replica.""" + encoded_images = (request.camera_high, request.camera_left_wrist, request.camera_right_wrist) + images = { + key: _decode_image(value, max_image_bytes=self.config.max_image_bytes) + for key, value in zip(ROBOTWIN_CAMERA_KEYS, encoded_images, strict=True) + } + observation = LingBotVlaV2Observation(task=request.task, state=request.state, images=images) + with self._inference_lock: + chunk = self.pipeline(observation, seed=request.seed) + return LingBotVlaV2ActionResponse( + canonical_normalized_actions=chunk.canonical_normalized_actions.tolist(), + horizon=chunk.horizon, + action_dim=chunk.action_dim, + checkpoint_variant=chunk.checkpoint_variant, + policy_verified=chunk.policy_verified, + verification_status=chunk.verification_status, + ) + + def close(self) -> None: + """Release model resources during application shutdown.""" + self.pipeline.close() + + +def create_lingbot_vla_v2_app( + config: LingBotVlaV2ServiceConfig, + *, + pipeline_factory: PipelineFactory = _default_pipeline_factory, +) -> FastAPI: + """Create a FastAPI application backed by exactly one policy replica.""" + + @asynccontextmanager + async def lifespan(app: FastAPI): + service = LingBotVlaV2Service(pipeline_factory(config), config) + app.state.lingbot_vla_v2_service = service + try: + yield + finally: + service.close() + + app = FastAPI(title="LingBot VLA v2", version="1", lifespan=lifespan) + + @app.get("/health", response_model=LingBotVlaV2HealthResponse) + async def health() -> LingBotVlaV2HealthResponse: + return LingBotVlaV2HealthResponse( + status="ready", + model="lingbot-vla-v2-6b-base", + device=config.device, + policy_verified=False, + ) + + @app.post("/v1/vla/actions", response_model=LingBotVlaV2ActionResponse) + async def predict(request: LingBotVlaV2ActionRequest) -> LingBotVlaV2ActionResponse: + service: LingBotVlaV2Service = app.state.lingbot_vla_v2_service + try: + return await run_in_threadpool(service.predict, request) + except (TypeError, ValueError) as error: + raise HTTPException(status_code=422, detail=str(error)) from error + + return app diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_service.py b/tests/unit/pipelines/lingbot_vla_v2/test_service.py new file mode 100644 index 00000000..a79c33d4 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_service.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import base64 +import io +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import torch +from PIL import Image +from fastapi.testclient import TestClient + +from telefuser.pipelines.lingbot_vla_v2.pipeline import LingBotVlaV2CanonicalActionChunk +from telefuser.pipelines.lingbot_vla_v2.robot_profile import ROBOTWIN_CAMERA_KEYS +from telefuser.pipelines.lingbot_vla_v2.service import ( + LingBotVlaV2ActionRequest, + LingBotVlaV2Service, + LingBotVlaV2ServiceConfig, + create_lingbot_vla_v2_app, +) + + +def _encoded_image(*, data_url: bool = False) -> str: + buffer = io.BytesIO() + Image.new("RGB", (8, 8), color=(10, 20, 30)).save(buffer, format="PNG") + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + return f"data:image/png;base64,{encoded}" if data_url else encoded + + +def _payload() -> dict: + image = _encoded_image() + return { + "task": "pick up the red block", + "state": [0.0] * 14, + "camera_high": image, + "camera_left_wrist": image, + "camera_right_wrist": image, + "seed": 7, + } + + +class _Pipeline: + def __init__(self, *, delay: float = 0.0) -> None: + self.delay = delay + self.closed = False + self.observations = [] + self.seeds = [] + self.active = 0 + self.max_active = 0 + self._counter_lock = threading.Lock() + + def __call__(self, observation, seed=None) -> LingBotVlaV2CanonicalActionChunk: + with self._counter_lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + try: + time.sleep(self.delay) + self.observations.append(observation) + self.seeds.append(seed) + return LingBotVlaV2CanonicalActionChunk( + canonical_normalized_actions=torch.zeros(2, 55), + horizon=2, + action_dim=55, + ) + finally: + with self._counter_lock: + self.active -= 1 + + def close(self) -> None: + self.closed = True + + +def _config(**kwargs) -> LingBotVlaV2ServiceConfig: + return LingBotVlaV2ServiceConfig( + model_root="/models/lingbot-vla-v2-6b", + qwen3vl_root="/models/Qwen3-VL-4B-Instruct", + **kwargs, + ) + + +def test_app_serves_health_and_normalized_action_contract() -> None: + pipeline = _Pipeline() + config = _config(device="cuda:3") + app = create_lingbot_vla_v2_app(config, pipeline_factory=lambda received: pipeline) + + with TestClient(app) as client: + health = client.get("/health") + response = client.post("/v1/vla/actions", json=_payload()) + + assert health.status_code == 200 + assert health.json() == { + "status": "ready", + "model": "lingbot-vla-v2-6b-base", + "device": "cuda:3", + "policy_verified": False, + } + assert response.status_code == 200 + body = response.json() + assert body["horizon"] == 2 + assert body["action_dim"] == 55 + assert body["checkpoint_variant"] == "base" + assert body["policy_verified"] is False + assert body["verification_status"] == "unverified_official_6b_base" + assert len(body["canonical_normalized_actions"]) == 2 + assert len(body["canonical_normalized_actions"][0]) == 55 + assert pipeline.seeds == [7] + assert tuple(pipeline.observations[0].images) == ROBOTWIN_CAMERA_KEYS + assert all(image.mode == "RGB" for image in pipeline.observations[0].images.values()) + assert pipeline.closed is True + + +def test_app_accepts_image_data_urls() -> None: + pipeline = _Pipeline() + payload = _payload() + payload["camera_high"] = _encoded_image(data_url=True) + app = create_lingbot_vla_v2_app(_config(), pipeline_factory=lambda received: pipeline) + + with TestClient(app) as client: + response = client.post("/v1/vla/actions", json=payload) + + assert response.status_code == 200 + + +def test_app_rejects_invalid_observations_without_running_policy() -> None: + pipeline = _Pipeline() + app = create_lingbot_vla_v2_app(_config(), pipeline_factory=lambda received: pipeline) + payload = _payload() + payload["camera_high"] = "not-base64" + + with TestClient(app) as client: + invalid_image = client.post("/v1/vla/actions", json=payload) + invalid_state = client.post("/v1/vla/actions", json={**_payload(), "state": [0.0] * 13}) + extra_field = client.post("/v1/vla/actions", json={**_payload(), "output_path": "/tmp/action"}) + + assert invalid_image.status_code == 422 + assert invalid_image.json()["detail"] == "image must be valid base64" + assert invalid_state.status_code == 422 + assert extra_field.status_code == 422 + assert pipeline.observations == [] + + +def test_service_serializes_policy_calls() -> None: + pipeline = _Pipeline(delay=0.02) + service = LingBotVlaV2Service(pipeline, _config()) + request = LingBotVlaV2ActionRequest.model_validate(_payload()) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(service.predict, request) for _ in range(2)] + responses = [future.result() for future in futures] + + assert [response.horizon for response in responses] == [2, 2] + assert pipeline.max_active == 1 + + +def test_service_config_rejects_non_positive_image_limit() -> None: + try: + _config(max_image_bytes=0) + except ValueError as error: + assert str(error) == "max_image_bytes must be positive" + else: + raise AssertionError("expected an invalid image size limit to be rejected") From 953ec3f2a11019f767c460bf438cfa4b84215bd7 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Wed, 5 Aug 2026 06:31:34 +0000 Subject: [PATCH 08/15] feat(vla): integrate native structured action serving Add a structured task contract and result path while preserving the existing media task API. Wire LingBot VLA v2 into the native scheduler, pipeline pool, status and metrics APIs, and expose action inference through TFClient. Add the native pipeline entrypoint, service documentation, and coverage for routing, validation, pool passthrough, client encoding, lifecycle cleanup, and JSON result handling. Verification: 168 focused service and VLA tests passed; Ruff check and format check passed; strict pipeline validation reported SAFE; real 6B checkpoint smoke returned a finite 50x55 action chunk. --- docs/en/service.md | 14 +- examples/lingbot_vla_v2/README.md | 59 +++++ .../lingbot_vla_v2_native_service.py | 114 +++++++++ telefuser/client/tf_client.py | 82 +++++- telefuser/entrypoints/cli/main.py | 2 +- telefuser/pipelines/lingbot_vla_v2/service.py | 40 +-- telefuser/service/api/__init__.py | 4 + telefuser/service/api/api_server.py | 5 +- telefuser/service/api/routers/tasks.py | 20 +- telefuser/service/api/schema.py | 24 ++ .../service/api/task_application_service.py | 30 ++- .../service/api/task_contract_runtime.py | 2 + telefuser/service/core/pipeline_contract.py | 5 + telefuser/service/core/pipeline_runner.py | 2 + telefuser/service/core/task_manager.py | 5 + telefuser/service/core/task_processor.py | 19 +- telefuser/service/core/task_service.py | 57 ++++- telefuser/service_types.py | 6 +- tests/unit/service/test_structured_tasks.py | 240 ++++++++++++++++++ 19 files changed, 702 insertions(+), 28 deletions(-) create mode 100644 examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py create mode 100644 tests/unit/service/test_structured_tasks.py diff --git a/docs/en/service.md b/docs/en/service.md index a0acc978..487e3a3d 100644 --- a/docs/en/service.md +++ b/docs/en/service.md @@ -156,7 +156,7 @@ telefuser serve /path/to/pipeline --task i2v [OPTIONS] | Parameter | Shortcut | Type | Default | Description | |-----------|----------|------|---------|-------------| | `pipe_path` | | string | **Required** | Positional path to the pipeline Python file | -| `--task` | `-t` | choice | `i2v` | Task type: t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr | +| `--task` | `-t` | choice | `i2v` | Task type: t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr, vla_action | | `--port` | `-p` | int | `8000` | Server port | | `--host` | | string | `127.0.0.1` | Server host address | | `--cache-dir` | `-c` | string | `work_dirs/server_cache` | Cache directory | @@ -262,6 +262,7 @@ telefuser serve --help | `i2i` | Image-to-Image: Generate image from input image and prompt | | `s2v` | Speech-to-Video: Generate video from speech | | `vsr` | Video Super-Resolution: Upscale an input video | +| `vla_action` | Structured VLA canonical action inference | ### Environment Variables @@ -1104,6 +1105,17 @@ telefuser serve ./pipeline.py --task t2v --- +### Structured JSON Tasks + +Pipelines that return JSON instead of image or video artifacts can declare a task contract with +`media_type="structured"`. Submit those tasks to `POST /v1/tasks/structured`; the existing status, cancellation, +queue, pool, and metrics endpoints remain unchanged. The pipeline entrypoint must return a finite JSON object. On +completion, `GET /v1/tasks/{task_id}/status` exposes that object under `result` together with +`inference_time_s` and the optional `peak_memory_mb`. + + +Media tasks continue to use `POST /v1/tasks/create` and artifact paths. The structured endpoint rejects media +contracts, and the media endpoint retains its existing request and response format. ## Client SDK ### Installation diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index b501da66..a11344fb 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -102,6 +102,65 @@ The response contains `canonical_normalized_actions`, `horizon`, `action_dim`, ` `policy_verified`, and `verification_status`. A successful HTTP response confirms service and model execution only; the normalized base-model output is not a physical robot command. +## Native TeleFuser Service + +The native service uses the shared `PIPELINE_CONTRACT`, asynchronous task scheduler, pipeline pool, status API, runtime +metrics, and `TFClient`. It keeps the standalone endpoint above as a small debugging path. + +The example resolves checkpoints under the existing `TF_MODEL_ZOO_PATH` layout: + +- `lingbot/lingbot-vla-v2-6b` +- `Qwen3-VL-4B-Instruct` + +Start one replica on one visible GPU: + +```bash +TF_MODEL_ZOO_PATH=/hhb-data/aigc/model_zoo \ + .venv-vla/bin/telefuser serve \ + examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py \ + --task vla_action \ + --parallelism 1 \ + --host 127.0.0.1 \ + --port 18080 +``` + +Submit `POST /v1/tasks/structured` with `task="vla_action"`, an `instruction`, the 14-dimensional `state`, and +the three Base64 camera fields. The creation response contains a task ID. Poll +`GET /v1/tasks/{task_id}/status`; a completed response contains the action payload under `result` and includes +`inference_time_s` and the optional `peak_memory_mb`. + +The unified client handles image encoding, submission, polling, and result extraction: + +```python +from telefuser.client import TFClient + +client = TFClient("http://127.0.0.1:18080") +actions = client.predict_vla_actions( + instruction="pick up the red block", + state=[0.0] * 14, + camera_high_path="/data/cam_high.png", + camera_left_wrist_path="/data/cam_left_wrist.png", + camera_right_wrist_path="/data/cam_right_wrist.png", + seed=7, +) +print(actions["horizon"], actions["action_dim"]) +``` + +For independent replicas, expose one GPU per replica through the existing pipeline pool: + +```bash +CUDA_VISIBLE_DEVICES=0,1 TF_MODEL_ZOO_PATH=/hhb-data/aigc/model_zoo \ + .venv-vla/bin/telefuser serve \ + examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py \ + --task vla_action \ + --parallelism 2 \ + --num-replicas 2 \ + --port 18080 +``` + +This is request-level replication, not tensor parallelism inside one policy replica. The response remains a normalized +base-model canonical action chunk and must not be treated as a physical robot command. + ## TeleFuser Regression Baseline The validation capture runs through the public loader and pipeline, then records preprocessing tensors, fixed initial diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py new file mode 100644 index 00000000..567dd456 --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py @@ -0,0 +1,114 @@ +"""Native TeleFuser service contract for LingBot-VLA v2 action inference.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from telefuser.pipelines.lingbot_vla_v2.pipeline import LingBotVlaV2Pipeline +from telefuser.pipelines.lingbot_vla_v2.runtime import get_lingbot_vla_v2_pipeline +from telefuser.pipelines.lingbot_vla_v2.service import ( + LingBotVlaV2ActionRequest, + predict_lingbot_vla_v2_action, +) + +TF_MODEL_ZOO_PATH = Path(os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo")).expanduser() + +PPL_CONFIG = { + "model_root": str(TF_MODEL_ZOO_PATH / "lingbot" / "lingbot-vla-v2-6b"), + "qwen3vl_root": str(TF_MODEL_ZOO_PATH / "Qwen3-VL-4B-Instruct"), + "device": "cuda:0", + "max_image_bytes": 10 * 1024 * 1024, +} + +PIPELINE_CONTRACT = { + "contract_version": "v1", + "pipeline_name": "lingbot_vla_v2_6b_base", + "supported_tasks": ["vla_action"], + "supported_media_types": ["structured"], + "execution_mode": "serial_single_pipeline", + "effective_max_concurrent_tasks": 1, + "entrypoints": { + "get_pipeline": "get_pipeline", + "run_with_file": "run_structured", + }, + "task_contracts": { + "vla_action": { + "media_type": "structured", + "required_inputs": ["camera_high", "camera_left_wrist", "camera_right_wrist"], + "optional_inputs": [], + "parameters": { + "instruction": { + "type": "string", + "required": True, + "description": "Robot instruction.", + }, + "state": { + "type": "array", + "required": True, + "description": "Raw 14-dimensional RobotWin state.", + }, + "camera_high": { + "type": "string", + "required": True, + "description": "Base64-encoded high camera image.", + }, + "camera_left_wrist": { + "type": "string", + "required": True, + "description": "Base64-encoded left wrist camera image.", + }, + "camera_right_wrist": { + "type": "string", + "required": True, + "description": "Base64-encoded right wrist camera image.", + }, + "seed": { + "type": "integer", + "required": False, + "default": None, + "description": "Optional deterministic inference seed.", + }, + }, + } + }, +} + + +def get_pipeline(parallelism: int = 1) -> LingBotVlaV2Pipeline: + """Load one policy replica for the native TeleFuser service.""" + if parallelism != 1: + raise ValueError("LingBot-VLA v2 supports parallelism=1 per replica; use --num-replicas for a pipeline pool") + return get_lingbot_vla_v2_pipeline( + PPL_CONFIG["model_root"], + PPL_CONFIG["qwen3vl_root"], + device=PPL_CONFIG["device"], + ) + + +def run_structured( + pipeline: LingBotVlaV2Pipeline, + instruction: str, + state: list[float], + camera_high: str, + camera_left_wrist: str, + camera_right_wrist: str, + seed: int | None = None, + **_: Any, +) -> dict[str, Any]: + """Return one JSON-serializable canonical normalized action chunk.""" + request = LingBotVlaV2ActionRequest( + task=instruction, + state=state, + camera_high=camera_high, + camera_left_wrist=camera_left_wrist, + camera_right_wrist=camera_right_wrist, + seed=seed, + ) + response = predict_lingbot_vla_v2_action( + pipeline, + request, + max_image_bytes=int(PPL_CONFIG["max_image_bytes"]), + ) + return response.model_dump(mode="json") diff --git a/telefuser/client/tf_client.py b/telefuser/client/tf_client.py index f71add4c..a44a06c2 100644 --- a/telefuser/client/tf_client.py +++ b/telefuser/client/tf_client.py @@ -95,10 +95,22 @@ TASK_S2V = "s2v" TASK_VSR = "vsr" TASK_EDIT = "edit" +TASK_VLA_ACTION = "vla_action" VIDEO_TASKS = (TASK_T2V, TASK_I2V, TASK_FL2V, TASK_VC, TASK_S2V, TASK_VSR) IMAGE_TASKS = (TASK_T2I, TASK_I2I, TASK_EDIT) -VALID_TASK_TYPES = (TASK_T2V, TASK_I2V, TASK_FL2V, TASK_VC, TASK_T2I, TASK_I2I, TASK_S2V, TASK_VSR, TASK_EDIT) +VALID_TASK_TYPES = ( + TASK_T2V, + TASK_I2V, + TASK_FL2V, + TASK_VC, + TASK_T2I, + TASK_I2I, + TASK_S2V, + TASK_VSR, + TASK_EDIT, + TASK_VLA_ACTION, +) # ── Constants: Aspect Ratios ──────────────────────────────────────────────── @@ -248,6 +260,74 @@ def create_task(self, task_type: str, **params: Any) -> Dict[str, Any]: except requests.RequestException as e: raise TaskCreationError(f"Task creation request failed: {e}") from e + def create_structured_task(self, task_type: str, **params: Any) -> Dict[str, Any]: + """Create a task whose pipeline contract returns a JSON result.""" + payload = {"task": task_type, **params} + try: + response = self._session.post( + f"{self.base_url}/v1/tasks/structured", + json=payload, + timeout=self.timeout, + ) + response.raise_for_status() + return response.json() + except requests.HTTPError as error: + raise TaskCreationError( + f"Structured task creation failed (HTTP {error.response.status_code}): {error.response.text}" + ) from error + except requests.RequestException as error: + raise TaskCreationError(f"Structured task creation request failed: {error}") from error + + def create_vla_action_task( + self, + instruction: str, + state: List[float], + camera_high_path: str, + camera_left_wrist_path: str, + camera_right_wrist_path: str, + seed: Optional[int] = None, + ) -> Dict[str, Any]: + """Create a LingBot-VLA v2 canonical action inference task.""" + return self.create_structured_task( + TASK_VLA_ACTION, + instruction=instruction, + state=state, + camera_high=self._encode_file_input(camera_high_path), + camera_left_wrist=self._encode_file_input(camera_left_wrist_path), + camera_right_wrist=self._encode_file_input(camera_right_wrist_path), + seed=seed, + ) + + def predict_vla_actions( + self, + instruction: str, + state: List[float], + camera_high_path: str, + camera_left_wrist_path: str, + camera_right_wrist_path: str, + seed: Optional[int] = None, + timeout: int = 300, + poll_interval: float = 0.5, + ) -> Dict[str, Any]: + """Run LingBot-VLA v2 inference and return its structured action result.""" + created = self.create_vla_action_task( + instruction=instruction, + state=state, + camera_high_path=camera_high_path, + camera_left_wrist_path=camera_left_wrist_path, + camera_right_wrist_path=camera_right_wrist_path, + seed=seed, + ) + status = self.wait_for_completion( + created["task_id"], + timeout=timeout, + poll_interval=poll_interval, + ) + result = status.get("result") + if not isinstance(result, dict): + raise TaskFailedError(f"Task {created['task_id']} completed without a structured result") + return result + # ── Video task creation methods ────────────────────────────────────────── def create_t2v_task( diff --git a/telefuser/entrypoints/cli/main.py b/telefuser/entrypoints/cli/main.py index c5d2adaf..c3e79ee1 100644 --- a/telefuser/entrypoints/cli/main.py +++ b/telefuser/entrypoints/cli/main.py @@ -31,7 +31,7 @@ def main(): "-t", default="i2v", type=click.Choice(TaskType.values(), case_sensitive=False), - help="Task type (t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr)", + help="Task type (t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr, vla_action)", ) @click.option("--port", "-p", default=8000, type=int, help="Server port") @click.option("--host", default="127.0.0.1", type=str, help="Server host") diff --git a/telefuser/pipelines/lingbot_vla_v2/service.py b/telefuser/pipelines/lingbot_vla_v2/service.py index 30c71e2e..e881b063 100644 --- a/telefuser/pipelines/lingbot_vla_v2/service.py +++ b/telefuser/pipelines/lingbot_vla_v2/service.py @@ -126,6 +126,30 @@ def _decode_image(value: str, *, max_image_bytes: int) -> Image.Image: raise ValueError("decoded payload must be a supported image") from error +def predict_lingbot_vla_v2_action( + pipeline: _Pipeline, + request: LingBotVlaV2ActionRequest, + *, + max_image_bytes: int, +) -> LingBotVlaV2ActionResponse: + """Decode one request and return the canonical normalized action chunk.""" + encoded_images = (request.camera_high, request.camera_left_wrist, request.camera_right_wrist) + images = { + key: _decode_image(value, max_image_bytes=max_image_bytes) + for key, value in zip(ROBOTWIN_CAMERA_KEYS, encoded_images, strict=True) + } + observation = LingBotVlaV2Observation(task=request.task, state=request.state, images=images) + chunk = pipeline(observation, seed=request.seed) + return LingBotVlaV2ActionResponse( + canonical_normalized_actions=chunk.canonical_normalized_actions.tolist(), + horizon=chunk.horizon, + action_dim=chunk.action_dim, + checkpoint_variant=chunk.checkpoint_variant, + policy_verified=chunk.policy_verified, + verification_status=chunk.verification_status, + ) + + class LingBotVlaV2Service: """Serialize requests through one loaded policy replica.""" @@ -136,22 +160,8 @@ def __init__(self, pipeline: _Pipeline, config: LingBotVlaV2ServiceConfig) -> No def predict(self, request: LingBotVlaV2ActionRequest) -> LingBotVlaV2ActionResponse: """Decode one request and run it on the process-local replica.""" - encoded_images = (request.camera_high, request.camera_left_wrist, request.camera_right_wrist) - images = { - key: _decode_image(value, max_image_bytes=self.config.max_image_bytes) - for key, value in zip(ROBOTWIN_CAMERA_KEYS, encoded_images, strict=True) - } - observation = LingBotVlaV2Observation(task=request.task, state=request.state, images=images) with self._inference_lock: - chunk = self.pipeline(observation, seed=request.seed) - return LingBotVlaV2ActionResponse( - canonical_normalized_actions=chunk.canonical_normalized_actions.tolist(), - horizon=chunk.horizon, - action_dim=chunk.action_dim, - checkpoint_variant=chunk.checkpoint_variant, - policy_verified=chunk.policy_verified, - verification_status=chunk.verification_status, - ) + return predict_lingbot_vla_v2_action(self.pipeline, request, max_image_bytes=self.config.max_image_bytes) def close(self) -> None: """Release model resources during application shutdown.""" diff --git a/telefuser/service/api/__init__.py b/telefuser/service/api/__init__.py index 3660471d..4da4d21a 100644 --- a/telefuser/service/api/__init__.py +++ b/telefuser/service/api/__init__.py @@ -17,6 +17,8 @@ OutputFormat, StopTaskResponse, StopTaskStatus, + StructuredTaskRequest, + StructuredTaskResponse, TaskRequest, TaskResponse, TaskStatus, @@ -29,6 +31,8 @@ "RateLimitMiddleware", "LoggingMiddleware", "setup_middleware", + "StructuredTaskRequest", + "StructuredTaskResponse", "TaskRequest", "TaskResponse", "StopTaskResponse", diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index f1398506..ea1ab886 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -18,7 +18,7 @@ from ..core.file_service import FileService from ..core.task_manager import TaskManager from ..core.task_processor import AsyncTaskProcessor -from ..core.task_service import MediaGenerationService +from ..core.task_service import MediaGenerationService, StructuredInferenceService from . import routers from .task_application_service import TaskApplicationService @@ -58,6 +58,7 @@ def __init__( self.file_service: FileService | None = None self.inference_service: PipelineService | None = None self.media_service: MediaGenerationService | None = None + self.structured_service: StructuredInferenceService | None = None self.task_app_service = TaskApplicationService(self) self.cache_service: Any | None = None self.max_queue_size = max_queue_size @@ -338,9 +339,11 @@ def initialize_services( cache_service=cache_service, cache_adapter=cache_adapter, ) + self.structured_service = StructuredInferenceService(inference_service) self.task_processor = AsyncTaskProcessor( task_manager=self.task_manager, media_service=self.media_service, + structured_service=self.structured_service, max_concurrent=self.max_concurrent_tasks, ) diff --git a/telefuser/service/api/routers/tasks.py b/telefuser/service/api/routers/tasks.py index 437264d7..e48beb62 100644 --- a/telefuser/service/api/routers/tasks.py +++ b/telefuser/service/api/routers/tasks.py @@ -19,7 +19,7 @@ from telefuser.service_types import MediaType, StopTaskStatus from telefuser.utils.logging import logger -from ..schema import StopTaskResponse, TaskRequest, TaskResponse +from ..schema import StopTaskResponse, StructuredTaskRequest, StructuredTaskResponse, TaskRequest, TaskResponse from ..task_contract_runtime import match_task_candidates if TYPE_CHECKING: @@ -40,6 +40,11 @@ async def create_task(message: TaskRequest) -> TaskResponse: """Create a new generation task.""" return await routes.create_task(message) + @new_router.post("/structured", response_model=StructuredTaskResponse) + async def create_structured_task(message: StructuredTaskRequest) -> StructuredTaskResponse: + """Create a task whose pipeline contract declares a structured result.""" + return await routes.create_structured_task(message) + @new_router.post("/form", response_model=TaskResponse) async def create_task_form( request: Request, @@ -99,6 +104,19 @@ async def check_image_path(image_name: str) -> None: logger.error(f"Failed to create task: {e}") raise HTTPException(status_code=500, detail=str(e)) + async def create_structured_task(self, message: StructuredTaskRequest) -> StructuredTaskResponse: + """Create a structured inference task without allocating an artifact path.""" + try: + return await self.api.task_app_service.submit_structured( + message, + explicit_fields=set(getattr(message, "model_fields_set", set())), + ) + except HTTPException: + raise + except Exception as error: + logger.error(f"Failed to create structured task: {error}") + raise HTTPException(status_code=500, detail=str(error)) from error + async def list_tasks(self) -> dict: """List all tasks.""" return self.api.task_manager.get_all_tasks() diff --git a/telefuser/service/api/schema.py b/telefuser/service/api/schema.py index fa23bbd7..e61fa25f 100644 --- a/telefuser/service/api/schema.py +++ b/telefuser/service/api/schema.py @@ -65,6 +65,23 @@ class TaskStatusMessage(BaseModel): task_id: str = Field(..., description="Task ID") +class StructuredTaskRequest(BaseModel): + """Request model for JSON-serializable inference results.""" + + model_config = ConfigDict(extra="allow") + + task_id: str = Field(default_factory=generate_task_id, description="Task ID (auto-generated)") + task: str = Field(..., description="Structured task type declared by the pipeline contract") + + @field_validator("task") + @classmethod + def validate_task(cls: type["StructuredTaskRequest"], value: str) -> str: + return validate_task_name_format(value) + + def get(self, key: str, default: Any = None) -> Any: + return getattr(self, key, default) + + class TaskResponse(BaseModel): """Response model for task creation.""" @@ -80,3 +97,10 @@ class StopTaskResponse(BaseModel): stop_status: StopTaskStatus reason: str + + +class StructuredTaskResponse(BaseModel): + """Response returned when a structured task is accepted.""" + + task_id: str + task_status: TaskStatus diff --git a/telefuser/service/api/task_application_service.py b/telefuser/service/api/task_application_service.py index 6d53c820..70b16d08 100644 --- a/telefuser/service/api/task_application_service.py +++ b/telefuser/service/api/task_application_service.py @@ -15,7 +15,7 @@ from telefuser.service.core.pipeline_contract import infer_media_type_for_task from telefuser.service_types import MediaType, TaskStatus -from .schema import TaskRequest, TaskResponse +from .schema import StructuredTaskRequest, StructuredTaskResponse, TaskRequest, TaskResponse from .task_contract_runtime import apply_task_contract_defaults, validate_required_task_parameters if TYPE_CHECKING: @@ -56,6 +56,34 @@ async def submit( except RuntimeError as exc: raise HTTPException(status_code=503, detail=str(exc)) + async def submit_structured( + self, + message: StructuredTaskRequest, + *, + explicit_fields: set[str], + ensure_processing: bool = True, + ) -> StructuredTaskResponse: + """Validate and enqueue a task whose result is returned as JSON.""" + try: + self.api.validate_task_supported(message.task) + contract = self.api.get_task_contract(message.task) + media_type = str((contract or {}).get("media_type") or infer_media_type_for_task(message.task)) + if media_type != MediaType.STRUCTURED.value: + raise HTTPException( + status_code=400, + detail=f"Task '{message.task}' does not declare a structured result contract", + ) + apply_task_contract_defaults(message, task_contract=contract, explicit_fields=explicit_fields) + validate_required_task_parameters(message, task_contract=contract) + + task_id = self.api.task_manager.create_task(message) + message.task_id = task_id + if ensure_processing: + await self.api.ensure_task_processor_running() + return StructuredTaskResponse(task_id=task_id, task_status=TaskStatus.PENDING) + except RuntimeError as exc: + raise HTTPException(status_code=503, detail=str(exc)) + def validate_output_path(self, message: TaskRequest) -> None: file_service = self.api.file_service if file_service is None: diff --git a/telefuser/service/api/task_contract_runtime.py b/telefuser/service/api/task_contract_runtime.py index 74e2a9c8..35fa8a42 100644 --- a/telefuser/service/api/task_contract_runtime.py +++ b/telefuser/service/api/task_contract_runtime.py @@ -132,6 +132,8 @@ def _build_default_output_path(message: Any) -> str: return "" media_type = infer_media_type_for_task(task) + if media_type == "structured": + return "" if media_type == "image": output_format = getattr(message, "output_format", "png") or "png" return f"{task_id}.{output_format}" diff --git a/telefuser/service/core/pipeline_contract.py b/telefuser/service/core/pipeline_contract.py index 747f2ffe..5a72de9f 100644 --- a/telefuser/service/core/pipeline_contract.py +++ b/telefuser/service/core/pipeline_contract.py @@ -10,6 +10,7 @@ VIDEO_TASKS = frozenset({"t2v", "i2v", "fl2v", "vc", "s2v", "vsr"}) IMAGE_TASKS = frozenset({"t2i", "i2i", "edit"}) +STRUCTURED_TASKS = frozenset({"vla_action"}) TASK_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]{1,31}$") @@ -257,6 +258,8 @@ def _derive_media_types_from_tasks(tasks: tuple[str, ...]) -> list[str]: media_types.append("video") if any(task in IMAGE_TASKS for task in tasks): media_types.append("image") + if any(task in STRUCTURED_TASKS for task in tasks): + media_types.append("structured") if not media_types: media_types.append("unknown") return media_types @@ -336,4 +339,6 @@ def infer_media_type_for_task(task: str) -> str: return "image" if is_video_task(task): return "video" + if task in STRUCTURED_TASKS: + return "structured" return "video" diff --git a/telefuser/service/core/pipeline_runner.py b/telefuser/service/core/pipeline_runner.py index dbc5dcda..09d5fddf 100644 --- a/telefuser/service/core/pipeline_runner.py +++ b/telefuser/service/core/pipeline_runner.py @@ -144,6 +144,8 @@ async def shutdown(self) -> None: await self._pipeline.astop() elif hasattr(self._pipeline, "stop"): await asyncio.to_thread(self._pipeline.stop) + elif hasattr(self._pipeline, "close"): + await asyncio.to_thread(self._pipeline.close) self._started = False diff --git a/telefuser/service/core/task_manager.py b/telefuser/service/core/task_manager.py index c04c1334..3258efe7 100644 --- a/telefuser/service/core/task_manager.py +++ b/telefuser/service/core/task_manager.py @@ -30,6 +30,7 @@ class TaskInfo: output_path: str | None = None peak_memory_mb: float | None = None inference_time_s: float | None = None + result: dict[str, Any] | None = None stop_event: threading.Event = field(default_factory=threading.Event) thread: threading.Thread | None = None @@ -127,6 +128,7 @@ def complete_task( *, peak_memory_mb: float | None = None, inference_time_s: float | None = None, + result: dict[str, Any] | None = None, ) -> None: """Mark task as completed with metrics.""" with self._lock: @@ -151,6 +153,7 @@ def complete_task( task.inference_time_s = inference_time_s if inference_time_s is not None else duration get_service_metrics().record_task_completed(duration) task.peak_memory_mb = peak_memory_mb + task.result = result def fail_task(self, task_id: str, error: str) -> None: """Mark task as failed with metrics.""" @@ -227,6 +230,8 @@ def get_task_status(self, task_id: str) -> dict[str, Any] | None: "peak_memory_mb": task.peak_memory_mb, "inference_time_s": task.inference_time_s, } + if task.result is not None: + status["result"] = task.result status.update(self._serialize_task_message(task.message)) return status diff --git a/telefuser/service/core/task_processor.py b/telefuser/service/core/task_processor.py index 16fe8ecf..5063144b 100644 --- a/telefuser/service/core/task_processor.py +++ b/telefuser/service/core/task_processor.py @@ -11,8 +11,9 @@ from telefuser.utils.logging import logger +from ..api.schema import StructuredTaskRequest from .task_manager import TaskManager, TaskStatus -from .task_service import MediaGenerationService +from .task_service import MediaGenerationService, StructuredInferenceService class AsyncTaskProcessor: @@ -27,10 +28,12 @@ def __init__( task_manager: TaskManager, media_service: MediaGenerationService, max_concurrent: int = 1, + structured_service: StructuredInferenceService | None = None, ) -> None: """Initialize the async task processor.""" self.task_manager = task_manager self.media_service = media_service + self.structured_service = structured_service self.max_concurrent = max_concurrent self._queue: asyncio.Queue = asyncio.Queue() @@ -137,9 +140,16 @@ async def _process_task(self, task_id: str) -> None: return try: - result = await self.media_service.generate_media_with_stop_event( - task_info.message, task_info.stop_event - ) + if isinstance(task_info.message, StructuredTaskRequest): + if self.structured_service is None: + raise RuntimeError("Structured inference service is not initialized") + result = await self.structured_service.execute_with_stop_event( + task_info.message, task_info.stop_event + ) + else: + result = await self.media_service.generate_media_with_stop_event( + task_info.message, task_info.stop_event + ) if result: self.task_manager.complete_task( @@ -147,6 +157,7 @@ async def _process_task(self, task_id: str) -> None: result.output_path, peak_memory_mb=result.peak_memory_mb, inference_time_s=result.inference_time_s, + result=getattr(result, "result", None), ) logger.info(f"Task {task_id} completed successfully") else: diff --git a/telefuser/service/core/task_service.py b/telefuser/service/core/task_service.py index 927f4cad..224a09d5 100644 --- a/telefuser/service/core/task_service.py +++ b/telefuser/service/core/task_service.py @@ -2,14 +2,16 @@ from __future__ import annotations +import json import threading +from dataclasses import dataclass from types import SimpleNamespace from typing import TYPE_CHECKING, Any from telefuser.service_types import MediaType, PipelineRunStatus, TaskStatus, TaskType from telefuser.utils.logging import logger -from ..api.schema import TaskRequest, TaskResponse +from ..api.schema import StructuredTaskRequest, TaskRequest, TaskResponse from ..media.media_base import AudioHandler, ImageHandler, VideoHandler from .file_service import FileService from .pipeline_contract import infer_media_type_for_task @@ -17,6 +19,59 @@ if TYPE_CHECKING: from .pipeline_service import PipelineService + +@dataclass(frozen=True) +class StructuredTaskExecutionResponse: + """Internal result returned to the shared task processor.""" + + task_id: str + result: dict[str, Any] + output_path: None = None + peak_memory_mb: float | None = None + inference_time_s: float | None = None + + +class StructuredInferenceService: + """Execute pipeline tasks that return JSON objects instead of artifacts.""" + + def __init__(self, inference_service: "PipelineService") -> None: + self.inference_service = inference_service + + async def execute_with_stop_event( + self, + message: StructuredTaskRequest, + stop_event: threading.Event, + ) -> StructuredTaskExecutionResponse | None: + """Run one structured task and retain its JSON result in task state.""" + if stop_event.is_set(): + logger.info(f"Task {message.task_id} cancelled before processing") + return None + + task_data = message.model_dump(mode="json") + result = await self.inference_service.run_task_with_stop_event(task_data, stop_event) + if result is None: + if stop_event.is_set(): + return None + raise RuntimeError("Task processing timeout") + if result.get("status") != PipelineRunStatus.SUCCESS: + raise RuntimeError(result.get("message") or "Inference failed") + + payload = result.get("raw") + if not isinstance(payload, dict): + raise RuntimeError("Structured pipeline entrypoint must return a JSON object") + try: + json.dumps(payload, allow_nan=False) + except (TypeError, ValueError) as error: + raise RuntimeError("Structured pipeline result must contain finite JSON-serializable values") from error + + return StructuredTaskExecutionResponse( + task_id=message.task_id, + result=payload, + peak_memory_mb=result.get("peak_memory_mb"), + inference_time_s=result.get("inference_time_s"), + ) + + # Media handlers _image_handler = ImageHandler() _video_handler = VideoHandler() diff --git a/telefuser/service_types.py b/telefuser/service_types.py index f2c98cc6..7dc13dbb 100644 --- a/telefuser/service_types.py +++ b/telefuser/service_types.py @@ -17,7 +17,7 @@ def values(cls) -> list[str]: class TaskType(_StringEnum): - """Supported media generation task types.""" + """Supported inference task types.""" T2V = "t2v" I2V = "i2v" @@ -27,6 +27,7 @@ class TaskType(_StringEnum): I2I = "i2i" S2V = "s2v" VSR = "vsr" + VLA_ACTION = "vla_action" class AspectRatio(_StringEnum): @@ -70,10 +71,11 @@ class StopTaskStatus(_StringEnum): class MediaType(_StringEnum): - """Generated media type.""" + """Inference result type.""" IMAGE = "image" VIDEO = "video" + STRUCTURED = "structured" class PipelineRunStatus(_StringEnum): diff --git a/tests/unit/service/test_structured_tasks.py b/tests/unit/service/test_structured_tasks.py new file mode 100644 index 00000000..41c7adb4 --- /dev/null +++ b/tests/unit/service/test_structured_tasks.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import asyncio +import base64 +import io +import threading +import time +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import torch +from PIL import Image +from fastapi.testclient import TestClient + +from examples.lingbot_vla_v2 import lingbot_vla_v2_native_service +from telefuser.client import TFClient, TaskFailedError +from telefuser.pipelines.lingbot_vla_v2.pipeline import LingBotVlaV2CanonicalActionChunk +from telefuser.service.api.api_server import ApiServer +from telefuser.service.api.schema import StructuredTaskRequest +from telefuser.service.core.task_manager import TaskManager +from telefuser.service.core.task_service import StructuredInferenceService + + +def _encoded_image() -> str: + buffer = io.BytesIO() + Image.new("RGB", (8, 8), color=(10, 20, 30)).save(buffer, format="PNG") + return base64.b64encode(buffer.getvalue()).decode("ascii") + + +def _payload() -> dict: + image = _encoded_image() + return { + "task": "vla_action", + "instruction": "pick up the red block", + "state": [0.0] * 14, + "camera_high": image, + "camera_left_wrist": image, + "camera_right_wrist": image, + "seed": 7, + } + + +class _StructuredPipelineService: + is_running = True + + def supported_tasks(self) -> tuple[str, ...]: + return ("vla_action",) + + def get_task_contract(self, task: str) -> dict: + assert task == "vla_action" + return lingbot_vla_v2_native_service.PIPELINE_CONTRACT["task_contracts"][task] + + async def run_task_with_stop_event(self, task_data, stop_event, **kwargs) -> dict: + assert task_data["instruction"] == "pick up the red block" + assert "output_path" not in task_data + return { + "status": "success", + "raw": { + "canonical_normalized_actions": [[0.0] * 55 for _ in range(2)], + "horizon": 2, + "action_dim": 55, + "checkpoint_variant": "base", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + }, + "peak_memory_mb": 128.0, + "inference_time_s": 0.25, + } + + +def test_structured_route_uses_scheduler_and_exposes_result_metrics(tmp_path: Path) -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.initialize_services(tmp_path, _StructuredPipelineService()) + + with TestClient(server.get_app()) as client: + created = client.post("/v1/tasks/structured", json=_payload()) + + assert created.status_code == 200 + created_body = created.json() + assert created_body["task_status"] == "pending" + assert "output_path" not in created_body + + deadline = time.monotonic() + 2.0 + while True: + status = client.get(f"/v1/tasks/{created_body['task_id']}/status") + assert status.status_code == 200 + body = status.json() + if body["status"] == "completed": + break + assert time.monotonic() < deadline + time.sleep(0.01) + + assert body["output_path"] is None + assert body["media_type"] == "structured" + assert body["peak_memory_mb"] == 128.0 + assert body["inference_time_s"] == 0.25 + assert body["result"]["horizon"] == 2 + assert len(body["result"]["canonical_normalized_actions"][0]) == 55 + assert "camera_high" not in body + + +def test_structured_route_rejects_media_contract(tmp_path: Path) -> None: + inference_service = _StructuredPipelineService() + inference_service.supported_tasks = lambda: ("t2i",) + inference_service.get_task_contract = lambda task: {"media_type": "image", "parameters": {}} + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.initialize_services(tmp_path, inference_service) + + with TestClient(server.get_app()) as client: + response = client.post("/v1/tasks/structured", json={"task": "t2i"}) + + assert response.status_code == 400 + assert "does not declare a structured result contract" in response.json()["detail"] + + +def test_structured_service_rejects_non_json_pipeline_result() -> None: + class InvalidService: + async def run_task_with_stop_event(self, task_data, stop_event): + return {"status": "success", "raw": {"value": torch.zeros(1)}} + + service = StructuredInferenceService(InvalidService()) + request = StructuredTaskRequest(task="vla_action") + + with pytest.raises(RuntimeError, match="finite JSON-serializable"): + asyncio.run(service.execute_with_stop_event(request, threading.Event())) + + +def test_native_vla_entrypoint_returns_action_contract() -> None: + class Pipeline: + def __call__(self, observation, seed=None): + assert observation.task == "pick up the red block" + assert seed == 7 + return LingBotVlaV2CanonicalActionChunk( + canonical_normalized_actions=torch.zeros(2, 55), + horizon=2, + action_dim=55, + ) + + result = lingbot_vla_v2_native_service.run_structured(Pipeline(), **_payload()) + + assert result["horizon"] == 2 + assert result["action_dim"] == 55 + assert len(result["canonical_normalized_actions"][0]) == 55 + + +def test_unified_client_encodes_vla_inputs_and_returns_result(tmp_path: Path) -> None: + image_path = tmp_path / "camera.png" + Image.new("RGB", (8, 8)).save(image_path) + client = TFClient("http://127.0.0.1:8000") + response = Mock() + response.raise_for_status.return_value = None + response.json.return_value = {"task_id": "task-1", "task_status": "pending"} + client._session.post = Mock(return_value=response) + client.wait_for_completion = Mock(return_value={"status": "completed", "result": {"horizon": 2}}) + + result = client.predict_vla_actions( + instruction="pick up the red block", + state=[0.0] * 14, + camera_high_path=str(image_path), + camera_left_wrist_path=str(image_path), + camera_right_wrist_path=str(image_path), + seed=7, + ) + + assert result == {"horizon": 2} + request = client._session.post.call_args + assert request.args[0].endswith("/v1/tasks/structured") + assert request.kwargs["json"]["task"] == "vla_action" + assert request.kwargs["json"]["camera_high"] == base64.b64encode(image_path.read_bytes()).decode("ascii") + + +def test_unified_client_rejects_missing_structured_result() -> None: + client = TFClient() + client.create_vla_action_task = Mock(return_value={"task_id": "task-1"}) + client.wait_for_completion = Mock(return_value={"status": "completed", "result": None}) + + with pytest.raises(TaskFailedError, match="without a structured result"): + client.predict_vla_actions( + instruction="pick", + state=[0.0] * 14, + camera_high_path="unused", + camera_left_wrist_path="unused", + camera_right_wrist_path="unused", + ) + + +def test_pipeline_pool_preserves_structured_result() -> None: + from telefuser.service.core.pipeline_pool import PipelinePool + + result = {"status": "success", "raw": {"horizon": 2}} + handle = SimpleNamespace( + _dead=False, + run_task=AsyncMock(return_value=result), + shutdown=Mock(), + ) + pool = PipelinePool( + num_replicas=1, + replica_device_ids=[["0"]], + security_level_name="NONE", + ) + pool._handles = [handle] + pool._instance_status = ["idle"] + pool._available.put_nowait(0) + + received = asyncio.run( + pool.run_task_with_stop_event( + {"task": "vla_action"}, + threading.Event(), + ) + ) + + assert received == result + assert pool._instance_status == ["idle"] + + +def test_pipeline_runner_closes_close_only_pipeline() -> None: + from telefuser.service.core.pipeline_runner import PipelineRunner + + class Pipeline: + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + pipeline = Pipeline() + + def run_structured(pipeline, **kwargs): + return {"value": 1} + + async def scenario() -> None: + runner = PipelineRunner(pipeline=pipeline, run_with_file=run_structured) + result = await runner.run(task_data={"task": "vla_action"}) + assert result.raw == {"value": 1} + await runner.shutdown() + + asyncio.run(scenario()) + assert pipeline.closed is True From 442106cf2c08793e52c8bbf268a99e8ae0ea1933 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Thu, 6 Aug 2026 01:55:15 +0000 Subject: [PATCH 09/15] perf(vla): optimize single-gpu serving latency Keep the LingBot VLA v2 policy resident on its target GPU, warm the fixed-shape inference path before service readiness, and avoid per-request allocator cache eviction. Add a reproducible single-GPU benchmark with latency distributions and memory metrics, and wake idle native-service workers immediately when tasks arrive. Verified with 178 focused VLA and shared-service tests, ruff checks, real H100 benchmarks, HTTP inference, and 38-layer strict repeatability. --- examples/lingbot_vla_v2/README.md | 25 ++ .../lingbot_vla_v2_native_service.py | 1 + .../pipelines/lingbot_vla_v2/pipeline.py | 32 +- telefuser/pipelines/lingbot_vla_v2/runtime.py | 5 + telefuser/pipelines/lingbot_vla_v2/service.py | 2 +- telefuser/service/api/api_server.py | 3 + telefuser/service/core/task_processor.py | 9 +- .../pipelines/lingbot_vla_v2/test_pipeline.py | 30 ++ tests/unit/service/test_service_smoke.py | 10 + tests/unit/service/test_task_runtime.py | 22 ++ .../test_lingbot_vla_v2_benchmark.py | 42 ++ .../benchmark_lingbot_vla_v2_service.py | 362 ++++++++++++++++++ 12 files changed, 536 insertions(+), 7 deletions(-) create mode 100644 tests/unit/validation/test_lingbot_vla_v2_benchmark.py create mode 100644 tools/validation/benchmark_lingbot_vla_v2_service.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index a11344fb..31d9a410 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -161,6 +161,31 @@ CUDA_VISIBLE_DEVICES=0,1 TF_MODEL_ZOO_PATH=/hhb-data/aigc/model_zoo \ This is request-level replication, not tensor parallelism inside one policy replica. The response remains a normalized base-model canonical action chunk and must not be treated as a physical robot command. +## Single-GPU Service Benchmark + +Use the VLA-specific benchmark to measure checkpoint construction, first-request latency, steady-state latency, +sequential throughput, process RSS, CUDA allocator peaks, and source-image-size overhead. The pipeline always converts +the three source images to the official `256x256` model input, so source size affects boundary and preprocessing cost, +not the model token shape. + +```bash +CUDA_VISIBLE_DEVICES=0 .venv-vla/bin/python \ + tools/validation/benchmark_lingbot_vla_v2_service.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --image examples/data/lingbot_world_fast/image.jpg \ + --image-sizes 256x256,640x480,1280x720 \ + --warmup 1 \ + --runs 20 \ + --output work_dirs/vla_service_benchmark/report.json +``` + +The native service moves the policy to its target GPU and runs one synthetic fixed-shape warmup before readiness. It +also keeps the allocator cache between requests. The report records construction and startup warmup separately, while +the first accepted request represents a ready replica. The default `service-thread` execution mode matches the native +service runner's fixed worker thread; use `--execution-mode direct` only to measure the in-process pipeline ceiling. +Shutdown still offloads the policy explicitly. + ## TeleFuser Regression Baseline The validation capture runs through the public loader and pipeline, then records preprocessing tensors, fixed initial diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py index 567dd456..deb6b3e2 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py @@ -84,6 +84,7 @@ def get_pipeline(parallelism: int = 1) -> LingBotVlaV2Pipeline: PPL_CONFIG["model_root"], PPL_CONFIG["qwen3vl_root"], device=PPL_CONFIG["device"], + warmup=True, ) diff --git a/telefuser/pipelines/lingbot_vla_v2/pipeline.py b/telefuser/pipelines/lingbot_vla_v2/pipeline.py index 4f941154..c7653075 100644 --- a/telefuser/pipelines/lingbot_vla_v2/pipeline.py +++ b/telefuser/pipelines/lingbot_vla_v2/pipeline.py @@ -12,7 +12,7 @@ from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation from .policy import LingBotVlaV2PolicyStage -from .robot_profile import RobotWinProfile +from .robot_profile import ROBOTWIN_CAMERA_KEYS, RobotWinProfile @dataclass @@ -40,6 +40,10 @@ class LingBotVlaV2CanonicalActionChunk: class LingBotVlaV2Pipeline(BasePipeline): """Single-replica LingBot-VLA v2 canonical action SDK.""" + # The service owns one fixed-shape resident policy. Per-request GC and + # allocator cache eviction add latency without releasing model weights. + clear_memory_after_call = False + def _get_stages(self) -> list: return [self.policy_stage] @@ -78,9 +82,7 @@ def predict( action_dim=int(canonical_actions.shape[1]), checkpoint_variant=str(getattr(policy_config, "checkpoint_variant", "base")), policy_verified=bool(getattr(policy_config, "policy_verified", False)), - verification_status=str( - getattr(policy_config, "verification_status", "unverified_official_6b_base") - ), + verification_status=str(getattr(policy_config, "verification_status", "unverified_official_6b_base")), ) @torch.inference_mode() @@ -92,7 +94,29 @@ def __call__( """Predict one normalized canonical action chunk.""" return self.predict(self.input_processor.prepare(observation), seed=seed) + def prepare_for_inference(self) -> None: + """Move the policy to its target device before the service becomes ready.""" + if not self.policy_stage.onload_models_flag: + self.policy_stage.onload_models() + self.policy_stage.onload_models_flag = True + + @torch.inference_mode() + def warmup(self) -> None: + """Initialize fixed-shape CUDA kernels before accepting service requests.""" + self.prepare_for_inference() + image_size = self.input_processor.image_size + image = torch.zeros(3, image_size, image_size, dtype=torch.uint8) + self( + LingBotVlaV2Observation( + task="warm up the policy", + state=[0.0] * 14, + images={key: image for key in ROBOTWIN_CAMERA_KEYS}, + ), + seed=0, + ) + def close(self) -> None: """Release policy device memory.""" if hasattr(self, "policy_stage"): self.policy_stage.offload_models() + self.policy_stage.onload_models_flag = False diff --git a/telefuser/pipelines/lingbot_vla_v2/runtime.py b/telefuser/pipelines/lingbot_vla_v2/runtime.py index 059eed65..6e3f93a2 100644 --- a/telefuser/pipelines/lingbot_vla_v2/runtime.py +++ b/telefuser/pipelines/lingbot_vla_v2/runtime.py @@ -16,6 +16,8 @@ def get_lingbot_vla_v2_pipeline( model_root: str, qwen3vl_root: str, device: str = "cuda:0", + *, + warmup: bool = False, ) -> LingBotVlaV2Pipeline: """Load one official 6B base checkpoint replica for inference.""" target_device = torch.device(device) @@ -44,4 +46,7 @@ def get_lingbot_vla_v2_pipeline( ), ), ) + pipeline.prepare_for_inference() + if warmup: + pipeline.warmup() return pipeline diff --git a/telefuser/pipelines/lingbot_vla_v2/service.py b/telefuser/pipelines/lingbot_vla_v2/service.py index e881b063..f0f78b38 100644 --- a/telefuser/pipelines/lingbot_vla_v2/service.py +++ b/telefuser/pipelines/lingbot_vla_v2/service.py @@ -101,7 +101,7 @@ def close(self) -> None: ... def _default_pipeline_factory(config: LingBotVlaV2ServiceConfig) -> _Pipeline: - return get_lingbot_vla_v2_pipeline(config.model_root, config.qwen3vl_root, device=config.device) + return get_lingbot_vla_v2_pipeline(config.model_root, config.qwen3vl_root, device=config.device, warmup=True) def _decode_image(value: str, *, max_image_bytes: int) -> Image.Image: diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index ea1ab886..76096e1f 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -189,6 +189,7 @@ async def ensure_task_processor_running(self) -> None: return if self.task_processor.is_running: + self.task_processor.notify_task_available() await self.ensure_artifact_cleanup_running() return @@ -197,8 +198,10 @@ async def ensure_task_processor_running(self) -> None: logger.warning("Task processor is not initialized; task will remain pending until services are ready") return if self.task_processor.is_running: + self.task_processor.notify_task_available() return await self.task_processor.start() + self.task_processor.notify_task_available() await self.ensure_artifact_cleanup_running() async def ensure_artifact_cleanup_running(self) -> None: diff --git a/telefuser/service/core/task_processor.py b/telefuser/service/core/task_processor.py index 5063144b..18577478 100644 --- a/telefuser/service/core/task_processor.py +++ b/telefuser/service/core/task_processor.py @@ -36,7 +36,7 @@ def __init__( self.structured_service = structured_service self.max_concurrent = max_concurrent - self._queue: asyncio.Queue = asyncio.Queue() + self._queue: asyncio.Queue[None] = asyncio.Queue() self._workers: list[asyncio.Task] = [] self._running = False self._stop_event = asyncio.Event() @@ -47,6 +47,10 @@ def is_running(self) -> bool: """Whether the processor workers are running.""" return self._running + def notify_task_available(self) -> None: + """Wake one idle worker after a task is added to the task manager.""" + self._queue.put_nowait(None) + async def start(self) -> None: """Start the task processor workers.""" if self._running: @@ -57,6 +61,7 @@ async def start(self) -> None: self._stop_event.clear() self._loop = asyncio.get_running_loop() + self._queue = asyncio.Queue() for i in range(self.max_concurrent): worker = asyncio.create_task(self._worker_loop(f"worker-{i}"), name=f"task-processor-{i}") self._workers.append(worker) @@ -101,7 +106,7 @@ async def _worker_loop(self, worker_name: str) -> None: task_id = self.task_manager.claim_next_pending_task() if task_id is None: - await asyncio.wait_for(self._stop_event.wait(), timeout=1.0) + await asyncio.wait_for(self._queue.get(), timeout=1.0) continue await self._process_task(task_id) diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py index e95e28f8..8889b11a 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py @@ -50,8 +50,10 @@ def __init__(self) -> None: policy_verified=False, verification_status="unverified_official_6b_base", ) + self.sample_count = 0 def sample_actions(self, **inputs) -> torch.Tensor: + self.sample_count += 1 assert inputs["state"].shape == (1, 55) assert inputs["images"].shape == (1, 3, 4, 6) return torch.zeros(1, self.config.n_action_steps, self.config.max_action_dim, device=self.anchor.device) @@ -88,3 +90,31 @@ def test_pipeline_returns_normalized_canonical_action_chunk() -> None: assert chunk.checkpoint_variant == "base" assert chunk.policy_verified is False assert chunk.verification_status == "unverified_official_6b_base" + + +def test_pipeline_prepares_resident_policy_and_disables_per_call_cache_eviction() -> None: + policy = _Policy() + processor = SimpleNamespace(image_processor=_ImageProcessor(), tokenizer=_Tokenizer()) + manager = ModuleManager(torch_dtype=torch.float32, device="cpu") + manager.add_module(policy, "lingbot_vla_v2") + manager.add_module(processor, "lingbot_vla_v2_processor") + pipeline = LingBotVlaV2Pipeline(device="cpu", torch_dtype=torch.float32) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + image_size=8, + ), + ) + + assert pipeline.clear_memory_after_call is False + assert pipeline.policy_stage.onload_models_flag is False + + pipeline.prepare_for_inference() + assert pipeline.policy_stage.onload_models_flag is True + + pipeline.warmup() + assert policy.sample_count == 1 + + pipeline.close() + assert pipeline.policy_stage.onload_models_flag is False diff --git a/tests/unit/service/test_service_smoke.py b/tests/unit/service/test_service_smoke.py index 18f635d6..661f545f 100644 --- a/tests/unit/service/test_service_smoke.py +++ b/tests/unit/service/test_service_smoke.py @@ -88,3 +88,13 @@ def test_openai_video_retrieve_includes_artifact_metadata_smoke(tmp_path: Path) assert data["artifact_id"] == f"local:tasks/{video_id}/outputs/videos/clip.mp4" assert data["artifact_metadata"]["backend"] == "local" assert data["artifact_metadata"]["size_bytes"] == 5 + + +def test_running_task_processor_is_notified_for_new_work(tmp_path: Path) -> None: + server = _make_smoke_server(tmp_path) + server.task_processor = Mock() + server.task_processor.is_running = True + + asyncio.run(server.ensure_task_processor_running()) + + server.task_processor.notify_task_available.assert_called_once_with() diff --git a/tests/unit/service/test_task_runtime.py b/tests/unit/service/test_task_runtime.py index dd0bda2c..5c181d08 100644 --- a/tests/unit/service/test_task_runtime.py +++ b/tests/unit/service/test_task_runtime.py @@ -114,6 +114,28 @@ async def wait_for_cancelled_status() -> None: asyncio.run(scenario()) +def test_async_task_processor_wakes_when_task_becomes_available() -> None: + """A newly submitted task should not wait for the idle polling timeout.""" + + async def scenario() -> None: + task_manager = TaskManager(max_queue_size=10) + media_service = _ControlledMediaService() + media_service.finish.set() + processor = AsyncTaskProcessor(task_manager=task_manager, media_service=media_service, max_concurrent=1) + + await processor.start() + try: + await asyncio.sleep(0) + task_manager.create_task(TaskRequest(task="t2i")) + processor.notify_task_available() + + await asyncio.wait_for(media_service.started.wait(), timeout=0.5) + finally: + await processor.stop() + + asyncio.run(scenario()) + + def test_claim_next_pending_task_atomic_single_winner() -> None: """Two PENDING tasks, single slot: only one is claimed, the second claim returns None.""" task_manager = TaskManager(max_queue_size=10) diff --git a/tests/unit/validation/test_lingbot_vla_v2_benchmark.py b/tests/unit/validation/test_lingbot_vla_v2_benchmark.py new file mode 100644 index 00000000..1a42b45f --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_benchmark.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import argparse + +import pytest +from PIL import Image + +from tools.validation.benchmark_lingbot_vla_v2_service import ( + encode_image, + parse_image_sizes, + percentile, + summarize, +) + + +def test_parse_image_sizes_deduplicates_and_preserves_order() -> None: + assert parse_image_sizes("256x256, 640X480,256x256") == ((256, 256), (640, 480)) + + +@pytest.mark.parametrize("value", ["", "256", "0x256", "axb"]) +def test_parse_image_sizes_rejects_invalid_values(value: str) -> None: + with pytest.raises(argparse.ArgumentTypeError): + parse_image_sizes(value) + + +def test_latency_summary_reports_interpolated_percentiles_and_throughput() -> None: + values = [1.0, 2.0, 3.0, 4.0] + result = summarize(values) + + assert percentile(values, 0.5) == 2.5 + assert result["count"] == 4 + assert result["mean_seconds"] == 2.5 + assert result["p95_seconds"] == pytest.approx(3.85) + assert result["throughput_requests_per_second"] == 0.4 + + +def test_encode_image_reports_decoded_jpeg_size() -> None: + encoded, encoded_bytes = encode_image(Image.new("RGB", (8, 8)), (32, 24), quality=90) + + assert encoded + assert encoded_bytes > 0 + assert len(encoded) >= encoded_bytes diff --git a/tools/validation/benchmark_lingbot_vla_v2_service.py b/tools/validation/benchmark_lingbot_vla_v2_service.py new file mode 100644 index 00000000..7db06807 --- /dev/null +++ b/tools/validation/benchmark_lingbot_vla_v2_service.py @@ -0,0 +1,362 @@ +"""Benchmark one in-process LingBot-VLA v2 service replica on a single GPU.""" + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import math +import statistics +import threading +import time +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any + +import psutil +import torch +from PIL import Image + +from telefuser.metrics.runtime import collect_runtime_environment +from telefuser.pipelines.lingbot_vla_v2.runtime import get_lingbot_vla_v2_pipeline +from telefuser.pipelines.lingbot_vla_v2.service import ( + LingBotVlaV2ActionRequest, + predict_lingbot_vla_v2_action, +) + +_MIB = 1024**2 + + +class PeakRssSampler: + """Sample process RSS while one benchmark phase is active.""" + + def __init__(self, process: psutil.Process, interval_s: float = 0.01) -> None: + self.process = process + self.interval_s = interval_s + self.peak_bytes = process.memory_info().rss + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + def __enter__(self) -> "PeakRssSampler": + self._thread = threading.Thread(target=self._sample, daemon=True) + self._thread.start() + return self + + def __exit__(self, *args: object) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=max(self.interval_s * 4, 0.1)) + self._record() + + def _record(self) -> None: + try: + self.peak_bytes = max(self.peak_bytes, self.process.memory_info().rss) + except psutil.Error: + pass + + def _sample(self) -> None: + while not self._stop.wait(self.interval_s): + self._record() + + +def parse_image_sizes(value: str) -> tuple[tuple[int, int], ...]: + """Parse a comma-separated WIDTHxHEIGHT list.""" + sizes: list[tuple[int, int]] = [] + for item in value.split(","): + parts = item.strip().lower().split("x", maxsplit=1) + if len(parts) != 2: + raise argparse.ArgumentTypeError(f"invalid image size {item!r}; expected WIDTHxHEIGHT") + try: + width, height = (int(part) for part in parts) + except ValueError as error: + raise argparse.ArgumentTypeError(f"invalid image size {item!r}; expected integers") from error + if width <= 0 or height <= 0: + raise argparse.ArgumentTypeError("image dimensions must be positive") + size = (width, height) + if size not in sizes: + sizes.append(size) + if not sizes: + raise argparse.ArgumentTypeError("at least one image size is required") + return tuple(sizes) + + +def percentile(values: Sequence[float], fraction: float) -> float: + """Return a linearly interpolated percentile for a non-empty sample.""" + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def summarize(values: Sequence[float]) -> dict[str, float | int]: + """Summarize a latency sample in seconds.""" + if not values: + raise ValueError("summary requires at least one value") + total = sum(values) + return { + "count": len(values), + "total_seconds": total, + "mean_seconds": statistics.fmean(values), + "stdev_seconds": statistics.pstdev(values), + "min_seconds": min(values), + "p50_seconds": percentile(values, 0.50), + "p90_seconds": percentile(values, 0.90), + "p95_seconds": percentile(values, 0.95), + "max_seconds": max(values), + "throughput_requests_per_second": len(values) / total, + } + + +def encode_image(source: Image.Image, size: tuple[int, int], *, quality: int) -> tuple[str, int]: + """Resize and JPEG-encode one service input outside measured request time.""" + image = source.resize(size, Image.Resampling.BICUBIC) + buffer = io.BytesIO() + image.save(buffer, format="JPEG", quality=quality, optimize=False) + payload = buffer.getvalue() + return base64.b64encode(payload).decode("ascii"), len(payload) + + +def _cuda_synchronize(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def _memory_snapshot(device: torch.device, process: psutil.Process) -> dict[str, float | None]: + result: dict[str, float | None] = {"cpu_rss_mib": process.memory_info().rss / _MIB} + if device.type != "cuda": + result.update( + gpu_allocated_mib=None, + gpu_reserved_mib=None, + gpu_peak_allocated_mib=None, + gpu_peak_reserved_mib=None, + ) + return result + result.update( + gpu_allocated_mib=torch.cuda.memory_allocated(device) / _MIB, + gpu_reserved_mib=torch.cuda.memory_reserved(device) / _MIB, + gpu_peak_allocated_mib=torch.cuda.max_memory_allocated(device) / _MIB, + gpu_peak_reserved_mib=torch.cuda.max_memory_reserved(device) / _MIB, + ) + return result + + +def measure( + operation: Callable[[], Any], + *, + device: torch.device, + process: psutil.Process, + synchronize_cuda: bool, +) -> tuple[Any, dict[str, Any]]: + """Measure wall time and process/device memory for one operation.""" + if synchronize_cuda: + _cuda_synchronize(device) + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + before = _memory_snapshot(device, process) + with PeakRssSampler(process) as rss_sampler: + started_at = time.perf_counter() + result = operation() + if synchronize_cuda: + _cuda_synchronize(device) + elapsed = time.perf_counter() - started_at + after = _memory_snapshot(device, process) + return result, { + "seconds": elapsed, + "cpu_rss_before_mib": before["cpu_rss_mib"], + "cpu_rss_after_mib": after["cpu_rss_mib"], + "cpu_rss_peak_mib": rss_sampler.peak_bytes / _MIB, + "gpu_allocated_after_mib": after["gpu_allocated_mib"], + "gpu_reserved_after_mib": after["gpu_reserved_mib"], + "gpu_peak_allocated_mib": after["gpu_peak_allocated_mib"], + "gpu_peak_reserved_mib": after["gpu_peak_reserved_mib"], + } + + +def _load_source_image(path: Path | None) -> Image.Image: + if path is not None: + with Image.open(path) as image: + return image.convert("RGB").copy() + return Image.new("RGB", (640, 480), color=(32, 96, 160)) + + +def run_benchmark(args: argparse.Namespace) -> dict[str, Any]: + """Load one replica and return a JSON-serializable benchmark report.""" + if args.warmup < 0 or args.runs < 1: + raise ValueError("--warmup must be non-negative and --runs must be positive") + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("LingBot VLA v2 service benchmarking requires one visible CUDA GPU") + process = psutil.Process() + source = _load_source_image(args.image) + encoded_by_size = {size: encode_image(source, size, quality=args.jpeg_quality) for size in args.image_sizes} + + pipeline, load_metrics = measure( + lambda: get_lingbot_vla_v2_pipeline( + str(args.model_root), + str(args.qwen3vl_root), + device=str(device), + ), + device=device, + process=process, + synchronize_cuda=False, + ) + startup_warmup = None + if args.startup_warmup: + _, startup_warmup = measure( + pipeline.warmup, + device=device, + process=process, + synchronize_cuda=True, + ) + + executor = ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="lingbot-vla-v2-benchmark") + if args.execution_mode == "service-thread" + else None + ) + active_phases: dict[str, float] = {} + original_prepare = pipeline.input_processor.prepare + original_predict = pipeline.predict + + def measured_prepare(observation: Any) -> Any: + started_at = time.perf_counter() + result = original_prepare(observation) + active_phases["preprocess_seconds"] = time.perf_counter() - started_at + return result + + def measured_predict(inputs: Any, seed: int | None = None) -> Any: + _cuda_synchronize(device) + started_at = time.perf_counter() + result = original_predict(inputs, seed=seed) + _cuda_synchronize(device) + active_phases["model_seconds"] = time.perf_counter() - started_at + return result + + pipeline.input_processor.prepare = measured_prepare + pipeline.predict = measured_predict + + def request_once(size: tuple[int, int]) -> tuple[dict[str, Any], dict[str, float]]: + active_phases.clear() + encoded, _ = encoded_by_size[size] + payload = { + "task": args.instruction, + "state": [0.0] * 14, + "camera_high": encoded, + "camera_left_wrist": encoded, + "camera_right_wrist": encoded, + "seed": args.seed, + } + request = LingBotVlaV2ActionRequest.model_validate(payload) + + def invoke_request() -> Any: + return predict_lingbot_vla_v2_action( + pipeline, + request, + max_image_bytes=args.max_image_bytes, + ) + + def invoke_service_thread() -> Any: + assert executor is not None + return executor.submit(invoke_request).result() + + operation: Callable[[], Any] = invoke_service_thread if executor is not None else invoke_request + response, metrics = measure( + operation, + device=device, + process=process, + synchronize_cuda=True, + ) + if response.horizon != 50 or response.action_dim != 55: + raise RuntimeError(f"unexpected action shape: {response.horizon}x{response.action_dim}") + if not all(math.isfinite(value) for row in response.canonical_normalized_actions for value in row): + raise RuntimeError("benchmark received non-finite actions") + phases = dict(active_phases) + phases["boundary_seconds"] = max( + metrics["seconds"] - phases.get("preprocess_seconds", 0.0) - phases.get("model_seconds", 0.0), + 0.0, + ) + return metrics, phases + + first_size = args.image_sizes[0] + try: + first_request, first_phases = request_once(first_size) + sizes_report: dict[str, Any] = {} + for size in args.image_sizes: + for _ in range(args.warmup): + request_once(size) + samples: list[dict[str, Any]] = [] + phases: list[dict[str, float]] = [] + for _ in range(args.runs): + sample, phase = request_once(size) + samples.append(sample) + phases.append(phase) + encoded_bytes = encoded_by_size[size][1] + sizes_report[f"{size[0]}x{size[1]}"] = { + "source_image_size": list(size), + "encoded_bytes_per_camera": encoded_bytes, + "total_latency": summarize([sample["seconds"] for sample in samples]), + "preprocess_latency": summarize([phase["preprocess_seconds"] for phase in phases]), + "model_latency": summarize([phase["model_seconds"] for phase in phases]), + "boundary_latency": summarize([phase["boundary_seconds"] for phase in phases]), + "cpu_rss_peak_mib": max(sample["cpu_rss_peak_mib"] for sample in samples), + "gpu_peak_allocated_mib": max(sample["gpu_peak_allocated_mib"] for sample in samples), + "gpu_peak_reserved_mib": max(sample["gpu_peak_reserved_mib"] for sample in samples), + } + report = { + "schema_version": 1, + "benchmark": "lingbot_vla_v2_single_gpu_service", + "model_root": str(args.model_root.resolve()), + "qwen3vl_root": str(args.qwen3vl_root.resolve()), + "device": str(device), + "seed": args.seed, + "instruction": args.instruction, + "internal_model_image_size": [pipeline.input_processor.image_size] * 2, + "warmup_runs_per_size": args.warmup, + "execution_mode": args.execution_mode, + "measured_runs_per_size": args.runs, + "environment": collect_runtime_environment([device], repo_root=Path(__file__).resolve().parents[2]), + "load": load_metrics, + "startup_warmup": startup_warmup, + "first_request": {**first_request, "source_image_size": list(first_size), "phases": first_phases}, + "steady_state_by_source_size": sizes_report, + "memory_after_benchmark": _memory_snapshot(device, process), + } + finally: + if executor is not None: + executor.shutdown(wait=True) + pipeline.close() + return report + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", required=True, type=Path) + parser.add_argument("--qwen3vl-root", required=True, type=Path) + parser.add_argument("--image", type=Path, help="Optional source image reused for all three camera inputs.") + parser.add_argument("--image-sizes", type=parse_image_sizes, default=parse_image_sizes("256x256,640x480,1280x720")) + parser.add_argument("--instruction", default="pick up the red block") + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--execution-mode", choices=("service-thread", "direct"), default="service-thread") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--runs", type=int, default=20) + parser.add_argument("--startup-warmup", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--jpeg-quality", type=int, choices=range(1, 101), default=95) + parser.add_argument("--max-image-bytes", type=int, default=10 * 1024 * 1024) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + report = run_benchmark(args) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() From 8ce0ea4c1a18f9d818e07f13ffc8c86c9f1623e1 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Thu, 6 Aug 2026 02:05:28 +0000 Subject: [PATCH 10/15] style(vla): clean imported model whitespace Remove trailing whitespace and the extra final blank line from the LingBot VLA v2 model sources so branch-level Git whitespace checks pass after syncing with main. No executable logic is changed. Verification: - 187 service and LingBot VLA v2 tests passed - 3 LingBot VLA v2 model tests passed - focused Ruff check and format check passed - git diff --check passed --- telefuser/models/lingbot_vla_v2_loader.py | 8 ++++---- telefuser/models/lingbot_vla_v2_moe.py | 15 +++++++-------- telefuser/models/lingbot_vla_v2_qwen.py | 6 +++--- 3 files changed, 14 insertions(+), 15 deletions(-) diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py index 95b0c160..7d72c0bd 100644 --- a/telefuser/models/lingbot_vla_v2_loader.py +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -871,7 +871,7 @@ def forward(self, x,pooled_text_embeds=None): class ResamplerXLIdentity(nn.Module): def __init__(self) -> None: super().__init__() - + def forward(self, x, pooled_text_embeds=None): return x, pooled_text_embeds @@ -911,13 +911,13 @@ def build_expand_mlp(in_hidden_size, hidden_size, out_size): class DepthHead(nn.Module): def __init__( - self, + self, proj_config=None, llm_hidden_size=4096, use_intermediate_depth=False, ): super(DepthHead, self).__init__() - + self.projector = Resampler( dim_in=llm_hidden_size, dim_mid=llm_hidden_size, @@ -935,7 +935,7 @@ def forward(self, llm_feats): class TaskTokenDepthHead(nn.Module): def __init__( - self, + self, proj_config=None, llm_hidden_size=4096, use_intermediate_depth=False, diff --git a/telefuser/models/lingbot_vla_v2_moe.py b/telefuser/models/lingbot_vla_v2_moe.py index a11e4024..ec3588d2 100644 --- a/telefuser/models/lingbot_vla_v2_moe.py +++ b/telefuser/models/lingbot_vla_v2_moe.py @@ -152,12 +152,12 @@ def forward(self, x): class Qwen2FusedExperts(nn.Module): """Fused expert module: stores E experts' weights as 3D tensors for group_gemm. - + Shape convention matches nn.Linear(in, out).weight = [out, in]: gate_proj: [E, intermediate_size, hidden_size] up_proj: [E, intermediate_size, hidden_size] down_proj: [E, hidden_size, intermediate_size] - + The forward() method runs the full fused_moe computation. This is critical for FSDP2: calling self.experts(...) triggers FSDP2's forward pre-hook to unshard the expert params on ep_fsdp_mesh BEFORE they are used by kernels. @@ -219,7 +219,7 @@ def _get_robby_moe_workspace(self, hidden_states, top_k): def forward(self, module, num_experts, routing_weights, selected_experts, hidden_states): """Run fused_moe_forward with FSDP2-managed weights. - + Must be called via self.experts(...) so FSDP2 unshards params first. """ return fused_moe_forward( @@ -411,7 +411,7 @@ def __init__(self, config: Qwen2Config, layer_idx: int): self.mlp = Qwen2MLP(config) self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": logger.warning_once( f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " @@ -450,7 +450,7 @@ def forward( value_state = self.self_attn.v_proj(hidden_states).view(hidden_shape) return query_state, key_state, value_state - + elif output_atten: if att_output.dtype != self.self_attn.o_proj.weight.dtype: att_output = att_output.to(self.self_attn.o_proj.weight.dtype) @@ -476,7 +476,7 @@ def forward( else: raise ValueError(f"Invaild Operation compute_kqv={compute_kqv} and output_atten={output_atten} with Qwen2DecoderLayer in LingBot-VLA") - + @auto_docstring class Qwen2PreTrainedModel(PreTrainedModel): config: Qwen2Config @@ -494,7 +494,7 @@ class Qwen2PreTrainedModel(PreTrainedModel): "hidden_states": Qwen2DecoderLayer, "attentions": Qwen2Attention, } - + def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): @@ -565,4 +565,3 @@ def apply_lingbot_qwen2_patch(): hf_qwen2.Qwen2PreTrainedModel = Qwen2PreTrainedModel hf_qwen2.Qwen2Model = Qwen2Model hf_qwen2.Qwen2ForCausalLM = Qwen2ForCausalLM - diff --git a/telefuser/models/lingbot_vla_v2_qwen.py b/telefuser/models/lingbot_vla_v2_qwen.py index 5ba74ccc..d2f0ec97 100644 --- a/telefuser/models/lingbot_vla_v2_qwen.py +++ b/telefuser/models/lingbot_vla_v2_qwen.py @@ -40,7 +40,7 @@ from flash_attn.layers.rotary import apply_rotary_emb from flash_attn.flash_attn_interface import flash_attn_varlen_func from transformers.modeling_flash_attention_utils import _flash_attention_forward -import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as hf_qwen25vl +import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as hf_qwen25vl from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2RMSNorm, Qwen2_5_VLMLP, @@ -265,7 +265,7 @@ def forward( key_state = self.k_layernorm(key_state) return query_state, key_state, value_state - + elif output_atten: if att_output.dtype != self.self_attn.o_proj.weight.dtype: att_output = att_output.to(self.self_attn.o_proj.weight.dtype) @@ -291,7 +291,7 @@ class Qwen2_5_VLTextModel(Qwen2_5_VLPreTrainedModel): get_input_embeddings = _Qwen2_5_VLTextModel.get_input_embeddings set_input_embeddings = _Qwen2_5_VLTextModel.set_input_embeddings forward = _Qwen2_5_VLTextModel.forward - + def __init__(self, config: Qwen2_5_VLConfig): super().__init__(config) self.padding_idx = config.pad_token_id From 2d40ee21d80fbc6f62ec9667f54827808aed7372 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Thu, 6 Aug 2026 02:32:33 +0000 Subject: [PATCH 11/15] refactor(vla): prune legacy LingBot model code Remove the unused V1 Qwen2.5-VL policy, legacy flow-matching implementation, dead resampler variants, demo code, and dormant global monkey-patch entry points. Keep the V2 inference path explicit through a compact FlowMatchingBase, retain checkpoint-compatible Qwen3-VL, action expert, alignment, and MoE modules, and preserve the models-to-ops-to-kernel boundary. Verification: - Ruff check and format check passed for all four model files - Python compilation passed - 190 service and LingBot VLA v2 tests passed - real 6B tensor replay produced 50x55 actions - strict 38-layer regression parity passed with max_abs=0 - git diff --check passed --- telefuser/models/lingbot_vla_v2.py | 1126 +++------------------ telefuser/models/lingbot_vla_v2_loader.py | 382 +------ telefuser/models/lingbot_vla_v2_moe.py | 117 +-- telefuser/models/lingbot_vla_v2_qwen.py | 442 +------- 4 files changed, 265 insertions(+), 1802 deletions(-) diff --git a/telefuser/models/lingbot_vla_v2.py b/telefuser/models/lingbot_vla_v2.py index cf78cde1..653b0d95 100644 --- a/telefuser/models/lingbot_vla_v2.py +++ b/telefuser/models/lingbot_vla_v2.py @@ -17,11 +17,53 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any, Dict, List, Literal, Optional, Union -from copy import deepcopy -from typing import Any, Dict, Literal, Optional +import einops +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from transformers import AutoConfig, AutoTokenizer, PreTrainedModel, PretrainedConfig +from transformers.cache_utils import Cache +from transformers.modeling_flash_attention_utils import is_flash_attn_available +from transformers.models.auto import CONFIG_MAPPING +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm +from transformers.models.qwen3_vl.modeling_qwen3_vl import apply_rotary_pos_emb +from transformers.utils import logging + +from telefuser.models.lingbot_vla_v2_loader import ( + LingBotVLAWeightLoader, + LingBotVlaV2StateDictConverter, + TaskTokenDepthHead, + block_suffix_to_fv_, + build_block_mask, + create_sinusoidal_pos_embedding, + flex_attention_forward, + flex_attention_with_block_mask, + make_att_2d_masks, + our_eager_attention_forward, + prefix_query_segments, + prefix_query_token_spans, +) +from telefuser.models.lingbot_vla_v2_moe import ( + FixQwen2RMSNorm, + Qwen2ForCausalLM, + Qwen2FusedExperts, + Qwen2TokenMoeBlock, +) +from telefuser.models.lingbot_vla_v2_qwen import ( + Qwen3VLForConditionalGeneration, + Qwen3VLPreTrainedModel, + Qwen3VLTextModel, +) + +try: + from dinov3.hub.backbones import dinov3_vitb16 +except Exception: + dinov3_vitb16 = None + +logger = logging.get_logger(__name__) -from transformers import AutoConfig, PretrainedConfig class LingbotVLAConfig(PretrainedConfig): """Configuration class for Lingbot-VLA. @@ -36,22 +78,18 @@ def __init__( vlm_repo_id: Optional[str] = None, expert_vision_path: Optional[str] = None, tokenizer_path: Optional[str] = None, - post_training: bool = False, adanorm_time: bool = False, split_gate_liner: bool = False, nosplit_gate_liner: bool = False, separate_time_proj: bool = False, final_norm_adanorm: bool = False, - enable_expert_vision: bool = False, expert_vision_type: Optional[str] = None, freeze_vision_encoder: bool = False, - incremental_training: bool = False, depth_incremental_training: bool = False, reinit_mismatched_weights: bool = False, - action_dim: int = 14, max_action_dim: int = 14, max_state_dim: int = 14, @@ -62,7 +100,6 @@ def __init__( norm_qkv: bool = False, align_params: Optional[Dict[str, Any]] = None, use_compile: bool = False, - use_moe: bool = False, token_moe_layers: Optional[list] = None, token_num_experts: int = 32, @@ -90,16 +127,13 @@ def __init__( qwen3vl_use_vision_boundaries: bool = False, precompute_grid_thw: bool = False, use_qwen3_fixed_grid_cache: bool = False, - use_lm_head: bool = False, vocab_size: int = 0, vit_attn_implementation: str = "flash_attention_2", attention_implementation: str = "flex", - train_expert_only: bool = False, train_state_proj: bool = True, - - **kwargs + **kwargs, ): super().__init__() if moe_implementation is None: @@ -112,8 +146,9 @@ def __init__( self.num_steps = 10 self.n_obs_steps = 1 - assert not (split_gate_liner and nosplit_gate_liner), \ + assert not (split_gate_liner and nosplit_gate_liner), ( "split_gate_liner and nosplit_gate_liner can not be both True." + ) self.vlm_repo_id = vlm_repo_id self.expert_vision_path = expert_vision_path @@ -175,9 +210,9 @@ def __init__( self.use_qwen3_fixed_grid_cache = use_qwen3_fixed_grid_cache self.use_lm_head = use_lm_head if vocab_size == 0: - if vlm_repo_id and 'paligemma' in vlm_repo_id.lower(): + if vlm_repo_id and "paligemma" in vlm_repo_id.lower(): self.vocab_size = 257216 - elif vlm_repo_id and 'qwen' in vlm_repo_id.lower(): + elif vlm_repo_id and "qwen" in vlm_repo_id.lower(): self.vocab_size = 151936 else: self.vocab_size = 257152 @@ -185,6 +220,7 @@ def __init__( self.vocab_size = vocab_size self.vit_attn_implementation = vit_attn_implementation + class LingbotVLAV2Config(LingbotVLAConfig): def __init__(self, **kwargs): kwargs.setdefault("attention_implementation", "flex_cached") @@ -206,225 +242,14 @@ def __init__(self, **kwargs): __all__ = ["LingbotVLAConfig", "LingbotVLAV2Config"] - -# Shared V1 flow-matching base retained by the V2 architecture. -from logging import raiseExceptions -import einops -import numpy as np -import torch -from torch import nn -import torch.nn.functional as F -from torch import Tensor, nn -from typing import Any, Callable, Dict, List, Optional, Tuple, Union -from functools import partial -import math -from transformers import ( - AutoConfig, - PretrainedConfig, - PreTrainedModel, -) -from transformers.models.auto import CONFIG_MAPPING -from transformers import AutoTokenizer -from transformers.cache_utils import Cache -from transformers.generation import GenerationMixin -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from transformers.utils import ( - is_torchdynamo_compiling, - logging, -) - - -from transformers.utils.deprecation import deprecate_kwarg -from transformers.activations import ACT2FN -from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, is_flash_attn_available -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from transformers.processing_utils import Unpack -from telefuser.models.lingbot_vla_v2_qwen import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLTextModel, Qwen2_5_VLPreTrainedModel - -from transformers.models.qwen2.modeling_qwen2 import ( - Qwen2RMSNorm, -) - -try: - from dinov3.hub.backbones import ( - dinov3_vits16, - dinov3_vits16plus, - dinov3_vitb16, - ) -except: pass -from telefuser.models.lingbot_vla_v2_loader import ( - create_sinusoidal_pos_embedding, - make_att_2d_masks, - resize_with_pad, -) -from telefuser.models.lingbot_vla_v2_loader import apply_rope, our_eager_attention_forward -from telefuser.models.lingbot_vla_v2_loader import flex_attention_forward -from telefuser.models.lingbot_vla_v2_loader import build_block_mask, flex_attention_with_block_mask - -from telefuser.models.lingbot_vla_v2_loader import LingBotVLAWeightLoader, TaskTokenDepthHead -from telefuser.models.lingbot_vla_v2_moe import ( - Qwen2ForCausalLM, - Qwen2FusedExperts, - Qwen2TokenMoeBlock, - FixQwen2RMSNorm, -) - -logger = logging.get_logger(__name__) - -class QwenvlWithExpertConfig(PretrainedConfig): - model_type = "QwenvlWithExpertModel" - sub_configs = {"qwenvl_config": AutoConfig, "qwen_expert_config": AutoConfig} - - def __init__( - self, - qwenvl_config: dict | None = None, - qwen_expert_config: dict | None = None, - freeze_vision_encoder: bool = False, - train_expert_only: bool = False, - vocab_size: int = 257152, - use_lm_head: bool = False, - attention_implementation: str = "eager", - tokenizer_path: str | None = None, - enable_expert_vision: bool = False, - expert_vision_type: str | None = None, - use_cache: bool = False, - expert_hidden_size: int = 768, - expert_intermediate_size: int = 2752, - **kwargs, - ): - self.freeze_vision_encoder = freeze_vision_encoder - self.train_expert_only = train_expert_only - self.attention_implementation = attention_implementation - self.tokenizer_path = tokenizer_path - self.enable_expert_vision = enable_expert_vision - self.expert_vision_type = expert_vision_type - self.vocab_size = vocab_size - self.use_lm_head = use_lm_head - if qwenvl_config is None: - self.qwenvl_config = CONFIG_MAPPING["qwen2_5_vl"]( - attention_dropout=0.0, - bos_token_id=151643, - eos_token_id=151645, - vision_start_token_id=151652, - vision_end_token_id=151653, - vision_token_id=151654, - image_token_id=151655, - video_token_id=151656, - hidden_act="silu", - hidden_size=2048, - initializer_range=0.02, - intermediate_size=11008, - max_position_embeddings=128000, - max_window_layers=70, - model_type="qwen2_5_vl", - num_attention_heads=16, - num_hidden_layers=36, - num_key_value_heads=2, - rms_norm_eps=1e-06, - rope_theta=1000000.0, - sliding_window=32768, - tie_word_embeddings=True, - torch_dtype="bfloat16", - transformers_version="4.41.2", - use_cache=True, - use_sliding_window=False, - vision_config={ - "depth": 32, - "hidden_act": "silu", - "hidden_size": 1280, - "intermediate_size": 3420, - "num_heads": 16, - "in_chans": 3, - "out_hidden_size": 2048, - "patch_size": 14, - "spatial_merge_size": 2, - "spatial_patch_size": 14, - "window_size": 112, - "fullatt_block_indexes": [ - 7, - 15, - 23, - 31 - ], - "tokens_per_second": 2, - "temporal_patch_size": 2 - }, - rope_scaling={ - "type": "mrope", - "mrope_section": [ - 16, - 24, - 24 - ] - }, - vocab_size=151936, - ) - elif isinstance(self.qwenvl_config, dict): - if "model_type" not in qwen_expert_config: - qwenvl_config["model_type"] = "qwen2_5_vl" - - cfg_cls = CONFIG_MAPPING[qwenvl_config["model_type"]] - self.qwenvl_config = cfg_cls(**qwenvl_config) - - if qwen_expert_config is None: - self.qwen_expert_config = CONFIG_MAPPING["qwen2"]( - attention_dropout=0.0, - bos_token_id=151643, - eos_token_id=151645, - hidden_act="silu", - hidden_size=expert_hidden_size, - head_dim=128, - initializer_range=0.02, - intermediate_size=expert_intermediate_size, - max_position_embeddings=32768, - max_window_layers=21, - model_type="qwen2", - num_attention_heads=16, - num_hidden_layers=36, - num_key_value_heads=2, - rms_norm_eps=1e-06, - rope_theta=1000000.0, - sliding_window=32768, - tie_word_embeddings=True, - torch_dtype="bfloat16", - transformers_version="4.43.1", - use_cache=use_cache, - use_sliding_window=False, - vocab_size=151936, - ) - elif isinstance(self.qwen_expert_config, dict): - if "model_type" not in qwen_expert_config: - qwen_expert_config["model_type"] = "qwen2" - - cfg_cls = CONFIG_MAPPING[qwenvl_config["model_type"]] - self.qwen_expert_config = cfg_cls(**qwen_expert_config) - - super().__init__(**kwargs) - - def __post_init__(self): - super().__post_init__() - if self.train_expert_only and not self.freeze_vision_encoder: - raise ValueError( - "You set `freeze_vision_encoder=False` and `train_expert_only=True` which are not compatible." - ) - - if self.attention_implementation not in ["eager", "fa2", "flex"]: - raise ValueError( - f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). Expected 'eager', 'fa2' or 'flex'." - ) - class AdaRMSNorm(nn.Module): def __init__(self, hidden_size, cond_dim, eps=1e-6): - """ - AdaRMSNorm: RMSNorm + FiLM - """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps self.gamma = nn.Linear(cond_dim, hidden_size) self.beta = nn.Linear(cond_dim, hidden_size) - # DiT style init: gamma.weight=0, gamma.bias=1; beta.weight=0, beta.bias=0 nn.init.zeros_(self.gamma.weight) nn.init.zeros_(self.gamma.bias) nn.init.zeros_(self.beta.weight) @@ -435,485 +260,80 @@ def forward(self, hidden_states, cond): hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - hidden_states = self.weight * hidden_states - # cond = cond.to(torch.float32) - gamma = self.gamma(cond).unsqueeze(1) # [B, 1, H] - beta = self.beta(cond).unsqueeze(1) # [B, 1, H] + gamma = self.gamma(cond).unsqueeze(1) + beta = self.beta(cond).unsqueeze(1) hidden_states = (1 + gamma.to(torch.float32)) * hidden_states + beta.to(torch.float32) return hidden_states.to(input_dtype) -class FixAdaRMSNorm(nn.Module): - def __init__(self, hidden_size, cond_dim, eps=1e-6): - """ - AdaRMSNorm: RMSNorm + FiLM - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - self.gamma = nn.Linear(cond_dim, hidden_size) - self.beta = nn.Linear(cond_dim, hidden_size) - - # DiT style init: gamma.weight=0, gamma.bias=1; beta.weight=0, beta.bias=0 - nn.init.zeros_(self.gamma.weight) - nn.init.zeros_(self.gamma.bias) - nn.init.zeros_(self.beta.weight) - nn.init.zeros_(self.beta.bias) +class FixAdaRMSNorm(AdaRMSNorm): def forward(self, hidden_states, cond): - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return super().forward(hidden_states, cond.float()) - hidden_states = self.weight * hidden_states - cond = cond.to(torch.float32) - gamma = self.gamma(cond).unsqueeze(1) # [B, 1, H] - beta = self.beta(cond).unsqueeze(1) # [B, 1, H] - hidden_states = (1 + gamma.to(torch.float32)) * hidden_states + beta.to(torch.float32) - return hidden_states.to(input_dtype) -# HACK: show directly use this norm during initialization -# TODO: clear the logics def replace_lnorm_with_adanorm(module, hidden_size, cond_dim, final_norm_adanorm): for name, child in module.named_children(): - if final_norm_adanorm: - if isinstance(child, Qwen2RMSNorm): - if 'q_layernorm' not in name and 'k_layernorm' not in name: - setattr(module, name, AdaRMSNorm(hidden_size, cond_dim)) - elif isinstance(child, FixQwen2RMSNorm): - if 'q_layernorm' not in name and 'k_layernorm' not in name: - setattr(module, name, FixAdaRMSNorm(hidden_size, cond_dim)) - else: - replace_lnorm_with_adanorm(child, hidden_size, cond_dim, final_norm_adanorm) - else: - if isinstance(child, Qwen2RMSNorm): - if 'q_layernorm' not in name and 'k_layernorm' not in name: - setattr(module, name, AdaRMSNorm(hidden_size, cond_dim)) - else: - replace_lnorm_with_adanorm(child, hidden_size, cond_dim, final_norm_adanorm) - -class QwenvlWithExpertModel(PreTrainedModel): - config_class = QwenvlWithExpertConfig - - def __init__(self, config: QwenvlWithExpertConfig, eval=False): - super().__init__(config=config) - self.config = config - vlm_config = AutoConfig.from_pretrained(self.config.tokenizer_path, local_files_only=True) - vlm_config.vision_config.initializer_range = 0.02 - print(f'=====Vocab_size in Config is {self.config.vocab_size}=====') - if self.config.vocab_size != 0 and self.config.vocab_size != 257152 and vlm_config.vocab_size != self.config.vocab_size: - vlm_config.vocab_size = self.config.vocab_size - print(f'====Vocabulary Size is {vlm_config.vocab_size}====') - vlm_config._attn_implementation = "flash_attention_2" - vlm_config.vision_config._attn_implementation = self.config.vit_attn_implementation - self.qwenvl = Qwen2_5_VLForConditionalGeneration._from_config(vlm_config) - if self.config.use_lm_head: - self.qwenvl.tie_weights() - self.config.qwen_expert_config._attn_implementation = "flash_attention_2" - self.qwen_expert = Qwen2ForCausalLM._from_config(self.config.qwen_expert_config, eval=eval) - - if getattr(self.config, 'adanorm_time', False): - replace_lnorm_with_adanorm(self.qwen_expert, self.config.qwen_expert_config.hidden_size, self.config.qwen_expert_config.hidden_size, config.final_norm_adanorm) - if getattr(self.config, 'use_moe', False): - bias_update_speed = getattr(self.config, 'bias_update_speed', 0.001) - hidden_size = self.config.qwen_expert_config.hidden_size # 768 - - token_moe_layers = getattr(self.config, 'token_moe_layers', None) or [] - - if token_moe_layers: - token_config = CONFIG_MAPPING['qwen2_moe']( - num_experts=getattr(self.config, 'token_num_experts', 32), - num_experts_per_tok=getattr(self.config, 'token_top_k', 1), - norm_topk_prob=True, - hidden_size=hidden_size, - moe_intermediate_size=getattr(self.config, 'token_moe_intermediate_size', 256), - shared_expert_intermediate_size=getattr(self.config, 'token_shared_intermediate_size', 256), - output_router_logits=False, - ) - token_config.bias_update_speed = bias_update_speed - token_config._moe_implementation = getattr(self.config, '_moe_implementation', None) - token_config.router_activation = getattr(self.config, 'router_activation', 'softmax') - token_config.routed_scaling_factor = getattr(self.config, 'routed_scaling_factor', 1.0) - token_config.use_shared_expert_gate = getattr(self.config, 'use_shared_expert_gate', True) - for idx in token_moe_layers: - self.qwen_expert.model.layers[idx].mlp = Qwen2TokenMoeBlock(token_config) - # Precomputed grid_thw cache (populated on first call when precompute_grid_thw=True) - self.rotary_pos_emb = None - self.window_index = None - self.cu_window_seqlens = None - self.cu_seqlens = None - - # Remove unused embed_tokens - del self.qwen_expert.model.embed_tokens - if self.config.enable_expert_vision: - if 'dinov3_vitb16' in self.config.expert_vision_type: - self.expert_visual = dinov3_vitb16(pretrained=False) - self.expert_visual_mlp = nn.Sequential( - nn.Linear(self.expert_visual.embed_dim, self.expert_visual.embed_dim * 2), - nn.GELU(), - nn.Linear(self.expert_visual.embed_dim * 2, self.config.qwen_expert_config.hidden_size), - ) - self.attention_interface = self.get_attention_interface() - - # self.to_bfloat16_like_physical_intelligence() - - - def to_bfloat16_like_physical_intelligence(self): - """casts the model to bfloat16. - - Modules not casted to bfloat16: - - .qwenvl.model.embed_tokens.weight - - .qwenvl.model.norm.weight - - qwen_expert.model.norm.weight - - qwen_expert.lm_head.weight - """ - self.qwenvl = self.qwenvl.to(dtype=torch.bfloat16) - - params_to_change_dtype = [ - ".qwenvl.model.layers", - "qwen_expert.model.layers", - "visual", - "multi_modal", - ] - for name, param in self.named_parameters(): - if any(selector in name for selector in params_to_change_dtype): - param.data = param.data.to(dtype=torch.bfloat16) - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None, precompute_grid_thw: bool = False): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - precompute_grid_thw (`bool`): If True, compute and cache rotary_pos_emb/window_index/cu_seqlens on first call. - """ - if precompute_grid_thw and self.rotary_pos_emb is None: - ( - self.rotary_pos_emb, - self.window_index, - self.cu_window_seqlens, - self.cu_seqlens - ) = self.qwenvl.visual.preprcess_grid_thw(grid_thw=image_grid_thw) - image_embeds = self.qwenvl.visual( - pixel_values, - grid_thw=image_grid_thw, - rotary_pos_emb=self.rotary_pos_emb, - window_index=self.window_index, - cu_window_seqlens=self.cu_window_seqlens, - cu_seqlens=self.cu_seqlens, - ) - split_sizes = (image_grid_thw.prod(-1) // self.qwenvl.visual.spatial_merge_size**2).tolist() - image_embeds = torch.split(image_embeds, split_sizes) - image_embeds = torch.stack(image_embeds, dim=0) - return image_embeds - - def embed_image(self, image: torch.Tensor, patch_size=14, temporal_patch_size=2, precompute_grid_thw=False): - h = w = int(image.shape[1] ** 0.5) - image_grid_thw = torch.tensor([[1, h, w]]*image.shape[0], device=image.device) - image_embeds = self.get_image_features(image, image_grid_thw=image_grid_thw, precompute_grid_thw=precompute_grid_thw) - return image_embeds - # return torch.randn(72, 64, 2048).to(device=image.device, dtype=torch.bfloat16) - - def embed_language_tokens(self, tokens: torch.Tensor): - return self.qwenvl.model.embed_tokens(tokens) - - def handle_kv_cache( - self, - key_states: torch.Tensor, - value_states: torch.Tensor, - layer_idx: int, - past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, - use_cache: Optional[bool] = None, - fill_kv_cache: Optional[bool] = None, - ): - if use_cache: - if past_key_values is None: - past_key_values = {} - - if fill_kv_cache: - past_key_values[layer_idx] = { - "key_states": key_states, - "value_states": value_states, - } - else: - key_states = torch.cat( - [past_key_values[layer_idx]["key_states"], key_states], dim=1 - ) - value_states = torch.cat( - [past_key_values[layer_idx]["value_states"], value_states], - dim=1, - ) - return key_states, value_states, past_key_values - - def forward( - self, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - vlm_position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, - inputs_embeds: List[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - fill_kv_cache: Optional[bool] = None, - ada_cond: List[torch.FloatTensor] = None, - ): - """ - Args: - attention_mask (Optional[torch.Tensor], optional): - Attention mask with shape (b, seq_len, seq_len). Defaults to None. - position_ids (Optional[torch.LongTensor], optional): - Position indices for applying RoPE. Defaults to None. - past_key_values (Optional[Union[List[torch.FloatTensor], Cache]], optional): - Optional kv cache. Defaults to None. - inputs_embeds (List[torch.FloatTensor], optional): - Input embeddings. Defaults to None. - use_cache (Optional[bool], optional): - Whether to use kv cache. Defaults to None. - fill_kv_cache (Optional[bool], optional): - Whether to return kv tensors in this forward pass as cache. Defaults to None. - - Returns: - outputs_embeds (torch.Tensor): Output embeddings. - past_key_values (Optional[Union[List[torch.FloatTensor], Cache]]): - Optional kv cache. - """ - models = [self.qwenvl.model, self.qwen_expert.model] # Qwen2_5_VLTextModel, Qwen2Model (We have re-writeen their forward as follows:) - - # RMSNorm - num_layers = self.qwenvl.config.num_hidden_layers # 36 - action_num_layers = self.config.qwen_expert_config.num_hidden_layers # 36 - assert action_num_layers == num_layers, ( - "Action expert and VLM must have the same number of layers " - f"(got action={action_num_layers}, vlm={num_layers})." - ) - - router_logits_list = [] - for layer_idx in range(num_layers): - query_states = [] - key_states = [] - value_states = [] - for i, hidden_states in enumerate(inputs_embeds): - if hidden_states is None: - continue - if i == 1: # For action expert - query_state, key_state, value_state = models[i].layers[layer_idx](hidden_states, compute_kqv=True, ada_cond = ada_cond) - else: # For VLM - query_state, key_state, value_state = models[i].layers[layer_idx](hidden_states, compute_kqv=True) - - if query_state.dtype != torch.float32: - query_state, key_state, value_state = query_state.to(torch.float32), key_state.to(torch.float32), value_state.to(torch.float32) - query_states.append(query_state) - key_states.append(key_state) - value_states.append(value_state) - - # B,L,H,D with L sequence length (img, lang, state, action), H number of heads, D head dim - # concatenate on the number of embeddings/tokens - query_states = torch.cat(query_states, dim=1) - key_states = torch.cat(key_states, dim=1) - value_states = torch.cat(value_states, dim=1) - - query_states = apply_rope(query_states, position_ids) - key_states = apply_rope(key_states, position_ids) - - key_states, value_states, past_key_values = self.handle_kv_cache( - key_states, - value_states, - layer_idx, - past_key_values=past_key_values, - use_cache=use_cache, - fill_kv_cache=fill_kv_cache, - ) - if self.config.attention_implementation == "flex_cached": - if layer_idx == 0: - _full_len = query_states.shape[1] - _full_block_mask = build_block_mask(attention_mask, self.qwenvl.config.num_attention_heads, _full_len, _full_len) - att_output = flex_attention_with_block_mask(query_states, key_states, value_states, _full_block_mask, query_states.shape[1]) - else: - att_output = self.attention_interface(query_states, key_states, value_states, attention_mask) - - # first part of att_output is prefix (up to sequence length, [:, 0:prefix_seq_len]) - outputs_embeds = [] - start = 0 - for i, hidden_states in enumerate(inputs_embeds): - if hidden_states is not None: - end = start + hidden_states.shape[1] - if i == 1: - out_emb, _router_logits = models[i].layers[layer_idx](hidden_states, att_output, start, end, output_atten=True, ada_cond = ada_cond) - if _router_logits is not None: - router_logits_list.append(_router_logits) - else: - out_emb = models[i].layers[layer_idx](hidden_states, att_output, start, end, output_atten=True) - outputs_embeds.append(out_emb) - start = end - else: - outputs_embeds.append(None) - - inputs_embeds = outputs_embeds - - # final norm - outputs_embeds = [] - for i, hidden_states in enumerate(inputs_embeds): - if hidden_states is not None: - if self.config.final_norm_adanorm: - if i == 1: - out_emb, _ = models[i].norm(hidden_states, ada_cond) - else: - out_emb = models[i].norm(hidden_states) - else: - out_emb = models[i].norm(hidden_states) - outputs_embeds.append(out_emb) - else: - outputs_embeds.append(None) - - return outputs_embeds, past_key_values, router_logits_list - - def get_attention_interface(self): - if self.config.attention_implementation == "fa2": - raise NotImplementedError("FA2 is not implemented (yet)") - elif self.config.attention_implementation == "flex": - print('=====Using Flex Attn=====') - attention_interface = flex_attention_forward - elif self.config.attention_implementation == "eager": - print('=====Using Eager Attn=====') - attention_interface = our_eager_attention_forward - elif self.config.attention_implementation == "flex_cached": - print('=====Using Flex Cached (prebuilt BlockMask) Attn=====') - attention_interface = flex_attention_forward # fallback - elif self.config.attention_implementation == "xformer": - # attention_interface = xformer_attention_forward - raise NotImplementedError("Xformer attention is not implemented (yet)") - else: - raise ValueError( - f"Invalid attention implementation: {self.config.attention_implementation}. " - "Expected one of ['fa2', 'flex', 'flex_cached', 'eager', 'xformer']." - ) - return attention_interface - -class FlowMatching(nn.Module): - def __init__(self, config, eval): - super().__init__() - self.config = config - - # qwenvl with action expert - qwenvl_with_export_config = QwenvlWithExpertConfig( - freeze_vision_encoder=self.config.freeze_vision_encoder, - train_expert_only=self.config.train_expert_only, - vocab_size=getattr(self.config,"vocab_size", 0), - use_lm_head=getattr(self.config,"use_lm_head", False), - attention_implementation=self.config.attention_implementation, - tokenizer_path=self.config.tokenizer_path, - enable_expert_vision=self.config.enable_expert_vision, - expert_vision_type=self.config.expert_vision_type, - use_cache=getattr(self.config,"use_cache", True), - expert_hidden_size=getattr(self.config, 'expert_hidden_size', 768), - expert_intermediate_size=getattr(self.config, 'expert_intermediate_size', 2752), - ) - qwenvl_with_export_config.adanorm_time = getattr(config, "adanorm_time", False) - qwenvl_with_export_config.final_norm_adanorm = getattr(config, "final_norm_adanorm", False) - qwenvl_with_export_config.vit_attn_implementation = getattr(config, "vit_attn_implementation", "flash_attention_2") - if getattr(config, "use_moe", False): - qwenvl_with_export_config.use_moe = config.use_moe - qwenvl_with_export_config.bias_update_speed = getattr(config, "bias_update_speed", 0.001) - qwenvl_with_export_config.token_moe_layers = getattr(config, "token_moe_layers", None) - qwenvl_with_export_config.token_num_experts = getattr(config, "token_num_experts", 32) - qwenvl_with_export_config.token_top_k = getattr(config, "token_top_k", 1) - qwenvl_with_export_config.token_moe_intermediate_size = getattr(config, "token_moe_intermediate_size", 256) - qwenvl_with_export_config.token_shared_intermediate_size = getattr(config, "token_shared_intermediate_size", 256) - # Pass _moe_implementation through for EP/fused support - qwenvl_with_export_config._moe_implementation = getattr(config, '_moe_implementation', None) - self.qwenvl_with_expert = QwenvlWithExpertModel( - qwenvl_with_export_config, eval - ) - self.config.proj_width = qwenvl_with_export_config.qwen_expert_config.hidden_size - self.config.initializer_range = getattr(qwenvl_with_export_config.qwen_expert_config, "initializer_range", None) - # projection layers - self.state_proj = nn.Linear(self.config.max_state_dim, self.config.proj_width) - self.action_in_proj = nn.Linear( - self.config.max_action_dim, self.config.proj_width - ) - self.action_out_proj = nn.Linear( - self.config.proj_width, self.config.max_action_dim - ) - self.action_time_mlp_in = nn.Linear( - self.config.proj_width * 2, self.config.proj_width - ) - self.action_time_mlp_out = nn.Linear( - self.config.proj_width, self.config.proj_width - ) - self.config.align_params = getattr(self.config, 'align_params', {}) - if self.config.align_params != {}: - self.steps=0 - self.use_depth_align = True - self.init_depth_heads(self.config.align_params) - self.use_future_video = self.config.align_params.get('use_future_video', False) - if self.use_future_video: - self.init_video_heads(self.config.align_params) + if isinstance(child, Qwen2RMSNorm) and "q_layernorm" not in name and "k_layernorm" not in name: + setattr(module, name, AdaRMSNorm(hidden_size, cond_dim)) + elif ( + final_norm_adanorm + and isinstance(child, FixQwen2RMSNorm) + and "q_layernorm" not in name + and "k_layernorm" not in name + ): + setattr(module, name, FixAdaRMSNorm(hidden_size, cond_dim)) else: - self.use_depth_align = False - self.use_future_video = False - self.use_future_video_patch = False - self.use_current_video_patch = False - self.use_current_shared_task_proj = False - self.use_future_video_cls = False - self.use_shared_future_task_proj = False - self.future_video_share_future_depth_query = False + replace_lnorm_with_adanorm(child, hidden_size, cond_dim, final_norm_adanorm) +class FlowMatchingBase(nn.Module): def init_depth_heads(self, config): - self.llm_image_token_size = config['llm']['image_token_size'] - self.llm_image_input_size = config['llm']['image_input_size'] - self.depth_token_size = config['depth']['token_size'] - self.depth_input_size = config['depth']['input_size'] - self.align_type = config.get('mode', None) - self.model_type = config['depth']['model_type'] + self.llm_image_token_size = config["llm"]["image_token_size"] + self.llm_image_input_size = config["llm"]["image_input_size"] + self.depth_token_size = config["depth"]["token_size"] + self.depth_input_size = config["depth"]["input_size"] + self.align_type = config.get("mode", None) + self.model_type = config["depth"]["model_type"] if self.align_type != "query": raise ValueError(f"Only query depth alignment is supported, got {self.align_type!r}.") if self.model_type != "MoRGBD": raise ValueError(f"Only MoRGBD depth distillation is supported, got {self.model_type!r}.") - self.use_future_depth = (config.get('depth') or {}).get('use_future_depth', False) - self.block_future_depth_to_action = (config.get('depth') or {}).get('block_future_depth_to_action', False) - self.detach_future_depth_image_feats = bool( - (config.get('depth') or {}).get('detach_future_image_feats', False) - ) - self.use_future_video = bool(config.get('use_future_video', False)) + self.use_future_depth = (config.get("depth") or {}).get("use_future_depth", False) + self.block_future_depth_to_action = (config.get("depth") or {}).get("block_future_depth_to_action", False) + self.detach_future_depth_image_feats = bool((config.get("depth") or {}).get("detach_future_image_feats", False)) + self.use_future_video = bool(config.get("use_future_video", False)) self.use_future_video_patch = False self.use_current_video_patch = False self.use_current_shared_task_proj = False self.use_future_video_cls = False self.use_shared_future_task_proj = False self.future_video_share_future_depth_query = False - self.num_task_tokens = config['num_task_tokens'] - assert config['depth']['num_backbone_tokens'] % self.num_task_tokens == 0 + self.num_task_tokens = config["num_task_tokens"] + assert config["depth"]["num_backbone_tokens"] % self.num_task_tokens == 0 self.depth_align_embs = nn.Parameter( - torch.randn( - config['depth']['num_backbone_tokens'], config['llm']['dim_out'] - ) + torch.randn(config["depth"]["num_backbone_tokens"], config["llm"]["dim_out"]) ) - self.depth_align_head = TaskTokenDepthHead(config['depth'], llm_hidden_size=config['llm']['dim_out']).to(dtype=torch.bfloat16) - + self.depth_align_head = TaskTokenDepthHead(config["depth"], llm_hidden_size=config["llm"]["dim_out"]).to( + dtype=torch.bfloat16 + ) if self.use_future_depth: self.future_depth_align_embs = nn.Parameter( - torch.randn( - config['depth']['num_backbone_tokens'], config['llm']['dim_out'] - ) + torch.randn(config["depth"]["num_backbone_tokens"], config["llm"]["dim_out"]) ) self.future_depth_align_head = TaskTokenDepthHead( - config['depth'], llm_hidden_size=config['llm']['dim_out'] + config["depth"], llm_hidden_size=config["llm"]["dim_out"] ).to(dtype=torch.bfloat16) - def init_video_heads(self, config): if self.align_type != "query": raise ValueError("future-video alignment is only supported for query align mode.") - video_config = dict(config.get('depth', {})) - video_config.update(config.get('video', {})) + video_config = dict(config.get("depth", {})) + video_config.update(config.get("video", {})) required_keys = ("num_backbone_tokens", "dim_out", "num_layers", "num_heads", "dim_head", "ff_mult") missing = [key for key in required_keys if key not in video_config] if missing: @@ -922,8 +342,7 @@ def init_video_heads(self, config): self.use_current_video_patch = bool(video_config.get("use_current_patch_loss", False)) if self.use_current_video_patch and not self.use_future_video_patch: raise ValueError( - "align_params.video.use_current_patch_loss=True requires " - "align_params.video.use_patch_loss=True." + "align_params.video.use_current_patch_loss=True requires align_params.video.use_patch_loss=True." ) self.use_current_shared_task_proj = bool( video_config.get("use_current_shared_task_proj", self.use_current_video_patch) @@ -934,16 +353,11 @@ def init_video_heads(self, config): "align_params.video.use_current_patch_loss=True." ) self.use_future_video_cls = bool(video_config.get("use_cls_loss", False)) - self.future_video_share_future_depth_query = bool( - video_config.get("share_future_depth_query", False) - ) - self.use_shared_future_task_proj = bool( - video_config.get("use_shared_future_task_proj", False) - ) + self.future_video_share_future_depth_query = bool(video_config.get("share_future_depth_query", False)) + self.use_shared_future_task_proj = bool(video_config.get("use_shared_future_task_proj", False)) if self.use_shared_future_task_proj and not self.use_future_video_patch: raise ValueError( - "align_params.video.use_shared_future_task_proj=True requires " - "align_params.video.use_patch_loss=True." + "align_params.video.use_shared_future_task_proj=True requires align_params.video.use_patch_loss=True." ) if self.use_shared_future_task_proj and not self.future_video_share_future_depth_query: raise ValueError( @@ -958,8 +372,7 @@ def init_video_heads(self, config): ) if int(video_config["num_backbone_tokens"]) != int(config["depth"]["num_backbone_tokens"]): raise ValueError( - "future-video shared query requires video.num_backbone_tokens " - "to match depth.num_backbone_tokens." + "future-video shared query requires video.num_backbone_tokens to match depth.num_backbone_tokens." ) self.block_suffix_to_future_video = bool(video_config.get("block_suffix_to_future_video", False)) @@ -972,42 +385,35 @@ def init_video_heads(self, config): if self.use_future_video_patch: if self.use_current_video_patch: self.current_video_align_embs = nn.Parameter( - torch.randn( - video_config['num_backbone_tokens'], config['llm']['dim_out'] - ) + torch.randn(video_config["num_backbone_tokens"], config["llm"]["dim_out"]) ) if self.use_current_shared_task_proj: self.current_shared_task_proj = nn.Linear( - config['llm']['dim_out'] * 2, - config['llm']['dim_out'], + config["llm"]["dim_out"] * 2, + config["llm"]["dim_out"], ) self.current_video_align_head = TaskTokenDepthHead( - video_config, llm_hidden_size=config['llm']['dim_out'] + video_config, llm_hidden_size=config["llm"]["dim_out"] ).to(dtype=torch.bfloat16) - if ( - not self.future_video_share_future_depth_query - or self.use_shared_future_task_proj - ): + if not self.future_video_share_future_depth_query or self.use_shared_future_task_proj: self.future_video_align_embs = nn.Parameter( - torch.randn( - video_config['num_backbone_tokens'], config['llm']['dim_out'] - ) + torch.randn(video_config["num_backbone_tokens"], config["llm"]["dim_out"]) ) if self.use_shared_future_task_proj: self.future_shared_task_proj = nn.Linear( - config['llm']['dim_out'] * 2, - config['llm']['dim_out'], + config["llm"]["dim_out"] * 2, + config["llm"]["dim_out"], ) self.future_video_align_head = TaskTokenDepthHead( - video_config, llm_hidden_size=config['llm']['dim_out'] + video_config, llm_hidden_size=config["llm"]["dim_out"] ).to(dtype=torch.bfloat16) if self.use_future_video_cls: - self.future_video_cls_align_emb = nn.Embedding(1, config['llm']['dim_out']) + self.future_video_cls_align_emb = nn.Embedding(1, config["llm"]["dim_out"]) self.future_video_cls_head = nn.Sequential( - nn.LayerNorm(config['llm']['dim_out']), - nn.Linear(config['llm']['dim_out'], video_config['dim_out']), + nn.LayerNorm(config["llm"]["dim_out"]), + nn.Linear(config["llm"]["dim_out"], video_config["dim_out"]), ).to(dtype=torch.bfloat16) def _future_depth_token_count(self): @@ -1017,9 +423,8 @@ def _future_video_own_token_count(self): if not getattr(self, "use_future_video", False): return 0 count = 1 if getattr(self, "use_future_video_cls", False) else 0 - if ( - getattr(self, "use_future_video_patch", True) - and not getattr(self, "future_video_share_future_depth_query", False) + if getattr(self, "use_future_video_patch", True) and not getattr( + self, "future_video_share_future_depth_query", False ): count += self.num_task_tokens return count @@ -1034,7 +439,7 @@ def _future_video_own_span(self, hidden_states): def _future_depth_task_tokens(self, hidden_states): if not getattr(self, "use_future_depth", False): raise ValueError("future-depth query tokens are not enabled.") - return hidden_states[:, -self.num_task_tokens:, :] + return hidden_states[:, -self.num_task_tokens :, :] def _future_video_cls_task_tokens(self, hidden_states): if not getattr(self, "use_future_video_cls", False): @@ -1089,7 +494,6 @@ def _block_suffix_to_future_video_if_enabled_( prefix_len=prefix_len, ) - def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, (nn.Linear, nn.Conv3d)): @@ -1115,91 +519,22 @@ def _init_weights(self, module): @staticmethod def _fp32_linear(module, x): """Compute linear layer in fp32 regardless of module's current parameter dtype.""" - return F.linear( - x.float(), - module.weight.float(), - module.bias.float() if module.bias is not None else None - ) + return F.linear(x.float(), module.weight.float(), module.bias.float() if module.bias is not None else None) - - def embed_prefix( - self, images, img_masks, lang_tokens, lang_masks, vlm_causal, precompute_grid_thw=False - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - bsize = images.shape[0] - device = images.device - dtype = images.dtype - - # embed image - if images.ndim == 5: - images = einops.rearrange(images, "b n c h w -> (b n) c h w") - elif images.ndim == 4: - images = einops.rearrange(images, "b n l d -> (b n) l d") - elif images.ndim == 3: # For inference bs=1 - bsize = 1 - img_emb = self.qwenvl_with_expert.embed_image(images, precompute_grid_thw=precompute_grid_thw) - num_patch = img_emb.shape[1] - img_emb = einops.rearrange(img_emb, "(b n) l d -> b (n l) d", b=bsize) # bsize = 24 - num_img_embs = img_emb.shape[1] - if img_masks.ndim ==1: # For inference bs=1 - img_masks = img_masks.unsqueeze(0) - if self.use_depth_align and self.align_type == "query": - align_masks = einops.repeat(img_masks, "b n -> b (n l)", l=self.num_task_tokens) - img_masks = einops.repeat(img_masks, "b n -> b (n l)", l=num_patch) - - # embed language - lang_emb = self.qwenvl_with_expert.embed_language_tokens(lang_tokens) - num_lang_embs = lang_emb.shape[1] - - if self.use_depth_align and self.align_type == "query": - def _get_align_tokens(tokens): - tk_weights = tokens.view(self.num_task_tokens, tokens.shape[0] // self.num_task_tokens, tokens.shape[1]) - tk_weights = tk_weights.mean(dim=1) - return tk_weights - - align_embs = _get_align_tokens(self.depth_align_embs).repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) - # align_masks = einops.rearrange(img_masks, "b (n l) -> b n l", n=3) - # align_masks = align_masks[:, :, 0] - # align_masks = einops.repeat(align_masks, "b n -> b (n l)", l=self.num_task_tokens) - embs = torch.cat([img_emb, align_embs, align_embs, align_embs, lang_emb], dim=1) - pad_masks = torch.cat([img_masks, align_masks, lang_masks], dim=1) - else: - # assemble embeddings - embs = torch.cat([img_emb, lang_emb], dim=1) - pad_masks = torch.cat([img_masks, lang_masks], dim=1) - - # (see `make_att_2d_masks` to understand why zeros means bidirection) - if not vlm_causal: - if self.use_depth_align and self.align_type == "query": - att_masks = torch.zeros( - (img_emb.size(0), num_img_embs + 3 * self.num_task_tokens + num_lang_embs), device=device, dtype=torch.bool - ) # 1, bs_img*(768+48) - else: - att_masks = torch.zeros( - (img_emb.size(0), num_img_embs + num_lang_embs), device=device, dtype=torch.bool - ) # 1, bs_img*(768+48) - else: - if self.use_depth_align and self.align_type == "query": - att_masks = torch.ones( - (img_emb.size(0), num_img_embs + 3 * self.num_task_tokens + num_lang_embs), device=device, dtype=torch.bool - ) # 1, bs_img*(768+48) - else: - att_masks = torch.ones( - (img_emb.size(0), num_img_embs + num_lang_embs), device=device, dtype=torch.bool - ) # 1, bs_img*(768+48) - return embs, pad_masks, att_masks - - def embed_suffix(self, state, noisy_actions, timestep): # (torch.Size([state_bs, 32]), torch.Size([1, state_bs*50, 32]), torch.Size([1])) - bsize = state.shape[0] # state_bs = img_bs + def embed_suffix( + self, state, noisy_actions, timestep + ): # (torch.Size([state_bs, 32]), torch.Size([1, state_bs*50, 32]), torch.Size([1])) + bsize = state.shape[0] # state_bs = img_bs device = state.device dtype = state.dtype - _fp32 = getattr(self.config, 'action_fp32', False) + _fp32 = getattr(self.config, "action_fp32", False) # embed state state_emb = self._fp32_linear(self.state_proj, state) if _fp32 else self.state_proj(state) # embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1] - time_emb = create_sinusoidal_pos_embedding( # 1, 1024 - timestep, # torch.Size([1])) - self.config.proj_width, # 1024 + time_emb = create_sinusoidal_pos_embedding( # 1, 1024 + timestep, # torch.Size([1])) + self.config.proj_width, # 1024 min_period=4e-3, max_period=4.0, device=device, @@ -1209,167 +544,36 @@ def embed_suffix(self, state, noisy_actions, timestep): # (torch.Size([state_bs, time_emb_ori = time_emb # Fuse timestep + action information using an MLP - action_emb = self._fp32_linear(self.action_in_proj, noisy_actions) if _fp32 else self.action_in_proj(noisy_actions) # torch.Size([1, state_bs*50, 1024]) - time_emb = einops.repeat(time_emb, "b d -> b n d", n=action_emb.shape[1]) # [1, 1024] -> [1, state_bs*50, 1024] - action_time_emb = torch.cat([action_emb, time_emb], dim=-1) # [1, state_bs*50, 2048] - - action_time_emb = self._fp32_linear(self.action_time_mlp_in, action_time_emb) if _fp32 else self.action_time_mlp_in(action_time_emb) + action_emb = ( + self._fp32_linear(self.action_in_proj, noisy_actions) if _fp32 else self.action_in_proj(noisy_actions) + ) # torch.Size([1, state_bs*50, 1024]) + time_emb = einops.repeat(time_emb, "b d -> b n d", n=action_emb.shape[1]) # [1, 1024] -> [1, state_bs*50, 1024] + action_time_emb = torch.cat([action_emb, time_emb], dim=-1) # [1, state_bs*50, 2048] + + action_time_emb = ( + self._fp32_linear(self.action_time_mlp_in, action_time_emb) + if _fp32 + else self.action_time_mlp_in(action_time_emb) + ) action_time_emb = F.silu(action_time_emb) # swish == silu - action_time_emb = self._fp32_linear(self.action_time_mlp_out, action_time_emb) if _fp32 else self.action_time_mlp_out(action_time_emb) # [1, state_bs*50, 1024] + action_time_emb = ( + self._fp32_linear(self.action_time_mlp_out, action_time_emb) + if _fp32 + else self.action_time_mlp_out(action_time_emb) + ) # [1, state_bs*50, 1024] action_time_dim = action_time_emb.shape[1] embs = torch.cat([state_emb[:, None], action_time_emb], dim=1) - pad_masks = torch.ones( - (bsize, action_time_dim + 1), device=device, dtype=torch.bool - ) + pad_masks = torch.ones((bsize, action_time_dim + 1), device=device, dtype=torch.bool) # Set attention masks for suffix tokens so that prefix tokens cannot attend to suffix tokens. # And state token cannot attend action tokens. # Action tokens use a bidirectional attention. - att_masks = torch.zeros( - (bsize, action_time_dim + 1), device=device, dtype=torch.bool - ) + att_masks = torch.zeros((bsize, action_time_dim + 1), device=device, dtype=torch.bool) att_masks[:, :2] = True return time_emb_ori, embs, pad_masks, att_masks - def forward(self, *args, **kwargs): - """Reject the upstream training API in the inference-only model.""" - del args, kwargs - raise RuntimeError("LingBot-VLA v2 is inference-only; use sample_actions()") - - def sample_actions( - self, images, img_masks, lang_tokens, lang_masks, state, vlm_causal=False, noise=None - ) -> Tensor: - """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)""" - bsize = state.shape[0] - device = state.device - dtype = state.dtype - - if noise is None: - actions_shape = ( - bsize, - self.config.n_action_steps, - self.config.max_action_dim, - ) - noise = torch.randn(actions_shape, device=device, dtype=dtype) - - prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( - images, img_masks, lang_tokens, lang_masks, vlm_causal - ) - prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) # bs, prefix_len, prefix_len - prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - - # Compute image and language key value cache - _, past_key_values, _ = self.qwenvl_with_expert.forward( - attention_mask=prefix_att_2d_masks, - position_ids=prefix_position_ids, - past_key_values=None, - inputs_embeds=[prefix_embs, None], - use_cache=self.config.use_cache, - fill_kv_cache=True, - ) - - dt = torch.tensor(-1.0 / self.config.num_steps, dtype=dtype, device=device) - x_t = noise - time = torch.tensor(1.0, dtype=dtype, device=device) - count = 0 - while time >= -dt / 2: - count += 1 - expanded_time = time.expand(bsize) - - v_t = self.predict_velocity( - state, prefix_pad_masks, past_key_values, x_t, expanded_time - ) - - # Euler step - x_t += dt * v_t - time += dt - print(f'Denoise {count} steps') - return x_t - - def predict_velocity(self, state, prefix_pad_masks, past_key_values, x_t, timestep): - """predict velocity at time t using the suffix model.""" - time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( - state, x_t, timestep - ) - - suffix_len = suffix_pad_masks.shape[1] - batch_size = prefix_pad_masks.shape[0] - prefix_len = prefix_pad_masks.shape[1] - prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand( - batch_size, suffix_len, prefix_len - ) - - suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) - - full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) # bs, suffix_len, prefix_len+suffix_len - - prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] - position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - - outputs_embeds, _, _ = self.qwenvl_with_expert.forward( - attention_mask=full_att_2d_masks, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=[None, suffix_embs], - use_cache=self.config.use_cache, - fill_kv_cache=False, - ada_cond = time_embs if getattr(self.config, 'adanorm_time', False) else None, - ) - suffix_out = outputs_embeds[1] - suffix_out = suffix_out[:, -self.config.n_action_steps :] - if getattr(self.config, 'action_fp32', False): - v_t = self._fp32_linear(self.action_out_proj, suffix_out) - else: - v_t = self.action_out_proj(suffix_out) - return v_t - - -# Qwen3-VL LingBot-VLA v2 policy. -FlowMatchingV1 = FlowMatching -import einops -import torch -from torch import Tensor, nn -import torch.nn.functional as F -from typing import List, Optional, Tuple, Union - -from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedModel -from transformers.models.auto import CONFIG_MAPPING -from transformers.cache_utils import Cache -from transformers.utils import logging - -from telefuser.models.lingbot_vla_v2_qwen import ( - Qwen3VLForConditionalGeneration, - Qwen3VLTextModel, - Qwen3VLPreTrainedModel, - apply_rotary_pos_emb, -) -from telefuser.models.lingbot_vla_v2_loader import ( - block_suffix_to_fv_, - create_sinusoidal_pos_embedding, - make_att_2d_masks, - our_eager_attention_forward, - prefix_query_segments, - prefix_query_token_spans, -) -from telefuser.models.lingbot_vla_v2_loader import build_block_mask, flex_attention_forward, flex_attention_with_block_mask -from telefuser.models.lingbot_vla_v2_loader import LingBotVLAWeightLoader -from telefuser.models.lingbot_vla_v2_moe import ( - Qwen2ForCausalLM, - Qwen2TokenMoeBlock, - Qwen2FusedExperts, - FixQwen2RMSNorm, -) - -try: - from dinov3.hub.backbones import dinov3_vitb16 -except Exception: - dinov3_vitb16 = None - - -logger = logging.get_logger(__name__) - class QwenvlWithExpertV2Config(PretrainedConfig): model_type = "QwenvlWithExpertV2Model" @@ -1565,14 +769,10 @@ def get_image_features( split_sizes = (image_grid_thw.prod(-1) // self.qwenvl.visual.spatial_merge_size**2).tolist() image_chunks = list(torch.split(image_embeds, split_sizes)) deepstack_chunks = [ - list(torch.split(deepstack_embeds, split_sizes)) - for deepstack_embeds in deepstack_image_embeds + list(torch.split(deepstack_embeds, split_sizes)) for deepstack_embeds in deepstack_image_embeds ] image_embeds = torch.stack(image_chunks, dim=0) - deepstack_image_embeds = [ - torch.stack(chunks, dim=0) - for chunks in deepstack_chunks - ] + deepstack_image_embeds = [torch.stack(chunks, dim=0) for chunks in deepstack_chunks] return image_embeds, deepstack_image_embeds def embed_image(self, image: torch.Tensor, image_grid_thw: torch.LongTensor): @@ -1589,8 +789,7 @@ def embed_special_token(self, token_id: int, batch: int, count: int, device, dty emb = self.embed_language_tokens(token).to(dtype=dtype) return emb.view(1, 1, 1, -1).expand(batch, count, 1, -1) - def build_prefix_position_ids(self, input_ids, attention_mask, - image_grid_thw=None, video_grid_thw=None): + def build_prefix_position_ids(self, input_ids, attention_mask, image_grid_thw=None, video_grid_thw=None): position_ids, _ = self.qwenvl.model.get_rope_index( input_ids=input_ids, image_grid_thw=image_grid_thw, @@ -1666,9 +865,7 @@ def forward( if hidden_states is None: continue if i == 1: - q, k, v = models[i].layers[layer_idx]( - hidden_states, compute_kqv=True, ada_cond=ada_cond - ) + q, k, v = models[i].layers[layer_idx](hidden_states, compute_kqv=True, ada_cond=ada_cond) else: q, k, v = models[i].layers[layer_idx](hidden_states, compute_kqv=True) query_states.append(q.float()) @@ -1721,9 +918,7 @@ def forward( if router_logits is not None: router_logits_list.append(router_logits) else: - out_emb = models[i].layers[layer_idx]( - hidden_states, att_output, start, end, output_atten=True - ) + out_emb = models[i].layers[layer_idx](hidden_states, att_output, start, end, output_atten=True) out_emb = self._apply_deepstack(out_emb, layer_idx, visual_pos_masks, deepstack_visual_embeds) outputs_embeds.append(out_emb) start = end @@ -1753,7 +948,7 @@ def get_attention_interface(self): raise ValueError(f"Invalid attention implementation: {self.config.attention_implementation}") -class FlowMatchingV2(FlowMatchingV1): +class FlowMatchingV2(FlowMatchingBase): def __init__(self, config, eval): nn.Module.__init__(self) self.config = config @@ -1822,7 +1017,6 @@ def __init__(self, config, eval): self.future_video_share_future_depth_query = False self.block_future_depth_to_action = False - def embed_prefix( self, images, @@ -1835,7 +1029,6 @@ def embed_prefix( raise ValueError("LingbotVlaV2Policy requires image_grid_thw from the Qwen3-VL image processor.") bsize = images.shape[0] device = images.device - dtype = images.dtype if images.ndim == 3: bsize = 1 num_images = images.shape[0] @@ -1857,10 +1050,7 @@ def embed_prefix( embed_dtype = img_emb.dtype num_patch = img_emb.shape[1] img_emb = einops.rearrange(img_emb, "(b n) l d -> b n l d", b=bsize, n=num_images) - deepstack_embs = [ - einops.rearrange(x, "(b n) l d -> b n l d", b=bsize, n=num_images) - for x in deepstack_embs - ] + deepstack_embs = [einops.rearrange(x, "(b n) l d -> b n l d", b=bsize, n=num_images) for x in deepstack_embs] if img_masks.ndim == 1: img_masks = img_masks.unsqueeze(0) @@ -1907,22 +1097,15 @@ def embed_prefix( lang_emb = self.qwenvl_with_expert.embed_language_tokens(lang_tokens).to(dtype=embed_dtype) if self.use_depth_align and self.align_type == "query": + def _get_align_tokens(tokens): tk_weights = tokens.view(self.num_task_tokens, tokens.shape[0] // self.num_task_tokens, tokens.shape[1]) tk_weights = tk_weights.mean(dim=1) return tk_weights - align_pad_masks = torch.ones( - bsize, - self.num_task_tokens, - device=device, - dtype=lang_masks.dtype - ) + align_pad_masks = torch.ones(bsize, self.num_task_tokens, device=device, dtype=lang_masks.dtype) fake_align_ids = torch.full( - (bsize, self.num_task_tokens), - cfg.text_config.eos_token_id, - dtype=torch.long, - device=device + (bsize, self.num_task_tokens), cfg.text_config.eos_token_id, dtype=torch.long, device=device ) current_task = _get_align_tokens(self.depth_align_embs) @@ -1932,9 +1115,7 @@ def _get_align_tokens(tokens): and getattr(self, "use_current_shared_task_proj", False) ): current_video_task = _get_align_tokens(self.current_video_align_embs) - current_task = self.current_shared_task_proj( - torch.cat([current_task, current_video_task], dim=-1) - ) + current_task = self.current_shared_task_proj(torch.cat([current_task, current_video_task], dim=-1)) align_embs = current_task.repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) parts = [img_emb] masks = [image_pad_masks] @@ -1964,9 +1145,7 @@ def _append( and getattr(self, "use_shared_future_task_proj", False) ): future_video_task = _get_align_tokens(self.future_video_align_embs) - future_task = self.future_shared_task_proj( - torch.cat([future_task, future_video_task], dim=-1) - ) + future_task = self.future_shared_task_proj(torch.cat([future_task, future_video_task], dim=-1)) future_align_embs = future_task.repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) if ( @@ -1974,9 +1153,7 @@ def _append( and getattr(self, "use_future_video", False) and getattr(self, "future_video_share_future_depth_query", False) ): - raise ValueError( - "share_future_depth_query=True requires depth.use_future_depth=True." - ) + raise ValueError("share_future_depth_query=True requires depth.use_future_depth=True.") for segment_name in prefix_query_segments( use_depth_align=True, @@ -2016,9 +1193,11 @@ def _append( ) _append(future_video_cls_align_emb, cls_align_pad_masks, fake_cls_align_ids) elif segment_name == "future_video": - future_video_align_embs = _get_align_tokens(self.future_video_align_embs).repeat( - img_emb.size(0), 1, 1 - ).to(img_emb.device, img_emb.dtype) + future_video_align_embs = ( + _get_align_tokens(self.future_video_align_embs) + .repeat(img_emb.size(0), 1, 1) + .to(img_emb.device, img_emb.dtype) + ) _append(future_video_align_embs, align_pad_masks, fake_align_ids) elif segment_name == "future_depth": _append(future_align_embs, align_pad_masks, fake_align_ids) @@ -2297,8 +1476,6 @@ def sample_actions(self, *args, **kwargs) -> Tensor: ] # __V2_END__ -from telefuser.models.lingbot_vla_v2_loader import LingBotVlaV2StateDictConverter - class LingBotVlaV2Model(LingbotVlaV2Policy): """TeleFuser-native entry point preserving official checkpoint key names.""" @@ -2312,4 +1489,5 @@ def __init__(self, config, eval=True): def state_dict_converter(**kwargs): return LingBotVlaV2StateDictConverter(**kwargs) + # __WRAPPER_END__ diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py index 7d72c0bd..8f74fbc0 100644 --- a/telefuser/models/lingbot_vla_v2_loader.py +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -10,8 +10,9 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch import Tensor from packaging.version import Version +from torch import Tensor + # from xformers.ops import memory_efficient_attention @@ -48,9 +49,7 @@ def create_sinusoidal_pos_embedding( if time.ndim != 1: raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") - fraction = torch.linspace( - 0.0, 1.0, dimension // 2, dtype=torch.float32, device=device - ) + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float32, device=device) period = min_period * (max_period / min_period) ** fraction # Compute the outer product @@ -60,8 +59,6 @@ def create_sinusoidal_pos_embedding( return pos_emb - - def make_att_2d_masks(pad_masks, att_masks): """Copied from big_vision. @@ -170,8 +167,10 @@ def fv_col_span(prefix_len, num_task_tokens, use_cls, use_patch): fv_len = (1 if use_cls else 0) + (num_task_tokens if use_patch else 0) return prefix_len - fv_len, prefix_len -def block_suffix_to_fv_(att_2d_masks, suffix_row_start, prefix_len, - num_task_tokens, use_cls=False, use_patch=True, drop_mask=None): + +def block_suffix_to_fv_( + att_2d_masks, suffix_row_start, prefix_len, num_task_tokens, use_cls=False, use_patch=True, drop_mask=None +): """In-place mask out the suffix-to-future-video attention edge. `make_att_2d_masks`' cumsum scheme cannot express "a query cannot see a @@ -201,6 +200,7 @@ def block_suffix_to_fv_(att_2d_masks, suffix_row_start, prefix_len, att_2d_masks[:, suffix_row_start:, fv_start:fv_end] = block & keep return att_2d_masks + def resize_with_pad(img, width, height, pad_value=-1): # assume no-op when width height fits already if img.ndim != 4: @@ -211,9 +211,7 @@ def resize_with_pad(img, width, height, pad_value=-1): ratio = max(cur_width / width, cur_height / height) resized_height = int(cur_height / ratio) resized_width = int(cur_width / ratio) - resized_img = F.interpolate( - img, size=(resized_height, resized_width), mode="bilinear", align_corners=False - ) + resized_img = F.interpolate(img, size=(resized_height, resized_width), mode="bilinear", align_corners=False) pad_height = max(0, int(height - resized_height)) pad_width = max(0, int(width - resized_width)) @@ -236,7 +234,8 @@ def our_eager_attention_forward( query_states: Query tensor of shape [batch_size, seq_len, num_attention_heads, head_dim]. key_states: Key tensor of shape [batch_size, seq_len, num_key_value_heads, head_dim]. value_states: Value tensor of shape [batch_size, seq_len, num_key_value_heads, head_dim]. - attention_mask: Attention mask tensor, typically [batch_size, 1, seq_len, seq_len] or [batch_size, seq_len, seq_len]. + attention_mask: Attention mask tensor, typically + [batch_size, 1, seq_len, seq_len] or [batch_size, seq_len, seq_len]. Returns: Output tensor of shape [batch_size, seq_len, num_attention_heads * head_dim]. @@ -245,33 +244,23 @@ def our_eager_attention_forward( num_key_value_heads = key_states.shape[2] num_key_value_groups = num_att_heads // num_key_value_heads - key_states = einops.repeat( - key_states, "b l h d -> b l (h g) d", g=num_key_value_groups - ) - value_states = einops.repeat( - value_states, "b l h d -> b l (h g) d", g=num_key_value_groups - ) + key_states = einops.repeat(key_states, "b l h d -> b l (h g) d", g=num_key_value_groups) + value_states = einops.repeat(value_states, "b l h d -> b l (h g) d", g=num_key_value_groups) query_states_permuted = torch.einsum("blhd->bhld", query_states) key_states_permuted = torch.einsum("blhd->bhld", key_states) - att_weights = torch.einsum( - "bhqd,bhkd->bhqk", query_states_permuted, key_states_permuted - ) + att_weights = torch.einsum("bhqd,bhkd->bhqk", query_states_permuted, key_states_permuted) att_weights *= head_dim**-0.5 big_neg = -2.3819763e38 - masked_att_weights = torch.where( - attention_mask[:, None, :, :], att_weights, big_neg - ) + masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg) probs = nn.functional.softmax(masked_att_weights, dim=-1) probs = probs.to(dtype=value_states.dtype) value_states_permuted = torch.einsum("blhd->bhld", value_states) # [B, H, L_v, D] - att_output = torch.einsum( - "bhqk,bhkv->bhqv", probs, value_states_permuted - ) # [B, H, L_q, D] + att_output = torch.einsum("bhqk,bhkv->bhqv", probs, value_states_permuted) # [B, H, L_q, D] att_output = torch.einsum("bhld->blhd", att_output) # [B, L, H, D] att_output = att_output.reshape(bsize, seq_len, num_att_heads * head_dim) @@ -286,7 +275,7 @@ def apply_rope( dtype: torch.dtype = torch.float32, ) -> torch.Tensor: """Applies RoPE positions [B, L] to x [B, L, H, D].""" - original_dtype = x.dtype # bf16 + original_dtype = x.dtype # bf16 d = x.shape[-1] d_half = d // 2 device = x.device @@ -297,19 +286,18 @@ def apply_rope( freq_exponents = (2.0 / d) * torch.arange(d_half, dtype=dtype, device=device) timescale = max_wavelength**freq_exponents - radians = torch.einsum("bl,h->blh", positions_casted, 1.0 / timescale) # fp32 -> bf16 + radians = torch.einsum("bl,h->blh", positions_casted, 1.0 / timescale) # fp32 -> bf16 radians = radians[..., None, :] # [B, L, 1, D_half] - sin = torch.sin(radians) # bf16 - cos = torch.cos(radians) # bf16 - - x1, x2 = x_casted.split(d_half, dim=-1) # fp32 + sin = torch.sin(radians) # bf16 + cos = torch.cos(radians) # bf16 - res = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) # fp32 + x1, x2 = x_casted.split(d_half, dim=-1) # fp32 - return res.to(original_dtype) # bf16 + res = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) # fp32 + return res.to(original_dtype) # bf16 # Copyright 2024 The HuggingFace Inc. team. All rights reserved. @@ -327,9 +315,6 @@ def apply_rope( # limitations under the License. import torch -import torch.nn.functional as F # noqa: N812 -from packaging.version import Version -import einops FLEX_SPARSE_BLOCK_SIZE = 128 FLEX_KERNEL_OPTIONS = {"BLOCK_M": 32, "BLOCK_N": 64, "num_warps": 4, "num_stages": 2} @@ -344,6 +329,7 @@ def apply_rope( flex_attention, ) + # @torch.compile(dynamic=False) def flex_attention_forward( query_states: torch.Tensor, @@ -357,9 +343,6 @@ def flex_attention_forward( """ batch_size, seq_len, num_att_heads, head_dim = query_states.shape original_dtype = query_states.dtype - num_key_value_heads = key_states.shape[2] - # num_key_value_groups = num_att_heads // num_key_value_heads # 16 // 2 = 8 - query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) @@ -456,7 +439,6 @@ def build_block_mask( This allocates the dense 4D mask once; the returned BlockMask can be reused across layers. """ from torch.nn.attention.flex_attention import ( - _mask_mod_signature, _round_up_to_multiple, create_block_mask, create_mask, @@ -475,19 +457,24 @@ def build_block_mask( def precomputed_mask_factory(precomputed_mask: torch.Tensor): def mask_mod(b, h, q_idx, kv_idx): return precomputed_mask[b][h][q_idx][kv_idx] + return mask_mod mask_4d = create_mask( mod_fn=precomputed_mask_factory(padded_mask), - B=b_mask, H=h_mask, - Q_LEN=q_len_rounded, KV_LEN=kv_len_rounded, + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, device=causal_mask.device, ) block_mask = create_block_mask( mask_mod=precomputed_mask_factory(mask_4d), - B=b_mask, H=h_mask, - Q_LEN=q_len_rounded, KV_LEN=kv_len_rounded, + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, BLOCK_SIZE=block_size, device=causal_mask.device, _compile=False, @@ -507,7 +494,6 @@ def flex_attention_with_block_mask( Run flex_attention with a pre-built BlockMask (no create_mask allocation per call). """ batch_size = query_states.shape[0] - num_att_heads = query_states.shape[2] head_dim = query_states.shape[3] original_dtype = query_states.dtype @@ -515,8 +501,8 @@ def flex_attention_with_block_mask( key_states = key_states.transpose(1, 2).to(torch.float32) value_states = value_states.transpose(1, 2).to(torch.float32) - q_len_rounded = block_mask.shape[-2] if hasattr(block_mask, 'shape') else query_states.shape[2] - kv_len_rounded = block_mask.shape[-1] if hasattr(block_mask, 'shape') else key_states.shape[2] + q_len_rounded = block_mask.shape[-2] if hasattr(block_mask, "shape") else query_states.shape[2] + kv_len_rounded = block_mask.shape[-1] if hasattr(block_mask, "shape") else key_states.shape[2] pad_q = q_len_rounded - query_states.shape[2] pad_k = kv_len_rounded - key_states.shape[2] @@ -543,12 +529,8 @@ def flex_attention_with_block_mask( return attn_output - # modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py -import math import torch -import torch.nn as nn -import torch.nn.functional as F # FFN @@ -564,7 +546,7 @@ def FeedForward(dim, mult=4): def reshape_tensor(x, heads): bs, length, width = x.shape - #(bs, length, width) --> (bs, length, n_heads, dim_per_head) + # (bs, length, width) --> (bs, length, n_heads, dim_per_head) x = x.view(bs, length, heads, -1) # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) x = x.transpose(1, 2) @@ -574,7 +556,6 @@ def reshape_tensor(x, heads): class PerceiverAttention(nn.Module): - def __init__(self, *, dim, dim_head=64, heads=8): super().__init__() self.scale = dim_head**-0.5 @@ -600,7 +581,7 @@ def forward(self, x, latents): x = self.norm1(x) latents = self.norm2(latents) - b, l, _ = latents.shape + batch_size, latent_length, _ = latents.shape q = self.to_q(latents) kv_input = torch.cat((x, latents), dim=-2) @@ -616,54 +597,12 @@ def forward(self, x, latents): weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) out = weight @ v - out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + out = out.permute(0, 2, 1, 3).reshape(batch_size, latent_length, -1) return self.to_out(out) -class AttentionPool2d(nn.Module): - - def __init__(self, seq_len: int, embed_dim: int, num_heads: int, output_dim: int = None): - super().__init__() - self.positional_embedding = nn.Parameter(torch.randn(seq_len + 1, embed_dim) / embed_dim**0.5) - self.k_proj = nn.Linear(embed_dim, embed_dim) - self.q_proj = nn.Linear(embed_dim, embed_dim) - self.v_proj = nn.Linear(embed_dim, embed_dim) - self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) - self.num_heads = num_heads - - def forward(self, x, return_all_tokens=False): - # x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1) # NCHW -> (HW)NC - x = x.permute(1, 0, 2) # (N(HW)C) => (HW)NC - x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC - x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC - x, _ = F.multi_head_attention_forward(query=x, - key=x, - value=x, - embed_dim_to_check=x.shape[-1], - num_heads=self.num_heads, - q_proj_weight=self.q_proj.weight, - k_proj_weight=self.k_proj.weight, - v_proj_weight=self.v_proj.weight, - in_proj_weight=None, - in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), - bias_k=None, - bias_v=None, - add_zero_attn=False, - dropout_p=0, - out_proj_weight=self.c_proj.weight, - out_proj_bias=self.c_proj.bias, - use_separate_proj_weight=True, - training=self.training, - need_weights=False) - if return_all_tokens: - return x - else: - return x[0] - - -class Resampler(nn.Module): - +class TaskTokenResampler(nn.Module): def __init__( self, dim_in=768, @@ -677,9 +616,9 @@ def __init__( ): super().__init__() - self.queries = nn.Parameter(torch.randn(1, num_queries, dim_in) / dim_mid ** 0.5) - - self.proj_in = nn.Linear(dim_in, dim_mid) + self.num_queries = num_queries + self.proj_in1 = nn.Linear(dim_in, dim_mid) + self.proj_in2 = nn.Linear(dim_in, dim_mid) self.proj_out = nn.Linear(dim_mid, dim_out) self.norm_out = nn.LayerNorm(dim_out) @@ -694,47 +633,6 @@ def __init__( ) ) - def forward(self, x): - queries = self.queries.repeat(x.size(0), 1, 1) - x = self.proj_in(x) - - for attn, ff in self.layers: - queries = attn(x, queries) + queries - queries = ff(queries) + queries - - queries = self.proj_out(queries) - queries = self.norm_out(queries) - return queries - -class TaskTokenResampler(nn.Module): - - def __init__( - self, - dim_in=768, - dim_mid=1024, - dim_head=64, - dim_out=1024, - num_layers=8, - num_queries=8, - num_heads=16, - ff_mult=4, - ): - super().__init__() - - self.num_queries = num_queries - self.proj_in1 = nn.Linear(dim_in, dim_mid) - self.proj_in2 = nn.Linear(dim_in, dim_mid) - self.proj_out = nn.Linear(dim_mid, dim_out) - self.norm_out = nn.LayerNorm(dim_out) - - self.layers = nn.ModuleList([]) - for _ in range(num_layers): - self.layers.append( - nn.ModuleList([ - PerceiverAttention(dim=dim_mid, dim_head=dim_head, heads=num_heads), - FeedForward(dim=dim_mid, mult=ff_mult), - ])) - def forward(self, x, queries): queries = self.proj_in1(queries) x = self.proj_in2(x) @@ -748,191 +646,6 @@ def forward(self, x, queries): return queries -class ResamplerXL(nn.Module): - - def __init__( - self, - dim=1024, - depth=8, - dim_head=64, - heads=16, - num_queries=8, - embedding_dim=768, - output1_dim=768, - output2_dim=1280, - ff_mult=4, - ): - super().__init__() - - self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) - - self.proj_in = nn.Linear(embedding_dim, dim) - - # self.proj_out = nn.Linear(dim, output_dim) - self.norm_out = nn.LayerNorm(dim) - - self.in_dim = dim - self.out_dim = output1_dim + output2_dim - - self.layers = nn.ModuleList([]) - for _ in range(depth): - self.layers.append( - nn.ModuleList([ - PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), - FeedForward(dim=dim, mult=ff_mult), - ])) - - self.unet_proj_1 = nn.Linear(self.in_dim, output1_dim) - self.unet_proj_2 = nn.Linear(self.in_dim, output2_dim) - self.unet_attnpool = AttentionPool2d(num_queries, self.in_dim, heads, output2_dim) - - def forward(self, x): - - latents = self.latents.repeat(x.size(0), 1, 1) - - x = self.proj_in(x) - - for attn, ff in self.layers: - latents = attn(x, latents) + latents - latents = ff(latents) + latents - - hidden_embeds = self.norm_out(latents) - - encoder_hidden_1 = self.unet_proj_1(hidden_embeds) # [bs, 256, 768] - encoder_hidden_2 = self.unet_proj_2(hidden_embeds) # [bs, 256, 1280] - prompt_embeds = torch.cat([encoder_hidden_1, encoder_hidden_2], dim=-1) # [bs, 256, 2048] - pooled_prompt_embeds = self.unet_attnpool(hidden_embeds) # [bs, 1280] - - return prompt_embeds, pooled_prompt_embeds - - -class ResamplerXLV2(nn.Module): - - def __init__( - self, - dim=1024, - depth=8, - dim_head=64, - heads=16, - num_queries=8, - embedding_dim=768, - output1_dim=768, - output2_dim=1280, - ff_mult=4, - normalize=True - ): - super().__init__() - - self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) - - self.normalize = normalize - self.proj_in = nn.Linear(embedding_dim, dim) - - # self.proj_out = nn.Linear(dim, output_dim) - self.norm_out = nn.LayerNorm(dim) - - self.in_dim = dim - self.out_dim = output1_dim + output2_dim - - self.layers = nn.ModuleList([]) - for _ in range(depth): - self.layers.append( - nn.ModuleList([ - PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), - FeedForward(dim=dim, mult=ff_mult), - ])) - - self.unet_proj_1 = nn.Linear(self.in_dim, output1_dim) - self.unet_proj_2 = nn.Linear(self.in_dim, output2_dim) - self.unet_attnpool = AttentionPool2d(num_queries, self.in_dim, heads, output2_dim) - - def forward(self, x,pooled_text_embeds=None): - - latents = self.latents.repeat(x.size(0), 1, 1) - - if self.normalize: - x = F.normalize(x) - - x = self.proj_in(x) - - for attn, ff in self.layers: - latents = attn(x, latents) + latents - latents = ff(latents) + latents - - hidden_embeds = self.norm_out(latents) - - encoder_hidden_1 = self.unet_proj_1(hidden_embeds) # [bs, 256, 768] - encoder_hidden_2 = self.unet_proj_2(hidden_embeds) # [bs, 256, 1280] - prompt_embeds = torch.cat([encoder_hidden_1, encoder_hidden_2], dim=-1) # [bs, 256, 2048] - pooled_prompt_embeds = self.unet_attnpool(hidden_embeds) # [bs, 1280] - - return prompt_embeds, pooled_prompt_embeds - -class ResamplerXLIdentity(nn.Module): - def __init__(self) -> None: - super().__init__() - - def forward(self, x, pooled_text_embeds=None): - return x, pooled_text_embeds - - -if __name__ == '__main__': - image_proj_model = Resampler(dim=1024, - depth=4, - dim_head=64, - heads=12, - num_queries=1024, - embedding_dim=1024, - output_dim=1024, - ff_mult=4) - numel = 0 - for name, param in image_proj_model.named_parameters(): - numel += param.numel() - - print(f'Total params: {numel}') - - - -import torch.nn as nn - -def build_mlp(in_hidden_size, hidden_size): - modules = [nn.Linear(in_hidden_size, hidden_size)] - modules.append(nn.ReLU()) - modules.append(nn.Linear(hidden_size, hidden_size)) - return nn.Sequential(*modules) - -def build_expand_mlp(in_hidden_size, hidden_size, out_size): - modules = [nn.Linear(in_hidden_size, hidden_size)] - modules.append(nn.ReLU()) - modules.append(nn.Linear(hidden_size, hidden_size)) - modules.append(nn.ReLU()) - modules.append(nn.Linear(hidden_size, out_size)) - return nn.Sequential(*modules) - -class DepthHead(nn.Module): - def __init__( - self, - proj_config=None, - llm_hidden_size=4096, - use_intermediate_depth=False, - ): - super(DepthHead, self).__init__() - - self.projector = Resampler( - dim_in=llm_hidden_size, - dim_mid=llm_hidden_size, - dim_head=proj_config["dim_head"], - dim_out=proj_config["dim_out"], - num_layers=proj_config["num_layers"], - num_heads=proj_config["num_heads"], - num_queries=proj_config["num_backbone_tokens"], - ff_mult=proj_config["ff_mult"], - ) - - def forward(self, llm_feats): - queries = self.projector(llm_feats) - return queries - class TaskTokenDepthHead(nn.Module): def __init__( self, @@ -954,8 +667,10 @@ def __init__( ) def forward(self, llm_feats, queries): - queries = self.projector(llm_feats, queries) - return queries + queries = self.projector(llm_feats, queries) + return queries + + import json from copy import deepcopy from pathlib import Path @@ -964,7 +679,6 @@ def forward(self, llm_feats, queries): from transformers import AutoConfig - class LingBotVLAWeightLoader: """Minimal native weight-name mapper retained for model compatibility.""" @@ -1168,9 +882,7 @@ def validate_official_6b_checkpoint(state_dict): if key not in state_dict: raise ValueError(f"Missing official LingBot-VLA v2 weight: {key}") if tuple(state_dict[key].shape) != expected: - raise ValueError( - f"Unexpected shape for {key}: expected {expected}, got {tuple(state_dict[key].shape)}" - ) + raise ValueError(f"Unexpected shape for {key}: expected {expected}, got {tuple(state_dict[key].shape)}") class LingBotVlaV2StateDictConverter: diff --git a/telefuser/models/lingbot_vla_v2_moe.py b/telefuser/models/lingbot_vla_v2_moe.py index ec3588d2..b8c8e91e 100644 --- a/telefuser/models/lingbot_vla_v2_moe.py +++ b/telefuser/models/lingbot_vla_v2_moe.py @@ -37,47 +37,23 @@ def fused_moe_forward( return output +from typing import Optional, Tuple -from logging import raiseExceptions -import einops -import numpy as np -import torch -from torch import nn import torch.nn.functional as F -from torch import Tensor, nn -from typing import List, Optional, Tuple -from transformers import AutoTokenizer -from dataclasses import dataclass -from transformers.models.qwen2.configuration_qwen2 import Qwen2Config -from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.cache_utils import Cache, SlidingWindowCache, StaticCache, DynamicCache -from transformers.generation import GenerationMixin -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from transformers.utils import ( - ModelOutput, - is_torchdynamo_compiling, - logging, - can_return_tuple, - auto_docstring -) -from transformers.utils.deprecation import deprecate_kwarg +from torch import nn from transformers.activations import ACT2FN -from transformers.modeling_attn_mask_utils import AttentionMaskConverter -from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, is_flash_attn_available -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.generation import GenerationMixin +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config from transformers.processing_utils import Unpack +from transformers.utils import auto_docstring, logging + -try: - from dinov3.hub.backbones import ( - dinov3_vits16, - dinov3_vits16plus, - dinov3_vitb16, - ) -except: pass def _update_moe_runtime_stats(block, routing_weights, selected_experts): """Update MoE runtime buffers outside torch.compile graphs.""" with torch.no_grad(): - if routing_weights is not None and hasattr(block, 'avg_topk_sigmoid_score'): + if routing_weights is not None and hasattr(block, "avg_topk_sigmoid_score"): avg_score = routing_weights.detach().float().mean() block.avg_topk_sigmoid_score.copy_( avg_score.reshape_as(block.avg_topk_sigmoid_score).to( @@ -86,7 +62,7 @@ def _update_moe_runtime_stats(block, routing_weights, selected_experts): ) ) - if hasattr(block, 'tokens_per_expert'): + if hasattr(block, "tokens_per_expert"): counts = F.one_hot( selected_experts.detach().reshape(-1), num_classes=block.num_experts, @@ -99,26 +75,24 @@ def _update_moe_runtime_stats(block, routing_weights, selected_experts): ) -import transformers.models.qwen2.modeling_qwen2 as hf_qwen2 from transformers.models.qwen2.modeling_qwen2 import ( - Qwen2MLP, - rotate_half, - apply_rotary_pos_emb, - repeat_kv, - eager_attention_forward, + PreTrainedModel, Qwen2Attention, + Qwen2MLP, Qwen2RMSNorm, Qwen2RotaryEmbedding, - PreTrainedModel, ) - from transformers.models.qwen2.modeling_qwen2 import ( - Qwen2Model as _Qwen2Model, Qwen2ForCausalLM as _Qwen2ForCausalLM, ) +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2Model as _Qwen2Model, +) + logger = logging.get_logger(__name__) # from transformers.models.mistral.modeling_mistral import MistralMLP + # Modified from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2Moe class Qwen2MoeRoutedExpertMLP(nn.Module): def __init__(self, config, intermediate_size=None): @@ -162,6 +136,7 @@ class Qwen2FusedExperts(nn.Module): for FSDP2: calling self.experts(...) triggers FSDP2's forward pre-hook to unshard the expert params on ep_fsdp_mesh BEFORE they are used by kernels. """ + def __init__(self, num_experts, hidden_size, intermediate_size, initializer_range=0.02): super().__init__() self.num_experts = num_experts @@ -260,6 +235,7 @@ def extra_repr(self): class Qwen2TokenMoeBlock(nn.Module): """Token-level routing MoE block with all-to-all computation for torch.compile compatibility.""" + def __init__(self, config): super().__init__() self.num_experts = config.num_experts @@ -270,19 +246,23 @@ def __init__(self, config): # equivalent to unbiased top-k selection; the optimizer pre-hook updates # the bias when bias_update_speed > 0. self.register_buffer( - "e_score_correction_bias", torch.zeros(config.num_experts), + "e_score_correction_bias", + torch.zeros(config.num_experts), persistent=True, ) self.register_buffer( - "tokens_per_expert", torch.zeros(config.num_experts, dtype=torch.float32), + "tokens_per_expert", + torch.zeros(config.num_experts, dtype=torch.float32), persistent=False, ) self.register_buffer( - "last_tokens_per_expert", torch.zeros(config.num_experts, dtype=torch.float32), + "last_tokens_per_expert", + torch.zeros(config.num_experts, dtype=torch.float32), persistent=False, ) self.register_buffer( - "avg_topk_sigmoid_score", torch.zeros(1, dtype=torch.float32), + "avg_topk_sigmoid_score", + torch.zeros(1, dtype=torch.float32), persistent=False, ) @@ -290,9 +270,9 @@ def __init__(self, config): self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) # EP/fused support: choose expert storage based on moe_implementation - self._moe_implementation = getattr(config, '_moe_implementation', None) or 'eager' + self._moe_implementation = getattr(config, "_moe_implementation", None) or "eager" self._use_robby_moe_kernel = bool(getattr(config, "use_robby_moe_kernel", False)) - if self._moe_implementation == 'fused': + if self._moe_implementation == "fused": self.experts = Qwen2FusedExperts( self.num_experts, config.hidden_size, @@ -301,20 +281,22 @@ def __init__(self, config): ) else: self.experts = nn.ModuleList( - [Qwen2MoeRoutedExpertMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)] + [ + Qwen2MoeRoutedExpertMLP(config, intermediate_size=config.moe_intermediate_size) + for _ in range(self.num_experts) + ] ) self.shared_expert = Qwen2MoeSharedExpertMLP(config, intermediate_size=config.shared_expert_intermediate_size) - self._router_activation = getattr(config, 'router_activation', 'softmax') - self.routed_scaling_factor = getattr(config, 'routed_scaling_factor', 1.0) - self._use_shared_expert_gate = getattr(config, 'use_shared_expert_gate', True) + self._router_activation = getattr(config, "router_activation", "softmax") + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self._use_shared_expert_gate = getattr(config, "use_shared_expert_gate", True) if self._use_shared_expert_gate: self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Token-level routing with all-to-all computation for torch.compile compatibility.""" batch_size, sequence_length, hidden_dim = hidden_states.shape - num_tokens = batch_size * sequence_length # Token-level routing: each token individually hidden_flat = hidden_states.reshape(-1, hidden_dim) # (B*T, D) @@ -324,7 +306,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: with torch.amp.autocast(hidden_flat.device.type, enabled=False): router_logits = F.linear(hidden_flat.float(), self.gate.weight.float()) # (B*T, num_experts) - if self._router_activation == 'sigmoid': + if self._router_activation == "sigmoid": routing_scores = router_logits.sigmoid() else: routing_scores = F.softmax(router_logits, dim=1, dtype=torch.float) @@ -341,7 +323,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: routing_weights = routing_weights.to(hidden_states.dtype) # Expert computation: fused (group_gemm) or eager (per-expert loop) - if self._moe_implementation == 'fused': + if self._moe_implementation == "fused": use_robby_moe = ( self._use_robby_moe_kernel and robby_moe_forward is not None @@ -385,11 +367,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: expert_outputs = torch.stack( [expert(hidden_flat) for expert in self.experts], dim=0 ) # (num_experts, B*T, D) - expert_mask = F.one_hot( - selected_experts, num_classes=self.num_experts - ).float() # (B*T, top_k, num_experts) - weights = (expert_mask * routing_weights.unsqueeze(-1).float()).sum(dim=1).to(hidden_states.dtype) # (B*T, num_experts) - final_hidden_states = torch.einsum('ebd,be->bd', expert_outputs, weights) # (B*T, D) + expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts).float() # (B*T, top_k, num_experts) + weights = ( + (expert_mask * routing_weights.unsqueeze(-1).float()).sum(dim=1).to(hidden_states.dtype) + ) # (B*T, num_experts) + final_hidden_states = torch.einsum("ebd,be->bd", expert_outputs, weights) # (B*T, D) # Shared expert: applied to all tokens (fixed shape) if final_hidden_states.dtype != hidden_flat.dtype: @@ -474,7 +456,10 @@ def forward( return out_emb, router_logits else: - raise ValueError(f"Invaild Operation compute_kqv={compute_kqv} and output_atten={output_atten} with Qwen2DecoderLayer in LingBot-VLA") + raise ValueError( + f"Invalid operation compute_kqv={compute_kqv} and output_atten={output_atten} " + "with Qwen2DecoderLayer in LingBot-VLA" + ) @auto_docstring @@ -517,6 +502,7 @@ class Qwen2Model(Qwen2PreTrainedModel): Args: config: Qwen2Config """ + get_input_embeddings = _Qwen2Model.get_input_embeddings set_input_embeddings = _Qwen2Model.set_input_embeddings forward = _Qwen2Model.forward @@ -551,6 +537,7 @@ class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): forward = _Qwen2ForCausalLM.forward set_decoder = _Qwen2ForCausalLM.set_decoder get_decoder = _Qwen2ForCausalLM.get_decoder + def __init__(self, config, eval): super().__init__(config) self.model = Qwen2Model(config, eval) @@ -559,9 +546,3 @@ def __init__(self, config, eval): # Initialize weights and apply final processing self.post_init() - -def apply_lingbot_qwen2_patch(): - hf_qwen2.Qwen2DecoderLayer = Qwen2DecoderLayer - hf_qwen2.Qwen2PreTrainedModel = Qwen2PreTrainedModel - hf_qwen2.Qwen2Model = Qwen2Model - hf_qwen2.Qwen2ForCausalLM = Qwen2ForCausalLM diff --git a/telefuser/models/lingbot_vla_v2_qwen.py b/telefuser/models/lingbot_vla_v2_qwen.py index d2f0ec97..47cc54cb 100644 --- a/telefuser/models/lingbot_vla_v2_qwen.py +++ b/telefuser/models/lingbot_vla_v2_qwen.py @@ -3,446 +3,44 @@ Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. """ -# Qwen2.5-VL compatibility is retained because the shared flow-matching base -# defines the legacy V1 policy alongside the V2 policy. - -import torch -from torch import nn -import torch.nn.functional as F -from torch.nn import CrossEntropyLoss -from torch import Tensor, nn -from typing import List, Optional, Tuple, Union, Callable, Dict, Any -import math -from transformers import ( - PreTrainedModel, -) -from dataclasses import dataclass -from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig -from transformers.cache_utils import Cache, SlidingWindowCache, StaticCache, DynamicCache -from transformers.generation import GenerationMixin -from transformers.modeling_outputs import ( - BaseModelOutputWithPast, -) -from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS -from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.utils import ( - ModelOutput, - logging, -) -from transformers.activations import ACT2FN -from transformers.modeling_attn_mask_utils import AttentionMaskConverter -from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, flash_attn_supports_top_left_mask, is_flash_attn_available -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from transformers.processing_utils import Unpack -import torch.distributed._tensor as dt - -if is_flash_attn_available(): - from flash_attn.layers.rotary import apply_rotary_emb - from flash_attn.flash_attn_interface import flash_attn_varlen_func - from transformers.modeling_flash_attention_utils import _flash_attention_forward -import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as hf_qwen25vl -from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( - Qwen2RMSNorm, - Qwen2_5_VLMLP, - Qwen2_5_VLAttention, - Qwen2MLP, - Qwen2_5_VisionTransformerPretrainedModel, - Qwen2_5_VLRotaryEmbedding, - apply_rotary_pos_emb_vision, - eager_attention_forward -) - -from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( - Qwen2_5_VLTextModel as _Qwen2_5_VLTextModel, - Qwen2_5_VLForConditionalGeneration as _Qwen2_5_VLForConditionalGeneration -) -logger = logging.get_logger(__name__) - - -class Qwen2_5_VLVisionAttention(nn.Module): - def __init__(self, config: Qwen2_5_VLVisionConfig) -> None: - super().__init__() - self.dim = config.hidden_size - self.num_heads = config.num_heads - self.head_dim = self.dim // self.num_heads - self.num_key_value_groups = 1 # needed for eager attention - self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) - self.proj = nn.Linear(self.dim, self.dim) - self.scaling = self.head_dim**-0.5 - self.config = config - self.attention_dropout = 0.0 - self.is_causal = False - # print(f"ViT Attention Type is {self.config._attn_implementation}") - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor, - rotary_pos_emb: Optional[torch.Tensor] = None, - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - **kwargs, - ) -> torch.Tensor: - seq_length = hidden_states.shape[0] - query_states, key_states, value_states = ( - self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) - ) - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) - - query_states = query_states.transpose(0, 1).unsqueeze(0) - key_states = key_states.transpose(0, 1).unsqueeze(0) - value_states = value_states.transpose(0, 1).unsqueeze(0) - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - if self.config._attn_implementation == "flash_attention_2": - # Flash Attention 2: Use cu_seqlens for variable length attention - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - out_fp32_atten = False - if key_states.dtype == torch.float32: - out_fp32_atten = True - query_states, key_states, value_states = query_states.to(torch.bfloat16), key_states.to(torch.bfloat16), value_states.to(torch.bfloat16) - attn_output, _ = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask=None, - scaling=self.scaling, - dropout=0.0 if not self.training else self.attention_dropout, - cu_seq_lens_q=cu_seqlens, - cu_seq_lens_k=cu_seqlens, - max_length_q=max_seqlen, - max_length_k=max_seqlen, - is_causal=False, - **kwargs, - ) - if out_fp32_atten: - attn_output = attn_output.to(torch.float32) - else: - # Other implementations: Process each chunk separately - lengths = cu_seqlens[1:] - cu_seqlens[:-1] - splits = [ - torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) - ] - - attn_outputs = [ - attention_interface( - self, - q, - k, - v, - attention_mask=None, - scaling=self.scaling, - dropout=0.0 if not self.training else self.attention_dropout, - is_causal=False, - **kwargs, - )[0] - for q, k, v in zip(*splits) - ] - attn_output = torch.cat(attn_outputs, dim=1) - - attn_output = attn_output.reshape(seq_length, -1).contiguous() - attn_output = self.proj(attn_output) - return attn_output - - -class Qwen2_5_VLVisionBlock(GradientCheckpointingLayer): - def __init__(self, config, attn_implementation: str = "flash_attention_2") -> None: - super().__init__() - self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) - self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) - self.attn = Qwen2_5_VLVisionAttention(config=config) - self.mlp = Qwen2_5_VLMLP(config, bias=True) - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor, - rotary_pos_emb: Optional[torch.Tensor] = None, - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - **kwargs, - ) -> torch.Tensor: - hidden_states = hidden_states + self.attn( - self.norm1(hidden_states), - cu_seqlens=cu_seqlens, - rotary_pos_emb=rotary_pos_emb, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) - return hidden_states - - -class Qwen2_5_VLPreTrainedModel(PreTrainedModel): - config_class = Qwen2_5_VLConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] - _skip_keys_device_placement = "past_key_values" - _supports_flash_attn_2 = True - _supports_sdpa = True - _supports_cache_class = True - _supports_static_cache = False # TODO (joao): fix. torch.compile failing probably due to `cache_positions` - - # def _init_weights(self, module): - # std = self.config.initializer_range - # if isinstance(module, (nn.Linear, nn.Conv3d)): - # module.weight.data.normal_(mean=0.0, std=std) - # if module.bias is not None: - # module.bias.data.zero_() - # elif isinstance(module, nn.Embedding): - # module.weight.data.normal_(mean=0.0, std=std) - # if module.padding_idx is not None: - # module.weight.data[module.padding_idx].zero_() - - -class Qwen2_5_VLDecoderLayer(GradientCheckpointingLayer): - def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - - if config.use_sliding_window and config._attn_implementation != "flash_attention_2": - logger.warning_once( - f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " - "unexpected results may be encountered." - ) - self.self_attn = Qwen2_5_VLAttention(config, layer_idx) - - self.mlp = Qwen2MLP(config) - self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - if config.norm_qkv: - self.q_layernorm = Qwen2RMSNorm(self.self_attn.head_dim, eps=config.rms_norm_eps) - self.k_layernorm = Qwen2RMSNorm(self.self_attn.head_dim, eps=config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - att_output: Optional[torch.Tensor] = None, - start: Optional[int] = 0, - end: Optional[int] = 0, - compute_kqv: bool = False, - norm_qkv: bool = False, - output_atten: bool = False, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: - """ - Args: - hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` - attention_mask (`torch.FloatTensor`, *optional*): attention mask of size - `(batch, sequence_length)` where padding elements are indicated by 0. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding - (see `past_key_values`). - past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states - cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): - Indices depicting the position of the input sequence tokens in the sequence. - position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): - Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, - with `head_dim` being the embedding dimension of each attention head. - kwargs (`dict`, *optional*): - Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code - into the model - """ - - if compute_kqv: - hidden_states = self.input_layernorm(hidden_states) - hidden_shape = (*hidden_states.shape[:-1], -1, self.self_attn.head_dim) - - query_state = self.self_attn.q_proj(hidden_states).view(hidden_shape) - key_state = self.self_attn.k_proj(hidden_states).view(hidden_shape) - value_state = self.self_attn.v_proj(hidden_states).view(hidden_shape) - - if norm_qkv: - query_state = self.q_layernorm(query_state) - key_state = self.k_layernorm(key_state) - - return query_state, key_state, value_state - - elif output_atten: - if att_output.dtype != self.self_attn.o_proj.weight.dtype: - att_output = att_output.to(self.self_attn.o_proj.weight.dtype) - out_emb = self.self_attn.o_proj(att_output[:, start:end]) - - # first residual - out_emb += hidden_states - after_first_residual = out_emb.clone() - - out_emb = self.post_attention_layernorm(out_emb) - out_emb = self.mlp(out_emb) - - # second residual - out_emb += after_first_residual - - return out_emb - - else: - raise ValueError(f"Invaild Operation compute_kqv={compute_kqv} and output_atten={output_atten} with Qwen2_5_VLDecoderLayer in LingBot-VLA") - - -class Qwen2_5_VLTextModel(Qwen2_5_VLPreTrainedModel): - get_input_embeddings = _Qwen2_5_VLTextModel.get_input_embeddings - set_input_embeddings = _Qwen2_5_VLTextModel.set_input_embeddings - forward = _Qwen2_5_VLTextModel.forward - - def __init__(self, config: Qwen2_5_VLConfig): - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList( - [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - self._attn_implementation = config._attn_implementation - self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) - - self.gradient_checkpointing = False - # Initialize weights and apply final processing - self._init_weights = lambda module: None - self.post_init() - - -class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] - config_class = Qwen2_5_VLConfig - _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] - get_input_embeddings = _Qwen2_5_VLForConditionalGeneration.get_input_embeddings - set_input_embeddings = _Qwen2_5_VLForConditionalGeneration.set_input_embeddings - get_output_embeddings = _Qwen2_5_VLForConditionalGeneration.get_output_embeddings - set_output_embeddings = _Qwen2_5_VLForConditionalGeneration.set_output_embeddings - get_decoder = _Qwen2_5_VLForConditionalGeneration.get_decoder - set_decoder = _Qwen2_5_VLForConditionalGeneration.set_decoder - forward = _Qwen2_5_VLForConditionalGeneration.forward - prepare_inputs_for_generation = _Qwen2_5_VLForConditionalGeneration.prepare_inputs_for_generation - def __init__(self, config): - super().__init__(config) - self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) - self.model = Qwen2_5_VLTextModel._from_config(config) - self.vocab_size = config.vocab_size - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - self.rope_deltas = None # cache rope_deltas here - - # Initialize weights and apply final processing - self.post_init() - - - -def qwen25_preprcess_grid_thw(self, grid_thw: torch.Tensor): - rotary_pos_emb = self.rot_pos_emb(grid_thw) - window_index, cu_window_seqlens = self.get_window_index(grid_thw) - cu_window_seqlens = torch.tensor( - cu_window_seqlens, - device=grid_thw.device, - dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, - ) - cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) - - cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( - dim=0, - dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, - ) - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) - - return rotary_pos_emb, window_index, cu_window_seqlens, cu_seqlens - - -def qwen25_forward_without_grid_thw( - self, - hidden_states: torch.Tensor, - grid_thw: torch.Tensor = None, - rotary_pos_emb = None, - window_index = None, - cu_window_seqlens = None, - cu_seqlens = None, - **kwargs -) -> torch.Tensor: - hidden_states = self.patch_embed(hidden_states) - - if rotary_pos_emb is None or window_index is None or cu_window_seqlens is None or cu_seqlens is None: - rotary_pos_emb, window_index, cu_window_seqlens, cu_seqlens = self.preprcess_grid_thw(grid_thw) - - seq_len, _ = hidden_states.size() - hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) - hidden_states = hidden_states[window_index, :, :] - hidden_states = hidden_states.reshape(seq_len, -1) - rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) - rotary_pos_emb = rotary_pos_emb[window_index, :, :] - rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) - emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - position_embeddings = (emb.cos(), emb.sin()) - - for layer_num, blk in enumerate(self.blocks): - if layer_num in self.fullatt_block_indexes: - cu_seqlens_now = cu_seqlens - else: - cu_seqlens_now = cu_window_seqlens - - hidden_states = blk( - hidden_states, - cu_seqlens=cu_seqlens_now, - position_embeddings=position_embeddings, - **kwargs, - ) - - hidden_states = self.merger(hidden_states) - reverse_indices = torch.argsort(window_index) - hidden_states = hidden_states[reverse_indices, :] - - return hidden_states - - -def apply_lingbot_qwen25_vl_patch(): - logger.info("apply Qwen2.5-VL LingBot patch") - hf_qwen25vl.Qwen2_5_VLPreTrainedModel = Qwen2_5_VLPreTrainedModel - hf_qwen25vl.Qwen2_5_VLDecoderLayer = Qwen2_5_VLDecoderLayer - hf_qwen25vl.Qwen2_5_VLTextModel = Qwen2_5_VLTextModel - hf_qwen25vl.Qwen2_5_VLForConditionalGeneration = Qwen2_5_VLForConditionalGeneration - hf_qwen25vl.Qwen2_5_VLVisionAttention = Qwen2_5_VLVisionAttention - hf_qwen25vl.Qwen2_5_VLVisionBlock = Qwen2_5_VLVisionBlock - hf_qwen25vl.Qwen2_5_VisionTransformerPretrainedModel.forward = qwen25_forward_without_grid_thw - hf_qwen25vl.Qwen2_5_VisionTransformerPretrainedModel.preprcess_grid_thw = qwen25_preprcess_grid_thw - # Qwen3-VL implementation used by LingBot-VLA v2. -import torch -from torch import nn -import torch.nn.functional as F from types import MethodType from typing import Callable, Optional, Tuple +import torch +import torch.nn.functional as F +from torch import nn from transformers.generation import GenerationMixin -from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel -from transformers.processing_utils import Unpack -from transformers.utils import logging from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig from transformers.models.qwen3_vl.modeling_qwen3_vl import ( Qwen3VLForConditionalGeneration as _Qwen3VLForConditionalGeneration, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( Qwen3VLModel as _Qwen3VLModel, - Qwen3VLTextModel as _Qwen3VLTextModel, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( Qwen3VLPreTrainedModel as _Qwen3VLPreTrainedModel, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( Qwen3VLTextAttention, Qwen3VLTextMLP, Qwen3VLTextRMSNorm, Qwen3VLTextRotaryEmbedding, - Qwen3VLVisionModel, Qwen3VLVisionMLP, - apply_rotary_pos_emb, + Qwen3VLVisionModel, apply_rotary_pos_emb_vision, eager_attention_forward, ) - +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLTextModel as _Qwen3VLTextModel, +) +from transformers.processing_utils import Unpack +from transformers.utils import logging logger = logging.get_logger(__name__) @@ -731,9 +329,3 @@ def forward_without_grid_thw( hidden_states = self.merger(hidden_states) return hidden_states, deepstack_feature_lists - - -def apply_lingbot_qwen3_vl_patch(): - logger.warning_once( - "apply_lingbot_qwen3_vl_patch is deprecated; LingBot-VLA v2 now installs Qwen3-VL changes per instance." - ) From 99d07d274f01525be1be97da601624343480456b Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Fri, 7 Aug 2026 04:16:57 +0000 Subject: [PATCH 12/15] feat(vla): add native structured service validator Add a VLA-specific real HTTP workload for fixed-count, concurrent, and duration validation. Verify the native service contract, task lifecycle, finite 50x55 canonical actions, latency, throughput, and bounded result artifacts without changing shared service interfaces.\n\nDocument single-replica, multi-replica, and soak usage, and add focused CPU coverage for contract validation and concurrent requests.\n\nVerification:\n- ruff check tools/validation/validate_lingbot_vla_v2_structured_service.py tests/unit/validation/test_lingbot_vla_v2_structured_service.py\n- 27 focused pytest tests passed\n- real 6B native HTTP smoke passed\n- 20/20 single-replica 256x256 requests completed successfully --- examples/lingbot_vla_v2/README.md | 58 ++ .../test_lingbot_vla_v2_structured_service.py | 157 ++++ ...idate_lingbot_vla_v2_structured_service.py | 688 ++++++++++++++++++ 3 files changed, 903 insertions(+) create mode 100644 tests/unit/validation/test_lingbot_vla_v2_structured_service.py create mode 100644 tools/validation/validate_lingbot_vla_v2_structured_service.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 31d9a410..f4f89748 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -186,6 +186,64 @@ the first accepted request represents a ready replica. The default `service-thre service runner's fixed worker thread; use `--execution-mode direct` only to measure the in-process pipeline ceiling. Shutdown still offloads the policy explicitly. +## Native Structured API Validation + +Use the VLA-specific HTTP validator after the native service reports ready. This is the structured-output counterpart +to the model-specific direct and AIPerf workloads used by the video and LingBot-World integrations: it exercises the +real TeleFuser HTTP boundary, asynchronous scheduler, task status polling, pipeline pool, and result serialization. +It emits raw request facts and aggregate latency distributions to a JSON artifact; it does not add a VLA-specific +service interface or change shared metric semantics. + +Run a single-replica smoke and latency check: + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_structured_service.py \ + --base-url http://127.0.0.1:18080 \ + --image examples/data/lingbot_world_fast/image.jpg \ + --warmup 1 \ + --requests 20 \ + --concurrency 1 \ + --output work_dirs/vla_service_validation/smoke_20.json +``` + +When the target was started with two independent replicas, validate request-level concurrency with: + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_structured_service.py \ + --base-url http://127.0.0.1:18080 \ + --image examples/data/lingbot_world_fast/image.jpg \ + --warmup 2 \ + --requests 100 \ + --concurrency 2 \ + --output work_dirs/vla_service_validation/two_replica_100.json +``` + +Use duration mode for a bounded soak. Workers use closed-loop scheduling: each worker submits its next request only +after its previous task reaches a terminal state. + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_structured_service.py \ + --base-url http://127.0.0.1:18080 \ + --camera-high /data/cam_high.png \ + --camera-left-wrist /data/cam_left_wrist.png \ + --camera-right-wrist /data/cam_right_wrist.png \ + --duration-seconds 7200 \ + --concurrency 1 \ + --output work_dirs/vla_service_validation/soak_2h.json +``` + +The command exits nonzero if readiness or contract checks fail, any measured request fails, task IDs are duplicated, +or the queue is not drained at the end. Each successful record validates the expected `50x55` finite action tensor +and retains only statistics and a float64 action fingerprint. Full actions and Base64 camera contents are deliberately +excluded from the artifact. `--max-records` bounds retained per-request samples during long runs while aggregate +latency and success counters still cover the complete run. + +This validation aligns VLA with the repository's existing deployment practice, but it is not yet an AIPerf workload: +AIPerf currently has maintained adapters for batch media and LingBot streaming transports, not the asynchronous +structured task API. The JSON report therefore keeps target inference time separate from client end-to-end time and +preserves target metadata and raw service metric snapshots so a future transport adapter can consume the same facts. +Passing this check proves serving and normalized action structure, not embodiment-specific control semantics. + ## TeleFuser Regression Baseline The validation capture runs through the public loader and pipeline, then records preprocessing tensors, fixed initial diff --git a/tests/unit/validation/test_lingbot_vla_v2_structured_service.py b/tests/unit/validation/test_lingbot_vla_v2_structured_service.py new file mode 100644 index 00000000..c8ca4745 --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_structured_service.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import argparse +import threading +from typing import Any + +import pytest + +from tools.validation import validate_lingbot_vla_v2_structured_service as validator + + +def _action_result(value: float = 0.25) -> dict[str, Any]: + return { + "canonical_normalized_actions": [[value] * 55 for _ in range(50)], + "horizon": 50, + "action_dim": 55, + "checkpoint_variant": "base", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + } + + +def _metadata() -> dict[str, Any]: + parameters = { + name: {"type": "string", "required": True} + for name in ( + "instruction", + "state", + "camera_high", + "camera_left_wrist", + "camera_right_wrist", + ) + } + return { + "declared_pipeline_contract": True, + "supported_tasks": ["vla_action"], + "supported_media_types": ["structured"], + "task_contracts": {"vla_action": {"media_type": "structured", "parameters": parameters}}, + } + + +def test_parse_state_json_requires_fourteen_finite_numbers() -> None: + assert validator.parse_state_json("[0,1,2,3,4,5,6,7,8,9,10,11,12,13]") == [float(index) for index in range(14)] + + with pytest.raises(argparse.ArgumentTypeError, match="exactly 14"): + validator.parse_state_json("[0, 1]") + with pytest.raises(argparse.ArgumentTypeError, match="finite numbers"): + validator.parse_state_json("[0,1,2,3,4,5,6,7,8,9,10,11,12,true]") + + +def test_validate_service_metadata_requires_native_structured_contract() -> None: + validator.validate_service_metadata(_metadata()) + + metadata = _metadata() + metadata["task_contracts"]["vla_action"]["media_type"] = "video" + with pytest.raises(validator.ValidationFailure, match="structured task contract"): + validator.validate_service_metadata(metadata) + + +def test_validate_action_result_reports_shape_stats_and_fingerprint() -> None: + summary = validator.validate_action_result(_action_result(), expected_horizon=50, expected_action_dim=55) + + assert summary["shape"] == [50, 55] + assert summary["value_count"] == 2750 + assert summary["minimum"] == 0.25 + assert summary["maximum"] == 0.25 + assert summary["policy_verified"] is False + assert len(summary["sha256_float64_le"]) == 64 + + +@pytest.mark.parametrize( + "mutation,match", + [ + (lambda result: result.update(horizon=49), "horizon field"), + (lambda result: result["canonical_normalized_actions"][0].pop(), "row 0"), + (lambda result: result["canonical_normalized_actions"][0].__setitem__(0, float("nan")), "non-finite"), + ], +) +def test_validate_action_result_rejects_invalid_contract(mutation, match: str) -> None: + result = _action_result() + mutation(result) + + with pytest.raises(validator.ValidationFailure, match=match): + validator.validate_action_result(result, expected_horizon=50, expected_action_dim=55) + + +class _Response: + def __init__(self, body: dict[str, Any]) -> None: + self._body = body + self.status_code = 200 + self.text = "" + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, Any]: + return self._body + + +class _Session: + def __init__(self, state: dict[str, Any]) -> None: + self.state = state + self.trust_env = False + + def __enter__(self) -> "_Session": + return self + + def __exit__(self, *args: object) -> None: + return None + + def request(self, method: str, url: str, *, json=None, timeout=None) -> _Response: + if method == "POST": + assert url.endswith("/v1/tasks/structured") + assert json["task"] == "vla_action" + with self.state["lock"]: + self.state["next_id"] += 1 + task_id = f"task-{self.state['next_id']}" + return _Response({"task_id": task_id, "task_status": "pending"}) + task_id = url.rsplit("/", maxsplit=2)[-2] + return _Response( + { + "task_id": task_id, + "status": "completed", + "inference_time_s": 0.25, + "peak_memory_mb": 128.0, + "result": _action_result(), + } + ) + + +def test_run_workload_exercises_concurrent_structured_requests(monkeypatch: pytest.MonkeyPatch) -> None: + state = {"lock": threading.Lock(), "next_id": 0} + monkeypatch.setattr(validator, "_new_session", lambda: _Session(state)) + config = validator.RequestConfig( + base_url="http://127.0.0.1:18080", + payload={"task": "vla_action"}, + http_timeout_seconds=1.0, + task_timeout_seconds=1.0, + poll_interval_seconds=0.001, + expected_horizon=50, + expected_action_dim=55, + ) + + report = validator.run_workload( + config, + request_count=4, + duration_seconds=None, + concurrency=2, + max_records=10, + ) + + assert report["requests"]["total"] == 4 + assert report["requests"]["succeeded"] == 4 + assert report["requests"]["failed"] == 0 + assert report["requests"]["unique_task_ids"] == 4 + assert report["latency_seconds"]["target_inference"]["mean"] == 0.25 + assert len(report["retained_records"]["successful"]) == 4 diff --git a/tools/validation/validate_lingbot_vla_v2_structured_service.py b/tools/validation/validate_lingbot_vla_v2_structured_service.py new file mode 100644 index 00000000..49912125 --- /dev/null +++ b/tools/validation/validate_lingbot_vla_v2_structured_service.py @@ -0,0 +1,688 @@ +"""Validate a real LingBot-VLA v2 native structured API service.""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import json +import math +import platform +import statistics +import struct +import subprocess +import sys +import threading +import time +from collections import Counter, deque +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timezone +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +import requests + +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) +_REQUIRED_PARAMETERS = frozenset( + { + "instruction", + "state", + "camera_high", + "camera_left_wrist", + "camera_right_wrist", + } +) + + +class ValidationFailure(RuntimeError): + """Raised when the target violates the VLA structured API contract.""" + + +@dataclass(frozen=True) +class RequestConfig: + """Immutable settings shared by validation workers.""" + + base_url: str + payload: dict[str, Any] + http_timeout_seconds: float + task_timeout_seconds: float + poll_interval_seconds: float + expected_horizon: int + expected_action_dim: int + + +def parse_state_json(value: str) -> list[float]: + """Parse and validate a finite 14-dimensional RobotWin state.""" + try: + raw = json.loads(value) + except json.JSONDecodeError as error: + raise argparse.ArgumentTypeError("state must be valid JSON") from error + if not isinstance(raw, list) or len(raw) != 14: + raise argparse.ArgumentTypeError("state must be a JSON array containing exactly 14 values") + state: list[float] = [] + for item in raw: + if isinstance(item, bool) or not isinstance(item, int | float) or not math.isfinite(float(item)): + raise argparse.ArgumentTypeError("state values must be finite numbers") + state.append(float(item)) + return state + + +def percentile(values: Sequence[float], fraction: float) -> float: + """Return a linearly interpolated percentile for a non-empty sample.""" + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def summarize(values: Sequence[float]) -> dict[str, float | int] | None: + """Summarize a possibly empty sample in seconds.""" + if not values: + return None + return { + "count": len(values), + "mean": statistics.fmean(values), + "min": min(values), + "p50": percentile(values, 0.50), + "p90": percentile(values, 0.90), + "p95": percentile(values, 0.95), + "p99": percentile(values, 0.99), + "max": max(values), + } + + +def validate_service_metadata(metadata: Any) -> None: + """Validate that the target exposes the native VLA structured contract.""" + if not isinstance(metadata, dict): + raise ValidationFailure("service metadata must be a JSON object") + if metadata.get("declared_pipeline_contract") is not True: + raise ValidationFailure("service does not expose a declared pipeline contract") + if "vla_action" not in metadata.get("supported_tasks", []): + raise ValidationFailure("service metadata does not declare the vla_action task") + if "structured" not in metadata.get("supported_media_types", []): + raise ValidationFailure("service metadata does not declare structured output") + task_contract = metadata.get("task_contracts", {}).get("vla_action") + if not isinstance(task_contract, dict) or task_contract.get("media_type") != "structured": + raise ValidationFailure("vla_action does not have a structured task contract") + parameters = task_contract.get("parameters") + if not isinstance(parameters, dict): + raise ValidationFailure("vla_action parameters are missing from service metadata") + missing = sorted(_REQUIRED_PARAMETERS.difference(parameters)) + if missing: + raise ValidationFailure(f"vla_action metadata is missing parameters: {', '.join(missing)}") + invalid = sorted(name for name in _REQUIRED_PARAMETERS if not isinstance(parameters[name], dict)) + if invalid: + raise ValidationFailure(f"vla_action parameter contracts are invalid: {', '.join(invalid)}") + not_required = sorted(name for name in _REQUIRED_PARAMETERS if parameters[name].get("required") is not True) + if not_required: + raise ValidationFailure(f"vla_action parameters are not required: {', '.join(not_required)}") + + +def validate_action_result(result: Any, *, expected_horizon: int, expected_action_dim: int) -> dict[str, Any]: + """Validate and summarize one canonical normalized action chunk.""" + if expected_horizon < 1 or expected_action_dim < 1: + raise ValueError("expected action dimensions must be positive") + if not isinstance(result, dict): + raise ValidationFailure("completed task result must be a JSON object") + actions = result.get("canonical_normalized_actions") + if not isinstance(actions, list) or len(actions) != expected_horizon: + observed = len(actions) if isinstance(actions, list) else type(actions).__name__ + raise ValidationFailure(f"expected action horizon {expected_horizon}, observed {observed}") + if result.get("horizon") != expected_horizon: + raise ValidationFailure(f"result horizon field is not {expected_horizon}") + if result.get("action_dim") != expected_action_dim: + raise ValidationFailure(f"result action_dim field is not {expected_action_dim}") + + flat: list[float] = [] + digest = hashlib.sha256() + for row_index, row in enumerate(actions): + if not isinstance(row, list) or len(row) != expected_action_dim: + observed = len(row) if isinstance(row, list) else type(row).__name__ + raise ValidationFailure(f"action row {row_index} has dimension {observed}, expected {expected_action_dim}") + for value in row: + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(float(value)): + raise ValidationFailure("action chunk contains a non-finite or non-numeric value") + number = float(value) + flat.append(number) + digest.update(struct.pack(" dict[str, Any]: + response = session.request(method, url, json=payload, timeout=timeout) + try: + response.raise_for_status() + except requests.HTTPError as error: + body = response.text[:1000] + raise ValidationFailure(f"{method} {url} returned HTTP {response.status_code}: {body}") from error + try: + body = response.json() + except ValueError as error: + raise ValidationFailure(f"{method} {url} did not return JSON") from error + if not isinstance(body, dict): + raise ValidationFailure(f"{method} {url} did not return a JSON object") + return body + + +def _new_session() -> requests.Session: + session = requests.Session() + session.trust_env = False + return session + + +def inspect_service(base_url: str, *, timeout_seconds: float) -> dict[str, Any]: + """Read and validate native service readiness and metadata.""" + with _new_session() as session: + ready = _request_json(session, "GET", f"{base_url}/v1/service/ready", timeout=timeout_seconds) + if ready.get("ready") is not True: + raise ValidationFailure("service readiness endpoint reports not ready") + metadata = _request_json(session, "GET", f"{base_url}/v1/service/metadata", timeout=timeout_seconds) + validate_service_metadata(metadata) + status = _request_json(session, "GET", f"{base_url}/v1/service/status", timeout=timeout_seconds) + metrics = _request_json(session, "GET", f"{base_url}/v1/service/metrics/json", timeout=timeout_seconds) + return {"ready": ready, "metadata": metadata, "status": status, "metrics": metrics} + + +def execute_request( + session: requests.Session, + config: RequestConfig, + *, + request_index: int, + worker_index: int, + run_started_at: float, +) -> dict[str, Any]: + """Submit, poll, validate, and summarize one real structured request.""" + record: dict[str, Any] = { + "request_index": request_index, + "worker_index": worker_index, + "start_offset_seconds": time.perf_counter() - run_started_at, + } + request_started_at = time.perf_counter() + try: + submit_started_at = time.perf_counter() + created = _request_json( + session, + "POST", + f"{config.base_url}/v1/tasks/structured", + timeout=config.http_timeout_seconds, + payload=config.payload, + ) + accepted_at = time.perf_counter() + record["submit_seconds"] = accepted_at - submit_started_at + task_id = created.get("task_id") + if not isinstance(task_id, str) or not task_id: + raise ValidationFailure("structured task creation response has no task_id") + record["task_id"] = task_id + + deadline = accepted_at + config.task_timeout_seconds + transitions: list[str] = [] + poll_count = 0 + while True: + if time.perf_counter() >= deadline: + raise ValidationFailure(f"task {task_id} exceeded {config.task_timeout_seconds:g}s timeout") + status = _request_json( + session, + "GET", + f"{config.base_url}/v1/tasks/{task_id}/status", + timeout=config.http_timeout_seconds, + ) + poll_count += 1 + task_status = status.get("status") or status.get("task_status") + if not isinstance(task_status, str): + raise ValidationFailure(f"task {task_id} status response has no status") + if not transitions or transitions[-1] != task_status: + transitions.append(task_status) + if task_status in _TERMINAL_STATUSES: + break + time.sleep(config.poll_interval_seconds) + + completed_at = time.perf_counter() + record.update( + end_to_end_seconds=completed_at - request_started_at, + accepted_to_terminal_seconds=completed_at - accepted_at, + poll_count=poll_count, + status_transitions=transitions, + terminal_status=task_status, + ) + inference_time = status.get("inference_time_s") + if inference_time is not None: + if isinstance(inference_time, bool) or not isinstance(inference_time, int | float): + raise ValidationFailure("inference_time_s must be numeric or null") + inference_time = float(inference_time) + if not math.isfinite(inference_time) or inference_time < 0: + raise ValidationFailure("inference_time_s must be finite and non-negative") + record["inference_time_seconds"] = inference_time + peak_memory = status.get("peak_memory_mb") + if peak_memory is not None: + if isinstance(peak_memory, bool) or not isinstance(peak_memory, int | float): + raise ValidationFailure("peak_memory_mb must be numeric or null") + peak_memory = float(peak_memory) + if not math.isfinite(peak_memory) or peak_memory < 0: + raise ValidationFailure("peak_memory_mb must be finite and non-negative") + record["peak_memory_mb"] = peak_memory + + if task_status != "completed": + raise ValidationFailure(f"task {task_id} reached terminal status {task_status}: {status.get('error')}") + record["action"] = validate_action_result( + status.get("result"), + expected_horizon=config.expected_horizon, + expected_action_dim=config.expected_action_dim, + ) + record["outcome"] = "succeeded" + except (requests.RequestException, ValidationFailure, ValueError) as error: + record["outcome"] = "failed" + record["error"] = str(error) + record.setdefault("end_to_end_seconds", time.perf_counter() - request_started_at) + return record + + +class RunAccumulator: + """Collect aggregate measurements while bounding retained request records.""" + + def __init__(self, max_records: int) -> None: + self._lock = threading.Lock() + self.max_records = max_records + self.total = 0 + self.succeeded = 0 + self.failed = 0 + self.task_ids: set[str] = set() + self.duplicate_task_ids: set[str] = set() + self.end_to_end: list[float] = [] + self.submit: list[float] = [] + self.accepted_to_terminal: list[float] = [] + self.inference: list[float] = [] + self.peak_memory: list[float] = [] + self.poll_counts: list[float] = [] + self.terminal_statuses: Counter[str] = Counter() + self.policy_statuses: Counter[str] = Counter() + self.failures: deque[dict[str, Any]] = deque(maxlen=max_records) + self.first_successes: list[dict[str, Any]] = [] + self.recent_successes: deque[dict[str, Any]] = deque(maxlen=max_records // 2) + + def add(self, record: dict[str, Any]) -> None: + with self._lock: + self.total += 1 + task_id = record.get("task_id") + if isinstance(task_id, str): + if task_id in self.task_ids: + self.duplicate_task_ids.add(task_id) + self.task_ids.add(task_id) + terminal_status = record.get("terminal_status") + if isinstance(terminal_status, str): + self.terminal_statuses[terminal_status] += 1 + if record["outcome"] == "failed": + self.failed += 1 + self.failures.append(record) + return + + self.succeeded += 1 + self.end_to_end.append(float(record["end_to_end_seconds"])) + self.submit.append(float(record["submit_seconds"])) + self.accepted_to_terminal.append(float(record["accepted_to_terminal_seconds"])) + self.poll_counts.append(float(record["poll_count"])) + if record.get("inference_time_seconds") is not None: + self.inference.append(float(record["inference_time_seconds"])) + if record.get("peak_memory_mb") is not None: + self.peak_memory.append(float(record["peak_memory_mb"])) + self.policy_statuses[str(record["action"]["verification_status"])] += 1 + first_capacity = self.max_records - self.recent_successes.maxlen + if len(self.first_successes) < first_capacity: + self.first_successes.append(record) + else: + self.recent_successes.append(record) + + def report(self, elapsed_seconds: float) -> dict[str, Any]: + retained_successes = self.first_successes + list(self.recent_successes) + retained_successes.sort(key=lambda record: int(record["request_index"])) + failures = sorted(self.failures, key=lambda record: int(record["request_index"])) + return { + "requests": { + "total": self.total, + "succeeded": self.succeeded, + "failed": self.failed, + "success_rate": self.succeeded / self.total if self.total else 0.0, + "unique_task_ids": len(self.task_ids), + "duplicate_task_ids": sorted(self.duplicate_task_ids), + "terminal_statuses": dict(sorted(self.terminal_statuses.items())), + "policy_statuses": dict(sorted(self.policy_statuses.items())), + }, + "elapsed_seconds": elapsed_seconds, + "throughput_requests_per_second": self.succeeded / elapsed_seconds if elapsed_seconds > 0 else 0.0, + "latency_seconds": { + "end_to_end": summarize(self.end_to_end), + "submission": summarize(self.submit), + "accepted_to_terminal": summarize(self.accepted_to_terminal), + "target_inference": summarize(self.inference), + }, + "poll_count": summarize(self.poll_counts), + "peak_memory_mb": summarize(self.peak_memory), + "retained_records": { + "limit_per_outcome": self.max_records, + "successful": retained_successes, + "failed": failures, + }, + } + + +def run_workload( + config: RequestConfig, + *, + request_count: int | None, + duration_seconds: float | None, + concurrency: int, + max_records: int, +) -> dict[str, Any]: + """Run a closed-loop fixed-count or duration workload.""" + accumulator = RunAccumulator(max_records) + counter = 0 + counter_lock = threading.Lock() + run_started_at = time.perf_counter() + stop_claiming_at = None if duration_seconds is None else run_started_at + duration_seconds + workers = concurrency if request_count is None else min(concurrency, request_count) + barrier = threading.Barrier(workers) + + def claim_request() -> int | None: + nonlocal counter + with counter_lock: + if request_count is not None and counter >= request_count: + return None + if stop_claiming_at is not None and time.perf_counter() >= stop_claiming_at: + return None + index = counter + counter += 1 + return index + + def worker(worker_index: int) -> None: + with _new_session() as session: + barrier.wait() + while (request_index := claim_request()) is not None: + accumulator.add( + execute_request( + session, + config, + request_index=request_index, + worker_index=worker_index, + run_started_at=run_started_at, + ) + ) + + with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="vla-structured-validator") as executor: + futures = [executor.submit(worker, worker_index) for worker_index in range(workers)] + for future in futures: + future.result() + return accumulator.report(time.perf_counter() - run_started_at) + + +def _encode_image(path: Path) -> str: + if not path.is_file(): + raise ValueError(f"camera image does not exist: {path}") + return base64.b64encode(path.read_bytes()).decode("ascii") + + +def _resolve_camera_paths(args: argparse.Namespace) -> tuple[Path, Path, Path]: + fallback = args.image + paths = tuple(path or fallback for path in (args.camera_high, args.camera_left_wrist, args.camera_right_wrist)) + if any(path is None for path in paths): + raise ValueError("provide --image or all three --camera-* paths") + return paths # type: ignore[return-value] + + +def _package_version() -> str: + try: + return version("telefuser") + except PackageNotFoundError: + return "source" + + +def _git_commit(repo_root: Path) -> str | None: + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo_root, + check=True, + capture_output=True, + text=True, + timeout=5, + ) + except (OSError, subprocess.SubprocessError): + return None + return result.stdout.strip() or None + + +def _metric_delta(before: dict[str, Any], after: dict[str, Any]) -> dict[str, Any]: + delta: dict[str, Any] = {} + for section in ("tasks",): + before_section = before.get(section, {}) + after_section = after.get(section, {}) + if isinstance(before_section, dict) and isinstance(after_section, dict): + delta[section] = { + key: after_section[key] - before_section.get(key, 0) + for key in after_section + if isinstance(after_section[key], int | float) and not isinstance(after_section[key], bool) + } + return delta + + +def run_validation(args: argparse.Namespace) -> dict[str, Any]: + """Validate the service and return a reproducible JSON report.""" + if args.concurrency < 1 or args.warmup < 0 or args.max_records < 2: + raise ValueError("concurrency must be positive, warmup non-negative, and max-records at least 2") + if args.requests is not None and args.requests < 1: + raise ValueError("requests must be positive") + if args.duration_seconds is not None and args.duration_seconds <= 0: + raise ValueError("duration-seconds must be positive") + if args.poll_interval_seconds <= 0 or args.http_timeout_seconds <= 0 or args.task_timeout_seconds <= 0: + raise ValueError("poll interval and HTTP/task timeouts must be positive") + if args.expected_horizon < 1 or args.expected_action_dim < 1: + raise ValueError("expected action dimensions must be positive") + base_url = args.base_url.rstrip("/") + camera_high, camera_left, camera_right = _resolve_camera_paths(args) + payload = { + "task": "vla_action", + "instruction": args.instruction, + "state": args.state_json, + "camera_high": _encode_image(camera_high), + "camera_left_wrist": _encode_image(camera_left), + "camera_right_wrist": _encode_image(camera_right), + "seed": args.seed, + } + config = RequestConfig( + base_url=base_url, + payload=payload, + http_timeout_seconds=args.http_timeout_seconds, + task_timeout_seconds=args.task_timeout_seconds, + poll_interval_seconds=args.poll_interval_seconds, + expected_horizon=args.expected_horizon, + expected_action_dim=args.expected_action_dim, + ) + before = inspect_service(base_url, timeout_seconds=args.http_timeout_seconds) + warmup_records: list[dict[str, Any]] = [] + warmup_started_at = time.perf_counter() + with _new_session() as session: + for index in range(args.warmup): + warmup_records.append( + execute_request( + session, + config, + request_index=index, + worker_index=0, + run_started_at=warmup_started_at, + ) + ) + if any(record["outcome"] != "succeeded" for record in warmup_records): + raise ValidationFailure("at least one warmup request failed") + + request_count = args.requests + if request_count is None and args.duration_seconds is None: + request_count = 1 + workload = run_workload( + config, + request_count=request_count, + duration_seconds=args.duration_seconds, + concurrency=args.concurrency, + max_records=args.max_records, + ) + after = inspect_service(base_url, timeout_seconds=args.http_timeout_seconds) + requests_report = workload["requests"] + checks = { + "service_ready_before": before["ready"].get("ready") is True, + "service_ready_after": after["ready"].get("ready") is True, + "warmup_succeeded": all(record["outcome"] == "succeeded" for record in warmup_records), + "all_measured_requests_succeeded": requests_report["failed"] == 0 and requests_report["total"] > 0, + "task_ids_unique": not requests_report["duplicate_task_ids"], + "queue_drained": ( + after["metrics"].get("queue", {}).get("pending") == 0 + and after["metrics"].get("queue", {}).get("processing") == 0 + ), + } + repo_root = Path(__file__).resolve().parents[2] + return { + "schema_version": 1, + "validation": "lingbot_vla_v2_native_structured_api", + "passed": all(checks.values()), + "checks": checks, + "created_at": datetime.now(timezone.utc).isoformat(), + "environment": { + "python": platform.python_version(), + "python_executable": sys.executable, + "platform": platform.platform(), + "telefuser_version": _package_version(), + "telefuser_commit": _git_commit(repo_root), + }, + "target": { + "base_url": base_url, + "transport": "HTTP native TeleFuser asynchronous structured task API", + "metadata": before["metadata"], + "status_before": before["status"], + "status_after": after["status"], + "health_before": before["ready"], + "health_after": after["ready"], + "metrics_before": before["metrics"], + "metrics_after": after["metrics"], + "metrics_delta": _metric_delta(before["metrics"], after["metrics"]), + }, + "workload": { + "mode": "duration" if args.duration_seconds is not None else "fixed_requests", + "requested_requests": request_count, + "requested_duration_seconds": args.duration_seconds, + "concurrency": args.concurrency, + "warmup_requests": args.warmup, + "instruction": args.instruction, + "state_dimension": len(args.state_json), + "seed": args.seed, + "camera_files": { + "high": str(camera_high.resolve()), + "left_wrist": str(camera_left.resolve()), + "right_wrist": str(camera_right.resolve()), + }, + "expected_action_shape": [args.expected_horizon, args.expected_action_dim], + "poll_interval_seconds": args.poll_interval_seconds, + "task_timeout_seconds": args.task_timeout_seconds, + }, + "warmup_records": warmup_records, + "result": workload, + "interpretation": ( + "This validates service transport, scheduling, and normalized canonical action structure. " + "It does not establish embodiment-specific robot control semantics." + ), + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080") + camera_group = parser.add_argument_group("camera inputs") + camera_group.add_argument("--image", type=Path, help="Fallback image reused for camera inputs not set explicitly.") + camera_group.add_argument("--camera-high", type=Path) + camera_group.add_argument("--camera-left-wrist", type=Path) + camera_group.add_argument("--camera-right-wrist", type=Path) + parser.add_argument("--instruction", default="pick up the red block") + parser.add_argument( + "--state-json", type=parse_state_json, default=parse_state_json("[0,0,0,0,0,0,0,0,0,0,0,0,0,0]") + ) + parser.add_argument("--seed", type=int, default=7) + workload_group = parser.add_mutually_exclusive_group() + workload_group.add_argument("--requests", type=int) + workload_group.add_argument("--duration-seconds", type=float) + parser.add_argument("--concurrency", type=int, default=1) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--poll-interval-seconds", type=float, default=0.1) + parser.add_argument("--http-timeout-seconds", type=float, default=30.0) + parser.add_argument("--task-timeout-seconds", type=float, default=300.0) + parser.add_argument("--expected-horizon", type=int, default=50) + parser.add_argument("--expected-action-dim", type=int, default=55) + parser.add_argument("--max-records", type=int, default=1000) + parser.add_argument("--output", required=True, type=Path) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + report: dict[str, Any] + exit_code = 0 + try: + report = run_validation(args) + if not report["passed"]: + exit_code = 1 + except Exception as error: + report = { + "schema_version": 1, + "validation": "lingbot_vla_v2_native_structured_api", + "passed": False, + "created_at": datetime.now(timezone.utc).isoformat(), + "fatal_error": str(error), + } + exit_code = 1 + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + summary = { + "passed": report["passed"], + "checks": report.get("checks"), + "requests": report.get("result", {}).get("requests"), + "latency_seconds": report.get("result", {}).get("latency_seconds"), + "fatal_error": report.get("fatal_error"), + "artifact": str(args.output), + } + print(json.dumps(summary, indent=2, sort_keys=True)) + raise SystemExit(exit_code) + + +if __name__ == "__main__": + main() From 6228096334e6049f1007f03683a566f393cfa8a2 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 10 Aug 2026 01:54:45 +0000 Subject: [PATCH 13/15] test(vla): harden structured service validation Freeze the LingBot VLA v2 structured request and action-result contracts, reject sensitive camera echoes, and report latency-window trends. Add opt-in process-tree RSS and per-GPU NVML sampling with bounded artifacts, document the validation workflow, and cover contract and sampler boundaries in unit tests. Verification: 33 focused pytest tests; ruff check; ruff format --check; git diff --check. --- examples/lingbot_vla_v2/README.md | 18 +- .../test_lingbot_vla_v2_structured_service.py | 114 +++++- ...idate_lingbot_vla_v2_structured_service.py | 361 ++++++++++++++++-- 3 files changed, 460 insertions(+), 33 deletions(-) diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index f4f89748..31d958d1 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -229,14 +229,30 @@ after its previous task reaches a terminal state. --camera-right-wrist /data/cam_right_wrist.png \ --duration-seconds 7200 \ --concurrency 1 \ + --service-pid \ + --gpu-indexes 0 \ + --resource-interval-seconds 1 \ --output work_dirs/vla_service_validation/soak_2h.json ``` +Resource sampling is opt-in and local-only. `--service-pid` must identify the parent `telefuser serve` process; its +replica descendants are discovered on every sample. RSS is summed across that process tree, while `nvidia-smi` +process memory is grouped by physical GPU index. For a two-replica service on physical GPUs 0 and 1, pass +`--gpu-indexes 0,1`. Omitting `--service-pid` keeps remote-service validation lightweight and does not invoke +`nvidia-smi`. Reports retain bounded raw samples plus distributions and first/last 10% trends for latency, RSS, and +per-GPU process memory. + +The validator freezes the current structured contract. Requests contain exactly `task`, `instruction`, `state`, the +three camera fields, and optional `seed`. Action results contain exactly `canonical_normalized_actions`, `horizon`, +`action_dim`, `checkpoint_variant`, `policy_verified`, and `verification_status`. Safe additive task-status metadata +remains allowed, but status responses must not echo the three Base64 camera fields. + The command exits nonzero if readiness or contract checks fail, any measured request fails, task IDs are duplicated, or the queue is not drained at the end. Each successful record validates the expected `50x55` finite action tensor and retains only statistics and a float64 action fingerprint. Full actions and Base64 camera contents are deliberately excluded from the artifact. `--max-records` bounds retained per-request samples during long runs while aggregate -latency and success counters still cover the complete run. +latency and success counters still cover the complete run. `--max-resource-samples` independently bounds retained +resource samples. This validation aligns VLA with the repository's existing deployment practice, but it is not yet an AIPerf workload: AIPerf currently has maintained adapters for batch media and LingBot streaming transports, not the asynchronous diff --git a/tests/unit/validation/test_lingbot_vla_v2_structured_service.py b/tests/unit/validation/test_lingbot_vla_v2_structured_service.py index c8ca4745..7b916373 100644 --- a/tests/unit/validation/test_lingbot_vla_v2_structured_service.py +++ b/tests/unit/validation/test_lingbot_vla_v2_structured_service.py @@ -2,6 +2,7 @@ import argparse import threading +import time from typing import Any import pytest @@ -22,20 +23,25 @@ def _action_result(value: float = 0.25) -> dict[str, Any]: def _metadata() -> dict[str, Any]: parameters = { - name: {"type": "string", "required": True} - for name in ( - "instruction", - "state", - "camera_high", - "camera_left_wrist", - "camera_right_wrist", - ) + "instruction": {"type": "string", "required": True}, + "state": {"type": "array", "required": True}, + "camera_high": {"type": "string", "required": True}, + "camera_left_wrist": {"type": "string", "required": True}, + "camera_right_wrist": {"type": "string", "required": True}, + "seed": {"type": "integer", "required": False}, } return { "declared_pipeline_contract": True, "supported_tasks": ["vla_action"], "supported_media_types": ["structured"], - "task_contracts": {"vla_action": {"media_type": "structured", "parameters": parameters}}, + "task_contracts": { + "vla_action": { + "media_type": "structured", + "required_inputs": ["camera_high", "camera_left_wrist", "camera_right_wrist"], + "optional_inputs": [], + "parameters": parameters, + } + }, } @@ -56,6 +62,11 @@ def test_validate_service_metadata_requires_native_structured_contract() -> None with pytest.raises(validator.ValidationFailure, match="structured task contract"): validator.validate_service_metadata(metadata) + metadata = _metadata() + del metadata["task_contracts"]["vla_action"]["parameters"]["seed"] + with pytest.raises(validator.ValidationFailure, match="parameter fields changed"): + validator.validate_service_metadata(metadata) + def test_validate_action_result_reports_shape_stats_and_fingerprint() -> None: summary = validator.validate_action_result(_action_result(), expected_horizon=50, expected_action_dim=55) @@ -68,6 +79,65 @@ def test_validate_action_result_reports_shape_stats_and_fingerprint() -> None: assert len(summary["sha256_float64_le"]) == 64 +def test_validate_action_result_rejects_additive_result_fields() -> None: + result = _action_result() + result["debug"] = "unstable" + + with pytest.raises(validator.ValidationFailure, match="result fields changed"): + validator.validate_action_result(result, expected_horizon=50, expected_action_dim=55) + + +def test_validate_task_status_rejects_sensitive_echo_and_missing_fields() -> None: + status = { + "task_id": "task-1", + "status": "completed", + "inference_time_s": 0.25, + "peak_memory_mb": None, + "result": _action_result(), + } + validator.validate_task_status(status, task_id="task-1") + + leaked = dict(status, camera_high="base64") + with pytest.raises(validator.ValidationFailure, match="sensitive image"): + validator.validate_task_status(leaked, task_id="task-1") + + +def test_compare_windows_reports_first_and_last_measurement_change() -> None: + result = validator.compare_windows([1.0, 1.0, 1.0, 1.2, 1.2, 1.2]) + + assert result is not None + assert result["window_count"] == 1 + assert result["change_percent"] == pytest.approx(20.0) + + +def test_parse_gpu_process_memory_filters_process_tree_and_physical_gpu() -> None: + output = "\n".join( + [ + "101, GPU-a, 1024", + "102, GPU-a, 512", + "101, GPU-b, 2048", + "999, GPU-a, 4096", + "malformed", + ] + ) + + result = validator._parse_gpu_process_memory( + output, + process_ids={101, 102}, + uuid_to_index={"GPU-a": "0", "GPU-b": "1"}, + gpu_indexes={"0"}, + ) + + assert result == {"0": 1536.0} + + +def test_parse_gpu_indexes_requires_integer_indexes() -> None: + assert validator.parse_gpu_indexes("0, 2") == {"0", "2"} + + with pytest.raises(argparse.ArgumentTypeError, match="comma-separated"): + validator.parse_gpu_indexes("0,GPU-a") + + @pytest.mark.parametrize( "mutation,match", [ @@ -155,3 +225,29 @@ def test_run_workload_exercises_concurrent_structured_requests(monkeypatch: pyte assert report["requests"]["unique_task_ids"] == 4 assert report["latency_seconds"]["target_inference"]["mean"] == 0.25 assert len(report["retained_records"]["successful"]) == 4 + + +def test_resource_sampler_bounds_samples_and_reports_trend() -> None: + state = {"value": 100.0} + ready = threading.Event() + + def sample() -> dict[str, Any]: + state["value"] += 10.0 + ready.set() + return {"process_ids": [123], "cpu_rss_mib": state["value"], "gpu_memory_mib": {"0": 2048.0}} + + sampler = validator.ResourceSampler( + 123, + interval_seconds=0.001, + max_samples=4, + sample_function=sample, + ) + sampler.start() + assert ready.wait(1.0) + time.sleep(0.01) + report = sampler.stop() + + assert report["sample_count"] >= 2 + assert len(report["retained_samples"]) <= 4 + assert report["cpu_rss_mib"]["trend"]["last_mean"] > report["cpu_rss_mib"]["trend"]["first_mean"] + assert report["gpu_memory_mib"]["0"]["distribution"]["mean"] == 2048.0 diff --git a/tools/validation/validate_lingbot_vla_v2_structured_service.py b/tools/validation/validate_lingbot_vla_v2_structured_service.py index 49912125..b35e7e28 100644 --- a/tools/validation/validate_lingbot_vla_v2_structured_service.py +++ b/tools/validation/validate_lingbot_vla_v2_structured_service.py @@ -15,7 +15,7 @@ import threading import time from collections import Counter, deque -from collections.abc import Sequence +from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime, timezone @@ -23,18 +23,30 @@ from pathlib import Path from typing import Any +import psutil import requests _TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) -_REQUIRED_PARAMETERS = frozenset( +_MIB = 1024**2 +_REQUEST_PARAMETER_CONTRACT = { + "instruction": ("string", True), + "state": ("array", True), + "camera_high": ("string", True), + "camera_left_wrist": ("string", True), + "camera_right_wrist": ("string", True), + "seed": ("integer", False), +} +_RESULT_FIELDS = frozenset( { - "instruction", - "state", - "camera_high", - "camera_left_wrist", - "camera_right_wrist", + "canonical_normalized_actions", + "horizon", + "action_dim", + "checkpoint_variant", + "policy_verified", + "verification_status", } ) +_SENSITIVE_REQUEST_FIELDS = frozenset({"camera_high", "camera_left_wrist", "camera_right_wrist"}) class ValidationFailure(RuntimeError): @@ -99,6 +111,24 @@ def summarize(values: Sequence[float]) -> dict[str, float | int] | None: } +def compare_windows(values: Sequence[float], fraction: float = 0.1) -> dict[str, float | int] | None: + """Compare the first and last windows of an ordered measurement series.""" + if not values: + return None + window_count = max(1, math.ceil(len(values) * fraction)) + first_mean = statistics.fmean(values[:window_count]) + last_mean = statistics.fmean(values[-window_count:]) + delta = last_mean - first_mean + return { + "sample_count": len(values), + "window_count": window_count, + "first_mean": first_mean, + "last_mean": last_mean, + "delta": delta, + "change_percent": delta / first_mean * 100.0 if first_mean else 0.0, + } + + def validate_service_metadata(metadata: Any) -> None: """Validate that the target exposes the native VLA structured contract.""" if not isinstance(metadata, dict): @@ -115,15 +145,23 @@ def validate_service_metadata(metadata: Any) -> None: parameters = task_contract.get("parameters") if not isinstance(parameters, dict): raise ValidationFailure("vla_action parameters are missing from service metadata") - missing = sorted(_REQUIRED_PARAMETERS.difference(parameters)) - if missing: - raise ValidationFailure(f"vla_action metadata is missing parameters: {', '.join(missing)}") - invalid = sorted(name for name in _REQUIRED_PARAMETERS if not isinstance(parameters[name], dict)) - if invalid: - raise ValidationFailure(f"vla_action parameter contracts are invalid: {', '.join(invalid)}") - not_required = sorted(name for name in _REQUIRED_PARAMETERS if parameters[name].get("required") is not True) - if not_required: - raise ValidationFailure(f"vla_action parameters are not required: {', '.join(not_required)}") + if set(parameters) != set(_REQUEST_PARAMETER_CONTRACT): + raise ValidationFailure( + "vla_action parameter fields changed: " + f"expected {sorted(_REQUEST_PARAMETER_CONTRACT)}, observed {sorted(parameters)}" + ) + for name, (expected_type, expected_required) in _REQUEST_PARAMETER_CONTRACT.items(): + parameter = parameters[name] + if not isinstance(parameter, dict): + raise ValidationFailure(f"vla_action parameter contract is invalid: {name}") + if parameter.get("type") != expected_type or parameter.get("required") is not expected_required: + raise ValidationFailure( + f"vla_action parameter {name} changed: expected type={expected_type}, required={expected_required}" + ) + if task_contract.get("required_inputs") != ["camera_high", "camera_left_wrist", "camera_right_wrist"]: + raise ValidationFailure("vla_action required_inputs changed") + if task_contract.get("optional_inputs") != []: + raise ValidationFailure("vla_action optional_inputs changed") def validate_action_result(result: Any, *, expected_horizon: int, expected_action_dim: int) -> dict[str, Any]: @@ -132,6 +170,8 @@ def validate_action_result(result: Any, *, expected_horizon: int, expected_actio raise ValueError("expected action dimensions must be positive") if not isinstance(result, dict): raise ValidationFailure("completed task result must be a JSON object") + if set(result) != set(_RESULT_FIELDS): + raise ValidationFailure(f"result fields changed: expected {sorted(_RESULT_FIELDS)}, observed {sorted(result)}") actions = result.get("canonical_normalized_actions") if not isinstance(actions, list) or len(actions) != expected_horizon: observed = len(actions) if isinstance(actions, list) else type(actions).__name__ @@ -178,6 +218,21 @@ def validate_action_result(result: Any, *, expected_horizon: int, expected_actio } +def validate_task_status(status: Any, *, task_id: str) -> None: + """Validate stable terminal task fields without rejecting safe additive metadata.""" + if not isinstance(status, dict): + raise ValidationFailure("task status must be a JSON object") + if status.get("task_id") != task_id: + raise ValidationFailure("task status returned a different task_id") + required = {"status", "inference_time_s", "peak_memory_mb", "result"} + missing = sorted(required.difference(status)) + if missing: + raise ValidationFailure(f"task status is missing fields: {', '.join(missing)}") + leaked = sorted(_SENSITIVE_REQUEST_FIELDS.intersection(status)) + if leaked: + raise ValidationFailure(f"task status echoed sensitive image fields: {', '.join(leaked)}") + + def _request_json( session: requests.Session, method: str, @@ -274,6 +329,7 @@ def execute_request( time.sleep(config.poll_interval_seconds) completed_at = time.perf_counter() + validate_task_status(status, task_id=task_id) record.update( end_to_end_seconds=completed_at - request_started_at, accepted_to_terminal_seconds=completed_at - accepted_at, @@ -313,6 +369,222 @@ def execute_request( return record +def _parse_gpu_process_memory( + output: str, + *, + process_ids: set[int], + uuid_to_index: dict[str, str], + gpu_indexes: set[str] | None, +) -> dict[str, float]: + """Parse nvidia-smi process memory rows for one service process tree.""" + memory_by_gpu: dict[str, float] = {} + for line in output.splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) != 3: + continue + try: + pid = int(fields[0]) + memory_mib = float(fields[2]) + except ValueError: + continue + gpu_index = uuid_to_index.get(fields[1], fields[1]) + if pid not in process_ids or (gpu_indexes is not None and gpu_index not in gpu_indexes): + continue + memory_by_gpu[gpu_index] = memory_by_gpu.get(gpu_index, 0.0) + memory_mib + return memory_by_gpu + + +def _query_gpu_index_map() -> dict[str, str]: + """Resolve stable GPU UUIDs to physical indexes once per validation run.""" + uuid_result = subprocess.run( + ["nvidia-smi", "--query-gpu=index,uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return { + fields[1].strip(): fields[0].strip() + for line in uuid_result.stdout.splitlines() + if len(fields := [field.strip() for field in line.split(",")]) == 2 + } + + +def _query_gpu_process_memory( + process_ids: set[int], + *, + uuid_to_index: dict[str, str], + gpu_indexes: set[str] | None, +) -> dict[str, float]: + """Read GPU memory used by the service process tree through nvidia-smi.""" + process_result = subprocess.run( + [ + "nvidia-smi", + "--query-compute-apps=pid,gpu_uuid,used_gpu_memory", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return _parse_gpu_process_memory( + process_result.stdout, + process_ids=process_ids, + uuid_to_index=uuid_to_index, + gpu_indexes=gpu_indexes, + ) + + +def _sample_local_resources( + root_pid: int, + *, + uuid_to_index: dict[str, str], + gpu_indexes: set[str] | None, +) -> dict[str, Any]: + """Sample process-tree RSS and GPU memory without touching model execution.""" + root = psutil.Process(root_pid) + processes = [root, *root.children(recursive=True)] + process_ids: set[int] = set() + rss_bytes = 0 + for process in processes: + try: + process_ids.add(process.pid) + rss_bytes += process.memory_info().rss + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + return { + "process_ids": sorted(process_ids), + "cpu_rss_mib": rss_bytes / _MIB, + "gpu_memory_mib": _query_gpu_process_memory( + process_ids, + uuid_to_index=uuid_to_index, + gpu_indexes=gpu_indexes, + ), + } + + +class ResourceSampler: + """Periodically sample local service resources in a background thread.""" + + def __init__( + self, + root_pid: int, + *, + interval_seconds: float, + max_samples: int, + gpu_indexes: set[str] | None = None, + sample_function: Callable[[], dict[str, Any]] | None = None, + ) -> None: + if root_pid < 1 or interval_seconds <= 0 or max_samples < 2: + raise ValueError("invalid resource sampler configuration") + self.root_pid = root_pid + self.interval_seconds = interval_seconds + self.max_samples = max_samples + self.gpu_indexes = gpu_indexes + if sample_function is None: + uuid_to_index = _query_gpu_index_map() + + def sample_function() -> dict[str, Any]: + return _sample_local_resources( + root_pid, + uuid_to_index=uuid_to_index, + gpu_indexes=gpu_indexes, + ) + + self.sample_function = sample_function + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._started_at = time.perf_counter() + self._sample_count = 0 + self._cpu_rss: list[float] = [] + self._gpu_memory: dict[str, list[float]] = {} + self._process_ids: set[int] = set() + self._errors: deque[str] = deque(maxlen=100) + self._first_samples: list[dict[str, Any]] = [] + recent_capacity = max(1, max_samples // 2) + self._recent_samples: deque[dict[str, Any]] = deque(maxlen=recent_capacity) + + def start(self) -> None: + """Start sampling and take an initial sample.""" + self._started_at = time.perf_counter() + self._record_once() + self._thread = threading.Thread(target=self._sample_loop, name="vla-resource-sampler", daemon=True) + self._thread.start() + + def stop(self) -> dict[str, Any]: + """Stop sampling, take a final sample, and return the report.""" + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=max(self.interval_seconds * 2, 1.0)) + self._record_once() + return self.report() + + def _sample_loop(self) -> None: + while not self._stop.wait(self.interval_seconds): + self._record_once() + + def _record_once(self) -> None: + offset = time.perf_counter() - self._started_at + try: + sample = self.sample_function() + cpu_rss = float(sample["cpu_rss_mib"]) + if not math.isfinite(cpu_rss) or cpu_rss < 0: + raise ValueError("invalid cpu_rss_mib sample") + gpu_memory = sample.get("gpu_memory_mib", {}) + if not isinstance(gpu_memory, dict): + raise ValueError("invalid gpu_memory_mib sample") + normalized_gpu = { + str(index): float(value) + for index, value in gpu_memory.items() + if math.isfinite(float(value)) and float(value) >= 0 + } + process_ids = {int(pid) for pid in sample.get("process_ids", [])} + self._sample_count += 1 + self._cpu_rss.append(cpu_rss) + self._process_ids.update(process_ids) + for index, value in normalized_gpu.items(): + self._gpu_memory.setdefault(index, []).append(value) + retained = { + "offset_seconds": offset, + "process_ids": sorted(process_ids), + "cpu_rss_mib": cpu_rss, + "gpu_memory_mib": normalized_gpu, + } + first_capacity = self.max_samples - self._recent_samples.maxlen + if len(self._first_samples) < first_capacity: + self._first_samples.append(retained) + else: + self._recent_samples.append(retained) + except Exception as error: # pragma: no cover - hardware errors are environment-dependent + self._errors.append(str(error)) + + def report(self) -> dict[str, Any]: + """Return bounded samples, resource distributions, and first/last trends.""" + retained = self._first_samples + list(self._recent_samples) + return { + "enabled": True, + "root_pid": self.root_pid, + "interval_seconds": self.interval_seconds, + "sample_count": self._sample_count, + "gpu_sample_count": sum(1 for sample in retained if sample["gpu_memory_mib"]), + "observed_process_ids": sorted(self._process_ids), + "errors": list(self._errors), + "cpu_rss_mib": { + "distribution": summarize(self._cpu_rss), + "trend": compare_windows(self._cpu_rss), + }, + "gpu_memory_mib": { + index: { + "distribution": summarize(values), + "trend": compare_windows(values), + } + for index, values in sorted(self._gpu_memory.items()) + }, + "retained_samples": retained, + } + + class RunAccumulator: """Collect aggregate measurements while bounding retained request records.""" @@ -391,6 +663,10 @@ def report(self, elapsed_seconds: float) -> dict[str, Any]: "accepted_to_terminal": summarize(self.accepted_to_terminal), "target_inference": summarize(self.inference), }, + "latency_trend": { + "end_to_end": compare_windows(self.end_to_end), + "target_inference": compare_windows(self.inference), + }, "poll_count": summarize(self.poll_counts), "peak_memory_mb": summarize(self.peak_memory), "retained_records": { @@ -464,6 +740,14 @@ def _resolve_camera_paths(args: argparse.Namespace) -> tuple[Path, Path, Path]: return paths # type: ignore[return-value] +def parse_gpu_indexes(value: str) -> set[str]: + """Parse a comma-separated set of physical GPU indexes.""" + indexes = {item.strip() for item in value.split(",") if item.strip()} + if not indexes or any(not item.isdigit() for item in indexes): + raise argparse.ArgumentTypeError("GPU indexes must be a comma-separated list of integers") + return indexes + + def _package_version() -> str: try: return version("telefuser") @@ -512,6 +796,10 @@ def run_validation(args: argparse.Namespace) -> dict[str, Any]: raise ValueError("poll interval and HTTP/task timeouts must be positive") if args.expected_horizon < 1 or args.expected_action_dim < 1: raise ValueError("expected action dimensions must be positive") + if args.resource_interval_seconds <= 0 or args.max_resource_samples < 2: + raise ValueError("resource interval must be positive and max-resource-samples at least 2") + if args.gpu_indexes is not None and args.service_pid is None: + raise ValueError("--gpu-indexes requires --service-pid") base_url = args.base_url.rstrip("/") camera_high, camera_left, camera_right = _resolve_camera_paths(args) payload = { @@ -552,13 +840,29 @@ def run_validation(args: argparse.Namespace) -> dict[str, Any]: request_count = args.requests if request_count is None and args.duration_seconds is None: request_count = 1 - workload = run_workload( - config, - request_count=request_count, - duration_seconds=args.duration_seconds, - concurrency=args.concurrency, - max_records=args.max_records, - ) + resource_sampler: ResourceSampler | None = None + resource_report: dict[str, Any] = {"enabled": False} + if args.service_pid is not None: + if not psutil.pid_exists(args.service_pid): + raise ValueError(f"service PID does not exist: {args.service_pid}") + resource_sampler = ResourceSampler( + args.service_pid, + interval_seconds=args.resource_interval_seconds, + max_samples=args.max_resource_samples, + gpu_indexes=args.gpu_indexes, + ) + resource_sampler.start() + try: + workload = run_workload( + config, + request_count=request_count, + duration_seconds=args.duration_seconds, + concurrency=args.concurrency, + max_records=args.max_records, + ) + finally: + if resource_sampler is not None: + resource_report = resource_sampler.stop() after = inspect_service(base_url, timeout_seconds=args.http_timeout_seconds) requests_report = workload["requests"] checks = { @@ -567,6 +871,8 @@ def run_validation(args: argparse.Namespace) -> dict[str, Any]: "warmup_succeeded": all(record["outcome"] == "succeeded" for record in warmup_records), "all_measured_requests_succeeded": requests_report["failed"] == 0 and requests_report["total"] > 0, "task_ids_unique": not requests_report["duplicate_task_ids"], + "resource_samples_collected": (not resource_report["enabled"] or resource_report["sample_count"] > 0), + "gpu_resource_samples_collected": (not resource_report["enabled"] or resource_report["gpu_sample_count"] > 0), "queue_drained": ( after["metrics"].get("queue", {}).get("pending") == 0 and after["metrics"].get("queue", {}).get("processing") == 0 @@ -618,6 +924,7 @@ def run_validation(args: argparse.Namespace) -> dict[str, Any]: }, "warmup_records": warmup_records, "result": workload, + "resources": resource_report, "interpretation": ( "This validates service transport, scheduling, and normalized canonical action structure. " "It does not establish embodiment-specific robot control semantics." @@ -649,6 +956,14 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--expected-horizon", type=int, default=50) parser.add_argument("--expected-action-dim", type=int, default=55) parser.add_argument("--max-records", type=int, default=1000) + parser.add_argument( + "--service-pid", + type=int, + help="Optional local TeleFuser parent PID; enables process-tree RSS and GPU memory sampling.", + ) + parser.add_argument("--gpu-indexes", type=parse_gpu_indexes, help="Optional physical GPU indexes to include.") + parser.add_argument("--resource-interval-seconds", type=float, default=1.0) + parser.add_argument("--max-resource-samples", type=int, default=10000) parser.add_argument("--output", required=True, type=Path) return parser.parse_args() From 41c02de18e6331628fb8d02f9976ba617eefdd02 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 10 Aug 2026 02:12:57 +0000 Subject: [PATCH 14/15] test(vla): register structured service contract Add the LingBot VLA v2 native service to the declared-contract coverage registry and validate the run entrypoint name declared by PIPELINE_CONTRACT instead of assuming run_with_file. Verification: 65 VLA, structured service, and registry tests; ruff check; ruff format; git diff --check. --- tests/unit/test_example_registry.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_example_registry.py b/tests/unit/test_example_registry.py index a56fb120..ff3ff87f 100644 --- a/tests/unit/test_example_registry.py +++ b/tests/unit/test_example_registry.py @@ -10,6 +10,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[2] EXAMPLES_ROOT = PROJECT_ROOT / "examples" SERVICE_PARITY_EXAMPLES = { + "lingbot_vla_v2/lingbot_vla_v2_native_service.py", "wan_video/wan21_14b_image_to_video_480p_service.py", "wan_video/wan22_14b_image_to_video_distill_h100.py", "lingbot_video/lingbot_video_dense_1_3b.py", @@ -38,6 +39,22 @@ def _declares_service_contract(path: Path) -> bool: return False +def _declared_run_entrypoint(path: Path) -> str: + tree = ast.parse(path.read_text(encoding="utf-8-sig"), filename=str(path)) + for node in tree.body: + targets = ( + node.targets if isinstance(node, ast.Assign) else [node.target] if isinstance(node, ast.AnnAssign) else [] + ) + if not any(isinstance(target, ast.Name) and target.id == "PIPELINE_CONTRACT" for target in targets): + continue + try: + contract = ast.literal_eval(node.value) + except (TypeError, ValueError): + break + return contract.get("entrypoints", {}).get("run_with_file", "run_with_file") + return "run_with_file" + + def test_example_regression_registry_has_runnable_entrypoints() -> None: config = load_config() @@ -78,6 +95,8 @@ def test_all_declared_service_examples_have_cpu_parity_coverage() -> None: assert declared_contract_examples == SERVICE_PARITY_EXAMPLES for script in declared_contract_examples: - symbols = _module_symbols(EXAMPLES_ROOT / script) + script_path = EXAMPLES_ROOT / script + symbols = _module_symbols(script_path) assert "get_pipeline" in symbols, f"{script} is missing get_pipeline()" - assert "run_with_file" in symbols, f"{script} is missing run_with_file()" + run_entrypoint = _declared_run_entrypoint(script_path) + assert run_entrypoint in symbols, f"{script} is missing {run_entrypoint}()" From 86278d4a22d35f7cd8606dddd80ae3a4637e396c Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 10 Aug 2026 03:51:12 +0000 Subject: [PATCH 15/15] feat(vla): add structured service resilience benchmarks Add a native LingBot VLA v2 AIPerf endpoint and HTTP polling transport with a pinned workload, bounded action summaries, contract assets, and documentation. Add real-service fault validation for malformed requests, cancellation, replica termination, GPU release, and graceful pool capacity degradation. Remove the replica cancel-forwarder completion delay and convert dead-process IPC failures into replica eviction without changing public service contracts. Verification: 93 focused VLA/service tests passed; 153 non-LiveKit service tests passed; 19 AIPerf adapter tests passed; real two-replica structured runs completed 100/100 requests; fault validation passed 5/5; ruff, formatting, config validation, shell syntax, and git diff checks passed. --- benchmarks/telefuser_aiperf/README.md | 31 +- .../configs/vla_structured_e2e.yaml | 49 ++ .../data/vla_structured.jsonl | 1 + .../scripts/run_vla_structured_bench.sh | 25 + .../telefuser_aiperf/__init__.py | 47 +- .../telefuser_aiperf/telefuser_aiperf/cli.py | 3 +- .../telefuser_aiperf/vla_structured.py | 385 +++++++++++++ .../tests/test_vla_structured.py | 123 +++++ .../vla_structured_contract.yaml | 39 ++ docs/en/benchmark_aiperf.md | 5 +- examples/lingbot_vla_v2/README.md | 31 +- telefuser/service/core/pipeline_pool.py | 7 +- telefuser/service/core/replica_worker.py | 21 +- tests/unit/service/test_pipeline_pool.py | 91 ++++ .../test_lingbot_vla_v2_service_faults.py | 131 +++++ .../validate_lingbot_vla_v2_service_faults.py | 505 ++++++++++++++++++ 16 files changed, 1476 insertions(+), 18 deletions(-) create mode 100644 benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml create mode 100644 benchmarks/telefuser_aiperf/data/vla_structured.jsonl create mode 100755 benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh create mode 100644 benchmarks/telefuser_aiperf/telefuser_aiperf/vla_structured.py create mode 100644 benchmarks/telefuser_aiperf/tests/test_vla_structured.py create mode 100644 benchmarks/telefuser_aiperf/vla_structured_contract.yaml create mode 100644 tests/unit/validation/test_lingbot_vla_v2_service_faults.py create mode 100644 tools/validation/validate_lingbot_vla_v2_service_faults.py diff --git a/benchmarks/telefuser_aiperf/README.md b/benchmarks/telefuser_aiperf/README.md index ae261644..93f93258 100644 --- a/benchmarks/telefuser_aiperf/README.md +++ b/benchmarks/telefuser_aiperf/README.md @@ -92,6 +92,31 @@ Available batch configs: | `configs/video_generation_rate.yaml` | Poisson-arrival load | | `configs/video_generation_wan21_i2v_480p_compare.yaml` | Fixed Wan2.1 I2V comparison | +## LingBot-VLA v2 Structured Actions + +Start the native VLA service from its isolated model environment, then run the AIPerf workload from the repository +root: + +```bash +bash benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh +``` + +The repository-owned `telefuser_vla_structured` endpoint and `telefuser_structured_http` transport submit +`POST /v1/tasks/structured`, poll `GET /v1/tasks/{task_id}/status`, and pass request latency, throughput, success, +trace, and server metric facts into AIPerf's normal warmup and aggregation pipeline. Defaults are two excluded warmup +requests followed by 20 measured requests at concurrency one. Override them without changing the checked-in config: + +```bash +TELEFUSER_AIPERF_REQUESTS=100 \ +TELEFUSER_AIPERF_CONCURRENCY=2 \ + bash benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh +``` + +Each terminal result is required to contain a finite `50x55` action chunk and the frozen structured result fields. +The adapter retains an action hash, bounds, dimensions, verification status, target inference time, and peak memory; +it does not copy full action arrays or Base64 cameras into AIPerf response records. This validates service execution +and normalized action structure, not physical robot control semantics. + ## LingBot-World v2 Streaming The v2 pipeline expects the following files below `TF_MODEL_ZOO_PATH`: @@ -264,10 +289,12 @@ AIPerf environment first, then run the checks from the repository root: PYTHONPATH=benchmarks/telefuser_aiperf \ .venv-aiperf/bin/python -m pytest \ benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py \ - benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py + benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py \ + benchmarks/telefuser_aiperf/tests/test_vla_structured.py bash -n \ scripts/setup_aiperf.sh \ benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh \ - benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh + benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh \ + benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh ``` diff --git a/benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml b/benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml new file mode 100644 index 00000000..3053385b --- /dev/null +++ b/benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml @@ -0,0 +1,49 @@ +# yaml-language-server: $schema=../../aiperf/src/aiperf/config/schema/aiperf-config.schema.json + +schemaVersion: "2.0" + +randomSeed: 42 + +benchmark: + model: lingbot-vla-v2-6b + + endpoint: + url: ${TELEFUSER_AIPERF_URL:http://127.0.0.1:18080} + type: telefuser_vla_structured + transport: telefuser_structured_http + timeout: ${TELEFUSER_AIPERF_TIMEOUT:120} + + tokenizer: + name: builtin + + dataset: + type: file + path: ./benchmarks/telefuser_aiperf/data/vla_structured.jsonl + format: single_turn + sampling: sequential + + warmup: + type: concurrency + concurrency: 1 + requests: ${TELEFUSER_AIPERF_WARMUP_REQUESTS:2} + excludeFromResults: true + + profiling: + type: concurrency + concurrency: ${TELEFUSER_AIPERF_CONCURRENCY:1} + requests: ${TELEFUSER_AIPERF_REQUESTS:20} + duration: ${TELEFUSER_AIPERF_DURATION:3600} + gracePeriod: ${TELEFUSER_AIPERF_GRACE_PERIOD:120} + + artifacts: + dir: ./artifacts/telefuser_aiperf/vla_structured + summary: [json] + records: [jsonl] + showTraceTiming: true + trace: true + + serverMetrics: + enabled: ${TELEFUSER_AIPERF_SERVER_METRICS:true} + urls: + - ${TELEFUSER_AIPERF_METRICS_URL:http://127.0.0.1:18080/v1/service/metrics} + formats: [json, csv] diff --git a/benchmarks/telefuser_aiperf/data/vla_structured.jsonl b/benchmarks/telefuser_aiperf/data/vla_structured.jsonl new file mode 100644 index 00000000..eb2fbaff --- /dev/null +++ b/benchmarks/telefuser_aiperf/data/vla_structured.jsonl @@ -0,0 +1 @@ +{"text":"pick up the object","image":"examples/data/101235-video-720_0.png","extra":{"state":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],"seed":7}} diff --git a/benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh b/benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh new file mode 100755 index 00000000..adf86b15 --- /dev/null +++ b/benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${ROOT_DIR}" + +CONFIG_PATH="${1:-benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml}" +SERVER_URL="${TELEFUSER_AIPERF_URL:-http://127.0.0.1:18080}" +HEALTH_URL="${TELEFUSER_AIPERF_HEALTH_URL:-${SERVER_URL}/v1/service/ready}" +DEFAULT_PYTHON="${ROOT_DIR}/.venv-aiperf/bin/python" +ADAPTER_ROOT="${ROOT_DIR}/benchmarks/telefuser_aiperf" +AIPERF_PYTHON="${TELEFUSER_AIPERF_PYTHON:-${DEFAULT_PYTHON}}" + +if [[ ! -x "${AIPERF_PYTHON}" ]]; then + echo "The isolated AIPerf environment is unavailable. Run: bash scripts/setup_aiperf.sh" >&2 + exit 1 +fi + +if command -v curl >/dev/null 2>&1; then + echo "Checking TeleFuser VLA readiness: ${HEALTH_URL}" + curl --noproxy '*' --fail --silent --show-error "${HEALTH_URL}" >/dev/null +fi + +export PYTHONPATH="${ADAPTER_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +exec "${AIPERF_PYTHON}" -m telefuser_aiperf.cli profile --config "${CONFIG_PATH}" diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py index 7834d47c..2a3b3473 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py @@ -6,6 +6,12 @@ from telefuser_aiperf.adapter import TeleFuserLiveKitAdapter from telefuser_aiperf.sglang_adapter import SGLangRealtimeAdapter +from telefuser_aiperf.vla_structured import ( + ENDPOINT_METADATA, + TRANSPORT_METADATA, + TeleFuserStructuredHttpTransport, + TeleFuserVlaStructuredEndpoint, +) def register_adapters(*, replace: bool = False) -> None: @@ -23,4 +29,43 @@ def register_adapters(*, replace: bool = False) -> None: ) -__all__ = ["SGLangRealtimeAdapter", "TeleFuserLiveKitAdapter", "register_adapters"] +def register_plugins(*, replace: bool = False) -> None: + """Register repository-owned AIPerf batch endpoint and transport plugins.""" + from aiperf.plugin import plugins + from aiperf.plugin.enums import EndpointType, TransportType + + if "telefuser_vla_structured" not in EndpointType: + EndpointType.register("TELEFUSER_VLA_STRUCTURED", "telefuser_vla_structured") + if "telefuser_structured_http" not in TransportType: + TransportType.register("TELEFUSER_STRUCTURED_HTTP", "telefuser_structured_http") + + definitions = ( + ( + "endpoint", + "telefuser_vla_structured", + TeleFuserVlaStructuredEndpoint, + ENDPOINT_METADATA, + ), + ( + "transport", + "telefuser_structured_http", + TeleFuserStructuredHttpTransport, + TRANSPORT_METADATA, + ), + ) + for category, name, plugin_class, metadata in definitions: + if plugins.has_entry(category, name): + if not replace: + continue + plugins.unregister(category, name) + plugins.register(category, name, plugin_class, metadata=metadata) + + +__all__ = [ + "SGLangRealtimeAdapter", + "TeleFuserLiveKitAdapter", + "TeleFuserStructuredHttpTransport", + "TeleFuserVlaStructuredEndpoint", + "register_adapters", + "register_plugins", +] diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py index 3ad1f4be..de77e0a8 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py @@ -6,9 +6,10 @@ def main() -> None: """Register TeleFuser adapters and delegate to the AIPerf CLI.""" - from telefuser_aiperf import register_adapters + from telefuser_aiperf import register_adapters, register_plugins register_adapters() + register_plugins() from aiperf.cli import app diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/vla_structured.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/vla_structured.py new file mode 100644 index 00000000..a7a64697 --- /dev/null +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/vla_structured.py @@ -0,0 +1,385 @@ +"""AIPerf endpoint and HTTP polling transport for TeleFuser VLA actions.""" + +from __future__ import annotations + +import asyncio +import base64 +import binascii +import hashlib +import math +import struct +import time +from dataclasses import dataclass +from typing import Any +from urllib.parse import quote, urlsplit, urlunsplit + +import orjson +from aiperf.common.exceptions import NotInitializedError +from aiperf.common.models import ( + BaseResponseData, + ErrorDetails, + InferenceServerResponse, + ParsedResponse, + RequestInfo, + RequestRecord, + TextResponse, +) +from aiperf.endpoints.base_endpoint import BaseEndpoint +from aiperf.plugin.schema.schemas import TransportMetadata +from aiperf.transports.aiohttp_transport import AioHttpTransport + +_RESULT_FIELDS = frozenset( + { + "canonical_normalized_actions", + "horizon", + "action_dim", + "checkpoint_variant", + "policy_verified", + "verification_status", + } +) +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) +_EXPECTED_HORIZON = 50 +_EXPECTED_ACTION_DIM = 55 +_POLL_INTERVAL_SECONDS = 0.05 + + +@dataclass(slots=True) +class VlaActionResponseData(BaseResponseData): + """Validated summary of one VLA action chunk.""" + + task_id: str + horizon: int + action_dim: int + value_count: int + sha256_float64_le: str + checkpoint_variant: str + policy_verified: bool + verification_status: str + inference_time_s: float | None = None + peak_memory_mb: float | None = None + + +def _finite_number(value: Any, *, name: str, allow_none: bool = False) -> float | None: + if value is None and allow_none: + return None + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError(f"{name} must be a finite number") + normalized = float(value) + if not math.isfinite(normalized): + raise ValueError(f"{name} must be a finite number") + return normalized + + +def validate_state(state: Any) -> list[float]: + """Validate the canonical 14-dimensional VLA state vector.""" + if not isinstance(state, list) or len(state) != 14: + raise ValueError("VLA state must contain exactly 14 values") + return [float(_finite_number(value, name="VLA state value")) for value in state] + + +def image_content_to_base64(content: str) -> str: + """Normalize an AIPerf image data URL to raw validated base64.""" + if content.lower().startswith(("http://", "https://")): + raise ValueError("TeleFuser VLA camera inputs must be inline image data, not URLs") + encoded = content + if content.startswith("data:"): + try: + header, encoded = content.split(",", 1) + except ValueError as error: + raise ValueError("VLA image data URL is missing a comma") from error + if ";base64" not in header.lower() or not header.lower().startswith("data:image/"): + raise ValueError("VLA camera input must be a base64 image data URL") + try: + decoded = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError) as error: + raise ValueError("VLA camera input is not valid base64") from error + if not decoded: + raise ValueError("VLA camera input is empty") + return encoded + + +def build_vla_payload( + instruction: str, + image_content: str, + *, + extra: dict[str, Any] | None, +) -> dict[str, Any]: + """Build the stable TeleFuser structured action request body.""" + if not instruction.strip(): + raise ValueError("VLA instruction must not be empty") + parameters = dict(extra or {}) + unsupported = sorted(set(parameters).difference({"state", "seed"})) + if unsupported: + raise ValueError(f"Unsupported VLA request fields: {', '.join(unsupported)}") + if "state" not in parameters: + raise ValueError("VLA dataset entry must provide state in extra") + state = validate_state(parameters["state"]) + seed = parameters.get("seed", 7) + if isinstance(seed, bool) or not isinstance(seed, int): + raise ValueError("VLA seed must be an integer") + image_base64 = image_content_to_base64(image_content) + return { + "task": "vla_action", + "instruction": instruction, + "state": state, + "camera_high": image_base64, + "camera_left_wrist": image_base64, + "camera_right_wrist": image_base64, + "seed": seed, + } + + +def summarize_action_result(result: Any) -> dict[str, Any]: + """Validate a 50 x 55 normalized action chunk and return bounded facts.""" + if not isinstance(result, dict) or set(result) != set(_RESULT_FIELDS): + observed = sorted(result) if isinstance(result, dict) else type(result).__name__ + raise ValueError(f"VLA result fields changed: {observed}") + actions = result.get("canonical_normalized_actions") + if not isinstance(actions, list) or len(actions) != _EXPECTED_HORIZON: + raise ValueError(f"VLA action horizon must be {_EXPECTED_HORIZON}") + if result.get("horizon") != _EXPECTED_HORIZON or result.get("action_dim") != _EXPECTED_ACTION_DIM: + raise ValueError("VLA action dimension metadata changed") + + digest = hashlib.sha256() + value_count = 0 + minimum = math.inf + maximum = -math.inf + for row_index, row in enumerate(actions): + if not isinstance(row, list) or len(row) != _EXPECTED_ACTION_DIM: + raise ValueError(f"VLA action row {row_index} must contain {_EXPECTED_ACTION_DIM} values") + for raw_value in row: + value = float(_finite_number(raw_value, name="VLA action value")) + digest.update(struct.pack(" str: + """Build the native task status URL on the same origin as submission.""" + parsed = urlsplit(submit_url) + path = f"/v1/tasks/{quote(task_id, safe='')}/status" + return urlunsplit((parsed.scheme, parsed.netloc, path, "", "")) + + +class TeleFuserVlaStructuredEndpoint(BaseEndpoint): + """Format and parse TeleFuser's native VLA structured API.""" + + def format_payload(self, request_info: RequestInfo) -> dict[str, Any]: + if not request_info.turns: + raise ValueError("TeleFuser VLA endpoint requires one dataset turn") + turn = request_info.turns[-1] + if not turn.texts or not turn.texts[0].contents: + raise ValueError("TeleFuser VLA endpoint requires one instruction") + if not turn.images or not turn.images[0].contents: + raise ValueError("TeleFuser VLA endpoint requires one camera image") + merged_extra = dict(request_info.model_endpoint.endpoint.extra or {}) + merged_extra.update(turn.extra_body or {}) + return build_vla_payload( + turn.texts[0].contents[0], + turn.images[0].contents[0], + extra=merged_extra, + ) + + def parse_response(self, response: InferenceServerResponse) -> ParsedResponse | None: + body = response.get_json() + if not isinstance(body, dict) or body.get("status") != "completed": + return None + summary = body.get("action_summary") + if not isinstance(summary, dict): + raise ValueError("completed VLA benchmark response has no action_summary") + data = VlaActionResponseData( + task_id=str(body["task_id"]), + horizon=int(summary["horizon"]), + action_dim=int(summary["action_dim"]), + value_count=int(summary["value_count"]), + sha256_float64_le=str(summary["sha256_float64_le"]), + checkpoint_variant=str(summary["checkpoint_variant"]), + policy_verified=bool(summary["policy_verified"]), + verification_status=str(summary["verification_status"]), + inference_time_s=body.get("inference_time_s"), + peak_memory_mb=body.get("peak_memory_mb"), + ) + return ParsedResponse(perf_ns=response.perf_ns, data=data, metadata={"media_type": "structured"}) + + +class TeleFuserStructuredHttpTransport(AioHttpTransport): + """HTTP JSON transport for TeleFuser submit/poll structured tasks.""" + + @classmethod + def metadata(cls) -> TransportMetadata: + return TransportMetadata(transport_type="telefuser_structured_http", url_schemes=[]) + + @staticmethod + def _parse_json_record(record: RequestRecord, context: str) -> tuple[dict[str, Any], TextResponse] | ErrorDetails: + if record.error: + return record.error + if not record.responses or not isinstance(record.responses[0], TextResponse): + return ErrorDetails(type="VlaStructuredError", message=f"No JSON response from {context}", code=500) + response = record.responses[0] + try: + body = orjson.loads(response.text) + except orjson.JSONDecodeError: + return ErrorDetails(type="VlaStructuredError", message=f"Invalid JSON from {context}", code=500) + if not isinstance(body, dict): + return ErrorDetails(type="VlaStructuredError", message=f"Non-object JSON from {context}", code=500) + return body, response + + async def send_request( + self, + request_info: RequestInfo, + payload: dict[str, Any], + *, + first_token_callback: Any = None, + ) -> RequestRecord: + """Submit one action task, poll terminal state, and retain bounded facts.""" + del first_token_callback + if self.aiohttp_client is None: + raise NotInitializedError("AioHttpClient not initialized") + start_ns = time.perf_counter_ns() + headers = self.build_headers(request_info) + responses: list[TextResponse] = [] + + def make_record(error: ErrorDetails | None = None, status: int | None = None) -> RequestRecord: + return RequestRecord( + request_info=request_info, + request_headers=headers, + start_perf_ns=start_ns, + end_perf_ns=time.perf_counter_ns(), + responses=responses, + error=error, + status=status, + ) + + try: + submit_url = self.build_url(request_info) + submitted = await self.aiohttp_client.post_request(submit_url, orjson.dumps(payload), headers) + parsed_submit = self._parse_json_record(submitted, "VLA task submission") + if isinstance(parsed_submit, ErrorDetails): + return make_record(error=parsed_submit, status=submitted.status) + submit_body, submit_response = parsed_submit + task_id = submit_body.get("task_id") + if not isinstance(task_id, str) or not task_id: + return make_record( + error=ErrorDetails( + type="VlaStructuredError", + message="VLA submission returned no task_id", + code=500, + ) + ) + responses.append( + TextResponse( + perf_ns=submit_response.perf_ns, + text=orjson.dumps({"task_id": task_id, "status": "pending"}).decode(), + content_type="application/json", + ) + ) + + status_url = build_task_status_url(submit_url, task_id) + timeout = request_info.model_endpoint.endpoint.timeout + deadline = time.monotonic() + timeout if timeout > 0 else math.inf + while time.monotonic() < deadline: + polled = await self.aiohttp_client.get_request(status_url, headers) + parsed_poll = self._parse_json_record(polled, "VLA task status") + if isinstance(parsed_poll, ErrorDetails): + return make_record(error=parsed_poll, status=polled.status) + status_body, status_response = parsed_poll + status = status_body.get("status") or status_body.get("task_status") + if status not in _TERMINAL_STATUSES: + await asyncio.sleep(_POLL_INTERVAL_SECONDS) + continue + if status != "completed": + return make_record( + error=ErrorDetails( + type="VlaStructuredError", + message=f"VLA task {task_id} ended with {status}: {status_body.get('error')}", + code=500, + ), + status=polled.status, + ) + action_summary = summarize_action_result(status_body.get("result")) + bounded_status = { + "task_id": task_id, + "status": "completed", + "inference_time_s": _finite_number( + status_body.get("inference_time_s"), name="inference_time_s", allow_none=True + ), + "peak_memory_mb": _finite_number( + status_body.get("peak_memory_mb"), name="peak_memory_mb", allow_none=True + ), + "action_summary": action_summary, + } + responses.append( + TextResponse( + perf_ns=status_response.perf_ns, + text=orjson.dumps(bounded_status).decode(), + content_type="application/json", + ) + ) + return make_record(status=200) + return make_record( + error=ErrorDetails( + type="TimeoutError", + message=f"VLA task {task_id} timed out after {timeout:g}s", + code=504, + ), + status=504, + ) + except asyncio.CancelledError: + raise + except Exception as error: + return make_record(error=ErrorDetails.from_exception(error)) + + +ENDPOINT_METADATA = { + "endpoint_path": "/v1/tasks/structured", + "supports_streaming": False, + "tokenizes_input": False, + "produces_tokens": False, + "supports_images": True, + "requires_polling": True, + "requires_form_data": False, + "metrics_title": "TeleFuser VLA Structured Metrics", + "service_kind": "telefuser_vla", +} + +TRANSPORT_METADATA = { + "transport_type": "telefuser_structured_http", + "url_schemes": [], +} + + +__all__ = [ + "ENDPOINT_METADATA", + "TRANSPORT_METADATA", + "TeleFuserStructuredHttpTransport", + "TeleFuserVlaStructuredEndpoint", + "VlaActionResponseData", + "build_task_status_url", + "build_vla_payload", + "image_content_to_base64", + "summarize_action_result", + "validate_state", +] diff --git a/benchmarks/telefuser_aiperf/tests/test_vla_structured.py b/benchmarks/telefuser_aiperf/tests/test_vla_structured.py new file mode 100644 index 00000000..1ca58d81 --- /dev/null +++ b/benchmarks/telefuser_aiperf/tests/test_vla_structured.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import base64 +import math + +import orjson +import pytest +from aiperf.common.models import TextResponse +from aiperf.plugin import plugins +from telefuser_aiperf import register_plugins +from telefuser_aiperf.vla_structured import ( + TeleFuserStructuredHttpTransport, + TeleFuserVlaStructuredEndpoint, + VlaActionResponseData, + build_task_status_url, + build_vla_payload, + summarize_action_result, +) + + +def _action_result(value: float = 0.25) -> dict: + return { + "canonical_normalized_actions": [[value] * 55 for _ in range(50)], + "horizon": 50, + "action_dim": 55, + "checkpoint_variant": "base", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + } + + +def test_registration_uses_aiperf_endpoint_and_transport_plugins() -> None: + register_plugins(replace=True) + + endpoint_class = plugins.get_class("endpoint", "telefuser_vla_structured") + transport_class = plugins.get_class("transport", "telefuser_structured_http") + + assert endpoint_class is TeleFuserVlaStructuredEndpoint + assert transport_class is TeleFuserStructuredHttpTransport + assert plugins.get_endpoint_metadata("telefuser_vla_structured").requires_polling is True + + +def test_build_vla_payload_reuses_inline_image_for_three_cameras() -> None: + encoded = base64.b64encode(b"image bytes").decode() + + payload = build_vla_payload( + "pick up the object", + f"data:image/png;base64,{encoded}", + extra={"state": [0.0] * 14, "seed": 7}, + ) + + assert payload["task"] == "vla_action" + assert payload["camera_high"] == encoded + assert payload["camera_left_wrist"] == encoded + assert payload["camera_right_wrist"] == encoded + assert payload["state"] == [0.0] * 14 + + +@pytest.mark.parametrize( + "extra,match", + [ + ({"state": [0.0] * 13}, "exactly 14"), + ({"state": [0.0] * 13 + [math.nan]}, "finite"), + ({"state": [0.0] * 14, "unknown": 1}, "Unsupported"), + ], +) +def test_build_vla_payload_rejects_contract_drift(extra: dict, match: str) -> None: + encoded = base64.b64encode(b"image bytes").decode() + + with pytest.raises(ValueError, match=match): + build_vla_payload("instruction", encoded, extra=extra) + + +def test_summarize_action_result_validates_shape_and_omits_full_actions() -> None: + summary = summarize_action_result(_action_result()) + + assert summary["horizon"] == 50 + assert summary["action_dim"] == 55 + assert summary["value_count"] == 2750 + assert summary["minimum"] == 0.25 + assert len(summary["sha256_float64_le"]) == 64 + assert "canonical_normalized_actions" not in summary + + +def test_summarize_action_result_rejects_non_finite_action() -> None: + result = _action_result() + result["canonical_normalized_actions"][0][0] = math.inf + + with pytest.raises(ValueError, match="finite"): + summarize_action_result(result) + + +def test_endpoint_parses_bounded_completed_response() -> None: + endpoint = object.__new__(TeleFuserVlaStructuredEndpoint) + summary = summarize_action_result(_action_result()) + response = TextResponse( + perf_ns=123, + content_type="application/json", + text=orjson.dumps( + { + "task_id": "task-1", + "status": "completed", + "inference_time_s": 0.65, + "peak_memory_mb": None, + "action_summary": summary, + } + ).decode(), + ) + + parsed = endpoint.parse_response(response) + + assert parsed is not None + assert isinstance(parsed.data, VlaActionResponseData) + assert parsed.data.task_id == "task-1" + assert parsed.data.value_count == 2750 + assert parsed.metadata == {"media_type": "structured"} + + +def test_task_status_url_uses_native_structured_route() -> None: + assert ( + build_task_status_url("http://127.0.0.1:18080/v1/tasks/structured", "task id") + == "http://127.0.0.1:18080/v1/tasks/task%20id/status" + ) diff --git a/benchmarks/telefuser_aiperf/vla_structured_contract.yaml b/benchmarks/telefuser_aiperf/vla_structured_contract.yaml new file mode 100644 index 00000000..81ca0cc8 --- /dev/null +++ b/benchmarks/telefuser_aiperf/vla_structured_contract.yaml @@ -0,0 +1,39 @@ +contract_version: v1 +name: telefuser_lingbot_vla_v2_structured +mode: structured_action +implementation: telefuser +model_family: lingbot_vla_v2 +model: lingbot-vla-v2-6b +supported_tasks: + - vla_action +transport: http_polling +endpoint: + submit_path: /v1/tasks/structured + status_path: /v1/tasks/{task_id}/status + protocol: telefuser_structured_task +request_encoding: + content_type: application/json + parameters: + instruction: text + state: extra.state + camera_high: image + camera_left_wrist: image + camera_right_wrist: image + seed: extra.seed +result_delivery: + terminal_status: completed + result_field: result + action_field: canonical_normalized_actions + expected_shape: [50, 55] +workload: + warmup_requests: 2 + profile_requests: 20 + concurrency: 1 +metrics: + - request_latency + - request_throughput + - success_rate + - server_metrics +artifacts: + config: benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml + dataset: benchmarks/telefuser_aiperf/data/vla_structured.jsonl diff --git a/docs/en/benchmark_aiperf.md b/docs/en/benchmark_aiperf.md index 94c1d143..d7dc2067 100644 --- a/docs/en/benchmark_aiperf.md +++ b/docs/en/benchmark_aiperf.md @@ -2,8 +2,8 @@ TeleFuser exposes raw target-side facts; AIPerf owns workload execution, aggregation, resource collection, artifacts, GreptimeDB history, and visualization. The checked-in integration covers batch video generation through the -OpenAI-compatible `/v1/videos` API, TeleFuser LingBot streaming through LiveKit, and SGLang LingBot streaming through -its native realtime WebSocket endpoint. +OpenAI-compatible `/v1/videos` API, LingBot-VLA structured actions through native HTTP task polling, TeleFuser LingBot +streaming through LiveKit, and SGLang LingBot streaming through its native realtime WebSocket endpoint. AIPerf's stream runner and result schema are transport-neutral. The LiveKit adapter is maintained by TeleFuser, loads from source at process startup, and produces AIPerf's standard session results. The contract records WebRTC as @@ -80,6 +80,7 @@ parity comparisons. See the benchmark README for model, GPU, port, and executabl |---|---|---| | TeleFuser runtime | TeleFuser | Emit synchronized phase, chunk, runtime, cache, and environment facts | | Batch target adapter | AIPerf | Convert `/v1/videos` HTTP events into the standard request timeline | +| VLA structured adapter | TeleFuser | Validate action results and convert native submit/poll events into bounded AIPerf records | | LiveKit source adapter | TeleFuser | Convert room, track, status, metrics, and control events into session results | | SGLang source adapter | TeleFuser | Convert MessagePack frames, chunk timings, and camera events into session results | | Aggregation and history | AIPerf | Apply warmup, percentiles, throughput, artifacts, GreptimeDB, and visualization | diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 31d958d1..918c0090 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -254,11 +254,32 @@ excluded from the artifact. `--max-records` bounds retained per-request samples latency and success counters still cover the complete run. `--max-resource-samples` independently bounds retained resource samples. -This validation aligns VLA with the repository's existing deployment practice, but it is not yet an AIPerf workload: -AIPerf currently has maintained adapters for batch media and LingBot streaming transports, not the asynchronous -structured task API. The JSON report therefore keeps target inference time separate from client end-to-end time and -preserves target metadata and raw service metric snapshots so a future transport adapter can consume the same facts. -Passing this check proves serving and normalized action structure, not embodiment-specific control semantics. +For fault handling, run the independent validator against a ready service: + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_service_faults.py \ + --base-url http://127.0.0.1:18080 \ + --image examples/data/lingbot_world_fast/image.jpg +``` + +It checks missing cameras, invalid state size, invalid Base64, and cancellation. Replica termination is opt-in and +requires a disposable two-replica service: add `--service-pid ` and +`--kill-replica-gpu-index `. The tool only selects a GPU compute process inside that parent process +tree, sends `SIGTERM`, and verifies one-replica capacity degradation plus a subsequent valid `50x55` response. It does +not promise automatic replica restart. + +The same structured API is available through the repository-owned AIPerf workload. Install the pinned isolated +AIPerf environment once, then run the workload while the native service is ready: + +```bash +bash scripts/setup_aiperf.sh +bash benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh +``` + +AIPerf excludes the configured warmup, aggregates request latency, throughput, success, traces, and server metrics, +and writes normal AIPerf artifacts. The adapter strictly validates the action contract but retains only bounded action +facts, not full arrays or Base64 inputs. Passing either validator proves serving and normalized action structure, not +embodiment-specific control semantics. ## TeleFuser Regression Baseline diff --git a/telefuser/service/core/pipeline_pool.py b/telefuser/service/core/pipeline_pool.py index e5081852..74b982fb 100644 --- a/telefuser/service/core/pipeline_pool.py +++ b/telefuser/service/core/pipeline_pool.py @@ -207,8 +207,11 @@ async def acquire(self) -> AsyncIterator[ReplicaHandle]: continue handle = self._handles[idx] - if handle._dead: - self._evict_replica(idx, "pre-existing dead state") + process = getattr(handle, "process", None) + process_dead = process is not None and not process.is_alive() + if handle._dead or process_dead: + reason = "process exited" if process_dead else "pre-existing dead state" + self._evict_replica(idx, reason) continue break diff --git a/telefuser/service/core/replica_worker.py b/telefuser/service/core/replica_worker.py index abb7de60..61945089 100644 --- a/telefuser/service/core/replica_worker.py +++ b/telefuser/service/core/replica_worker.py @@ -197,13 +197,16 @@ def _forward_cancel_fn( cancel_event: mp_stdlib.Event, stop_event: threading.Event, forwarder_exit: threading.Event, + forwarder_wake: threading.Event, forwarder_done: threading.Event, ) -> None: """Forward main-process stop_event to subprocess cancel_event (polled).""" while not forwarder_exit.is_set(): - if stop_event.wait(timeout=0.5): + if stop_event.is_set(): cancel_event.set() break + forwarder_wake.wait(timeout=0.5) + forwarder_wake.clear() forwarder_done.set() @@ -236,24 +239,32 @@ async def run_task( self.cancel_event.clear() forwarder_exit = threading.Event() + forwarder_wake = threading.Event() forwarder_done = threading.Event() forwarder = threading.Thread( target=_forward_cancel_fn, - args=(self.cancel_event, stop_event, forwarder_exit, forwarder_done), + args=(self.cancel_event, stop_event, forwarder_exit, forwarder_wake, forwarder_done), daemon=True, ) forwarder.start() loop = asyncio.get_running_loop() - self.conn.send(("task", task_data, timeout_s, output_root)) - ipc_timeout = (timeout_s or 600) + _TASK_IPC_MARGIN_S try: - result = await loop.run_in_executor(None, self._recv_with_health_check, ipc_timeout) + if not self.process.is_alive(): + self._dead = True + raise ReplicaDeadError(f"Replica {self.replica_id} process is not alive") + try: + self.conn.send(("task", task_data, timeout_s, output_root)) + result = await loop.run_in_executor(None, self._recv_with_health_check, ipc_timeout) + except (EOFError, OSError) as error: + self._dead = True + raise ReplicaDeadError(f"Replica {self.replica_id} IPC failed: {error}") from error finally: forwarder_exit.set() + forwarder_wake.set() forwarder_done.wait(2.0) if result is None: diff --git a/tests/unit/service/test_pipeline_pool.py b/tests/unit/service/test_pipeline_pool.py index ff34e06e..e9c7f1cb 100644 --- a/tests/unit/service/test_pipeline_pool.py +++ b/tests/unit/service/test_pipeline_pool.py @@ -16,11 +16,102 @@ from telefuser.platforms import current_platform from telefuser.service.api.schema import TaskRequest from telefuser.service.core.config import ServerConfig +from telefuser.service.core.pipeline_pool import PipelinePool +from telefuser.service.core.replica_worker import ReplicaDeadError, ReplicaHandle, _forward_cancel_fn from telefuser.service.core.task_manager import TaskManager, TaskStatus _DEVICE_ENV_VAR = current_platform.device_control_env_var +def test_cancel_forwarder_wakes_immediately_on_normal_completion() -> None: + cancel_event = threading.Event() + stop_event = threading.Event() + forwarder_exit = threading.Event() + forwarder_wake = threading.Event() + forwarder_done = threading.Event() + forwarder = threading.Thread( + target=_forward_cancel_fn, + args=(cancel_event, stop_event, forwarder_exit, forwarder_wake, forwarder_done), + daemon=True, + ) + forwarder.start() + + forwarder_exit.set() + forwarder_wake.set() + + assert forwarder_done.wait(0.2) + assert not cancel_event.is_set() + assert not stop_event.is_set() + + +def test_cancel_forwarder_preserves_request_cancellation() -> None: + cancel_event = threading.Event() + stop_event = threading.Event() + stop_event.set() + forwarder_done = threading.Event() + + _forward_cancel_fn( + cancel_event, + stop_event, + threading.Event(), + threading.Event(), + forwarder_done, + ) + + assert cancel_event.is_set() + assert forwarder_done.is_set() + + +def test_replica_handle_converts_broken_pipe_to_dead_replica() -> None: + process = MagicMock() + process.is_alive.return_value = True + connection = MagicMock() + connection.send.side_effect = BrokenPipeError("closed") + handle = ReplicaHandle( + replica_id=0, + process=process, + conn=connection, + cancel_event=threading.Event(), + metadata={}, + ) + + with pytest.raises(ReplicaDeadError, match="IPC failed"): + asyncio.run(handle.run_task({}, threading.Event(), None, None)) + + assert handle._dead is True + + +def test_pipeline_pool_evicts_exited_replica_and_uses_remaining_capacity() -> None: + task_manager = MagicMock() + pool = PipelinePool( + num_replicas=2, + replica_device_ids=[["0"], ["1"]], + security_level_name="NONE", + task_manager=task_manager, + ) + dead_handle = MagicMock() + dead_handle._dead = False + dead_handle.process.is_alive.return_value = False + live_handle = MagicMock() + live_handle._dead = False + live_handle.process.is_alive.return_value = True + pool._handles = [dead_handle, live_handle] + pool._instance_status = ["idle", "idle"] + pool._available.put_nowait(0) + pool._available.put_nowait(1) + + async def scenario() -> None: + async with pool.acquire() as handle: + assert handle is live_handle + + asyncio.run(scenario()) + + assert pool._live_count == 1 + assert pool._instance_status == ["dead", "idle"] + dead_handle.shutdown.assert_called_once() + task_manager.set_max_concurrent_processing.assert_called_once_with(1) + + # ============================================================================ # Test 1: TaskManager atomic claim — concurrent claim, single winner # ============================================================================ diff --git a/tests/unit/validation/test_lingbot_vla_v2_service_faults.py b/tests/unit/validation/test_lingbot_vla_v2_service_faults.py new file mode 100644 index 00000000..cd05f32d --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_service_faults.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from tools.validation import validate_lingbot_vla_v2_service_faults as validator + + +class _Response: + def __init__(self, status_code: int, body: dict[str, Any]) -> None: + self.status_code = status_code + self._body = body + + def json(self) -> dict[str, Any]: + return self._body + + +class _Session: + def __init__(self, responses: list[_Response]) -> None: + self.responses = responses + + def request(self, *args: object, **kwargs: object) -> _Response: + return self.responses.pop(0) + + +def test_expect_invalid_request_accepts_synchronous_rejection() -> None: + result = validator._expect_rejected_or_failed( + _Session([_Response(422, {"detail": "invalid"})]), + "http://127.0.0.1:18080", + {"task": "vla_action"}, + case_name="invalid", + http_timeout_seconds=1.0, + task_timeout_seconds=1.0, + poll_interval_seconds=0.001, + ) + + assert result == {"name": "invalid", "passed": True, "handling": "rejected", "http_status": 422} + + +def test_expect_invalid_request_accepts_asynchronous_failure() -> None: + session = _Session( + [ + _Response(200, {"task_id": "task-1", "task_status": "pending"}), + _Response(200, {"task_id": "task-1", "status": "failed", "error": "bad input"}), + ] + ) + + result = validator._expect_rejected_or_failed( + session, + "http://127.0.0.1:18080", + {"task": "vla_action"}, + case_name="invalid", + http_timeout_seconds=1.0, + task_timeout_seconds=1.0, + poll_interval_seconds=0.001, + ) + + assert result["handling"] == "asynchronous_failure" + assert result["terminal_status"] == "failed" + assert result["error"] == "bad input" + + +def test_select_replica_process_filters_unrelated_and_root_processes() -> None: + rows = "\n".join( + [ + "100, GPU-a", + "101, GPU-a", + "102, GPU-b", + "999, GPU-a", + ] + ) + + selected = validator.select_replica_process( + rows, + service_process_ids={101, 102}, + gpu_uuid_to_index={"GPU-a": "0", "GPU-b": "1"}, + gpu_index="0", + service_pid=100, + ) + + assert selected == 101 + + +def test_select_replica_process_requires_unambiguous_target() -> None: + with pytest.raises(validator.FaultValidationFailure, match="exactly one"): + validator.select_replica_process( + "101, GPU-a\n102, GPU-a", + service_process_ids={101, 102}, + gpu_uuid_to_index={"GPU-a": "0"}, + gpu_index="0", + service_pid=100, + ) + + +def test_gpu_compute_process_ids_ignores_malformed_rows() -> None: + rows = "101, GPU-a\ninvalid, GPU-b\n102, GPU-c\n" + + assert validator.gpu_compute_process_ids(rows) == {101, 102} + + +def test_validate_pool_degradation_requires_one_dead_replica_and_reduced_capacity() -> None: + before = { + "effective_max_concurrent_tasks": 2, + "pool": [{"id": 0, "status": "idle"}, {"id": 1, "status": "idle"}], + } + after = { + "effective_max_concurrent_tasks": 1, + "pool": [{"id": 0, "status": "dead"}, {"id": 1, "status": "idle"}], + } + + result = validator.validate_pool_degradation(before, after) + + assert result["before_capacity"] == 2 + assert result["after_capacity"] == 1 + assert result["dead_replica_ids"] == [0] + assert result["recovery_semantics"] == "graceful_capacity_degradation_without_automatic_restart" + + +def test_validate_pool_degradation_rejects_unchanged_capacity() -> None: + before = { + "effective_max_concurrent_tasks": 2, + "pool": [{"id": 0, "status": "idle"}, {"id": 1, "status": "idle"}], + } + after = { + "effective_max_concurrent_tasks": 2, + "pool": [{"id": 0, "status": "dead"}, {"id": 1, "status": "idle"}], + } + + with pytest.raises(validator.FaultValidationFailure, match="capacity 1"): + validator.validate_pool_degradation(before, after) diff --git a/tools/validation/validate_lingbot_vla_v2_service_faults.py b/tools/validation/validate_lingbot_vla_v2_service_faults.py new file mode 100644 index 00000000..489b0853 --- /dev/null +++ b/tools/validation/validate_lingbot_vla_v2_service_faults.py @@ -0,0 +1,505 @@ +"""Validate fault handling of a real LingBot-VLA v2 structured service.""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import signal +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import psutil +import requests + +try: + from tools.validation import validate_lingbot_vla_v2_structured_service as structured_validator +except ModuleNotFoundError as error: + if error.name != "tools": + raise + import validate_lingbot_vla_v2_structured_service as structured_validator + +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) + + +class FaultValidationFailure(RuntimeError): + """Raised when the service violates an expected fault-handling behavior.""" + + +def _new_session() -> requests.Session: + session = requests.Session() + session.trust_env = False + return session + + +def _response_body(response: requests.Response) -> dict[str, Any]: + try: + body = response.json() + except ValueError as error: + raise FaultValidationFailure(f"HTTP {response.status_code} response is not JSON") from error + if not isinstance(body, dict): + raise FaultValidationFailure(f"HTTP {response.status_code} response is not a JSON object") + return body + + +def _request_json( + session: requests.Session, + method: str, + url: str, + *, + timeout: float, + payload: dict[str, Any] | None = None, +) -> tuple[int, dict[str, Any]]: + response = session.request(method, url, json=payload, timeout=timeout) + return response.status_code, _response_body(response) + + +def _wait_terminal( + session: requests.Session, + base_url: str, + task_id: str, + *, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> dict[str, Any]: + deadline = time.monotonic() + task_timeout_seconds + while time.monotonic() < deadline: + status_code, body = _request_json( + session, + "GET", + f"{base_url}/v1/tasks/{task_id}/status", + timeout=http_timeout_seconds, + ) + if status_code != 200: + raise FaultValidationFailure(f"task status returned HTTP {status_code}: {body}") + status = body.get("status") or body.get("task_status") + if status in _TERMINAL_STATUSES: + return body + time.sleep(poll_interval_seconds) + raise FaultValidationFailure(f"task {task_id} did not become terminal within {task_timeout_seconds:g}s") + + +def _expect_rejected_or_failed( + session: requests.Session, + base_url: str, + payload: dict[str, Any], + *, + case_name: str, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> dict[str, Any]: + status_code, created = _request_json( + session, + "POST", + f"{base_url}/v1/tasks/structured", + timeout=http_timeout_seconds, + payload=payload, + ) + if 400 <= status_code < 500: + return {"name": case_name, "passed": True, "handling": "rejected", "http_status": status_code} + if status_code != 200: + raise FaultValidationFailure(f"{case_name} returned unexpected HTTP {status_code}: {created}") + task_id = created.get("task_id") + if not isinstance(task_id, str) or not task_id: + raise FaultValidationFailure(f"{case_name} accepted without a task_id") + terminal = _wait_terminal( + session, + base_url, + task_id, + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + if terminal.get("status") != "failed": + raise FaultValidationFailure(f"{case_name} reached unexpected terminal status {terminal.get('status')}") + return { + "name": case_name, + "passed": True, + "handling": "asynchronous_failure", + "http_status": status_code, + "task_id": task_id, + "terminal_status": "failed", + "error": str(terminal.get("error") or "")[:1000], + } + + +def validate_request_faults( + session: requests.Session, + base_url: str, + payload: dict[str, Any], + *, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> list[dict[str, Any]]: + """Validate required fields, payload validation, and request cancellation.""" + cases: list[dict[str, Any]] = [] + + missing_camera = dict(payload) + missing_camera.pop("camera_high", None) + cases.append( + _expect_rejected_or_failed( + session, + base_url, + missing_camera, + case_name="missing_required_camera", + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + ) + + invalid_state = dict(payload) + invalid_state["state"] = [0.0] * 13 + cases.append( + _expect_rejected_or_failed( + session, + base_url, + invalid_state, + case_name="invalid_state_dimension", + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + ) + + invalid_image = dict(payload) + invalid_image["camera_high"] = "not-valid-base64" + cases.append( + _expect_rejected_or_failed( + session, + base_url, + invalid_image, + case_name="invalid_camera_base64", + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + ) + + status_code, created = _request_json( + session, + "POST", + f"{base_url}/v1/tasks/structured", + timeout=http_timeout_seconds, + payload=payload, + ) + if status_code != 200 or not isinstance(created.get("task_id"), str): + raise FaultValidationFailure(f"cancellation case was not accepted: HTTP {status_code}: {created}") + task_id = created["task_id"] + cancel_status, cancellation = _request_json( + session, + "DELETE", + f"{base_url}/v1/tasks/{task_id}", + timeout=http_timeout_seconds, + ) + if cancel_status != 200 or cancellation.get("stop_status") not in {"success", "do_nothing"}: + raise FaultValidationFailure(f"task cancellation failed: HTTP {cancel_status}: {cancellation}") + terminal = _wait_terminal( + session, + base_url, + task_id, + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + terminal_status = terminal.get("status") + if terminal_status not in {"cancelled", "completed"}: + raise FaultValidationFailure(f"cancelled task reached unexpected terminal status {terminal_status}") + cases.append( + { + "name": "client_cancellation", + "passed": True, + "task_id": task_id, + "stop_status": cancellation.get("stop_status"), + "terminal_status": terminal_status, + "race_with_completion": terminal_status == "completed", + } + ) + return cases + + +def _gpu_uuid_to_index() -> dict[str, str]: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=index,uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return { + fields[1]: fields[0] + for line in result.stdout.splitlines() + if len(fields := [field.strip() for field in line.split(",")]) == 2 + } + + +def select_replica_process( + compute_rows: str, + *, + service_process_ids: set[int], + gpu_uuid_to_index: dict[str, str], + gpu_index: str, + service_pid: int, +) -> int: + """Select exactly one descendant compute process on a physical GPU.""" + candidates: set[int] = set() + for line in compute_rows.splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) < 2: + continue + try: + pid = int(fields[0]) + except ValueError: + continue + observed_index = gpu_uuid_to_index.get(fields[1], fields[1]) + if pid != service_pid and pid in service_process_ids and observed_index == gpu_index: + candidates.add(pid) + if len(candidates) != 1: + raise FaultValidationFailure( + f"expected exactly one service descendant compute process on GPU {gpu_index}, observed {sorted(candidates)}" + ) + return candidates.pop() + + +def discover_replica_process(service_pid: int, gpu_index: str) -> int: + """Discover one replica process without considering unrelated system processes.""" + root = psutil.Process(service_pid) + service_process_ids = {process.pid for process in root.children(recursive=True)} + result = subprocess.run( + ["nvidia-smi", "--query-compute-apps=pid,gpu_uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return select_replica_process( + result.stdout, + service_process_ids=service_process_ids, + gpu_uuid_to_index=_gpu_uuid_to_index(), + gpu_index=gpu_index, + service_pid=service_pid, + ) + + +def gpu_compute_process_ids(compute_rows: str) -> set[int]: + """Parse compute PIDs from nvidia-smi rows.""" + process_ids: set[int] = set() + for line in compute_rows.splitlines(): + try: + process_ids.add(int(line.split(",", 1)[0].strip())) + except (ValueError, IndexError): + continue + return process_ids + + +def wait_for_replica_exit(replica_pid: int, *, timeout_seconds: float = 10.0) -> None: + """Wait until a replica is exited or zombie and no longer owns GPU memory.""" + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + try: + process_exited = psutil.Process(replica_pid).status() == psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + process_exited = True + result = subprocess.run( + ["nvidia-smi", "--query-compute-apps=pid,gpu_uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + if process_exited and replica_pid not in gpu_compute_process_ids(result.stdout): + return + time.sleep(0.1) + raise FaultValidationFailure(f"replica process {replica_pid} did not exit and release GPU memory after SIGTERM") + + +def validate_pool_degradation(before: dict[str, Any], after: dict[str, Any]) -> dict[str, Any]: + """Validate one-replica capacity reduction without requiring automatic restart.""" + before_pool = before.get("pool") + after_pool = after.get("pool") + if not isinstance(before_pool, list) or len(before_pool) < 2: + raise FaultValidationFailure("replica termination requires service status with at least two pool entries") + if not isinstance(after_pool, list) or len(after_pool) != len(before_pool): + raise FaultValidationFailure("pool status disappeared or changed size after replica termination") + before_capacity = before.get("effective_max_concurrent_tasks") + after_capacity = after.get("effective_max_concurrent_tasks") + if not isinstance(before_capacity, int) or not isinstance(after_capacity, int): + raise FaultValidationFailure("service status has no integer effective capacity") + dead = [replica for replica in after_pool if replica.get("status") == "dead"] + live = [replica for replica in after_pool if replica.get("status") != "dead"] + if len(dead) != 1 or not live or after_capacity != before_capacity - 1: + raise FaultValidationFailure( + f"expected one dead replica and capacity {before_capacity - 1}, " + f"observed dead={len(dead)}, capacity={after_capacity}" + ) + return { + "before_capacity": before_capacity, + "after_capacity": after_capacity, + "dead_replica_ids": [replica.get("id") for replica in dead], + "live_replica_ids": [replica.get("id") for replica in live], + "recovery_semantics": "graceful_capacity_degradation_without_automatic_restart", + } + + +def validate_replica_exit( + session: requests.Session, + base_url: str, + payload: dict[str, Any], + *, + service_pid: int, + gpu_index: str, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> dict[str, Any]: + """Terminate one explicitly selected replica and validate graceful degradation.""" + status_code, before = _request_json(session, "GET", f"{base_url}/v1/service/status", timeout=http_timeout_seconds) + if status_code != 200 or before.get("execution_mode") != "concurrent_pipeline_pool": + raise FaultValidationFailure("replica termination requires a ready concurrent pipeline pool") + replica_pid = discover_replica_process(service_pid, gpu_index) + os.kill(replica_pid, signal.SIGTERM) + wait_for_replica_exit(replica_pid) + + failed_attempts: list[str] = [] + successful_request: dict[str, Any] | None = None + after: dict[str, Any] | None = None + config = structured_validator.RequestConfig( + base_url=base_url, + payload=payload, + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + expected_horizon=50, + expected_action_dim=55, + ) + for attempt in range(3): + record = structured_validator.execute_request( + session, + config, + request_index=attempt, + worker_index=0, + run_started_at=time.perf_counter(), + ) + if record.get("outcome") == "succeeded": + successful_request = record + else: + failed_attempts.append(str(record.get("error") or "unknown request failure")) + _, observed = _request_json(session, "GET", f"{base_url}/v1/service/status", timeout=http_timeout_seconds) + if any(replica.get("status") == "dead" for replica in observed.get("pool", [])): + after = observed + if successful_request is not None and after is not None: + break + if successful_request is None or after is None: + raise FaultValidationFailure( + f"service did not degrade cleanly after replica exit; request_errors={failed_attempts}" + ) + degradation = validate_pool_degradation(before, after) + return { + "name": "replica_exit", + "passed": True, + "terminated_pid": replica_pid, + "physical_gpu_index": gpu_index, + "failed_attempts": failed_attempts, + "successful_action": successful_request["action"], + **degradation, + } + + +def _image_base64(path: Path) -> str: + return base64.b64encode(path.read_bytes()).decode("ascii") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080") + parser.add_argument("--image", type=Path, required=True, help="Image reused for all three camera inputs.") + parser.add_argument("--instruction", default="pick up the object") + parser.add_argument("--state", type=structured_validator.parse_state_json, default=[0.0] * 14) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--http-timeout-seconds", type=float, default=10.0) + parser.add_argument("--task-timeout-seconds", type=float, default=120.0) + parser.add_argument("--poll-interval-seconds", type=float, default=0.05) + parser.add_argument("--service-pid", type=int) + parser.add_argument("--kill-replica-gpu-index") + parser.add_argument("--output", type=Path, default=Path("work_dirs/vla_service_fault_validation.json")) + args = parser.parse_args() + if (args.service_pid is None) != (args.kill_replica_gpu_index is None): + parser.error("--service-pid and --kill-replica-gpu-index must be provided together") + if args.http_timeout_seconds <= 0 or args.task_timeout_seconds <= 0 or args.poll_interval_seconds <= 0: + parser.error("timeout and polling values must be positive") + if not args.image.is_file(): + parser.error(f"image does not exist: {args.image}") + return args + + +def main() -> int: + args = parse_args() + encoded_image = _image_base64(args.image) + payload = { + "task": "vla_action", + "instruction": args.instruction, + "state": args.state, + "camera_high": encoded_image, + "camera_left_wrist": encoded_image, + "camera_right_wrist": encoded_image, + "seed": args.seed, + } + report: dict[str, Any] = { + "schema_version": 1, + "validation": "lingbot_vla_v2_structured_service_faults", + "created_at": datetime.now(timezone.utc).isoformat(), + "target": args.base_url.rstrip("/"), + "checks": [], + "passed": False, + } + try: + with _new_session() as session: + structured_validator.inspect_service(report["target"], timeout_seconds=args.http_timeout_seconds) + report["checks"].extend( + validate_request_faults( + session, + report["target"], + payload, + http_timeout_seconds=args.http_timeout_seconds, + task_timeout_seconds=args.task_timeout_seconds, + poll_interval_seconds=args.poll_interval_seconds, + ) + ) + if args.service_pid is not None: + report["checks"].append( + validate_replica_exit( + session, + report["target"], + payload, + service_pid=args.service_pid, + gpu_index=args.kill_replica_gpu_index, + http_timeout_seconds=args.http_timeout_seconds, + task_timeout_seconds=args.task_timeout_seconds, + poll_interval_seconds=args.poll_interval_seconds, + ) + ) + report["passed"] = all(check.get("passed") is True for check in report["checks"]) + except ( + FaultValidationFailure, + structured_validator.ValidationFailure, + requests.RequestException, + OSError, + ) as error: + report["error"] = str(error) + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps({"passed": report["passed"], "checks": len(report["checks"]), "output": str(args.output)})) + return 0 if report["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main())