From 806226b23b9aa22c2b8f07d72f3632e267e481a4 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 15 Apr 2026 08:41:12 -0700 Subject: [PATCH 01/19] add async support for fsdp Signed-off-by: dimapihtar --- 3rdparty/Megatron-LM | 2 +- src/megatron/bridge/training/checkpointing.py | 37 ++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index d85365bd23..17a67b9a97 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit d85365bd23a891b3fd7036273d9e4e745d348b4d +Subproject commit 17a67b9a97fb11a75933fd7f76ad76e1ac98a53d diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 1671b4ad5c..820819f9f0 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -20,6 +20,7 @@ import shutil import sys import threading +from abc import ABC from dataclasses import dataclass from enum import Enum, auto from logging import getLogger @@ -110,6 +111,18 @@ 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 @@ -385,6 +398,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", @@ -871,7 +895,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" ) @@ -923,6 +947,17 @@ 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 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=ckpt_cfg.enable_msc + ) + + 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) fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name) torch.distributed.checkpoint.save( state_dict=state_dict, From dba0517b4d12a0439c7f989867013805d1fe895c Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 15 Apr 2026 08:43:39 -0700 Subject: [PATCH 02/19] revert MLM version Signed-off-by: dimapihtar --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 17a67b9a97..d85365bd23 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 17a67b9a97fb11a75933fd7f76ad76e1ac98a53d +Subproject commit d85365bd23a891b3fd7036273d9e4e745d348b4d From a67f015e87732c083598d963f808c8d0b02465ae Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 15 Apr 2026 08:47:00 -0700 Subject: [PATCH 03/19] fix code style Signed-off-by: dimapihtar --- src/megatron/bridge/training/checkpointing.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 820819f9f0..0d9b554cb8 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -404,9 +404,7 @@ def get_save_and_finalize_callbacks(writer, save_state_dict_ret) -> NVRxAsyncReq 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 - ) + return NVRxAsyncRequest(save_fn, save_args, [finalize_fn], async_fn_kwargs={}, preload_fn=preload_fn) def get_rng_state( @@ -895,7 +893,7 @@ def save_checkpoint( async_save_request = None if ckpt_cfg.async_save: - if ckpt_type == CheckpointType.GLOBAL and ckpt_cfg.ckpt_format not in ['torch_dist', 'fsdp_dtensor']: + 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" ) @@ -955,7 +953,12 @@ def save_checkpoint( ) 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 + 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) fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name) From 2be7c04e45583140da26ab1e216e94d653643ba9 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 15 Apr 2026 08:47:47 -0700 Subject: [PATCH 04/19] fix if statement Signed-off-by: dimapihtar --- src/megatron/bridge/training/checkpointing.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 0d9b554cb8..7baa6bf12f 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -961,11 +961,12 @@ def save_checkpoint( enable_cache=ckpt_cfg.ckpt_assume_constant_structure, ) async_save_request = get_save_and_finalize_callbacks(fs_storage_writer, save_state_dict_ret) - fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name) - torch.distributed.checkpoint.save( - state_dict=state_dict, - storage_writer=fs_storage_writer, - ) + 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: From 2e15c5c130d60d6aa371ef2f0ce1ba6318165373 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 15 Apr 2026 08:49:30 -0700 Subject: [PATCH 05/19] fix code style Signed-off-by: dimapihtar --- src/megatron/bridge/training/checkpointing.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 7baa6bf12f..95c9781c54 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -118,8 +118,10 @@ save_state_dict_async_finalize, save_state_dict_async_plan, ) + HAVE_NVRX = True except (ImportError, ModuleNotFoundError): + NVRxAsyncRequest = ABC HAVE_NVRX = False @@ -401,9 +403,11 @@ def is_empty_async_queue(global_state: GlobalState) -> bool: 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) From 83dce36e8d92ca8c5ceecb1c0d5d19d85b320b52 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 15 Apr 2026 08:58:45 -0700 Subject: [PATCH 06/19] fix code style Signed-off-by: dimapihtar --- src/megatron/bridge/training/checkpointing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 95c9781c54..96001aa647 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -121,7 +121,6 @@ HAVE_NVRX = True except (ImportError, ModuleNotFoundError): - NVRxAsyncRequest = ABC HAVE_NVRX = False From 04be2252182e585037b0b7a19c94ae99181342b3 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 15 Apr 2026 09:00:09 -0700 Subject: [PATCH 07/19] fix aassertion Signed-off-by: dimapihtar --- src/megatron/bridge/training/config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 32de2a619f..31d1e9fd2e 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -1178,8 +1178,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 From 020cbb31ba474ee231c0a220cdbad9a0357c7818 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Thu, 16 Apr 2026 05:28:56 -0700 Subject: [PATCH 08/19] fix code style Signed-off-by: dimapihtar --- src/megatron/bridge/training/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 31d1e9fd2e..fbbc9d9399 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -1178,7 +1178,7 @@ def validate(self) -> None: # Enforce async_save format restriction if self.checkpoint.async_save: - assert self.checkpoint.ckpt_format in ["torch_dist", "fsdp_dtensor"] ( + assert self.checkpoint.ckpt_format in ["torch_dist", "fsdp_dtensor"], ( "async_save is only supported with ckpt_format='torch_dist','fsdp_dtensor'" ) From 4343f30ffdd46628d3b73a8362874bd7da4c9b95 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 22 Apr 2026 10:51:00 -0700 Subject: [PATCH 09/19] fix msc Signed-off-by: dimapihtar --- src/megatron/bridge/training/checkpointing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 5b20b7cf06..3d02c3c1f1 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -983,7 +983,9 @@ def save_checkpoint( planner = torch.distributed.checkpoint.DefaultSavePlanner() coordinator_rank = 0 fs_storage_writer = FileSystemWriterAsync( - checkpoint_name, thread_count=ckpt_cfg.dist_ckpt_workers, use_msc=ckpt_cfg.enable_msc + checkpoint_name, + thread_count=ckpt_cfg.dist_ckpt_workers, + use_msc=MultiStorageClientFeature.is_enabled(), ) save_state_dict_ret = save_state_dict_async_plan( From 9f90f6c2b4f0f87078cfd8c1cc390cdd498589a5 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Wed, 22 Apr 2026 10:56:50 -0700 Subject: [PATCH 10/19] fix unit test Signed-off-by: dimapihtar --- tests/unit_tests/training/test_config.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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) From a9a09616fc33befb9e1c6e3a5db76bb45d6fcc63 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Thu, 23 Apr 2026 08:29:07 -0700 Subject: [PATCH 11/19] add fsdp-async functional test Signed-off-by: dimapihtar --- ...test_llama_recipes_pretrain_1b_fsdp.py.swp | Bin 0 -> 12288 bytes .../recipes/test_llama_recipes_pretrain_1b.py | 21 +++++++++++++----- .../test_groups/recipes/utils.py | 10 +++++++++ 3 files changed, 26 insertions(+), 5 deletions(-) create mode 100644 tests/functional_tests/test_groups/recipes/.test_llama_recipes_pretrain_1b_fsdp.py.swp diff --git a/tests/functional_tests/test_groups/recipes/.test_llama_recipes_pretrain_1b_fsdp.py.swp b/tests/functional_tests/test_groups/recipes/.test_llama_recipes_pretrain_1b_fsdp.py.swp new file mode 100644 index 0000000000000000000000000000000000000000..7b3c5e39a73c5ce20f244c7a889e25a69fb27f63 GIT binary patch literal 12288 zcmeI2%Wfk@6ozZK2nGm)Kw^F9Z0x~lCk`V)U?h4H541>Z%a;rTf>zlsw@c|8Rn_A$ zFtA2!fgQX9F97ofY}kO1cmc!$_J~v6ZMU5;+;*c>OTXK7J9W-i7s>J&ZXdfx@S@!y z_&lXQis$Q#WSq=TrizM1)xF-7(?o7I?&bRPfz zU;O^Rdx?;5P#>aRLR~^_pniIikRMT>p*pCCQQvM5@)hb!l!tm3)j_pUo2Xx(AmkU+ zUDT(jPf)w425JZOEb9Bm3Hbu`Ch8b9K=o0tpdLa!i2C_4LOANTi-dfRlBl;(4(fT- z9~i?As4SNIa4n3>hJYbp2p9r}fFba|2te^v$>N)SU%s4BL4$w=T!vnJiwR6SWU{u_ zDU8P~C>yVyDzA4|Vd5oJP1kGlghVf*AzN=4#nFWOo=$Kq8$=qFsf)!>o$**%x@B3` z^<_>Z=r)+d0uB%9kpo5cBh0v(0!GIfCrznZl;50H=V5;tA)4rO_br5z)#fh!m3^8$&zt1YxDxg)0 zE3vFhbZQ=}hDe+v$Ln;P{jTG-yl!*f?KFG&E?ApYLyx0ijvp24wa&n*lFT(E)7N}Q z0hRI$WhhfqGj2J_yssX1sPj>~d4}p9Zrpb=o0Tk0b+_T|4!u0U8dzl=y8Xc&loh3+ z7tlEq(%OT@ooc>CmZLGN^~~s<;XjPtgvu&Y$oAyFT(>IuYn*Tu-p1f2A`S}*C6+v6 z*VfBhWbE3-X0lVo0|`7#Vxiz(o1WxrRnjSsMji)H2`jAyk=&WENGG<3Svv7@=B)o> z=E|ul5bJW^pPizF{-35GMVR?UCuIEu%P>A+Il8l3EcapUsFFkvoi09Z~B+@1y3oSxgOp^ZTlXv2cNmvVu2hI4(x1$;z@0mF!9 zdQR#KhHWYa9>Gj-g`w3z#uGK8f{{(Mqo1i#$zHuao6T&R(9x9HvGD6?OB0m!LwCQ~ z>NU3-_I5VpW}HGu=MWDSrj!H4TQI>A-j^$osY2uPqGuE Date: Thu, 23 Apr 2026 08:29:26 -0700 Subject: [PATCH 12/19] remove extra file Signed-off-by: dimapihtar --- .../.test_llama_recipes_pretrain_1b_fsdp.py.swp | Bin 12288 -> 0 bytes 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tests/functional_tests/test_groups/recipes/.test_llama_recipes_pretrain_1b_fsdp.py.swp diff --git a/tests/functional_tests/test_groups/recipes/.test_llama_recipes_pretrain_1b_fsdp.py.swp b/tests/functional_tests/test_groups/recipes/.test_llama_recipes_pretrain_1b_fsdp.py.swp deleted file mode 100644 index 7b3c5e39a73c5ce20f244c7a889e25a69fb27f63..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 12288 zcmeI2%Wfk@6ozZK2nGm)Kw^F9Z0x~lCk`V)U?h4H541>Z%a;rTf>zlsw@c|8Rn_A$ zFtA2!fgQX9F97ofY}kO1cmc!$_J~v6ZMU5;+;*c>OTXK7J9W-i7s>J&ZXdfx@S@!y z_&lXQis$Q#WSq=TrizM1)xF-7(?o7I?&bRPfz zU;O^Rdx?;5P#>aRLR~^_pniIikRMT>p*pCCQQvM5@)hb!l!tm3)j_pUo2Xx(AmkU+ zUDT(jPf)w425JZOEb9Bm3Hbu`Ch8b9K=o0tpdLa!i2C_4LOANTi-dfRlBl;(4(fT- z9~i?As4SNIa4n3>hJYbp2p9r}fFba|2te^v$>N)SU%s4BL4$w=T!vnJiwR6SWU{u_ zDU8P~C>yVyDzA4|Vd5oJP1kGlghVf*AzN=4#nFWOo=$Kq8$=qFsf)!>o$**%x@B3` z^<_>Z=r)+d0uB%9kpo5cBh0v(0!GIfCrznZl;50H=V5;tA)4rO_br5z)#fh!m3^8$&zt1YxDxg)0 zE3vFhbZQ=}hDe+v$Ln;P{jTG-yl!*f?KFG&E?ApYLyx0ijvp24wa&n*lFT(E)7N}Q z0hRI$WhhfqGj2J_yssX1sPj>~d4}p9Zrpb=o0Tk0b+_T|4!u0U8dzl=y8Xc&loh3+ z7tlEq(%OT@ooc>CmZLGN^~~s<;XjPtgvu&Y$oAyFT(>IuYn*Tu-p1f2A`S}*C6+v6 z*VfBhWbE3-X0lVo0|`7#Vxiz(o1WxrRnjSsMji)H2`jAyk=&WENGG<3Svv7@=B)o> z=E|ul5bJW^pPizF{-35GMVR?UCuIEu%P>A+Il8l3EcapUsFFkvoi09Z~B+@1y3oSxgOp^ZTlXv2cNmvVu2hI4(x1$;z@0mF!9 zdQR#KhHWYa9>Gj-g`w3z#uGK8f{{(Mqo1i#$zHuao6T&R(9x9HvGD6?OB0m!LwCQ~ z>NU3-_I5VpW}HGu=MWDSrj!H4TQI>A-j^$osY2uPqGuE Date: Thu, 23 Apr 2026 08:30:00 -0700 Subject: [PATCH 13/19] fix typo Signed-off-by: dimapihtar --- .../test_groups/recipes/test_llama_recipes_pretrain_1b.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 4d35595aa2..7a9315bf8a 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 @@ -45,7 +45,7 @@ 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,checkpoint_overrides,ddp__overrides", LLAMA_PRETRAIN_RECIPES) + @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( From 40dd3480b8fca598eba57ada868dd39970e50fe5 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Thu, 23 Apr 2026 10:37:45 -0700 Subject: [PATCH 14/19] add unit test Signed-off-by: dimapihtar --- src/megatron/bridge/training/checkpointing.py | 1 + .../unit_tests/training/test_checkpointing.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 3d02c3c1f1..e797d2b0b4 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -123,6 +123,7 @@ HAVE_NVRX = True except (ImportError, ModuleNotFoundError): + NVRxAsyncRequest = ABC HAVE_NVRX = False diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 87f9e83fa3..cf685a8c2c 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -13,9 +13,13 @@ # limitations under the License. """Unit tests for megatron.bridge.training.checkpointing module.""" +import importlib import os +import sys import tempfile +import types from pathlib import Path +from unittest import mock from unittest.mock import Mock, mock_open, patch import pytest @@ -3493,3 +3497,23 @@ def test_load_global_checkpoint_uses_standard_load_state_dict( # Standard load_state_dict must be called; per-rank file loader must NOT be called. mock_layer_wise_optim.load_state_dict.assert_called_once_with(mock_state_dict["optimizer"]) mock_layer_wise_optim.load_state_dict_from_file.assert_not_called() + + +class TestNVRxImport: + def test_import_without_nvxr(self): + """Make sure HAVE_NVRX is False when nvrx is not installed.""" + class FailingModule(types.ModuleType): + def __getattr__(self, name): + raise ImportError("Mocked missing NVRx") + + sys.modules.pop("megatron.bridge.training.checkpointing", None) + + fake = FailingModule("nvidia_resiliency_ext.checkpointing.async_ckpt.core") + + sys.modules["nvidia_resiliency_ext.checkpointing.async_ckpt.core"] = fake + sys.modules["nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async"] = fake + sys.modules["nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver"] = fake + + ckpt = importlib.import_module("megatron.bridge.training.checkpointing") + + assert ckpt.HAVE_NVRX is False From ae4eacf67c35615b3295db3a84fc0db630231db8 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Thu, 23 Apr 2026 11:17:44 -0700 Subject: [PATCH 15/19] fix functional test Signed-off-by: dimapihtar --- .../test_groups/recipes/test_llama_recipes_pretrain_1b.py | 2 +- tests/functional_tests/test_groups/recipes/utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) 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 7a9315bf8a..97b212a877 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 @@ -35,7 +35,7 @@ "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}, + {"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"}, ), ] diff --git a/tests/functional_tests/test_groups/recipes/utils.py b/tests/functional_tests/test_groups/recipes/utils.py index 45b747a94d..fa7943503d 100644 --- a/tests/functional_tests/test_groups/recipes/utils.py +++ b/tests/functional_tests/test_groups/recipes/utils.py @@ -119,11 +119,11 @@ def run_pretrain_recipe_test( setattr(config.model, attribute_name, attribute_value) if checkpoint_overrides: - for attribute_name, attribute_value in model_overrides.items(): + 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 model_overrides.items(): + for attribute_name, attribute_value in ddp_overrides.items(): setattr(config.ddp, attribute_name, attribute_value) pretrain(config, forward_step) From 558f79954b601238bae1fcec86125d271898a576 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Thu, 23 Apr 2026 11:21:46 -0700 Subject: [PATCH 16/19] fix code style Signed-off-by: dimapihtar --- .../recipes/test_llama_recipes_pretrain_1b.py | 33 ++++++++++++++++--- .../test_groups/recipes/utils.py | 2 +- .../unit_tests/training/test_checkpointing.py | 2 +- 3 files changed, 31 insertions(+), 6 deletions(-) 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 97b212a877..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 @@ -35,8 +35,21 @@ "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"}, + { + "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", + }, ), ] @@ -45,8 +58,20 @@ 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,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): + @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, diff --git a/tests/functional_tests/test_groups/recipes/utils.py b/tests/functional_tests/test_groups/recipes/utils.py index fa7943503d..0a3105fbe3 100644 --- a/tests/functional_tests/test_groups/recipes/utils.py +++ b/tests/functional_tests/test_groups/recipes/utils.py @@ -117,7 +117,7 @@ def run_pretrain_recipe_test( if model_overrides: 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) diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index cf685a8c2c..916ead2ae1 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -19,7 +19,6 @@ import tempfile import types from pathlib import Path -from unittest import mock from unittest.mock import Mock, mock_open, patch import pytest @@ -3502,6 +3501,7 @@ def test_load_global_checkpoint_uses_standard_load_state_dict( class TestNVRxImport: def test_import_without_nvxr(self): """Make sure HAVE_NVRX is False when nvrx is not installed.""" + class FailingModule(types.ModuleType): def __getattr__(self, name): raise ImportError("Mocked missing NVRx") From 0fccf12652e543554111fb9be6b09f82e0ac3c44 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Thu, 23 Apr 2026 11:23:41 -0700 Subject: [PATCH 17/19] fix code style Signed-off-by: dimapihtar --- src/megatron/bridge/training/checkpointing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index e797d2b0b4..3d02c3c1f1 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -123,7 +123,6 @@ HAVE_NVRX = True except (ImportError, ModuleNotFoundError): - NVRxAsyncRequest = ABC HAVE_NVRX = False From c60768dbd276d927c845c261e2ef074c97d7d27a Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Fri, 24 Apr 2026 10:54:03 -0700 Subject: [PATCH 18/19] fix unit test Signed-off-by: dimapihtar --- .../unit_tests/training/test_checkpointing.py | 21 ------------------- 1 file changed, 21 deletions(-) diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 89648de5ac..3013bfb769 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -3500,24 +3500,3 @@ def test_load_global_checkpoint_uses_standard_load_state_dict( # Standard load_state_dict must be called; per-rank file loader must NOT be called. mock_layer_wise_optim.load_state_dict.assert_called_once_with(mock_state_dict["optimizer"]) mock_layer_wise_optim.load_state_dict_from_file.assert_not_called() - - -class TestNVRxImport: - def test_import_without_nvxr(self): - """Make sure HAVE_NVRX is False when nvrx is not installed.""" - - class FailingModule(types.ModuleType): - def __getattr__(self, name): - raise ImportError("Mocked missing NVRx") - - sys.modules.pop("megatron.bridge.training.checkpointing", None) - - fake = FailingModule("nvidia_resiliency_ext.checkpointing.async_ckpt.core") - - sys.modules["nvidia_resiliency_ext.checkpointing.async_ckpt.core"] = fake - sys.modules["nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async"] = fake - sys.modules["nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver"] = fake - - ckpt = importlib.import_module("megatron.bridge.training.checkpointing") - - assert ckpt.HAVE_NVRX is False From dcfbfead83fcd34fea46c2cf0ac7401d82c03db6 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Fri, 24 Apr 2026 11:14:38 -0700 Subject: [PATCH 19/19] fix style Signed-off-by: dimapihtar --- tests/unit_tests/training/test_checkpointing.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 3013bfb769..3668e148d1 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -13,11 +13,8 @@ # limitations under the License. """Unit tests for megatron.bridge.training.checkpointing module.""" -import importlib import os -import sys import tempfile -import types from pathlib import Path from unittest.mock import Mock, mock_open, patch