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

[JESUSROYETH]

Follow-up to the hand-off ceiling found while measuring #1392. I open it separately because it is not about WebDataset — it sits downstream of the dataset and applies the same way to the existing loose-file loader. It doesn't compete with #1396, it composes with it: the change is in collate_fn and transfer_batch_to_device, so it doesn't care where the bytes come from, and it works the same for tar shards. Everything below is measured on develop, with no WebDataset in the tree.

TL;DR — Every tensor a DataLoader worker returns gets moved into its own shared-memory segment and handed to the parent as a file descriptor. rf-detr returns 7 tensors per detection sample, so a batch of 16 crosses as 114 objects. Of those, 112 carry 3,592 bytes in total — 0.005% of a 65.81 MB batch — but they pay most of the per-object cost anyway. Concatenating each field into one tensor takes that down to 9 objects, same payload: 1.25x-1.67x loader throughput depending on worker count, and it removes a file-descriptor exhaustion crash the current loader hits at 16+ workers. Output is bit-identical and COCO metrics don't move, to ten decimal places. Off by default.

The cost, isolated

Only the collate output changes here, everything upstream stays the same. 16 workers, unpinned, three repetitions, median:

what crosses	img/s	ms/batch	added by this stage
nothing (count)	553.4	28.914	— upstream: read, decode, augment, build targets
+ image and padding mask	507.0	31.556	+2.641
+ packed targets (9 objects)	504.6	31.708	+0.152
+ per-sample targets (114 objects, today)	312.3	51.233	+19.525

Shipping the targets as 112 separate objects costs 19.5 ms per batch — 128x what the packed form costs, and 7x the cost of the entire 65.81 MB image payload. So the cost is per object, not per byte. This ablation uses a standalone packer, so the only thing that varies is the collate output — the throughput table below uses the collate this PR actually ships.

Loader throughput, as shipped

COCO train2017, RFDETRNano, resolution 576, batch_size=16, one g2-standard-32 (32 vCPU, L4). make_collate_fn(pack=True) against today's collate, 300 timed batches after 20 warm-up, three repetitions, median [min-max]:

workers	pin_memory	today	packed	ratio
8	on	228.7 [219.0-230.6]	285.6 [284.9-288.0]	1.25x
8	off	286.3 [280.4-292.0]	284.0 [281.6-284.8]	0.99x (ranges overlap)
16	on	218.4 [213.6-225.9]	365.6 [362.9-366.4]	1.67x
16	off	313.8 [310.2-314.4]	443.9 [429.2-445.7]	1.42x
32	on	212.9 [211.4-221.8]	337.0 [331.4-341.7]	1.58x
32	off	299.7 [297.5-304.9]	465.9 [463.3-483.9]	1.56x

At 8 workers unpinned there is no gain: the ranges overlap, and the workers, not the hand-off, are the constraint there. Every other cell is non-overlapping. Pinned throughput dips from 365.6 at 16 workers to 337.0 at 32 — my guess is the single pin-memory thread turns into the serialiser once enough workers feed it. That would also explain why the unpinned column keeps scaling over the same step (443.9 → 465.9), but I did not instrument that thread to prove it.

It also removes a crash

At the stock soft limit of 1024 file descriptors, the current loader dies at 16 and 32 workers when pin_memory=False:

RuntimeError: received 0 items of ancdata

in rebuild_storage_fd → recvfds. One descriptor per tensor, 114 per batch, several batches in flight. Reproduced 6 out of 6 attempts (three at each worker count), and the packed path completed every time in the same setup. So the unpinned rows above were measured with ulimit -n 65536, otherwise both arms could not even run. This is also why set_sharing_strategy("file_system") helps — it is the usual workaround for this exact failure — and packing gets there without touching the strategy. The pinned path is not affected: the same setup at ulimit -n 1024 with pin_memory=True runs clean (0 ancdata errors). So this is a hazard of the unpinned loader, not of high worker counts in general.

pin_memory

