perf(training): group optimizer parameters by their hyperparameters

CHANGELOG.md

@@ -22,6 +22,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Changed
 
+- Optimizer parameter groups are now one per distinct learning-rate/weight-decay combination instead of one per parameter. `get_param_dict` built a separate group for every trainable tensor, so `torch.optim`'s foreach and fused AdamW kernels — which batch a group's parameters into one multi-tensor launch — degenerated to one single-tensor launch per parameter, plus one Python iteration of the optimizer's per-group loop. `rfdetr-nano` went from 465 groups to 28. Layer-wise backbone LR decay and per-parameter weight decay are unchanged: parameters are bucketed by the hyperparameters they were already assigned, so each keeps exactly the learning rate and weight decay it had, and the resulting AdamW steps are bit-identical (verified fused and foreach). Checkpoints written with the previous layout still resume: `on_load_checkpoint` regroups their optimizer and LR-scheduler state onto the merged groups (including a `SequentialLR` warmup wrapper's own nested scheduler state), which `Optimizer.load_state_dict` would otherwise reject as a different number of parameter groups. An explicit `lr_scheduler` whose `lr_scheduler_kwargs` pass a list sized to a specific parameter-group count (e.g. `LambdaLR`'s per-group `lr_lambda`) needs that list resized to the new group count.
+
 - `RFDETR.predict()` now converts PIL and uint8 NumPy inputs from HWC byte storage to contiguous CHW floating-point storage with one dtype/layout allocation and in-place scaling. The default source-image path reuses its already-materialized PIL array, while non-uint8 NumPy inputs retain torchvision's conversion path. The `[0, 1]` range-scan skip for those same two input types is unaffected: the fused conversion divides `uint8` storage by 255 and carries the same guarantee `to_tensor` did.
 
 - The deformable-attention core now reuses its sampled tensor directly for one-level inputs instead of stacking and flattening a one-element list. Multi-level packing is unchanged. This primarily removes allocation work from keypoint cross-attention, whose current configuration uses one feature level and many more queries than the detection decoder.

src/rfdetr/training/module_model.py

@@ -34,7 +34,7 @@
 from rfdetr.models.lwdetr import build_criterion_from_config, build_model_from_config
 from rfdetr.models.weights import apply_lora, interpolate_position_embeddings, load_pretrain_weights
 from rfdetr.training.callbacks.coco_eval import _get_ema_inner_module
-from rfdetr.training.param_groups import get_param_dict
+from rfdetr.training.param_groups import get_param_dict, regroup_unmerged_optimizer_state
 from rfdetr.utilities.logger import get_logger
 
 logger = get_logger()
@@ -1200,7 +1200,6 @@ def configure_optimizers(self) -> OptimizerLRSchedulerConfig:
         # name-prefix mismatches that put the same tensor in multiple groups.
         model_for_params = getattr(self.model, "_orig_mod", self.model)
         param_dicts = get_param_dict(ns, model_for_params)
-        param_dicts = [param_group for param_group in param_dicts if param_group["params"].requires_grad]
 
         optimizer_cfg = tc.optimizer
         optimizer: torch.optim.Optimizer
@@ -1473,6 +1472,11 @@ def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
                 self.model_config.positional_encoding_size,
             )
 
+        # Optimizer/scheduler state saved before parameters were grouped by hyperparameters carries
+        # one parameter group per parameter, a layout the optimizer no longer has. Regroup it so
+        # resuming such a run keeps its momentum and LR schedule instead of failing to load.
+        regroup_unmerged_optimizer_state(checkpoint)
+
         # Stash legacy EMA weights for RFDETREMACallback.setup(), which restores
         # them into AveragedModel when resuming from converted legacy checkpoints.
         if "legacy_ema_state_dict" in checkpoint:

src/rfdetr/training/param_groups.py

@@ -58,7 +58,56 @@ def get_vit_weight_decay_rate(name: str, weight_decay_rate: float = 1.0) -> floa
     return weight_decay_rate
 
 
