Fit planner: deterministic per-file chunk budgets (device_memory_budget)

.github/workflows/test-paddle.yaml

@@ -50,6 +50,7 @@ jobs:
           COVERAGE_FILE=.coverage_0 pytest -s --cov=${LIBDIR} test_fastsafetensors.py > /tmp/pytest-log/0.log 2>&1
           COVERAGE_FILE=.coverage_ep pytest -s --cov=${LIBDIR} test_ep_slice.py > /tmp/pytest-log/ep_slice.log 2>&1
           COVERAGE_FILE=.coverage_config pytest -s --cov=${LIBDIR} test_config.py > /tmp/pytest-log/config.log 2>&1
+          COVERAGE_FILE=.coverage_planner pytest -s --cov=${LIBDIR} test_planner.py > /tmp/pytest-log/planner.log 2>&1
           COVERAGE_FILE=.coverage_auto pytest -s --cov=${LIBDIR} test_auto_loader.py > /tmp/pytest-log/auto_loader.log 2>&1
           COVERAGE_FILE=.coverage_1 WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 0 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/1.log 2>&1 & \
           COVERAGE_FILE=.coverage_2 WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 1 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/2.log 2>&1 && \

.github/workflows/test-torch.yaml

@@ -49,6 +49,7 @@ jobs:
           COVERAGE_FILE=.coverage_ep pytest -s --cov=${LIBDIR} test_ep_slice.py > /tmp/pytest-log/ep_slice.log 2>&1
           COVERAGE_FILE=.coverage_rob pytest -s --cov=${LIBDIR} test_robustness.py > /tmp/pytest-log/robustness.log 2>&1
           COVERAGE_FILE=.coverage_config pytest -s --cov=${LIBDIR} test_config.py > /tmp/pytest-log/config.log 2>&1
+          COVERAGE_FILE=.coverage_planner pytest -s --cov=${LIBDIR} test_planner.py > /tmp/pytest-log/planner.log 2>&1
           COVERAGE_FILE=.coverage_auto pytest -s --cov=${LIBDIR} test_auto_loader.py > /tmp/pytest-log/auto_loader.log 2>&1
           COVERAGE_FILE=.coverage_3fs pytest -s --cov=${LIBDIR} threefs/ > /tmp/pytest-log/threefs.log 2>&1
           COVERAGE_FILE=.coverage_1 torchrun --nnodes=1 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/1.log 2>&1

fastsafetensors/_planner.py

@@ -1,12 +1,55 @@
 # SPDX-License-Identifier: Apache-2.0
 
-"""Internal planning helpers for sub-file chunked loading.
+"""Internal planners for bounded-memory loading.
 
