From 412aa1788f791b7c9bab78d516dfd03350d78dc0 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Tue, 28 Oct 2025 09:28:01 -0700 Subject: [PATCH 1/3] zarr soft deprecation Signed-off-by: dimapihtar --- megatron/core/dist_checkpointing/strategies/zarr.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/megatron/core/dist_checkpointing/strategies/zarr.py b/megatron/core/dist_checkpointing/strategies/zarr.py index 9b01403bf2c..59a7f46c13c 100644 --- a/megatron/core/dist_checkpointing/strategies/zarr.py +++ b/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): From 620635cef937d10267b04e0948b42ae89523d4b0 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Tue, 28 Oct 2025 10:06:29 -0700 Subject: [PATCH 2/3] fix style Signed-off-by: dimapihtar --- megatron/core/dist_checkpointing/strategies/zarr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/dist_checkpointing/strategies/zarr.py b/megatron/core/dist_checkpointing/strategies/zarr.py index 59a7f46c13c..ac6e04e647f 100644 --- a/megatron/core/dist_checkpointing/strategies/zarr.py +++ b/megatron/core/dist_checkpointing/strategies/zarr.py @@ -196,7 +196,7 @@ 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( From 5ad362ff34a4ad3ee086d34487183ee53ff2e5dc Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Tue, 28 Oct 2025 10:18:16 -0700 Subject: [PATCH 3/3] fix style Signed-off-by: dimapihtar --- megatron/core/dist_checkpointing/strategies/zarr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/dist_checkpointing/strategies/zarr.py b/megatron/core/dist_checkpointing/strategies/zarr.py index ac6e04e647f..cea21ddea16 100644 --- a/megatron/core/dist_checkpointing/strategies/zarr.py +++ b/megatron/core/dist_checkpointing/strategies/zarr.py @@ -317,7 +317,7 @@ def pad_to_expected_shape(x: torch.Tensor, expected_sharded_ten: ShardedTensor): return torch.nn.functional.pad(x, pad_args) # unsqueeze and squeeze to get shapes supported by cudnn - print(f"Replicating last row for {expected_sharded_ten.key}") + logger.info(f"Replicating last row for {expected_sharded_ten.key}") if x.dtype == torch.bfloat16: return ( torch.nn.functional.pad(x.float().unsqueeze(0), pad_args, mode="replicate")