feat(datasets): stream training input from WebDataset tar shards

[JESUSROYETH]

TL;DR — No, it doesn't bring the speed-up #1392 expected, and I'd rather say that plainly. On local NVMe streaming shards match loose files within noise (0.92-1.03x) because that pipeline was never I/O-bound. Where it does pay off is object storage with a constrained worker budget: on GCS via gcsfuse it's 2.00x at 4 workers and 1.23x at 8, converging at 16. And it removes per-rank startup cost outright: 19.9 s -> 0.012 s and 2,810 MB -> 0.9 MB of resident memory. The ceiling that remains is the worker-to-main hand-off, which sits downstream of the dataset and no input format can move.

Why there's no speed-up on local NVMe

The two paths only differ up to the moment an image is decoded. After that — augmentation, collate, the hand-off, pin_memory, Kornia — it's the same code on byte-identical tensors (verified: 500 real COCO images, 0 pixel and 0 box mismatches). So the most packing can win is the share of time spent getting bytes off the disk, and on this disk that share measured under 1%: iowait never exceeded 0.91%, one fully page-cached run read 0 MB and ran at the same speed as its cold twin, and a read-only benchmark sustained ~1,230 img/s against the ~210 the loader consumes. Sequential and random access are also equally fast here (0.99x), both converging on ~192 MB/s against a ~250 MB/s single-stream dd ceiling — a bandwidth limit, not an IOPS one. COCO's ~160 KB JPEGs are already large enough that there's no per-open() latency left to remove.

So the null result isn't packing failing to work. It's packing optimising a stage that wasn't costing anything — on that hardware.

Implements the WebDataset pipeline from #1392: a packer, a streaming reader, the existing Albumentations functions mapped over the decoded stream in the CPU workers, and batched tensors handed to pin_memory and the existing Kornia GPU stage.

It's opt-in through dataset_file="webdataset". coco, roboflow and yolo stay untouched.

Changes

• src/rfdetr/datasets/webdataset_io.py — pack_coco_to_shards() plus a python -m rfdetr.datasets.webdataset_io CLI that writes ~100 MB .tar shards and a small JSON index; WebDatasetDetection, an IterableDataset that streams them; build_webdataset_loader().
• datasets/__init__.py, config.py — route and accept the new dataset_file.
• training/module_data.py — the four dataloader methods return the streaming loader for a streaming split, aligned to grad_accum_steps; a packed test split is used when present; class_names comes from the shard index, indexed by the same labels the loader emits under either category_ids policy; the private sample-grid helper rejects a split it cannot index.
• pyproject.toml — rfdetr[webdataset] extra, reader only. .github/workflows/ci-tests-cpu.yml installs it, otherwise the streaming half of the new tests sits behind importorskip and silently skips.
• tests/datasets/test_webdataset_io.py, docs/learn/train/dataset-formats.md, CHANGELOG.md.

Two decisions worth flagging

Packing uses the standard library's tarfile, not wds.ShardWriter. A WebDataset shard is just an ordinary tar file whose members are grouped by basename, so writing one doesn't need anything from the package. That means the packer and its tests run in the default CI job instead of an extra-gated one, only the reader needs the dependency.

The loader is a plain DataLoader, not wds.WebLoader. WebLoader builds exactly that DataLoader and then wraps it in a DataPipeline, which drops __len__ and the DataLoader type. len() on one raises TypeError. RF-DETR needs that length: trainer.estimated_stepping_batches drives the LR schedule and the drop schedule, and WebLoader only offers with_length(), which asserts a length without guaranteeing it. On top of that, WebLoader adds a post-loader unbatched().shuffle().batched() across workers, and that would break the same guarantee. The streaming behaviour comes from the pipeline inside the dataset, not from the loader class, so the result is the same either way. I'm happy to switch if you'd rather have the wrapper.

Epoch sizing

Training gives each worker a fixed sample count, floored to a whole number of accumulation windows (batch_size × grad_accum_steps). So the reported length equals the batches produced and no partial window reaches PTL — this is the streaming counterpart of the padding that GradAccumAlignedDataset already gives the map-style loader for the same known PTL issue. Validation and test are different: each worker drains its shards once, so the split gets scored exactly once and those loaders report no length.

