diffct 2.3.0__py3-none-any.whl → 2.4.0__py3-none-any.whl

diffct/__init__.py

@@ -30,6 +30,8 @@ from .reconstruction import (
 from .solvers import SolverResult, cgls, sirt, tv_reconstruction, cgls_into, total_variation
 
 from .geometry import (
+    transform_trajectory,
+    circular_trajectory_3d_parallel,
     circular_trajectory_3d,
     random_trajectory_3d,
     spiral_trajectory_3d,
@@ -62,7 +64,7 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.3.0'
+__version__ = '2.4.0'
 
 __all__ = [
     'Projector',
@@ -88,6 +90,8 @@ __all__ = [
     'FanBackprojectorFunction',
     'ConeProjectorFunction',
     'ConeBackprojectorFunction',
+    'transform_trajectory',
+    'circular_trajectory_3d_parallel',
     'circular_trajectory_3d',
     'random_trajectory_3d',
     'spiral_trajectory_3d',

diffct/_solver_state.py

@@ -11,6 +11,7 @@ import torch.distributed as dist
 
 from .pipeline import _fragments
 from .storage import NpyStore, TensorStore
+from .distribution import control_status, preflight
 
 
 def count(value, name, positive=False):
@@ -25,47 +26,56 @@ def finite(value, name, positive=False):
     return float(value)
 
 
-def agreement(operator, error=None, token=None):
+def agreement(operator, error=None, token=None, *, cache_key=None):
     """Every rank rejects a bad request before numerical execution starts."""
     world = getattr(operator, "world_size", 1)
     if world == 1:
         if error is not None:
             raise error
         return
-    records = [None] * world
-    dist.all_gather_object(records, (None if error is None else str(error), token),
-                           group=operator.process_group)
-    failures = [record[0] for record in records if record[0] is not None]
-    if failures or any(record[1] != records[0][1] for record in records):
-        raise ValueError("distributed solver agreement failed: " + (failures[0] if failures else "inconsistent checkpoint iteration or controls"))
+    error = None if error is None else ValueError("distributed solver agreement failed: " + str(error))
+    mismatch = "distributed solver agreement failed: inconsistent checkpoint iteration or controls"
+    partition = getattr(operator, "_partition", None)
+    if partition is not None:
+        partition.preflight(error, token, cache_key=cache_key, mismatch=mismatch)
+    else:
+        preflight(error, token, process_group=operator.process_group, world_size=world,
+                  mismatch=mismatch)
 
 
-def scalar_sum(operator, value, kind):
-    """Count physical rays/voxels once, including view/space replicas."""
-    world = getattr(operator, "world_size", 1)
+def owns_values(operator, kind):
+    """Select one copy of each physical ray or voxel for scalar reductions."""
     partition = getattr(operator, "partition", "views")
     rank = getattr(operator, "rank", 0)
-    owns = (getattr(operator, "view_rank", 0) == 0 if kind == "volume"
+    return (getattr(operator, "view_rank", 0) == 0 if kind == "volume"
             else getattr(operator, "space_rank", 0) == 0) if partition == "hybrid" else (
         rank == 0 if (kind == "volume") == (partition == "views") else True
     )
-    value = float(value) if owns else 0.
+
+
+def scalar_sum(operator, value, kind):
+    """Count physical rays/voxels once, including view/space replicas."""
+    world = getattr(operator, "world_size", 1)
+    value = value if owns_values(operator, kind) else 0.
     if world > 1:
         group = operator.process_group
         device = torch.device("cuda", torch.cuda.current_device()) if dist.get_backend(group) == "nccl" else torch.device("cpu")
-        accumulator = torch.tensor(value, dtype=torch.float64, device=device)
+        accumulator = torch.as_tensor(value, dtype=torch.float64, device=device).clone()
         dist.all_reduce(accumulator, group=group)
-        value = accumulator.item()
-    return value
+        return accumulator.item()
+    return float(value)
 
 
 def squared_norm(operator, value, kind, elements=1048576):
     store = TensorStore(value) if isinstance(value, torch.Tensor) else value
     total = 0.
-    for selected in _fragments(tuple(store.shape), elements):
-        # Clone even an FP64 view: norms must never alter caller measurements.
-        block = store.read(selected).to(dtype=torch.float64).clone()
-        total += block.square_().sum().item()
+    if owns_values(operator, kind):
+        total = torch.zeros((), dtype=torch.float64, device=getattr(store, "device", "cpu"))
+        for selected in _fragments(tuple(store.shape), elements):
+            # Clone even an FP64 view: norms must never alter caller measurements.
+            block = store.read(selected).to(dtype=torch.float64).clone()
+            # Preserve the block reductions and their sequential FP64 order.
+            total.add_(block.square_().sum().to(device=total.device))
     return scalar_sum(operator, total, kind)
 
 
@@ -270,15 +280,22 @@ def save_checkpoint(operator, directory, states, metadata, iteration, history, m
     agreement(operator, error, iteration)
 
 
+def callbacks_enabled(operator, callback):
+    """Check callback presence once, including callbacks on only some ranks."""
+    return bool(control_status((int(callback is not None),),
+                               process_group=getattr(operator, "process_group", None),
+                               world_size=getattr(operator, "world_size", 1))[0])
+
+
 def callback_stop(operator, callback, iteration, residual, volume):
     error, stop = None, False
     try:
         stop = callback is not None and callback(iteration, residual, volume) is False
     except Exception as caught:
         error = caught
-    agreement(operator, error)
-    if getattr(operator, "world_size", 1) > 1:
-        records = [None] * operator.world_size
-        dist.all_gather_object(records, stop, group=operator.process_group)
-        stop = any(records)
-    return stop
+    failed, stop = control_status((int(error is not None), int(stop)),
+                                  process_group=getattr(operator, "process_group", None),
+                                  world_size=getattr(operator, "world_size", 1))
+    if failed:
+        agreement(operator, error)
+    return bool(stop)

diffct/backends.py

@@ -61,6 +61,37 @@ def _rectangles(shape, chunks):
                     for start, step, size in zip(starts, chunks, shape))
 
 
+def _flat_admission_batches(owner, geometry, detector_chunk_shape):
+    """Bound flat admission with corners and exact ambiguous-view fallback."""
+    views = owner.view_slice
+    if views.start == views.stop:
+        return
+    # Flat point component arithmetic is monotone in each detector
+    # coordinate; finite corner points and source differences bound the grid.
+    endpoints = tuple((slice(0, 1),) if size == 1 else (slice(0, 1), slice(size - 1, size))
+                      for size in owner.detector_shape)
+    for pixels in itertools.product(*endpoints):
+        yield views, pixels
+    if owner.beam in ("parallel", "parallel3d"):
+        return
+    source, lower = (value[views].to(dtype=owner.dtype) for value in geometry[:2])
+    upper = lower
+    for size, pitch, axis in zip(owner.detector_shape, owner.detector_spacing, geometry[2:]):
+        coordinates = (torch.tensor([.5 - size / 2, size - .5 - size / 2], dtype=torch.float64)
+                       * pitch).to(device=lower.device, dtype=owner.dtype)
+        products = coordinates[None, :, None] * axis[views, None, :].to(dtype=owner.dtype)
+        lower = lower + products.amin(dim=1)
+        upper = upper + products.amax(dim=1)
+    # Corners cannot exclude an interior pixel rounded onto the source.
+    # A staged source outside this AABB cannot coincide; inconclusive views
+    # retain the exact bounded pixel check rather than changing admission.
+    ambiguous = torch.all((source >= lower) & (source <= upper), dim=-1)
+    for index in torch.nonzero(ambiguous, as_tuple=False).flatten().tolist():
+        selected = slice(views.start + index, views.start + index + 1)
+        for pixels in _rectangles(owner.detector_shape, detector_chunk_shape):
+            yield selected, pixels
+
+
 class _DistributedApply(torch.autograd.Function):
     """Distributed linear map whose adjoint SUMs replicated cotangents."""
 
@@ -132,7 +163,7 @@ class ReferenceProjector:
     every voxel. All data operations remain ordinary differentiable PyTorch.
     """
 
-    def __init__(self, owner, *, validate=True):
+    def __init__(self, owner, *, validate=True, synchronize=True):
         self.owner = owner
         self.spacing = voxel_spacing_tuple(owner.voxel_spacing, len(owner.global_volume_shape))
         self.chunk_shape, self.view_chunk_size, self.detector_chunk_shape = self._work_chunks()
@@ -140,25 +171,25 @@ class ReferenceProjector:
             # Native owners already validated geometry before transform dispatch.
             return
         with torch.enable_grad():
-            geometry = self._snapshot_geometry()
+            geometry = self._snapshot_geometry(synchronize=synchronize)
         if (owner.detector_surface is not None
                 and not isinstance(owner.detector_surface, ParameterizedSurface)):
             trainable = any(value.requires_grad for value in geometry)
-            if owner.world_size > 1:
+            if synchronize and owner.world_size > 1:
                 device = (torch.device("cuda", torch.cuda.current_device())
                           if dist.get_backend(owner.process_group) == "nccl" else torch.device("cpu"))
                 flag = torch.tensor(int(trainable), device=device)
                 dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=owner.process_group)
                 trainable = bool(flag.item())
             owner._geometry_trainable = owner._geometry_trainable or trainable
-        for views, pixels in self._batches():
-            origins, directions, _ = self._rays(views, pixels, torch.device("cpu"), validate=True,
-                                              geometry=geometry)
-            if (not torch.isfinite(origins.to(dtype=owner.dtype)).all().item()
-                    or not torch.isfinite(directions.to(dtype=owner.dtype)).all().item()):
-                raise ValueError("detector surface world coordinates must be finite")
-            if owner.beam != "parallel" and torch.any(torch.all(directions == 0, dim=-1)).item():
-                raise ValueError("source and detector surface pixels must differ")
+        self._validate_geometry(geometry, geometry[0].device, admission=True,
+                                synchronize=synchronize)
+
+    def _admission_batches(self, geometry):
+        if self.owner.detector_surface is not None:
+            yield from self._batches()
+        else:
+            yield from _flat_admission_batches(self.owner, geometry, self.detector_chunk_shape)
 
     def _work_chunks(self, batch_elements=1):
         owner = self.owner
@@ -192,41 +223,96 @@ class ReferenceProjector:
             for pixels in _rectangles(owner.detector_shape, detector_chunk_shape or self.detector_chunk_shape):
                 yield views, pixels
 
-    def _snapshot_geometry(self):
+    def _snapshot_geometry(self, *, synchronize=True):
         owner = self.owner
         frames = owner._learnable or owner._trajectory
         surface = owner.detector_surface
-        if isinstance(surface, ParameterizedSurface):
-            return tuple(value.clone() for value in (*frames, *surface.parameters))
         if surface is None:
             return tuple(value.clone() for value in frames)
-        axes = [(torch.arange(size, device="cpu", dtype=torch.float64) + .5 - size / 2) * pitch
-                for size, pitch in zip(owner.detector_shape, owner.detector_spacing)]
-        u, v = (torch.meshgrid(*axes, indexing="ij") if len(axes) == 2
-                else (axes[0], torch.zeros_like(axes[0])))
-
-        def pack(value):
-            return value.detach().to(device="cpu").clone(), value.device
-
-        def unpack(saved):
-            value, device = saved
-            return value.to(device=device)
-
-        context = (nullcontext() if torch._C._are_functorch_transforms_active()
-                   else torch.autograd.graph.saved_tensors_hooks(pack, unpack))
-        with context:
-            offsets = surface(u, v)
-        if not isinstance(offsets, torch.Tensor) or not torch.is_floating_point(offsets):
-            raise TypeError("detector_surface must return a floating-point tensor")
-        if tuple(offsets.shape) not in ((*owner.detector_shape, 3),
-                                       (owner._n_views, *owner.detector_shape, 3)):
-            raise ValueError("detector_surface output shape must match the complete detector grids")
-        return (*tuple(value.clone() for value in frames), offsets.clone())
-
-    def _rays(self, views, pixels, device, validate=False, stats=None, geometry=None):
+        error, metadata, snapshots = None, None, ()
+        try:
+            if isinstance(surface, ParameterizedSurface):
+                # Parameters are sampled on CPU. Preserve their autograd state
+                # when a sampler differentiates a potential during admission.
+                with torch.enable_grad():
+                    parameters = tuple(value.to(device="cpu", copy=True) for value in surface.parameters)
+                snapshots = (*tuple(value.clone() for value in frames), *parameters)
+            else:
+                axes = [(torch.arange(size, device="cpu", dtype=torch.float64) + .5 - size / 2) * pitch
+                        for size, pitch in zip(owner.detector_shape, owner.detector_spacing)]
+                u, v = (torch.meshgrid(*axes, indexing="ij") if len(axes) == 2
+                        else (axes[0], torch.zeros_like(axes[0])))
+
+                def pack(value):
+                    return value.detach().to(device="cpu").clone(), value.device
+
+                def unpack(saved):
+                    value, device = saved
+                    return value.to(device=device)
+
+                context = (nullcontext() if torch._C._are_functorch_transforms_active()
+                           else torch.autograd.graph.saved_tensors_hooks(pack, unpack))
+                with context:
+                    offsets = surface(u, v)
+                if not isinstance(offsets, torch.Tensor) or not torch.is_floating_point(offsets):
+                    raise TypeError("detector_surface must return a floating-point tensor")
+                if tuple(offsets.shape) not in ((*owner.detector_shape, 3),
+                                               (owner._n_views, *owner.detector_shape, 3)):
+                    raise ValueError("detector_surface output shape must match the complete detector grids")
+                metadata = (tuple(offsets.shape), str(offsets.dtype))
+                snapshots = (*tuple(value.clone() for value in frames), offsets.to(device="cpu", copy=True))
+        except (TypeError, ValueError, RuntimeError) as failure:
+            error = failure
+        if synchronize:
+            owner._partition.preflight(error, metadata, cache_key="portable surface snapshot",
+                                       revision=owner._agreement_revision())
+        elif error is not None:
+            raise error
+        return snapshots
+
+    def _validate_geometry(self, geometry, fallback, *, admission=False, synchronize=True):
+        """Finish every local geometry check before any rank starts numerics."""
+        owner = self.owner
+        if not admission and (owner.detector_surface is None or
+                              (owner.world_size == 1 and isinstance(owner.detector_surface, ParameterizedSurface))):
+            return
+        devices = tuple(torch.device(device) for device in (owner._devices or (fallback,)))
+        error = None
+        try:
+            frame_count = 4 if owner.beam in ("cone", "parallel3d") else 3
+            parameters = geometry[frame_count:]
+            if isinstance(owner.detector_surface, ParameterizedSurface):
+                with torch.enable_grad():
+                    parameters = tuple(value.to(device="cpu") for value in parameters)
+            with torch.no_grad():
+                # Reuse frame transfers across detector chunks. Offset batches
+                # remain bounded and move only when a ray batch needs them.
+                staged = {device: (*tuple(value.to(device=device) for value in geometry[:frame_count]),
+                                    *parameters) for device in set(devices)}
+                batches = self._admission_batches(staged[devices[0]]) if admission else self._batches()
+                for batch, (views, pixels) in enumerate(batches):
+                    device = devices[batch % len(devices)]
+                    origins, directions, _ = self._rays(
+                        views, pixels, device, validate=True, geometry=staged[device],
+                        admission=admission,
+                    )
+                    if (not torch.isfinite(origins.to(dtype=owner.dtype)).all().item()
+                            or not torch.isfinite(directions.to(dtype=owner.dtype)).all().item()):
+                        raise ValueError("detector surface world coordinates must be finite")
+                    if (owner.beam not in ("parallel", "parallel3d")
+                            and torch.any(torch.all(directions == 0, dim=-1)).item()):
+                        raise ValueError("source and detector surface pixels must differ")
+        except (TypeError, ValueError, RuntimeError) as failure:
+            error = failure
+        if synchronize:
+            owner._partition.preflight(error)
+        elif error is not None:
+            raise error
+
+    def _rays(self, views, pixels, device, validate=False, stats=None, geometry=None, *, admission=True):
         owner = self.owner
         # Admission compares original coordinates before FP32 rounding.
-        dtype = torch.float64 if validate else owner.dtype
+        dtype = torch.float64 if validate and admission else owner.dtype
         axes = [(torch.arange(pixel.start, pixel.stop, dtype=torch.float64, device="cpu")
                  + .5 - size / 2) * pitch
                 for pixel, size, pitch in zip(pixels, owner.detector_shape, owner.detector_spacing)]
@@ -234,7 +320,7 @@ class ReferenceProjector:
                 else (axes[0], torch.zeros_like(axes[0])))
         ids = torch.arange(views.start, views.stop, device="cpu")
         surface = owner.detector_surface
-        frame_count = 4 if owner.beam == "cone" else 3
+        frame_count = 4 if owner.beam in ("cone", "parallel3d") else 3
         if isinstance(surface, ParameterizedSurface):
             parameters = tuple(self._transfer(value, torch.device("cpu"), stats)
                                for value in (surface.parameters if geometry is None else geometry[frame_count:]))
@@ -247,16 +333,16 @@ class ReferenceProjector:
             offsets = torch.stack((u, v, torch.zeros_like(u)), dim=-1)
         if not isinstance(offsets, torch.Tensor) or not torch.is_floating_point(offsets):
             raise TypeError("detector_surface must return a floating-point tensor")
-        if validate:
-            if not torch.isfinite(offsets).all().item():
-                raise ValueError("detector_surface offsets must be finite")
-            if owner.beam != "cone" and torch.any(offsets[..., 1] != 0).item():
-                raise ValueError("2D detector_surface middle offsets must be zero")
         shared_shape = (*u.shape, 3)
         if tuple(offsets.shape) == (owner._n_views, *shared_shape):
             offsets = offsets[views]
         elif tuple(offsets.shape) not in (shared_shape, (len(ids), *shared_shape)):
             raise ValueError("detector_surface output shape must match the selected view/pixel batch")
+        if validate:
+            if not torch.isfinite(offsets).all().item():
+                raise ValueError("detector_surface offsets must be finite")
+            if owner.beam not in ("cone", "parallel3d") and torch.any(offsets[..., 1] != 0).item():
+                raise ValueError("2D detector_surface middle offsets must be zero")
         frames = (owner._learnable or owner._trajectory) if geometry is None else geometry[:frame_count]
         native_surface = owner.backend == "numba" and surface is not None
         native_legacy = native_surface and not isinstance(surface, ParameterizedSurface)
@@ -279,7 +365,7 @@ class ReferenceProjector:
         def world_points(offsets, components):
             axis_u = frame(components[2])
             points = frame(components[1]) + offsets[..., 0, None] * axis_u
-            if owner.beam == "cone":
+            if owner.beam in ("cone", "parallel3d"):
                 axis_v = frame(components[3])
                 points = points + offsets[..., 1, None] * axis_v
                 normal_dtype = torch.promote_types(axis_u.dtype, axis_v.dtype)
@@ -290,7 +376,7 @@ class ReferenceProjector:
             return points + offsets[..., 2, None] * normal
 
         points = world_points(local, selected)
-        if (validate and owner.beam != "parallel" and owner.dtype == torch.float32
+        if (validate and owner.beam not in ("parallel", "parallel3d") and owner.dtype == torch.float32
                 and owner.detector_surface is not None):
             # Raw FP64 coordinates preserve the exact admission boundaries.
             sources = frame(selected[0]).expand_as(points)
@@ -317,7 +403,7 @@ class ReferenceProjector:
             # FP32 component arithmetic. Reuse the same sampled offsets.
             selected = tuple(value.to(dtype=owner.dtype) for value in selected)
             points = world_points(local.to(dtype=owner.dtype), selected)
-        if owner.beam == "parallel":
+        if owner.beam in ("parallel", "parallel3d"):
             origins, directions = points, frame(selected[0]).expand_as(points)
         else:
             sources = frame(selected[0]).expand_as(points)
@@ -368,7 +454,7 @@ class ReferenceProjector:
         first, last = (lower - origins) / safe_direction, (upper - origins) / safe_direction
         enter = torch.where(constant, -torch.inf, torch.minimum(first, last)).amax(dim=-1)
         leave = torch.where(constant, torch.inf, torch.maximum(first, last)).amin(dim=-1)
-        if owner.beam != "parallel":
+        if owner.beam not in ("parallel", "parallel3d"):
             enter, leave = enter.clamp_min(0), torch.minimum(leave, ray_length)
         inside = (~constant | ((origins >= lower) & (origins < upper))).all(dim=-1)
         valid = inside & (leave > enter)
@@ -475,7 +561,9 @@ class ReferenceProjector:
             error = ValueError(f"input trailing shape must be {shape}")
         else:
             leading = tuple(tensor.shape[:-len(shape)])
-        owner._partition.preflight(error, (is_project, reduce_cotangent, leading))
+        owner._partition.preflight(error, (is_project, reduce_cotangent, leading),
+                                   cache_key="portable input",
+                                   revision=owner._agreement_revision())
         batch_elements = math.prod(leading)
         plan = self._work_chunks(batch_elements) if _plan is None else _plan
         stats = self._stats(is_project, plan)
@@ -569,6 +657,7 @@ class ReferenceProjector:
 
     def _apply(self, tensor, is_project):
         geometry = self._snapshot_geometry()
+        self._validate_geometry(geometry, tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu"))
         if self.owner.world_size > 1 and torch.is_grad_enabled():
             return _DistributedApply.apply(tensor, self, is_project, False, self.owner._participation(), *geometry)
         return self._tensor_apply(tensor, is_project, geometry=geometry)
@@ -622,6 +711,7 @@ class ReferenceProjector:
         fallback = torch.device(getattr(source, "device", "cpu"))
         output_device = torch.device(getattr(sink, "device", "cpu"))
         geometry = self._snapshot_geometry()
+        self._validate_geometry(geometry, fallback)
 
         def read(slices, device):
             value = source.read(slices)

diffct/chunking.py

@@ -63,6 +63,40 @@ def _memory_budget(device):
     return budget, allocation_limit
 
 
+def _memory_lower_bounds(device):
+    """Lower bounds of `_memory_budget` and `_reservation_budget` without an allocator snapshot.
+
+    Returns (budget, allocation_limit, reservation budget), each at most the value the two exact
+    functions return at the same allocator state. The snapshot is what makes those functions
+    slow; it only serves two corrections, and both are bounded by allocator counters: free splits
+    pinned by other streams are inactive reserved bytes, at most reserved - allocated, and pinned
+    segments hold at most every reserved byte.
+    """
+    free, total = torch.cuda.mem_get_info(device)
+    # One nested statistics query; memory_allocated/memory_reserved each flatten every statistic.
+    stats = torch.cuda.memory_stats_as_nested_dict(device)
+    if stats:
+        allocated = stats["allocated_bytes"]["all"]["current"]
+        reserved = stats["reserved_bytes"]["all"]["current"]
+    else:                                   # allocator not initialised: nothing allocated or reserved
+        allocated = reserved = 0
+    reusable = max(0, reserved - allocated)
+    getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
+    fraction = getter(device) if getter is not None else 1.0
+    allocator_room = max(0, int(total * fraction) - allocated)
+    backend = getattr(torch.cuda, "get_allocator_backend", lambda: "native")()
+    foreign_splits = reusable if backend == "native" else 0
+    available = max(0, min(free + reusable, allocator_room) - foreign_splits)
+    allocation_limit = (1 << 63) - 1
+    if backend == "native" and available < 20 * 1024 ** 2:
+        allocation_limit = 1024 ** 2
+    if backend != "native":
+        reservation = (1 << 63) - 1
+    else:
+        reservation = max(0, min(free, int(total * fraction) - reserved))
+    return available * 3 // 4, allocation_limit, reservation
+
+
 def _native_segments(device):
     index = torch.device(device).index if not isinstance(device, int) else device
     index = torch.cuda.current_device() if index is None else index

diffct/distribution.py

@@ -1,9 +1,12 @@
 """Shared deterministic ownership for view and first-axis space partitions."""
 
+import hashlib
 import math
 import numbers
+from copy import deepcopy
 from weakref import WeakKeyDictionary
 
+import torch
 import torch.distributed as dist
 
 
@@ -52,20 +55,73 @@ def balanced_slice(size, rank, world_size):
     return slice(start, start + base + int(rank < extra))
 
 
-def preflight(error=None, metadata=None, *, process_group=None, world_size=1):
-    """Reject invalid/disagreeing requests collectively before numerical I/O."""
+def control_status(values, *, process_group=None, world_size=1):
+    """Combine fixed-size control values without serializing Python objects."""
+    if world_size == 1:
+        return values
+    with torch._C._DisableFuncTorch():
+        device = (torch.device("cuda", torch.cuda.current_device())
+                  if dist.get_backend(process_group) == "nccl" else torch.device("cpu"))
+        status = torch.tensor(values, dtype=torch.int64, device=device)
+        dist.all_reduce(status, op=dist.ReduceOp.MAX, group=process_group)
+        return status.tolist()
+
+
+def _canonical_metadata(value):
+    """Keep container types while making dictionary order deterministic."""
+    if isinstance(value, dict):
+        return ("dict", tuple(sorted(
+            ((_canonical_metadata(key), _canonical_metadata(item)) for key, item in value.items()),
+            key=repr,
+        )))
+    if isinstance(value, (tuple, list)):
+        return ("tuple" if isinstance(value, tuple) else "list",
+                tuple(_canonical_metadata(item) for item in value))
+    return value
+
+
+def preflight(error=None, metadata=None, *, process_group=None, world_size=1,
+              cache=None, cache_key=None, revision=None,
+              mismatch="distributed ranks must agree on request metadata"):
+    """Reject invalid/disagreeing requests collectively before numerical I/O.
+
+    Cache only local digest computation. Every call compares current metadata
+    across ranks, including calls through different warmed Partition objects.
+    """
     result = (None if error is None else (type(error).__name__, str(error)), metadata)
     results = [result]
     if world_size > 1:
-        results = [None] * world_size
-        dist.all_gather_object(results, result, group=process_group)
+        # Control messages do not participate in data transforms or autograd.
+        with torch._C._DisableFuncTorch():
+            try:
+                changed = (cache is None or cache_key not in cache
+                           or cache[cache_key][:2] != (revision, metadata))
+                digest = (int.from_bytes(hashlib.blake2b(
+                    repr(_canonical_metadata(metadata)).encode("utf-8"), digest_size=8,
+                ).digest(), "big", signed=True) if changed else cache[cache_key][2])
+                exchange = error is not None
+            except Exception:
+                # Metadata that cannot be digested here (e.g. self-referential) must not leave
+                # this rank outside the collective; the object exchange below reports it.
+                changed, digest, exchange = False, 0, True
+            # MAX(digest) and MAX(~digest) encode both signed int64 bounds.
+            # Complementing also handles -2**63 without negation overflow.
+            failed, upper, inverse_lower = control_status(
+                (int(exchange), digest, ~digest),
+                process_group=process_group, world_size=world_size,
+            )
+            if failed or upper != ~inverse_lower:
+                results = [None] * world_size
+                dist.all_gather_object(results, result, group=process_group)
     failure = next((entry[0] for entry in results if entry[0] is not None), None)
     if failure is not None:
         classes = {kind.__name__: kind for kind in (TypeError, ValueError, RuntimeError,
                                                    PermissionError, FileNotFoundError)}
         raise classes.get(failure[0], RuntimeError)(failure[1])
     if any(entry[1] != metadata for entry in results):
-        raise ValueError("distributed ranks must agree on request metadata")
+        raise ValueError(mismatch)
+    if cache is not None and world_size > 1 and changed:
+        cache[cache_key] = (revision, deepcopy(metadata), digest)
 
 
 class Partition:
@@ -82,6 +138,7 @@ class Partition:
         if process_group is not None and not distributed:
             raise ValueError("process_group requires distributed=True")
         self.process_group = process_group
+        self._preflight_cache = {}
         self.rank, self.world_size = 0, 1
         if distributed:
             if not dist.is_available() or not dist.is_initialized():
@@ -132,6 +189,9 @@ class Partition:
     def global_rank(self, rank):
         return self.ranks[rank]
 
-    def preflight(self, error=None, metadata=None):
+    def preflight(self, error=None, metadata=None, *, cache_key=None, revision=None,
+                  mismatch="distributed ranks must agree on request metadata"):
         """Reject invalid/disagreeing requests on every parent rank before I/O."""
-        return preflight(error, metadata, process_group=self.process_group, world_size=self.world_size)
+        return preflight(error, metadata, process_group=self.process_group, world_size=self.world_size,
+                         cache=self._preflight_cache if cache_key is not None else None,
+                         cache_key=cache_key, revision=revision, mismatch=mismatch)