PackedTargets implements pin_memory() and on purpose does not subclass collections.abc.Sequence. PyTorch's pin-memory worker checks that ABC before it looks for a pin_memory method (torch/utils/data/_utils/pin_memory.py), so a Sequence subclass gets taken apart and pinned tensor by tensor — which rebuilds, in the main process, exactly what packing was meant to avoid, in the default CUDA setup. As written, the batch stays packed through pinning, and transfer_batch_to_device moves one tensor per field. There is a regression test for this, because the coupling is not obvious just from reading the class.

Against set_sharing_strategy("file_system"), which you asked about

Same harness, 16 workers, batch_size=16, three repetitions, median:

	file_descriptor (default)	file_system
today	218.4	353.9
packed	365.6	417.9

Pinned, which is the CUDA default, packing is the better single change — 365.6 against 353.9 for the strategy switch — and the two together reach 417.9, 1.91x over the untouched baseline.

Unpinned, the answer flips, and I'd rather say so here than have it found in review: file_system alone reaches 488.3, better than packing on file_descriptor (443.9), and packing on top of file_system drops to 427.6, 0.88x. For an unpinned loader the one-line strategy switch is the better change, and this patch should not be stacked on it. My guess is file_system already removes the descriptor round-trip that packing exists to amortise, leaving only its bookkeeping — but I did not check that, so take it as an observation, not a mechanism.

The saving is per-object, so it scales with batch size

Same worker count and pinning, varying only batch_size:

batch	today	packed	saved per batch	objects removed	per object
8	229.2	352.1	12.19 ms	49	0.249 ms
16	218.4	365.6	29.50 ms	105	0.281 ms
32	219.9	352.5	54.74 ms	217	0.252 ms

The per-object cost stays flat across a 4x range of batch size, and that's what makes this structural instead of a tuning artifact: a batch carries 7 x batch_size target tensors no matter how many annotations each image holds, so the object count removed from the worker-to-main boundary follows from batch size alone. This run only sweeps batch size — annotation density, image content and dataset size were not varied. Packing's own compute (torch.cat, and .clone() in as_list()) does scale with the total number of boxes, just several orders below the ~19.5 ms/batch IPC saving measured above, at COCO's typical annotation density. What changes with upstream work — bigger images, heavier augmentation — is the ratio, because that grows the denominator.

The operational effect: fewer workers needed

configuration	img/s
today, 32 workers, pinned	212.9
packed, 8 workers, pinned	285.6
today, 32 workers, unpinned	299.7
packed, 16 workers, unpinned	443.9

Eight packed workers beat thirty-two unpacked ones in the pinned default, and sixteen beat thirty-two unpinned. It's the same lock contention and thread thrashing from #1392, just showing up as free cores this time instead of extra throughput.

End to end, this changes nothing

A RFDETRNano fine-tune (3 epochs, 2,000 images, batch_size=8, 8 workers, two repetitions):

	wall clock (s)	steady-state GPU
today	252.0 / 252.9	62.3% / 61.6%
packed	249.9 / 252.2	63.0% / 62.1%

Identical within noise, and that's the expected result: this run only consumes about 24 img/s while the loader already supplies 220+, so the loader was never the bottleneck and making it faster can't help here. What this table is really showing is that there is no regression. The change pays off where the loader is the bottleneck — high worker counts, object storage, several ranks sharing a host's cores — and in removing the descriptor crash.

Correctness

