Fix: stabilize test after distributed training completion

[stop1one]

Description

Related Issues

Fixes #316 #374

Summary of Changes

This PR fixes two problems for distributed training when run_test=True:

1. Synchronization before testing

   • Added torch.distributed.barrier() when args.distributed is enabled, right after training and before testing.
   • Ensures all processes wait until rank 0 has finished saving checkpoint_best_total.pth before moving on to testing.

2. Correct checkpoint loading in DDP

   • Changed from model.load_state_dict(best_state_dict) to model_without_ddp.load_state_dict(best_state_dict)
   • Prevents key mismatch errors when loading checkpoints in distributed environments.

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)

How has this change been tested, please provide a testcase or example of how you tested the change?

train.py :

from rfdetr import RFDETRLarge

model = RFDETRLarge()

model.train(
    dataset_dir=<DATASET_PATH>,
    epochs=30,
    batch_size=8
    grad_accum_steps=2,
    lr=1e-4,
    output_dir=<OUTPUT_PATH>,
)

python -m torch.distributed.launch --nproc_per_node=4 --use_env train.py

[CLAassistant]

All committers have signed the CLA.

[stop1one]

Hi @probicheaux ,
I just realized I hadn’t signed the CLA earlier — that’s done now ✅
Would you mind taking another look at this PR when you have a moment?
Thanks again for your time!

[Copilot]

Pull request overview

This PR fixes distributed training synchronization issues when running tests after training completion. The changes ensure proper checkpoint loading and process synchronization in multi-GPU environments.

Changes:

• Added a distributed barrier before test execution to ensure all processes wait for rank 0 to complete checkpoint saving
• Corrected checkpoint loading to use model_without_ddp instead of model to avoid DDP wrapper key mismatches

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.