perf(training): pack per-sample targets across the DataLoader boundary (#1399)

* perf(training): pack per-sample targets across the DataLoader boundary
* fix(data): harden packed-target worker handoff

---------

Co-authored-by: Jesús Royeth <JESUSROYETH@users.noreply.github.com>
Co-authored-by: jirka <6035284+borda@users.noreply.github.com>
Co-authored-by: Codex <codex@openai.com>

Author: 3 people

CHANGELOG.md

@@ -8,6 +8,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Added
 
+- Added `TrainConfig.pack_targets` (default `True`), which concatenates each batch's per-sample target dicts into one tensor per field before they cross the DataLoader worker-to-main boundary, and rebuilds them in `transfer_batch_to_device`. Every tensor a worker returns is moved into its own shared-memory segment and passed to the parent as a file descriptor, so a batch of 16 crosses as 114 objects of which 112 carry a few kilobytes in total; packing takes that to 9 without changing a byte of payload. The rebuilt targets are bit-identical, dtypes included, and the training step receives the same plain list of dicts as before. The maintainer selected the default after submitted detection-loader measurements favored packing. All `RFDETRDataModule` loaders yield batches whose targets are `PackedTargets` rather than a tuple of dicts whenever they pack losslessly, which is visible to direct consumers; a batch that cannot pack losslessly falls back to the original tuple of dicts. Representative detection and segmentation host/device-memory and transfer benchmarks have not yet been collected, so that default's full memory envelope remains a follow-up boundary.
+
 - Added `TrainConfig.eval_batch_size`, decoupling the validation, test and predict dataloaders from the training micro-batch size. The three eval loaders previously always reused the resolved `batch_size`, so lowering it to fit an optimizer step also shrank evaluation batches. Evaluation runs under `no_grad`, which avoids autograd activation storage, but in-fit validation still shares device memory with the model and optimizer state and needs memory for its own forward outputs. The default `None` inherits `batch_size` exactly as before. Unlike `batch_size` it accepts no `"auto"`: an explicit `eval_batch_size` is never probed and stays usable on the `batch_size="auto"` path, while leaving it unset keeps the existing "auto was never resolved" error for eval loaders too. The training dataloader, including its `grad_accum_steps` alignment padding, is unaffected.
 
 - `deploy_to_roboflow()`'s `version` argument is now optional: when omitted, the highest existing dataset version of the target project is resolved automatically via the Roboflow API (falling back to version `1` for a project with no generated versions, where the Roboflow SDK then raises its usual "Version number 1 is not found."). Passing an explicit `version` behaves exactly as before, with no extra API call. ([#1116](https://github.com/roboflow/rf-detr/issues/1116))

docs/learn/train/customization.md

@@ -57,7 +57,7 @@ module = RFDETRModelModule(model_config, train_config)
 | -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | `on_fit_start`             | Seeds RNGs when `train_config.seed` is set.                                                                                                                                                                                     |
 | `on_train_batch_start`     | Applies multi-scale random resize when `train_config.multi_scale=True`.                                                                                                                                                         |
-| `transfer_batch_to_device` | Moves `NestedTensor` batches to the target device.                                                                                                                                                                              |
+| `transfer_batch_to_device` | Moves `NestedTensor` batches and targets to the target device. With the default `TrainConfig.pack_targets=True`, it materializes `PackedTargets` into a plain per-sample dict list.                                             |
 | `training_step`            | Computes loss and logs `train/loss` plus per-term losses. Keypoint models use manual optimization with box-normalized accumulation across microbatches; detection and segmentation use Lightning's automatic optimization path. |
 | `validation_step`          | Runs forward pass and postprocessing; returns `{results, targets}` for `COCOEvalCallback`.                                                                                                                                      |
 | `test_step`                | Same as `validation_step`, logs under `test/`.                                                                                                                                                                                  |

docs/learn/train/training-parameters.md

@@ -294,11 +294,14 @@ The parameters below are available for fine-grained control over training behavi
 
 ### DataLoader Tuning
 
-| Parameter            | Type   | Default | Description                                                                                               |
-| -------------------- | ------ | ------- | --------------------------------------------------------------------------------------------------------- |
-| `pin_memory`         | `bool` | `None`  | Pin host memory in the DataLoader for faster GPU transfers. `None` defers to PyTorch Lightning's default. |
-| `persistent_workers` | `bool` | `None`  | Keep DataLoader worker processes alive between epochs. `None` defers to PyTorch Lightning's default.      |
-| `prefetch_factor`    | `int`  | `None`  | Number of batches to prefetch per DataLoader worker. `None` uses PyTorch's built-in default.              |
+| Parameter            | Type   | Default | Description                                                                                                              |
+| -------------------- | ------ | ------- | ------------------------------------------------------------------------------------------------------------------------ |
+| `pin_memory`         | `bool` | `None`  | Pin host memory in the DataLoader for faster GPU transfers. `None` defers to PyTorch Lightning's default.                |
+| `persistent_workers` | `bool` | `None`  | Keep DataLoader worker processes alive between epochs. `None` defers to PyTorch Lightning's default.                     |
+| `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.
 
 ## Complete Parameter Reference
 
@@ -354,3 +357,4 @@ Below is a summary table of all training parameters:
 | `pin_memory`                 | bool                   | None           | Pin DataLoader memory. None defers to PyTorch Lightning's default.                                                                                                                                                        |
 | `persistent_workers`         | bool                   | None           | Keep DataLoader workers alive between epochs. None uses PTL default.                                                                                                                                                      |
 | `prefetch_factor`            | int                    | None           | Number of batches prefetched per worker. None uses PyTorch default.                                                                                                                                                       |
+| `pack_targets`               | bool                   | True           | Concatenate target dicts before crossing the DataLoader worker boundary. See DataLoader Tuning; set False to opt out.                                                                                                     |

src/rfdetr/config.py

@@ -1267,6 +1267,7 @@ def _coerce_amp_dtype(cls, value: Any) -> Any:
     pin_memory: bool | None = None
     persistent_workers: bool | None = None
     prefetch_factor: int | None = None
+    pack_targets: bool = True
 
     @field_validator("batch_size", mode="after")
     @classmethod

src/rfdetr/training/module_data.py

@@ -22,7 +22,7 @@
 from rfdetr.datasets.yolo import YoloSplitUnavailableError
 from rfdetr.utilities.box_ops import box_xyxy_to_cxcywh
 from rfdetr.utilities.logger import get_logger
-from rfdetr.utilities.tensors import make_collate_fn
+from rfdetr.utilities.tensors import PackedTargets, make_collate_fn
 
 logger = get_logger()
 
@@ -153,6 +153,7 @@ def __init__(self, model_config: ModelConfig, train_config: TrainConfig) -> None
             )
         self._collate_fn = make_collate_fn(
             block_size=block_size,
+            pack=train_config.pack_targets,
         )
 
         self._dataset_train: torch.utils.data.Dataset[Any] | None = None
@@ -362,7 +363,8 @@ def train_dataloader(self) -> DataLoader[Any]:
         so that PTL can auto-inject ``DistributedSampler`` in DDP mode.
 
         Returns:
-            DataLoader for the training dataset.
+            DataLoader for the training dataset. With ``TrainConfig.pack_targets=True`` (the default), its collated
+            batches contain ``PackedTargets`` for losslessly packable target batches.
         """
         dataset: torch.utils.data.Dataset[Any] = self._require_dataset(self._dataset_train, "fit")
         batch_size = self._resolve_batch_size()
@@ -417,7 +419,8 @@ def val_dataloader(self) -> DataLoader[Any]:
         """Return the validation DataLoader.
 
         Returns:
-            DataLoader for the validation dataset with sequential sampling.
+            DataLoader for the validation dataset with sequential sampling. With ``TrainConfig.pack_targets=True``
+            (the default), its collated batches contain ``PackedTargets`` for losslessly packable target batches.
         """
         dataset = self._require_dataset(self._dataset_val, "validate")
         return DataLoader(
@@ -437,7 +440,8 @@ def test_dataloader(self) -> DataLoader[Any]:
         """Return the test DataLoader.
 
         Returns:
-            DataLoader for the test dataset with sequential sampling.
+            DataLoader for the test dataset with sequential sampling. With ``TrainConfig.pack_targets=True`` (the
+            default), its collated batches contain ``PackedTargets`` for losslessly packable target batches.
         """
         dataset = self._require_dataset(self._dataset_test, "test")
         return DataLoader(
@@ -457,7 +461,8 @@ def predict_dataloader(self) -> DataLoader[Any]:
         """Return the predict DataLoader (reuses the validation dataset, no augmentation).
 
         Returns:
-            DataLoader for the validation dataset with sequential sampling.
+            DataLoader for the validation dataset with sequential sampling. With ``TrainConfig.pack_targets=True``
+            (the default), its collated batches contain ``PackedTargets`` for losslessly packable target batches.
         """
         dataset = self._require_dataset(self._dataset_val, "predict")
         return DataLoader(
@@ -796,7 +801,9 @@ def transfer_batch_to_device(
         ``NestedTensor`` must be moved explicitly.
 
         Args:
-            batch: Tuple of (NestedTensor samples, list of target dicts).
+            batch: Tuple of ``NestedTensor`` samples and targets that are either a collated ``PackedTargets`` object
+                or an unpacked tuple of target dicts. Packed targets are materialized to a plain list of dicts after
+                transfer.
             device: Target device.
             dataloader_idx: Index of the dataloader providing this batch.
 
@@ -806,5 +813,10 @@ def transfer_batch_to_device(
         samples, targets = batch
         non_blocking = device.type == "cuda"
         samples = samples.to(device, non_blocking=non_blocking)
-        targets = [{k: v.to(device, non_blocking=non_blocking) for k, v in t.items()} for t in targets]
+        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()
+        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/__init__.py

@@ -25,9 +25,11 @@
 from rfdetr.utilities.state_dict import clean_state_dict, strip_checkpoint
 from rfdetr.utilities.tensors import (
     NestedTensor,
+    PackedTargets,
     collate_fn,
     make_collate_fn,
     nested_tensor_from_tensor_list,
+    pack_targets,
 )
 
 __all__ = [
@@ -41,9 +43,11 @@
     "save_on_master",
     # tensors
     "NestedTensor",
+    "PackedTargets",
     "collate_fn",
     "make_collate_fn",
     "nested_tensor_from_tensor_list",
+    "pack_targets",
     # box_ops (submodule)
     "box_ops",
     # logger