[optim] bucket checkpoint save to avoid CPU memory spike during ckpt saving

megatron/core/dist_checkpointing/strategies/filesystem_async.py

@@ -83,10 +83,12 @@ def __init__(
         *args,
         separation_hint: Optional[str] = None,
         use_msc: bool = False,
+        staging_max_buckets: int = 0,
         **kwargs,
     ):
         self.checkpoint_dir = path
         self.use_msc = use_msc
+        self.staging_max_buckets = staging_max_buckets
 
         super().__init__(path, *args, **kwargs)
         if not self.single_file_per_rank:
@@ -202,9 +204,20 @@ def get_save_function_and_args(self) -> Tuple[Optional[Callable], Optional[Calla
         if not self.write_buckets:
             return None, None, []
         transform_list = [self.transforms] if hasattr(self, "transforms") else []
+        if self.staging_max_buckets:
+            # Staging happens wave-by-wave inside write_preloaded_data_multiproc,
+            # which then must run in the training process (sync save only).
+            preload_fn = None
+        else:
+            preload_fn = partial(self.preload_tensors, self.write_buckets, True)
         return (
-            partial(self.write_preloaded_data_multiproc, transform_list, self.use_msc),
-            partial(self.preload_tensors, self.write_buckets, True),
+            partial(
+                self.write_preloaded_data_multiproc,
+                transform_list,
+                self.use_msc,
+                staging_max_buckets=self.staging_max_buckets,
+            ),
+            preload_fn,
             [torch.distributed.get_rank(), self.write_buckets, self.results_queue],
         )
 
@@ -238,6 +251,7 @@ def write_preloaded_data_multiproc(
         rank: int,
         write_buckets: List[WriteBucket],
         global_results_queue: mp.Queue,
+        staging_max_buckets: int = 0,
     ) -> None:
         """
         Performs saving data to storage with multiple processes.
@@ -258,11 +272,41 @@ def write_preloaded_data_multiproc(
             write_buckets (List[WriteBucket]): write plan
             global_results_queue (mp.Queue): mp.Queue to collect Dict[List[WriteResults]]
                 (or an Exception) from parallel write processes to the main training process
+            staging_max_buckets (int): when > 0, `write_buckets` still holds GPU tensors
+                and is staged to host and written at most this many buckets at a time,
+                bounding host memory. Requires running in the training process (sync save).
         Returns: None
         """
         logger = logging.getLogger(__name__)
         w_start = time()
         write_results_or_exc: Union[dict, Exception] = dict()
+        wave_size = staging_max_buckets or len(write_buckets)
+        for wave_start in range(0, len(write_buckets), wave_size):
+            wave = write_buckets[wave_start : wave_start + wave_size]
+            if staging_max_buckets:
+                wave = FileSystemWriterAsync.preload_tensors(wave, non_blocking=True)
+            wave_results = FileSystemWriterAsync._fork_and_write_wave(transform_list, use_msc, wave)
+            del wave
+            if isinstance(wave_results, Exception):
+                write_results_or_exc = wave_results
+                break
+            for local_proc_idx, local_results in wave_results.items():
+                write_results_or_exc[wave_start + local_proc_idx] = local_results
+
+        global_results_queue.put(write_results_or_exc)
+
+        w_end = time()
+        logger.debug(f"{w_end}, rank: {rank}, write(sync,parallel): {w_end - w_start}")
+
+    @staticmethod
+    def _fork_and_write_wave(
+        transform_list: List[_StorageWriterTransforms],
+        use_msc: bool,
+        write_buckets: List[WriteBucket],
+    ) -> Union[dict, Exception]:
+        """Fork a writer process per bucket, join them and collect their results."""
+        logger = logging.getLogger(__name__)
+        write_results_or_exc: Union[dict, Exception] = dict()
         ctx = mp.get_context("fork")
         local_results_queue = ctx.Queue()
         count_queue = ctx.JoinableQueue()
@@ -332,10 +376,7 @@ def write_preloaded_data_multiproc(
 
             logger.debug("FileSystemWriterAsync: collected worker results successfully")
 
-        global_results_queue.put(write_results_or_exc)
-
-        w_end = time()
-        logger.debug(f"{w_end}, rank: {rank}, write(sync,parallel): {w_end - w_start}")
+        return write_results_or_exc
 
     @staticmethod
     @_disable_gc()

megatron/core/dist_checkpointing/strategies/torch.py

@@ -626,6 +626,7 @@ def __init__(
         super().__init__(backend, version)
         self.keep_only_main_replica = keep_only_main_replica
         self.thread_count = thread_count
+        self.staging_max_buckets = 0
 
         # Cached SavePlans to skip plan in `save_state_dict_async_plan`
         # cached outcome of `SavePlan.prepare_global_plan`,
@@ -670,6 +671,7 @@ def async_save(
             separation_hint=self.separation_hint,
             thread_count=self.thread_count,
             use_msc=MultiStorageClientFeature.is_enabled(),
+            staging_max_buckets=self.staging_max_buckets,
         )
         # This should be set differently if we run in a smaller process group than the default
         coordinator = 0

megatron/training/arguments.py

@@ -2210,6 +2210,14 @@ def _add_checkpointing_args(parser):
     group.add_argument('--ckpt-fully-parallel-save', action='store_true',
                        dest='ckpt_fully_parallel_save_deprecated',
                        help='Deprecated: see --no-ckpt-fully-parallel-save.')
+    group.add_argument('--ckpt-save-thread-count', type=int, default=2,
+                       help='Number of write buckets (and files) per rank for torch_dist '
+                       'checkpoint save.')
+    group.add_argument('--ckpt-save-staging-buckets', type=int, default=0,
+                       help='Stage and write at most this many buckets at a time during '
+                       'synchronous torch_dist checkpoint save, bounding host memory to '
+                       'roughly this fraction of the rank shard. 0 stages all buckets '
+                       'upfront. Incompatible with --async-save.')
     return parser
 
 

megatron/training/checkpointing.py

@@ -614,6 +614,12 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati
             else:
                 validate_sharding_integrity = True
                 save_strategy = get_default_save_sharded_strategy(args.ckpt_format)
+                save_strategy.thread_count = args.ckpt_save_thread_count
+                if args.ckpt_save_staging_buckets:
+                    assert not args.async_save, \
+                        '--ckpt-save-staging-buckets stages tensors inside the write function, ' \
+                        'which must run in the training process'
+                    save_strategy.staging_max_buckets = args.ckpt_save_staging_buckets
                 if args.ckpt_assume_constant_structure and args.ckpt_format == 'torch_dist':
                     save_strategy.use_cached_ckpt_structure = args.ckpt_assume_constant_structure
                     if checkpointing_context is not None and 'load_strategy' in checkpointing_context: