diff --git a/megatron/core/dist_checkpointing/exchange_utils.py b/megatron/core/dist_checkpointing/exchange_utils.py index 79f906b237a..7c7863532f6 100644 --- a/megatron/core/dist_checkpointing/exchange_utils.py +++ b/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) diff --git a/megatron/core/dist_checkpointing/serialization.py b/megatron/core/dist_checkpointing/serialization.py index a0426ec5f80..2ee7970f143 100644 --- a/megatron/core/dist_checkpointing/serialization.py +++ b/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 diff --git a/megatron/core/dist_checkpointing/strategies/base.py b/megatron/core/dist_checkpointing/strategies/base.py index eb20e145ffb..c438382ed14 100644 --- a/megatron/core/dist_checkpointing/strategies/base.py +++ b/megatron/core/dist_checkpointing/strategies/base.py @@ -2,18 +2,22 @@ """ 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' @@ -21,53 +25,34 @@ class StrategyAction(Enum): 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 diff --git a/megatron/core/dist_checkpointing/strategies/fully_parallel.py b/megatron/core/dist_checkpointing/strategies/fully_parallel.py index a85efdaa10a..6638f215cd4 100644 --- a/megatron/core/dist_checkpointing/strategies/fully_parallel.py +++ b/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: diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py index 000aa4d9265..31b9175f70b 100644 --- a/megatron/core/dist_checkpointing/strategies/torch.py +++ b/megatron/core/dist_checkpointing/strategies/torch.py @@ -49,12 +49,6 @@ is_main_replica, ) from .async_utils import AsyncRequest -from .base import ( - AsyncSaveShardedStrategy, - LoadShardedStrategy, - StrategyAction, - register_default_strategy, -) from .checkpointable import CheckpointableShardedTensor, LocalShardsContainer try: @@ -62,10 +56,14 @@ from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( CheckpointMetadataCache, ) + + HAVE_NVRX = True except (ImportError, ModuleNotFoundError): CheckpointMetadataCache = ABC NVRxAsyncRequest = ABC + HAVE_NVRX = False + try: if not torch.cuda.is_available(): raise ImportError @@ -103,16 +101,6 @@ class MCoreSavePlan: pass -def register_default_torch_strategies(): - """Register default strategies related to PyT Distributed backend.""" - register_default_strategy( - StrategyAction.LOAD_SHARDED, 'torch_dist', 1, TorchDistLoadShardedStrategy() - ) - register_default_strategy( - StrategyAction.SAVE_SHARDED, 'torch_dist', 1, TorchDistSaveShardedStrategy() - ) - - logger = getLogger(__name__) @@ -596,7 +584,7 @@ def commit_tensor(self, read_item: ReadItem, tensor: torch.Tensor) -> None: return super().commit_tensor(read_item, tensor) -class TorchDistSaveShardedStrategy(AsyncSaveShardedStrategy): +class TorchDistSaveShardedStrategy: """Async save strategy for the PyT Distributed format. The idea is to translate MCore ShardedTensors into PyT ShardedTensors @@ -627,7 +615,8 @@ def __init__( separation_hint(str, optional): If provided, all tensors whose keys have this prefix will be saved to a separate file. """ - super().__init__(backend, version) + self.backend = backend + self.version = version self.keep_only_main_replica = keep_only_main_replica self.thread_count = thread_count @@ -654,6 +643,13 @@ def __init__( self.validated_loaded_metadata_reuse = False + def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path): + """Each async strategy can be trivially used as a sync strategy.""" + strategy = "nvrx" if HAVE_NVRX else "mcore" + async_request = self.async_save(sharded_state_dict, checkpoint_dir, async_strategy=strategy) + async_request.execute_sync() + del async_request + def async_save( self, sharded_state_dict: ShardedStateDict, @@ -805,9 +801,6 @@ def finalize_fn(): return async_request(save_fn, save_args, [finalize_fn], preload_fn=preload_fn) - def can_handle_sharded_objects(self): - return True - def _get_filesystem_reader( checkpoint_dir: Union[str, Path], cache_metadata: bool = False, async_strategy: str = "nvrx" @@ -823,13 +816,12 @@ def _get_filesystem_reader( return FileSystemReader(checkpoint_dir) -class TorchDistLoadShardedStrategy(LoadShardedStrategy): +class TorchDistLoadShardedStrategy: """Basic load strategy for the PyT Distributed format.""" def __init__(self, cache_metadata: bool = False): self.cached_global_metadata: Optional[Metadata] = None self.cache_metadata = cache_metadata - super().__init__() def load( self, @@ -1012,15 +1004,6 @@ def remove_sharded_tensors(self, checkpoint_dir: str, key_prefix: str): else: fs_writer.fs.rm_file(old_path) - def can_handle_sharded_objects(self): - return True - - def check_backend_compatibility(self, loaded_version): - pass # TODO - - def check_version_compatibility(self, loaded_version): - pass # TODO - def get_async_strategy(async_strategy: str = "nvrx", module: str = None) -> tuple: """Returns async strategy and related async imported modules""" diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_gdn/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_gdn/model_config.yaml index 25bb4077a02..cc7211c9967 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_gdn/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_gdn/model_config.yaml @@ -79,6 +79,4 @@ MODEL_ARGS: --bf16: true --attention-backend: unfused --log-memory-to-tensorboard: true - --async-save: true - --use-persistent-ckpt-worker: true TEST_TYPE: ckpt-resume diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml index 7e5ca833000..8c655bc135c 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml @@ -55,6 +55,4 @@ MODEL_ARGS: --bf16: true --attention-backend: unfused --log-memory-to-tensorboard: true - --async-save: true - --use-persistent-ckpt-worker: true TEST_TYPE: regular diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml index 4b9f7b03ace..329639ad0c9 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml @@ -56,6 +56,4 @@ MODEL_ARGS: --disable-bias-linear: true --no-bias-gelu-fusion: true --log-memory-to-tensorboard: true - --async-save: true - --use-persistent-ckpt-worker: true TEST_TYPE: ckpt-resume diff --git a/tests/unit_tests/dist_checkpointing/conftest.py b/tests/unit_tests/dist_checkpointing/conftest.py index 301e78d9733..1237c4441e0 100644 --- a/tests/unit_tests/dist_checkpointing/conftest.py +++ b/tests/unit_tests/dist_checkpointing/conftest.py @@ -4,7 +4,7 @@ import pytest -from megatron.core.dist_checkpointing.strategies.base import StrategyAction, get_default_strategy +from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy from megatron.core.msc_utils import MultiStorageClientFeature @@ -13,21 +13,12 @@ def pytest_sessionfinish(session, exitstatus): session.exitstatus = 0 -@pytest.fixture(scope="class") -def tmp_dir_per_class(tmp_path_factory): - return tmp_path_factory.mktemp("data") - - @pytest.fixture(scope='session', autouse=True) -def set_default_dist_ckpt_strategy(): - # Disable MSC for tests +def disable_msc(): MultiStorageClientFeature.disable() + yield - def get_pyt_dist_save_sharded_strategy(): - return get_default_strategy(StrategyAction.SAVE_SHARDED, 'torch_dist', 1) - with mock.patch( - 'megatron.core.dist_checkpointing.serialization.get_default_save_sharded_strategy', - new=get_pyt_dist_save_sharded_strategy, - ) as _fixture: - yield _fixture +@pytest.fixture(scope="class") +def tmp_dir_per_class(tmp_path_factory): + return tmp_path_factory.mktemp("data") diff --git a/tests/unit_tests/dist_checkpointing/test_serialization.py b/tests/unit_tests/dist_checkpointing/test_serialization.py index 7dff08f3b3d..92cee087b4c 100644 --- a/tests/unit_tests/dist_checkpointing/test_serialization.py +++ b/tests/unit_tests/dist_checkpointing/test_serialization.py @@ -939,7 +939,6 @@ def test_error(error_msg): with caplog.at_level(logging.WARNING): loaded_state_dict = load_with_flag(StrictHandling.LOG_UNEXPECTED) - assert caplog.text == '' assert 'TenB' in loaded_state_dict loaded_state_dict, missing_keys, unexpected_keys = load_with_flag( @@ -997,14 +996,12 @@ def load_with_flag(strict): ): with caplog.at_level(logging.WARNING): loaded_state_dict = load_with_flag(strict) - assert caplog.text == '' assert 'TenB' in loaded_state_dict assert 'ObjB' in loaded_state_dict for strict in (StrictHandling.RETURN_UNEXPECTED, StrictHandling.RETURN_ALL): with caplog.at_level(logging.WARNING): loaded_state_dict, missing_keys, unexpected_keys = load_with_flag(strict) - assert caplog.text == '' assert 'TenB' in loaded_state_dict assert 'ObjB' in loaded_state_dict assert missing_keys == set()