feat(core): implement optimized serialization tensor saving (#12)

- Implement `_save_tensor_optimized` in `CheckpointSaver` to support
writing tensors directly to mmap buffers using `MLF_TENS` format.
- Update `CheckpointLoader` to transparently detect and load both the
new `MLF_TENS` format and legacy `torch.save` format.
- Enable `use_optimized_save` by default in NeMo `wrapper_util` to
leverage performance improvements.

---------

Co-authored-by: g-husam <husameldawi@google.com>

Author: Leahlijuan

src/ml_flashpoint/adapter/nemo/wrapper_util.py

@@ -43,6 +43,7 @@ def wrap_trainer_and_auto_resume_with_mlflashpoint(
     always_save_context: bool = False,
     write_thread_count: int = 1,
     initial_write_buffer_size_bytes: int = DEFAULT_INITIAL_BUFFER_SIZE_BYTES,
+    use_optimized_save: bool = True,
 ) -> MLFlashpointAutoResume:
     """Wraps the trainer and creates an MLFlashpointAutoResume instance wrapping `default_auto_resume`.
 
@@ -87,6 +88,7 @@ def wrap_trainer_and_auto_resume_with_mlflashpoint(
         always_save_context=always_save_context,
         write_thread_count=write_thread_count,
         initial_write_buffer_size_bytes=initial_write_buffer_size_bytes,
+        use_optimized_save=use_optimized_save,
     )
 
     default_auto_resume_args = vars(default_auto_resume) if default_auto_resume else {}
@@ -107,6 +109,7 @@ def wrap_trainer_checkpoint_io_with_mlflashpoint(
     always_save_context: bool = False,
     write_thread_count: int = 1,
     initial_write_buffer_size_bytes: int = DEFAULT_INITIAL_BUFFER_SIZE_BYTES,
+    use_optimized_save: bool = True,
 ):
     """Wraps the trainer's checkpoint I/O with ML Flashpoint capabilities.
 
@@ -202,6 +205,7 @@ def wrap_trainer_checkpoint_io_with_mlflashpoint(
                 ckpt_obj_manager=ckpt_obj_manager,
                 replication_manager=replication_manager,
                 initial_buffer_size_bytes=initial_write_buffer_size_bytes,
+                use_optimized_save=use_optimized_save,
             ),
             mp_manager=torch_mp.Manager(),
             thread_count=write_thread_count,

src/ml_flashpoint/checkpoint_object_manager/buffer_io.py

@@ -323,6 +323,40 @@ def tell(self) -> int:
         self._check_validity()
         return self._pos
 
+    def next_buffer_slice(self, size: int) -> memoryview:
+        """Returns a writable memoryview slice of the buffer at the current position.
+
+        This allows for zero-copy operations into the buffer (e.g., direct tensor copy).
+        The stream position is advanced by `size` bytes.
+
+        Args:
+            size: The size of the slice in bytes.
+
+        Returns:
+            A writable memoryview slice.
+        """
+        self._check_validity("write")
+        if size < 0:
+            raise ValueError(f"Size must be non-negative, got {size}")
+
+        actual_start = METADATA_SIZE + self._pos
+        actual_end = actual_start + size
+
+        if actual_end > len(self._mv):
+            raise ValueError(
+                f"Requested slice (size={size}) exceeds buffer capacity "
+                f"(pos={self._pos}, cap={len(self._mv) - METADATA_SIZE})"
+            )
+
+        # Create the slice
+        slice_mv = self._mv[actual_start:actual_end]
+
+        # Advance position
+        self._pos += size
+        self._update_written_data_length(self._pos)
+
+        return slice_mv
+
     def close(self, truncate: bool = True) -> None:
         """Closes the BufferIO stream and the underlying C++ BufferObject.
 
@@ -433,3 +467,24 @@ def flush(self):
         # Check validity first, as flush() is still an I/O operation.
         self._check_validity()
         pass
+
+    @property
+    def format_signature(self) -> bytes:
+        """Returns the format signature stored in the buffer metadata.
+
+        Returns:
+            The format signature bytes.
+        """
+        self._check_validity()
+        return self._metadata.format_signature
+
+    def set_format_signature(self, signature: bytes) -> None:
+        """Sets the format signature in the buffer metadata.
+
+        Args:
+            signature: The signature bytes to set. Must be at most 8 bytes.
+        """
+        self._check_validity("write")
+        if len(signature) > 8:
+            raise ValueError(f"Format signature must be at most 8 bytes, got {len(signature)}")
+        self._metadata.format_signature = signature

src/ml_flashpoint/checkpoint_object_manager/buffer_metadata.py

@@ -27,10 +27,12 @@ class BufferMetadataType(ctypes.LittleEndianStructure):
     _fields_ = [
         # 8 bytes for the length of valid data written *after* the metadata block
         ("len_written_data", ctypes.c_uint64),
+        # 8 bytes for checkpoint format signature to identify the file format version
+        ("format_signature", ctypes.c_char * 8),
         # Pad the rest of the structure to reach METADATA_SIZE
         (
             "reserved",
-            ctypes.c_uint8 * (METADATA_SIZE - ctypes.sizeof(ctypes.c_uint64)),
+            ctypes.c_uint8 * (METADATA_SIZE - ctypes.sizeof(ctypes.c_uint64) - 8),
         ),
     ]
 

src/ml_flashpoint/core/checkpoint_loader.py

@@ -19,6 +19,7 @@
 import os
 import pickle
 import re
+import struct
 from collections import defaultdict
 from pathlib import Path
 from typing import IO, List, Optional, Tuple, TypeVar, cast
@@ -42,6 +43,7 @@
     COMMON_STATE_FNAME,
     DIRTY_MARKER_SUFFIX,
     GLOBAL_RANK_PATTERN,
+    CheckpointFormat,
     default_metadata_object_name,
 )
 from ml_flashpoint.core.mlf_logging import get_logger
@@ -158,20 +160,48 @@ def read_metadata(
             _LOGGER.exception("Error reading metadata from '%s'", metadata_path)
             raise
 
-    def read_tensor(self, buffer_slice: IO[bytes], req: ReadItem) -> torch.Tensor:
+    def read_tensor(self, buffer_slice: IO[bytes], req: ReadItem, use_optimized_loader: bool = False) -> torch.Tensor:
         """Read tensor from file slice.
 
         Args:
             buffer_slice (IO[bytes]): file slice to read from.
             req (ReadItem): read item.
+            use_optimized_loader (bool): whether to use optimized loader.
 
         Returns:
             torch.Tensor: read tensor.
         """
-        tensor = cast(
-            torch.Tensor,
-            torch.load(cast(IO[bytes], buffer_slice), map_location="cpu", weights_only=True),
-        )
+        pos = buffer_slice.tell()
+        tensor: Optional[torch.Tensor] = None
+
+        if use_optimized_loader:
+            # Read as optimized format (TensorHeader)
+            # First read 4 bytes for length
+            len_bytes = buffer_slice.read(4)
+            if len(len_bytes) == 4:
+                header_len = struct.unpack("<I", len_bytes)[0]
+                # stored header length should be reasonable, if it's too large, it might be legacy format
+                if header_len < 1024 * 1024:
+                    pickle_bytes = buffer_slice.read(header_len)
+
+                    try:
+                        tensor_header = pickle.loads(pickle_bytes)
+
+                        tensor_dtype = tensor_header.dtype
+                        tensor_shape = tensor_header.shape
+                        data_bytes = buffer_slice.read()
+                        tensor = torch.frombuffer(data_bytes, dtype=tensor_dtype)
+                        tensor = tensor.reshape(tensor_shape)
+                    except Exception:
+                        _LOGGER.exception("Failed to parse tensor header")
+                        raise
+        # Fallback to torch.load if optimized loader fails.
+        if tensor is None:
+            buffer_slice.seek(pos)
+            tensor = cast(
+                torch.Tensor,
+                torch.load(cast(IO[bytes], buffer_slice), map_location="cpu", weights_only=True),
+            )
         return narrow_tensor_by_index(tensor, req.storage_offsets, req.lengths)
 
     def _try_retrieve_object_if_missing(self, checkpoint_object_id: CheckpointObjectId) -> bool:
@@ -270,6 +300,11 @@ def read_data(
             raise FileNotFoundError(error_msg)
 
         with self._checkpoint_object_manager.get_buffer(checkpoint_object_id) as stream:
+            use_optimized_loader = False
+            if stream.format_signature == CheckpointFormat.MLF_FORMAT:
+                use_optimized_loader = True
+                _LOGGER.debug("Using optimized loader for '%s'", checkpoint_object_id.data)
+
             for req in read_items:
                 item_md = storage_data[req.storage_index]
                 buffer_slice = cast(IO[bytes], _create_file_view(stream, item_md.offset, item_md.length))
@@ -278,7 +313,7 @@ def read_data(
                     read_bytes.seek(0)
                     planner.load_bytes(req, read_bytes)
                 else:
-                    tensor = self.read_tensor(buffer_slice, req)
+                    tensor = self.read_tensor(buffer_slice, req, use_optimized_loader=use_optimized_loader)
                     target_tensor = planner.resolve_tensor(req).detach()
                     assert target_tensor.size() == tensor.size(), (
                         f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}"

src/ml_flashpoint/core/checkpoint_saver.py

@@ -33,8 +33,9 @@
 from ml_flashpoint.checkpoint_object_manager.checkpoint_object_manager import CheckpointObjectManager
 from ml_flashpoint.checkpoint_object_manager.object_manager import object_manager_ext
 from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId, CheckpointObjectId
-from ml_flashpoint.core.defaults import DIRTY_MARKER_SUFFIX, default_metadata_object_name
+from ml_flashpoint.core.defaults import DIRTY_MARKER_SUFFIX, CheckpointFormat, default_metadata_object_name
 from ml_flashpoint.core.mlf_logging import get_logger
+from ml_flashpoint.core.tensor_header import TensorHeader
 from ml_flashpoint.core.utils import log_execution_time
 from ml_flashpoint.replication.replication_manager import ReplicationManager
 
@@ -294,6 +295,7 @@ def __init__(
         ckpt_obj_manager: CheckpointObjectManager,
         replication_manager: ReplicationManager,
         initial_buffer_size_bytes: int = DEFAULT_INITIAL_BUFFER_SIZE_BYTES,
+        use_optimized_save: bool = True,
     ):
         """Initializes the DefaultMLFlashpointCheckpointSaver.
 
@@ -307,13 +309,16 @@ def __init__(
                 across nodes.
             initial_buffer_size_bytes: The initial buffer size in bytes to use
                 for writing data.
+            use_optimized_save: Whether to use the optimized zero-copy tensor saving.
+                Defaults to True.
         """
         self._global_rank_getter = global_rank_getter
         self._local_rank_getter = local_rank_getter
         self._barrier_func = global_barrier_func
         self._chkpt_obj_manager = ckpt_obj_manager
         self._replication_manager = replication_manager
         self._initial_buffer_size_bytes = initial_buffer_size_bytes
+        self._use_optimized_save = use_optimized_save
 
     @override
     @log_execution_time(logger=_LOGGER, name="initialize_checkpoint")
@@ -443,14 +448,16 @@ def write_data(
         for i in range(1, thread_count):
             thread = threading.Thread(
                 target=self._write_to_buffer_from_queue_worker,
-                args=(object_items_queue, results_from_threads, replicate_after_write),
+                args=(object_items_queue, results_from_threads, replicate_after_write, self._use_optimized_save),
                 name=f"{self.__class__.__name__}-Thread-{i}",
             )
             threads.append(thread)
             thread.start()
 
         # Main thread execution.
-        self._write_to_buffer_from_queue_worker(object_items_queue, results_from_threads, replicate_after_write)
+        self._write_to_buffer_from_queue_worker(
+            object_items_queue, results_from_threads, replicate_after_write, self._use_optimized_save
+        )
 
         for thread in threads:
             thread.join()
@@ -581,13 +588,15 @@ def _write_to_buffer_from_queue_worker(
         object_write_bucket_queue: queue.Queue,
         results_from_threads: queue.Queue,
         replicate_after_write: bool,
+        use_optimized_write: bool,
     ):
         """Worker function for writing data from a queue to buffer objects.
 
         Args:
             object_write_bucket_queue: A queue containing `ObjectWriteBucket` instances to process.
             results_from_threads: A queue to put `(List[WriteResult], Exception)` tuples into.
             replicate_after_write: Whether to trigger async replication of each object after it is written.
+            use_optimized_write: Whether to use optimized write.
         """
         while not object_write_bucket_queue.empty():
             try:
@@ -614,11 +623,19 @@ def _write_to_buffer_from_queue_worker(
                     self._initial_buffer_size_bytes,
                     overwrite=True,
                 ) as buffer_io_writer:
+                    # Set the format signature
+                    if use_optimized_write:
+                        buffer_io_writer.set_format_signature(CheckpointFormat.MLF_FORMAT)
+                    else:
+                        buffer_io_writer.set_format_signature(CheckpointFormat.TORCH_SAVE)
+
                     # First write tensors.
                     for tensor_item, tensor in tensor_tuples:
                         write_start_offset = buffer_io_writer.tell()
-
-                        torch.save(tensor, buffer_io_writer)
+                        if use_optimized_write:
+                            self._save_tensor_optimized(tensor, buffer_io_writer)
+                        else:
+                            torch.save(tensor, buffer_io_writer)
 
                         num_bytes_written = buffer_io_writer.tell() - write_start_offset
                         item_storage_data = _StorageInfo(
@@ -690,3 +707,41 @@ def _remove_older_checkpoints(
                     siblings_to_delete.add(full_path)
 
         return object_manager_ext.delete_directories_async(list(siblings_to_delete))
+
+    def _save_tensor_optimized(self, tensor: torch.Tensor, buffer_io_writer):
+        """Saves a tensor to the buffer using a zero-copy approach where possible.
+
+        NOTE: This method saves the tensor's data in a C-contiguous format,
+        regardless of its original memory layout (stride).
+        The stride information is not saved.
+
+        Format:
+        [4 bytes HEADER_LEN] [HEADER_BYTES (JSON)] [RAW_BYTES]
+
+        Args:
+            tensor: The tensor to save.
+            buffer_io_writer: The BufferIO instance to write to.
+        """
+        # Metadata
+        tensor_header = TensorHeader(dtype=tensor.dtype, shape=tensor.shape)
+
+        # Write Header (Len + JSON)
+        header_data = tensor_header.to_bytes()
+        buffer_io_writer.write(header_data)
+
+        # Write Data (Zero Copy)
+        num_bytes = tensor.numel() * tensor.element_size()
+
+        # Get a writable slice of the underlying C++ buffer
+        if num_bytes > 0:
+            try:
+                dest_mv = buffer_io_writer.next_buffer_slice(num_bytes)
+            except AttributeError:
+                _LOGGER.exception("BufferIO does not support next_buffer_slice, try to disable use_optimized_save.")
+                raise
+
+            # Create a tensor wrapper around the buffer slice
+            dest_tensor = torch.frombuffer(dest_mv, dtype=tensor.dtype, count=tensor.numel()).reshape(tensor.shape)
+
+            # Perform the actual copy.
+            dest_tensor.copy_(tensor)

src/ml_flashpoint/core/defaults.py

@@ -13,12 +13,20 @@
 # limitations under the License.
 
 import re
+from enum import Enum
 
 DIRTY_MARKER_SUFFIX = "unfinished"
 GLOBAL_RANK_PATTERN = re.compile(r"src(\d+)")
 COMMON_STATE_FNAME = "common.pt"
 
 
+class CheckpointFormat(bytes, Enum):
+    # Standard PyTorch save format
+    TORCH_SAVE = b"TORCH___"
+    # Our custom optimized format
+    MLF_FORMAT = b"MLF_TENS"
+
+
 def default_metadata_object_name() -> str:
     """Returns the default object name for metadata files (i.e. filename).