base strategy simplification

megatron/core/dist_checkpointing/exchange_utils.py

@@ -11,7 +11,6 @@
 import numpy as np
 import torch
 
-from ..utils import get_pg_rank, get_pg_size, log_single_rank
 from .core import CheckpointingException
 from .dict_utils import nested_values
 from .mapping import ShardedStateDict, ShardedTensor, is_main_replica
@@ -197,6 +196,8 @@ def determine_main_replica_uniform_distribution(
         parallelization. Returns None if the process_group is trivial (1 rank)
 
     """
+    from ..utils import get_pg_size
+
     if parallelization_group is None:
         parallelization_group = torch.distributed.group.WORLD
     group_size = get_pg_size(group=parallelization_group)
@@ -285,6 +286,8 @@ def exchange_loaded_tensors_gather_rounds(
             needed by this rank to load a given state dict. Includes
             previously loaded tensors (from `loaded_tensors` input)
     """
+    from ..utils import get_pg_rank, get_pg_size
+
     if parallelization_group is None:
         parallelization_group = torch.distributed.group.WORLD
     main_rank_for_shard, _, shard_to_metadata, all_ranks_for_shard = shard_distribution
@@ -398,6 +401,8 @@ def exchange_loaded_tensors_gather_object(
             previously loaded tensors (from `loaded_tensors` input)
 
     """
+    from ..utils import log_single_rank
+
     all_loaded_tensors_list = [None] * torch.distributed.get_world_size(group=parallelization_group)
     torch.distributed.all_gather_object(
         all_loaded_tensors_list, loaded_tensors, group=parallelization_group
@@ -431,6 +436,8 @@ def exchange_loaded_objects_gather_object(
         Dict[_ShardId, Any]: dictionary mapping shard ids to objects needed by this rank to
          load a given state dict.
     """
+    from ..utils import log_single_rank
+
     all_loaded_objects_list = [None] * torch.distributed.get_world_size()
     torch.distributed.all_gather_object(all_loaded_objects_list, loaded_objects, group=None)
     all_loaded_objects_list = cast(List[Dict[_ShardId, Any]], all_loaded_objects_list)

megatron/core/dist_checkpointing/serialization.py

@@ -21,7 +21,6 @@
 from .core import CheckpointingConfig, save_config
 from .dict_utils import merge
 from .mapping import (
-    CheckpointingException,
     CommonStateDict,
     ShardedObject,
     ShardedStateDict,
@@ -30,7 +29,6 @@
 )
 from .state_dict_utils import load_preprocess, save_preprocess
 from .strategies.async_utils import AsyncRequest
-from .strategies.base import AsyncSaveShardedStrategy
 from .strategies.common import load_common, save_common
 from .strategies.torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy
 from .utils import extract_sharded_base, force_all_tensors_to_non_fp8
@@ -393,10 +391,6 @@ def metadata_finalize_fn():
         metadata_finalize_fn()
         return None
 
-    if not isinstance(sharded_strategy, AsyncSaveShardedStrategy):
-        raise CheckpointingException(
-            f'Cannot apply async_save to non-async strategy {sharded_strategy}'
-        )
     async_request = sharded_strategy.async_save(sharded_state_dict, checkpoint_dir, async_strategy)
     async_request.finalize_fns.append(metadata_finalize_fn)
     return async_request

megatron/core/dist_checkpointing/strategies/base.py

@@ -2,72 +2,57 @@
 
 """ Strategies base interfaces. """
 
+import logging
 from abc import ABC, abstractmethod
-from collections import defaultdict
 from enum import Enum
 from pathlib import Path
-from typing import Any, DefaultDict, Union
+from typing import Union
 
-from ..mapping import CheckpointingException, ShardedStateDict
-from .async_utils import AsyncCallsQueue, AsyncRequest
+from ..mapping import ShardedStateDict
+from .async_utils import AsyncRequest
+from .torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy
+
+logger = logging.getLogger(__name__)
 
 
 class StrategyAction(Enum):
-    """Specifies save vs load action."""
+    """Specifies save vs load and sharded vs common action.
+    To be removed in future releases."""
 
     LOAD_COMMON = 'load_common'
     LOAD_SHARDED = 'load_sharded'
     SAVE_COMMON = 'save_common'
     SAVE_SHARDED = 'save_sharded'
 
 
-default_strategies: DefaultDict[str, dict[tuple, Any]] = defaultdict(dict)
-
-async_calls = AsyncCallsQueue()
-
-
 def get_default_strategy(action: StrategyAction, backend: str, version: int):
     """Retrieves a default strategy for a given action, backend and version."""
-    error_hint: str = ""
-    try:
-        error_hint = ' Please use PyTorch version >=2.1'
-        from .torch import register_default_torch_strategies
-
-        register_default_torch_strategies()
-    except ImportError as e:
-        raise CheckpointingException(
-            f'Cannot import a default strategy for: {(action.value, backend, version)}. '
-            f'Error: {e}. Hint: {error_hint}'
-        ) from e
-    try:
-        return default_strategies[action.value][(backend, version)]
-    except KeyError as e:
-        raise CheckpointingException(
-            f'Cannot find a default strategy for: {(action.value, backend, version)}'
-        ) from e
-
-
-def register_default_strategy(
-    action: StrategyAction,
-    backend: str,
-    version: int,
-    strategy: Union['SaveStrategyBase', 'LoadStrategyBase'],
-):
-    """Adds a given strategy to the registry of default strategies.
-
-    Args:
-        action (StrategyAction): specifies save/load and sharded
-        backend (str): backend that the strategy becomes a default for
-        version (int): version that the strategy becomes a default for
-        strategy (SaveStrategyBase, LoadStrategyBase): strategy to register
-    """
-    default_strategies[action.value][(backend, version)] = strategy
+
+    logger.warning(
+        'megatron.core.dist_checkpointing.strategies.base.get_default_strategy'
+        ' is deprecated and will be removed in the future releases. Please use'
+        ' TorchDistLoadShardedStrategy() and TorchDistSaveShardedStrategy()'
+        ' to get the default load and save sharded strategies.'
+    )
+    if backend != 'torch_dist':
+        logger.warning(f'{backend} is not supported. `torch_dist` backend will be used.')
+    if action == StrategyAction.LOAD_SHARDED:
+        return TorchDistLoadShardedStrategy()
+    else:
+        assert action == StrategyAction.SAVE_SHARDED, f'{action} is not supported'
+        return TorchDistSaveShardedStrategy()
 
 
 class LoadStrategyBase(ABC):
     """Base class for a load strategy. Requires implementing checks for compatibility with a
     given checkpoint version."""
 
+    def __init__(self):
+        logger.warning(
+            "LoadStrategyBase & LoadShardedStrategy are deprecated "
+            "and will be removed in future releases."
+        )
+
     @abstractmethod
     def check_backend_compatibility(self, loaded_backend):
         """Verifies if this strategy is compatible with `loaded_backend`."""
@@ -89,6 +74,10 @@ class SaveStrategyBase(ABC):
     version of the saved format."""
 
     def __init__(self, backend: str, version: int):
+        logger.warning(
+            "SaveStrategyBase & SaveShardedStrategy are deprecated "
+            "and will be removed in future releases."
+        )
         self.backend = backend
         self.version = version
 
@@ -102,7 +91,7 @@ def __str__(self):
 
 
 class LoadShardedStrategy(LoadStrategyBase):
-    """Load strategy for sharded tensors"""
+    """Base class for load strategies to be removed in future releases."""
 
     @abstractmethod
     def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
@@ -145,7 +134,7 @@ def remove_sharded_tensors(self, checkpoint_dir: Union[str, Path], key_prefix: s
 
 
 class SaveShardedStrategy(SaveStrategyBase):
-    """Save strategy for sharded tensors"""
+    """Base class for save strategies to be removed in future releases."""
 
     @abstractmethod
     def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
@@ -154,7 +143,7 @@ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str,
 
 
 class AsyncSaveShardedStrategy(SaveShardedStrategy):
-    """Save strategy suitable for async save."""
+    """Save strategy suitable for async save. To be removed in future releases."""
 
     @abstractmethod
     def async_save(
@@ -174,6 +163,9 @@ def async_save(
 
     def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
         """Each async strategy can be trivially used as a sync strategy."""
+        logger.warning(
+            "AsyncSaveShardedStrategy is deprecated and will be removed in future releases."
+        )
         async_request = self.async_save(sharded_state_dict, checkpoint_dir)
         async_request.execute_sync()
         del async_request

megatron/core/dist_checkpointing/strategies/fully_parallel.py

@@ -23,10 +23,9 @@
     exchange_loaded_objects_gather_object,
 )
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict, StateDict, is_main_replica
-from megatron.core.dist_checkpointing.strategies.base import (
-    AsyncSaveShardedStrategy,
-    LoadShardedStrategy,
-    SaveShardedStrategy,
+from megatron.core.dist_checkpointing.strategies.torch import (
+    TorchDistLoadShardedStrategy,
+    TorchDistSaveShardedStrategy,
 )
 from megatron.core.dist_checkpointing.utils import (
     _sharded_object_id,
@@ -38,14 +37,13 @@
     determine_global_metadata,
     validate_sharding_integrity,
 )
-from megatron.core.utils import get_pg_rank, get_pg_size
 
 logger = logging.getLogger(__name__)
 
 T = TypeVar('T', ShardedObject, ShardedTensor)
 
 
-class FullyParallelSaveStrategyWrapper(AsyncSaveShardedStrategy):
+class FullyParallelSaveStrategyWrapper:
     """Wraps arbitrary strategy and distributes the save during `save`.
 
     The save distribution happens without any *data* communication.
@@ -60,7 +58,7 @@ class FullyParallelSaveStrategyWrapper(AsyncSaveShardedStrategy):
     described in `distribute_shards_to_ranks`.
 
     Args:
-        strategy (SaveShardedStrategy): base strategy to wrap
+        strategy (TorchDistSaveShardedStrategy): base strategy to wrap
         parallelization_group (ProcessGroup, optional): process group to use for save
             distribution. Note that this doesn't have to match exactly the
             data distribution, but should cover the replication pattern
@@ -72,16 +70,20 @@ class FullyParallelSaveStrategyWrapper(AsyncSaveShardedStrategy):
 
     def __init__(
         self,
-        strategy: SaveShardedStrategy,
+        strategy: TorchDistSaveShardedStrategy,
         parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
         do_cache_distribution: bool = False,
+        backend: str = "torch_dist",
+        version: int = 1,
     ):
-        super().__init__(strategy.backend, strategy.version)
+        """ """
         self.base_strategy = strategy
         if parallelization_group is None:
             parallelization_group = torch.distributed.group.WORLD
         self.parallelization_group = parallelization_group
         self.do_cache_distribution = do_cache_distribution
+        self.backend = backend
+        self.version = version
 
         self.cached_distribution: Optional[ShardDistribution] = None
 
@@ -92,10 +94,6 @@ def async_save(
         async_strategy: str = "nvrx",
     ):
         """ """
-        if not isinstance(self.base_strategy, AsyncSaveShardedStrategy):
-            raise CheckpointingException(
-                f'Cannot apply async_save to non-async base strategy {self.base_strategy}'
-            )
         self.apply_saving_parallelization(sharded_state_dict)
         return self.base_strategy.async_save(sharded_state_dict, checkpoint_dir, async_strategy)
 
@@ -140,19 +138,14 @@ def apply_saving_parallelization(self, sharded_state_dict: ShardedStateDict) ->
         end = time()
         logger.debug(f"parallel save sharding, time: {end - start}")
 
-    @property
-    def can_handle_sharded_objects(self):
-        """ """
-        return self.base_strategy.can_handle_sharded_objects
-
 
-class FullyParallelLoadStrategyWrapper(LoadShardedStrategy):
+class FullyParallelLoadStrategyWrapper:
     """Wraps arbitrary load strategy and distributes the load during `load`.
 
     See `load` method docs for details.
 
     Args:
-        strategy (LoadShardedStrategy): base strategy to wrap
+        strategy (TorchDistLoadShardedStrategy): base strategy to wrap
         parallelization_group (ProcessGroup, optional): process group to use for load
             distribution. Note that this doesn't have to match exactly the
             data distribution, but should cover the replication pattern
@@ -174,12 +167,11 @@ class FullyParallelLoadStrategyWrapper(LoadShardedStrategy):
 
     def __init__(
         self,
-        strategy: LoadShardedStrategy,
+        strategy: TorchDistLoadShardedStrategy,
         parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
         do_cache_distribution: bool = False,
         exchange_algo: str = 'broadcast',
     ):
-        super().__init__()
         self.base_strategy = strategy
         if parallelization_group is None:
             parallelization_group = (
@@ -227,6 +219,7 @@ def load(
             a state dict that would be loaded with the underlying strategy
             without this wrapper.
         """
+        from megatron.core.utils import get_pg_size
 
         loaded_state_dict = {}
 
@@ -403,11 +396,6 @@ def apply_loading_parallelization(
 
         return precomputed_distribution
 
-    @property
-    def can_handle_sharded_objects(self):
-        """ """
-        return self.base_strategy.can_handle_sharded_objects
-
     def load_tensors_metadata(self, checkpoint_dir: Path):
         """ """
         return self.base_strategy.load_tensors_metadata(checkpoint_dir)
@@ -416,14 +404,6 @@ def load_sharded_metadata(self, checkpoint_dir: Path):
         """ """
         return self.base_strategy.load_sharded_metadata(checkpoint_dir)
 
-    def check_backend_compatibility(self, loaded_version):
-        """ """
-        return self.base_strategy.check_backend_compatibility(loaded_version)
-
-    def check_version_compatibility(self, loaded_version):
-        """ """
-        return self.base_strategy.check_version_compatibility(loaded_version)
-
 
 def distribute_main_replicas_with_precomputed_distribution(
     sharded_state_dict: ShardedStateDict,
@@ -455,6 +435,8 @@ def distribute_main_replicas_with_precomputed_distribution(
     rank1: A: 1, B: 0, C: 1
     rank2: A: 1, B: 1, C: 0
     """
+    from megatron.core.utils import get_pg_rank, get_pg_size
+
     if parallelization_group is None:
         parallelization_group = torch.distributed.group.WORLD
     if get_pg_size(group=parallelization_group) <= 1: