diff --git a/.github/workflows/test-paddle.yaml b/.github/workflows/test-paddle.yaml
index 1d31183..f6a56c4 100644
--- a/.github/workflows/test-paddle.yaml
+++ b/.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 && \
diff --git a/.github/workflows/test-torch.yaml b/.github/workflows/test-torch.yaml
index 4bd171c..e074229 100644
--- a/.github/workflows/test-torch.yaml
+++ b/.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
diff --git a/fastsafetensors/_planner.py b/fastsafetensors/_planner.py
index 253a63f..bb7f543 100644
--- a/fastsafetensors/_planner.py
+++ b/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
diff --git a/fastsafetensors/config.py b/fastsafetensors/config.py
index 89672ae..10a5052 100644
--- a/fastsafetensors/config.py
+++ b/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),
diff --git a/fastsafetensors/copier/__init__.py b/fastsafetensors/copier/__init__.py
index 8f19c71..b7690db 100644
--- a/fastsafetensors/copier/__init__.py
+++ b/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
diff --git a/fastsafetensors/copier/base.py b/fastsafetensors/copier/base.py
index 7604e76..fd80482 100644
--- a/fastsafetensors/copier/base.py
+++ b/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
diff --git a/fastsafetensors/copier/dstorage.py b/fastsafetensors/copier/dstorage.py
index 618474f..03db52d 100644
--- a/fastsafetensors/copier/dstorage.py
+++ b/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)
diff --git a/fastsafetensors/copier/gds.py b/fastsafetensors/copier/gds.py
index c9164e5..97a98d1 100644
--- a/fastsafetensors/copier/gds.py
+++ b/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,
diff --git a/fastsafetensors/copier/nogds.py b/fastsafetensors/copier/nogds.py
index 739636c..1ca0e20 100644
--- a/fastsafetensors/copier/nogds.py
+++ b/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,
diff --git a/fastsafetensors/copier/registry.py b/fastsafetensors/copier/registry.py
index d819644..4304782 100644
--- a/fastsafetensors/copier/registry.py
+++ b/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:
diff --git a/fastsafetensors/copier/threefs.py b/fastsafetensors/copier/threefs.py
index 4151e49..2abb0e5 100644
--- a/fastsafetensors/copier/threefs.py
+++ b/fastsafetensors/copier/threefs.py
@@ -70,7 +70,7 @@ def wait_io(
         )
 
 
-@register_copier_constructor("3fs")
+@register_copier_constructor("3fs", ThreeFSFileCopier)
 def new_threefs_file_copier(
     device: Device,
     mount_point: str,
diff --git a/fastsafetensors/copier/unified.py b/fastsafetensors/copier/unified.py
index 2079f07..db4c779 100644
--- a/fastsafetensors/copier/unified.py
+++ b/fastsafetensors/copier/unified.py
@@ -113,6 +113,26 @@ 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.
+
+        The O_DIRECT reader (dma_load_runs) reads runs straight into the device
+        buffer: each live chunk costs ~1x its span plus a small fixed thread
+        pool (measured +~150 MB on GB10 regardless of chunk size; not charged
+        here). The mmap+pin_memory fallback additionally pins the chunk's file
+        pages for the copy's lifetime; on unified-memory systems both draws
+        come from one physical pool, so each live chunk costs ~2x its span.
+        Mirrors submit_io's own path selection.
+        """
+        if getattr(fstcpp, "dma_load_runs", None) is None:
+            return 2
+        if int(os.environ.get("FASTSAFETENSORS_DMA_THREADS", "8")) <= 0:
+            return 2
+        if any(not _odirect_ok(p) for p in paths):
+            return 2
+        return 1
+
     def submit_io(
         self, use_buf_register: bool, max_copy_block_size: int
     ) -> fstcpp.gds_device_buffer:
@@ -236,7 +256,7 @@ def is_unified_memory_system(framework: Optional[FrameworkOpBase] = None) -> boo
     return "gb10" in framework.get_device_name(0).lower()
 
 
-@register_copier_constructor("unified")
+@register_copier_constructor("unified", UnifiedMemCopier)
 def new_unified_copier(device: Device, **kwargs) -> CopierConstructFunc:
     """Factory function for UnifiedMemCopier.
 
diff --git a/fastsafetensors/loader.py b/fastsafetensors/loader.py
index c7ea4f9..937c160 100644
--- a/fastsafetensors/loader.py
+++ b/fastsafetensors/loader.py
@@ -12,6 +12,7 @@
     OrderedDict,
     Set,
     Tuple,
+    Type,
     Union,
 )
 
@@ -22,7 +23,13 @@
     get_device_numa_node,
     init_logger,
 )
-from .copier import CopierConstructFunc, CopierType, create_copier_constructor
+from .copier import (
+    CopierConstructFunc,
+    CopierInterface,
+    CopierType,
+    copier_class_of,
+    create_copier_constructor,
+)
 from .copier.unified import is_unified_memory_system
 from .file_buffer import FilesBufferOnDevice
 from .frameworks import TensorBase, get_framework_op
@@ -87,6 +94,14 @@ def __init__(
             framework=self.framework,
             **kwargs,
         )
+        # The class behind copier_constructor, for policy the planner needs
+        # before any copier exists (chunk_transient_multiplier). Read off the
+        # constructor rather than copier_type: asking for "gds" on a host
+        # without cuFile hands back a nogds/unified constructor, and the plan
+        # has to reflect the copier that will really run.
+        self.copier_class: Type[CopierInterface] = copier_class_of(
+            self.copier_constructor
+        )
 
     def init_numa(self, set_numa: bool = True):
         global gl_set_numa
diff --git a/fastsafetensors/parallel_loader.py b/fastsafetensors/parallel_loader.py
index 5ceb8bb..e6b348d 100644
--- a/fastsafetensors/parallel_loader.py
+++ b/fastsafetensors/parallel_loader.py
@@ -136,11 +136,14 @@ def __init__(
         # queue_size semantics:
         #   -1 : fully serial — copy_files → broadcast → copy_files (1 batch in GPU mem)
         #    0 : unbuffered pipeline — 1 copying + 1 broadcasting concurrently (2 batches)
-        #   >0 : buffered pipeline — up to (queue_size+1) batches in GPU mem
+        #   >0 : buffered pipeline — up to (queue_size+2) batches in GPU mem
+        #        (queue_size queued + 1 being produced + 1 being consumed)
         queue_size: int = 0,
         use_tqdm_on_load: bool = True,
         tensor_filter: Optional[Callable[[str], bool]] = None,
         max_batch_bytes: Optional[int] = None,
+        device_memory_budget: Optional[int] = None,
+        accumulate_resident: bool = True,
         **kwargs,
     ):
 
@@ -174,8 +177,22 @@ def __init__(
         # bytes) so peak device buffer per rank is bounded regardless of shard
         # size. See _planner.plan_chunks / CopierInterface.set_chunk.
         self.max_batch_bytes = max_batch_bytes
-
-        # Batch files (or, with max_batch_bytes, sub-file chunk-batches)
+        # When set (bytes), bound the load's TOTAL device footprint (resident
+        # tensors + transient buffers) via a static fit plan: whole-file loads
+        # while headroom is ample, per-file chunk budgets declining as the
+        # device fills. The caller picks the number -- it knows what else lives
+        # on the device and, under broadcast loading, must pass the same value
+        # on every rank (e.g. all-reduce(MIN) of each rank's free memory) so
+        # the plan stays identical across ranks. See fastsafetensors._planner.
+        self.device_memory_budget = device_memory_budget
+        # True: the consumer keeps every yielded tensor (resident grows by
+        # cumulative kept bytes). False: tensors are copied into preallocated
+        # destinations (e.g. model params) so resident growth is 0 and the fit
+        # plan degenerates to a uniform per-file budget.
+        self.accumulate_resident = accumulate_resident
+
+        # Batch files (or, with max_batch_bytes / device_memory_budget,
+        # sub-file chunk-batches)
         self.weight_files_batches = self._create_batches(pg)
 
         # Producer-consumer communication
@@ -222,30 +239,72 @@ def _create_batches(self, pg) -> List[List[Any]]:
             self.hf_weights_files[i : i + batch_size]
             for i in range(0, len(self.hf_weights_files), batch_size)
         ]
-        if self.max_batch_bytes is None:
+        if self.max_batch_bytes is None and self.device_memory_budget is None:
             return file_batches
-        # Sub-file chunking: expand each file-batch (one file per rank) into
-        # aligned chunk-batches. Chunk-batch j holds rank r's j-th chunk (or
-        # None once that rank's file runs out), so every rank issues the same
-        # broadcast sequence in lockstep. Header reads are deterministic, so all
-        # ranks build identical batches. Each shard stays owned by one rank and
-        # is loaded in chunks over successive batches -> peak buffer bounded by
-        # max_batch_bytes per rank.
+
         keep = self.loader._tensor_filter
         fw = self.loader.framework
