perf(training): synchronize keypoint DDP once per optimizer step

CHANGELOG.md

@@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Changed
 
+- Multi-GPU keypoint training with `grad_accum_steps > 1` now synchronizes gradients once per optimizer step instead of once per microbatch, avoiding redundant DDP reductions.
+
 ### Deprecated
 
 ### Fixed

docs/learn/train/advanced.md

@@ -249,7 +249,7 @@ model.train(
     dataset_dir="path/to/keypoint-dataset",
     epochs=100,
     batch_size=2,  # per-GPU batch size
-    grad_accum_steps=1,  # recommended on multi-GPU — see note below
+    grad_accum_steps=1,  # see note below for multi-GPU accumulation behavior
     lr=1e-4,
     output_dir="output",
     devices="auto",  # or devices=8
@@ -260,9 +260,9 @@ model.train(
 torchrun --nproc_per_node=8 train_pose.py
 ```
 
-!!! note "Prefer `grad_accum_steps=1` on multi-GPU for keypoints"
+!!! note "Gradient accumulation on multi-GPU keypoint training"
 
-    Keypoint models use **manual optimization** so the per-step box-count loss normalization is computed over the full accumulated batch. As a result, gradients synchronize on **every** microbatch rather than only at the end of an accumulation window. Training with `grad_accum_steps > 1` on multiple GPUs is still numerically correct, but performs one `all_reduce` per microbatch (i.e. `grad_accum_steps`× the necessary communication). For best throughput, scale with more GPUs / a larger per-GPU `batch_size` and keep `grad_accum_steps=1`.
+    Keypoint models use **manual optimization** so the per-step box-count loss normalization is computed over the full accumulated batch. Intermediate microbatches accumulate gradients locally on each rank. The backward pass that closes the accumulation window synchronizes the full accumulated gradient before the optimizer step, avoiding redundant DDP reductions while preserving full-effective-batch normalization.
 
     Sharded strategies (FSDP / DeepSpeed) are **not** supported for keypoint models — use `ddp` (or `strategy="auto"` with `devices > 1`).
 

src/rfdetr/training/module_model.py

@@ -12,6 +12,7 @@
 import math
 import random
 import warnings
+from contextlib import nullcontext
 from typing import Any, Callable, cast
 
 import torch
@@ -622,8 +623,17 @@ def training_step(self, batch: tuple[Any, Any], batch_idx: int) -> Tensor | dict
             # loss_for_backward is only None in the automatic-optimization branch above,
             # which is mutually exclusive with _use_manual_optimization.
             assert loss_for_backward is not None
-            self.manual_backward(loss_for_backward)
-            if self._should_step_optimizer(batch_idx):
+            should_step = self._should_step_optimizer(batch_idx)
+            # LightningOptimizer maps sync_grad=False to DDP's no_sync context. Intermediate
+            # microbatches accumulate locally; the closing backward reduces the whole window.
+            sync_context = (
+                optimizer.toggle_model(sync_grad=should_step)
+                if isinstance(optimizer, LightningOptimizer)
+                else nullcontext()
+            )
+            with sync_context:
+                self.manual_backward(loss_for_backward)
+            if should_step:
                 self._step_optimizer(optimizer)
         if self.train_config.compute_train_metrics:
             with torch.no_grad():

src/rfdetr/training/trainer.py

@@ -680,9 +680,7 @@ def _resolve_precision() -> str:
                 )
         _logger.info(
             "Keypoint model + distributed execution (strategy=%r, devices=%r, num_nodes=%r) → "
-            "DDP with manual optimization. For best throughput on multi-GPU keep grad_accum_steps=1: "
-            "the manual-optimization path synchronizes gradients on every microbatch, so "
-            "grad_accum_steps>1 is correct but performs redundant all-reduces.",
+            "DDP with manual optimization. Accumulated gradients synchronize only when the optimizer steps.",
             strategy,
             devices,
             num_nodes,

tests/training/test_module_model.py

@@ -8,12 +8,14 @@
 import logging
 import random
 import warnings
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
 import torch
 from pytorch_lightning import Callback, Trainer
+from pytorch_lightning.core.optimizer import LightningOptimizer
 from torch import nn
 
 from rfdetr.config import RFDETRBaseConfig, RFDETRSmallConfig, TrainConfig
@@ -1023,6 +1025,32 @@ def test_loss_backward_uses_box_normalizer_contract(self, tmp_path):
         backward_loss = module.manual_backward.call_args.args[0]
         assert backward_loss.item() == pytest.approx(1.0)
 
+    @pytest.mark.parametrize(
+        "grad_accum_steps,num_training_batches,batch_idx,sync_grad",
+        [(2, 2, 0, False), (2, 4, 1, True), (4, 2, 1, True)],
+    )
+    def test_keypoint_accumulation_syncs_only_when_optimizer_steps(
+        self, tmp_path, grad_accum_steps, num_training_batches, batch_idx, sync_grad
+    ):
+        """Keypoint DDP must skip gradient synchronization until an accumulation window closes."""
+        keypoint_config = _base_model_config(use_grouppose_keypoints=True, num_keypoints_per_class=[17])
+        module, samples, targets, _, _ = self._run_step(
+            tmp_path,
+            accumulate_grad_batches=grad_accum_steps,
+            model_config=keypoint_config,
+        )
+        optimizer = MagicMock(spec=LightningOptimizer)
+        optimizer.param_groups = [{"lr": 1e-3}]
+        optimizer.toggle_model.return_value = nullcontext()
+        optimizer.zero_grad = MagicMock()
+        module.optimizers.return_value = optimizer
+        module._trainer.num_training_batches = num_training_batches
+
+        module.training_step((samples, targets), batch_idx=batch_idx)
+
+        optimizer.toggle_model.assert_called_once_with(sync_grad=sync_grad)
+        module.manual_backward.assert_called_once()
+
     def test_detection_loss_uses_lightning_grad_accum_scaling(self, tmp_path):
         """Detection (automatic optimization) divides loss by ``trainer.accumulate_grad_batches`` so the returned loss
         matches the legacy non-manual training path."""