From 63cd2111ade80173e40094ccff3d868937831359 Mon Sep 17 00:00:00 2001 From: Hakjin Lee Date: Tue, 21 Oct 2025 21:26:31 +0900 Subject: [PATCH] [Fix] Device Error on SmolVLA Multi-GPU Training (#2270) Co-authored-by: Steven Palma --- src/lerobot/policies/smolvla/modeling_smolvla.py | 1 + src/lerobot/policies/smolvla/smolvlm_with_expert.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/lerobot/policies/smolvla/modeling_smolvla.py b/src/lerobot/policies/smolvla/modeling_smolvla.py index 23fc3ca4f..6e54d3ea5 100644 --- a/src/lerobot/policies/smolvla/modeling_smolvla.py +++ b/src/lerobot/policies/smolvla/modeling_smolvla.py @@ -485,6 +485,7 @@ class VLAFlowMatching(nn.Module): num_vlm_layers=self.config.num_vlm_layers, self_attn_every_n_layers=self.config.self_attn_every_n_layers, expert_width_multiplier=self.config.expert_width_multiplier, + device=self.config.device, ) self.state_proj = nn.Linear( self.config.max_state_dim, self.vlm_with_expert.config.text_config.hidden_size diff --git a/src/lerobot/policies/smolvla/smolvlm_with_expert.py b/src/lerobot/policies/smolvla/smolvlm_with_expert.py index f3d1a693a..555c40773 100644 --- a/src/lerobot/policies/smolvla/smolvlm_with_expert.py +++ b/src/lerobot/policies/smolvla/smolvlm_with_expert.py @@ -70,13 +70,14 @@ class SmolVLMWithExpertModel(nn.Module): num_vlm_layers: int = -1, self_attn_every_n_layers: int = -1, expert_width_multiplier: float = 0.5, + device: str = "auto", ): super().__init__() if load_vlm_weights: print(f"Loading {model_id} weights ...") self.vlm = AutoModelForImageTextToText.from_pretrained( model_id, - device_map="auto", + device_map=device, torch_dtype="bfloat16", low_cpu_mem_usage=True, )