Merge branch 'main' into cye/te-dcp-test

Author: cspades

megatron/core/dist_checkpointing/serialization.py

@@ -453,6 +453,14 @@ def get_default_save_common_strategy(
     return get_default_strategy(StrategyAction.SAVE_COMMON, backend, version)
 
 
-def get_default_load_sharded_strategy(checkpoint_dir: str) -> LoadShardedStrategy:
-    """Get default load sharded strategy."""
-    return verify_checkpoint_and_load_strategy(checkpoint_dir)[0]
+def get_default_load_sharded_strategy(
+    checkpoint_dir: str, cache_metadata: bool = False
+) -> LoadShardedStrategy:
+    """Get default load sharded strategy.
+
+    Args:
+        checkpoint_dir: Path to the checkpoint directory.
+        cache_metadata: If True and checkpoint format is torch_dist, use a strategy that caches
+            metadata (e.g. when ckpt_assume_constant_structure is enabled).
+    """
+    return verify_checkpoint_and_load_strategy(checkpoint_dir, cache_metadata=cache_metadata)[0]

megatron/core/dist_checkpointing/strategies/cached_metadata_filesystem_reader.py

@@ -3,7 +3,7 @@
 """ FS Reader with metadata cached support. """
 
 import os
-from typing import Union
+from typing import Dict, Union
 
 from torch.distributed.checkpoint import FileSystemReader, Metadata
 
@@ -12,27 +12,42 @@ class CachedMetadataFileSystemReader(FileSystemReader):
     """
     Extends FileSystemReader to cache metadata for improved performance.
 
+    Metadata is shared across all reader instances that use the same checkpoint
+    directory (same path), since the loaded metadata is identical.
+
     Attributes:
-        _cached_metadata (Metadata or None): Cached metadata from the file system.
+        _metadata_cache (Dict[str, Metadata]): Class-level cache keyed by checkpoint path.
     """
 
-    def __init__(self, path: Union[str, os.PathLike]) -> None:
+    _metadata_cache: Dict[str, Metadata] = {}
+
+    def __init__(self, path: Union[str, os.PathLike], cache_metadata: bool = True) -> None:
         """
         Initialize with file system path.
 
         Args:
             path (Union[str, os.PathLike]): Path to the checkpoint directory or file.
         """
         super().__init__(path=path)
-        self._cached_metadata = None
+        self._cache_key = os.path.abspath(os.fspath(path)) if cache_metadata else None
 
     def read_metadata(self) -> Metadata:
         """
         Read metadata from file system, caching for subsequent calls.
+        Shared across instances when the checkpoint directory is the same.
 
         Returns:
             Metadata: Checkpoint metadata.
         """
-        if self._cached_metadata is None:
-            self._cached_metadata = super().read_metadata()
-        return self._cached_metadata
+        if self._cache_key not in CachedMetadataFileSystemReader._metadata_cache:
+            CachedMetadataFileSystemReader._metadata_cache[self._cache_key] = (
+                super().read_metadata()
+            )
+        return CachedMetadataFileSystemReader._metadata_cache[self._cache_key]
+
+    @classmethod
+    def clear_metadata_cache(cls):
+        """
+        Clear the metadata cache.
+        """
+        cls._metadata_cache.clear()

megatron/core/dist_checkpointing/strategies/torch.py

@@ -763,16 +763,17 @@ def _get_filesystem_reader(
         return msc.torch.MultiStorageFileSystemReader(checkpoint_dir, thread_count=2)
 
     if cache_metadata:
-        return CachedMetadataFileSystemReader(checkpoint_dir)
+        return CachedMetadataFileSystemReader(checkpoint_dir, cache_metadata=cache_metadata)
 
     return FileSystemReader(checkpoint_dir)
 
 
 class TorchDistLoadShardedStrategy(LoadShardedStrategy):
     """Basic load strategy for the PyT Distributed format."""
 
-    def __init__(self):
+    def __init__(self, cache_metadata: bool = False):
         self.cached_global_metadata: Optional[Metadata] = None
+        self.cache_metadata = cache_metadata
         super().__init__()
 
     def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> StateDict:
@@ -803,7 +804,7 @@ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> St
         )
         pyt_state_dict = mcore_to_pyt_state_dict(sharded_state_dict, True)
         # Load PyT Distributed format
-        fsr = _get_filesystem_reader(checkpoint_dir, cache_metadata=True)
+        fsr = _get_filesystem_reader(checkpoint_dir, cache_metadata=self.cache_metadata)
         checkpoint.load_state_dict(
             pyt_state_dict,
             fsr,
@@ -815,9 +816,10 @@ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> St
             ),
         )
 
