From 3f2179f3b69708b6ad009b2e7685dd9d05269ee1 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Wed, 15 Jul 2026 17:02:01 +0200 Subject: [PATCH] refactor(evo1): use transformers flash attention probe (#4013) Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com> --- .../policies/evo1/internvl3_embedder.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/lerobot/policies/evo1/internvl3_embedder.py b/src/lerobot/policies/evo1/internvl3_embedder.py index d47105d96..295a4c343 100644 --- a/src/lerobot/policies/evo1/internvl3_embedder.py +++ b/src/lerobot/policies/evo1/internvl3_embedder.py @@ -27,9 +27,11 @@ from lerobot.utils.import_utils import _transformers_available, require_package if TYPE_CHECKING or _transformers_available: from transformers import AutoModel, AutoTokenizer + from transformers.utils import is_flash_attn_2_available else: AutoModel = None AutoTokenizer = None + is_flash_attn_2_available = None IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) @@ -135,9 +137,13 @@ class InternVL3Embedder(nn.Module): raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc self.model_dtype = model_dtype - attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager" + attn_implementation = ( + "flash_attention_2" if (use_flash_attn and is_flash_attn_2_available()) else "eager" + ) if use_flash_attn and attn_implementation == "eager": - logger.warning("flash_attn is not installed. Falling back to eager attention.") + logger.warning( + "Flash Attention 2 is unavailable on this runtime. Falling back to eager attention." + ) self.model = AutoModel.from_pretrained( model_name, @@ -359,11 +365,3 @@ class InternVL3Embedder(nn.Module): @property def device(self) -> torch.device: return next(self.model.parameters()).device - - -def _flash_attn_available() -> bool: - try: - import flash_attn # noqa: F401 - except ModuleNotFoundError: - return False - return True