fix: address review comments - PEFT cloud path, optional dependencies, and unit test fixes

Signed-off-by: Edison <edisonggacc@gmail.com>

git add .
git rebase --continue
git push origin feat/s3-dcp-support# interactive rebase in progress; onto ec20197

Signed-off-by: Edison <edisonggacc@gmail.com>

Author: edjson

nemo_automodel/components/checkpoint/checkpointing.py

@@ -40,7 +40,7 @@
     HF_HUB_CACHE = None
 
 from packaging.version import parse
-from safetensors.torch import load_file, save_file
+from safetensors.torch import load_file, save_file, load as safetensors_load
 from torch import nn
 from torch.distributed.device_mesh import DeviceMesh
 
@@ -76,7 +76,11 @@ def _ensure_msc_available() -> None:
     """Raise an error if MSC is not installed but a cloud path is used."""
     if not MSC_AVAILABLE:
         raise ImportError(
+<<<<<<< HEAD
             "multistorageclient is required for cloud storage paths. "
+=======
+            "multistorageclient is required for cloud storage paths."
+>>>>>>> 91a223c (fix: address review comments - PEFT cloud path, optional dependencies, and unit test fixes)
             "Install it with: pip install multi-storage-client "
             "--index-url https://pypi.nvidia.com"
         )
@@ -700,7 +704,14 @@ def _do_load(
         is_model = True if "/model" in path else False
         # PEFT loading is broadcasted from rank0 so it is a special case
         if self.config.is_peft and is_model and (not is_init_step):
-            state_dict = load_file(os.path.join(path, "adapter_model.safetensors"))
+            if is_cloud_path(path):
+                _ensure_msc_available()
+                adapter_path = path.rstrip("/") + "/adapter_model.safetensors"
+                with msc.open(adapter_path, "rb") as f:
+                    data = f.read()
+                state_dict = safetensors_load(data)
+            else:
+                state_dict = load_file(os.path.join(path, "adapter_model.safetensors"))
         else:
             if is_cloud_path(path):
                 _ensure_msc_available()
@@ -730,7 +741,13 @@ def _do_save(
         # PEFT saving is done on rank0 so it is a special case
         if self.config.is_peft and is_model:
             if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
-                save_file(state_dict, os.path.join(path, "adapter_model.safetensors"))
+                if is_cloud_path(path):
+                    _ensure_msc_available()
+                    adapter_path = path.rstrip("/") + "/adapter_model.safetensors"
+                    with msc.open(adapter_path, "wb") as f:
+                        save_file(state_dict, f)
+                else:
+                    save_file(state_dict, os.path.join(path, "adapter_model.safetensors"))
             if torch.distributed.is_initialized():
                 torch.distributed.barrier()
             return

pyproject.toml

@@ -87,11 +87,11 @@ dependencies = [
     "torchao",
     "mlflow",
     "flashoptim>=0.1.3",
-    "localstack>=2026.3.0",
-    "multistorageclient[aws]",
 ]
 
 [project.optional-dependencies]
+cloud = ["multi-storage-client",]
+dev = ["localstack>=2026.3.0",]
 diffusion = [
     "diffusers>=0.36.0",
     "ftfy",

tests/unit_tests/checkpoint/test_checkpointing.py

@@ -904,7 +904,7 @@ def test_cloud_path_uses_msc_reader(self):
              patch("nemo_automodel.components.checkpoint.checkpointing.dcp"):
             Checkpointer._do_load(ckptr, state_dict, "msc://bucket/step-100")
 
-        mock_msc.torch. MultiStorageFileSystemReader.assert_called_once_with("msc://bucket/step-100")
+        mock_msc.torch.MultiStorageFileSystemReader.assert_called_once_with("msc://bucket/step-100")
 
     def test_local_path_does_not_use_msc_reader(self, tmp_path):
         ckptr = self._make_checkpointer()
@@ -914,17 +914,25 @@ def test_local_path_does_not_use_msc_reader(self, tmp_path):
              patch("nemo_automodel.components.checkpoint.checkpointing.dcp"):
             Checkpointer._do_load(ckptr, state_dict, str(tmp_path / "step-100"))
 
-        mock_msc.torch. MultiStorageFileSystemReader.assert_not_called()
+        mock_msc.open.assert_not_called()
 
     def test_peft_cloud_load_still_routes_through_msc_reader(self):
         ckptr = self._make_checkpointer(is_peft=True)
         state_dict = {"weight": torch.zeros(4)}
+        mock_file = MagicMock()
+        mock_file.read.return_value= b"fake bytes"
 
         with patch("nemo_automodel.components.checkpoint.checkpointing.msc") as mock_msc, \
-             patch("nemo_automodel.components.checkpoint.checkpointing.dcp"):
-            Checkpointer._do_load(ckptr, state_dict, "msc://bucket/step-100")
+             patch("nemo_automodel.components.checkpoint.checkpointing.dcp"), \
+             patch("nemo_automodel.components.checkpoint.checkpointing.safetensors_load") as mock_load:
+            mock_msc.open.return_value.__enter__=MagicMock(return_value=mock_file)
+            mock_msc.open.return_value.__exit__=MagicMock(return_value=False)
+            mock_load.return_value = state_dict
+            Checkpointer._do_load(ckptr, state_dict, "msc://bucket/step-100/model")
+        
+        mock_msc.open.assert_called_once()
 
-        mock_msc.torch. MultiStorageFileSystemReader.assert_called_once_with("msc://bucket/step-100")
+       
 
     def test_save_and_load_use_same_path(self):
         config = MagicMock()
@@ -941,4 +949,4 @@ def test_save_and_load_use_same_path(self):
             Checkpointer._do_load(ckptr, state_dict, path)
 
         mock_msc.torch.MultiStorageFileSystemWriter.assert_called_once_with(path)
-        mock_msc.torch. MultiStorageFileSystemReader.assert_called_once_with(path)
+        mock_msc.torch.MultiStorageFileSystemReader.assert_called_once_with(path)