diff --git a/megatron/core/dist_checkpointing/strategies/filesystem_async.py b/megatron/core/dist_checkpointing/strategies/filesystem_async.py index b23c4e9893d..fe89e7e5af7 100644 --- a/megatron/core/dist_checkpointing/strategies/filesystem_async.py +++ b/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() diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py index fb0bcaa4b3c..cf5bc5fc2ee 100644 --- a/megatron/core/dist_checkpointing/strategies/torch.py +++ b/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 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 28fea46195d..0c05786bc39 100644 --- a/megatron/training/arguments.py +++ b/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 diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 7c87eca191a..0f6e4c5e5b5 100644 --- a/megatron/training/checkpointing.py +++ b/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: