Merge branch 'main' into rename-thread-count-to-files-per-rank

Author: ronaldw07

docs/user-guide.md

@@ -92,6 +92,8 @@ auto_resume = wrap_trainer_and_auto_resume_with_mlflashpoint(
     # always_save_context=False, # Optional, defaults to False
     # write_files_per_rank=1, # Optional, defaults to 1
     # initial_write_buffer_size_bytes=DESIRED_NUM_BYTES, # Optional, defaults to 16 GB
+    # use_optimized_save=True, # Optional, defaults to True. Uses the optimized save method to reduce write time.
+    # use_cached_ckpt_structure=True, # Optional, defaults to False. Caches the checkpoint structure after identifying 2 consecutive save plan structures that are equal.
 )
 ```
 
@@ -126,6 +128,7 @@ from ml_flashpoint.adapter.megatron.save_strategies import (
 )
 
 # Loading
+import torch.distributed as dist
 from ml_flashpoint.adapter.megatron.load_strategies import MLFlashpointMegatronLoadStrategy
 from ml_flashpoint.checkpoint_object_manager.checkpoint_object_manager import CheckpointObjectManager
 from ml_flashpoint.core.checkpoint_loader import DefaultMLFlashpointCheckpointLoader
@@ -148,6 +151,7 @@ memory_storage_writer = MemoryStorageWriter(...)
 # Use it to instantiate the Save Strategy
 megatron_save_strategy = MLFlashpointMegatronAsyncSaveStrategy(
     storage_writer=memory_storage_writer,
+    # use_cached_ckpt_structure=True, # Optional, defaults to False. Caches the checkpoint structure after identifying 2 consecutive save plan structures that are equal.
 )
 ```
 
@@ -167,7 +171,7 @@ async_request = save_local_aware_megatron_checkpoint(
 
 !!! note
 
-    Make sure to specify the checkpoint ID/path when saving based on the current step using: 
+    Make sure to specify the checkpoint ID/path when saving based on the current step using:
     `CheckpointContainerId.create_child(base_container, CheckpointContainerId.format_version_container(current_step))`
     where `base_container` is the base path CheckpointContainerId used for all checkpoints for the current job, e.g. `"/tmp/mlf-checkpoints/job123"`.
 
@@ -188,6 +192,11 @@ replication_manager.initialize(checkpoint_object_manager)
 checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
     checkpoint_object_manager=checkpoint_object_manager,
     replication_manager=replication_manager,
+    global_rank_getter=dist.get_rank,
+    local_rank_getter=torch.distributed.get_node_local_rank,
+    broadcast_object_list_func=dist.broadcast_object_list,
+    all_gather_object_func=dist.all_gather_object,
+    world_size_getter=dist.get_world_size,
 )
 
 # Instantiate the Load Strategy with the dependencies
@@ -229,11 +238,12 @@ Code: See the [`ml_flashpoint.adapter.pytorch`](https://github.com/google/ml-fla
 To use directly with PyTorch DCP, use the provided `StorageWriter` and `StorageReader` implementations.
 You can use whatever `Planner` implementations work for your use case, or resort to the defaults.
 
-If your per-rank checkpoint data exceeds the default buffer size (16 GB as of this writing), you can increase it using the optional `initial_buffer_size_bytes` parameter. 
+If your per-rank checkpoint data exceeds the default buffer size (16 GB as of this writing), you can increase it using the optional `initial_buffer_size_bytes` parameter.
 
 #### Imports
 ```python
 import torch
+import torch.distributed as dist
 from torch import multiprocessing as torch_mp
 import torch.distributed.checkpoint as dcp
 
@@ -262,6 +272,7 @@ memory_storage_writer = MemoryStorageWriter(
         ckpt_obj_manager=checkpoint_object_manager,
         replication_manager=replication_manager,
         # initial_buffer_size_bytes=initial_write_buffer_size_bytes, # Optional - increase for larger checkpoint sizes per rank
+        # use_optimized_save=True, # Optional, defaults to True. Uses the optimized save method to reduce write time.
     ),
     mp_manager=torch_mp.Manager(),
 )
@@ -270,6 +281,11 @@ memory_storage_writer = MemoryStorageWriter(
 checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
     checkpoint_object_manager=checkpoint_object_manager,
     replication_manager=replication_manager,
+    global_rank_getter=dist.get_rank,
+    local_rank_getter=torch.distributed.get_node_local_rank,
+    broadcast_object_list_func=dist.broadcast_object_list,
+    all_gather_object_func=dist.all_gather_object,
+    world_size_getter=dist.get_world_size,
 )
 memory_storage_reader = MemoryStorageReader(
     path=checkpoint_dir,

src/ml_flashpoint/adapter/nemo/nemo_checkpoint_loader.py

@@ -14,7 +14,7 @@
 
 import os
 from pathlib import Path
-from typing import List, Set
+from typing import Callable, List, Set
 
 from typing_extensions import override
 
@@ -33,6 +33,12 @@ def __init__(
         self,
         checkpoint_object_manager: CheckpointObjectManager,
         replication_manager: ReplicationManager,
+        *,
+        global_rank_getter: Callable[[], int],
+        local_rank_getter: Callable[[], int],
+        broadcast_object_list_func: Callable[..., None],
+        all_gather_object_func: Callable[..., None],
+        world_size_getter: Callable[[], int],
         recover_context: bool = False,
     ):
         """Initializes the NeMoMLFlashpointCheckpointLoader.
@@ -42,9 +48,24 @@ def __init__(
                 reading data.
             replication_manager: The replication manager to use for retrieving
                 missing checkpoint objects from peer nodes.
+            global_rank_getter: A callable that returns the global rank.
+            local_rank_getter: A callable that returns the node-local rank.
+            broadcast_object_list_func: A callable with the same signature as
+                ``torch.distributed.broadcast_object_list``.
+            all_gather_object_func: A callable with the same signature as
+                ``torch.distributed.all_gather_object``.
+            world_size_getter: A callable that returns the world size.
             recover_context: Whether to recover the context directory if missing.
         """
-        super().__init__(checkpoint_object_manager, replication_manager)
+        super().__init__(
+            checkpoint_object_manager,
+            replication_manager,
+            global_rank_getter=global_rank_getter,
+            local_rank_getter=local_rank_getter,
+            broadcast_object_list_func=broadcast_object_list_func,
+            all_gather_object_func=all_gather_object_func,
+            world_size_getter=world_size_getter,
+        )
         self._recover_context = recover_context
 
     @override

src/ml_flashpoint/adapter/nemo/wrapper_util.py

@@ -15,6 +15,7 @@
 from typing import Union
 
 import torch
+import torch.distributed as dist
 from nemo import lightning as nl
 from nemo.lightning.io.pl import MegatronCheckpointIO
 from nemo.lightning.pytorch import strategies as nl_strategies
@@ -79,6 +80,11 @@ def wrap_trainer_and_auto_resume_with_mlflashpoint(
     ckpt_loader = NeMoMLFlashpointCheckpointLoader(
         checkpoint_object_manager=ckpt_obj_manager,
         replication_manager=replication_manager,
+        global_rank_getter=dist.get_rank,
+        local_rank_getter=dist.get_node_local_rank,
+        broadcast_object_list_func=dist.broadcast_object_list,
+        all_gather_object_func=dist.all_gather_object,
+        world_size_getter=dist.get_world_size,
         recover_context=always_save_context,
     )
 

src/ml_flashpoint/core/checkpoint_loader.py

@@ -22,10 +22,9 @@
 import struct
 from collections import defaultdict
 from pathlib import Path
-from typing import IO, List, Optional, Set, Tuple, TypeVar, cast
+from typing import IO, Callable, List, Optional, Set, Tuple, TypeVar, cast
 
 import torch
-import torch.distributed as dist
 from torch.distributed._shard._utils import narrow_tensor_by_index
 from torch.distributed.checkpoint import Metadata
 from torch.distributed.checkpoint.filesystem import _StorageInfo
@@ -128,6 +127,12 @@ def __init__(
         self,
         checkpoint_object_manager: CheckpointObjectManager,
         replication_manager: ReplicationManager,
+        *,
+        global_rank_getter: Callable[[], int],
+        local_rank_getter: Callable[[], int],
+        broadcast_object_list_func: Callable[..., None],
+        all_gather_object_func: Callable[..., None],
+        world_size_getter: Callable[[], int],
     ):
         """Initializes the DefaultMLFlashpointCheckpointLoader.
 
@@ -136,9 +141,21 @@ def __init__(
                 reading data.
             replication_manager: The replication manager to use for retrieving
                 missing checkpoint objects from peer nodes.
+            global_rank_getter: A callable that returns the global rank.
+            local_rank_getter: A callable that returns the node-local rank.
+            broadcast_object_list_func: A callable with the same signature as
+                ``torch.distributed.broadcast_object_list``.
+            all_gather_object_func: A callable with the same signature as
+                ``torch.distributed.all_gather_object``.
+            world_size_getter: A callable that returns the world size.
         """
         self._checkpoint_object_manager = checkpoint_object_manager
         self._replication_manager = replication_manager
+        self._global_rank_getter = global_rank_getter
+        self._local_rank_getter = local_rank_getter
+        self._broadcast_object_list_func = broadcast_object_list_func
+        self._all_gather_object_func = all_gather_object_func
+        self._world_size_getter = world_size_getter
         # Cache for available objects: CheckpointContainerId -> dict[object_path, list[rank]]
         self._available_objects_cache: dict[CheckpointContainerId, dict[str, List[int]]] = {}
 
@@ -337,8 +354,7 @@ def get_latest_complete_checkpoint(
             else continue to the next candidate checkpoint
             - return the checkpoint container id of the latest complete checkpoint
         """
-        # TODO: use global_rank_getter and local_rank_getter.
-        rank = dist.get_rank()
+        rank = self._global_rank_getter()
         _LOGGER.debug(
             "Rank %s: Getting latest complete checkpoint for '%s'",
             rank,
@@ -382,7 +398,7 @@ def get_latest_complete_checkpoint(
                 retrieval_plan = self._compute_retrieval_plan(checkpoint, available_objects_by_rank)
             # Broadcast the retrieval plan to all ranks.
             plan_container = [retrieval_plan]
-            dist.broadcast_object_list(plan_container, src=planner_rank)
+            self._broadcast_object_list_func(plan_container, src=planner_rank)
             retrieval_plan = plan_container[0]
 
             if retrieval_plan is None:
@@ -451,7 +467,7 @@ def _compute_retrieval_plan(
 
         objects_needed_by_local_rank_0.update(self._get_extra_needed_objects(checkpoint, available_objects_by_rank))
 
-        world_size = dist.get_world_size()
+        world_size = self._world_size_getter()
         num_nodes = get_num_of_nodes()
         ranks_per_node = world_size // num_nodes
 
@@ -507,8 +523,8 @@ def get_candidate_checkpoints(
 
         # Scan locally only on the first rank of each node
         base_path = Path(checkpoint_base_container.data)
-        rank = dist.get_rank()
-        local_rank = dist.get_node_local_rank()
+        rank = self._global_rank_getter()
+        local_rank = self._local_rank_getter()
 
         local_candidate_ckpt_ids = []
 
@@ -532,8 +548,8 @@ def get_candidate_checkpoints(
         else:
             _LOGGER.debug("Rank %s: Base path '%s' is not a directory or does not exist.", rank, base_path)
 
-        all_checkpoint_container_path_lists = [None for _ in range(dist.get_world_size())]
-        dist.all_gather_object(all_checkpoint_container_path_lists, local_candidate_ckpt_ids)
+        all_checkpoint_container_path_lists = [None for _ in range(self._world_size_getter())]
+        self._all_gather_object_func(all_checkpoint_container_path_lists, local_candidate_ckpt_ids)
         _LOGGER.debug(
             "Rank %s: Gathered checkpoint container paths from all ranks: '%s'",
             rank,
@@ -589,8 +605,8 @@ def get_checkpoint_objects_by_rank(
 
             local_objects.extend(self._get_extra_local_objects(container_path))
 
-        all_objects_by_rank_paths = [None for _ in range(dist.get_world_size())]
-        dist.all_gather_object(all_objects_by_rank_paths, local_objects)
+        all_objects_by_rank_paths = [None for _ in range(self._world_size_getter())]
+        self._all_gather_object_func(all_objects_by_rank_paths, local_objects)
 
         result = {}
         object_locations = defaultdict(list)
@@ -620,7 +636,7 @@ def retrieve_checkpoint(
                             If empty for this rank, no retrieval is needed.
         """
 
-        rank = dist.get_rank()
+        rank = self._global_rank_getter()
         all_success = True
 
         # Only proceed with retrieval if we have items to retrieve
@@ -656,8 +672,8 @@ def retrieve_checkpoint(
 
         # Gather success status from all ranks
         _LOGGER.debug("Gathering success status from all ranks")
-        all_success_list = [None for _ in range(dist.get_world_size())]
-        dist.all_gather_object(all_success_list, all_success)
+        all_success_list = [None for _ in range(self._world_size_getter())]
+        self._all_gather_object_func(all_success_list, all_success)
         _LOGGER.debug("All success list: '%s'", all_success_list)
         return all(all_success_list)
 

tests/adapter/megatron/test_load_strategies.py

@@ -276,6 +276,11 @@ def test_load_metadata_with_default_loader(checkpoint_directory, mocker):
     loader = DefaultMLFlashpointCheckpointLoader(
         checkpoint_object_manager=CheckpointObjectManager(),
         replication_manager=mock_replication_manager,
+        global_rank_getter=lambda: 0,
+        local_rank_getter=lambda: 0,
+        broadcast_object_list_func=lambda *args, **kwargs: None,
+        all_gather_object_func=lambda *args, **kwargs: None,
+        world_size_getter=lambda: 1,
     )
     strategy = MLFlashpointMegatronLoadStrategy(checkpoint_loader=loader, replication_manager=mock_replication_manager)
 

tests/adapter/nemo/test_auto_resume.py

@@ -64,6 +64,11 @@ def test_initializer(self, mocker):
         checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
             checkpoint_object_manager=chkpt_obj_manager,
             replication_manager=replication_manager,
+            global_rank_getter=lambda: 0,
+            local_rank_getter=lambda: 0,
+            broadcast_object_list_func=lambda *args, **kwargs: None,
+            all_gather_object_func=lambda *args, **kwargs: None,
+            world_size_getter=lambda: 1,
         )
         base_container = CheckpointContainerId("/tmp/ml_flashpoint_checkpoints")
 
@@ -83,6 +88,11 @@ def test_initializer_superclass_properties_are_correct(self, mocker):
         checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
             checkpoint_object_manager=CheckpointObjectManager(),
             replication_manager=ReplicationManager(),
+            global_rank_getter=lambda: 0,
+            local_rank_getter=lambda: 0,
+            broadcast_object_list_func=lambda *args, **kwargs: None,
+            all_gather_object_func=lambda *args, **kwargs: None,
+            world_size_getter=lambda: 1,
         )
         base_container = CheckpointContainerId("/tmp/ml_flashpoint_checkpoints")
 
@@ -103,6 +113,11 @@ def test_initializer_propagates_true_params(self):
         checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
             checkpoint_object_manager=CheckpointObjectManager(),
             replication_manager=ReplicationManager(),
+            global_rank_getter=lambda: 0,
+            local_rank_getter=lambda: 0,
+            broadcast_object_list_func=lambda *args, **kwargs: None,
+            all_gather_object_func=lambda *args, **kwargs: None,
+            world_size_getter=lambda: 1,
         )
         base_container = CheckpointContainerId("/tmp/ml_flashpoint_checkpoints")
 
@@ -122,7 +137,13 @@ def test_initializer_respects_params(self):
         """Tests that init respects the passed parameters for resume flags."""
         # Arrange
         checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
-            checkpoint_object_manager=CheckpointObjectManager(), replication_manager=ReplicationManager()
+            checkpoint_object_manager=CheckpointObjectManager(),
+            replication_manager=ReplicationManager(),
+            global_rank_getter=lambda: 0,
+            local_rank_getter=lambda: 0,
+            broadcast_object_list_func=lambda *args, **kwargs: None,
+            all_gather_object_func=lambda *args, **kwargs: None,
+            world_size_getter=lambda: 1,
         )
         base_container = CheckpointContainerId("/tmp/ml_flashpoint_checkpoints")
 
@@ -142,7 +163,13 @@ def test_initializer_passes_kwargs_to_super(self):
         """Tests that kwargs (like restore_config) are passed to the superclass."""
         # Arrange
         checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
-            checkpoint_object_manager=CheckpointObjectManager(), replication_manager=ReplicationManager()
+            checkpoint_object_manager=CheckpointObjectManager(),
+            replication_manager=ReplicationManager(),
+            global_rank_getter=lambda: 0,
+            local_rank_getter=lambda: 0,
+            broadcast_object_list_func=lambda *args, **kwargs: None,
+            all_gather_object_func=lambda *args, **kwargs: None,
+            world_size_getter=lambda: 1,
         )
         base_container = CheckpointContainerId("/tmp/ml_flashpoint_checkpoints")
         restore_config = RestoreConfig(path="nemo://some-model")