From ef32c04f445caeb5ec4c0b0cb32d9234cefd8b59 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Thu, 9 Jul 2026 12:09:46 +0000 Subject: [PATCH] avoid loading the model to rank 0 to avoid a big vram spike which can OOM --- src/lerobot/policies/pretrained.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 702569b8c..812e7826e 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -29,6 +29,7 @@ from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, sa from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from huggingface_hub.errors import HfHubHTTPError from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor +import torch from torch import Tensor, nn from lerobot.__version__ import __version__ @@ -221,6 +222,14 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): @classmethod def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T: + # safetensors' load_file maps the bare string "cuda" to cuda:0 regardless of the current + # device (unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under + # multi-GPU accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it + # before sharding. Resolve "cuda" to the concrete current-device index so each rank loads + # onto its own GPU. + if map_location == "cuda" and torch.cuda.is_available(): + map_location = f"cuda:{torch.cuda.current_device()}" + # Create base kwargs kwargs = {"strict": strict}