fix(safetensors): expand bare "cuda" to current device for safetensors loads (#4042)

src/lerobot/policies/pretrained.py

@@ -32,6 +32,7 @@ from torch import Tensor, nn
 from lerobot.__version__ import __version__
 from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
+from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
 
 from .utils import log_model_loading_keys
@@ -220,7 +221,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
         missing_keys, unexpected_keys = load_model_as_safetensor(
-            model, model_file, strict=strict, device=map_location
+            model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
         )
         log_model_loading_keys(missing_keys, unexpected_keys)
         return model

src/lerobot/rewards/pretrained.py

@@ -28,6 +28,7 @@ from safetensors.torch import load_model as load_model_as_safetensor, save_model
 from torch import Tensor, nn
 
 from lerobot.configs.rewards import RewardModelConfig
+from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
 
 if TYPE_CHECKING:
@@ -128,7 +129,7 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
         missing_keys, unexpected_keys = load_model_as_safetensor(
-            model, model_file, strict=strict, device=map_location
+            model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
         )
         if missing_keys:
             logging.warning(f"Missing key(s) when loading model: {missing_keys}")

src/lerobot/utils/device_utils.py

@@ -59,6 +59,20 @@ def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
     return device
 
 
+def resolve_safetensors_device(map_location: str | torch.device) -> str:
+    """Resolve a device string for a safetensors load, working around a device-mapping quirk.
+
+    safetensors' load 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.
+    """
+    map_location = str(map_location)
+    if map_location == "cuda" and torch.cuda.is_available():
+        return f"cuda:{torch.cuda.current_device()}"
+    return map_location
+
+
 def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
     """
     mps is currently not compatible with float64