diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 6b8338e143..6cb8c36d96 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -21,6 +21,7 @@ import shutil import sys import threading +from abc import ABC from dataclasses import dataclass, replace from enum import Enum, auto from logging import getLogger @@ -114,6 +115,19 @@ except ImportError: handle_gdn_in_state_dict = None +try: + from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest + from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import FileSystemWriterAsync + from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( + save_state_dict_async_finalize, + save_state_dict_async_plan, + ) + + HAVE_NVRX = True +except (ImportError, ModuleNotFoundError): + NVRxAsyncRequest = ABC + HAVE_NVRX = False + TRACKER_PREFIX = "latest" _CHECKPOINT_VERSION = None @@ -389,6 +403,17 @@ def is_empty_async_queue(global_state: GlobalState) -> bool: return async_queue.get_num_unfinalized_calls() == 0 +def get_save_and_finalize_callbacks(writer, save_state_dict_ret) -> NVRxAsyncRequest: + """Creates an async save request for fsdp_dtensor & torch_dcp with a finalize function.""" + save_fn, preload_fn, save_args = writer.get_save_function_and_args() + + def finalize_fn(): + """Finalizes async checkpointing and synchronizes processes.""" + save_state_dict_async_finalize(*save_state_dict_ret) + + return NVRxAsyncRequest(save_fn, save_args, [finalize_fn], async_fn_kwargs={}, preload_fn=preload_fn) + + def get_rng_state( data_parallel_random_init: bool, ckpt_format: str = "torch_dist", @@ -875,7 +900,7 @@ def save_checkpoint( async_save_request = None if ckpt_cfg.async_save: - if ckpt_type == CheckpointType.GLOBAL and ckpt_cfg.ckpt_format != "torch_dist": + if ckpt_type == CheckpointType.GLOBAL and ckpt_cfg.ckpt_format not in ["torch_dist", "fsdp_dtensor"]: raise NotImplementedError( f"Async checkpoint save not implemented for {ckpt_cfg.ckpt_format} distributed checkpoint format" ) @@ -962,16 +987,35 @@ def save_checkpoint( state_dict = preprocess_fsdp_dtensor_state_dict(cfg, state_dict, model[0]) # FSDP DTensor checkpoint save path using PyTorch Distributed Checkpointing - if MultiStorageClientFeature.is_enabled(): - from multistorageclient.contrib.torch.filesystem import MultiStorageFileSystemWriter + if ckpt_cfg.async_save and HAVE_NVRX: + planner = torch.distributed.checkpoint.DefaultSavePlanner() + coordinator_rank = 0 + fs_storage_writer = FileSystemWriterAsync( + checkpoint_name, + thread_count=ckpt_cfg.dist_ckpt_workers, + use_msc=MultiStorageClientFeature.is_enabled(), + ) - fs_storage_writer = MultiStorageFileSystemWriter(checkpoint_name) + save_state_dict_ret = save_state_dict_async_plan( + state_dict, + fs_storage_writer, + None, + coordinator_rank, + planner=planner, + enable_cache=ckpt_cfg.ckpt_assume_constant_structure, + ) + async_save_request = get_save_and_finalize_callbacks(fs_storage_writer, save_state_dict_ret) else: - fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name) - torch.distributed.checkpoint.save( - state_dict=state_dict, - storage_writer=fs_storage_writer, - ) + if MultiStorageClientFeature.is_enabled(): + from multistorageclient.contrib.torch.filesystem import MultiStorageFileSystemWriter + + fs_storage_writer = MultiStorageFileSystemWriter(checkpoint_name) + else: + fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name) + torch.distributed.checkpoint.save( + state_dict=state_dict, + storage_writer=fs_storage_writer, + ) else: # torch_dist and other formats using MCore distributed checkpointing if checkpointing_context is not None and "save_strategy" in checkpointing_context: diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 2cb543dd69..532ec174e0 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -1199,8 +1199,8 @@ def validate(self) -> None: # Enforce async_save format restriction if self.checkpoint.async_save: - assert self.checkpoint.ckpt_format == "torch_dist", ( - "async_save is only supported with ckpt_format='torch_dist'" + assert self.checkpoint.ckpt_format in ["torch_dist", "fsdp_dtensor"], ( + "async_save is only supported with ckpt_format='torch_dist','fsdp_dtensor'" ) # Set defaults for tensor inspect callback diff --git a/tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_1b.py b/tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_1b.py index 18fcb847d1..70aba4f160 100644 --- a/tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_1b.py +++ b/tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_1b.py @@ -26,9 +26,31 @@ LLAMA_PRETRAIN_RECIPES = [ - # (config_func, name, parallelism_overrides, model_overrides) - (llama32_1b_config, "llama32_1b", {}, {"num_layers": 2}), - (llama32_3b_config, "llama32_3b", {}, {"num_layers": 2}), + # (config_func, name, parallelism_overrides, model_overrides, checkpoint_overrides, ddp_overrides) + (llama32_1b_config, "llama32_1b", {}, {"num_layers": 2}, {}, {}), + (llama32_3b_config, "llama32_3b", {}, {"num_layers": 2}, {}, {}), + # FSDP-async test case + ( + llama32_1b_config, + "llama32_1b", + {}, + {"num_layers": 2}, + { + "ckpt_format": "fsdp_dtensor", + "strict_fsdp_dtensor_load": True, + "async_save": True, + "async_strategy": "nvrx", + "use_persistent_ckpt_worker": True, + "dist_ckpt_workers": 1, + }, + { + "use_megatron_fsdp": True, + "use_distributed_optimizer": True, + "grad_reduce_in_fp32": True, + "average_in_collective": True, + "data_parallel_sharding_strategy": "optim_grads_params", + }, + ), ] @@ -36,13 +58,27 @@ class TestLlamaRecipes: """Test class for LLaMA recipe functional tests.""" @pytest.mark.run_only_on("GPU") - @pytest.mark.parametrize("config_func,recipe_name,parallelism_overrides,model_overrides", LLAMA_PRETRAIN_RECIPES) - def test_llama_pretrain_recipes(self, config_func, recipe_name, parallelism_overrides, model_overrides, tmp_path): + @pytest.mark.parametrize( + "config_func,recipe_name,parallelism_overrides,model_overrides,checkpoint_overrides,ddp_overrides", + LLAMA_PRETRAIN_RECIPES, + ) + def test_llama_pretrain_recipes( + self, + config_func, + recipe_name, + parallelism_overrides, + model_overrides, + checkpoint_overrides, + ddp_overrides, + tmp_path, + ): """Functional test for LLaMA recipes with appropriate parallelism configurations.""" run_pretrain_recipe_test( config_func, recipe_name, tmp_path, model_overrides=model_overrides, + checkpoint_overrides=checkpoint_overrides, + ddp_overrides=ddp_overrides, **parallelism_overrides, ) diff --git a/tests/functional_tests/test_groups/recipes/utils.py b/tests/functional_tests/test_groups/recipes/utils.py index 330e8d04ed..0a3105fbe3 100644 --- a/tests/functional_tests/test_groups/recipes/utils.py +++ b/tests/functional_tests/test_groups/recipes/utils.py @@ -36,6 +36,8 @@ def run_pretrain_recipe_test( pipeline_model_parallel_size: Optional[int] = None, expert_model_parallel_size: Optional[int] = None, model_overrides: Optional[dict] = None, + checkpoint_overrides: Optional[dict] = None, + ddp_overrides: Optional[dict] = None, ): """ Common test implementation for pretrain recipe configurations. @@ -116,6 +118,14 @@ def run_pretrain_recipe_test( for attribute_name, attribute_value in model_overrides.items(): setattr(config.model, attribute_name, attribute_value) + if checkpoint_overrides: + for attribute_name, attribute_value in checkpoint_overrides.items(): + setattr(config.checkpoint, attribute_name, attribute_value) + + if ddp_overrides: + for attribute_name, attribute_value in ddp_overrides.items(): + setattr(config.ddp, attribute_name, attribute_value) + pretrain(config, forward_step) # Basic verification that training completed successfully diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index 257bb14cf5..71a3e6f979 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -1985,9 +1985,7 @@ def test_async_save_format_validation_fsdp_dtensor_fails(self, monkeypatch): dist_config=dist_cfg, ) try: - # Should raise error - async_save with fsdp_dtensor is not allowed - with pytest.raises(AssertionError, match="async_save is only supported with ckpt_format='torch_dist'"): - container.validate() + container.validate() finally: restore_get_world_size_safe(og_ws, cfg_mod)