-def get_param_dict(args: Any, model_without_ddp: nn.Module) -> list[dict[str, Any]]:
+def _hyperparameter_key(param_group: dict[str, Any]) -> tuple[tuple[str, str], ...]:
+    """Return the hyperparameter overrides of ``param_group``, without its parameters.
+
+    Values are compared by ``repr`` so a group carrying an unhashable setting (a third-party
+    optimizer's list-valued option, say) still yields a key. Distinct floats keep distinct ``repr``,
+    and anything without a value-based ``repr`` lands in its own bucket, which under-merges rather
+    than merging two differently configured parameters.
+
+    Args:
+        param_group: A single optimizer parameter group.
+
+    Returns:
+        The group's non-``params`` items, sorted by key so groups configured identically compare equal.
+    """
+    return tuple(sorted((key, repr(value)) for key, value in param_group.items() if key != "params"))
+
+
+def _merge_buckets(param_dicts: list[dict[str, Any]]) -> list[list[int]]:
+    """Bucket ``param_dicts`` indices by hyperparameter overrides, in first-appearance order.
+
+    Args:
+        param_dicts: Single-parameter groups as built by :func:`_build_param_dicts`.
+
+    Returns:
+        One list of ``param_dicts`` indices per distinct hyperparameter combination.
+    """
+    buckets: dict[tuple[tuple[str, str], ...], list[int]] = {}
+    ordered: list[list[int]] = []
+    for index, param_group in enumerate(param_dicts):
+        key = _hyperparameter_key(param_group)
+        bucket = buckets.get(key)
+        if bucket is None:
+            bucket = buckets[key] = []
+            ordered.append(bucket)
+        bucket.append(index)
+    return ordered
+
+
+def _build_param_dicts(args: Any, model_without_ddp: nn.Module) -> list[dict[str, Any]]:
+    """Build one single-parameter group per trainable parameter, with its LR/weight-decay overrides.
+
+    Args:
+        args: Namespace supplying the learning-rate and weight-decay knobs.
+        model_without_ddp: The model whose parameters the optimizer will own.
+
+    Returns:
+        One group per trainable parameter, ordered head/neck parameters first, then backbone, then
+        decoder. :func:`get_param_dict` merges these; the order is also the parameter order of
+        checkpoints written before that merge.
+    """
     assert isinstance(model_without_ddp.backbone, Joiner)
     backbone = cast("Any", model_without_ddp.backbone[0])
     backbone_named_param_lr_pairs = backbone.get_named_param_lr_pairs(args, prefix="backbone.0")
