Sub-file chunk loading bounded by max_batch_bytes

fastsafetensors/common.py

@@ -6,7 +6,7 @@
 import sys
 from collections import OrderedDict
 from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Set, Tuple
 
 from . import cpp as fstcpp
 from .dlpack import from_cuda_buffer
@@ -37,6 +37,31 @@ def is_gpu_found():
     return fstcpp.is_cuda_found() or fstcpp.is_hip_found()
 
 
+def get_fs_type(path: str, mounts_file: str = "/proc/mounts") -> str:
+    """Best-effort filesystem type for *path* (longest mount-point prefix
+    match). Returns "" when it cannot be determined (non-Linux, unreadable
+    mounts table). Used to pick I/O strategies -- e.g. O_DIRECT is a win on
+    local block devices but loses kernel readahead on network filesystems.
+    """
+    try:
+        real = os.path.realpath(path)
+        best_mnt, best_type = "", ""
+        with open(mounts_file) as f:
+            for line in f:
+                parts = line.split()
+                if len(parts) < 3:
+                    continue
+                # /proc/mounts octal-escapes spaces and tabs in mount points
+                mnt = parts[1].replace("\\040", " ").replace("\\011", "\t")
+                fstype = parts[2]
+                if real == mnt or real.startswith(mnt.rstrip("/") + "/"):
+                    if len(mnt) > len(best_mnt):
+                        best_mnt, best_type = mnt, fstype
+        return best_type
+    except OSError:
+        return ""
+
+
 def get_device_numa_node(device: Optional[int]) -> Optional[int]:
     if device is None or not sys.platform.startswith("linux"):
         return None
@@ -340,9 +365,15 @@ def get_tensors(
         device: Device,
         copy_start_offset: int,
         dtype: DType = DType.AUTO,
+        names: Optional[Set[str]] = None,
     ) -> Dict[str, TensorBase]:
+        # ``names`` restricts instantiation to that subset of tensors. Required
+        # when ``gbuf`` is a compacted chunk buffer (holding only some tensors'
+        # bytes): other tensors' computed offsets would fall outside the buffer.
         ret = {}
         for tensor_name, t in self.tensors.items():