• COCO metrics don't move. Same published RFDETRNano checkpoint evaluated twice over COCO val, differing only in pack_targets: all 7 metrics identical to ten decimal places. At resolution=576 (what repro/eval_ab.py passes): mAP@50:95 = 0.5181497931, mAP@50 = 0.7134777308, F1 = 0.6934525371. The absolute level depends on the eval resolution, so the same comparison run at evaluate()'s own default of 384 lands at mAP@50:95 = 0.4801311492919922 — also exact, every metric differing by 0.0. So parity holds at both resolutions, but quote the number together with the resolution or the two won't match. build_trainer() sets deterministic=False for evaluation, same as for training, so nothing in the trainer path forces bit-exact reproducibility — the exact 0.0 diff above is what I observed running fixed weights through it twice, not a guarantee from the code.
• Bit-identical on real data, through the real RFDETRDataModule: 25 batches, 200 images, 1,400 target tensors compared field by field after transfer_batch_to_device — 0 value, 0 dtype and 0 shape mismatches. Object count confirmed 58 → 9 at batch_size=8.
• Mixed dtypes fall back instead of promote. torch.cat promotes across dtypes instead of failing, and COCO hits that case: a sample with no annotations builds area/iscrowd from an empty list and gets float32, where a populated sample gets int64. Packing a field like that would silently change its dtype and can lose integer precision (2**54+1 does not survive float32). The packer catches it and returns the batch untouched.
• Non-strided and quantized values fall back too, for the same reason. A sparse value has no linear element order for reshape(-1) to preserve, so torch.cat raises RuntimeError instead of falling back on its own. A quantized value is trickier: it can share torch.dtype with another sample while using a different (scale, zero_point), and in that case torch.cat silently requantizes it to the first sample's parameters instead of failing — so dtype equality alone does not catch either case. Neither is reachable through this repo's own COCO/YOLO loaders, but pack_targets is documented to fall back for any custom dataset it cannot pack losslessly. Both are now covered by that guard and a regression test.
• TrainConfig.pack_targets is tested through RFDETRDataModule, not only through pack_targets() directly. train_dataloader().collate_fn is asserted to return a PackedTargets when the config sets it, and the original tuple of dicts when it does not — that covers the wiring between the config field and the DataLoader it configures.
• PackedTargets/pack_targets are exported from rfdetr.utilities, matching NestedTensor/nested_tensor_from_tensor_list — the same public surface the image half of the collated batch already has, so a caller doing isinstance(targets, PackedTargets) against train_dataloader()'s output doesn't need to reach into the private module. The transfer_batch_to_device row in docs/learn/train/customization.md's lifecycle-hooks table now also names the packed-targets branch.
• There are regression tests for each of those, plus mutation testing: collapsing dtypes, dropping empty samples, raising instead of falling back, returning one shared dict, disabling the dtype guard, removing the materialisation in transfer_batch_to_device, and reverting the as_list copy. Each mutation fails exactly the test that claims to cover it, and the suite goes back to green on revert.
• Full CPU suite as ci-tests-cpu.yml runs it, on this branch and on pristine develop: 2 failures on both, the same two, 4312 passed here against 4296 — the difference is the tests this branch adds.
• mypy src/rfdetr/ produces byte-identical output to pristine develop.

Not covered

• Segmentation and keypoints. The 114 → 9 count is for detection targets. Those heads add a field, making it 130 → 10, and the packer handles any same-keyed tensor field, but neither was measured — no throughput number and no parity run for them. transfer_batch_to_device always materialises the packed batch back into a plain list (as_list(), one .clone() per sample) before handing it off, so the packed and materialised copies of a field sit in memory at the same time for the duration of that call. For a small boxes/labels field this is negligible, but for a segmentation masks field (bytes scale with instances × resolution²) it is a real, unmeasured transient doubling of that field's memory. A mask field would also make as_list() hold both the packed buffer and the materialised views at once. The obvious fix is to release each packed field as it gets materialised, which keeps memory bounded, but that's a follow-up, not implemented here.
• Multi-GPU DDP. The change is rank-agnostic, since it lives inside collate and the device hook, but I did not run it under DDP.
• How the remaining upstream cost splits between JPEG decode, augmentation and target construction. The ablation bounds the boundary at 2.79 ms of a 31.7 ms batch, so 28.9 ms is upstream, but profiling attempts failed on this host (ptrace_scope), so the internal split stays unmeasured.

Where the remaining headroom is

After packing, the boundary is 2.79 ms of a 31.7 ms batch (8.8%). That bounds the structural follow-ups: a pre-allocated buffer pool, or packing the image into the same buffer, can only attack that 2.79 ms — and that means replacing the loader's IPC, not just changing a collate_fn. Cutting image bytes does not stack — uint8 on top of packing measured 0.983x, with overlapping ranges. The only lever that would reach the upstream 28.9 ms is moving the boundary itself, crossing compressed bytes and decoding on GPU — but that conflicts with keeping Albumentations usable, since it works on the CPU-decoded image.

Default