@@ -79,3 +128,109 @@ def get_param_dict(args: Any, model_without_ddp: nn.Module) -> list[dict[str, An
     final_param_dicts = other_param_dicts + backbone_param_lr_pairs + decoder_param_lr_pairs
 
     return final_param_dicts
+
+
+def get_param_dict(args: Any, model_without_ddp: nn.Module) -> list[dict[str, Any]]:
+    """Build optimizer parameter groups with layer-wise LR (and backbone weight-decay) overrides.
+
+    Parameters that end up configured identically share one group: ``torch.optim``'s foreach and
+    fused kernels batch a group's parameters into a single multi-tensor launch, so one group per
+    parameter would run ~500 single-tensor launches (and ~500 Python iterations of the optimizer's
+    per-group loop) per step instead of one launch per distinct configuration.
+
+    Args:
+        args: Namespace supplying ``lr``, ``lr_encoder``, ``lr_component_decay``,
+            ``lr_vit_layer_decay``, ``weight_decay``, and ``out_feature_indexes``.
+        model_without_ddp: The model whose parameters the optimizer will own.
+
+    Returns:
+        Optimizer parameter groups, each holding every trainable parameter that shares its
+        hyperparameter overrides.
+    """
+    param_dicts = _build_param_dicts(args, model_without_ddp)
+    return [
+        {
+            **{key: value for key, value in param_dicts[bucket[0]].items() if key != "params"},
+            "params": [param_dicts[index]["params"] for index in bucket],
+        }
+        for bucket in _merge_buckets(param_dicts)
+    ]
+
+
+def regroup_unmerged_optimizer_state(checkpoint: dict[str, Any]) -> None:
+    """Rewrite one-group-per-parameter optimizer/scheduler state onto the merged parameter groups.
+
+    :func:`get_param_dict` used to emit one parameter group per parameter, so ``torch.optim`` numbered
+    the saved optimizer state by each parameter's position in that layout, and every per-group
+    scheduler list (``base_lrs``, ``_last_lr``, ``lr_lambdas``, ``min_lrs``) had one entry per
+    parameter. Groups now hold every parameter sharing their hyperparameters — a group count
+    ``Optimizer.load_state_dict`` rejects outright — so reindex the saved state rather than fail the
+    resume.
+
+    Each optimizer's scheduler state is collapsed alongside it, matched by position the way
+    PyTorch Lightning stores the two lists.
+
+    The merged layout is derived from the saved groups themselves: bucketing them by the
+    hyperparameters they recorded, in the order they were saved, reproduces the buckets
+    :func:`get_param_dict` builds for the same run, since those are the same parameters in the same
+    order. State already saved in the merged layout has groups holding several parameters and is left
+    untouched.
+
+    Args:
+        checkpoint: Checkpoint dict carrying ``optimizer_states`` (and optionally ``lr_schedulers``),
+            mutated in-place.
+    """
+    scheduler_states = checkpoint.get("lr_schedulers") or []
+    for index, optimizer_state in enumerate(checkpoint.get("optimizer_states") or []):
+        saved_groups = optimizer_state.get("param_groups") or []
+        if not saved_groups or any(len(saved_group["params"]) != 1 for saved_group in saved_groups):
+            continue
+        buckets = _merge_buckets(saved_groups)
+        saved_state = optimizer_state.get("state", {})
+        merged_state: dict[int, Any] = {}
+        merged_groups: list[dict[str, Any]] = []
+        slot = 0
+        for bucket in buckets:
+            merged_group = {key: value for key, value in saved_groups[bucket[0]].items() if key != "params"}
+            slots = []
+            for unmerged_index in bucket:
+                saved_id = saved_groups[unmerged_index]["params"][0]
+                if saved_id in saved_state:
+                    merged_state[slot] = saved_state[saved_id]
+                slots.append(slot)
+                slot += 1
+            merged_group["params"] = slots
+            merged_groups.append(merged_group)
+        optimizer_state["state"] = merged_state
+        optimizer_state["param_groups"] = merged_groups
+        if index < len(scheduler_states):
+            _regroup_scheduler_lists(scheduler_states[index], len(saved_groups), buckets)
+        if len(merged_groups) < len(saved_groups):
+            logger.info(
+                "Regrouped resumed optimizer state from %d single-parameter groups onto %d merged groups.",
+                len(saved_groups),
+                len(merged_groups),
+            )
+
+
+def _regroup_scheduler_lists(scheduler_state: dict[str, Any], unmerged_count: int, buckets: list[list[int]]) -> None:
+    """Collapse a scheduler's per-parameter-group lists the way its optimizer's groups were collapsed.
+
+    A composite scheduler (``SequentialLR``, ``ChainedScheduler``) nests its wrapped schedulers' own
+    state dicts under a ``_schedulers`` list rather than holding per-group lists at the top level, so
+    those nested dicts need the same collapse applied recursively.
+
+    Args:
+        scheduler_state: Saved scheduler state, mutated in-place.
+        unmerged_count: Number of parameter groups the scheduler state was saved with.
+        buckets: Saved-group indices per merged group, as produced by :func:`_merge_buckets`.
+    """
+    # A bucket's parameters all shared one group's hyperparameters before the merge, so its first
+    # entry is the value the merged group inherits.
+    for key, value in scheduler_state.items():
+        if isinstance(value, list) and len(value) == unmerged_count:
+            scheduler_state[key] = [value[bucket[0]] for bucket in buckets]
+        elif key == "_schedulers" and isinstance(value, list):
+            for nested_state in value:
+                if isinstance(nested_state, dict):
+                    _regroup_scheduler_lists(nested_state, unmerged_count, buckets)

tests/training/test_detr_shim.py

@@ -1008,6 +1008,33 @@ def test_second_call_without_ema_leaves_first_stash(self, patch_lit):
         RFDETRModelModule.on_load_checkpoint(fake, {"state_dict": {}})
         assert fake._pending_legacy_ema_state is first_ema
 
+    def test_pre_merge_optimizer_state_is_regrouped(self, patch_lit):
+        """on_load_checkpoint itself regroups one-group-per-parameter optimizer state.
+
+        tests/training/test_param_groups.py exercises regroup_unmerged_optimizer_state() directly; this proves the hook
+        actually calls it, since a wiring mistake there (wrong argument, wrong order relative to the other normalisation
+        steps, a silently swallowed exception) would not be caught by calling the helper on its own.
+        """
+        fake = _FakeModule()
+        ckpt = {
+            "state_dict": {"model.w": torch.zeros(1)},
+            "optimizer_states": [
+                {
+                    "state": {0: {"exp_avg": torch.ones(1)}, 1: {"exp_avg": torch.full((1,), 2.0)}},
+                    "param_groups": [
+                        {"lr": 0.1, "weight_decay": 0.0, "params": [0]},
+                        {"lr": 0.1, "weight_decay": 0.0, "params": [1]},
+                    ],
+                }
+            ],
+        }
+
+        RFDETRModelModule.on_load_checkpoint(fake, ckpt)
+
+        groups = ckpt["optimizer_states"][0]["param_groups"]
+        assert len(groups) == 1
+        assert groups[0]["params"] == [0, 1]
+
 
 # ---------------------------------------------------------------------------
 # 5. Public API exports