zarr soft deprecation

Signed-off-by: dimapihtar <[email protected]>

Author: dimapihtar

megatron/core/dist_checkpointing/strategies/zarr.py

@@ -77,9 +77,9 @@ class ZarrSaveShardedStrategy(SaveShardedStrategy):
 
     def __init__(self, backend: str, version: int):
         super().__init__(backend, version)
-        logger.warning(
-            f"`zarr` distributed checkpoint backend is deprecated."
-            " Please switch to PyTorch Distributed format (`torch_dist`)."
+        raise CheckpointingException(
+            "`zarr` distributed checkpoint backend is no longer supported. "
+            "Please switch to PyTorch Distributed format (`torch_dist`)."
         )
 
     def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
@@ -196,6 +196,13 @@ def _create_zarr_array(sharded_tensor: ShardedTensor, checkpoint_dir: Path):
 
 class ZarrLoadShardedStrategy(LoadShardedStrategy):
     """Load strategy for the Zarr backend."""
+    
+    def __init__(self, backend: str, version: int):
+        super().__init__(backend, version)
+        raise CheckpointingException(
+            "`zarr` distributed checkpoint backend is no longer supported. "
+            "Please switch to PyTorch Distributed format (`torch_dist`)."
+        )
 
     def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
         if isinstance(checkpoint_dir, str):