From 8600f88d9fc793e167a3bb53368e2b40fba6a795 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Tue, 21 Jul 2026 15:24:53 +0000 Subject: [PATCH] device autocast --- .../policies/vla_jepa/modeling_vla_jepa.py | 9 ++++---- src/lerobot/utils/device_utils.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index 962f2b6fb..991fcbc7a 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -26,6 +26,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 get_autocast_context from lerobot.utils.import_utils import _transformers_available, require_package if TYPE_CHECKING or _transformers_available: @@ -183,7 +184,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) @@ -246,8 +247,8 @@ class VLAJEPAModel(nn.Module): ) if reduction == "none": # Per-sample loss (B,): mean over all non-batch dims (tokens, feature). - l = F.l1_loss(predicted_states, gt_states.float(), reduction="none") - return l.mean(dim=tuple(range(1, l.ndim))) + elementwise = F.l1_loss(predicted_states, gt_states.float(), reduction="none") + return elementwise.mean(dim=tuple(range(1, elementwise.ndim))) return F.l1_loss(predicted_states, gt_states.float(), reduction="mean") def _action_loss( @@ -264,7 +265,7 @@ class VLAJEPAModel(nn.Module): independent noise draws are averaged back per original sample — for RA-BC weighting. """ device_type = next(self.parameters()).device.type - with torch.autocast(device_type=device_type, dtype=torch.float32): + with get_autocast_context(device_type, torch.float32): r = self.config.repeated_diffusion_steps horizon = self.config.chunk_size b = embodied_action_tokens.shape[0] diff --git a/src/lerobot/utils/device_utils.py b/src/lerobot/utils/device_utils.py index 3f9b58773..13623ce1b 100644 --- a/src/lerobot/utils/device_utils.py +++ b/src/lerobot/utils/device_utils.py @@ -15,6 +15,7 @@ # limitations under the License. import logging +from contextlib import nullcontext import torch @@ -121,3 +122,25 @@ def is_amp_available(device: str): return False else: raise ValueError(f"Unknown device '{device}.") + + +def get_autocast_context(device_type: str, dtype: torch.dtype = torch.bfloat16): + """Return a device-safe autocast context manager. + + Hardcoding `torch.autocast(dtype=torch.bfloat16)` breaks on backends without AMP + (MPS) and silently misbehaves on pre-Ampere CUDA GPUs that lack bf16 support. This + picks a safe context per device: + - no AMP support (e.g. mps): `nullcontext()` (run in the tensors' native dtype) + - CUDA requesting bf16 on compute capability < 8.0 (pre-Ampere): fall back to fp16 + - otherwise: `torch.autocast(device_type, dtype)` + """ + if not is_amp_available(device_type): + return nullcontext() + if ( + device_type == "cuda" + and dtype == torch.bfloat16 + and torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] < 8 + ): + dtype = torch.float16 + return torch.autocast(device_type=device_type, dtype=dtype) \ No newline at end of file