TrainConfig.pack_targets defaults to False, so nothing changes unless you ask for it. My proposal is to default it to True for the pinned path, which is the CUDA default and where it's faster at every worker count I measured, and composes with file_system. I would not turn it on for an unpinned loader already on file_system, where it measured 0.88x. At 16 or more workers the current path doesn't just run slower, it crashes on a system with an untouched descriptor limit, and the packed path is faster in five of six setups and no worse in the sixth. The cost is a visible API surface — a caller consuming datamodule.train_dataloader() directly gets a PackedTargets instead of a tuple of dicts whenever a batch packs losslessly — which is why I did not flip it on my own. All six setups above are batch_size=16 with num_workers in {8, 16, 32}. TrainConfig's own shipped defaults are batch_size=4 and num_workers=2, a regime none of this throughput grid covers, and where neither the crash nor the gain is measured. If defaulting it on is too much for one PR, the smaller step is to warn when num_workers >= 16 with pin_memory=False and pack_targets=False, which is the exact combination that dies. The full suite above ran with it on, so the evidence for the flip exists either way — the call is yours.

Happy to re-run any of this on your 48-core box if you want the numbers from the machine that motivated #1392 .. the scenarios where this pays off end to end (many ranks, constrained cores, object storage) are argued from the mechanism here, and your machine is the one that would actually settle it.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86%. Comparing base (05786ce) to head (7a8b804).
⚠️ Report is 8 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop   #1399    +/-   ##
========================================
  Coverage       86%     86%            
========================================
  Files          112     113     +1     
  Lines        14671   14791   +120     
========================================
+ Hits         12629   12748   +119     
- Misses        2042    2043     +1     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Borda]

Not covered> - Segmentation and keypoints. The 114 → 9 count is for detection targets. Those heads add a field, making it 130 → 10, and the packer handles any same-keyed tensor field, but neither was measured — no throughput number and no parity run for them. transfer_batch_to_device always materialises the packed batch back into a plain list (as_list(), one .clone() per sample) before handing it off, so the packed and materialised copies of a field sit in memory at the same time for the duration of that call. For a small boxes/labels field this is negligible, but for a segmentation masks field (bytes scale with instances × resolution²) it is a real, unmeasured transient doubling of that field's memory. A mask field would also make as_list() hold both the packed buffer and the materialised views at once. The obvious fix is to release each packed field as it gets materialised, which keeps memory bounded, but that's a follow-up, not implemented here.

Could a good follow-up to try in two separate PRs - segm & kp

[Copilot]

Pull request overview

Adds opt-in target packing to reduce DataLoader IPC overhead while preserving downstream target semantics.

Changes:

• Adds PackedTargets, lossless packing, pinning, and device transfer.
• Wires packing through TrainConfig and RFDETRDataModule.
• Adds tests and user documentation.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
src/rfdetr/utilities/tensors.py	Implements target packing and materialization.
src/rfdetr/utilities/__init__.py	Exports the new public utilities.
src/rfdetr/training/module_data.py	Integrates packing into DataLoaders and transfers.
src/rfdetr/config.py	Adds the opt-in configuration field.
tests/utilities/test_tensors.py	Tests packing behavior and fidelity.
tests/training/test_module_data.py	Tests DataModule integration.
docs/learn/train/training-parameters.md	Documents the configuration option.
docs/learn/train/customization.md	Documents transfer-hook behavior.
CHANGELOG.md	Records the feature.

Suppressed comments (1)

docs/learn/train/training-parameters.md:358

• This statement mentions only train_dataloader(), but the configured collate function is also passed to validation, test, and predict loaders. Update the complete reference so callers know that enabling this option changes the direct output type of every RFDETRDataModule loader.

| `pack_targets`               | bool                   | False          | Concatenate each batch's target dicts into one tensor per field before crossing the DataLoader worker boundary. `train_dataloader()` yields `PackedTargets` instead of a tuple of dicts whenever a batch packs losslessly. |

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Borda]

crossing compressed bytes and decoding on GPU — but that conflicts with keeping Albumentations usable

I think we already have preferable use of Kornia on the GPU, so if we provide meaningful gain, I would be fine to make a bonus for GPU augmentation, but if the user, for whatever reason, insists on using Albumentations, we won't use it...

[Borda]

• How the remaining upstream cost splits between JPEG decode

This may be addressed by #1391