Training refuses to start when world_size × num_workers is bigger than the shard count. It also warns when an uneven shard split leaves the worst-served worker more than 5% short. That check is now measured from each shard's real sample count instead of assuming shards are uniform in sample count — they are cut by byte size, so that assumption can miss a real skew when image sizes vary (see the accuracy section for why this one matters).

Scope of the distributed testing: wds.split_by_node is exercised, but by two gloo ranks on CPU, not a real NCCL multi-GPU job. Both ranks completed training and validation; a split with fewer shards than ranks now fails fast with an actionable error instead of deadlocking the process group (previously only the training loader checked this, so evaluation could hang).

Shard order and the in-stream reservoir buffer get reseeded every epoch from a per-worker epoch counter, not from torch.initial_seed() alone. The reason: under persistent_workers=True (the DataModule's own default whenever num_workers > 0), a worker process is never restarted between epochs, so its torch seed stays the same, and a seed derived only from it would replay the first epoch's order forever. The counter lives on the dataset instance instead, which does survive across epochs inside a persistent worker, so it stays exact no matter the worker-restart policy.

Throughput

On local NVMe there is no I/O bottleneck to remove, so these numbers report what the pipeline actually does rather than a speed-up claim. The object-storage case, where the bottleneck is real, is in the next section.

One GCP g2-standard-32: 32 vCPU, L4, 125 GB RAM, 200 GB NVMe-attached persistent disk. COCO 2017 train2017, 118,287 images, 19 GB loose, packed into 189 shards. RFDETRNano at resolution 576, batch_size=16. Each cell is one process: 20 warm-up batches, then 400 timed batches, three repetitions, median with [min-max]. cold drops the page cache before every repetition.

augment	page cache	workers	loose files (img/s)	tar shards (img/s)	shards / loose
default	cold	8	226.3 [221.5-226.3]	224.4 [222.5-226.9]	0.99x
default	cold	32	210.4 [199.8-211.0]	192.9 [190.5-209.0]	0.92x
default	warm	8	225.2 [221.5-226.5]	220.4 [217.9-228.8]	0.98x
default	warm	32	198.9 [191.2-199.3]	202.8 [192.3-211.2]	1.02x
albumentations	cold	32	203.1 [200.3-207.2]	209.6 [206.2-215.8]	1.03x
albumentations	warm	32	200.6 [199.3-206.1]	199.2 [198.7-209.2]	0.99x

Three checks from the same runs say the disk is not the limit. iowait never goes above 0.91%. One fully page-cached configuration read 0 MB from disk and matched its cold twin. And a read-only benchmark on the same disk (no decode, no augmentation) sustains ~1,230 img/s against the ~210 the loader actually consumes, shards again landing at 0.99x loose.

Both access patterns land around ~192 MB/s against a ~250 MB/s single-stream dd ceiling, so this disk is bandwidth-limited, not IOPS-limited. COCO's ~160 KB JPEGs are already big enough that there is no per-open() latency left for packing to remove.

What is limiting them is the worker-to-main batch hand-off. Everything upstream stays identical; the only thing that shrinks is what crosses the process boundary (a collate that just returns the batch size, a measurement device, not something I'm proposing to ship), num_workers=32:

what crosses the boundary	loose (img/s)	CPU busy	shards (img/s)	CPU busy
full batch, pinned (the real pipeline)	198.9	33.8%	202.8	34.7%
full batch, not pinned	282.3 [263.2-291.4]	41.6%	278.4 [277.6-286.8]	42.1%
batch size only (diagnostic)	778.9 [759.3-799.7]	95.2%	822.2 [611.7-883.3]	97.9%

Decode and augmentation parallelize just fine. With the hand-off removed, going from 8 to 32 workers takes throughput from 412.7 to 778.9 img/s and CPU from 33% to 95%. With the hand-off in place, CPU stays near 33% no matter how many workers you add, which is the symptom in the issue, and it shows up on both input paths. (pin_memory=True is still the right default here. This benchmark never touches the GPU, so it only sees the pinning cost, not the async host-to-device copy that pays for it.)

Where the streaming path does win is construction time, and the gap scales with the annotation file, not with the images:

Three repetitions per arm, page cache dropped before each, median with [min-max]:

	build seconds	RSS added by setup("fit") (MB)	peak RSS (MB)
dataset_file="coco"	19.920 [19.805-19.930]	2,810.4 [2,810.4-2,810.4]	3,837.8 [3,837.8-3,837.9]
dataset_file="webdataset"	0.012 [0.011-0.012]	0.9 [0.9-0.9]	795.4 [795.1-795.6]

That's per process, so per rank under DDP — on 8 ranks, roughly 22 GB of host RAM and 160 s of startup that stop being spent. Packing pays off where per-file access is genuinely the constraint; that is measured below rather than assumed.

Loader throughput is not training throughput, so I also ran a real fine-tune on both paths: RFDETRNano, 2 epochs over 2,000 train / 500 val images, batch_size=8, 8 workers, on the same L4. Wall clock was 172.7 s on loose files against 173.3 s on shards, and steady-state GPU utilisation 61.6% against 61.3% — no end-to-end regression on a local dataset. Both arms leave the GPU idle about 39% of the time, and that idle does not move when the input format changes: it is the same hand-off ceiling the loader table runs into.

Object storage, where the bottleneck is real

Same 2,000-image split uploaded to a GCS bucket in the same region and mounted with gcsfuse, so the only difference from the rows above is the filesystem. Page cache dropped and gcsfuse's 60 s metadata TTL waited out before every run, so each is genuinely cold. Median of 3, [min-max]:

num_workers	loose files (img/s)	tar shards (img/s)	shards / loose
4	86.3 [75.0-88.3]	172.7 [160.0-173.8]	2.00x
8	174.3 [163.3-176.9]	214.5 [208.9-218.7]	1.23x
16	220.5 [210.1-226.3]	222.6 [207.2-227.1]	1.01x

The ranges don't overlap at 4 or 8 workers. Both arms converge at 16 on ~220 img/s, which is the same ceiling the local-NVMe rows hit — so shards don't raise the ceiling, they let you reach it with fewer workers. At 4 workers the loose arm sits at 10.4% CPU: starved, waiting on per-file latency that parallelism would otherwise hide.

That worker count is not a corner case. This box has 32 vCPU, so a DDP job with 8 ranks leaves about 4 workers per rank — precisely the column where the gap is 2x. The construction saving above is also per rank, so both effects land on the same scenario rather than being separate curiosities.

On batching inside WebDataset

I tested .batched(batch_size) with DataLoader(batch_size=None) as suggested: 1.02-1.06x across num_workers 4/8/16/32, with and without pinning — positive in all eight cells, with non-overlapping ranges in five of them, so the gain is real but small. What it does not do is move the ceiling, because it does not change what crosses the process boundary. PyTorch already calls collate_fn inside the worker, so both arms hand the parent the same structure — I counted it: 114 tensors either way (the image, the padding mask, and 7 per-sample target tensors x 16). .batched() moves where the collate runs, not what gets shipped.

So the integration isn't hiding WebDataset's benefit. The hand-off cost is real, but it is downstream of the dataset and independent of input format: holding the payload at 65.81 MB and sending it as 4 objects instead of 114 goes from 203.5 to 426.9 img/s, while cutting the bytes 3.25x at constant object count buys only 1.21x. That looks like its own issue rather than something to load onto this PR — happy to open it with the measurements if useful.

Accuracy

RFDETRNano fine-tuned for 3 epochs on the same 2,000-image split, once from loose files and once from shards, five seeds each, num_workers=8. Final-epoch val/mAP_50_95:

	median	range	spread
loose files	0.3681	[0.3658, 0.3704]	0.0046
tar shards, 39 shards	0.3574	[0.3556, 0.3694]	0.0138
tar shards, 290 shards (n=3)	0.3685	[0.3600, 0.3708]	0.0108

The middle row is a shard-count artefact, and finding it is exactly why the warning above exists. 39 shards over 8 workers splits 5,5,5,5,5,5,5,4, so the short worker is ~18% under the samples its fixed-length epoch asks for and wraps to fill, while better-supplied workers leave some unseen. Re-packing the identical split into 290 shards (0.6% under) moves the median from 0.3574 to 0.3685, onto the map-style 0.3681.

What that doesn't fix is the variance: the streaming arm runs 2-3x the seed-to-seed spread of the map-style one in both configurations. So, with an adequate shard count the medians match on this proxy. With a skewed one it costs ~0.011. The streaming path stays noisier across seeds either way, and parity on a full-length run over a real dataset is not verified here. A 3-epoch fine-tune on 2,000 images is a noisy proxy, and I'm not going to dress it up as more than that.

Validation

• New test file: 84 passed, 8 skipped (all 8 are # doctest: +SKIP helpers that need pytest fixture injection, a @pytest.fixture or a real tmp_path, to run at all).
• Streamed samples are identical to loose-file ones through the same pipeline, on 500 real COCO images: 0 pixel mismatches, 0 box mismatches, identical cat2label, class_names and per-image label sets. Both val loaders yield 500 samples and 500 unique image ids. Segmentation masks are covered the same way on a synthetic split with polygon annotations, include_masks=True on both sides: 0 mask mismatches.
• Three mutations of the implementation (dropping the batch-multiple flooring, re-encoding instead of copying image bytes, skipping a missing image instead of raising) each break exactly the tests that claim to cover them, and the suite goes green again on revert.
• Full CPU suite with the command from ci-tests-cpu.yml: 4351 passed, 24 failed, 79 skipped, 3 errors. Pristine develop at 74bfc449 gives 4259 passed, 24 failed, 71 skipped, 3 errors, and the 27 FAILED/ERROR lines are identical in both: optional packages missing from my environment (roboflow, onnxscript, jsonargparse[signatures], onnxruntime), none of them in a file this PR touches.
• Resuming from a checkpoint behaves exactly as the loose-file path does. Recording the image_id order of every training epoch, a run interrupted after epoch 1 and resumed from last.ckpt replays the sample order of epoch 0 at epoch 2 and of epoch 1 at epoch 3, rather than continuing the sequence — and the map-style COCO path does the identical thing on the same probe. So this is existing behaviour of the repo's resume, not something the shard path introduces; I mention it only so it isn't attributed to this PR.
• pre-commit run --all-files: 19/19 hooks, mypy --strict included. mkdocs build succeeds with the new section.
• uv lock and uv sync --no-default-groups --extra webdataset --python <V> --dry-run for V in 3.10-3.14 all succeed.

More defects were caught by running the code rather than reading it, each now with a regression test:

• shardshuffle=True makes webdataset warn from every worker on every epoch (it wants an integer).
• A fixed shuffle seed replays one sample order every epoch, because the pipeline is rebuilt per epoch and webdataset seeds its shuffle per generator call.
• A planned epoch with more workers than shards leaves a worker empty and silently shortens training.
• The per-epoch reseed above derived its seed from torch.initial_seed() alone, which doesn't change across epochs when persistent_workers=True (the DataModule's own default whenever num_workers > 0). So shard order and the reservoir buffer replayed the first epoch's draw forever under that config, instead of reshuffling. PyTorch still calls __iter__ fresh on every epoch for a persistent worker, just without restarting the process, so a plain counter on the dataset instance survives across epochs too. Add it to whichever seed source doesn't already change on its own, and it fixes both cases.
• The fixed epoch length floored only to a multiple of batch_size, not batch_size × grad_accum_steps, so a partial accumulation window at the tail could still reach PTL. That's the exact bug GradAccumAlignedDataset exists to prevent on the map-style path.
• A packing run that failed partway through (a missing image, an unlabelled annotation) wrote its shards directly into the destination directory, so re-packing over an existing valid split left it corrupted, with some shards overwritten by the new, incomplete run and others left stale, all still referenced by the old, untouched index. The packer now stages a run in a temp directory and publishes shards and the index together only once every image has been read successfully.
• A minority of annotations whose image_id matched no image were dropped silently. The packer only rejected the case where all image ids were orphaned.
• The shard-skew warning approximated the worst-served worker's share from shard count alone. Shards are cut by byte size, not sample count, so a split with wide per-image size variation can read as perfectly balanced under that approximation while one worker actually gets a small fraction of the samples. The packer now records each shard's real sample count, and the warning uses it when available.
• class_names under category_ids="raw" returned a list sorted by id but indexed by position, not by id. So class_names[raw_label] could name the wrong category whenever the id range has a gap below the label being looked up. It's now indexed by id, matching how the remap branch already worked (empty string at every unused slot). The docs' num_classes formula for raw (len(categories)) undercounts real COCO for the same reason: ids run 1-90 with gaps for 80 categories, so the model needs 91 output slots, not 80. Now it's stated as max(category_id) + 1, matching dataset_file="coco"'s existing convention.
• split and an annotation's file_name flowed into shard/index file paths and image reads with no validation, so split="../x" or a crafted file_name could write or read outside the intended directory. Both are now checked to resolve inside the shard/image directory.

Not covered

Two gaps worth naming rather than leaving for review to find. The streaming path still shows 2-3x the seed-to-seed spread of the map-style loader (0.0138 and 0.0108 against 0.0046 over five seeds), and that extra variance was not isolated to a cause. And there is no verified accuracy parity on a full-length run over a real dataset — what exists is a 3-epoch, 2,000-image proxy. Distributed testing is gloo/CPU with two ranks, not NCCL multi-GPU.

Keypoint training rejects this format explicitly: its label space gets inferred from a whole parsed COCO annotation file, and a shard index doesn't carry that. num_classes has to be set on the model by hand (RFDETRSmall(num_classes=...)), because the auto-detection the other formats use looks for train/_annotations.coco.json or data.yaml and finds neither. Under category_ids="remap" that number is the categories left after grouping nodes are dropped; under "raw" it's the highest category_id plus one. Class names need no such help. Shard paths resolve against the local dataset_dir; streaming straight from object storage is the obvious next step, but it's not wired up here.

---

closes #1392

[socket-security]

Review the following changes in direct dependencies. Learn more about Socket for GitHub.

Diff	Package	Supply Chain Security	Vulnerability	Quality	Maintenance	License
	pypi/​webdataset@​1.0.2

View full report

[codecov]

Codecov Report

❌ Patch coverage is 94.32990% with 22 lines in your changes missing coverage. Please review.
✅ Project coverage is 86%. Comparing base (7b9cba6) to head (7918c18).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop   #1396    +/-   ##
========================================
  Coverage       86%     86%            
========================================
  Files          114     115     +1     
  Lines        14877   15264   +387     
========================================
+ Hits         12832   13201   +369     
- Misses        2045    2063    +18     

🚀 New features to boost your workflow:

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

[Copilot]

Pull request overview

Adds opt-in WebDataset shard packing and streaming for detection and segmentation training.

Changes:

• Adds COCO-to-tar packing, shard indexing, streaming datasets, and loaders.
• Integrates WebDataset into configuration, data modules, CI, and dependencies.
• Adds comprehensive 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/datasets/webdataset_io.py	Implements packing, streaming, indexing, and loader construction.
src/rfdetr/datasets/__init__.py	Routes WebDataset creation.
src/rfdetr/training/module_data.py	Integrates streaming loaders into training and evaluation.
src/rfdetr/config.py	Adds the WebDataset format option.
tests/datasets/test_webdataset_io.py	Tests packing, streaming, transforms, and integration.
pyproject.toml	Adds the optional WebDataset dependency.
.github/workflows/ci-tests-cpu.yml	Installs WebDataset during CPU CI.
docs/learn/train/dataset-formats.md	Documents packing and training workflows.
CHANGELOG.md	Records the feature and its limitations.

---

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

[Borda]

@JESUSROYETH, this is a great, detailed report. Could you pls alos inlcude TLDR section? Also, from the numbers, it seems that this did not bring any expected speedup?

[Borda]

I think it may be worth exploring a more native WebDataset pipeline rather than constraining the implementation to the current DataLoader pattern.

Constraints

• No performance regression on local datasets vs. the existing loader.
• Keep augmentations flexible: Kornia, Albumentations, or torchvision should remain usable.
• Compatible with PyTorch Lightning, directly or with only small/localized adjustments.
• Preserve existing dataset paths; WebDataset remains opt-in.

Why explore this?

The current experiments suggest WebDataset itself may not be the bottleneck:

• full pipeline with pinned batches: ~203 img/s
• without pinning: ~278 img/s
• diagnostic avoiding image-batch IPC: ~822 img/s

This suggests the workers can produce data substantially faster than the full loader delivers it, with the likely bottleneck around:

WebDataset workers
    ↓
multiprocessing IPC / batching
    ↓
main process
    ↓
pin_memory
    ↓
GPU

So the ~1× result on local NVMe may not represent the performance ceiling of WebDataset. Local storage also doesn't appear I/O-bound, meaning faster tar reads alone wouldn't necessarily improve end-to-end throughput.

Try more native WebDataset batching

In particular, I'd test batching inside WebDataset:

dataset = (
    WebDataset(...)
    .decode(...)
    .map(...)
    .batched(batch_size)
)

loader = DataLoader(
    dataset,
    batch_size=None,
    ...
)

instead of sending individual samples through the normal DataLoader batching path.

This could reduce worker → main-process IPC/batching overhead and tell us whether the current integration is leaving performance on the table.

Lightning / dataset-size constraint

If the main blocker to a more native WebDataset pipeline is that Lightning/RF-DETR needs a reliable len(dataset) / len(loader) for scheduling, I'd prefer solving that explicitly rather than shaping the whole pipeline around this limitation.

For example, wrap/subclass the WebDataset pipeline and expose the known size:

class SizedWebDataset(...):
    def __init__(self, ..., size: int):
        ...
        self.size = size

    def __len__(self):
        return self.size

We already know the number of samples when generating shards / reading dataset metadata, so this seems like a relatively small compatibility layer.

Suggested experiment matrix

1. Current implementation:
   WebDataset samples → DataLoader batching
2. WebDataset .batched():
   WebDataset batches → DataLoader(batch_size=None)
3. Both with pin_memory=True/False.
4. Sweep num_workers to locate the actual saturation point.
5. Measure both loader throughput and actual GPU utilization/training throughput.
6. Repeat on local NVMe and a storage setup where file-open/IOPS overhead matters.
7. Verify the best pipeline with Lightning + DDP.
8. Verify augmentations remain interchangeable between Kornia / Albumentations / torchvision.

The goal isn't necessarily to make WebDataset faster than loose files on local NVMe—the current results suggest that case isn't I/O-bound. Rather, we should make sure the integration isn't introducing a PyTorch IPC/pinning bottleneck that hides WebDataset's potential benefits, while keeping zero/minimal local performance cost.

[JESUSROYETH]

@Borda Sounds good, I'll try the suggestion and get back to you.

[JESUSROYETH]

@Borda :

Ran the full matrix. One g2-standard-32 (32 vCPU, L4), COCO train2017 in 189 shards, resolution 576, batch_size=16, 300 timed batches after 20 warm-up, three repetitions, medians. Tables and harness are in the PR body.

Constraints

• No performance regression on local datasets ✅ — loader 0.92-1.03x, and end-to-end a wash: 172.7 s loose vs 173.3 s shards, GPU utilisation 61.6% vs 61.3%.
• Kornia / Albumentations / torchvision stay usable ✅ — all three bit-identical to the loose-file pipeline, 0 pixel and 0 box mismatches over 32 images.
• Compatible with PyTorch Lightning ✅ — no training-loop changes; __len__ exposed exactly as you suggested, verified under DDP.
• Existing paths preserved, WebDataset opt-in ✅ — through dataset_file="webdataset".

Experiment matrix

1. Current implementation — baseline: 160 / 236 / 232 / 226 img/s at 4 / 8 / 16 / 32 workers pinned; 163 / 315 / 329 / 303 unpinned.
2. .batched() — 1.02-1.06x, positive in all eight cells. Real but small, and it doesn't move the ceiling: PyTorch already runs collate_fn in the worker, so 114 tensors cross the boundary either way. It changes where the collate runs, not what gets shipped.
3. pin_memory on/off — the larger effect: unpinned is 1.33x / 1.42x / 1.34x at 8 / 16 / 32 workers. Your 203 vs 278 img/s reproduces. A wash at 4 workers, where the workers are the constraint.
4. num_workers sweep — unpinned saturates at 16 (329 img/s) and regresses at 32 (303). Pinned saturates earlier, at 8.
5. GPU utilisation / training throughput — 61.6% vs 61.3% steady state, ~39% of the GPU idle in both arms; the idle doesn't move when the input format changes.
6. NVMe vs IOPS-bound storage — NVMe 0.92-1.03x, nothing to remove. GCS via gcsfuse: 2.00x at 4 workers, 1.23x at 8, 1.01x at 16. Since 8 DDP ranks on 32 vCPU leaves ~4 workers per rank, that first column is the normal case, not a corner one.
7. Lightning + DDP ✅ — two gloo ranks, disjoint shard subsets, epoch length agreed across ranks.
8. Augmentation backends ✅ — as above.

Also checked, since a surprise there would cost you review time: resume from checkpoint replays epoch 0's sample order at epoch 2 instead of continuing the sequence — and the map-style COCO path does the identical thing on the same probe, so it's existing resume behaviour rather than anything the shard path introduces.

Where the headroom is

The hand-off, not the input format. The 112 per-sample target tensors are 3,592 bytes of a 65.81 MB batch — 0.005% of the payload — but each pays its own shared-memory allocation and fd round trip. Packing them into 7 takes the same harness from 329 to 565 img/s at 16 workers unpinned (1.72x), bit-identical output. It's downstream of the dataset and applies to the loose-file path too, so it isn't strictly about WebDataset — do you want it folded into this PR, or opened as a separate one?

[Borda]

The hand-off, not the input format. The 112 per-sample target tensors are 3,592 bytes of a 65.81 MB batch — 0.005% of the payload — but each pays its own shared-memory allocation and fd round trip. Packing them into 7 takes the same harness from 329 to 565 img/s at 16 workers unpinned (1.72x), bit-identical output. It's downstream of the dataset and applies to the loose-file path too, so it isn't strictly about WebDataset — do you want it folded into this PR, or opened as a separate one?

Seems the target works are quite independent with possible high yield, so let's focus on that one in another PR, and when it lands, we could integrate it here and finish it 🦝

[JESUSROYETH]

@Borda :

Revalidated this branch after #1399 and current develop (7b9cba61).

• [MEASURED] Full COCO detection parity was exact for all 5,000 validation images with loose files and WebDataset shards, packing on/off. The corresponding targets, boxes and model metrics matched within each source; minimum matched-box IoU was 1.0.
• [MEASURED] Full segmentation parity was exact across 5,000 images, 36,335 instances and 3,536,994,240 mask bytes, for both loose files and WebDataset shards with packing on/off.
• [MEASURED] The segmentation memory concern is confirmed. On real COCO batches, packed transfer increased peak CUDA memory by almost exactly one additional mask field: +5,455,360 bytes for a 5,451,264-byte p50 mask field, +9,350,656 for a 9,345,024-byte p95 field, and +12,760,576 for a 12,752,064-byte maximum field. Across 10 repetitions, isolated packed transfer was also 33–42% slower for these batches.
• [MEASURED] Full loose-COCO keypoint parity was exact across 5,000 images, 10,777 people and 549,627 keypoint values. However, only 13/625 batches packed. The other 612 fell back solely because empty iscrowd tensors are float32, while populated tensors are int64 (src/rfdetr/datasets/coco.py:689). The lossless mixed-dtype fallback in pack_targets is therefore behaving correctly (src/rfdetr/utilities/tensors.py:550-556).
• [MEASURED] The CPU suite completed with 4,232 passed and 168 skipped. The latest-base focused suite completed with 302 passed, 1 skipped and 2 deselected; the matcher suite added by current develop completed with 142 passed, 5 skipped and 8 deselected. All pre-commit hooks, including strict mypy, passed.
• [MEASURED] All 66 serial GPU test nodes passed after resolving an ONNX-environment-only missing ml_dtypes dependency. Twelve real two-epoch training arms also completed successfully. Their metric variation was not deterministic: three identical unpacked same-seed baselines ranged from 0.31652 to 0.32394 mAP50-95, and exact batch tracing showed stable image order but different augmented image/target bytes from batch 0 even between unpacked repeats.

I agree that segmentation and keypoints should be follow-ups, but I suggest separating them by cause:

1. A segmentation materialization PR that prevents the packed mask buffer and all materialized mask clones from coexisting. It should carry real p50/p95/max peak-memory measurements, 10-repetition timing spreads and full-COCO mask parity.
2. A small generic COCO target-dtype PR that makes iscrowd explicitly torch.int64 and adds real keypoint packability/parity coverage. This is not a keypoint-specific packer change.

WebDataset keypoint support should remain a separate feature PR. It is currently explicitly unsupported because its label schema requires information that the shard index does not carry (src/rfdetr/datasets/webdataset_io.py:1095-1100).

Actual multi-GPU/DDP execution remains unverified here; this validation was single-GPU and serial on an NVIDIA L4.

[Borda]

Sounds like the file lock penalize too many workers, so wondering if we could have nb packets as batch size so we each worker could sample only one basket as only customer...