-        self.cached_global_metadata = (
-            fsr.read_metadata()
-        )  # no storage interaction thanks to caching
+        if self.cache_metadata:
+            self.cached_global_metadata = (
+                fsr.read_metadata()
+            )  # no storage interaction thanks to caching
 
         pyt_state_dict = cast(
             Dict[str, Union[TorchShardedTensor, List[io.BytesIO]]], pyt_state_dict

megatron/core/dist_checkpointing/validation.py

@@ -203,6 +203,7 @@ def verify_checkpoint_and_load_strategy(
     checkpoint_dir: str,
     sharded_strategy: Union[LoadShardedStrategy, Tuple[str, int], None] = None,
     common_strategy: Union[LoadCommonStrategy, Tuple[str, int], None] = None,
+    cache_metadata: bool = False,
 ) -> Tuple[LoadShardedStrategy, LoadCommonStrategy]:
     """Verifies if checkpoint metadata exists and matches given strategies.
 
@@ -216,6 +217,8 @@ def verify_checkpoint_and_load_strategy(
         common_strategy (LoadCommonStrategy, Tuple[str, int], optional): common load strategy to be verified
             if compatible with the checkpoint content. If None, the default common load strategy
             for the checkpoint backend will be returned.
+        cache_metadata (bool): if True and checkpoint backend is torch_dist, use a load strategy that caches
+            metadata (e.g. when ckpt_assume_constant_structure is enabled). Ignored if sharded_strategy is set.
     """
     isdir = True
     if MultiStorageClientFeature.is_enabled():
@@ -231,11 +234,18 @@ def verify_checkpoint_and_load_strategy(
         raise CheckpointingException(f"{checkpoint_dir} is not a distributed checkpoint")
 
     if sharded_strategy is None:
-        sharded_strategy = get_default_strategy(
-            StrategyAction.LOAD_SHARDED,
-            saved_config.sharded_backend,
-            saved_config.sharded_backend_version,
-        )
+        if cache_metadata and saved_config.sharded_backend == 'torch_dist':
+            from megatron.core.dist_checkpointing.strategies.torch import (
+                TorchDistLoadShardedStrategy,
+            )
+
+            sharded_strategy = TorchDistLoadShardedStrategy(cache_metadata=True)
+        else:
+            sharded_strategy = get_default_strategy(
+                StrategyAction.LOAD_SHARDED,
+                saved_config.sharded_backend,
+                saved_config.sharded_backend_version,
+            )
     elif isinstance(sharded_strategy, tuple):
         sharded_strategy = get_default_strategy(StrategyAction.LOAD_SHARDED, *sharded_strategy)
 

megatron/training/async_utils.py

@@ -7,6 +7,9 @@
 import logging
 
 from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest
+from megatron.core.dist_checkpointing.strategies.cached_metadata_filesystem_reader import (
+    CachedMetadataFileSystemReader,
+)
 from megatron.core.dist_checkpointing.strategies.filesystem_async import _results_queue
 from megatron.training import get_args
 from megatron.training.utils import print_rank_0
@@ -76,3 +79,4 @@ def reset_persistent_async_worker():
         del _results_queue
     _results_queue = None
     _async_calls_queue = None
+    CachedMetadataFileSystemReader.clear_metadata_cache()

megatron/training/checkpointing.py

@@ -1130,7 +1130,9 @@ def _load_global_dist_base_checkpoint(
         )
 
     checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=True)
-    load_strategy = get_default_load_sharded_strategy(checkpoint_name)
+    load_strategy = get_default_load_sharded_strategy(
+        checkpoint_name, cache_metadata=args.ckpt_assume_constant_structure
+    )
     # NOTE: `args.ckpt_fully_parallel_load` applies to both persistent and non-persistent checkpoints.
     if args.ckpt_fully_parallel_load:
         if args.ckpt_fully_parallel_load_process_group == 'dp':

tests/unit_tests/dist_checkpointing/utils.py

@@ -6,6 +6,9 @@
 
 import torch
 
+from megatron.core.dist_checkpointing.strategies.cached_metadata_filesystem_reader import (
+    CachedMetadataFileSystemReader,
+)
 from megatron.core.models.gpt import GPTModel
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_layer_local_spec,
@@ -167,6 +170,9 @@ def init_checkpointing_mock_args(args, ckpt_dir, fully_parallel=False):
     args.dist_ckpt_optim_fully_reshardable = False
     args.distrib_optim_fully_reshardable_mem_efficient = False
     args.phase_transition_iterations = None
+    # Clear the metadata cache to avoid contamination between tests
+
+    CachedMetadataFileSystemReader.clear_metadata_cache()
 
 
 def setup_model_and_optimizer(
@@ -224,7 +230,7 @@ def setup_model_and_optimizer(
             opt.init_state_fn(opt)
 
     optimizer.reload_model_params()
-
+    CachedMetadataFileSystemReader.clear_metadata_cache()
     return unwrap_model(model), optimizer
 
 
@@ -322,5 +328,5 @@ def setup_moe_model_and_optimizer(
             opt.init_state_fn(opt)
 
     optimizer.reload_model_params()
-
+    CachedMetadataFileSystemReader.clear_metadata_cache()
     return unwrap_model(model), optimizer