fix(training): avoid duplicate packed mask transfers

CHANGELOG.md

@@ -62,6 +62,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Fixed
 
+- Packed targets are now materialized directly into independent per-sample tensors on the destination device instead of first moving every concatenated field and then cloning it. This removes a transient CUDA allocation equal to the mask field's size, present whenever segmentation transfer held both the packed and the materialized masks at once; DataLoader worker transport remains packed. This is a mechanism-level fix, not the representative detection/segmentation host-device memory benchmark suite noted as a follow-up boundary above.
+
 - Empty COCO targets now keep `iscrowd` as `int64` and `area` as `float32`, matching populated targets and allowing mixed empty/populated batches to use lossless packed-target worker transport.
 
 ---

docs/learn/train/training-parameters.md

@@ -301,7 +301,7 @@ The parameters below are available for fine-grained control over training behavi
 | `prefetch_factor`    | `int`  | `None`  | Number of batches to prefetch per DataLoader worker. `None` uses PyTorch's built-in default.                             |
 | `pack_targets`       | `bool` | `True`  | Concatenate target dicts before crossing the DataLoader worker boundary. See the contract below; set `False` to opt out. |
 
-With `pack_targets=True`, train, validation, test, and predict loaders yield batches whose target element is `PackedTargets` whenever packing is lossless. The Lightning `transfer_batch_to_device` hook accepts those batches or an unpacked tuple of target dicts. It moves packed fields to the target device, then materializes them into the same plain per-sample dict list that training, validation, test, and prediction hooks receive on the unpacked path. Batches that cannot be packed losslessly retain their original tuple of dicts.
+With `pack_targets=True`, train, validation, test, and predict loaders yield batches whose target element is `PackedTargets` whenever packing is lossless. The Lightning `transfer_batch_to_device` hook accepts those batches or an unpacked tuple of target dicts. It materializes each packed field directly into its own independently owned per-sample tensor on the target device, producing the same plain per-sample dict list that training, validation, test, and prediction hooks receive on the unpacked path. Batches that cannot be packed losslessly retain their original tuple of dicts.
 
 ## Complete Parameter Reference
 

src/rfdetr/training/module_data.py