+            if names is not None and tensor_name not in names:
+                continue
             dst_dev_ptr = (
                 gbuf.get_base_address()
                 + self.header_length
@@ -428,6 +459,62 @@ def select_byte_ranges(
                 merged.append([s, e])
         return [(s, e) for s, e in merged]
 
+    def plan_chunks(
+        self,
+        max_batch_bytes: int,
+        keep_tensor: Optional[Callable[[str], bool]] = None,
+        merge_gap: int = 4096,
+    ) -> List[Tuple[Set[str], List[Tuple[int, int]]]]:
+        """Partition this shard's tensors into byte-budgeted sub-file chunks.
+
+        Walks the (optionally ``keep_tensor``-filtered) tensors in file-offset
+        order and greedily packs consecutive tensors into a chunk while the
+        chunk's byte span stays within ``max_batch_bytes``. Returns a list of
+        ``(names, byte_ranges)`` pairs; feed each to a copier's ``set_chunk`` to
+        load only that chunk, bounding peak device memory to ~``max_batch_bytes``
+        (the device buffer covers the chunk's span). The byte ranges are the
+        kept tensors' runs within the chunk (kept tensors separated by at most
+        ``merge_gap`` bytes are coalesced), so holes left by a filter are not
+        read even though the buffer spans them.
+
+        A tensor is the atomic load unit, so ``max_batch_bytes`` must be at least
+        the largest kept tensor -- otherwise that tensor fits in no chunk and a
+        ``ValueError`` is raised naming it.
+        """
+        chunks: List[Tuple[Set[str], List[Tuple[int, int]]]] = []
+        cur: List[str] = []
+        cur_runs: List[List[int]] = []  # merged kept runs within the chunk
+        cur_start = 0
+
+        def _push(name: str, s: int, e: int) -> None:
+            if cur_runs and s - cur_runs[-1][1] <= merge_gap:
+                cur_runs[-1][1] = max(cur_runs[-1][1], e)
+            else:
+                cur_runs.append([s, e])
+
+        for name, frame in self.tensors.items():  # insertion order == offset order
+            if keep_tensor is not None and not keep_tensor(name):
+                continue
+            s = self.header_length + frame.data_offsets[0]
+            e = self.header_length + frame.data_offsets[1]
+            tensor_bytes = frame.data_offsets[1] - frame.data_offsets[0]
+            if tensor_bytes > max_batch_bytes:
+                raise ValueError(
+                    f"max_batch_bytes={max_batch_bytes} is smaller than tensor "
+                    f"'{name}' ({tensor_bytes} bytes); it must be at least the "
+                    f"largest kept tensor so every tensor fits in one chunk."
+                )
+            if cur and (e - cur_start) > max_batch_bytes:
+                chunks.append((set(cur), [(s0, e0) for s0, e0 in cur_runs]))
+                cur, cur_runs = [], []
+            if not cur:
+                cur_start = s
+            cur.append(name)
+            _push(name, s, e)
+        if cur:
+            chunks.append((set(cur), [(s0, e0) for s0, e0 in cur_runs]))
+        return chunks
+
     def __repr__(self) -> str:
         return str({"__metadata__": self.metadata, "tensors": self.tensors})
 

fastsafetensors/config.py

@@ -3,7 +3,7 @@
 import json
 import os
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from .common import init_logger
 
@@ -33,6 +33,11 @@ class LoaderConfig:
     queue_size: int = 0
     use_tqdm_on_load: bool = True
 
+    # Cap peak device-buffer bytes per rank by loading each shard in sub-file
+    # chunks (load -> [broadcast] -> release). None keeps whole-shard loading.
+    # Must be >= the largest single tensor. See SafeTensorsMetadata.plan_chunks.
+    max_batch_bytes: Optional[int] = None
+
     _extensions: Dict[str, Dict[str, Any]] = field(default_factory=dict)
 
     def __post_init__(self):
@@ -104,14 +109,19 @@ def from_file(cls, path: str) -> "LoaderConfig":
         return cls._from_json(path)
 
     def create_parallel_kwargs(self) -> Dict[str, Any]:
+        # Memory knobs apply with or without pipelining.
+        common: Dict[str, Any] = {
+            "max_batch_bytes": self.max_batch_bytes,
+        }
         if not self.use_pipeline:
             # queue_size=-1: fully serial (copy_files → broadcast → copy_files),
             # only 1 batch in GPU memory at a time.
-            return {"queue_size": -1}
+            return {"queue_size": -1, **common}
         return {
             "max_concurrent_producers": self.max_concurrent_producers,
             "queue_size": self.queue_size,
             "use_tqdm_on_load": self.use_tqdm_on_load,
+            **common,
         }
 
 

fastsafetensors/copier/nogds.py

@@ -2,7 +2,7 @@
 
 import os
 import sys
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Set, Tuple
 
 from .. import cpp as fstcpp
 from ..common import (
@@ -40,6 +40,8 @@ def __init__(
         self.device = device
         self.reqs: List[int] = []
         self.byte_ranges: Optional[List[Tuple[int, int]]] = None
+        self._chunk_names: Optional[Set[str]] = None
+        self._base_off = metadata.header_length
 
     def set_byte_ranges(self, byte_ranges: Optional[List[Tuple[int, int]]]) -> None:
         """Restrict reads to these ``[start, end)`` absolute file-offset runs.
@@ -51,26 +53,44 @@ def set_byte_ranges(self, byte_ranges: Optional[List[Tuple[int, int]]]) -> None:
         """
         self.byte_ranges = validated_byte_ranges(self.metadata, byte_ranges)
 
+    def set_chunk(self, byte_ranges: List[Tuple[int, int]], names: Set[str]) -> None:
+        """Load only ``names`` into a device buffer sized to those tensors' span.
+
+        Unlike ``set_byte_ranges`` (which still allocates the whole data section
+        and leaves it sparsely filled), this allocates only
+        ``max_end - min_start`` of the runs and instantiates just ``names`` --
+        so peak device memory is bounded by the chunk, not the shard. This is
+        the building block for ``max_batch_bytes`` sub-file batching.
+        """
+        self.byte_ranges = byte_ranges
+        self._chunk_names = names
+
     def submit_io(
         self, use_buf_register: bool, max_copy_block_size: int
     ) -> fstcpp.gds_device_buffer:
         header_length = self.metadata.header_length
-        total_length = self.metadata.size_bytes - header_length
-        gbuf = self.framework.alloc_tensor_memory(total_length, self.device)
         # Default to a single run spanning the whole data section, which
         # reproduces the original full-file read.
         runs = self.byte_ranges
         if runs is None:
             runs = [(header_length, self.metadata.size_bytes)]
+        if self._chunk_names is not None:
+            # Compact chunk: allocate only the runs' span and map gbuf[0] to the
+            # first run's start, so peak memory tracks the chunk, not the shard.
+            base_off = min(s for s, _ in runs)
+            alloc_length = max(e for _, e in runs) - base_off
+        else:
+            base_off = header_length
+            alloc_length = self.metadata.size_bytes - header_length
+        self._base_off = base_off
+        gbuf = self.framework.alloc_tensor_memory(alloc_length, self.device)
         for start, end in runs:
             count = start
             while count < end:
                 l = end - count
                 if max_copy_block_size < l:
                     l = max_copy_block_size
-                req = self.reader.submit_read(
-                    self.fd, gbuf, count, l, count - header_length
-                )
+                req = self.reader.submit_read(self.fd, gbuf, count, l, count - base_off)
                 if req < 0:
                     raise Exception(f"submit_io: submit_nogds_read failed, err={req}")
                 self.reqs.append(req)
@@ -96,7 +116,7 @@ def wait_io(
         if len(failed) > 0:
             raise Exception(f"wait_io: wait_nogds_read failed, reqs={failed}")
         return self.metadata.get_tensors(
-            gbuf, self.device, self.metadata.header_length, dtype=dtype
+            gbuf, self.device, self._base_off, dtype=dtype, names=self._chunk_names
         )