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

[JESUSROYETH]

What is wrong

Keypoint training uses manual optimization to normalise the loss over the full accumulated box count. Under DDP, manual_backward() currently runs with gradient synchronisation enabled on every microbatch, then _should_step_optimizer() decides if the optimizer can step. So with grad_accum_steps=4, each accumulation window performs four gradient reductions when one is enough.

This is also the limitation documented when keypoint DDP support landed in #1232, and it is the DDP training question left open in #1410.

What changes

The backward pass now uses Lightning's own LightningOptimizer.toggle_model(sync_grad=...) context. Intermediate microbatches run under DDP no_sync(), while the backward that closes a full or partial accumulation window synchronises the accumulated gradient before the optimizer step.

There is no new option or fallback. Detection and segmentation keep the automatic-optimization path, and grad_accum_steps=1 still synchronises every backward. The warning in build_trainer() and the multi-GPU training note are updated to match the new behaviour.

Two-L4 result

Measured with RFDETRKeypointPreview on real COCO 2017 person-keypoint images, two NVIDIA L4 GPUs, BF16, batch size 1 per rank, two-rank NCCL DDP, grad_accum_steps=4, torch 2.10.0+cu130 / CUDA 13.0 / driver 580.173.02. Each trial ran 64 microbatches per rank, discarded the first 8 and timed the remaining 56 with torch.cuda.synchronize() at both boundaries. The arms were counterbalanced.

	develop median [range], n=5	this PR median [range], n=5	change
steady-state seconds / 56 microbatches	20.4620 [20.2320, 20.5213]	19.4276 [19.2765, 19.7102]	-5.06%
seconds / optimizer step	1.46157 [1.44514, 1.46581]	1.38769 [1.37689, 1.40787]	-5.06%
full Trainer.fit seconds	26.2290 [25.8738, 26.3033]	25.0748 [24.7415, 25.2746]	-4.40%
peak CUDA allocation	3,446,091,264 bytes	3,446,091,264 bytes	unchanged

All ten performance trials completed and the steady-state ranges do not overlap.

Correctness and validation

The new parametrised test covers the three synchronisation decisions: an intermediate microbatch, a complete accumulation window, and the partial final-window flush. It failed on develop because toggle_model() was never called, then passed with the change. The existing large-effective-batch tests continue to check the accumulated gradient against one large batch across six box-count/accumulation cases.

A deterministic two-device Trainer.fit check also ran two optimizer steps with exactly representable float64 gradients. Two baseline and two candidate runs produced the same parameter bytes (bfaed4915cfeef3f). The real model's full state was not bitwise repeatable even baseline against itself, so I am not using that noisy comparison as a parity claim ..

• pytest tests/training/test_module_model.py -q: 197 passed, 1 skipped
• affected training shards on the L4 VM: 350 passed, 6 skipped
• pre-commit run --all-files: all hooks passed

This does not include training to convergence or an mAP claim. It also does not generalise the 5.06% number beyond two L4 GPUs and grad_accum_steps=4; more GPUs, multi-node DDP and other interconnects were not measured.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87%. Comparing base (c301e47) to head (001e2dc).
⚠️ Report is 1 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff           @@
##           develop   #1443   +/-   ##
=======================================
  Coverage       87%     87%           
=======================================
  Files          117     117           
  Lines        15260   15264    +4     
=======================================
+ Hits         13214   13218    +4     
  Misses        2046    2046           

🚀 New features to boost your workflow:

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