+
+        # Per-file chunk budget. Uniform (max_batch_bytes) by default; with
+        # device_memory_budget, a static fit plan chooses declining budgets so
+        # resident + transient stays within the budget (see planner module).
+        per_file_budget: Optional[Dict[str, int]] = None
+        meta_by_path: Dict[str, SafeTensorsMetadata] = {}
+        if self.device_memory_budget is not None:
+            from ._planner import (
+                collect_file_stats,
+                pipeline_depth,
+                plan_file_budgets,
+            )
+
+            metas = [
+                (f, SafeTensorsMetadata.from_file(f, fw)) for f in self.hf_weights_files
+            ]
+            # Broadcast mode adds one in-flight receive tensor (<= one chunk
+            # budget) on top of the live gbufs; the caller passing the same
+            # budget on every rank keeps the plan deterministic across ranks.
+            # batch_size also sets the group width: those files load together,
+            # one per rank, and every rank keeps all of them.
+            depth = pipeline_depth(self.queue_size) + (1 if batch_size > 1 else 0)
+            # How much transient device memory a live chunk costs is the
+            # copier's own business (e.g. the unified copier's mmap+pin
+            # fallback pins the chunk's pages alongside the device buffer,
+            # costing 2x span on a shared physical pool), so ask it.
+            copier = self.loader.copier_class
+            multiplier = copier.chunk_transient_multiplier([f for f, _ in metas])
+            budgets = plan_file_budgets(
+                collect_file_stats(metas, keep),
+                self.device_memory_budget,
+                depth,
+                max_batch_bytes=self.max_batch_bytes,
+                accumulate_resident=self.accumulate_resident,
+                transient_multiplier=multiplier,
+                group_size=batch_size,
+            )
+            per_file_budget = {f: b for (f, _), b in zip(metas, budgets)}
+            meta_by_path = dict(metas)
+
+        # Expand each file-batch (one file per rank) into aligned chunk-batches.
+        # Chunk-batch j holds rank r's j-th chunk (or None once that rank's file
+        # runs out), so every rank issues the same broadcast sequence in
+        # lockstep. Header reads are deterministic, so all ranks build identical
+        # batches. Each shard stays owned by one rank and is loaded in chunks
+        # over successive batches -> peak buffer bounded by its budget per rank.
+        def _plan_file(f: str) -> List[Tuple[Set[str], List[Tuple[int, int]]]]:
+            if per_file_budget is not None:
+                return plan_chunks(
+                    meta_by_path[f], per_file_budget[f], keep_tensor=keep
+                )
+            assert self.max_batch_bytes is not None
+            return plan_chunks(
+                SafeTensorsMetadata.from_file(f, fw),
+                self.max_batch_bytes,
+                keep_tensor=keep,
+            )
+
         chunk_batches: List[List[Any]] = []
         for group in file_batches:
-            planned = [
-                (
-                    f,
-                    plan_chunks(
-                        SafeTensorsMetadata.from_file(f, fw),
-                        self.max_batch_bytes,
-                        keep_tensor=keep,
-                    ),
-                )
-                for f in group
-            ]
+            planned = [(f, _plan_file(f)) for f in group]
             maxn = max((len(chunks) for _, chunks in planned), default=0)
             for j in range(maxn):
                 spec: List[Any] = []
@@ -312,11 +371,12 @@ def _drain_queue(self):
     def _spec_to_maps(self, spec: List[Any]):
         """Turn a batch spec into (rank_file_map, chunk_plan).
 
-        Without max_batch_bytes the spec is a list of files (one per rank). With
-        it, the spec is a chunk-batch: per rank, either (file, names, ranges) or
-        None (that rank has no chunk this batch).
+        Without sub-file chunking (neither max_batch_bytes nor
+        device_memory_budget set) the spec is a list of files (one per rank).
+        With it, the spec is a chunk-batch: per rank, either
+        (file, names, ranges) or None (that rank has no chunk this batch).
         """
-        if self.max_batch_bytes is None:
+        if self.max_batch_bytes is None and self.device_memory_budget is None:
             return {i: [f] for i, f in enumerate(spec)}, None
         rank_file_map: Dict[int, List[str]] = {}
         chunk_plan: Dict[str, Tuple[Set[str], List[Tuple[int, int]]]] = {}
