feat(adapter/nemo): adapt to NeMo's new CheckpointIO interface and fix metadata persistence (#99)

This commit adapts ML Flashpoint to recent NeMo and Megatron-Core
updates where `load_content_metadata` is now required during the
checkpoint loading sequence, while also fixing several related
persistence and tensor unwrapping issues.

Author: kkkapu

src/ml_flashpoint/adapter/megatron/load_strategies.py

@@ -95,8 +95,7 @@ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str,
         torch_dist_checkpoint.load(state_dict=pyt_state_dict, storage_reader=storage_reader, planner=planner)
 
         mlf_state_dict: dict[str, Union[TorchShardedTensor, list[io.BytesIO]]] = {
-            k: v if not isinstance(v, TorchShardedTensor) else _unwrap_pyt_sharded_tensor(v)
-            for k, v in pyt_state_dict.items()
+            k: _unwrap_pyt_sharded_tensor(v) for k, v in pyt_state_dict.items()
         }
         mlf_state_dict = _replace_sharded_keys_with_state_dict_keys(mlf_state_dict, flat_mapping, rename_mapping)
         # Need to restore dict key types to handle str<->int conversions for later merging/processing.

src/ml_flashpoint/adapter/nemo/checkpoint_io.py

@@ -26,7 +26,7 @@
 from megatron.core.dist_checkpointing.strategies.async_utils import (
     AsyncRequest as MegatronAsyncRequest,
 )
-from megatron.core.dist_checkpointing.strategies.common import TorchCommonLoadStrategy
+from megatron.core.dist_checkpointing.strategies.common import COMMON_STATE_FNAME, TorchCommonLoadStrategy
 from nemo.lightning.io.pl import MegatronCheckpointIO, TrainerContext, _fix_tensors_device
 from nemo.lightning.pytorch.trainer import Trainer
 from nemo.utils.callbacks.dist_ckpt_io import AsyncCompatibleCheckpointIO, AsyncFinalizableCheckpointIO