-Not part of the public API: the entry point is
-``ParallelLoader(max_batch_bytes=...)``, which plans chunks here and feeds
-them to partial-read copiers via the loader's chunk plan.
+Not part of the public API: the entry points are
+``ParallelLoader(max_batch_bytes=..., device_memory_budget=...)``.
+
+plan_chunks partitions one shard into byte-budgeted sub-file chunks; the
+fit planner below turns a whole-load device memory budget into per-file
+chunk budgets.
+
+
+Given the byte sizes of every (kept) tensor -- all known upfront from
+safetensors headers -- and a total ``device_memory_budget`` the load may
+occupy at peak, compute a per-file chunk budget so that
+
+    resident_bytes + live_transient_buffers <= budget
+
+holds at every moment of the load. Budgets are large while cumulative
+resident bytes are small (whole-file loads, no chunking overhead) and
+decline only as the device fills, so chunking cost is paid only where the
+fit actually requires it. The plan is precomputed and deterministic: same
+files, same filter, same budget -> same plan. There is no runtime feedback.
+
+Why per-file budgets are safe (the bound): let ``G(i)`` be the last file of
+file ``i``'s batch group -- the ``group_size`` files loaded concurrently, one
+per rank. While any chunk of file ``i`` is alive, resident bytes are at most
+``R[G(i)+1]`` (only tensors of files ``<= G(i)`` have been materialized; the
+group's kept bytes are charged up front, because broadcast leaves every rank
+holding every file of the group) and every live transient buffer belongs to
+file ``i`` or a later file ``j > i``. The per-file budget ``B`` declines
+monotonically with ``R``, so every live buffer span is ``<= B[i]``. With at
+most ``depth`` buffers alive on any one rank,
+
+    peak <= R[G(i)+1] + depth * B[i] <= budget   (by choice of B[i]).
+
+With ``group_size == 1``, ``G(i) == i`` and this reduces to ``R[i+1]``.
+
+The budget itself is the caller's to choose -- only the caller knows what
+else will live on the device. A caller sizing it from free memory should
+keep a reserve for allocator rounding and the copier's fixed pools (5% or
+1 GiB, whichever is larger, is a reasonable starting point), e.g.::
+
+    free, _ = torch.cuda.mem_get_info(dev)
+    budget = free - max(free // 20, 1 << 30)
+
+and, under broadcast loading, all-reduce(MIN) that value before passing it:
+per-rank readings diverge, and differing budgets would give ranks different
+plans and deadlock the lockstep broadcast sequence.
 """
 
+from dataclasses import dataclass
 from typing import Callable, List, Optional, Set, Tuple
 
 from .common import SafeTensorsMetadata
@@ -67,3 +110,135 @@ def _push(name: str, s: int, e: int) -> None:
     if cur:
         chunks.append((set(cur), [(s0, e0) for s0, e0 in cur_runs]))
     return chunks
+
+
+class BudgetInfeasibleError(ValueError):
+    """The model cannot be loaded within the given device memory budget."""
+
+
+@dataclass(frozen=True)
+class FileWeightStats:
+    """Per-file byte accounting for the fit plan (kept tensors only)."""
+
+    path: str
+    kept_bytes: int  # sum of kept tensor bytes: resident growth
+    span_bytes: int  # last kept byte - first kept byte: single-chunk buffer size
+    largest_tensor: int  # chunk floor: a tensor is the atomic load unit
+
+
+def pipeline_depth(queue_size: int) -> int:
+    """Worst-case number of concurrently live device buffers for a queue size.
+
+    queue_size -1: fully serial, 1 buffer. 0: unbuffered pipeline, the producer
+    materializes the next buffer while the consumer holds one -> 2. n > 0: n in
+    the queue + 1 being produced + 1 being consumed -> n + 2.
+    """
+    if queue_size < 0:
+        return 1
+    return queue_size + 2
+
+
+def collect_file_stats(
+    metas: List[Tuple[str, SafeTensorsMetadata]],
+    keep_tensor: Optional[Callable[[str], bool]] = None,
+) -> List[FileWeightStats]:
+    """Byte accounting per file from already-parsed headers."""
+    stats = []
+    for path, meta in metas:
+        kept = span_start = span_end = largest = 0
+        first = True
+        for name, frame in meta.tensors.items():
+            if keep_tensor is not None and not keep_tensor(name):
+                continue
+            s, e = frame.data_offsets[0], frame.data_offsets[1]
+            kept += e - s
+            largest = max(largest, e - s)
+            if first:
+                span_start, first = s, False
+            span_end = max(span_end, e)
+        stats.append(
+            FileWeightStats(path, kept, span_end - span_start if kept else 0, largest)
+        )
+    return stats
+
+
+def plan_file_budgets(
+    stats: List[FileWeightStats],
+    device_memory_budget: int,
+    depth: int,
+    max_batch_bytes: Optional[int] = None,
+    accumulate_resident: bool = True,
+    transient_multiplier: int = 1,
+    group_size: int = 1,
+) -> List[int]:
+    """Per-file chunk budgets satisfying the peak-memory bound.
+
+    ``accumulate_resident=True`` models consumers that keep every yielded
+    tensor (resident grows by cumulative kept bytes). ``False`` models
+    consumers whose destination memory is already allocated before the load
+    (e.g. copying into preallocated model parameters): resident growth is 0
+    and the plan degenerates to a uniform budget of ``budget / depth``.
+
+    ``transient_multiplier`` scales the per-buffer transient cost: how many
+    times over the copier stages each in-flight chunk. Only the copier knows
+    (its reader path decides), so callers pass
+    ``CopierInterface.chunk_transient_multiplier(paths)``. Fixed overheads
+    that do not scale with chunk size -- bounce-buffer pools, the O_DIRECT
+    reader's thread pool (measured on GB10 unified memory: +~150 MB regardless
+    of chunk size) -- are not modelled here and must be left outside the
+    budget the caller passes.
+
+    ``group_size`` is the number of files loaded concurrently, one per rank
+    (``pg.size()`` under broadcast loading, 1 otherwise). Every rank ends up
+    with every tensor of the group, so a whole group's kept bytes become
+    resident together and each file in it is charged the group total rather
+    than its own prefix -- otherwise the bound below under-counts by up to
+    ``group_size - 1`` files whenever shard sizes are uneven.
+
+    Returns one budget per file; feed each to
+    ``SafeTensorsMetadata.plan_chunks``. A budget >= the file's span yields a
+    single whole-span chunk (no splitting). Raises ``BudgetInfeasibleError``
+    at plan time when some file's largest tensor cannot fit.
+    """
+    if device_memory_budget <= 0:
+        raise BudgetInfeasibleError(
+            f"device_memory_budget={device_memory_budget} must be positive"
+        )
+    if depth < 1:
+        raise ValueError(f"depth must be >= 1, got {depth}")
+    if transient_multiplier < 1:
+        raise ValueError(
+            f"transient_multiplier must be >= 1, got {transient_multiplier}"
+        )
+    if group_size < 1:
+        raise ValueError(f"group_size must be >= 1, got {group_size}")
+    eff_depth = depth * transient_multiplier
+    # Cumulative kept bytes through the end of each file's batch group. With
+    # group_size == 1 this is just R[i+1]; under broadcast the whole group is
+    # in flight at once, so every file in it is charged the group's total.
+    kept_through_group = []
+    running = 0
+    for end in range(len(stats)):
+        running += stats[end].kept_bytes
+        kept_through_group.append(running)
+    group_resident = [
+        kept_through_group[min((i // group_size + 1) * group_size, len(stats)) - 1]
+        for i in range(len(stats))
+    ]
+    budgets = []
+    for i, st in enumerate(stats):
+        resident = group_resident[i] if accumulate_resident else 0
+        b = (device_memory_budget - resident) // eff_depth
+        if max_batch_bytes is not None:
+            b = min(b, max_batch_bytes)
+        if b < st.largest_tensor:
+            required = resident + eff_depth * st.largest_tensor
+            raise BudgetInfeasibleError(
+                f"Model does not fit device_memory_budget: loading '{st.path}' "
+                f"needs >= {required} bytes ({resident} resident + {eff_depth} x "
+                f"{st.largest_tensor} transient), budget is {device_memory_budget}. "
+                f"Reduce pipeline depth (queue_size), free device memory, or pass "
+                f"a larger explicit budget."
+            )
+        budgets.append(b)
+    return budgets

fastsafetensors/config.py

@@ -38,6 +38,13 @@ class LoaderConfig:
     # Must be >= the largest single tensor. See _planner.plan_chunks.
     max_batch_bytes: Optional[int] = None
 
+    # Bound the load's total device footprint in bytes (resident tensors +
+    # transient buffers) per rank via a static fit plan: whole-file loads while
+    # headroom is ample, chunking only where the fit requires it. The caller
+    # picks the number; under broadcast loading it must be identical on every
+    # rank. See fastsafetensors._planner.
+    device_memory_budget: Optional[int] = None
+
     _extensions: Dict[str, Dict[str, Any]] = field(default_factory=dict)
 
     def __post_init__(self):
@@ -112,6 +119,7 @@ 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,
+            "device_memory_budget": self.device_memory_budget,
         }
         if not self.use_pipeline:
             # queue_size=-1: fully serial (copy_files → broadcast → copy_files),

fastsafetensors/copier/__init__.py

@@ -11,7 +11,9 @@
 from .registry import (
     CopierConstructFunc,
     CopierType,
+    copier_class_of,
     create_copier_constructor,
+    get_copier_class,
     register_copier_constructor,
 )
 from .threefs import ThreeFSFileCopier

fastsafetensors/copier/base.py

@@ -90,6 +90,25 @@ def set_chunk(self, byte_ranges: List[Tuple[int, int]], names: Set[str]) -> None
             f"Use the nogds or unified copier, or unset max_batch_bytes."
         )
 
+    @classmethod
+    def chunk_transient_multiplier(cls, paths: List[str]) -> int:
+        """Transient device bytes this copier holds per in-flight chunk, as a
+        multiple of the chunk's span, when loading *paths*.
+
+        The fit planner (``ParallelLoader(device_memory_budget=...)``) charges
+        every live buffer this multiple of its budget, so a copier that stages
+        a chunk twice must say so or the plan under-counts and OOMs. Fixed
+        overheads that do not scale with chunk size (bounce-buffer pools,
+        reader thread pools) are not counted here. Like ``set_chunk``, the
+        default refuses rather than guessing: chunking copiers override it.
+        """
+        raise NotImplementedError(
+            f"device_memory_budget needs a copier that overrides "
+            f"chunk_transient_multiplier; {cls.__name__} does not implement "
+            f"sub-file chunking. Use the nogds or unified copier, or unset "
+            f"device_memory_budget."
+        )
+
     @abstractmethod
     def submit_io(
         self, use_buf_register: bool, max_copy_block_size: int

fastsafetensors/copier/dstorage.py

@@ -177,7 +177,7 @@ def wait_io(self, gbuf, dtype=DType.AUTO, noalign=False):
         )
 
 
-@register_copier_constructor("dstorage")
+@register_copier_constructor("dstorage", DStorageFileCopier)
 def new_dstorage_copier(device: Device, **kwargs) -> CopierConstructFunc:
     """Factory for DirectStorage file copier."""
     init_dstorage(device.index if device.index is not None else 0)

fastsafetensors/copier/gds.py

@@ -225,7 +225,7 @@ def init_gds(framework: Optional[FrameworkOpBase] = None):
         _inited_gds = True
 
 
-@register_copier_constructor("gds")
+@register_copier_constructor("gds", GdsFileCopier)
 def new_gds_file_copier(
     device: Device,
     bbuf_size_kb: int = 16 * 1024,

fastsafetensors/copier/nogds.py

@@ -65,6 +65,17 @@ def set_chunk(self, byte_ranges: List[Tuple[int, int]], names: Set[str]) -> None
         self.byte_ranges = byte_ranges
         self._chunk_names = names
 
+    @classmethod
+    def chunk_transient_multiplier(cls, paths: List[str]) -> int:
+        """Per in-flight-chunk transient cost, as a multiple of chunk span: 1.
+
+        Reads land in the reader's fixed pool of host bounce buffers
+        (``bbuf_size_kb`` x ``max_threads``, sized independently of the chunk),
+        so the only device-side allocation that scales with a chunk is the
+        chunk buffer itself.
+        """
+        return 1
+
     def submit_io(
         self, use_buf_register: bool, max_copy_block_size: int
     ) -> fstcpp.gds_device_buffer:
@@ -142,7 +153,7 @@ def load_library_func(framework=None):
     _loaded_library = True
 
 
-@register_copier_constructor("nogds")
+@register_copier_constructor("nogds", NoGdsFileCopier)
 def new_nogds_file_copier(
     device: Device,
     bbuf_size_kb: int = 16 * 1024,

fastsafetensors/copier/registry.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Callable, Dict
+import functools
+from typing import Any, Callable, Dict, Optional, Type
 
 from ..common import SafeTensorsMetadata
 from ..frameworks import FrameworkOpBase
@@ -14,16 +15,68 @@
 CopierConstructorFactory = Callable[..., CopierConstructFunc]
 
 _copier_registry: Dict[CopierType, CopierConstructorFactory] = {}
+_copier_class_registry: Dict[CopierType, Type[CopierInterface]] = {}
 
 
-def register_copier_constructor(copier_type: CopierType):
+def register_copier_constructor(
+    copier_type: CopierType, copier_class: Optional[Type[CopierInterface]] = None
+):
+    """Register a factory for *copier_type*.
+
+    *copier_class* is the ``CopierInterface`` subclass the factory builds.
+    Registering it lets callers reach the copier's class-level policy (e.g.
+    ``chunk_transient_multiplier``) before any instance exists; see
+    ``copier_class_of``.
+    """
+
     def decorator(factory_func: CopierConstructorFactory) -> CopierConstructorFactory:
-        _copier_registry[copier_type] = factory_func
-        return factory_func
+        @functools.wraps(factory_func)
+        def factory(*args: Any, **kwargs: Any) -> CopierConstructFunc:
+            construct = factory_func(*args, **kwargs)
+            # A factory may delegate to another copier's factory -- gds hands
+            # off to nogds/unified when cuFile is unavailable. The delegate
+            # runs first and tags the constructor it returns, so the innermost
+            # factory wins and the tag names the copier that will actually be
+            # built, not the one that was asked for.
+            if getattr(construct, "copier_class", None) is None:
+                try:
+                    construct.copier_class = copier_class  # type: ignore[attr-defined]
+                except AttributeError:
+                    pass  # exotic callable that rejects attributes; tag is optional
+            return construct
+
+        _copier_registry[copier_type] = factory
+        if copier_class is not None:
+            _copier_class_registry[copier_type] = copier_class
+        return factory
 
     return decorator
 
 
+def get_copier_class(copier_type: CopierType) -> Type[CopierInterface]:
+    """The ``CopierInterface`` subclass registered for *copier_type*.
+
+    This is the statically requested copier. When a constructor is already in
+    hand, prefer ``copier_class_of`` -- a factory may have fallen back to a
+    different copier than the type name suggests.
+
+    Falls back to ``CopierInterface`` itself when a factory was registered
+    without its class, so class-level policy hits the interface defaults
+    (which raise ``NotImplementedError``) rather than a wrong guess.
+    """
+    return _copier_class_registry.get(copier_type, CopierInterface)
+
+
+def copier_class_of(construct: CopierConstructFunc) -> Type[CopierInterface]:
+    """The ``CopierInterface`` subclass *construct* will actually build.
+
+    Reads the tag applied by ``register_copier_constructor``, so it survives a
+    factory delegating to another copier's factory. Unregistered constructors
+    fall back to ``CopierInterface`` (whose class-level policy refuses).
+    """
+    return getattr(construct, "copier_class", None) or CopierInterface
+
+
 def create_copier_constructor(
     copier_type: CopierType, device: Device, **kwargs
 ) -> CopierConstructFunc: