perf(training): group optimizer parameters by their hyperparameters

[JESUSROYETH]

What

get_param_dict built one optimizer parameter group per trainable tensor. torch.optim's foreach and fused AdamW kernels batch a group's parameters into a single multi-tensor launch, so one group per parameter kills that batching completely. Every step ran one single-tensor kernel launch per parameter, plus one Python iteration of the optimizer's per-group loop. On rfdetr-nano that is 465 groups covering 28 distinct (lr, weight_decay) combinations.

The fix groups parameters by the hyperparameters they already had, so each keeps exactly the learning rate and weight decay it had before, and AdamW.step() issues one launch per distinct configuration instead of one per tensor.

Model	Trainable tensors	Parameter groups before	After
rfdetr-nano	465	465	28
rfdetr-small	487	487	28
rfdetr-medium	509	509	28

Measurements

One L4 (g2-standard-8, torch 2.9.1+cu129), develop@6674d858 against this branch, arms interleaved run by run.

AdamW.step() in isolation, real model parameters and gradients resident, 51 interleaved repetitions each, torch.cuda.synchronize() on both sides of every step:

Model	fused	Before (median)	After (median)	Change
rfdetr-nano	True	58.838 ms [58.801, 58.924]	8.786 ms [8.747, 8.906]	−85.1%
rfdetr-nano	False	58.856 ms [58.800, 58.887]	8.783 ms [8.742, 8.928]	−85.1%
rfdetr-small	True	61.523 ms [61.468, 61.566]	9.149 ms [9.108, 9.313]	−85.1%
rfdetr-small	False	61.512 ms [61.461, 61.570]	9.182 ms [9.137, 9.295]	−85.1%

Full training step through RFDETR.train() on a 512-image COCO subset, 100 timed steps after 20 warm-up steps, one process per arm:

Model	batch	Before (median)	After (median)	Change	Absolute
rfdetr-nano	4	223.17 / 228.04 / 225.18 ms	170.77 / 176.65 / 172.09 ms	−23.5% / −22.5% / −23.6%	≈ −52 ms
rfdetr-small	4	247.54 / 249.32 / 247.31 ms	189.86 / 190.70 / 192.00 ms	−23.3% / −23.5% / −22.4%	≈ −57 ms
rfdetr-small	8	357.89 / 359.75 ms	294.54 / 311.24 ms	−17.7% / −13.5%	≈ −56 ms
rfdetr-small	16	625.10 / 632.31 ms	568.24 / 567.00 ms	−9.1% / −10.3%	≈ −61 ms

The saving is a fixed cost per optimizer step, not a fraction of it. That's why the percentage falls as the batch grows, while the paired absolute deltas above (before − after, per row) stay in a 48.5–65.3 ms range. That matches the 50.1 ms (nano) and 52.4 ms (small) the isolated measurement predicts. At the default batch_size=4 it is the larger number. At batch_size=16 it is still a real 9–10%. Interquartile ranges do not overlap at batch_size 4 or 8 (rfdetr-nano, batch 4, round 1: before p25–p75 219.9–229.8 ms, after 168.4–176.4 ms). At batch_size=16 they do overlap, and only the medians separate the two arms.

Equivalence

AdamW works elementwise per parameter, and grouping only decides which hyperparameters apply and how the kernels batch. I started from identical weights and ran 12 steps with identical pseudo-gradients on a real rfdetr-nano: all 465 parameter tensors come out bit-identical between the two groupings, with fused=True and with fused=False.

test_each_parameter_keeps_its_own_hyperparameters checks the same property statically: it compares the (lr, weight_decay) each parameter ends up with, parameter by parameter, against the pre-merge assignment. That way backbone layer-wise LR decay and the zero weight decay on norms/biases get checked instead of just assumed. test_each_parameter_keeps_its_layerwise_learning_rate repeats that check against the optimizer configure_optimizers actually builds.

The managed LambdaLR presets are unaffected: a single lr_lambda scales every group from its own initial_lr, and merged groups share the lr they were constructed with.

Resuming a run started before this change

Optimizer.load_state_dict refuses a state dict whose parameter-group count differs from the optimizer's, so a checkpoint written with the old layout would otherwise fail to resume. on_load_checkpoint now regroups it. regroup_unmerged_optimizer_state recognises the one-group-per-parameter layout, reindexes the per-parameter state onto the merged groups, and collapses the paired LR-scheduler lists that hold one entry per parameter group (base_lrs, _last_lr, lr_lambdas, min_lrs). The merged layout comes straight from the saved groups, so it needs no model traversal, and state already saved in the current layout just passes through untouched.

A scheduler wrapped in a linear warmup ramp (SequentialLR, built by _wrap_with_warmup whenever warmup_epochs > 0 with a non-managed scheduler) nests each wrapped scheduler's own state inside a _schedulers list, instead of holding its per-group lists at the top level. The collapse recurses into that nesting too. Left alone, the wrapped scheduler's own stale per-group list raises once it becomes active, because its get_lr() zips that list against the now-shorter optimizer.param_groups. I reproduced this directly against torch.optim.lr_scheduler.SequentialLR before adding the recursive collapse.

An explicit lr_scheduler constructed from lr_scheduler_kwargs (rather than the managed "step"/"cosine" presets) that passes a value per parameter group — think LambdaLR's or MultiplicativeLR's list-valued lr_lambda — needs that list sized to the new group count. torch.optim raises ValueError at scheduler construction if the length no longer matches.

Memory

Peak allocated device memory over the timed steps is unchanged: 6388.59 MiB in both arms at batch_size=8, and 12150.94 MiB in both arms at batch_size=16 (one baseline run at batch 16 peaked at 12041.06 MiB and the other at 12150.94 MiB, so that spread at batch 16 comes from the baseline itself, not from this change).

I measured that directly rather than inferring it. The transient device memory of one AdamW.step(), over 11 repetitions per arm, is 0.000 MiB with fused=True in both groupings and 7.617 → 7.677 MiB with fused=False — a 0.060 MiB difference, with min = median = max in every arm on both rfdetr-nano and rfdetr-small, so there is no run-to-run spread to argue about. A merged group's foreach temporary spans the whole group, so the ceiling does rise from the largest single tensor (0.998M parameters, 3.8 MiB fp32) to the largest merged group (5.50M, 21.0 MiB fp32), but the measured cost at that ceiling is those 0.060 MiB. The auto-batch probe stays a valid upper bound either way — it already models the optimizer with one group covering every parameter.

Tests

tests/training/test_param_groups.py adds 17 tests (22 collected items counting their helper doctests). They cover parameter coverage and non-duplication, group count against the distinct hyperparameter combinations, per-parameter hyperparameter preservation, and frozen-parameter exclusion, plus the optimizer configure_optimizers actually builds. The checkpoint regrouping gets its own coverage too: loadability, state staying with its own parameter, scheduler-list collapse and values, a SequentialLR-wrapped scheduler's nested state and its activation past the warmup milestone, and already-merged or unrecognised state left untouched. Reverting get_param_dict to the pre-merge behaviour fails 12 of them.

tests/training/test_detr_shim.py::TestOnLoadCheckpoint adds one more test that drives the real on_load_checkpoint hook end to end, not regroup_unmerged_optimizer_state directly. That way a wiring mistake in the hook itself (wrong argument, wrong order relative to the other normalisation steps) gets caught too. It fails if the hook's call to regroup_unmerged_optimizer_state is removed.

The CPU suite (pytest src/ tests/ -n 1 -m "not gpu" --ignore=tests/run_smoke_all_models.py --ignore=tests/legacy/test_checkpoint_compat.py --ignore=tests/benchmarks --timeout=240) gives 4402 passed, 101 skipped, and one failure: the _export_tiny_model doctest in tests/export/test_onnx_notes.py. I checked the same doctest against unmodified develop@6674d858 and got the same failure there too, so it's not from this change. tests/benchmarks/ is excluded here because its COCO download-and-infer cases go over the 240 s per-test limit on this machine, unrelated to this change.

configure_optimizers no longer re-filters the groups on requires_grad: every source inside get_param_dict already filters. test_frozen_parameters_are_excluded and test_optimizer_covers_every_trainable_parameter_once check that what reaches the optimizer stays the same.

Not measured

Multi-GPU/DDP, MPS, XLA/TPU, and non-AdamW optimizers were not exercised. The change itself is optimizer-agnostic (it only changes the parameter groups handed to whatever optimizer is configured), but the numbers above are single-GPU CUDA AdamW only. mAP is not reported because the optimizer steps are bit-identical, not just approximately equal.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86%. Comparing base (6674d85) to head (7cf2f98).

Additional details and impacted files

@@           Coverage Diff           @@
##           develop   #1409   +/-   ##
=======================================
  Coverage       86%     86%           
=======================================
  Files          114     114           
  Lines        14880   14932   +52     
=======================================
+ Hits         12835   12887   +52     
  Misses        2045    2045           

🚀 New features to boost your workflow:

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