From d648e308d2703df49218d950720823d0846e2210 Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Sat, 13 Jun 2026 20:56:04 -0400 Subject: [PATCH] fix(vla_jepa): use device-safe autocast instead of hardcoded bfloat16 VLA-JEPA hardcodes torch.autocast with dtype=torch.bfloat16, which crashes on MPS (no AMP support) and silently misbehaves on pre-Ampere CUDA GPUs (no bf16). Add a _get_autocast_context() helper that reuses the existing is_amp_available() utility to pick a safe strategy per device, matching the pattern used by pi05 and molmoact2. Fixes #3744 --- .../policies/vla_jepa/modeling_vla_jepa.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index 9c689d3c7..a9c5290eb 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -16,6 +16,7 @@ from __future__ import annotations import logging from collections import deque +from contextlib import nullcontext from pathlib import Path from typing import TYPE_CHECKING, Any @@ -26,6 +27,7 @@ from torch import Tensor, nn from lerobot.policies.pretrained import PreTrainedPolicy, T from lerobot.policies.utils import populate_queues from lerobot.utils.constants import ACTION, OBS_STATE +from lerobot.utils.device_utils import is_amp_available from lerobot.utils.import_utils import _transformers_available, require_package if TYPE_CHECKING or _transformers_available: @@ -39,6 +41,21 @@ from .configuration_vla_jepa import VLAJEPAConfig from .qwen_interface import Qwen3VLInterface from .world_model import ActionConditionedVideoPredictor + +def _get_autocast_context(device_type: str, dtype: torch.dtype = torch.bfloat16): + """Return an autocast context appropriate for the device. + + MPS does not support ``torch.autocast`` at all. On CUDA devices + without bfloat16 support (compute capability < 8.0) we fall back to + float16. + """ + if not is_amp_available(device_type): + return nullcontext() + if device_type == "cuda" and dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported(): + dtype = torch.float16 + return torch.autocast(device_type=device_type, dtype=dtype) + + # ============================================================================ # Native VLA-JEPA Model - follows original starVLA VLA_JEPA.py implementation # ============================================================================ @@ -183,7 +200,7 @@ class VLAJEPAModel(nn.Module): action_idx = action_mask.nonzero(as_tuple=True) device_type = next(self.parameters()).device.type - with torch.autocast(device_type=device_type, dtype=torch.bfloat16): + with _get_autocast_context(device_type, torch.bfloat16): last_hidden = self._qwen_last_decoder_hidden(qwen_inputs) # [B, seq_len, H] b, _, h = last_hidden.shape embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h)