diff --git a/nemo_aligner/models/nlp/gpt/gpt_sft_model.py b/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
index d3a615500..15fa86d83 100644
--- a/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
+++ b/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
@@ -59,6 +59,13 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
                 inference_params["strategy"] = hydra.utils.instantiate(inference_params["strategy"], model=self)
         self.set_inference_params(**inference_params)
 
+    def configure_sharded_model(self):
+        """Configure the model for sharded training.
+        This method is required by PyTorch Lightning to properly handle model sharding.
+        """
+        # Call the parent class's implementation
+        return super().configure_sharded_model()
+
     def set_inference_params(self, length_params=None, sampling_params=None, strategy=None):
         # TODO (igitman): the name self._inference_params is very similar to self.inference_params
         #    that's used by the base model for another purpose. There is also self._inference_config
diff --git a/nemo_aligner/utils/utils.py b/nemo_aligner/utils/utils.py
index 8eb73d295..757ed2c3e 100644
--- a/nemo_aligner/utils/utils.py
+++ b/nemo_aligner/utils/utils.py
@@ -88,7 +88,8 @@ def load_from_nemo(
     """load a model using nemo checkpoint
     """
     connector = CustomSaveRestoreConnector(load_base_model_only=load_base_model_only)
-    assert os.path.exists(restore_path), f"tried to load from {restore_path=} but it does not exist"
+    if "msc" not in restore_path:
+        assert os.path.exists(restore_path), f"tried to load from {restore_path=} but it does not exist"
 
     # if we gave it a directory, then load as if it was extracted already
     if os.path.isdir(restore_path):
@@ -114,7 +115,8 @@ def load_checkpoint_model_config(restore_path):
     """load only the model config from a checkpoint
     """
     config_name_in_ckpt = NLPSaveRestoreConnector()._model_config_yaml
-    assert os.path.exists(restore_path), f"tried to load from {restore_path=} but it does not exist"
+    if "msc" not in restore_path:
+        assert os.path.exists(restore_path), f"tried to load from {restore_path=} but it does not exist"
 
     # if we gave it a directory, then load the cfg directly
     if os.path.isdir(restore_path):