@@ -131,9 +131,13 @@ def save_checkpoint(
         """
         if not _is_ml_flashpoint_checkpoint(self.flashpoint_base_dir, path):
             _LOGGER.info("Fallback to alternative checkpoint io.")
-            return self.fallback_checkpoint_io.save_checkpoint(checkpoint, path)
+            return self.fallback_checkpoint_io.save_checkpoint(checkpoint, path, storage_options=storage_options)
         _LOGGER.info("Use ML Flashpoint checkpoint io. Async_save: '%s'", self.async_save)
 
+        content_metadata = (storage_options or {}).get("content_metadata")
+        if content_metadata is not None:
+            checkpoint["content_metadata"] = content_metadata
+
         # Use the helper for local-aware megatron save
         optional_async_request = save_local_aware_megatron_checkpoint(
             checkpoint=checkpoint,
@@ -265,6 +269,50 @@ def remove_checkpoint(self, path: _PATH) -> None:
         else:
             self.fallback_checkpoint_io.remove_checkpoint(path)
 
+    @override
+    @log_execution_time(logger=_LOGGER, name="MLFlashpointCheckpointIO.load_content_metadata", level=logging.INFO)
+    def load_content_metadata(self, path: Optional[_PATH] = None, preloaded_state_dict: Optional[dict] = None) -> dict:
+        """Loads checkpoint content metadata, handling ML Flashpoint checkpoints specifically.
+
+        This implementation is a specialized version of the standard logic found in NeMo's
+        MegatronCheckpointIO (see: https://sourcegraph.com/r/github.com/NVIDIA-NeMo/NeMo@v2.5.0/-/blob/nemo/lightning/io/pl.py?L115),
+        but is tailored to ML Flashpoint.
+
+        Standard helpers like `dist_checkpointing.load_content_metadata` are designed to read
+        from the standard distributed metadata format (e.g., the .metadata file). However,
+        ML Flashpoint checkpoints utilize a stub `metadata.json` containing only
+        {"sharded_backend": ""} to satisfy Megatron's internal validation checks. Because this
+        stub does not contain the actual training state, the standard helper would not be
+        able to retrieve the correct metadata.
+
+        Therefore, this implementation persists the `content_metadata` inside the `common.pt`
+        file and bypasses the stub to explicitly load the metadata from `common.pt`.
+
+        Args:
+            path: The path to the checkpoint directory.
+            preloaded_state_dict: Optional preloaded state dictionary.
+
+        Returns:
+            A dictionary containing the metadata of the checkpoint.
+        """
+        if not _is_ml_flashpoint_checkpoint(self.flashpoint_base_dir, path):
+            _LOGGER.info("Fallback to alternative checkpoint io for load_content_metadata.")
+            return self.fallback_checkpoint_io.load_content_metadata(path, preloaded_state_dict)
+
+        if preloaded_state_dict is not None:
+            return preloaded_state_dict.get("content_metadata")
+
+        common_pt_path = os.path.join(path, COMMON_STATE_FNAME)
+        if os.path.exists(common_pt_path):
+            common_state_dict = torch.load(common_pt_path, map_location="cpu", weights_only=False)
+            if "content_metadata" in common_state_dict:
+                return common_state_dict["content_metadata"]
+
+        # Log a warning if the file exists but doesn't have the expected metadata key
+        _LOGGER.warning("Checkpoint at %s exists but does not contain 'content_metadata'.", path)
+
+        return None
+
 
 class MLFlashpointAsyncFinalizableCheckpointIO(AsyncFinalizableCheckpointIO):
     """CheckpointIO wrapper for async checkpoint saving and synchronous finalization

tests/adapter/megatron/test_load_strategies.py

@@ -435,8 +435,14 @@ def _setup_load_mocks(mocker, global_rank):
         "ml_flashpoint.adapter.megatron.load_strategies.mcore_to_pyt_state_dict",
         return_value=mock_pyt_state_dict,
     )
+
+    def mock_unwrap(sh_ten):
+        if isinstance(sh_ten, list):
+            return sh_ten
+        return [torch.empty(10, 20, dtype=torch.float32)]
+
     mocker.patch(
         "ml_flashpoint.adapter.megatron.load_strategies._unwrap_pyt_sharded_tensor",
-        return_value=[torch.empty(10, 20, dtype=torch.float32)],
+        side_effect=mock_unwrap,
     )
     return load_patched_fn

tests/adapter/nemo/test_checkpoint_io.py

@@ -170,14 +170,18 @@ def test_save_checkpoint_fallback(self, checkpoint_io_components, tmp_path, mock
         checkpoint = {"model": torch.nn.Linear(2, 2)}
         ckpt_version_path = str(tmp_path / "diff_path")
 
+        storage_options = {"content_metadata": {"version": 1}}
+
         expected_return = mocker.MagicMock()
         alt_checkpoint_io.save_checkpoint.return_value = expected_return
 
         # When
-        result = checkpoint_io.save_checkpoint(checkpoint, ckpt_version_path)
+        result = checkpoint_io.save_checkpoint(checkpoint, ckpt_version_path, storage_options=storage_options)
 
         # Then
-        alt_checkpoint_io.save_checkpoint.assert_called_once_with(checkpoint, ckpt_version_path)
+        alt_checkpoint_io.save_checkpoint.assert_called_once_with(
+            checkpoint, ckpt_version_path, storage_options=storage_options
+        )
         assert result is expected_return
 
     def test_save_ml_flashpoint_checkpoint_writes_common_state_dict(self, checkpoint_io_components, mocker):
@@ -222,6 +226,120 @@ def test_save_ml_flashpoint_checkpoint_writes_common_state_dict(self, checkpoint
         loaded_common_state_dict = torch.load(common_state_file_path)
         assert loaded_common_state_dict == common_state_dict
 
+    def test_save_ml_flashpoint_checkpoint_writes_metadata(self, checkpoint_io_components, mocker):
+        """Tests that content_metadata is injected into the checkpoint before saving."""
+        # Given
+        mocker.patch("ml_flashpoint.adapter.megatron.save_utils.torch.distributed.get_node_local_rank", return_value=0)
+        checkpoint_io = checkpoint_io_components["checkpoint_io"]
+        base_path = checkpoint_io_components["base_path"]
+        ckpt_version_path = base_path + "/checkpoint1"
+
+        checkpoint = {"some_state": 123}
+        storage_options = {"content_metadata": {"is_mlf": True}}
+
+        mock_save_preprocess = mocker.patch(
+            "ml_flashpoint.adapter.megatron.save_utils.mcore_state_dict_utils.save_preprocess", return_value=({}, {})
+        )
+        mocker.patch("ml_flashpoint.adapter.megatron.save_utils.torch.save")
+        mocker.patch.object(checkpoint_io, "_save_context")
+
+        # When
+        checkpoint_io.save_checkpoint(checkpoint, ckpt_version_path, storage_options)
+
+        # Then
+        mock_save_preprocess.assert_called_once()
+        modified_checkpoint = mock_save_preprocess.call_args[0][0]
+        assert "content_metadata" in modified_checkpoint
+        assert modified_checkpoint["content_metadata"] == {"is_mlf": True}
+
+    def test_save_ml_flashpoint_checkpoint_does_not_overwrite_existing_metadata(self, checkpoint_io_components, mocker):
+        """Tests that existing content_metadata in the checkpoint is not overwritten
+        if storage_options doesn't provide it."""
+        # Given
+        mocker.patch("ml_flashpoint.adapter.megatron.save_utils.torch.distributed.get_node_local_rank", return_value=0)
+        checkpoint_io = checkpoint_io_components["checkpoint_io"]
+        base_path = checkpoint_io_components["base_path"]
+        ckpt_version_path = base_path + "/checkpoint_no_overwrite"
+
+        # Prepare a checkpoint that already contains metadata
+        original_metadata = {"existing_key": "original_value"}
+        checkpoint = {"model_state": [1, 2, 3], "content_metadata": original_metadata}
+
+        mocker.patch(
+            "ml_flashpoint.adapter.megatron.save_utils.mcore_state_dict_utils.save_preprocess", return_value=({}, {})
+        )
+
+        mocker.patch("ml_flashpoint.adapter.megatron.save_utils.torch.save")
+        mocker.patch.object(checkpoint_io, "_save_context")
+
+        # Scenario 1: storage_options is None
+        # When
+        checkpoint_io.save_checkpoint(checkpoint, ckpt_version_path, storage_options=None)
+        # Then: Verify metadata was not modified or removed
+        assert checkpoint["content_metadata"] == original_metadata
+
+        # Scenario 2: storage_options is an empty dictionary {}
+        # When
+        checkpoint_io.save_checkpoint(checkpoint, ckpt_version_path, storage_options={})
+        # Then: Verify metadata still remains unchanged
+        assert checkpoint["content_metadata"] == original_metadata
+
+    def test_load_content_metadata_fallback(self, checkpoint_io_components, tmp_path):
+        """Tests load_content_metadata falls back to alternative IO for non-MLF paths."""
+        # Given
+        checkpoint_io = checkpoint_io_components["checkpoint_io"]
+        alt_checkpoint_io = checkpoint_io_components["alt_checkpoint_io"]
+        ckpt_version_path = str(tmp_path / "diff_path")
+
+        expected_metadata = {"meta": "fallback"}
+        alt_checkpoint_io.load_content_metadata.return_value = expected_metadata
+
+        # When
+        result = checkpoint_io.load_content_metadata(ckpt_version_path)
+
+        # Then
+        alt_checkpoint_io.load_content_metadata.assert_called_once_with(ckpt_version_path, None)
+        assert result == expected_metadata
+
+    def test_load_content_metadata_from_preloaded(self, checkpoint_io_components):
+        """Tests load_content_metadata prioritizes preloaded_state_dict."""
+        # Given
+        checkpoint_io = checkpoint_io_components["checkpoint_io"]
+        ckpt_version_path = checkpoint_io.flashpoint_base_dir.data + "/checkpoint1"
+
+        expected_metadata = {"from_memory": True}
+        preloaded = {"content_metadata": expected_metadata}
+
+        # When
+        result = checkpoint_io.load_content_metadata(ckpt_version_path, preloaded_state_dict=preloaded)
+
+        # Then
+        assert result == expected_metadata
+
+    def test_load_content_metadata_from_disk(self, checkpoint_io_components, mocker):
+        """Tests load_content_metadata loads from common.pt."""
+        # Given
+        checkpoint_io = checkpoint_io_components["checkpoint_io"]
+        ckpt_version_path = checkpoint_io.flashpoint_base_dir.data + "/checkpoint1"
+
+        mocker.patch("ml_flashpoint.adapter.nemo.checkpoint_io.os.path.exists", return_value=True)
+
+        expected_metadata = {"from_disk": True}
+        # Mock torch.load to return a dictionary containing our expected content_metadata
+        mock_torch_load = mocker.patch(
+            "ml_flashpoint.adapter.nemo.checkpoint_io.torch.load", return_value={"content_metadata": expected_metadata}
+        )
+
+        # When
+        result = checkpoint_io.load_content_metadata(ckpt_version_path)
+
+        # Then
+        # It should load the common state dict from disk safely (CPU, weights_only=False) and extract the metadata
+        mock_torch_load.assert_called_once()
+        assert mock_torch_load.call_args[1]["map_location"] == "cpu"
+        assert mock_torch_load.call_args[1]["weights_only"] is False
+        assert result == expected_metadata
+
     def test_save_ml_flashpoint_checkpoint_async_success(self, checkpoint_io_components, mocker):
         """Tests a successful asynchronous MLF save."""
         # Given