@@ -604,6 +664,8 @@ def __init__(
         tensor_filter: Optional[Callable[[str], bool]] = None,
         all_local: bool = False,
         max_batch_bytes: Optional[int] = None,
+        device_memory_budget: Optional[int] = None,
+        accumulate_resident: bool = True,
         **kwargs,
     ):
         """Initialize PipelineParallelLoader with a pre-configured SafeTensorsFileLoader.
@@ -648,5 +710,7 @@ def __init__(
             use_tqdm_on_load,
             tensor_filter=tensor_filter,
             max_batch_bytes=max_batch_bytes,
+            device_memory_budget=device_memory_budget,
+            accumulate_resident=accumulate_resident,
             **kwargs,
         )
diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py
index 6190859..c77fb5c 100644
--- a/tests/unit/test_config.py
+++ b/tests/unit/test_config.py
@@ -365,6 +365,7 @@ def test_create_parallel_kwargs_pipeline_enabled(self):
             "queue_size": 2,
             "use_tqdm_on_load": False,
             "max_batch_bytes": None,
+            "device_memory_budget": None,
         }
 
     def test_create_parallel_kwargs_pipeline_disabled(self):
@@ -378,6 +379,7 @@ def test_create_parallel_kwargs_pipeline_disabled(self):
         assert kwargs == {
             "queue_size": -1,
             "max_batch_bytes": None,
+            "device_memory_budget": None,
         }
 
     def test_max_concurrent_producers_validation(self):
diff --git a/tests/unit/test_planner.py b/tests/unit/test_planner.py
new file mode 100644
index 0000000..84dd6fe
--- /dev/null
+++ b/tests/unit/test_planner.py
@@ -0,0 +1,442 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the static fit planner (per-file declining chunk budgets)."""
+
+import random
+
+import pytest
+
+from fastsafetensors import SafeTensorsMetadata
+from fastsafetensors import cpp as fstcpp
+from fastsafetensors._planner import (
+    BudgetInfeasibleError,
+    FileWeightStats,
+    collect_file_stats,
+    pipeline_depth,
+    plan_file_budgets,
+)
+
+GiB = 1024**3
+MiB = 1024**2
+
+
+def _st(path, kept, span=None, largest=None):
+    return FileWeightStats(
+        path,
+        kept,
+        span if span is not None else kept,
+        largest if largest is not None else kept // 4 or kept,
+    )
+
+
+# ---- pipeline depth ----
+
+
+def test_pipeline_depth_mapping():
+    assert pipeline_depth(-1) == 1
+    assert pipeline_depth(0) == 2
+    assert pipeline_depth(1) == 3
+    assert pipeline_depth(4) == 6
+
+
+# ---- copier-supplied transient multiplier ----
+
+
+def test_chunk_transient_multiplier_default_refuses():
+    # Like set_chunk, the interface default refuses rather than guessing, so a
+    # copier that cannot chunk can never silently under-count the plan.
+    from fastsafetensors.copier import CopierInterface, GdsFileCopier
+
+    for cls in (CopierInterface, GdsFileCopier):
+        with pytest.raises(NotImplementedError, match="chunk_transient_multiplier"):
+            cls.chunk_transient_multiplier(["f0"])
+
+
+def test_chunking_copiers_declare_their_transient_cost():
+    # The two overrides have to travel together: a copier that implements
+    # set_chunk but inherits the refusing default would be rejected by
+    # device_memory_budget despite being able to chunk, and one that declares
+    # a cost without implementing set_chunk would be planned for and then fail.
+    from fastsafetensors.copier import CopierInterface, get_copier_class
+    from fastsafetensors.copier.registry import _copier_class_registry
+
+    # Iterate the registry rather than a literal list: dropping the class
+    # argument from a @register_copier_constructor would make a hardcoded
+    # lookup return CopierInterface and satisfy the invariant vacuously, while
+    # breaking device_memory_budget for every user of that copier.
+    assert {"gds", "nogds", "unified", "3fs"} <= set(_copier_class_registry)
+    for name, registered in sorted(_copier_class_registry.items()):
+        cls = get_copier_class(name)
+        assert cls is registered is not CopierInterface, name
+        chunks = cls.set_chunk is not CopierInterface.set_chunk
+        # chunk_transient_multiplier is a classmethod: attribute access builds
+        # a fresh bound method every time, so compare the underlying functions.
+        declares = (
+            cls.chunk_transient_multiplier.__func__
+            is not CopierInterface.chunk_transient_multiplier.__func__
+        )
+        assert chunks == declares, (
+            f"{name}: set_chunk override={chunks} but "
+            f"chunk_transient_multiplier override={declares}"
+        )
+        if declares:
+            assert cls.chunk_transient_multiplier(["f0", "f1"]) >= 1
+    # Exact, not just >= 1: silently bumping nogds to 2 would halve every
+    # budget on the default CPU path without failing anything else.
+    assert get_copier_class("nogds").chunk_transient_multiplier(["f0"]) == 1
+
+
+def test_copier_class_follows_factory_fallback():
+    """A factory may delegate to another copier's factory -- gds hands off to
+    nogds/unified on a host without cuFile (any GPU box without GDS, i.e. the
+    common case on the default nogds=False path). The class the planner reads
+    must be the delegate's, not the requested type's: resolving it from the
+    type name instead makes device_memory_budget die with 'GdsFileCopier does
+    not implement sub-file chunking' while max_batch_bytes, which never
+    consults the class, keeps working on the very same loader."""
+    from fastsafetensors.copier import CopierInterface, copier_class_of
+    from fastsafetensors.copier.registry import (
+        _copier_class_registry,
+        _copier_registry,
+        create_copier_constructor,
+        register_copier_constructor,
+    )
+
+    class _Delegate(CopierInterface):
+        @classmethod
+        def chunk_transient_multiplier(cls, paths):
+            return 1
+
+    class _Front(CopierInterface):
+        pass
+
+    saved = dict(_copier_registry), dict(_copier_class_registry)
+    try:
+
+        @register_copier_constructor("_test_delegate", _Delegate)
+        def _delegate_factory(device, **kwargs):
+            def construct(metadata, device, framework):
+                raise AssertionError("not constructed in this test")
+
+            return construct
+
+        @register_copier_constructor("_test_front", _Front)
+        def _front_factory(device, **kwargs):
+            return _delegate_factory(device, **kwargs)  # hand off, like gds
+
+        assert copier_class_of(create_copier_constructor("_test_delegate", None)) is (
+            _Delegate
+        )
+        # The delegate tags first, so it wins over the front's own class.
+        front = create_copier_constructor("_test_front", None)
+        assert copier_class_of(front) is _Delegate
+        assert copier_class_of(front).chunk_transient_multiplier(["f"]) == 1
+    finally:
+        _copier_registry.clear()
+        _copier_registry.update(saved[0])
+        _copier_class_registry.clear()
+        _copier_class_registry.update(saved[1])
+
+
+def test_chunk_transient_multiplier_unified_tracks_reader_path(monkeypatch):
+    from fastsafetensors.copier import UnifiedMemCopier
+
+    if getattr(fstcpp, "dma_load_runs", None) is None:
+        pytest.skip("built without the O_DIRECT reader; only the fallback exists")
+
+    # O_DIRECT reader available and usable -> chunk staged once.
+    monkeypatch.setenv("FASTSAFETENSORS_ODIRECT", "1")
+    monkeypatch.setenv("FASTSAFETENSORS_DMA_THREADS", "8")
+    assert UnifiedMemCopier.chunk_transient_multiplier(["f0"]) == 1
+
+    # Reader disabled -> mmap+pin fallback pins the chunk alongside the buffer.
+    monkeypatch.setenv("FASTSAFETENSORS_DMA_THREADS", "0")
+    assert UnifiedMemCopier.chunk_transient_multiplier(["f0"]) == 2
+
+    # Network filesystem -> reader skipped for the same reason.
+    monkeypatch.setenv("FASTSAFETENSORS_DMA_THREADS", "8")
+    monkeypatch.setenv("FASTSAFETENSORS_ODIRECT", "0")
+    assert UnifiedMemCopier.chunk_transient_multiplier(["f0"]) == 2
+
+
+# ---- planner math ----
+
+
+def test_budgets_decline_monotonically():
+    stats = [_st(f"f{i}", 4 * GiB) for i in range(8)]
+    budgets = plan_file_budgets(stats, 40 * GiB, depth=2)
+    assert budgets == sorted(budgets, reverse=True)
+    # first file: (40 - 4) / 2 = 18 GiB; last: (40 - 32) / 2 = 4 GiB
+    assert budgets[0] == 18 * GiB
+    assert budgets[-1] == 4 * GiB
+
+
+def test_whole_file_when_ample():
+    # budget >> everything: every per-file budget exceeds its span, so
+    # plan_chunks would return a single whole-span chunk per file.
+    stats = [_st(f"f{i}", 2 * GiB) for i in range(4)]
+    budgets = plan_file_budgets(stats, 100 * GiB, depth=2)
+    assert all(b >= st.span_bytes for b, st in zip(budgets, stats))
+
+
+def test_min_with_max_batch_bytes():
+    stats = [_st("f0", 1 * GiB)]
+    budgets = plan_file_budgets(stats, 100 * GiB, depth=1, max_batch_bytes=512 * MiB)
+    assert budgets == [512 * MiB]
+
+
+def test_group_resident_charges_the_whole_batch_group():
+    # Broadcast leaves every rank holding every file of the batch group, so a
+    # group's kept bytes go resident together: all its files are charged the
+    # group total and share one budget, instead of each file being charged
+    # only its own prefix and leaving the rest of its group unbudgeted.
+    stats = [
+        _st("f0", 8 * GiB, largest=1 * GiB),
+        _st("f1", 2 * GiB, largest=1 * GiB),
+        _st("f2", 8 * GiB, largest=1 * GiB),
+        _st("f3", 2 * GiB, largest=1 * GiB),
+    ]
+    grouped = plan_file_budgets(stats, 40 * GiB, depth=2, group_size=2)
+    assert grouped[0] == grouped[1] and grouped[2] == grouped[3]
+
+    # A group is charged exactly as if it were one file of the group's total
+    # kept bytes -- an independent characterization of the rule.
+    merged = plan_file_budgets(
+        [_st("g0", 10 * GiB, largest=1 * GiB), _st("g1", 10 * GiB, largest=1 * GiB)],
+        40 * GiB,
+        depth=2,
+    )
+    assert [grouped[0], grouped[2]] == merged
+
+    # Charging per file instead would hand f0 more than its group can afford.
+    per_file = plan_file_budgets(stats, 40 * GiB, depth=2, group_size=1)
+    assert per_file[0] > grouped[0]
+    assert per_file[1] == grouped[1]  # last file of a group is charged alike
+
+
+def test_accumulate_resident_false_is_uniform():
+    stats = [_st(f"f{i}", 8 * GiB, largest=1 * GiB) for i in range(6)]
+    budgets = plan_file_budgets(stats, 12 * GiB, depth=2, accumulate_resident=False)
+    assert budgets == [6 * GiB] * 6
+
+
+def test_infeasible_raises_with_details():
+    # 10 files x 2 GiB resident vs 12 GiB budget: infeasible partway through
+    stats = [_st(f"f{i}", 2 * GiB, largest=1 * GiB) for i in range(10)]
+    with pytest.raises(BudgetInfeasibleError) as ei:
+        plan_file_budgets(stats, 12 * GiB, depth=2)
+    msg = str(ei.value)
+    # Name the file that broke the fit and the bytes it needs -- that number is
+    # what a user resizes their budget by, so pin both rather than a substring
+    # that the message contains no matter which file failed.
+    assert "'f5'" in msg, msg
+    assert str(12 * GiB + 2 * GiB) in msg, msg
+    assert "queue_size" in msg
+
+
+def test_nonpositive_budget_and_bad_depth():
+    with pytest.raises(BudgetInfeasibleError):
+        plan_file_budgets([_st("f", 1)], 0, depth=1)
+    # BudgetInfeasibleError subclasses ValueError, so match the message too --
+    # otherwise the wrong exception type would satisfy this.
+    with pytest.raises(ValueError, match="depth"):
+        plan_file_budgets([_st("f", 1)], 1, depth=0)
+    with pytest.raises(ValueError, match="group_size"):
+        plan_file_budgets([_st("f", 1)], 1, depth=1, group_size=0)
+
+
+def test_empty_kept_file_contributes_nothing():
+    stats = [_st("f0", 4 * GiB), FileWeightStats("f1", 0, 0, 0), _st("f2", 4 * GiB)]
+    budgets = plan_file_budgets(stats, 20 * GiB, depth=1)
+    # empty file consumes no resident: f2's budget only reflects f0 + f2
+    assert budgets[2] == 20 * GiB - 8 * GiB
+
+
+# ---- simulation property test: replay the plan, assert peak <= budget ----
+
+
+def _chunk_sizes(span, budget):
+    """The chunk spans plan_chunks would produce for a file of `span` bytes."""
+    if budget >= span:
+        return [span]
+    sizes, rem = [], span
+    while rem > 0:
+        sizes.append(min(budget, rem))
+        rem -= min(budget, rem)
+    return sizes
+
+
+def _simulate_peak(
+    stats, budgets, base_depth, group_size=1, accumulate_resident=True, multiplier=1
+):
+    """Replay the load and return the worst peak on any single rank.
+
+    Files are processed in groups of `group_size`, one per rank; chunk-batch j
+    carries each rank's j-th chunk (nothing once its file runs out). Broadcast
+    leaves every rank holding every tensor, so a batch's bytes become resident
+    on every rank once it is consumed. A rank holds up to `base_depth` of its
+    own chunk buffers (each costing `multiplier` x its span) plus, under
+    broadcast, one in-flight receive tensor.
+
+    Resident is accumulated from the replay itself rather than from the
+    planner's R[G(i)+1] formula, so a wrong formula cannot cancel out here.
+    """
+    batches = []  # list of {rank: (file_idx, size)}
+    for start in range(0, len(stats), group_size):
+        group = list(range(start, min(start + group_size, len(stats))))
+        per_rank = {
+            r: _chunk_sizes(stats[i].span_bytes, budgets[i])
+            for r, i in enumerate(group)
+            if stats[i].kept_bytes > 0
+        }
+        for j in range(max((len(v) for v in per_rank.values()), default=0)):
+            batches.append(
+                {r: (group[r], s[j]) for r, s in per_rank.items() if j < len(s)}
+            )
+
+    # bytes materialized once batch k has been consumed
+    batch_bytes = [sum(sz for _, sz in b.values()) for b in batches]
+    peak = 0
+    for k in range(len(batches)):
+        start = max(0, k - base_depth + 1)
+        resident = sum(batch_bytes[:start]) if accumulate_resident else 0
+        recv = max((sz for _, sz in batches[k].values()), default=0)
+        for r in batches[k]:
+            own = sum(b[r][1] for b in batches[start : k + 1] if r in b)
+            live = multiplier * own + (recv if group_size > 1 else 0)
+            peak = max(peak, resident + live)
+    return peak
+
+
+def test_simulation_peak_within_budget():
+    rng = random.Random(1234)
+    trials = 400
+    multi_chunk = infeasible = grouped = 0
+    for trial in range(trials):
+        n = rng.randint(1, 12)
+        stats = []
+        for i in range(n):
+            kept = rng.randint(1, 64) * MiB
+            largest = max(1, kept // rng.randint(2, 8))
+            stats.append(FileWeightStats(f"f{i}", kept, kept, largest))
+        qs = rng.choice([-1, 0, 1, 3])
+        # The replay's depth model is written out here rather than taken from
+        # pipeline_depth(): the planner below calls the real function, so if it
+        # ever disagrees with these semantics the planner hands out budgets
+        # sized for the wrong number of live buffers and the replay catches it.
+        # Sharing one function would let the error cancel on both sides.
+        base_depth = 1 if qs < 0 else qs + 2
+        group_size = rng.choice([1, 1, 2, 4])
+        mult = rng.choice([1, 1, 2])
+        depth = pipeline_depth(qs) + (1 if group_size > 1 else 0)
+        acc = rng.random() < 0.5
+        max_largest = max(s.largest_tensor for s in stats)
+        # Headroom is sometimes too small for the largest-tensor floor, so the
+        # planner has to refuse rather than hand back an unusable budget.
+        budget = (
+            sum(s.kept_bytes for s in stats)
+            + rng.randint(0, 32) * MiB
+            + rng.randint(0, depth * mult) * max_largest
+        )
+        try:
+            budgets = plan_file_budgets(
+                stats,
+                budget,
+                depth,
+                accumulate_resident=acc,
+                transient_multiplier=mult,
+                group_size=group_size,
+            )
+        except BudgetInfeasibleError:
+            infeasible += 1
+            continue
+        assert all(b >= st.largest_tensor for b, st in zip(budgets, stats))
+        if any(
+            len(_chunk_sizes(st.span_bytes, b)) >= 3 for st, b in zip(stats, budgets)
+        ):
+            multi_chunk += 1
+        if group_size > 1 and n > group_size:
+            grouped += 1
+        peak = _simulate_peak(
+            stats,
+            budgets,
+            base_depth,
+            group_size=group_size,
+            accumulate_resident=acc,
+            multiplier=mult,
+        )
+        # acc=True: resident + transient <= budget. acc=False: the plan bounds
+        # the transient side only (destinations preallocated), and the sim
+        # models exactly that side -> same assertion.
+        assert peak <= budget, (trial, peak, budget, acc, group_size, mult, qs)
+    # Guard against the corpus quietly degenerating: the bound says nothing
+    # where nothing is split, the infeasibility floor is untested if no plan is
+    # ever refused, and the group-resident rule is untested without multi-rank
+    # groups. Failing these means the generator drifted, not the planner.
+    assert multi_chunk > 40, f"only {multi_chunk}/{trials} trials split a file 3+ ways"
+    assert infeasible > 20, f"only {infeasible}/{trials} trials hit the budget floor"
+    assert grouped > 40, f"only {grouped}/{trials} trials used multi-rank groups"
+
+
+# ---- collect_file_stats on real files ----
+
+
+def test_collect_file_stats_matches_headers(input_files, framework):
+    meta = SafeTensorsMetadata.from_file(input_files[0], framework)
+    (st,) = collect_file_stats([(input_files[0], meta)])
+    total = sum(f.data_offsets[1] - f.data_offsets[0] for f in meta.tensors.values())
+    largest = max(f.data_offsets[1] - f.data_offsets[0] for f in meta.tensors.values())
+    assert st.kept_bytes == total
+    assert st.largest_tensor == largest
+    assert st.span_bytes >= largest
+
+    # filtered: keep every other tensor -> kept < total, span may include holes
+    names = sorted(meta.tensors.keys())
+    keep = set(names[::2])
+    (fst,) = collect_file_stats([(input_files[0], meta)], lambda n: n in keep)
+    assert 0 < fst.kept_bytes < total
+    assert fst.span_bytes >= fst.kept_bytes
+
+
+# ---- end-to-end on CPU: budgeted load is byte-identical to a plain load ----
+
+
+def test_parallel_loader_device_memory_budget_cpu(input_files, framework):
+    if framework.get_name() != "pytorch":
+        pytest.skip("pytorch-only integration test")
+    import torch
+    from safetensors.torch import load_file
+
+    from fastsafetensors import ParallelLoader
+
+    expected = load_file(input_files[0])
+    meta = SafeTensorsMetadata.from_file(input_files[0], framework)
+    (st,) = collect_file_stats([(input_files[0], meta)])
+    # tight budget: forces chunking (several chunks) but stays feasible
+    budget = st.kept_bytes + pipeline_depth(0) * st.largest_tensor + 4096
+
+    pl = ParallelLoader(
+        pg=None,
+        hf_weights_files=[input_files[0]],
+        device="cpu",
+        nogds=True,
+        use_tqdm_on_load=False,
+        device_memory_budget=budget,
+    )
+    got = dict(pl.iterate_weights())
+    assert set(got.keys()) == set(expected.keys())
+    for k in expected:
+        assert torch.equal(got[k], expected[k]), k
+
+
+def test_transient_multiplier_infeasible():
+    # a plan feasible at 1x must fail at 2x when the budget is tight
+    stats = [_st("f0", 2 * GiB)]
+    budget = stats[0].kept_bytes + 2 * stats[0].largest_tensor  # depth=2, k=1 fits
+    plan_file_budgets(stats, budget, depth=2, transient_multiplier=1)
+    with pytest.raises(BudgetInfeasibleError):
+        plan_file_budgets(stats, budget, depth=2, transient_multiplier=2)
+
+
+def test_transient_multiplier_validation():
+    with pytest.raises(ValueError):
+        plan_file_budgets([_st("f0", GiB)], GiB, depth=1, transient_multiplier=0)
diff --git a/tests/unit/test_robustness.py b/tests/unit/test_robustness.py
index fac7172..db5dd29 100644
--- a/tests/unit/test_robustness.py
+++ b/tests/unit/test_robustness.py
@@ -220,3 +220,81 @@ def test_gds_fallback_warns_once_and_shares_reader(
     assert len(cache) == 1  # one shared nogds constructor for the whole loader
     framework.free_tensor_memory(g1, device)
     framework.free_tensor_memory(g2, device)
+
+
+# ---- device_memory_budget in broadcast mode ----
+
+
+class _FakePG:
+    def size(self):
+        return 2
+
+    def rank(self):
+        return 0
+
+
+def _make_loader(framework):
+    from fastsafetensors import SafeTensorsFileLoader
+
+    return SafeTensorsFileLoader(None, "cpu", nogds=True, framework="pytorch")
+
+
+def _tight_budget(path, framework):
+    """The smallest budget that still admits two copies of *path* under
+    broadcast: resident for both shards plus one in-flight chunk per unit of
+    pipeline depth. Anything larger stops forcing sub-file chunking."""
+    from fastsafetensors import SafeTensorsMetadata
+    from fastsafetensors._planner import collect_file_stats, pipeline_depth
+
+    meta = SafeTensorsMetadata.from_file(path, framework)
+    (st,) = collect_file_stats([(path, meta)])
+    # broadcast adds one depth unit for the in-flight receive tensor
+    return 2 * st.kept_bytes + (pipeline_depth(0) + 1) * st.largest_tensor
+
+
+def test_broadcast_explicit_budget_allowed(input_files, framework, tmp_path):
+    if framework.get_name() != "pytorch":
+        pytest.skip("pytorch-only")
+    import shutil
+
+    from fastsafetensors.parallel_loader import PipelineParallel
+
+    f2 = str(tmp_path / "copy.safetensors")
+    shutil.copy(input_files[0], f2)
+    pp = PipelineParallel(
+        _FakePG(),
+        _make_loader(framework),
+        [input_files[0], f2],
+        queue_size=0,
+        use_tqdm_on_load=False,
+        device_memory_budget=1 << 30,  # explicit int: deterministic across ranks
+    )
+    # one file per rank per group -> chunk-batch specs of width 2
+    assert pp.weight_files_batches
+    assert all(len(spec) == 2 for spec in pp.weight_files_batches)
+
+
+def test_broadcast_budget_chunks_in_lockstep(input_files, framework, tmp_path):
+    """Under broadcast, a budget tight enough to split shards must still give
+    every rank the same number of chunk-batches, each full width -- a ragged
+    sequence would leave some rank broadcasting while another has finished."""
+    if framework.get_name() != "pytorch":
+        pytest.skip("pytorch-only")
+    import shutil
+
+    from fastsafetensors.parallel_loader import PipelineParallel
+
+    f2 = str(tmp_path / "copy2.safetensors")
+    shutil.copy(input_files[0], f2)
+    pp = PipelineParallel(
+        _FakePG(),
+        _make_loader(framework),
+        [input_files[0], f2],
+        queue_size=0,
+        use_tqdm_on_load=False,
+        device_memory_budget=_tight_budget(input_files[0], framework),
+    )
+    specs = pp.weight_files_batches
+    # More batches than files => the budget actually forced sub-file chunking.
+    assert len(specs) > 1, specs
+    assert all(len(spec) == 2 for spec in specs), specs