chore: import upstream snapshot with attribution

tests/training/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/training/callbacks/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/training/callbacks/test_drop_path_callback.py

@@ -0,0 +1,293 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for :class:`rfdetr.training.callbacks.drop_schedule.DropPathCallback`."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from rfdetr.training.callbacks.drop_schedule import DropPathCallback
+from rfdetr.training.drop_schedule import drop_scheduler
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_trainer(global_step: int = 0, estimated_stepping_batches: int = 50) -> MagicMock:
+    """Create a minimal mock Trainer with controllable step metadata."""
+    trainer = MagicMock()
+    trainer.global_step = global_step
+    trainer.estimated_stepping_batches = estimated_stepping_batches
+    return trainer
+
+
+def _make_mock_pl_module(epochs: int = 5) -> MagicMock:
+    """Create a minimal mock RFDETRModule with ``train_config.epochs``."""
+    pl_module = MagicMock()
+    pl_module.train_config.epochs = epochs
+    return pl_module
+
+
+# ---------------------------------------------------------------------------
+# TestDropPathCallbackInit
+# ---------------------------------------------------------------------------
+
+
+class TestDropPathCallbackInit:
+    """Verify constructor defaults."""
+
+    def test_default_args(self) -> None:
+        """Default rates are zero and vit_encoder_num_layers is 12."""
+        cb = DropPathCallback()
+        assert cb._drop_path == 0.0
+        assert cb._dropout == 0.0
+        assert cb._vit_encoder_num_layers == 12
+        assert cb._dp_schedule is None
+        assert cb._do_schedule is None
+
+
+# ---------------------------------------------------------------------------
+# TestOnTrainStart
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainStart:
+    """Verify schedule arrays built in ``on_train_start``."""
+
+    def test_dp_schedule_matches_drop_scheduler_standard(self) -> None:
+        """drop_path schedule matches ``drop_scheduler`` for standard mode."""
+        cb = DropPathCallback(drop_path=0.3)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10)
+        assert cb._dp_schedule is not None
+        np.testing.assert_array_equal(cb._dp_schedule, expected)
+
+    def test_do_schedule_matches_drop_scheduler_standard(self) -> None:
+        """Dropout schedule matches ``drop_scheduler`` for standard mode."""
+        cb = DropPathCallback(dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.1, 5, 10)
+        assert cb._do_schedule is not None
+        np.testing.assert_array_equal(cb._do_schedule, expected)
+
+    def test_no_dp_schedule_when_rate_zero(self) -> None:
+        """drop_path=0.0 leaves ``_dp_schedule`` as None."""
+        cb = DropPathCallback(drop_path=0.0)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        assert cb._dp_schedule is None
+
+    def test_dp_schedule_early_mode(self) -> None:
+        """Early mode: rates at step 0 and step 30 match ``drop_scheduler``."""
+        cb = DropPathCallback(drop_path=0.3, cutoff_epoch=2, mode="early")
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10, 2, "early")
+        assert cb._dp_schedule is not None
+        assert cb._dp_schedule[0] == expected[0]
+        assert cb._dp_schedule[30] == expected[30]
+
+    def test_dp_schedule_late_mode(self) -> None:
+        """Late mode: rates at step 0 and step 30 match ``drop_scheduler``."""
+        cb = DropPathCallback(drop_path=0.3, cutoff_epoch=2, mode="late")
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10, 2, "late")
+        assert cb._dp_schedule is not None
+        assert cb._dp_schedule[0] == expected[0]
+        assert cb._dp_schedule[30] == expected[30]
+
+
+# ---------------------------------------------------------------------------
+# TestOnTrainBatchStart
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainBatchStart:
+    """Verify model update calls in ``on_train_batch_start``."""
+
+    def test_update_drop_path_called_with_correct_rate(self) -> None:
+        """``update_drop_path`` is called with the schedule value at step 0."""
+        cb = DropPathCallback(drop_path=0.3, vit_encoder_num_layers=6)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 0
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._dp_schedule is not None
+        pl_module.model.update_drop_path.assert_called_once_with(cb._dp_schedule[0], 6)
+
+    def test_update_dropout_called_with_correct_rate(self) -> None:
+        """``update_dropout`` is called with the schedule value at step 0."""
+        cb = DropPathCallback(dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 0
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._do_schedule is not None
+        pl_module.model.update_dropout.assert_called_once_with(cb._do_schedule[0])
+
+    def test_no_update_when_step_out_of_bounds(self) -> None:
+        """No model updates when ``global_step`` exceeds schedule length."""
+        cb = DropPathCallback(drop_path=0.3, dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 9999
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        pl_module.model.update_drop_path.assert_not_called()
+        pl_module.model.update_dropout.assert_not_called()
+
+    @pytest.mark.parametrize(
+        "step",
+        [
+            pytest.param(0, id="first_step"),
+            pytest.param(5, id="mid_step"),
+            pytest.param(9, id="last_of_first_epoch"),
+        ],
+    )
+    def test_drop_rates_at_multiple_steps_match_schedule(self, step: int) -> None:
+        """Each step uses the correct value from the pre-computed schedule."""
+        cb = DropPathCallback(drop_path=0.3, vit_encoder_num_layers=6)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = step
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._dp_schedule is not None
+        pl_module.model.update_drop_path.assert_called_once_with(cb._dp_schedule[step], 6)
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerValidation
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerValidation:
+    """Verify drop_scheduler raises for invalid inputs."""
+
+    @pytest.mark.parametrize(
+        "cutoff_epoch",
+        [
+            pytest.param(6, id="above_epochs"),
+            pytest.param(-1, id="negative"),
+        ],
+    )
+    def test_raises_for_invalid_cutoff_epoch(self, cutoff_epoch: int) -> None:
+        """drop_scheduler raises ValueError when cutoff_epoch is outside [0, epochs]."""
+        with pytest.raises(ValueError, match="cutoff_epoch must be in"):
+            drop_scheduler(0.3, 5, 10, cutoff_epoch=cutoff_epoch, mode="early")
+
+    @pytest.mark.parametrize(
+        ("epochs", "niter_per_ep", "match"),
+        [
+            pytest.param(0, 10, "epochs must be >= 1", id="epochs_zero"),
+            pytest.param(5, 0, "niter_per_ep must be >= 1", id="niter_per_ep_zero"),
+        ],
+    )
+    def test_raises_for_invalid_epoch_counts(self, epochs: int, niter_per_ep: int, match: str) -> None:
+        """drop_scheduler raises ValueError when epochs or niter_per_ep is less than 1."""
+        with pytest.raises(ValueError, match=match):
+            drop_scheduler(0.3, epochs, niter_per_ep)
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerBoundary
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerBoundary:
+    """Verify drop_scheduler with cutoff_epoch at the inclusive boundaries 0 and epochs."""
+
+    @pytest.mark.parametrize(
+        ("cutoff_epoch", "mode", "expected_first", "expected_last"),
+        [
+            pytest.param(0, "early", 0.0, 0.0, id="early_cutoff_zero_all_zeros"),
+            pytest.param(5, "early", 0.3, 0.3, id="early_cutoff_full_all_rate"),
+            pytest.param(0, "late", 0.3, 0.3, id="late_cutoff_zero_all_rate"),
+            pytest.param(5, "late", 0.0, 0.0, id="late_cutoff_full_all_zeros"),
+        ],
+    )
+    def test_boundary_cutoff_epoch(
+        self,
+        cutoff_epoch: int,
+        mode: str,
+        expected_first: float,
+        expected_last: float,
+    ) -> None:
+        """Boundary cutoff_epoch values (0 and epochs) produce correct first and last rates."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=cutoff_epoch, mode=mode)
+        assert schedule[0] == expected_first
+        assert schedule[-1] == expected_last
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerLinear
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerLinear:
+    """Verify drop_scheduler with schedule='linear' in early mode."""
+
+    def test_linear_early_starts_at_drop_rate(self) -> None:
+        """Linear early schedule first value equals drop_rate."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert schedule[0] == pytest.approx(0.3, abs=1e-9)
+
+    def test_linear_early_ends_early_phase_at_zero(self) -> None:
+        """Linear early schedule last value of the early phase equals 0."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert schedule[19] == pytest.approx(0.0, abs=1e-9)
+
+    def test_linear_early_late_phase_is_zero(self) -> None:
+        """Linear early schedule: all values after cutoff_epoch are zero."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        np.testing.assert_array_equal(schedule[20:], 0.0)
+
+    def test_linear_early_decreases_monotonically(self) -> None:
+        """Linear early schedule values decrease monotonically during the early phase."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert np.all(np.diff(schedule[:20]) <= 0)
+
+    def test_linear_same_shape_as_constant(self) -> None:
+        """Schedule='linear' output has the same length as schedule='constant'."""
+        linear = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        constant = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="constant")
+        assert linear.shape == constant.shape

tests/training/callbacks/test_ema_callback.py

@@ -0,0 +1,260 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit and parity tests for RFDETREMACallback."""
+
+from __future__ import annotations
+
+import math
+import warnings
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from torch import nn
+from torch.optim.swa_utils import AveragedModel
+
+from rfdetr.training.callbacks.ema import RFDETREMACallback
+from rfdetr.training.model_ema import ModelEma
+
+
+class _EMAContainerModule(nn.Module):
+    """Minimal module with `.model` to mirror RFDETRModelModule shape."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.model = nn.Linear(4, 2)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+
+class TestAvgFnDecayFormula:
+    """Verify the tau / no-tau decay formula matches ModelEma."""
+
+    @pytest.mark.parametrize(
+        "num_averaged",
+        [
+            pytest.param(0, id="step-0"),
+            pytest.param(5, id="step-5"),
+            pytest.param(99, id="step-99"),
+        ],
+    )
+    def test_tau_zero_uses_fixed_decay(self, num_averaged: int) -> None:
+        """With tau=0 the effective decay equals the base decay at every step."""
+        decay = 0.99
+        cb = RFDETREMACallback(decay=decay, tau=0)
+        ema_val = torch.tensor(1.0)
+        model_val = torch.tensor(2.0)
+
+        result = cb._avg_fn(ema_val, model_val, num_averaged)
+
+        expected = ema_val * decay + model_val * (1.0 - decay)
+        assert torch.allclose(result, expected, atol=1e-7)
+
+    def test_tau_warmup_at_step_1(self) -> None:
+        """At the first call (num_averaged=0) with tau>0 the effective decay uses updates=1 matching ModelEma's
+        1-indexed counter."""
+        decay = 0.993
+        tau = 100
+        cb = RFDETREMACallback(decay=decay, tau=tau)
+        ema_val = torch.tensor(1.0)
+        model_val = torch.tensor(2.0)
+
+        result = cb._avg_fn(ema_val, model_val, num_averaged=0)
+
+        updates = 1  # num_averaged + 1
+        effective_decay = decay * (1 - math.exp(-updates / tau))
+        expected = ema_val * effective_decay + model_val * (1.0 - effective_decay)
+        assert torch.allclose(result, expected, atol=1e-7)
+
+
+class TestModelEmaParity:
+    """Ensure N-step EMA weights match ModelEma exactly."""
+
+    def test_avg_fn_matches_modelema_weight_parity(self) -> None:
+        """Simulate 500 update steps and compare final EMA weights with ModelEma.module to confirm numerical parity."""
+        torch.manual_seed(42)
+        n_steps = 500
+        decay = 0.993
+        tau = 100
+
+        model = nn.Linear(4, 4)
+        model_ema = ModelEma(model, decay=decay, tau=tau)
+        cb = RFDETREMACallback(decay=decay, tau=tau)
+
+        # Initialise manual EMA state from model (same as ModelEma deepcopy)
+        ema_weights: dict[str, torch.Tensor] = {name: p.clone() for name, p in model.named_parameters()}
+
+        for step in range(n_steps):
+            # Perturb model parameters
+            with torch.no_grad():
+                for p in model.parameters():
+                    p.add_(torch.randn_like(p) * 0.01)
+
+            # Update legacy ModelEma
+            model_ema.update(model)
+
+            # Replicate update via callback avg_fn
+            model_weights = {name: p.clone() for name, p in model.named_parameters()}
+            for name in ema_weights:
+                ema_weights[name] = cb._avg_fn(ema_weights[name], model_weights[name], step)
+
+        # Compare
+        legacy_state = dict(model_ema.module.named_parameters())
+        for name, cb_val in ema_weights.items():
+            assert torch.allclose(cb_val, legacy_state[name], atol=1e-5), (
+                f"Parity failed for {name}: max diff = {(cb_val - legacy_state[name]).abs().max().item()}"
+            )
+
+
+class TestShouldUpdate:
+    """Verify should_update triggers on steps and epochs."""
+
+    def test_should_update_on_step(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update(step_idx=42) is True
+
+    def test_should_update_on_epoch(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update(epoch_idx=3) is True
+
+    def test_should_update_neither(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update() is False
+
+
+class TestInit:
+    """Construction and EMA-state access behavior."""
+
+    def test_init_emits_no_user_warning(self) -> None:
+        """Instantiation should not emit runtime UserWarnings."""
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            RFDETREMACallback()
+        user_warns = [w for w in caught if issubclass(w.category, UserWarning)]
+        assert not user_warns
+
+    def test_get_ema_model_state_dict_none_before_setup(self) -> None:
+        """EMA state accessor returns None before averaged model is created."""
+        cb = RFDETREMACallback()
+        assert cb.get_ema_model_state_dict() is None
+
+    def test_get_ema_model_state_dict_returns_model_weights(self) -> None:
+        """EMA state accessor returns the wrapped `.model` state dict."""
+
+        class _Container(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.model = nn.Linear(4, 2)
+
+        cb = RFDETREMACallback()
+        container = _Container()
+        cb._average_model = AveragedModel(container, avg_fn=cb._avg_fn)
+
+        state = cb.get_ema_model_state_dict()
+
+        assert state is not None
+        assert "weight" in state
+        assert "bias" in state
+
+
+class TestUpdateInterval:
+    """Verify update_interval_steps throttles EMA updates on step hooks."""
+
+    def test_updates_only_on_interval_steps(self) -> None:
+        """update_interval_steps=2 updates on steps 2, 4, ...
+
+        only.
+        """
+        cb = RFDETREMACallback(update_interval_steps=2)
+        cb._average_model = MagicMock()
+
+        trainer = MagicMock()
+        pl_module = MagicMock()
+
+        for step in (1, 2, 3, 4):
+            trainer.global_step = step
+            cb.on_train_batch_end(trainer, pl_module, outputs=None, batch=None, batch_idx=step - 1)
+
+        assert cb._average_model.update_parameters.call_count == 2
+
+
+class TestLegacyEMAResume:
+    """Legacy checkpoint EMA payload is consumed by the callback setup path."""
+
+    def test_setup_loads_pending_legacy_ema_state_into_average_model(self) -> None:
+        """`_pending_legacy_ema_state` must initialize EMA weights at fit setup."""
+        cb = RFDETREMACallback()
+        pl_module = _EMAContainerModule()
+        trainer = MagicMock()
+
+        legacy_ema_state = {k: torch.full_like(v, 2.0) for k, v in pl_module.model.state_dict().items()}
+        pl_module._pending_legacy_ema_state = legacy_ema_state
+
+        cb.setup(trainer, pl_module, stage="fit")
+
+        assert cb._average_model is not None
+        restored = cb._average_model.module.model.state_dict()
+        for key, expected in legacy_ema_state.items():
+            assert torch.allclose(restored[key], expected)
+        assert not hasattr(pl_module, "_pending_legacy_ema_state")
+
+
+class TestSuppressTestSwap:
+    """suppress_test_swap must disable the test-time EMA weight swap while leaving defaults unchanged."""
+
+    @staticmethod
+    def _make_swap_scenario() -> tuple[RFDETREMACallback, _EMAContainerModule]:
+        """Build a module at weight 7.0 with an EMA average model captured at weight 5.0."""
+        cb = RFDETREMACallback()
+        pl_module = _EMAContainerModule()
+        with torch.no_grad():
+            for p in pl_module.parameters():
+                p.fill_(5.0)
+        cb._average_model = AveragedModel(model=pl_module, use_buffers=True, avg_fn=cb._avg_fn)
+        with torch.no_grad():
+            for p in pl_module.parameters():
+                p.fill_(7.0)
+        return cb, pl_module
+
+    def test_default_flag_is_false(self) -> None:
+        """The suppression flag defaults to False so standalone trainer.test() keeps EMA evaluation."""
+        cb = RFDETREMACallback()
+        assert cb.suppress_test_swap is False
+
+    def test_on_test_epoch_start_swaps_by_default(self) -> None:
+        """Without suppression, the test hooks swap live weights (7.0) for EMA weights (5.0)."""
+        cb, pl_module = self._make_swap_scenario()
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 5.0))
+
+    def test_on_test_epoch_start_suppressed_keeps_live_weights(self) -> None:
+        """With suppress_test_swap=True the live weights (7.0) must stay in place during test."""
+        cb, pl_module = self._make_swap_scenario()
+        cb.suppress_test_swap = True
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 7.0))
+
+    def test_on_test_epoch_end_suppressed_does_not_swap(self) -> None:
+        """With suppression active, on_test_epoch_end must not swap EMA weights in unpaired."""
+        cb, pl_module = self._make_swap_scenario()
+        cb.suppress_test_swap = True
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+        cb.on_test_epoch_end(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 7.0))

tests/training/conftest.py

@@ -0,0 +1,132 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Package-level pytest fixtures for tests/training/.
+
+Provides cross-test cleanup that prevents class-level state from leaking between individual tests in the training/ test
+package, plus shared config factory fixtures used across multiple test modules.
+"""
+
+import pytest
+
+from rfdetr.config import RFDETRBaseConfig, SegmentationTrainConfig, TrainConfig
+
+# ---------------------------------------------------------------------------
+# Shared config factory fixtures (used by test_module, test_datamodule,
+# test_args — avoids duplicate fixture definitions across files)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def base_model_config():
+    """Factory fixture — call with **overrides to get a minimal RFDETRBaseConfig."""
+
+    def _make(**overrides):
+        defaults = dict(pretrain_weights=None, device="cpu", num_classes=5)
+        defaults.update(overrides)
+        return RFDETRBaseConfig(**defaults)
+
+    return _make
+
+
+@pytest.fixture
+def base_train_config(tmp_path):
+    """Factory fixture — call with **overrides to get a minimal TrainConfig.
+
+    tmp_path is injected automatically so test methods do not need to declare it.
+    """
+
+    def _make(**overrides):
+        defaults = dict(
+            dataset_dir=str(tmp_path / "dataset"),
+            output_dir=str(tmp_path / "output"),
+            epochs=10,
+            lr=1e-4,
+            lr_encoder=1.5e-4,
+            batch_size=2,
+            weight_decay=1e-4,
+            lr_drop=8,
+            warmup_epochs=1.0,
+            drop_path=0.0,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            grad_accum_steps=1,
+            num_workers=0,
+            tensorboard=False,
+        )
+        defaults.update(overrides)
+        return TrainConfig(**defaults)
+
+    return _make
+
+
+@pytest.fixture
+def seg_train_config(tmp_path):
+    """Factory fixture — call with **overrides to get a minimal SegmentationTrainConfig.
+
+    tmp_path is injected automatically so test methods do not need to declare it.
+    """
+
+    def _make(**overrides):
+        defaults = dict(
+            dataset_dir=str(tmp_path / "dataset"),
+            output_dir=str(tmp_path / "output"),
+            epochs=10,
+            batch_size=2,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            grad_accum_steps=1,
+            drop_path=0.0,
+            num_workers=0,
+            tensorboard=False,
+        )
+        defaults.update(overrides)
+        return SegmentationTrainConfig(**defaults)
+
+    return _make
+
+
+# ---------------------------------------------------------------------------
+# Class-level isolation
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def _restore_rfdetr_module_trainer_property():
+    """Restore RFDETRModelModule.trainer to the LightningModule parent property after each test.
+
+    Several unit tests in test_module_model.py patch the ``trainer`` property directly on the ``RFDETRModelModule``
+    class (``type(module).trainer = property(...)``). Without cleanup this mutates the class for the remainder of the
+    session and breaks ``Trainer.fit()`` calls in smoke tests (PTL cannot set ``.trainer`` on the module because the
+    patched property has no setter).
+
+    This fixture deletes any class-level override from ``RFDETRModelModule.__dict__`` after every test, so the next test
+    starts with a clean class that inherits PTL's read/write ``trainer`` descriptor from ``LightningModule``.
+    """
+    yield
+    # Lazy import so the fixture does not force module import at collection time.
+    from rfdetr.training.module_model import RFDETRModelModule
+
+    if "trainer" in RFDETRModelModule.__dict__:
+        delattr(RFDETRModelModule, "trainer")
+
+
+@pytest.fixture(autouse=True)
+def _restore_rfdetr_datamodule_trainer_property():
+    """Restore RFDETRDataModule.trainer to the LightningDataModule parent property after each test.
+
+    Tests that mock the ``trainer`` property on ``RFDETRDataModule`` (e.g. for ``on_after_batch_transfer`` tests) patch
+    it at the class level.  Without cleanup this mutates the class for the remainder of the session.
+
+    This fixture deletes any class-level override from ``RFDETRDataModule.__dict__`` after every test, mirroring the
+    ``_restore_rfdetr_module_trainer_property`` pattern above.
+    """
+    yield
+    from rfdetr.training.module_data import RFDETRDataModule
+
+    if "trainer" in RFDETRDataModule.__dict__:
+        delattr(RFDETRDataModule, "trainer")

tests/training/helpers.py

@@ -0,0 +1,126 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Shared test helpers for the rfdetr.training test suite.
+
+Plain classes and functions (not pytest fixtures) shared across multiple test modules to avoid verbatim duplication.
+Import with a relative import::
+
+    from .helpers import _FakeCriterion, _FakeDataset, _TinyModel
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.utils.data
+
+
+class _TinyModel(nn.Module):
+    """Minimal real nn.Module satisfying the RFDETRModule model contract.
+
+    Has a single trainable parameter so the optimizer has something to update and the loss has a gradient path back
+    through the model.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.dummy = nn.Parameter(torch.zeros(1))
+
+    def forward(self, samples, targets=None):
+        return {"dummy": self.dummy}
+
+    def update_drop_path(self, *args, **kwargs) -> None:
+        pass
+
+    def update_dropout(self, *args, **kwargs) -> None:
+        pass
+
+    def reinitialize_detection_head(self, *args, **kwargs) -> None:
+        pass
+
+
+class _FakeCriterion:
+    """Callable criterion that returns a loss connected to the model output.
+
+    Keeps a gradient path from the loss back to _TinyModel.dummy so that ``loss.backward()`` does not error when the
+    Trainer calls it.
+    """
+
+    weight_dict = {"loss_ce": 1.0}
+
+    def num_boxes_for_targets(self, outputs, targets):
+        dummy = outputs.get("dummy", torch.zeros(1))
+        return torch.ones((), dtype=dummy.dtype, device=dummy.device)
+
+    def __call__(self, outputs, targets, num_boxes=None):
+        dummy = outputs.get("dummy", torch.zeros(1))
+        denominator = self.num_boxes_for_targets(outputs, targets) if num_boxes is None else num_boxes
+        return {"loss_ce": dummy.mean() / denominator}
+
+
+class _FakeDataset(torch.utils.data.Dataset):
+    """Dataset with ``(image, target)`` pairs for detection.
+
+    The image is a ``(3, 32, 32)`` float tensor; the target dict includes the fields expected by RFDETRModule:
+    ``boxes``, ``labels``, ``image_id``, ``orig_size``, ``size``.
+    """
+
+    def __init__(self, length: int = 20) -> None:
+        self._length = length
+
+    def __len__(self) -> int:
+        return self._length
+
+    def __getitem__(self, idx):
+        image = torch.randn(3, 32, 32)
+        target = {
+            "boxes": torch.tensor([[0.5, 0.5, 0.1, 0.1]]),
+            "labels": torch.tensor([1]),
+            "image_id": torch.tensor(idx),
+            "orig_size": torch.tensor([32, 32]),
+            "size": torch.tensor([32, 32]),
+        }
+        return image, target
+
+
+class _FakeDatasetWithMasks(_FakeDataset):
+    """Like _FakeDataset but includes binary instance masks (for segmentation)."""
+
+    def __getitem__(self, idx):
+        image, target = super().__getitem__(idx)
+        target["masks"] = torch.zeros(1, 32, 32, dtype=torch.bool)
+        return image, target
+
+
+class _FakePostProcess:
+    """Picklable postprocessor for ddp_spawn tests.
+
+    ``MagicMock`` is not picklable and cannot survive the subprocess boundary that ``ddp_spawn`` creates.  This plain
+    class is a drop-in replacement.
+
+    Delegates to ``_fake_postprocess``; keep both in sync if the fake output format changes.
+    """
+
+    def __call__(self, outputs, orig_sizes):
+        return _fake_postprocess(outputs, orig_sizes)
+
+
+def _fake_postprocess(outputs, orig_sizes):
+    """Return one non-empty prediction per image so COCOEvalCallback has something to score."""
+    n = orig_sizes.shape[0]
+    return [
+        {
+            "boxes": torch.tensor([[5.0, 5.0, 20.0, 20.0]]),
+            "scores": torch.tensor([0.9]),
+            "labels": torch.tensor([1]),
+        }
+        for _ in range(n)
+    ]
+
+
+def _make_param_dicts(model: nn.Module) -> list[dict]:
+    """Build a minimal param-dict list for AdamW from all trainable parameters."""
+    return [{"params": p, "lr": 1e-4} for p in model.parameters() if p.requires_grad]

tests/training/test_args.py

@@ -0,0 +1,156 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for _namespace_from_configs — the canonical config-to-Namespace mapping."""
+
+import pytest
+
+from rfdetr._namespace import _namespace_from_configs
+
+
+class TestNamespaceFromConfigs:
+    """_namespace_from_configs maps ModelConfig + TrainConfig to a Namespace."""
+
+    def test_forwards_model_config_fields(self, base_model_config, base_train_config):
+        """All key ModelConfig fields are faithfully mapped."""
+        mc = base_model_config(num_classes=7)
+        args = _namespace_from_configs(mc, base_train_config())
+
+        assert args.encoder == mc.encoder
+        assert args.num_classes == 7
+        assert args.hidden_dim == mc.hidden_dim
+        assert args.resolution == mc.resolution
+        assert args.patch_size == mc.patch_size
+        assert args.num_windows == mc.num_windows
+        assert args.segmentation_head == mc.segmentation_head
+        assert args.positional_encoding_size == mc.positional_encoding_size
+
+    def test_forwards_train_config_fields(self, base_model_config, base_train_config):
+        """All key TrainConfig fields are faithfully mapped."""
+        tc = base_train_config(
+            lr=3e-4,
+            epochs=20,
+            weight_decay=5e-5,
+            batch_size=4,
+            num_workers=0,
+            eval_interval=3,
+            log_per_class_metrics=False,
+            train_log_sync_dist=True,
+            train_log_on_step=True,
+            compute_val_loss=False,
+            compute_test_loss=False,
+            ema_update_interval=2,
+            prefetch_factor=4,
+        )
+        args = _namespace_from_configs(base_model_config(), tc)
+
+        assert args.lr == pytest.approx(3e-4)
+        assert args.epochs == 20
+        assert args.weight_decay == pytest.approx(5e-5)
+        assert args.batch_size == 4
+        assert args.num_workers == 0
+        assert args.eval_interval == 3
+        assert args.log_per_class_metrics is False
+        assert args.train_log_sync_dist is True
+        assert args.train_log_on_step is True
+        assert args.compute_val_loss is False
+        assert args.compute_test_loss is False
+        assert args.ema_update_interval == 2
+        assert args.prefetch_factor == 4
+
+    def test_forwards_promoted_train_fields(self, base_model_config, base_train_config):
+        """Promoted TrainConfig fields are forwarded to the namespace."""
+        tc = base_train_config(clip_max_norm=0.35, seed=123, sync_bn=True, fp16_eval=True)
+        args = _namespace_from_configs(base_model_config(), tc)
+
+        assert args.clip_max_norm == pytest.approx(0.35)
+        assert args.seed == 123
+        assert args.sync_bn is True
+        assert args.fp16_eval is True
+
+    def test_seed_falls_back_to_legacy_default_when_unset(self, base_model_config, base_train_config):
+        """Seed defaults to 42 in the namespace when TrainConfig.seed is None."""
+        tc = base_train_config(seed=None)
+        args = _namespace_from_configs(base_model_config(), tc)
+        assert args.seed == 42
+
+    def test_forwards_dataset_fields(self, base_model_config, base_train_config):
+        """Dataset-routing fields are forwarded to the Namespace."""
+        tc = base_train_config(multi_scale=True, expanded_scales=True, dataset_file="coco")
+        args = _namespace_from_configs(base_model_config(), tc)
+
+        assert args.multi_scale is True
+        assert args.expanded_scales is True
+        assert args.dataset_file == "coco"
+
+    def test_num_queries_from_subclass_config(self, base_model_config, base_train_config):
+        """num_queries is read from subclass config attributes."""
+        mc = base_model_config()  # RFDETRBaseConfig has num_queries=300
+        args = _namespace_from_configs(mc, base_train_config())
+        assert args.num_queries == 300
+
+    def test_resume_none_becomes_empty_string(self, base_model_config, base_train_config):
+        """Resume=None (the default) is converted to '' for the Namespace."""
+        tc = base_train_config()
+        assert tc.resume is None
+        args = _namespace_from_configs(base_model_config(), tc)
+        assert args.resume == ""
+
+    def test_segmentation_extras_forwarded_from_seg_config(self, base_model_config, seg_train_config):
+        """SegmentationTrainConfig mask loss coefficients are forwarded."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config()
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.mask_ce_loss_coef == pytest.approx(5.0)
+        assert args.mask_dice_loss_coef == pytest.approx(5.0)
+
+    def test_segmentation_cls_loss_default_matches_pre_1_7_effective_weight(self, base_model_config, seg_train_config):
+        """Default segmentation classification loss weight must stay at the pre-1.7 effective value."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config()
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.cls_loss_coef == pytest.approx(1.0)
+
+    def test_segmentation_cls_loss_explicit_override_is_forwarded(self, base_model_config, seg_train_config):
+        """Explicit segmentation classification loss weight overrides are preserved."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config(cls_loss_coef=5.0)
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.cls_loss_coef == pytest.approx(5.0)
+
+    def test_segmentation_num_select_none_falls_back_to_model_config(self, base_model_config, seg_train_config) -> None:
+        """SegmentationTrainConfig(num_select=None) must not overwrite ModelConfig.num_select."""
+        mc = base_model_config(segmentation_head=True, num_select=200)
+        # Explicitly passing num_select=None triggers the deprecation warning (Item #3).
+        with pytest.warns(DeprecationWarning, match="TrainConfig.num_select is deprecated"):
+            tc = seg_train_config(num_select=None)
+
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.num_select == 200
+
+    def test_segmentation_extras_default_for_plain_config(self, base_model_config, base_train_config):
+        """mask_* attributes default to 5.0 for a plain TrainConfig (not segmentation)."""
+        args = _namespace_from_configs(base_model_config(), base_train_config())
+        assert args.mask_ce_loss_coef == pytest.approx(5.0)
+        assert args.mask_dice_loss_coef == pytest.approx(5.0)
+
+    def test_segmentation_head_flag_forwarded(self, base_model_config, base_train_config):
+        """segmentation_head=True from ModelConfig reaches the Namespace."""
+        mc = base_model_config(segmentation_head=True)
+        args = _namespace_from_configs(mc, base_train_config())
+        assert args.segmentation_head is True
+
+    def test_build_namespace_emits_deprecation_warning(
+        self, base_model_config, base_train_config, reset_build_namespace_warning_state
+    ):
+        """build_namespace() must emit a DeprecationWarning on every call."""
+        from rfdetr._namespace import build_namespace
+
+        with pytest.warns(FutureWarning, match="build_namespace"):
+            build_namespace(base_model_config(), base_train_config())

tests/training/test_auto_batch.py

@@ -0,0 +1,228 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from rfdetr.detr import RFDETR
+from rfdetr.training import auto_batch
+from rfdetr.training.auto_batch import AutoBatchResult
+
+
+class _TinyModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.w = torch.nn.Parameter(torch.ones(1))
+
+
+def test_recommend_grad_accum_steps_rounds_up():
+    assert auto_batch.recommend_grad_accum_steps(3, 16) == 6
+
+
+def test_probe_max_micro_batch_uses_exponential_then_binary_search():
+    model = _TinyModule()
+    criterion = _TinyModule()
+    threshold = 7
+
+    def _fake_probe(*args, **kwargs):
+        micro_batch_size = args[2]
+        return micro_batch_size <= threshold
+
+    with (
+        patch("rfdetr.training.auto_batch._probe_step", side_effect=_fake_probe),
+        patch("rfdetr.training.auto_batch.torch.cuda.empty_cache"),
+    ):
+        safe = auto_batch.probe_max_micro_batch(
+            model=model,
+            criterion=criterion,
+            resolution=64,
+            device=torch.device("cuda"),
+            num_classes=5,
+            amp=False,
+            safety_margin=1.0,
+            max_micro_batch=32,
+        )
+    assert safe == threshold
+
+
+def test_probe_max_micro_batch_raises_if_one_is_not_safe():
+    model = _TinyModule()
+    criterion = _TinyModule()
+
+    with (
+        patch("rfdetr.training.auto_batch._probe_step", return_value=False),
+        patch("rfdetr.training.auto_batch.torch.cuda.empty_cache"),
+        pytest.raises(RuntimeError, match="micro_batch_size=1"),
+    ):
+        auto_batch.probe_max_micro_batch(
+            model=model,
+            criterion=criterion,
+            resolution=64,
+            device=torch.device("cuda"),
+            num_classes=5,
+            amp=False,
+        )
+
+
+def test_probe_step_raises_when_loss_keys_do_not_overlap_weight_keys():
+    """_probe_step must fail fast when weighted loss would be empty."""
+
+    class _DummyCriterion(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight_dict = {"loss_bbox": 1.0}
+
+        def forward(self, outputs, targets):
+            return {"loss_ce": torch.tensor(1.0)}
+
+    class _DummyModel(torch.nn.Module):
+        def forward(self, samples, targets):
+            return {}
+
+    model = _DummyModel()
+    criterion = _DummyCriterion()
+
+    with (
+        patch(
+            "rfdetr.training.auto_batch._make_synthetic_batch",
+            return_value=(MagicMock(), []),
+        ),
+        pytest.raises(RuntimeError, match="no overlap between criterion loss_dict and weight_dict keys"),
+    ):
+        auto_batch._probe_step(
+            model=model,
+            criterion=criterion,
+            micro_batch_size=1,
+            resolution=64,
+            device=torch.device("cpu"),
+            num_classes=5,
+            amp=False,
+        )
+
+
+def test_resolve_auto_batch_config_requires_cuda():
+    model_context = SimpleNamespace(device=torch.device("cpu"), model=MagicMock())
+    model_config = SimpleNamespace(resolution=64, num_classes=5, amp=False, segmentation_head=False)
+    train_config = SimpleNamespace(batch_size="auto", auto_batch_target_effective=16)
+
+    with (
+        patch("rfdetr.training.auto_batch.torch.cuda.is_available", return_value=False),
+        pytest.raises(RuntimeError, match="requires a CUDA device"),
+    ):
+        auto_batch.resolve_auto_batch_config(model_context, model_config, train_config)
+
+
+def test_resolve_auto_batch_config_returns_expected_values():
+    model_context = SimpleNamespace(device=torch.device("cuda"), model=MagicMock())
+    model_config = SimpleNamespace(resolution=64, num_classes=5, amp=False, segmentation_head=True)
+    train_config = SimpleNamespace(batch_size="auto", auto_batch_target_effective=16)
+    criterion = MagicMock()
+    criterion.to.return_value = criterion
+
+    with (
+        patch("rfdetr.training.auto_batch.torch.cuda.is_available", return_value=True),
+        patch("rfdetr.training.auto_batch.build_criterion_from_config", return_value=(criterion, None)),
+        patch("rfdetr.training.auto_batch.probe_max_micro_batch", return_value=5),
+        patch("rfdetr.training.auto_batch.torch.cuda.get_device_name", return_value="Fake GPU"),
+    ):
+        result = auto_batch.resolve_auto_batch_config(model_context, model_config, train_config)
+
+    assert isinstance(result, AutoBatchResult)
+    assert result.safe_micro_batch == 5
+    assert result.recommended_grad_accum_steps == 4
+    assert result.effective_batch_size == 20
+    assert result.device_name == "Fake GPU"
+
+
+@patch("rfdetr.detr.is_main_process", return_value=False)
+@patch("rfdetr.training.auto_batch.resolve_auto_batch_config")
+@patch("rfdetr.training.build_trainer")
+@patch("rfdetr.training.RFDETRDataModule")
+@patch("rfdetr.training.RFDETRModelModule")
+@patch("rfdetr.detr._move_model_context_to_device")
+def test_train_auto_batch_ensures_model_on_device_before_resolve(
+    mock_move: MagicMock,
+    _mock_module: MagicMock,
+    _mock_data_module: MagicMock,
+    _mock_build_trainer: MagicMock,
+    mock_resolve: MagicMock,
+    _mock_is_main: MagicMock,
+) -> None:
+    """Model weights must be moved before resolve_auto_batch_config when batch_size='auto'."""
+    auto_result = SimpleNamespace(safe_micro_batch=4, recommended_grad_accum_steps=1, effective_batch_size=4)
+    call_order: list[str] = []
+
+    def _move_side_effect(model: object) -> None:
+        call_order.append("ensure")
+
+    def _resolve_side_effect(**_kwargs: object) -> object:
+        call_order.append("resolve")
+        return auto_result
+
+    mock_move.side_effect = _move_side_effect
+    mock_resolve.side_effect = _resolve_side_effect
+
+    train_config = SimpleNamespace(
+        batch_size="auto",
+        grad_accum_steps=99,
+        dataset_dir=None,
+        resume=None,
+        class_names=None,
+        save_dataset_grids=False,
+    )
+    mock_self = MagicMock()
+    mock_self.model_config = SimpleNamespace(model_name=None)
+    mock_self.get_train_config.return_value = train_config
+
+    RFDETR.train(mock_self)
+
+    assert train_config.batch_size == 4
+    assert train_config.grad_accum_steps == 1
+    mock_move.assert_called_once_with(mock_self.model)
+    mock_resolve.assert_called_once_with(
+        model_context=mock_self.model,
+        model_config=mock_self.model_config,
+        train_config=train_config,
+    )
+    assert call_order == ["ensure", "resolve"]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for segmentation probe")
+def test_probe_step_with_real_segmentation_criterion(tmp_path):
+    """Run one probe step with real segmentation model and criterion so loss_masks and t['masks'] are exercised."""
+    from rfdetr._namespace import _namespace_from_configs
+    from rfdetr.config import RFDETRSegNanoConfig, SegmentationTrainConfig
+    from rfdetr.models.lwdetr import build_criterion_and_postprocessors, build_model
+
+    mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cuda", num_classes=2)
+    tc = SegmentationTrainConfig(
+        dataset_dir=str(tmp_path / "ds"),
+        output_dir=str(tmp_path / "out"),
+        batch_size=2,
+        grad_accum_steps=1,
+        tensorboard=False,
+    )
+    args = _namespace_from_configs(mc, tc)
+    model = build_model(args)
+    criterion, _ = build_criterion_and_postprocessors(args)
+    device = torch.device("cuda")
+    model = model.to(device)
+    criterion = criterion.to(device)
+
+    ok = auto_batch._probe_step(
+        model=model,
+        criterion=criterion,
+        micro_batch_size=1,
+        resolution=mc.resolution,
+        device=device,
+        num_classes=mc.num_classes,
+        amp=False,
+        segmentation_head=True,
+    )
+    assert ok is True

tests/training/test_cli.py

@@ -0,0 +1,117 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for RFDETRCli — PTL Ch4/T4.
+
+Verifies that the CLI module is correctly structured: importable, subclasses LightningCLI, overrides
+add_arguments_to_parser, and exposes a callable main() entry point.  CLI integration / smoke tests (--help subprocess,
+YAML roundtrip) live in T4-7.
+"""
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Structure and importability
+# ---------------------------------------------------------------------------
+
+
+class TestRFDETRCliStructure:
+    """RFDETRCli is correctly structured and importable."""
+
+    def test_cli_module_importable(self):
+        """rfdetr.training.cli imports without error."""
+        import rfdetr.training.cli  # noqa: F401
+
+    def test_rfdetr_cli_importable(self):
+        """RFDETRCli can be imported from rfdetr.training.cli."""
+        from rfdetr.training.cli import RFDETRCli  # noqa: F401
+
+    def test_main_importable(self):
+        """Main() can be imported from rfdetr.training.cli."""
+        from rfdetr.training.cli import main  # noqa: F401
+
+    def test_rfdetr_cli_is_lightning_cli_subclass(self):
+        """RFDETRCli must subclass pytorch_lightning LightningCLI."""
+        from pytorch_lightning.cli import LightningCLI
+
+        from rfdetr.training.cli import RFDETRCli
+
+        assert issubclass(RFDETRCli, LightningCLI)
+
+    def test_main_is_callable(self):
+        """Main must be a callable (function, not e.g. a string)."""
+        from rfdetr.training.cli import main
+
+        assert callable(main)
+
+    def test_add_arguments_to_parser_is_overridden(self):
+        """RFDETRCli overrides add_arguments_to_parser from LightningCLI."""
+        from pytorch_lightning.cli import LightningCLI
+
+        from rfdetr.training.cli import RFDETRCli
+
+        assert RFDETRCli.add_arguments_to_parser is not LightningCLI.add_arguments_to_parser
+
+    def test_exported_from_lit_package(self):
+        """RFDETRCli is exported from rfdetr.training (appears in __all__)."""
+        import rfdetr.training as lit
+
+        assert hasattr(lit, "RFDETRCli")
+        assert "RFDETRCli" in lit.__all__
+
+
+# ---------------------------------------------------------------------------
+# Argument linking
+# ---------------------------------------------------------------------------
+
+
+class TestRFDETRCliArgumentLinking:
+    """add_arguments_to_parser registers the expected argument links."""
+
+    def _collect_links(self):
+        """Instantiate a minimal parser and collect registered link sources."""
+        import unittest.mock as mock
+
+        from rfdetr.training.cli import RFDETRCli
+
+        captured = []
+
+        class _FakeParser:
+            def link_arguments(self, source, target, **kwargs):
+                captured.append({"source": source, "target": target, **kwargs})
+
+            # LightningArgumentParser methods that may be called during setup
+            def __getattr__(self, name):
+                return mock.MagicMock()
+
+        cli = RFDETRCli.__new__(RFDETRCli)
+        cli.add_arguments_to_parser(_FakeParser())
+        return captured
+
+    def test_model_config_link_registered(self):
+        """model.model_config is linked to data.model_config."""
+        links = self._collect_links()
+        sources = [lnk["source"] for lnk in links]
+        assert "model.model_config" in sources
+
+    def test_train_config_link_registered(self):
+        """model.train_config is linked to data.train_config."""
+        links = self._collect_links()
+        sources = [lnk["source"] for lnk in links]
+        assert "model.train_config" in sources
+
+    @pytest.mark.parametrize(
+        "source, expected_target",
+        [
+            pytest.param("model.model_config", "data.model_config", id="model_config"),
+            pytest.param("model.train_config", "data.train_config", id="train_config"),
+        ],
+    )
+    def test_link_target(self, source, expected_target):
+        """Each link points to the correct data.* target."""
+        links = self._collect_links()
+        match = next((lnk for lnk in links if lnk["source"] == source), None)
+        assert match is not None, f"No link registered for source {source!r}"
+        assert match["target"] == expected_target

tests/training/test_drop_path.py

@@ -0,0 +1,203 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import logging
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+import torch.nn as nn
+
+import rfdetr.models.backbone.dinov2 as dinov2_module
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import RFDETRNanoConfig, TrainConfig
+from rfdetr.models import build_model
+from rfdetr.models.backbone.dinov2 import DinoV2
+from rfdetr.models.backbone.dinov2_with_windowed_attn import (
+    Dinov2WithRegistersDropPath,
+    WindowedDinov2WithRegistersBackbone,
+)
+from rfdetr.models.lwdetr import LWDETR
+
+
+@pytest.fixture
+def model_with_drop_path(monkeypatch: pytest.MonkeyPatch) -> LWDETR:
+    """Create RF-DETR Nano LWDETR with drop_path enabled."""
+    monkeypatch.setattr(
+        WindowedDinov2WithRegistersBackbone,
+        "from_pretrained",
+        classmethod(lambda cls, name, config: cls(config)),
+    )
+    mc = RFDETRNanoConfig(num_classes=3, pretrain_weights=None)
+    tc = TrainConfig(
+        dataset_dir=".",
+        output_dir=".",
+        drop_path=0.1,
+    )
+    args = _namespace_from_configs(mc, tc)
+    return build_model(args)
+
+
+@pytest.fixture
+def model_without_drop_path(monkeypatch: pytest.MonkeyPatch) -> LWDETR:
+    """Create RF-DETR Nano LWDETR without drop_path."""
+    monkeypatch.setattr(
+        WindowedDinov2WithRegistersBackbone,
+        "from_pretrained",
+        classmethod(lambda cls, name, config: cls(config)),
+    )
+    mc = RFDETRNanoConfig(num_classes=3, pretrain_weights=None)
+    tc = TrainConfig(
+        dataset_dir=".",
+        output_dir=".",
+        drop_path=0.0,
+    )
+    args = _namespace_from_configs(mc, tc)
+    return build_model(args)
+
+
+def test_get_backbone_encoder_layers_dinov2(model_with_drop_path: LWDETR) -> None:
+    """Verify _get_backbone_encoder_layers() returns encoder.encoder.layer for DinoV2."""
+    model = model_with_drop_path
+
+    layers = model._get_backbone_encoder_layers()
+    assert layers is not None
+
+    enc = model.backbone[0].encoder
+    assert hasattr(enc, "encoder"), "DinoV2 encoder should have encoder attribute"
+    assert hasattr(enc.encoder, "encoder"), "DinoV2 encoder.encoder should have encoder attribute"
+    assert hasattr(enc.encoder.encoder, "layer"), "DinoV2 encoder.encoder.encoder should have layer attribute"
+    assert layers is enc.encoder.encoder.layer, "Should return encoder.encoder.encoder.layer"
+
+    assert len(layers) > 0, "Should have at least one layer"
+    for layer in layers:
+        assert hasattr(layer, "drop_path"), "Each layer should have drop_path attribute"
+
+
+def test_update_drop_path_dinov2(model_with_drop_path: LWDETR) -> None:
+    """Verify update_drop_path() sets drop_prob values correctly with linear schedule."""
+    model = model_with_drop_path
+
+    layers = model._get_backbone_encoder_layers()
+    assert layers is not None
+
+    num_layers = len(layers)
+    drop_path_rate = 0.1
+
+    model.update_drop_path(drop_path_rate, num_layers)
+
+    # All layers must be Dinov2WithRegistersDropPath (drop_path_rate=0.1 > 0 at model build time).
+    expected_rates = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]
+    for i, layer in enumerate(layers):
+        assert isinstance(layer.drop_path, Dinov2WithRegistersDropPath), (
+            f"Layer {i} drop_path should be Dinov2WithRegistersDropPath, got {type(layer.drop_path)}"
+        )
+        actual_prob = layer.drop_path.drop_prob
+        assert abs(actual_prob - expected_rates[i]) < 1e-6, (
+            f"Layer {i} drop_prob should be {expected_rates[i]}, got {actual_prob}"
+        )
+
+    assert abs(layers[0].drop_path.drop_prob - 0.0) < 1e-6, "First layer should have drop_prob = 0"
+    assert abs(layers[-1].drop_path.drop_prob - drop_path_rate) < 1e-6, (
+        f"Last layer should have drop_prob = {drop_path_rate}"
+    )
+
+
+def test_drop_path_initialization(model_with_drop_path: LWDETR, model_without_drop_path: LWDETR) -> None:
+    """Verify drop_path initialization: Dinov2WithRegistersDropPath vs Identity based on rate."""
+    layers_with_dp = model_with_drop_path._get_backbone_encoder_layers()
+    layers_without_dp = model_without_drop_path._get_backbone_encoder_layers()
+
+    assert layers_with_dp is not None
+    assert layers_without_dp is not None
+
+    # drop_path_rate=0.1 -> every layer initialised as Dinov2WithRegistersDropPath
+    for i, layer in enumerate(layers_with_dp):
+        assert hasattr(layer, "drop_path"), "Layer should have drop_path attribute"
+        assert isinstance(layer.drop_path, Dinov2WithRegistersDropPath), (
+            f"Layer {i}: expected Dinov2WithRegistersDropPath, got {type(layer.drop_path)}"
+        )
+
+    # drop_path_rate=0.0 -> every layer initialised as nn.Identity
+    for i, layer in enumerate(layers_without_dp):
+        assert hasattr(layer, "drop_path"), "Layer should have drop_path attribute"
+        assert isinstance(layer.drop_path, torch.nn.Identity), (
+            f"Layer {i}: expected nn.Identity for zero drop_path, got {type(layer.drop_path)}"
+        )
+
+
+def test_update_drop_path_handles_missing_layers(model_with_drop_path: LWDETR, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Verify update_drop_path() handles models without recognizable layer structure gracefully."""
+    model = model_with_drop_path
+
+    monkeypatch.setattr(model, "_get_backbone_encoder_layers", lambda: None)
+
+    # Should not raise an error, just return early
+    model.update_drop_path(0.1, 12)
+
+
+def test_update_drop_path_partial_layers(model_with_drop_path: LWDETR) -> None:
+    """Verify min() guard prevents IndexError when vit_encoder_num_layers > len(layers)."""
+    model = model_with_drop_path
+
+    layers = model._get_backbone_encoder_layers()
+    assert layers is not None
+    actual_num_layers = len(layers)
+
+    # Request more layers than exist in the backbone
+    requested_num_layers = actual_num_layers + 4
+    drop_path_rate = 0.2
+
+    # Should not raise IndexError
+    model.update_drop_path(drop_path_rate, requested_num_layers)
+
+    # Each updated layer gets a rate from 0 to drop_path_rate (shorter, capped linspace)
+    expected_rates = [x.item() for x in torch.linspace(0, drop_path_rate, actual_num_layers)]
+    for i in range(actual_num_layers):
+        actual_prob = layers[i].drop_path.drop_prob
+        assert abs(actual_prob - expected_rates[i]) < 1e-6, (
+            f"Layer {i} drop_prob should be {expected_rates[i]}, got {actual_prob}"
+        )
+
+
+def test_non_windowed_drop_path_warns(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Verify a warning is emitted when drop_path_rate > 0 with non-windowed backbone."""
+    mock_backbone = MagicMock()
+    monkeypatch.setattr(dinov2_module, "AutoBackbone", MagicMock(from_pretrained=MagicMock(return_value=mock_backbone)))
+
+    # The rf-detr logger sets propagate=False, so intercept warning() directly.
+    warning_messages: list[str] = []
+    rf_detr_logger = logging.getLogger("rf-detr")
+    monkeypatch.setattr(rf_detr_logger, "warning", lambda msg, *args, **kwargs: warning_messages.append(msg))
+
+    DinoV2(size="base", use_windowed_attn=False, drop_path_rate=0.1)
+
+    assert any("drop_path_rate" in msg and "ignored" in msg for msg in warning_messages), (
+        "Expected warning about drop_path_rate being ignored for non-windowed backbone"
+    )
+
+
+def test_get_backbone_encoder_layers_blocks_path() -> None:
+    """Verify _get_backbone_encoder_layers() returns enc.blocks for standard ViT backbones."""
+    mock_blocks = nn.ModuleList([nn.Linear(1, 1) for _ in range(3)])
+    # SimpleNamespace gives only the attributes we define, so hasattr checks work correctly.
+    mock_encoder = SimpleNamespace(blocks=mock_blocks)
+    mock_self = SimpleNamespace(backbone=[SimpleNamespace(encoder=mock_encoder)])
+
+    result = LWDETR._get_backbone_encoder_layers(mock_self)  # type: ignore[arg-type]
+    assert result is mock_blocks, "Should return enc.blocks for standard ViT backbone"
+
+
+def test_get_backbone_encoder_layers_trunk_blocks_path() -> None:
+    """Verify _get_backbone_encoder_layers() returns enc.trunk.blocks for aimv2 backbones."""
+    mock_blocks = nn.ModuleList([nn.Linear(1, 1) for _ in range(3)])
+    mock_trunk = SimpleNamespace(blocks=mock_blocks)
+    mock_encoder = SimpleNamespace(trunk=mock_trunk)  # no 'blocks' at top level
+    mock_self = SimpleNamespace(backbone=[SimpleNamespace(encoder=mock_encoder)])
+
+    result = LWDETR._get_backbone_encoder_layers(mock_self)  # type: ignore[arg-type]
+    assert result is mock_blocks, "Should return enc.trunk.blocks for aimv2 backbone"

tests/training/test_load_pretrain_weights.py

@@ -0,0 +1,870 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for fine-tuned checkpoint weight destruction.
+
+When a user loads a fine-tuned N-class checkpoint but has ``num_classes`` configured to a LARGER value (e.g. default
+90), the second reinit in ``load_pretrain_weights`` (models/weights.py) must NOT erroneously resize the detection head
+to ``num_classes + 1``, destroying the loaded weights.
+
+The fix changes the second reinit condition from:
+    ``checkpoint_num_classes != args.num_classes + 1``
+to the user-override-aware logic that auto-aligns to the checkpoint when the user did not explicitly set
+``num_classes``.
+
+These tests exercise ``rfdetr.models.weights.load_pretrain_weights`` directly, which is the unified function that
+replaced the two prior separate implementations (``detr.py:_load_pretrain_weights_into`` and
+``module_model.py:RFDETRModelModule._load_pretrain_weights``).
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, call
+
+import pytest
+import torch
+
+from rfdetr.config import (
+    RFDETRBaseConfig,
+    RFDETRKeypointPreviewConfig,
+    RFDETRLargeConfig,
+    RFDETRMediumConfig,
+    RFDETRNanoConfig,
+    RFDETRSeg2XLargeConfig,
+    RFDETRSegLargeConfig,
+    RFDETRSegMediumConfig,
+    RFDETRSegNanoConfig,
+    RFDETRSegPreviewConfig,
+    RFDETRSegSmallConfig,
+    RFDETRSegXLargeConfig,
+    RFDETRSmallConfig,
+    TrainConfig,
+)
+from rfdetr.models.weights import load_pretrain_weights
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_checkpoint(num_classes=91, num_queries=300, group_detr=13):
+    """Build a minimal checkpoint dict with the given class count.
+
+    Args:
+        num_classes: Total classes including background (bias shape).
+        num_queries: Number of object queries per group.
+        group_detr: Number of groups.
+    """
+    total_queries = num_queries * group_detr
+    state = {
+        "class_embed.weight": torch.randn(num_classes, 256),
+        "class_embed.bias": torch.randn(num_classes),
+        "refpoint_embed.weight": torch.randn(total_queries, 4),
+        "query_feat.weight": torch.randn(total_queries, 256),
+        "other_layer.weight": torch.randn(10, 10),
+    }
+    ckpt_args = SimpleNamespace(
+        segmentation_head=False,
+        patch_size=14,
+        class_names=[],
+    )
+    return {"model": state, "args": ckpt_args}
+
+
+def _make_train_config():
+    """Return a minimal TrainConfig for use in load_pretrain_weights.
+
+    Returns:
+        Minimal TrainConfig with placeholder dataset and output dirs.
+    """
+    return TrainConfig(
+        dataset_dir="/nonexistent/dataset",
+        output_dir="/nonexistent/output",
+        epochs=10,
+        lr=1e-4,
+        lr_encoder=1.5e-4,
+        batch_size=2,
+        weight_decay=1e-4,
+        lr_drop=8,
+        warmup_epochs=1.0,
+        drop_path=0.0,
+        multi_scale=False,
+        expanded_scales=False,
+        do_random_resize_via_padding=False,
+        grad_accum_steps=1,
+        tensorboard=False,
+    )
+
+
+def _suppress_pretrain_io(monkeypatch) -> None:
+    """Suppress download/validate/file-existence side effects on the canonical load path."""
+    monkeypatch.setattr("rfdetr.models.weights.download_pretrain_weights", lambda *a, **kw: None)
+    monkeypatch.setattr("rfdetr.models.weights.validate_pretrain_weights", lambda *a, **kw: None)
+    monkeypatch.setattr("rfdetr.models.weights.validate_checkpoint_compatibility", lambda *a, **kw: None)
+    monkeypatch.setattr("rfdetr.models.weights.os.path.isfile", lambda _: True)
+
+
+# ---------------------------------------------------------------------------
+# Regression tests: load_pretrain_weights (models/weights.py)
+# ---------------------------------------------------------------------------
+
+
+class TestLoadPretrainWeightsSecondReinit:
+    """Regression tests for ``load_pretrain_weights`` in ``rfdetr.models.weights``.
+
+    Validates that the second reinitialize_detection_head call only fires when the checkpoint has MORE classes than
+    configured (backbone pretrain scenario), not when it has fewer (fine-tuned checkpoint scenario).
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_download(self, monkeypatch):
+        """Suppress all download and file-existence side effects."""
+        _suppress_pretrain_io(monkeypatch)
+
+    def test_finetune_checkpoint_preserves_weights(self, monkeypatch):
+        """Fine-tuned checkpoint (fewer classes) must NOT trigger second reinit.
+
+        Scenario: 2-class fine-tuned checkpoint (bias shape [3]) loaded with
+        default num_classes=90. The first reinit correctly resizes the head to 3 so load_state_dict works. The second
+        reinit must NOT resize to 91 — that would destroy the loaded fine-tuned weights.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        checkpoint = _make_checkpoint(num_classes=3)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert calls[0] == call(3), f"First reinit should resize to checkpoint size 3, got {calls[0]}"
+        assert len(calls) == 1, (
+            f"Expected exactly 1 reinit call (to checkpoint size), but got {len(calls)}: "
+            f"{calls}. The second reinit to 91 destroys loaded weights."
+        )
+        assert mc.num_classes == 2, (
+            f"mc.num_classes must be auto-aligned to 2 (checkpoint_logits - 1), got {mc.num_classes}"
+        )
+
+    def test_no_mismatch_no_reinit(self, monkeypatch):
+        """Checkpoint class count matches config — no reinit at all.
+
+        Scenario: COCO checkpoint (91 classes) with num_classes=90. 91 == 90 + 1, so no reinit should fire.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=90)
+        checkpoint = _make_checkpoint(num_classes=91)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        fake_model.reinitialize_detection_head.assert_not_called()
+
+    def test_backbone_pretrain_still_reinits(self, monkeypatch):
+        """Backbone pretrain (more classes in checkpoint) must still reinit.
+
+        Scenario: COCO 91-class checkpoint loaded for 2-class fine-tuning
+        (num_classes=2). Both reinits are correct here: first to 91 for load_state_dict, second to 3 for the configured
+        class count.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=2)
+        checkpoint = _make_checkpoint(num_classes=91)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert calls == [call(91), call(3)], f"Expected reinit to [91, 3] (expand then trim), got {calls}"
+
+    def test_user_override_larger_than_checkpoint_reexpands_head(self, monkeypatch):
+        """Explicit larger num_classes must be restored after checkpoint load.
+
+        Scenario: 91-class checkpoint loaded with explicit num_classes=93.
+        Loader must temporarily match checkpoint size for load_state_dict, then expand to 94 logits and keep
+        args.num_classes unchanged.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=93)
+        checkpoint = _make_checkpoint(num_classes=91)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert calls == [call(91), call(94)], f"Expected reinit to [91, 94] (load then expand), got {calls}"
+        assert mc.num_classes == 93, "Explicitly configured num_classes must not be overwritten."
+
+    # All non-deprecated model configs (RFDETRLargeDeprecatedConfig and
+    # RFDETRBaseConfig are excluded; the former is deprecated, the latter
+    # serves as the base class for the concrete variants below).
+    @pytest.mark.parametrize(
+        "config_cls",
+        [
+            pytest.param(RFDETRNanoConfig, id="nano"),
+            pytest.param(RFDETRSmallConfig, id="small"),
+            pytest.param(RFDETRMediumConfig, id="medium"),
+            pytest.param(RFDETRLargeConfig, id="large"),
+            pytest.param(RFDETRSegNanoConfig, id="seg_nano"),
+            pytest.param(RFDETRSegSmallConfig, id="seg_small"),
+            pytest.param(RFDETRSegMediumConfig, id="seg_medium"),
+            pytest.param(RFDETRSegLargeConfig, id="seg_large"),
+            pytest.param(RFDETRSegXLargeConfig, id="seg_xlarge"),
+            pytest.param(RFDETRSeg2XLargeConfig, id="seg_2xlarge"),
+        ],
+    )
+    def test_eight_class_finetune_checkpoint_auto_aligns_num_classes_and_reinits_once(self, monkeypatch, config_cls):
+        """Auto-align ``mc.num_classes`` and avoid a second reinit for 8-class checkpoints.
+
+        Scenario (from user bug report): user trains on 8 categories (IDs 0–7).
+        The checkpoint stores ``class_embed.bias`` with shape [9] (8 user classes + 1 background). Loading without
+        specifying ``num_classes`` must NOT trigger a second reinit to 91 after temporarily matching the checkpoint size
+        for ``load_state_dict``.
+
+        This test asserts the loader auto-aligns ``mc.num_classes`` to 8 (9 - 1) and fires exactly one reinit call — to
+        9 (the checkpoint size).
+        """
+        # 8 dataset categories → training builds a model with 8+1=9 logits.
+        checkpoint = _make_checkpoint(num_classes=9)
+        mc = config_cls(pretrain_weights="/fake/weights.pth", device="cpu")
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert len(calls) == 1, (
+            f"Expected exactly 1 reinit call (to checkpoint size 9), but got {len(calls)}: "
+            f"{calls}. A second reinit to 91 would produce OOB class IDs like 73."
+        )
+        assert calls[0] == call(9), f"Reinit must resize to checkpoint's 9 logits, got {calls[0]}"
+        assert mc.num_classes == 8, (
+            f"mc.num_classes must be auto-aligned to 8 (checkpoint_logits - 1), got {mc.num_classes}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Regression #960: PE interpolation for custom resolution
+# ---------------------------------------------------------------------------
+
+PE_KEY = "backbone.0.encoder.encoder.embeddings.position_embeddings"
+
+
+class TestLoadPretrainWeightsPEInterpolation:
+    """Regression tests for #960 — PE must be interpolated when resolution changes.
+
+    ``load_pretrain_weights`` must bicubic-interpolate the checkpoint's DINOv2 positional embeddings to match the
+    model's ``positional_encoding_size`` before calling ``load_state_dict``.  Without this, any custom ``resolution``
+    that changes the PE grid size causes a ``RuntimeError: size mismatch``.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_download(self, monkeypatch):
+        """Suppress all download and file-existence side effects."""
+        _suppress_pretrain_io(monkeypatch)
+
+    @pytest.mark.parametrize(
+        "src_pe_size, tgt_resolution, patch_size, expected_tgt_pe_size",
+        [
+            pytest.param(24, 640, 16, 40, id="nano_24x24_upscale_to_40x40"),
+            pytest.param(40, 384, 16, 24, id="nano_40x40_downscale_to_24x24"),
+            pytest.param(32, 640, 16, 40, id="small_32x32_upscale_to_40x40"),
+        ],
+    )
+    def test_pe_in_checkpoint_is_interpolated_to_model_resolution(
+        self, monkeypatch, src_pe_size, tgt_resolution, patch_size, expected_tgt_pe_size
+    ):
+        """Checkpoint PE is bicubic-interpolated to match model_config.positional_encoding_size.
+
+        Regression for #960: ``load_pretrain_weights`` must not raise ``RuntimeError`` when model resolution differs
+        from checkpoint resolution.  The PE tensor in the checkpoint must be resized in-place before ``load_state_dict``
+        is called.
+        """
+        mc = RFDETRNanoConfig(
+            pretrain_weights="/fake/weights.pth",
+            device="cpu",
+            resolution=tgt_resolution,
+            patch_size=patch_size,
+        )
+        assert mc.positional_encoding_size == tgt_resolution // patch_size
+
+        dim = 384
+        src_n = src_pe_size * src_pe_size + 1  # patches + class token
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = torch.randn(1, src_n, dim).half()  # float16 to verify dtype round-trip
+
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        expected_n = expected_tgt_pe_size * expected_tgt_pe_size + 1
+        assert pe.shape == torch.Size([1, expected_n, dim]), (
+            f"Expected PE shape [1, {expected_n}, {dim}], got {tuple(pe.shape)}. "
+            f"PE was not interpolated from {src_pe_size}x{src_pe_size} "
+            f"to {expected_tgt_pe_size}x{expected_tgt_pe_size}."
+        )
+        assert pe.dtype == torch.float16, f"Dtype must be preserved after interpolation, got {pe.dtype}"
+
+    def test_matching_pe_shape_is_not_modified(self, monkeypatch):
+        """When checkpoint PE matches model expectations, the tensor is not changed.
+
+        Ensures PE interpolation is a no-op for same-resolution checkpoints so that normal weight loading is unaffected.
+        """
+        mc = RFDETRNanoConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        # Default: positional_encoding_size=24 → PE = [1, 24*24+1, 384] = [1, 577, 384]
+
+        dim = 384
+        original_pe = torch.randn(1, 577, dim)
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = original_pe.clone()
+
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        assert pe.shape == torch.Size([1, 577, dim]), "Matching PE shape must not be modified."
+        assert torch.equal(pe, original_pe), "Matching PE tensor values must not be modified."
+
+    def test_base_config_non_formula_pe_is_interpolated_from_smaller_checkpoint(self, monkeypatch):
+        """RFDETRBaseConfig PE=37 (not formula-derived) is interpolated when checkpoint differs.
+
+        RFDETRBaseConfig.positional_encoding_size=37 is not updated by ``_sync_pe_with_resolution`` because 37 ≠
+        560//16=35 (not formula-derived). Loading a checkpoint with a smaller PE grid (e.g., 24×24) must still trigger
+        interpolation to the model's fixed PE=37×37 target.
+        """
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        assert mc.positional_encoding_size == 37, "RFDETRBaseConfig PE must remain 37 (not formula-derived)"
+
+        dim = 384
+        src_pe_size = 24
+        src_n = src_pe_size * src_pe_size + 1
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = torch.randn(1, src_n, dim)
+
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        expected_n = 37 * 37 + 1
+        assert pe.shape == torch.Size([1, expected_n, dim]), (
+            f"Expected PE shape [1, {expected_n}, {dim}] (37×37 grid), got {tuple(pe.shape)}. "
+            "BaseConfig's non-formula-derived PE must be the interpolation target."
+        )
+
+    def test_non_square_source_pe_logs_warning_and_is_not_modified(self, monkeypatch):
+        """Non-square source PE grids are skipped with a warning and left unchanged.
+
+        When ``n_source`` is not a perfect square the interpolation is skipped to avoid producing malformed embeddings.
+        The tensor must remain untouched and a warning must be emitted via the weights module logger.
+        """
+        mc = RFDETRNanoConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        # positional_encoding_size=24 → n_target=576 (perfect square, so the
+        # target-side guard does not trigger; only the source-side guard fires)
+
+        dim = 384
+        # 17 is not a perfect square: isqrt(17)=4, 4*4=16 ≠ 17
+        non_square_n_source = 17
+        original_pe = torch.randn(1, non_square_n_source + 1, dim)
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = original_pe.clone()
+
+        warning_calls: list[tuple] = []
+        monkeypatch.setattr("rfdetr.models.weights.logger.warning", lambda *a, **kw: warning_calls.append(a))
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        assert torch.equal(pe, original_pe), "Non-square source PE must not be modified."
+        assert any("not a perfect square" in str(args) for args in warning_calls), (
+            f"Expected a 'not a perfect square' warning; got calls: {warning_calls}"
+        )
+
+
+class TestL1FacadePEInterpolationEndToEnd:
+    """Regression for instantiating an RF-DETR L1 facade variant with a custom ``resolution`` and a checkpoint trained
+    at the variant's default resolution must not raise ``RuntimeError`` from a PE shape mismatch.
+
+    In v1.6.5 the L1 facade (``RFDETRLarge``, ``RFDETRNano``, ...) used a private ``_load_pretrain_weights_into`` helper
+    in ``detr.py`` that bypassed the PE bicubic-interpolation added to ``models.weights.load_pretrain_weights`` Code
+    that wired the L1 facade through the unified loader landed later (``inference._build_model_context`` calling
+    ``load_pretrain_weights`` from ``models.weights``).  This test pins that wiring so a future refactor cannot
+    reintroduce a divergent loader path that silently skips PE interpolation.
+
+    Current coverage: ``RFDETRNano`` (detection) and ``RFDETRSegNano`` (segmentation), upward-interpolation only.  When
+    a third L1 facade variant is added, collapse both methods to a single ``@pytest.mark.parametrize`` over
+    ``(variant_class, default_pe_grid, patch_size, new_resolution)``. Downward-interpolation (high-res checkpoint →
+    lower-res model) is not currently exercised; add a reverse-direction parametrize row when refactoring.
+    """
+
+    def test_rfdetr_nano_loads_default_pe_checkpoint_at_custom_resolution(self, tmp_path):
+        """Saving an RFDETRNano state_dict at default resolution and loading at a higher resolution must succeed via PE
+        interpolation in the L1 facade.
+
+        Mirrors the user-reported scenario in https://github.com/roboflow/rf-detr/issues/990 (PE size mismatch ``[1,
+        1937, 384]`` vs ``[1, 6401, 384]`` raised from ``LWDETR.load_state_dict``), reduced to RFDETRNano for test
+        speed.
+        """
+        from rfdetr import RFDETRNano
+
+        # 1. Build a default-resolution RFDETRNano (no pretrain, on CPU) so that
+        #    it produces a state_dict with the variant's native PE grid.
+        default_model = RFDETRNano(pretrain_weights=None, num_classes=2, device="cpu")
+        default_pe_grid = default_model.model_config.positional_encoding_size
+        assert default_pe_grid == 24, "RFDETRNano default PE grid must be 24×24"
+        patch_size = default_model.model_config.patch_size
+        default_state = default_model.model.model.state_dict()
+        default_pe = default_state[PE_KEY]
+        pe_dim = default_pe.shape[-1]
+        assert default_pe.shape == torch.Size([1, default_pe_grid * default_pe_grid + 1, pe_dim])
+
+        # 2. Persist as a checkpoint that mimics what `model.train()` saves —
+        #    a top-level "model" key plus a SimpleNamespace "args" block.
+        ckpt_path = tmp_path / "user_finetuned.pth"
+        torch.save(
+            {
+                "model": dict(default_state),
+                "args": SimpleNamespace(class_names=["a", "b"], patch_size=patch_size),
+            },
+            ckpt_path,
+        )
+
+        # 3. Re-instantiate at a NEW resolution.  Without PE interpolation in
+        #    the L1 facade path this raises ``RuntimeError: size mismatch for
+        #    backbone.0.encoder.encoder.embeddings.position_embeddings`` from
+        #    LWDETR.load_state_dict — exactly the user-reported failure.
+        new_resolution = 640
+        loaded = RFDETRNano(
+            pretrain_weights=str(ckpt_path),
+            resolution=new_resolution,
+            num_classes=2,
+            device="cpu",
+        )
+
+        # 4. The model_config validator must update PE proportionally to the
+        #    new resolution, and the loaded backbone PE parameter must have the
+        #    interpolated target shape (40 × 40 + 1 = 1601 tokens).
+        expected_pe_grid = new_resolution // patch_size
+        assert expected_pe_grid == 40
+        assert loaded.model_config.positional_encoding_size == expected_pe_grid
+        loaded_pe = loaded.model.model.state_dict()[PE_KEY]
+        assert loaded_pe.shape == torch.Size([1, expected_pe_grid * expected_pe_grid + 1, pe_dim]), (
+            f"Backbone PE was not interpolated to the requested resolution; "
+            f"got shape {tuple(loaded_pe.shape)}, expected [1, {expected_pe_grid**2 + 1}, {pe_dim}]."
+        )
+
+    def test_rfdetr_seg_nano_loads_default_pe_checkpoint_at_custom_resolution(self, tmp_path):
+        """Saving an RFDETRSegNano state_dict at default resolution and loading at a higher resolution must succeed via
+        PE interpolation in the L1 facade.
+
+        Regression for https://github.com/roboflow/rf-detr/issues/1023 — the segmentation model variant
+        (``RFDETRSegNano``) raised ``RuntimeError: size mismatch for
+        backbone.0.encoder.encoder.embeddings.position_embeddings`` when instantiated with a non-default ``resolution``
+        because the L1 facade's checkpoint-loading path did not interpolate positional embeddings for segmentation
+        models.
+        """
+        from rfdetr import RFDETRSegNano
+
+        # 1. Build a default-resolution RFDETRSegNano (no pretrain, on CPU).
+        #    Uses 90 classes to mimic an official COCO-pretrained checkpoint so
+        #    the load path at step 3 exercises both head-reinit (90 → 2 classes)
+        #    and PE interpolation simultaneously.
+        default_model = RFDETRSegNano(pretrain_weights=None, num_classes=90, device="cpu")
+        default_pe_grid = default_model.model_config.positional_encoding_size
+        assert default_pe_grid == 26, "RFDETRSegNano default PE grid must be 26×26 (312 // 12)"
+        patch_size = default_model.model_config.patch_size
+        assert patch_size == 12, "RFDETRSegNano patch_size must be 12"
+        default_state = default_model.model.model.state_dict()
+        default_pe = default_state[PE_KEY]
+        pe_dim = default_pe.shape[-1]
+        assert default_pe.shape == torch.Size([1, default_pe_grid * default_pe_grid + 1, pe_dim])
+
+        # 2. Persist as a checkpoint that mimics the official pretrain weights
+        #    format.  Saved as .pth (not .pt) so the tmp_path fixture path does
+        #    not trigger ModelWeights registry / MD5 lookup.  Top-level keys
+        #    match the real checkpoint: "model" (state_dict) and "args" with
+        #    segmentation_head=True and patch_size=12.
+        ckpt_path = tmp_path / "rf-detr-seg-nano.pth"
+        torch.save(
+            {
+                "model": dict(default_state),
+                "args": SimpleNamespace(
+                    class_names=[],
+                    patch_size=patch_size,
+                    segmentation_head=True,
+                ),
+            },
+            ckpt_path,
+        )
+
+        # 3. Re-instantiate at a custom resolution with fewer classes.  Without
+        #    PE interpolation this raises
+        #    ``RuntimeError: size mismatch for
+        #    backbone.0.encoder.encoder.embeddings.position_embeddings``
+        #    from LWDETR.load_state_dict — exactly the user-reported failure.
+        # resolution=1008 (user-reported in #1023, 84×84=7057 tokens) is deferred to follow-up parametrization.
+        new_resolution = 624  # 2× the default 312; divisible by patch_size=12
+        loaded = RFDETRSegNano(
+            pretrain_weights=str(ckpt_path),
+            resolution=new_resolution,
+            num_classes=2,
+            device="cpu",
+        )
+
+        # 4. The model_config validator must update PE proportionally to the
+        #    new resolution, and the loaded backbone PE parameter must have the
+        #    interpolated target shape (52 × 52 + 1 = 2705 tokens).
+        expected_pe_grid = new_resolution // patch_size
+        assert expected_pe_grid == 52
+        assert loaded.model_config.positional_encoding_size == expected_pe_grid
+        loaded_pe = loaded.model.model.state_dict()[PE_KEY]
+        assert loaded_pe.shape == torch.Size([1, expected_pe_grid * expected_pe_grid + 1, pe_dim]), (
+            f"Backbone PE was not interpolated to the requested resolution; "
+            f"got shape {tuple(loaded_pe.shape)}, expected [1, {expected_pe_grid**2 + 1}, {pe_dim}]."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Deprecation: train_config argument
+# ---------------------------------------------------------------------------
+
+
+class TestLoadPretrainWeightsDeprecation:
+    """Passing train_config must emit a DeprecationWarning."""
+
+    def test_emits_deprecation_warning_when_train_config_passed(self, monkeypatch):
+        """Any non-None train_config triggers a DeprecationWarning."""
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights=None, device="cpu")
+        tc = _make_train_config()
+
+        with pytest.warns(FutureWarning, match="train_config.*deprecated"):
+            load_pretrain_weights(MagicMock(), mc, tc)
+
+
+# ---------------------------------------------------------------------------
+# Regression #1038: PE interpolation for custom resolution — training path
+# ---------------------------------------------------------------------------
+
+
+class TestModuleLoadPretrainWeightsPEInterpolationCustomResolution:
+    """Regression for #1038 — PE interpolation must fire through ``RFDETRModelModule.__init__``.
+
+    The L2 training entry path (``RFDETRSegLarge(resolution=1008).train(...)``) constructs an
+    :class:`~rfdetr.training.module_model.RFDETRModelModule` whose ``__init__`` delegates to
+    :func:`~rfdetr.models.weights.load_pretrain_weights`. That helper must bicubic-interpolate the checkpoint's DINOv2
+    positional embeddings to match ``model_config.positional_encoding_size`` before calling ``load_state_dict``.
+    Without this, any ``model.train()`` call with a custom ``resolution`` that changes the PE grid raises::
+
+        RuntimeError: Error(s) in loading state_dict for LWDETR:
+            size mismatch for backbone.0.encoder.encoder.embeddings.position_embeddings
+
+    These tests exercise the construction path end-to-end (mocking only the heavy ``build_model_from_config`` /
+    ``build_criterion_from_config`` calls and disk I/O), so the regression cannot reappear if the in-init delegation to
+    ``load_pretrain_weights`` is removed.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_download(self, monkeypatch):
+        """Suppress download/validate side effects on the canonical load path."""
+        _suppress_pretrain_io(monkeypatch)
+
+    def _construct_module(self, mc, checkpoint, monkeypatch, tmp_path):
+        """Construct an RFDETRModelModule with all heavy work mocked.
+
+        Returns the constructed module and the fake nn_model whose ``load_state_dict`` receives the (now-interpolated)
+        state dict.
+        """
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        # Pretend the head was already aligned so the canonical loader does not
+        # try to introspect the MagicMock's internals.
+        fake_model.num_classes = mc.num_classes
+        monkeypatch.setattr(
+            "rfdetr.training.module_model.build_model_from_config",
+            lambda *a, **kw: fake_model,
+        )
+        monkeypatch.setattr(
+            "rfdetr.training.module_model.build_criterion_from_config",
+            lambda *a, **kw: (MagicMock(), MagicMock()),
+        )
+
+        tc = TrainConfig(
+            dataset_dir=str(tmp_path / "dataset"),
+            output_dir=str(tmp_path / "output"),
+            epochs=1,
+            lr=1e-4,
+            lr_encoder=1.5e-4,
+            batch_size=2,
+            weight_decay=1e-4,
+            lr_drop=1,
+            warmup_epochs=0.0,
+            drop_path=0.0,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            grad_accum_steps=1,
+            tensorboard=False,
+        )
+        from rfdetr.training.module_model import RFDETRModelModule
+
+        module = RFDETRModelModule(mc, tc)
+        return module, fake_model
+
+    @pytest.mark.parametrize(
+        "config_cls, src_pe_size, tgt_pe_size",
+        [
+            # All 7 segmentation variants listed in #1038, plus detection up/downscale.
+            pytest.param(RFDETRSegNanoConfig, 26, 84, id="seg_nano_upscale_26_to_84"),
+            pytest.param(RFDETRSegSmallConfig, 32, 84, id="seg_small_upscale_32_to_84"),
+            pytest.param(RFDETRSegMediumConfig, 36, 84, id="seg_medium_upscale_36_to_84"),
+            pytest.param(RFDETRSegLargeConfig, 42, 84, id="seg_large_upscale_42_to_84"),
+            pytest.param(RFDETRSegPreviewConfig, 24, 84, id="seg_preview_upscale_24_to_84"),
+            pytest.param(RFDETRSegXLargeConfig, 52, 84, id="seg_xlarge_upscale_52_to_84"),
+            pytest.param(RFDETRSeg2XLargeConfig, 64, 84, id="seg_2xlarge_upscale_64_to_84"),
+            pytest.param(RFDETRNanoConfig, 24, 40, id="nano_upscale_24_to_40"),
+            pytest.param(RFDETRNanoConfig, 40, 24, id="nano_downscale_40_to_24"),
+        ],
+    )
+    def test_pe_interpolated_in_training_path(self, monkeypatch, config_cls, src_pe_size, tgt_pe_size, tmp_path):
+        """Module construction interpolates PE to match ``positional_encoding_size``.
+
+        Regression for #1038 — ``RFDETRModelModule.__init__`` must trigger PE interpolation through the canonical loader
+        so ``load_state_dict`` does not raise ``RuntimeError: size mismatch`` at custom training resolutions.
+        """
+        mc = config_cls(
+            pretrain_weights="/fake/weights.pth",
+            device="cpu",
+            positional_encoding_size=tgt_pe_size,
+        )
+
+        dim = 384
+        src_n = src_pe_size * src_pe_size + 1
+        checkpoint = _make_checkpoint(num_classes=mc.num_classes + 1)
+        checkpoint["model"][PE_KEY] = torch.randn(1, src_n, dim)
+
+        _, fake_model = self._construct_module(mc, checkpoint, monkeypatch, tmp_path)
+
+        pe = checkpoint["model"][PE_KEY]
+        expected_n = tgt_pe_size * tgt_pe_size + 1
+        assert pe.shape == torch.Size([1, expected_n, dim]), (
+            f"Expected PE shape [1, {expected_n}, {dim}] after interpolation from "
+            f"{src_pe_size}x{src_pe_size} to {tgt_pe_size}x{tgt_pe_size}, got {tuple(pe.shape)}. "
+            "PE interpolation must fire during RFDETRModelModule.__init__ via canonical load_pretrain_weights."
+        )
+        # load_state_dict was called on the model with the interpolated state dict.
+        fake_model.load_state_dict.assert_called_once()
+
+    def test_matching_pe_not_modified_in_training_path(self, monkeypatch, tmp_path):
+        """Same-resolution checkpoint PE is untouched in the training path."""
+        pe_size = 24
+        mc = RFDETRNanoConfig(
+            pretrain_weights="/fake/weights.pth",
+            device="cpu",
+            positional_encoding_size=pe_size,
+        )
+
+        dim = 384
+        original_pe = torch.randn(1, pe_size * pe_size + 1, dim)
+        checkpoint = _make_checkpoint(num_classes=mc.num_classes + 1)
+        checkpoint["model"][PE_KEY] = original_pe.clone()
+
+        self._construct_module(mc, checkpoint, monkeypatch, tmp_path)
+
+        pe = checkpoint["model"][PE_KEY]
+        assert pe.shape == torch.Size([1, pe_size * pe_size + 1, dim]), "Matching PE must not be modified."
+        assert torch.equal(pe, original_pe), "Matching PE values must not be modified."
+
+
+# ---------------------------------------------------------------------------
+# Regression: keypoint schema auto-align from checkpoint _kp_active_mask
+# ---------------------------------------------------------------------------
+
+
+def _make_kp_checkpoint(kp_schema: list[int], num_queries: int = 300, group_detr: int = 13) -> dict:
+    """Build a minimal keypoint checkpoint encoding *kp_schema* via ``_kp_active_mask``.
+
+    Args:
+        kp_schema: Keypoints-per-class list, e.g. ``[0, 17]`` for bg-first or ``[17]`` for active-first.
+        num_queries: Number of queries per group.
+        group_detr: Number of decoder groups.
+
+    Returns:
+        Checkpoint dict with ``model``, ``args``, and schema-derived ``_kp_active_mask``.
+    """
+    num_classes = len(kp_schema)
+    total_queries = num_queries * group_detr
+    max_kp = max(kp_schema) if kp_schema else 0
+    mask = torch.zeros(num_classes, max_kp, dtype=torch.bool)
+    for idx, n_kp in enumerate(kp_schema):
+        mask[idx, :n_kp] = True
+    state = {
+        "class_embed.weight": torch.randn(num_classes + 1, 256),
+        "class_embed.bias": torch.randn(num_classes + 1),
+        "refpoint_embed.weight": torch.randn(total_queries, 4),
+        "query_feat.weight": torch.randn(total_queries, 256),
+        "_kp_active_mask": mask,
+    }
+    ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14, class_names=[])
+    return {"model": state, "args": ckpt_args}
+
+
+def _make_kp_fake_model(initial_schema: list[int]) -> MagicMock:
+    """Build a fake nn_model mock that mimics the GroupPose keypoint model interface.
+
+    Args:
+        initial_schema: Keypoints-per-class schema the model was built with.
+
+    Returns:
+        Configured MagicMock with keypoint schema methods and ``_kp_active_mask`` state.
+    """
+    max_kp = max(initial_schema) if initial_schema else 0
+    current_schema = list(initial_schema)
+    mask = torch.zeros(len(initial_schema), max_kp, dtype=torch.bool)
+    for idx, n in enumerate(initial_schema):
+        mask[idx, :n] = True
+
+    fake_model = MagicMock()
+    fake_model.state_dict.return_value = {"_kp_active_mask": mask.clone()}
+
+    def _reinit_kp(schema):
+        current_schema[:] = schema
+        new_max = max(schema) if schema else 0
+        new_mask = torch.zeros(len(schema), new_max, dtype=torch.bool)
+        for i, n in enumerate(schema):
+            new_mask[i, :n] = True
+        fake_model.state_dict.return_value = {"_kp_active_mask": new_mask}
+        fake_model.get_num_keypoints_per_class.return_value = list(schema)
+
+    fake_model.reinitialize_keypoint_head.side_effect = _reinit_kp
+    fake_model.get_num_keypoints_per_class.return_value = list(initial_schema)
+    fake_model.get_num_keypoints_per_class_from_checkpoint = lambda sd: (
+        [int(n) for n in sd["_kp_active_mask"].sum(dim=1).tolist()] if "_kp_active_mask" in sd else None
+    )
+    return fake_model
+
+
+class TestLoadPretrainWeightsKeypointSchemaAutoAlign:
+    """Regression for AP=0 when loading bg-first checkpoint into active-first default config.
+
+    Before the fix, ``load_pretrain_weights`` would restore ``mc.num_keypoints_per_class`` to the config default
+    ``[17]`` after loading a ``[0, 17]`` pretrained checkpoint.  The detection-head trimming kept rows 0 (background)
+    and 1 (person) from the checkpoint but the active-first schema treats row 0 as person, producing AP≈0.
+
+    The fix auto-aligns ``mc.num_keypoints_per_class`` from the checkpoint ``_kp_active_mask`` when the user did not
+    explicitly set the field, mirroring the existing ``num_classes`` auto-align pattern.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_io(self, monkeypatch):
+        """Suppress all download and file-existence side effects."""
+        _suppress_pretrain_io(monkeypatch)
+
+    def test_bg_first_checkpoint_auto_aligns_active_first_config(self, monkeypatch):
+        """Loading bg-first [0,17] checkpoint into active-first [17] config auto-aligns schema."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        assert mc.num_keypoints_per_class == [17], "Precondition: config default is active-first [17]."
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [0, 17], (
+            f"expected auto-align to [0, 17], got {mc.num_keypoints_per_class}"
+        )
+        assert mc.num_classes == 2, (
+            f"mc.num_classes must be auto-aligned to 2 (checkpoint has 3 logit slots), got {mc.num_classes}"
+        )
+
+    def test_matching_schema_no_change(self, monkeypatch):
+        """Config and checkpoint with same schema leaves mc.num_keypoints_per_class unchanged."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        mc.num_keypoints_per_class = [17]
+
+        checkpoint = _make_kp_checkpoint([17])
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [17]
+
+    def test_user_explicit_schema_not_overridden(self, monkeypatch):
+        """Explicit num_keypoints_per_class from user survives even when checkpoint differs."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        # Simulate user explicitly providing num_keypoints_per_class
+        mc.num_keypoints_per_class = [0, 33]
+        mc.model_fields_set.add("num_keypoints_per_class")
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([0, 33])
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [0, 33], (
+            "Explicit user schema must not be overridden by checkpoint auto-align."
+        )
+
+    def test_1d_kp_active_mask_skips_auto_align_with_warning(self, monkeypatch):
+        """Malformed 1-D _kp_active_mask skips auto-align and emits a logger.warning."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        assert mc.num_keypoints_per_class == [17], "Precondition: config default is active-first [17]."
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        # Replace the 2-D mask with a malformed 1-D tensor.
+        checkpoint["model"]["_kp_active_mask"] = torch.ones(17, dtype=torch.bool)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+        # The fake model's get_num_keypoints_per_class_from_checkpoint calls sum(dim=1),
+        # which raises IndexError on a 1-D tensor; return None to skip that code path.
+        fake_model.get_num_keypoints_per_class_from_checkpoint = lambda sd: None
+
+        warned: list[str] = []
+        monkeypatch.setattr("rfdetr.models.weights.logger.warning", lambda msg, *args: warned.append(msg % args))
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [17], (
+            "Auto-align must not fire for a 1-D _kp_active_mask; config default should be unchanged."
+        )
+        assert any("unexpected shape" in msg for msg in warned), (
+            f"Expected a warning mentioning 'unexpected shape'; got: {warned}"
+        )
+
+    def test_all_zero_kp_active_mask_skips_auto_align_with_warning(self, monkeypatch):
+        """All-zero 2-D _kp_active_mask (degenerate) skips auto-align and emits a logger.warning."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        assert mc.num_keypoints_per_class == [17], "Precondition: config default is active-first [17]."
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        # Replace mask with a valid 2-D shape but all zeros (no active slots).
+        checkpoint["model"]["_kp_active_mask"] = torch.zeros(2, 17, dtype=torch.bool)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+
+        warned: list[str] = []
+        monkeypatch.setattr("rfdetr.models.weights.logger.warning", lambda msg, *args: warned.append(msg % args))
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [17], (
+            "Auto-align must not fire for an all-zero _kp_active_mask; config default should be unchanged."
+        )
+        assert any("no active slots" in msg for msg in warned), (
+            f"Expected a warning mentioning 'no active slots'; got: {warned}"
+        )

tests/training/test_metrics_csv.py

@@ -0,0 +1,159 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Integration tests: metrics.csv contains all columns used by plot_metrics().
+
+Runs a minimal PTL training loop (1 epoch, 2 batches each) using mocked model internals so no real dataset or GPU is
+required.  After training, reads the CSVLogger output and asserts that every metric column that ``plot_metrics()`` needs
+is present and has at least one non-NaN value.
+
+Also verifies that ``train/loss`` is logged at the same scale as ``val/loss`` (i.e. NOT divided by ``grad_accum_steps``
+before logging).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pandas as pd
+import torch
+
+from rfdetr.config import RFDETRBaseConfig, TrainConfig
+from rfdetr.training import build_trainer
+from rfdetr.training.module_data import RFDETRDataModule
+from rfdetr.training.module_model import RFDETRModelModule
+
+from .helpers import _fake_postprocess, _FakeCriterion, _FakeDataset, _make_param_dicts, _TinyModel
+
+# ---------------------------------------------------------------------------
+# Helpers local to this module
+# ---------------------------------------------------------------------------
+
+
+def _fit_and_read_csv(mc: RFDETRBaseConfig, tc: TrainConfig, criterion=None) -> pd.DataFrame:
+    """Run 1 epoch (2 train + 2 val batches) and return the resulting metrics.csv."""
+    fake_criterion = criterion or _FakeCriterion()
+    with (
+        patch("rfdetr.training.module_model.build_model_from_config", return_value=_TinyModel()),
+        patch(
+            "rfdetr.training.module_model.build_criterion_from_config",
+            return_value=(fake_criterion, MagicMock(side_effect=_fake_postprocess)),
+        ),
+        patch("rfdetr.training.module_data.build_dataset", return_value=_FakeDataset(length=20)),
+        patch(
+            "rfdetr.training.module_model.get_param_dict",
+            side_effect=lambda args, model: _make_param_dicts(model),
+        ),
+    ):
+        module = RFDETRModelModule(mc, tc)
+        datamodule = RFDETRDataModule(mc, tc)
+        trainer = build_trainer(
+            tc,
+            mc,
+            accelerator="cpu",
+            max_epochs=1,
+            limit_train_batches=2,
+            limit_val_batches=2,
+            log_every_n_steps=1,
+        )
+        trainer.fit(module, datamodule=datamodule)
+
+    csv_path = Path(tc.output_dir) / "metrics.csv"
+    assert csv_path.exists(), "CSVLogger must write metrics.csv to output_dir"
+    return pd.read_csv(csv_path)
+
+
+# ---------------------------------------------------------------------------
+# Expected columns (must exist and have ≥1 non-NaN row after one epoch)
+# ---------------------------------------------------------------------------
+
+_REQUIRED_DETECTION = frozenset(
+    {
+        "train/loss",
+        "train/lr",
+        "val/loss",
+        "val/mAP_50",
+        "val/mAP_50_95",
+        "val/mAR",
+    }
+)
+
+_REQUIRED_DETECTION_EMA = _REQUIRED_DETECTION | frozenset(
+    {
+        "val/ema_mAP_50",
+        "val/ema_mAP_50_95",
+        "val/ema_mAR",
+    }
+)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestDetectionMetricsCSV:
+    """metrics.csv contains all columns that plot_metrics() needs for detection."""
+
+    def test_base_metrics_present_without_ema(self, base_model_config, base_train_config):
+        """Without EMA all core val/* columns must appear in metrics.csv with non-NaN data."""
+        mc = base_model_config()
+        tc = base_train_config(use_ema=False, run_test=False)
+        df = _fit_and_read_csv(mc, tc)
+
+        missing = _REQUIRED_DETECTION - set(df.columns)
+        assert not missing, f"Missing columns in metrics.csv: {sorted(missing)}"
+
+        all_nan = {c for c in _REQUIRED_DETECTION if df[c].isna().all()}
+        assert not all_nan, f"Columns with all-NaN values: {sorted(all_nan)}"
+
+    def test_ema_metrics_present_with_ema_enabled(self, base_model_config, base_train_config):
+        """With use_ema=True the ema_* aliases must also appear in metrics.csv."""
+        mc = base_model_config()
+        tc = base_train_config(use_ema=True, run_test=False)
+        df = _fit_and_read_csv(mc, tc)
+
+        missing = _REQUIRED_DETECTION_EMA - set(df.columns)
+        assert not missing, f"Missing EMA columns in metrics.csv: {sorted(missing)}"
+
+        all_nan = {c for c in _REQUIRED_DETECTION_EMA if df[c].isna().all()}
+        assert not all_nan, f"EMA columns with all-NaN values: {sorted(all_nan)}"
+
+    def test_train_loss_is_unscaled(self, base_model_config, base_train_config):
+        """Train/loss must be logged at the raw criterion scale, not divided by grad_accum_steps.
+
+        With grad_accum_steps=4 the old code divided the logged value by 4, making train/loss ~4× smaller than val/loss.
+        After the fix the logged value equals the raw weighted criterion output so both losses are on the same scale.
+        """
+        fixed_loss_value = 5.0
+        grad_accum_steps = 4
+
+        class _FixedCriterion:
+            weight_dict = {"loss_ce": 1.0}
+
+            def num_boxes_for_targets(self, outputs, targets):
+                dummy = outputs.get("dummy", torch.zeros(1))
+                return torch.ones((), dtype=dummy.dtype, device=dummy.device)
+
+            def __call__(self, outputs, targets, num_boxes=None):
+                # Loss is always fixed_loss_value, connected to model params for gradient.
+                dummy = outputs.get("dummy", torch.zeros(1))
+                denominator = self.num_boxes_for_targets(outputs, targets) if num_boxes is None else num_boxes
+                return {"loss_ce": (dummy.mean() * 0 + fixed_loss_value) / denominator}
+
+        mc = base_model_config()
+        tc = base_train_config(use_ema=False, run_test=False, grad_accum_steps=grad_accum_steps)
+        df = _fit_and_read_csv(mc, tc, criterion=_FixedCriterion())
+
+        logged = df["train/loss"].dropna().mean()
+        expected_unscaled = fixed_loss_value
+        expected_if_divided = fixed_loss_value / grad_accum_steps
+
+        assert abs(logged - expected_unscaled) < abs(logged - expected_if_divided), (
+            f"train/loss={logged:.4f} is closer to the grad-accum-divided value "
+            f"({expected_if_divided:.4f}) than the raw criterion output "
+            f"({expected_unscaled:.4f}). The division must have been removed."
+        )

tests/training/test_train_resume.py

@@ -0,0 +1,88 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for resuming training from checkpoint."""
+
+import warnings
+from pathlib import Path
+from unittest.mock import patch
+
+from rfdetr import RFDETRNano
+
+
+def test_resume_with_completed_epochs_returns_early(tmp_path: Path) -> None:
+    """Passing start_epoch emits DeprecationWarning and still reaches trainer.fit().
+
+    In the legacy engine.py path, ``start_epoch=epochs`` caused the training loop to be skipped (``range(start_epoch,
+    epochs)`` was empty), which triggered an ``UnboundLocalError`` when accessing ``test_stats``.
+
+    In the PTL path, ``start_epoch`` is a deprecated kwarg that is absorbed and ignored (PTL resumes automatically via
+    ``ckpt_path``). The shim should emit the warning and still call ``trainer.fit(...)`` without raising.
+
+    Args:
+        tmp_path: Pytest temporary directory.
+    """
+    output_dir = tmp_path / "train_output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    model = RFDETRNano(pretrain_weights=None, num_classes=3, device="cpu")
+
+    with (
+        patch("rfdetr.training.RFDETRModelModule"),
+        patch("rfdetr.training.RFDETRDataModule"),
+        patch("rfdetr.training.build_trainer") as mock_build_trainer,
+        warnings.catch_warnings(record=True) as caught,
+    ):
+        warnings.simplefilter("always")
+        model.train(
+            dataset_dir=str(tmp_path),
+            epochs=1,
+            start_epoch=1,
+            batch_size=1,
+            grad_accum_steps=1,
+            output_dir=str(output_dir),
+            device="cpu",
+        )
+
+    depr = [w for w in caught if issubclass(w.category, DeprecationWarning)]
+    assert any("start_epoch" in str(w.message) for w in depr), "Expected a DeprecationWarning mentioning start_epoch"
+    mock_build_trainer.return_value.fit.assert_called_once()
+
+
+def test_resume_with_completed_epochs_calls_on_train_end_callback(tmp_path: Path) -> None:
+    """Old-style on_train_end callbacks are not forwarded to PTL.
+
+    In the legacy engine.py path, callbacks added to ``model.callbacks["on_train_end"]`` were invoked at the end of
+    training (including when the loop was skipped). In the PTL path the old-style callback dict on the model instance is
+    not consulted; use PTL ``Callback`` objects via ``build_trainer()`` instead.
+    """
+    output_dir = tmp_path / "train_output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    callback_calls = 0
+
+    def _callback() -> None:
+        nonlocal callback_calls
+        callback_calls += 1
+
+    model = RFDETRNano(pretrain_weights=None, num_classes=3, device="cpu")
+    model.callbacks["on_train_end"].append(_callback)
+
+    with (
+        patch("rfdetr.training.RFDETRModelModule"),
+        patch("rfdetr.training.RFDETRDataModule"),
+        patch("rfdetr.training.build_trainer"),
+    ):
+        model.train(
+            dataset_dir=str(tmp_path),
+            epochs=1,
+            batch_size=1,
+            grad_accum_steps=1,
+            output_dir=str(output_dir),
+            device="cpu",
+        )
+
+    # Old-style callbacks on model.callbacks are no longer invoked in the PTL path.
+    assert callback_calls == 0

tests/training/test_trainer.py

@@ -0,0 +1,79 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for build_trainer() — callback stack and config coercion."""
+
+import pytest
+from pytorch_lightning.callbacks import RichProgressBar, TQDMProgressBar
+
+from rfdetr.training import build_trainer
+
+# ---------------------------------------------------------------------------
+# TestProgressBarCallbacks — verifies the correct callback is installed
+# ---------------------------------------------------------------------------
+
+
+class TestProgressBarCallbacks:
+    """build_trainer() must install the right progress bar callback for each mode."""
+
+    def test_rich_progress_bar_installed_for_rich(self, base_model_config, base_train_config):
+        """progress_bar='rich' must add RichProgressBar and not TQDMProgressBar."""
+        mc = base_model_config()
+        tc = base_train_config(progress_bar="rich")
+        trainer = build_trainer(tc, mc, accelerator="cpu")
+        cb_types = [type(cb) for cb in trainer.callbacks]
+        assert RichProgressBar in cb_types
+        assert TQDMProgressBar not in cb_types
+
+    def test_tqdm_progress_bar_installed_for_tqdm(self, base_model_config, base_train_config):
+        """progress_bar='tqdm' must add TQDMProgressBar and not RichProgressBar."""
+        mc = base_model_config()
+        tc = base_train_config(progress_bar="tqdm")
+        trainer = build_trainer(tc, mc, accelerator="cpu")
+        cb_types = [type(cb) for cb in trainer.callbacks]
+        assert TQDMProgressBar in cb_types
+        assert RichProgressBar not in cb_types
+
+    def test_progress_bar_refresh_rate_is_five(self, base_model_config, base_train_config):
+        """The installed progress bar callback should refresh every five batches."""
+        mc = base_model_config()
+        tc = base_train_config(progress_bar="tqdm")
+        trainer = build_trainer(tc, mc, accelerator="cpu")
+        progress_bar = next(cb for cb in trainer.callbacks if isinstance(cb, TQDMProgressBar))
+
+        assert progress_bar.refresh_rate == 5
+
+    def test_no_progress_bar_callback_for_none(self, base_model_config, base_train_config):
+        """progress_bar=None must not add any progress bar callback."""
+        mc = base_model_config()
+        tc = base_train_config(progress_bar=None)
+        trainer = build_trainer(tc, mc, accelerator="cpu")
+        cb_types = [type(cb) for cb in trainer.callbacks]
+        assert RichProgressBar not in cb_types
+        assert TQDMProgressBar not in cb_types
+
+
+# ---------------------------------------------------------------------------
+# TestCoerceLegacyProgressBar — backward-compat validator on TrainConfig
+# ---------------------------------------------------------------------------
+
+
+class TestCoerceLegacyProgressBar:
+    """_coerce_legacy_progress_bar must normalise legacy bool values."""
+
+    @pytest.mark.parametrize(
+        "value, expected",
+        [
+            pytest.param(True, "tqdm", id="True->tqdm"),
+            pytest.param(False, None, id="False->None"),
+            pytest.param("rich", "rich", id="rich_passthrough"),
+            pytest.param("tqdm", "tqdm", id="tqdm_passthrough"),
+            pytest.param(None, None, id="None_passthrough"),
+        ],
+    )
+    def test_coerce(self, base_train_config, value, expected):
+        """progress_bar field normalises legacy bool and passes through string/None."""
+        tc = base_train_config(progress_bar=value)
+        assert tc.progress_bar == expected

tests/training/test_trainer_smoke.py

@@ -0,0 +1,432 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Smoke tests: Trainer(fast_dev_run=2).fit(module, datamodule) — T7.
+
+Verifies that the PTL training loop runs end-to-end without error for both detection and segmentation configurations.
+All heavy operations (build_model, build_criterion_and_postprocessors, build_dataset, get_param_dict) are patched so no
+real dataset or GPU is required.
+
+Chapter 1 gate: these must pass before Chapter 2 begins.
+"""
+
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from pytorch_lightning import Trainer
+
+from rfdetr.config import SegmentationTrainConfig
+from rfdetr.training import build_trainer
+from rfdetr.training.module_data import RFDETRDataModule
+from rfdetr.training.module_model import RFDETRModelModule
+
+from .helpers import (
+    _fake_postprocess,
+    _FakeCriterion,
+    _FakeDataset,
+    _FakeDatasetWithMasks,
+    _FakePostProcess,
+    _make_param_dicts,
+    _TinyModel,
+)
+
+# ---------------------------------------------------------------------------
+# Private helpers unique to smoke tests
+# ---------------------------------------------------------------------------
+
+
+def _make_trainer() -> Trainer:
+    """Create a Trainer configured for minimal smoke-test runs."""
+    return Trainer(
+        fast_dev_run=2,
+        accelerator="cpu",
+        enable_progress_bar=False,
+        enable_model_summary=False,
+        logger=False,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Smoke test classes
+# ---------------------------------------------------------------------------
+
+
+class TestDetectionSmoke:
+    """Trainer(fast_dev_run=2).fit() must complete without error for detection."""
+
+    def test_fit_runs_without_error(self, base_model_config, base_train_config):
+        """Full PTL fit loop runs 2 train + 2 val batches without raising."""
+        mc = base_model_config()
+        tc = base_train_config()
+
+        tiny_model = _TinyModel()
+        fake_criterion = _FakeCriterion()
+        fake_postprocess = MagicMock(side_effect=_fake_postprocess)
+        fake_dataset = _FakeDataset(length=20)
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=tiny_model),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(fake_criterion, fake_postprocess),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=fake_dataset),
+            patch(
+                "rfdetr.training.module_model.get_param_dict",
+                side_effect=lambda args, model: _make_param_dicts(model),
+            ),
+        ):
+            module = RFDETRModelModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+            _make_trainer().fit(module, datamodule)
+
+    def test_training_step_called_expected_times(self, base_model_config, base_train_config):
+        """fast_dev_run=2 must run exactly 2 training steps."""
+        mc = base_model_config()
+        tc = base_train_config()
+
+        tiny_model = _TinyModel()
+        fake_criterion = _FakeCriterion()
+        fake_postprocess = MagicMock(side_effect=_fake_postprocess)
+        fake_dataset = _FakeDataset(length=20)
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=tiny_model),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(fake_criterion, fake_postprocess),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=fake_dataset),
+            patch(
+                "rfdetr.training.module_model.get_param_dict",
+                side_effect=lambda args, model: _make_param_dicts(model),
+            ),
+        ):
+            module = RFDETRModelModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+
+            original_training_step = module.training_step
+            call_count = []
+
+            def _counting_training_step(batch, batch_idx):
+                call_count.append(1)
+                return original_training_step(batch, batch_idx)
+
+            module.training_step = _counting_training_step
+            _make_trainer().fit(module, datamodule)
+
+        assert sum(call_count) == 2
+
+    def test_val_step_called_expected_times(self, base_model_config, base_train_config):
+        """fast_dev_run=2 must run exactly 2 validation steps."""
+        mc = base_model_config()
+        tc = base_train_config()
+
+        tiny_model = _TinyModel()
+        fake_criterion = _FakeCriterion()
+        fake_postprocess = MagicMock(side_effect=_fake_postprocess)
+        fake_dataset = _FakeDataset(length=20)
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=tiny_model),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(fake_criterion, fake_postprocess),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=fake_dataset),
+            patch(
+                "rfdetr.training.module_model.get_param_dict",
+                side_effect=lambda args, model: _make_param_dicts(model),
+            ),
+        ):
+            module = RFDETRModelModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+
+            original_validation_step = module.validation_step
+            call_count = []
+
+            def _counting_val_step(batch, batch_idx):
+                call_count.append(1)
+                return original_validation_step(batch, batch_idx)
+
+            module.validation_step = _counting_val_step
+            _make_trainer().fit(module, datamodule)
+
+        assert sum(call_count) == 2
+
+    def test_loss_decreases_or_is_finite(self, base_model_config, base_train_config):
+        """Training loss must be finite (not NaN/inf) for the run to be valid."""
+        mc = base_model_config()
+        tc = base_train_config()
+
+        tiny_model = _TinyModel()
+        fake_postprocess = MagicMock(side_effect=_fake_postprocess)
+        fake_dataset = _FakeDataset(length=20)
+
+        losses = []
+
+        def _recording_criterion(outputs, targets, num_boxes=None):
+            dummy = outputs.get("dummy", torch.zeros(1))
+            denominator = fake_criterion.num_boxes_for_targets(outputs, targets) if num_boxes is None else num_boxes
+            loss = dummy.mean() / denominator
+            losses.append(loss.detach().item())
+            return {"loss_ce": loss}
+
+        fake_criterion = MagicMock(side_effect=_recording_criterion)
+        fake_criterion.weight_dict = {"loss_ce": 1.0}
+        fake_criterion.num_boxes_for_targets.return_value = torch.tensor(1.0)
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=tiny_model),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(fake_criterion, fake_postprocess),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=fake_dataset),
+            patch(
+                "rfdetr.training.module_model.get_param_dict",
+                side_effect=lambda args, model: _make_param_dicts(model),
+            ),
+        ):
+            module = RFDETRModelModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+            _make_trainer().fit(module, datamodule)
+
+        assert len(losses) > 0
+        assert all(torch.isfinite(torch.tensor(v)) for v in losses)
+
+
+class TestSegmentationSmoke:
+    """Trainer(fast_dev_run=2).fit() must complete without error for segmentation."""
+
+    def test_fit_runs_without_error(self, base_model_config, seg_train_config):
+        """Full PTL fit loop runs 2 train + 2 val batches without raising."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config()
+
+        tiny_model = _TinyModel()
+        fake_criterion = _FakeCriterion()
+        fake_postprocess = MagicMock(side_effect=_fake_postprocess)
+        fake_dataset = _FakeDatasetWithMasks(length=20)
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=tiny_model),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(fake_criterion, fake_postprocess),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=fake_dataset),
+            patch(
+                "rfdetr.training.module_model.get_param_dict",
+                side_effect=lambda args, model: _make_param_dicts(model),
+            ),
+        ):
+            module = RFDETRModelModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+            _make_trainer().fit(module, datamodule)
+
+    def test_segmentation_config_accepted(self, base_model_config, seg_train_config):
+        """SegmentationTrainConfig must be accepted by both module and datamodule."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config()
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=_TinyModel()),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(_FakeCriterion(), MagicMock(side_effect=_fake_postprocess)),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=_FakeDatasetWithMasks()),
+            patch(
+                "rfdetr.training.module_model.get_param_dict",
+                side_effect=lambda args, model: _make_param_dicts(model),
+            ),
+        ):
+            module = RFDETRModelModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+
+            assert isinstance(module.train_config, SegmentationTrainConfig)
+            assert isinstance(datamodule.train_config, SegmentationTrainConfig)
+
+
+class TestBuildTrainerSmoke:
+    """Smoke tests for the ``build_trainer()`` public factory.
+
+    Verifies that the full callback stack wired by ``build_trainer`` runs end-to-end with ``fast_dev_run``, using mocked
+    internals so no real dataset or GPU is required.
+    """
+
+    def test_fit_via_build_trainer(self, base_model_config, base_train_config):
+        """build_trainer() + trainer.fit(module, datamodule=datamodule) must not raise."""
+        mc = base_model_config()
+        tc = base_train_config(use_ema=False, run_test=False)
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=_TinyModel()),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(_FakeCriterion(), MagicMock(side_effect=_fake_postprocess)),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=_FakeDataset(length=20)),
+            patch(
+                "rfdetr.training.module_model.get_param_dict",
+                side_effect=lambda args, model: _make_param_dicts(model),
+            ),
+        ):
+            module = RFDETRModelModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+            trainer = build_trainer(tc, mc, accelerator="cpu", fast_dev_run=2)
+            trainer.fit(module, datamodule=datamodule)
+
+
+class _DDPModule(RFDETRModelModule):
+    """RFDETRModelModule subclass for ddp_spawn smoke tests.
+
+    Overrides ``configure_optimizers`` so ``get_param_dict`` is never called in child processes.  ``ddp_spawn`` forks
+    child processes that unpack a pickled copy of this module; patches applied in the parent process are not visible in
+    children, so the real ``get_param_dict`` would be called and would fail on ``_TinyModel`` (no ``.backbone``
+    attribute).
+
+    Must be defined at module level so ``pickle`` can look up the class by qualified name when deserialising in the
+    child process.
+    """
+
+    def configure_optimizers(self):
+        """Minimal single-group AdamW — bypasses get_param_dict."""
+        return torch.optim.AdamW(self.parameters(), lr=1e-4)
+
+
+class _MultiScaleCheckDDPModule(RFDETRModelModule):
+    """DDP-safe module that asserts on_train_batch_start mutation reaches training_step.
+
+    With multi_scale=True and _FakeDataset's 32×32 images, on_train_batch_start interpolates samples.tensors to a multi-
+    scale resolution (≥392 for RFDETRBaseConfig resolution=560).  This module raises AssertionError in training_step if
+    the tensor height is still 32, meaning the in-place NestedTensor mutation did not propagate through the PTL batch-
+    hook chain.
+
+    Must be defined at module level so pickle can look up the class by qualified name when ddp_spawn deserialises it in
+    the child process.
+
+    Regression guard for issue #952.
+    """
+
+    def configure_optimizers(self):
+        """Minimal single-group AdamW — bypasses get_param_dict."""
+        return torch.optim.AdamW(self.parameters(), lr=1e-4)
+
+    def training_step(self, batch, batch_idx):
+        """Assert resize from on_train_batch_start propagated before calling super."""
+        samples, _ = batch
+        h = samples.tensors.shape[2]
+        if h == 32:
+            raise AssertionError(
+                f"training_step received images at original 32-px height (h={h}). "
+                "on_train_batch_start's in-place NestedTensor mutation did not "
+                "propagate through the PTL hook chain. "
+                "Regression of issue #952: resize bypass in DDP batch-hook chain."
+            )
+        return super().training_step(batch, batch_idx)
+
+
+# ---------------------------------------------------------------------------
+# Multi-scale hook propagation tests (issue #952 regression)
+# ---------------------------------------------------------------------------
+
+
+class TestMultiScaleHookPropagation:
+    """on_train_batch_start resize must propagate to training_step via NestedTensor mutation.
+
+    _FakeDataset emits 32×32 images.  With multi_scale=True and RFDETRBaseConfig(resolution=560, patch_size=14,
+    num_windows=4) the computed scales start at 392, so none equal 32.  _MultiScaleCheckDDPModule raises AssertionError
+    in training_step if h==32, making trainer.fit() fail when the in-place mutation does not propagate.
+    """
+
+    def test_mutation_persists_to_training_step(self, base_model_config, base_train_config):
+        """Single-process: training_step must see resized tensors, not original 32×32."""
+        mc = base_model_config()
+        tc = base_train_config(multi_scale=True, use_ema=False, run_test=False)
+        fake_dataset = _FakeDataset(length=20)
+
+        with (
+            patch("rfdetr.training.module_model.build_model_from_config", return_value=_TinyModel()),
+            patch(
+                "rfdetr.training.module_model.build_criterion_from_config",
+                return_value=(_FakeCriterion(), _FakePostProcess()),
+            ),
+            patch("rfdetr.training.module_data.build_dataset", return_value=fake_dataset),
+        ):
+            module = _MultiScaleCheckDDPModule(mc, tc)
+            datamodule = RFDETRDataModule(mc, tc)
+            trainer = build_trainer(tc, mc, accelerator="cpu", fast_dev_run=2)
+            trainer.fit(module, datamodule=datamodule)
+
+
+# Windows CI currently cannot run this smoke test because gloo DDP spawn fails
+# with makeDeviceForHostname unsupported-device errors.
+@pytest.mark.skipif(sys.platform == "win32", reason="gloo DDP spawn unsupported on Windows CI")
+def test_ddp_spawn_fit_runs_without_error(base_model_config, base_train_config):
+    """ddp_spawn with 2 CPU workers must run fast_dev_run=2 without error.
+
+    ``ddp_spawn`` forks child processes, so all objects passed to ``trainer.fit()`` must be picklable.  ``MagicMock`` is
+    NOT picklable; this test uses ``_FakePostProcess``, plain dataset instances, and ``_DDPModule`` (module-level class)
+    instead.
+    """
+    mc = base_model_config()
+    tc = base_train_config(use_ema=False, run_test=False, devices=2, strategy="ddp_spawn")
+
+    fake_dataset = _FakeDataset(length=20)
+
+    with (
+        patch("rfdetr.training.module_model.build_model_from_config", return_value=_TinyModel()),
+        patch(
+            "rfdetr.training.module_model.build_criterion_from_config",
+            return_value=(_FakeCriterion(), _FakePostProcess()),
+        ),
+    ):
+        module = _DDPModule(mc, tc)
+
+    datamodule = RFDETRDataModule(mc, tc)
+    # Pre-set datasets: build_dataset mock doesn't survive the spawn boundary.
+    datamodule._dataset_train = fake_dataset
+    datamodule._dataset_val = fake_dataset
+
+    trainer = build_trainer(tc, mc, accelerator="cpu", fast_dev_run=2)
+    trainer.fit(module, datamodule=datamodule)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="gloo DDP spawn unsupported on Windows CI")
+def test_ddp_spawn_multi_scale_mutation_propagates(base_model_config, base_train_config):
+    """ddp_spawn with multi_scale=True must propagate on_train_batch_start resize to training_step.
+
+    _MultiScaleCheckDDPModule raises AssertionError in training_step when the NestedTensor height is still 32 (original
+    _FakeDataset size).  If trainer.fit() completes without error the PTL batch-hook reference chain is intact in DDP,
+    i.e. the in-place mutation in on_train_batch_start is visible in training_step on both workers.
+
+    Regression test for issue #952 on CPU DDP (non-Windows): confirms the transforms/resize propagation is not a
+    Windows-only concern.
+    """
+    mc = base_model_config()
+    tc = base_train_config(multi_scale=True, use_ema=False, run_test=False, devices=2, strategy="ddp_spawn")
+
+    fake_dataset = _FakeDataset(length=20)
+
+    with (
+        patch("rfdetr.training.module_model.build_model_from_config", return_value=_TinyModel()),
+        patch(
+            "rfdetr.training.module_model.build_criterion_from_config",
+            return_value=(_FakeCriterion(), _FakePostProcess()),
+        ),
+    ):
+        module = _MultiScaleCheckDDPModule(mc, tc)
+
+    datamodule = RFDETRDataModule(mc, tc)
+    # Pre-set datasets: build_dataset mock doesn't survive the spawn boundary.
+    datamodule._dataset_train = fake_dataset
+    datamodule._dataset_val = fake_dataset
+
+    trainer = build_trainer(tc, mc, accelerator="cpu", fast_dev_run=2)
+    trainer.fit(module, datamodule=datamodule)