@@ -814,9 +814,9 @@ def transfer_batch_to_device(
         non_blocking = device.type == "cuda"
         samples = samples.to(device, non_blocking=non_blocking)
         if isinstance(targets, PackedTargets):
-            # One transfer per field instead of one per field per sample, then materialised so that callers can
-            # mutate the dicts exactly as they do on the unpacked path.
-            targets = targets.to(device, non_blocking=non_blocking).as_list()
+            # Materialise directly on the destination so the packed and per-sample copies of a large field such as
+            # segmentation masks never coexist there. Each returned tensor owns its storage, matching the unpacked path.
+            targets = targets.to_list(device, non_blocking=non_blocking)
         else:
             targets = [{k: v.to(device, non_blocking=non_blocking) for k, v in t.items()} for t in targets]
         return samples, targets

src/rfdetr/utilities/tensors.py

@@ -389,8 +389,9 @@ class PackedTargets:
     Sequence subclass would be taken apart and pinned tensor by tensor, rebuilding in the main process exactly
     what the packing avoided. Indexing rebuilds a dict of views: reassigning a key on the returned dict does not
     write back into the packed storage, but an in-place write into one of its tensors does, since the view shares
-    the packed field's storage. Call :meth:`as_list` for a materialised copy that is safe to mutate in place, the
-    way the unpacked batch already is.
+    the packed field's storage. Call :meth:`as_list` for a same-device materialised copy, or :meth:`to_list` to
+    materialise directly on another device; both return independent tensors that are safe to mutate in place, the way
+    the unpacked batch already is.
 
     Args:
         values: One flat tensor per field, holding that field concatenated across the batch.
@@ -450,8 +451,8 @@ def pin_memory(self) -> PackedTargets:
         """Pin every packed field, keeping the batch packed.
 
         Called by PyTorch's pin-memory worker. Pinning the concatenated fields costs one pinned allocation
-        per field rather than one per field per sample, and leaves the batch in the packed form that
-        ``transfer_batch_to_device`` moves in a single copy per field.
+        per field rather than one per field per sample, and leaves the batch packed until
+        ``transfer_batch_to_device`` materialises its independent destination tensors.
 
         Returns:
             A packed batch whose fields are in pinned memory.
@@ -483,6 +484,35 @@ def as_list(self) -> list[dict[str, Tensor]]:
         """
         return [{key: value.clone() for key, value in self[position].items()} for position in range(self._batch)]
 
+    def to_list(self, device: torch.device | str, non_blocking: bool = False) -> list[dict[str, Tensor]]:
+        """Materialise independent per-sample tensors directly on *device*.
+
+        Moving the packed object and then calling :meth:`as_list` makes the concatenated destination fields coexist
+        with all of their per-sample clones. That transient duplicate is negligible for boxes and labels but can be
+        large for segmentation masks. Transferring each packed view directly into its final independently-owned tensor
+        avoids the full-field destination copy while preserving the unpacked path's mutation semantics.
+
+        Args:
+            device: Destination device.
+            non_blocking: Whether to request asynchronous copies.
+
+        Returns:
+            One dict per sample, ordered as packed, with independent tensors on *device*.
+        """
+        materialised: list[dict[str, Tensor]] = [{} for _ in range(self._batch)]
+        for key, flat in self._values.items():
+            offset = 0
+            for target, shape in zip(materialised, self._shapes[key], strict=True):
+                elements = _numel(shape)
+                value = flat[offset : offset + elements].reshape(shape)
+                target[key] = value.to(
+                    device=device,
+                    non_blocking=non_blocking,
+                    copy=True,
+                )
+                offset += elements
+        return materialised
+
     def to(self, device: torch.device | str, non_blocking: bool = False) -> PackedTargets:
         """Move every packed field to *device* in one transfer per field.
 

tests/training/test_module_data.py

@@ -1325,6 +1325,24 @@ def test_packed_targets_are_unpacked_to_a_plain_list_on_target_device(self, fixt
             for key, value in original.items():
                 assert torch.equal(rebuilt[key], value)
 
+    def test_packed_targets_are_materialised_without_a_whole_batch_device_copy(self, fixture_training_setup) -> None:
+        """Packed transfer must not create a device copy of every field before constructing per-sample tensors."""
+        _, _, dm = fixture_training_setup
+        samples, plain_targets = _make_batch()
+        packed_targets = pack_targets(plain_targets)
+        assert isinstance(packed_targets, PackedTargets)
+
+        with patch.object(
+            PackedTargets,
+            "to",
+            side_effect=AssertionError("whole packed batch copied to the target device"),
+        ):
+            _, result_targets = dm.transfer_batch_to_device(
+                (samples, packed_targets), torch.device("cpu"), dataloader_idx=0
+            )
+
+        assert isinstance(result_targets, list)
+
 
 # ---------------------------------------------------------------------------
 # TestBackendResolution — validates augmentation_backend logic in setup("fit")

tests/utilities/test_tensors.py

@@ -868,6 +868,90 @@ def test_as_list_tensors_do_not_alias_the_packed_storage(self) -> None:
         assert packed.fields["boxes"][0].item() == 1.0
         assert torch.equal(packed.as_list()[0]["labels"], torch.tensor([3, 7]))
 
+    def test_to_list_same_device_does_not_alias_packed_storage(self) -> None:
+        """Direct materialisation must preserve the unpacked path's independent tensor ownership on CPU."""
+        packed = pack_targets(self._batch())
+        assert isinstance(packed, PackedTargets)
+
+        materialised = packed.to_list(torch.device("cpu"))
+        materialised[0]["labels"][0] = 999
+        materialised[0]["boxes"][0, 0] = 999.0
+
+        assert torch.equal(packed.fields["labels"], torch.tensor([3, 7, 5]))
+        assert packed.fields["boxes"][0].item() == 1.0
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_to_list_transfers_pinned_views_directly_to_cuda(self) -> None:
+        """Direct materialisation must preserve values and ownership across a non-blocking CUDA transfer."""
+        packed = pack_targets(self._batch())
+        assert isinstance(packed, PackedTargets)
+        pinned = packed.pin_memory()
+
+        materialised = pinned.to_list(torch.device("cuda"), non_blocking=True)
+        torch.cuda.synchronize()
+
+        assert materialised[0]["labels"].device.type == "cuda"
+        assert torch.equal(materialised[0]["labels"].cpu(), torch.tensor([3, 7]))
+        materialised[0]["labels"][0] = 999
+        assert torch.equal(pinned.fields["labels"], torch.tensor([3, 7, 5]))
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_to_list_bounds_the_transient_cuda_peak_for_a_mask_field(self) -> None:
+        """``to_list``'s docstring claims the avoided duplicate "can be large for segmentation masks" -- pin that claim
+        on the actual CUDA peak with a mask field, not just on values or ownership."""
+        pinned = pack_targets(
+            [
+                {"labels": torch.tensor([3, 7]), "masks": torch.ones((2, 256, 256), dtype=torch.bool)},
+                {"labels": torch.tensor([5]), "masks": torch.ones((1, 256, 256), dtype=torch.bool)},
+            ]
+        ).pin_memory()
+        mask_bytes = pinned.fields["masks"].numel()
+
+        def peak_extra(materialise):
+            torch.cuda.synchronize()
+            torch.cuda.reset_peak_memory_stats()
+            out = materialise()
+            torch.cuda.synchronize()
+            extra = torch.cuda.max_memory_allocated() - torch.cuda.memory_allocated()
+            del out
+            return extra
+
+        old_extra = peak_extra(lambda: pinned.to(torch.device("cuda"), non_blocking=True).as_list())
+        new_extra = peak_extra(lambda: pinned.to_list(torch.device("cuda"), non_blocking=True))
+
+        assert old_extra >= mask_bytes
+        assert new_extra < mask_bytes // 2
+
+    def test_to_returns_self_when_already_on_the_target_device(self) -> None:
+        """``to()`` keeps the batch packed instead of materialising it, unlike ``to_list()``.
+
+        Losing its only caller in ``transfer_batch_to_device`` (replaced by ``to_list()``) must not leave it untested: a
+        no-op device request has to return the same instance, matching every no-op ``Tensor.to()`` call underneath it.
+        """
+        packed = pack_targets(self._batch())
+        assert isinstance(packed, PackedTargets)
+
+        same_device = packed.to(torch.device("cpu"))
+
+        assert same_device is packed
+        assert torch.equal(same_device.fields["labels"], torch.tensor([3, 7, 5]))
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_to_moves_every_field_and_keeps_the_batch_packed_on_cuda(self) -> None:
+        """A genuine device change must move every field and return a new packed batch, not a materialised list."""
+        packed = pack_targets(self._batch())
+        assert isinstance(packed, PackedTargets)
+
+        moved = packed.to(torch.device("cuda"))
+
+        assert isinstance(moved, PackedTargets)
+        assert moved is not packed
+        assert all(tensor.device.type == "cuda" for tensor in moved.fields.values())
+        assert torch.equal(moved.fields["labels"].cpu(), torch.tensor([3, 7, 5]))
+
     def test_collate_packs_only_when_asked(self) -> None:
         """The collate contract only changes for callers that opt in."""
         images = [torch.zeros(3, 8, 8), torch.zeros(3, 8, 8)]