From 91e0e5057bf65c905678d3bf32d32f6fcb9416da Mon Sep 17 00:00:00 2001
From: PinkTulips139 <242140955+PinkTulips139@users.noreply.github.com>
Date: Wed, 29 Jul 2026 08:55:37 +0800
Subject: [PATCH] feat(optim): add auditable parameter-group coverage

---
 tests/test_optimizer_group_audit.py | 313 ++++++++++++++++++++++++++++
 ultralytics/engine/trainer.py       |  43 +++-
 ultralytics/optim/__init__.py       |   3 +-
 ultralytics/optim/audit.py          | 266 +++++++++++++++++++++++
 4 files changed, 617 insertions(+), 8 deletions(-)
 create mode 100644 tests/test_optimizer_group_audit.py
 create mode 100644 ultralytics/optim/audit.py

diff --git a/tests/test_optimizer_group_audit.py b/tests/test_optimizer_group_audit.py
new file mode 100644
index 00000000..32f0eabc
--- /dev/null
+++ b/tests/test_optimizer_group_audit.py
@@ -0,0 +1,313 @@
+"""CPU-only tests for optimizer parameter-group auditing."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from ultralytics.engine.trainer import BaseTrainer
+from ultralytics.optim import OptimizerGroupAuditError, audit_optimizer_param_groups
+
+
+class _AuditModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(2, 3))
+        self.bias = nn.Parameter(torch.ones(2))
+        self.frozen = nn.Parameter(torch.ones(4), requires_grad=False)
+
+
+class _FakeOptimizer:
+    def __init__(self, param_groups, defaults=None):
+        self.param_groups = param_groups
+        self.defaults = defaults or {}
+
+
+def _optimizer(*groups, defaults=None):
+    return _FakeOptimizer(list(groups), defaults=defaults)
+
+
+def test_all_trainable_parameters_are_covered_exactly_once():
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight, model.bias], "group_name": "trainable"})
+
+    audit = audit_optimizer_param_groups(model, optimizer)
+
+    assert audit["trainable_coverage_complete"]
+    assert audit["trainable_coverage_exactly_once"]
+    assert not audit["has_missing_trainable"]
+    assert not audit["has_duplicates"]
+    assert audit["trainable_parameter_count"] == 2
+    assert audit["trainable_element_count"] == 8
+
+
+def test_missing_trainable_is_reported_without_strict_mode():
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight], "group_name": "partial"})
+
+    audit = audit_optimizer_param_groups(model, optimizer, strict=False)
+
+    assert audit["has_missing_trainable"]
+    assert not audit["trainable_coverage_complete"]
+    assert [item["name"] for item in audit["missing_trainable"]] == ["bias"]
+
+
+def test_missing_trainable_raises_in_strict_mode():
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight], "group_name": "partial"})
+
+    with pytest.raises(OptimizerGroupAuditError, match=r"missing_trainable=1.*bias"):
+        audit_optimizer_param_groups(model, optimizer, strict=True)
+
+
+def test_duplicate_across_groups_is_reported_and_strict():
+    model = _AuditModel()
+    optimizer = _optimizer(
+        {"params": [model.weight, model.bias], "group_name": "first"},
+        {"params": [model.weight], "group_name": "second"},
+    )
+
+    audit = audit_optimizer_param_groups(model, optimizer)
+
+    assert audit["has_duplicates"]
+    assert not audit["trainable_coverage_exactly_once"]
+    assert audit["duplicated"][0]["name"] == "weight"
+    assert [group["name"] for group in audit["duplicated"][0]["groups"]] == ["first", "second"]
+    with pytest.raises(OptimizerGroupAuditError, match=r"duplicated=1.*first,second"):
+        audit_optimizer_param_groups(model, optimizer, strict=True)
+
+
+def test_duplicate_within_one_group_is_reported():
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight, model.weight, model.bias], "group_name": "repeated"})
+
+    audit = audit_optimizer_param_groups(model, optimizer)
+
+    assert audit["duplicated_count"] == 1
+    assert audit["optimizer_parameter_occurrence_count"] == 3
+    assert [group["name"] for group in audit["duplicated"][0]["groups"]] == ["repeated", "repeated"]
+
+
+def test_frozen_parameter_is_reported_but_nonfatal_in_strict_mode():
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight, model.bias, model.frozen], "group_name": "mixed"})
+
+    audit = audit_optimizer_param_groups(model, optimizer, strict=True)
+
+    assert audit["has_frozen_in_optimizer"]
+    assert audit["frozen_in_optimizer_count"] == 1
+    assert audit["frozen_in_optimizer"][0]["name"] == "frozen"
+    assert audit["trainable_coverage_exactly_once"]
+
+
+def test_unknown_optimizer_parameter_is_reported_and_strict():
+    model = _AuditModel()
+    external = nn.Parameter(torch.ones(5))
+    optimizer = _optimizer({"params": [model.weight, model.bias, external], "group_name": "external"})
+
+    audit = audit_optimizer_param_groups(model, optimizer)
+
+    assert audit["has_unknown_parameters"]
+    assert audit["unknown_optimizer_parameter_count"] == 1
+    assert audit["unknown_optimizer_parameters"][0]["numel"] == 5
+    with pytest.raises(OptimizerGroupAuditError, match=r"unknown_optimizer_parameters=1.*external"):
+        audit_optimizer_param_groups(model, optimizer, strict=True)
+
+
+def test_group_tensor_and_element_counts_are_accurate():
+    model = _AuditModel()
+    optimizer = _optimizer(
+        {"params": [model.weight, model.frozen], "group_name": "mixed"},
+        {"params": [model.bias], "group_name": "bias"},
+    )
+
+    audit = audit_optimizer_param_groups(model, optimizer)
+
+    mixed, bias = audit["groups"]
+    assert (mixed["tensor_count"], mixed["total_element_count"]) == (2, 10)
+    assert (mixed["trainable_element_count"], mixed["frozen_element_count"]) == (6, 4)
+    assert (bias["tensor_count"], bias["total_element_count"]) == (1, 2)
+    assert audit["optimizer_unique_parameter_count"] == 3
+    assert audit["optimizer_unique_element_count"] == 12
+
+
+def test_group_hyperparameters_and_lr_scale_are_reported():
+    model = _AuditModel()
+    optimizer = _optimizer(
+        {
+            "params": [model.weight, model.bias],
+            "group_name": "explicit",
+            "lr": 0.02,
+            "initial_lr": 0.025,
+            "weight_decay": 0.001,
+        },
+        defaults={"lr": 0.01},
+    )
+
+    group = audit_optimizer_param_groups(model, optimizer)["groups"][0]
+
+    assert group["name"] == "explicit"
+    assert group["lr"] == pytest.approx(0.02)
+    assert group["initial_lr"] == pytest.approx(0.025)
+    assert group["weight_decay"] == pytest.approx(0.001)
+    assert group["base_lr"] == pytest.approx(0.01)
+    assert group["lr_scale"] == pytest.approx(2.0)
+    assert group["lr_scale_source"] == "lr/base_lr"
+
+
+def test_explicit_lr_scale_takes_precedence_over_inferred_scale():
+    model = _AuditModel()
+    optimizer = _optimizer(
+        {"params": [model.weight, model.bias], "lr": 0.02, "lr_scale": 3.0},
+        defaults={"lr": 0.01},
+    )
+
+    group = audit_optimizer_param_groups(model, optimizer)["groups"][0]
+
+    assert group["lr_scale"] == pytest.approx(3.0)
+    assert group["lr_scale_source"] == "group.lr_scale"
+
+
+@pytest.mark.parametrize(
+    ("metadata", "expected"),
+    [
+        ({"group_name": "named"}, "named"),
+        ({"name": "optimizer-name"}, "optimizer-name"),
+        ({"role": "role-name"}, "role-name"),
+        ({"param_group": "legacy-explicit"}, "legacy-explicit"),
+    ],
+)
+def test_explicit_group_semantics_are_preserved(metadata, expected):
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight, model.bias], **metadata})
+
+    assert audit_optimizer_param_groups(model, optimizer)["groups"][0]["name"] == expected
+
+
+def test_missing_group_semantics_use_index_fallback():
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight]}, {"params": [model.bias]})
+
+    groups = audit_optimizer_param_groups(model, optimizer)["groups"]
+
+    assert [group["name"] for group in groups] == ["group_0", "group_1"]
+    assert groups[0]["initial_lr"] is None
+    assert groups[0]["lr_scale"] is None
+
+
+def test_parameter_name_does_not_create_inferred_adapter_role():
+    class _OrdinaryModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.adapter_projection = nn.Parameter(torch.ones(2))
+
+    model = _OrdinaryModel()
+    optimizer = _optimizer({"params": [model.adapter_projection]})
+
+    audit = audit_optimizer_param_groups(model, optimizer)
+
+    assert audit["groups"][0]["name"] == "group_0"
+    assert audit["groups"][0]["parameter_names"] == ["adapter_projection"]
+
+
+def test_audit_result_is_json_serializable():
+    model = _AuditModel()
+    optimizer = _optimizer({"params": [model.weight, model.bias], "group_name": "ordinary"})
+
+    json.dumps(audit_optimizer_param_groups(model, optimizer))
+
+
+def test_scheduler_update_is_visible_in_a_new_snapshot():
+    model = nn.Linear(2, 1)
+    optimizer = torch.optim.SGD(
+        [{"params": list(model.parameters()), "group_name": "ordinary", "lr": 0.1}],
+        lr=0.1,
+    )
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda epoch: 1.0 if epoch == 0 else 0.5)
+    before = audit_optimizer_param_groups(model, optimizer)
+
+    optimizer.step()
+    scheduler.step()
+    after = audit_optimizer_param_groups(model, optimizer)
+
+    assert after["groups"][0]["lr"] == pytest.approx(optimizer.param_groups[0]["lr"])
+    assert after["groups"][0]["lr"] != before["groups"][0]["lr"]
+    assert after["groups"][0]["initial_lr"] == pytest.approx(0.1)
+
+
+def test_trainer_audit_is_read_only_for_peft_named_group():
+    model = nn.Linear(2, 1)
+    optimizer = torch.optim.AdamW(
+        [{"params": list(model.parameters()), "param_group": "adapter", "lr": 0.003, "weight_decay": 0.0}]
+    )
+    trainer = object.__new__(BaseTrainer)
+    trainer.model = model
+    trainer.optimizer = optimizer
+    before_groups = [
+        {
+            "parameter_ids": [id(parameter) for parameter in group["params"]],
+            "lr": group["lr"],
+            "weight_decay": group["weight_decay"],
+        }
+        for group in optimizer.param_groups
+    ]
+
+    audit = trainer._audit_optimizer_groups()
+
+    after_groups = [
+        {
+            "parameter_ids": [id(parameter) for parameter in group["params"]],
+            "lr": group["lr"],
+            "weight_decay": group["weight_decay"],
+        }
+        for group in optimizer.param_groups
+    ]
+    assert before_groups == after_groups
+    assert trainer.optimizer_group_audit is audit
+    assert audit["groups"][0]["name"] == "adapter"
+
+
+def test_train_pipeline_audits_after_adapter_configuration_and_before_scheduler():
+    trainer = object.__new__(BaseTrainer)
+    trainer.model = nn.Linear(2, 1)
+    trainer.batch_size = 2
+    trainer.world_size = 0
+    trainer.epochs = 1
+    trainer.data = {"train": "train", "val": "val"}
+    trainer.args = SimpleNamespace(
+        task="detect",
+        nbs=2,
+        weight_decay=0.001,
+        optimizer="AdamW",
+        lr0=0.01,
+        momentum=0.9,
+    )
+    loader = SimpleNamespace(dataset=[0, 1])
+    trainer.get_dataloader = lambda *args, **kwargs: loader
+    events = []
+
+    class _Controller:
+        @staticmethod
+        def prepare_optimizer(iterations):
+            events.append("prepare")
+
+        @staticmethod
+        def configure_optimizer(optimizer):
+            events.append("configure")
+
+    trainer.adapter_controller = _Controller()
+
+    def _build_optimizer(**kwargs):
+        events.append("build")
+        return torch.optim.AdamW(trainer.model.parameters(), lr=0.01)
+
+    trainer.build_optimizer = _build_optimizer
+    trainer._audit_optimizer_groups = lambda: events.append("audit")
+    trainer._save_run_args = lambda: events.append("save")
+    trainer._setup_scheduler = lambda: events.append("scheduler")
+
+    trainer._build_train_pipeline()
+
+    assert events == ["prepare", "build", "configure", "audit", "save", "scheduler"]
diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py
index 9e3309ab..9c98f22f 100644
--- a/ultralytics/engine/trainer.py
+++ b/ultralytics/engine/trainer.py
@@ -37,7 +37,7 @@
 from ultralytics.nn.distill_model import DistillationModel
 from ultralytics.nn.mixture_loss import has_routed_modules
 from ultralytics.nn.tasks import load_checkpoint
-from ultralytics.optim import MuSGD
+from ultralytics.optim import MuSGD, audit_optimizer_param_groups
 from ultralytics.utils import (
     DEFAULT_CFG,
     LOCAL_RANK,
@@ -135,12 +135,7 @@ def _optimizer_state_family(state_dict) -> str | None:
     groups = state_dict.get("param_groups", ())
     if any("use_muon" in group for group in groups):
         return "musgd"
-    state_keys = {
-        key
-        for state in state_dict.get("state", {}).values()
-        if isinstance(state, dict)
-        for key in state
-    }
+    state_keys = {key for state in state_dict.get("state", {}).values() if isinstance(state, dict) for key in state}
     if {"exp_avg", "exp_avg_sq"} <= state_keys:
         return "adam"
     if "square_avg" in state_keys:
@@ -365,6 +360,39 @@ def _setup_scheduler(self):
             self.lf = lambda x: max(1 - x / self.epochs, 0) * (1.0 - self.args.lrf) + self.args.lrf  # linear
         self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)
 
+    def _audit_optimizer_groups(self):
+        """Capture and log a read-only snapshot of finalized optimizer parameter groups."""
+        audit = audit_optimizer_param_groups(self.model, self.optimizer, strict=False)
+        self.optimizer_group_audit = audit
+        if RANK not in {-1, 0}:
+            return audit
+        group_summary = ", ".join(
+            f"{group['index']}:{group['name']}"
+            f"(tensors={group['tensor_count']}, elements={group['total_element_count']}, "
+            f"lr={group['lr']}, decay={group['weight_decay']})"
+            for group in audit["groups"]
+        )
+        LOGGER.info(
+            f"{colorstr('optimizer audit:')} groups={audit['group_count']}, "
+            f"coverage={'complete' if audit['trainable_coverage_complete'] else 'incomplete'}, "
+            f"missing={audit['missing_trainable_count']}, duplicates={audit['duplicated_count']}, "
+            f"frozen={audit['frozen_in_optimizer_count']}, "
+            f"unknown={audit['unknown_optimizer_parameter_count']}; {group_summary}"
+        )
+        issue_samples = []
+        for issue_name in (
+            "missing_trainable",
+            "duplicated",
+            "frozen_in_optimizer",
+            "unknown_optimizer_parameters",
+        ):
+            if audit[issue_name]:
+                names = ", ".join(item["name"] for item in audit[issue_name][:5])
+                issue_samples.append(f"{issue_name}=[{names}]")
+        if issue_samples:
+            LOGGER.warning(f"{colorstr('optimizer audit:')} " + "; ".join(issue_samples))
+        return audit
+
     def _setup_ddp(self):
         """Initialize and set the DistributedDataParallel parameters for training."""
         index = int(self.args.device.split(",")[LOCAL_RANK])  # world_size > 1 guarantees a multi-device string
@@ -404,6 +432,7 @@ def _build_train_pipeline(self):
             iterations=iterations,
         )
         self.adapter_controller.configure_optimizer(self.optimizer)
+        self._audit_optimizer_groups()
         self.args.effective_optimizer = type(self.optimizer).__name__
         self.args.effective_optimizer_lrs = [float(group["lr"]) for group in self.optimizer.param_groups]
         self._save_run_args()
diff --git a/ultralytics/optim/__init__.py b/ultralytics/optim/__init__.py
index 84751ed8..2e37f2f6 100644
--- a/ultralytics/optim/__init__.py
+++ b/ultralytics/optim/__init__.py
@@ -1,5 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from .audit import OptimizerGroupAuditError, audit_optimizer_param_groups
 from .muon import Muon, MuSGD
 
-__all__ = ["MuSGD", "Muon"]
+__all__ = ["MuSGD", "Muon", "OptimizerGroupAuditError", "audit_optimizer_param_groups"]
diff --git a/ultralytics/optim/audit.py b/ultralytics/optim/audit.py
new file mode 100644
index 00000000..80528c32
--- /dev/null
+++ b/ultralytics/optim/audit.py
@@ -0,0 +1,266 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+"""Read-only optimizer parameter-group auditing utilities."""
+
+from __future__ import annotations
+
+from numbers import Real
+from typing import Any
+
+
+class OptimizerGroupAuditError(ValueError):
+    """Raised when strict optimizer parameter-group validation fails."""
+
+    def __init__(self, message: str, audit: dict[str, Any]) -> None:
+        """Initialize the error and retain the serializable audit result."""
+        super().__init__(message)
+        self.audit = audit
+
+
+def _number_or_none(value: Any) -> float | None:
+    """Convert scalar optimizer metadata to a JSON-serializable float."""
+    if value is None or isinstance(value, bool):
+        return None
+    if isinstance(value, Real):
+        return float(value)
+    item = getattr(value, "item", None)
+    if callable(item):
+        try:
+            scalar = item()
+        except (RuntimeError, TypeError, ValueError):
+            return None
+        if isinstance(scalar, Real) and not isinstance(scalar, bool):
+            return float(scalar)
+    return None
+
+
+def _parameter_numel(parameter: Any) -> int:
+    """Return a parameter's element count without assuming a concrete tensor type."""
+    numel = getattr(parameter, "numel", None)
+    if not callable(numel):
+        return 0
+    try:
+        return int(numel())
+    except (RuntimeError, TypeError, ValueError):
+        return 0
+
+
+def _group_name(group: dict[str, Any], index: int) -> str:
+    """Return explicit group semantics without inferring roles from parameter names."""
+    # ``param_group`` is the existing Trainer field and is treated as a legacy
+    # explicit semantic label. No optimizer group is mutated by this helper.
+    for key in ("group_name", "name", "role", "param_group"):
+        value = group.get(key)
+        if value is not None and str(value).strip():
+            return str(value)
+    return f"group_{index}"
+
+
+def _parameter_record(
+    parameter_id: int,
+    model_parameters: dict[int, dict[str, Any]],
+    occurrences: dict[int, list[dict[str, Any]]],
+) -> dict[str, Any]:
+    """Build a serializable issue record for one parameter identity."""
+    model_record = model_parameters.get(parameter_id)
+    name = model_record["name"] if model_record is not None else f"<unknown:{parameter_id}>"
+    return {
+        "name": name,
+        "parameter_id": parameter_id,
+        "numel": model_record["numel"] if model_record is not None else 0,
+        "requires_grad": model_record["requires_grad"] if model_record is not None else None,
+        "groups": [
+            {
+                "index": occurrence["group_index"],
+                "name": occurrence["group_name"],
+                "position": occurrence["position"],
+            }
+            for occurrence in occurrences.get(parameter_id, [])
+        ],
+    }
+
+
+def _format_strict_error(audit: dict[str, Any], sample_size: int = 5) -> str:
+    """Format strict validation failures with short parameter and group samples."""
+    parts = []
+    for issue_name in ("missing_trainable", "duplicated", "unknown_optimizer_parameters"):
+        records = audit[issue_name]
+        if not records:
+            continue
+        samples = []
+        for record in records[:sample_size]:
+            groups = ",".join(group["name"] for group in record.get("groups", []))
+            samples.append(f"{record['name']}[{groups or 'no optimizer group'}]")
+        parts.append(f"{issue_name}={len(records)} ({'; '.join(samples)})")
+    return "Optimizer parameter-group audit failed: " + "; ".join(parts)
+
+
+def audit_optimizer_param_groups(model: Any, optimizer: Any, *, strict: bool = False) -> dict[str, Any]:
+    """Audit model-to-optimizer parameter coverage without mutating either object.
+
+    The returned dictionary is a point-in-time snapshot of the optimizer. Calling
+    the function again after scheduler or resume updates reports the then-current
+    ``param_groups`` values. Missing ``initial_lr`` metadata is represented as
+    ``None``.
+
+    Args:
+        model: Model exposing ``named_parameters()``.
+        optimizer: Optimizer exposing ``param_groups`` and optional ``defaults``.
+        strict: Raise for missing, duplicated, or unknown parameters. Frozen
+            parameters in the optimizer are always reported but remain non-fatal.
+
+    Returns:
+        JSON-serializable optimizer parameter-group audit dictionary.
+
+    Raises:
+        OptimizerGroupAuditError: If ``strict=True`` and a fatal coverage issue is
+            detected.
+    """
+    model_parameters: dict[int, dict[str, Any]] = {}
+    for name, parameter in model.named_parameters():
+        parameter_id = id(parameter)
+        if parameter_id not in model_parameters:
+            model_parameters[parameter_id] = {
+                "name": name,
+                "parameter_id": parameter_id,
+                "requires_grad": bool(getattr(parameter, "requires_grad", False)),
+                "numel": _parameter_numel(parameter),
+            }
+
+    optimizer_defaults = getattr(optimizer, "defaults", {}) or {}
+    default_lr = _number_or_none(optimizer_defaults.get("lr")) if isinstance(optimizer_defaults, dict) else None
+    occurrences: dict[int, list[dict[str, Any]]] = {}
+    optimizer_parameters: dict[int, Any] = {}
+    group_summaries = []
+    optimizer_parameter_occurrence_count = 0
+
+    for group_index, group in enumerate(getattr(optimizer, "param_groups", ())):
+        semantic_name = _group_name(group, group_index)
+        parameters = list(group.get("params", ()))
+        parameter_names = []
+        total_element_count = 0
+        trainable_element_count = 0
+        frozen_element_count = 0
+
+        for position, parameter in enumerate(parameters):
+            parameter_id = id(parameter)
+            optimizer_parameters.setdefault(parameter_id, parameter)
+            optimizer_parameter_occurrence_count += 1
+            occurrences.setdefault(parameter_id, []).append(
+                {"group_index": group_index, "group_name": semantic_name, "position": position}
+            )
+            model_record = model_parameters.get(parameter_id)
+            parameter_names.append(model_record["name"] if model_record is not None else f"<unknown:{parameter_id}>")
+            numel = _parameter_numel(parameter)
+            total_element_count += numel
+            if bool(getattr(parameter, "requires_grad", False)):
+                trainable_element_count += numel
+            else:
+                frozen_element_count += numel
+
+        lr = _number_or_none(group.get("lr"))
+        initial_lr = _number_or_none(group.get("initial_lr"))
+        weight_decay = _number_or_none(group.get("weight_decay"))
+        explicit_lr_scale = _number_or_none(group.get("lr_scale"))
+        base_lr = _number_or_none(group.get("base_lr"))
+        if base_lr is None:
+            base_lr = default_lr
+        if explicit_lr_scale is not None:
+            lr_scale = explicit_lr_scale
+            lr_scale_source = "group.lr_scale"
+        elif lr is not None and base_lr not in {None, 0.0}:
+            lr_scale = lr / base_lr
+            lr_scale_source = "lr/base_lr"
+        else:
+            lr_scale = None
+            lr_scale_source = None
+
+        group_summaries.append(
+            {
+                "index": group_index,
+                "name": semantic_name,
+                "tensor_count": len(parameters),
+                "total_element_count": total_element_count,
+                "trainable_element_count": trainable_element_count,
+                "frozen_element_count": frozen_element_count,
+                "lr": lr,
+                "initial_lr": initial_lr,
+                "weight_decay": weight_decay,
+                "base_lr": base_lr,
+                "lr_scale": lr_scale,
+                "lr_scale_source": lr_scale_source,
+                "parameter_names": parameter_names,
+            }
+        )
+
+    trainable_ids = {parameter_id for parameter_id, record in model_parameters.items() if record["requires_grad"]}
+    optimizer_ids = set(optimizer_parameters)
+    missing_ids = sorted(trainable_ids - optimizer_ids, key=lambda item: model_parameters[item]["name"])
+    duplicated_ids = sorted(
+        (parameter_id for parameter_id, locations in occurrences.items() if len(locations) > 1),
+        key=lambda item: model_parameters.get(item, {}).get("name", f"<unknown:{item}>"),
+    )
+    frozen_ids = sorted(
+        (
+            parameter_id
+            for parameter_id in optimizer_ids & set(model_parameters)
+            if not model_parameters[parameter_id]["requires_grad"]
+        ),
+        key=lambda item: model_parameters[item]["name"],
+    )
+    unknown_ids = sorted(optimizer_ids - set(model_parameters))
+
+    missing_trainable = [
+        {
+            **model_parameters[parameter_id],
+            "groups": [],
+        }
+        for parameter_id in missing_ids
+    ]
+    duplicated = [_parameter_record(parameter_id, model_parameters, occurrences) for parameter_id in duplicated_ids]
+    frozen_in_optimizer = [
+        _parameter_record(parameter_id, model_parameters, occurrences) for parameter_id in frozen_ids
+    ]
+    unknown_optimizer_parameters = [
+        {
+            **_parameter_record(parameter_id, model_parameters, occurrences),
+            "numel": _parameter_numel(optimizer_parameters[parameter_id]),
+            "requires_grad": bool(getattr(optimizer_parameters[parameter_id], "requires_grad", False)),
+        }
+        for parameter_id in unknown_ids
+    ]
+
+    duplicated_trainable_ids = trainable_ids & set(duplicated_ids)
+    audit = {
+        "group_count": len(group_summaries),
+        "groups": group_summaries,
+        "model_parameters": list(model_parameters.values()),
+        "missing_trainable": missing_trainable,
+        "duplicated": duplicated,
+        "frozen_in_optimizer": frozen_in_optimizer,
+        "unknown_optimizer_parameters": unknown_optimizer_parameters,
+        "missing_trainable_count": len(missing_trainable),
+        "duplicated_count": len(duplicated),
+        "frozen_in_optimizer_count": len(frozen_in_optimizer),
+        "unknown_optimizer_parameter_count": len(unknown_optimizer_parameters),
+        "trainable_parameter_count": len(trainable_ids),
+        "trainable_element_count": sum(model_parameters[item]["numel"] for item in trainable_ids),
+        "optimizer_parameter_occurrence_count": optimizer_parameter_occurrence_count,
+        "optimizer_unique_parameter_count": len(optimizer_ids),
+        "optimizer_unique_element_count": sum(
+            _parameter_numel(parameter) for parameter in optimizer_parameters.values()
+        ),
+        "trainable_coverage_complete": not missing_trainable,
+        "trainable_coverage_exactly_once": not missing_trainable and not duplicated_trainable_ids,
+        "has_duplicates": bool(duplicated),
+        "has_missing_trainable": bool(missing_trainable),
+        "has_frozen_in_optimizer": bool(frozen_in_optimizer),
+        "has_unknown_parameters": bool(unknown_optimizer_parameters),
+    }
+
+    if strict and (missing_trainable or duplicated or unknown_optimizer_parameters):
+        raise OptimizerGroupAuditError(_format_strict_error(audit), audit)
+    return audit
+
+
+__all__ = ["OptimizerGroupAuditError", "audit_optimizer_param_groups"]