perf: disable find_unused_parameters for faster DDP training

[Luodian]

Summary

• Disable find_unused_parameters in DDP wrapper for performance improvement
• Add documentation explaining the reasoning and how to debug if issues arise

Problem

find_unused_parameters=True adds overhead:

• Extra autograd graph traversal after every forward pass
• Reduced comm/compute overlap
• CPU overhead for large models

Analysis

Code review confirms this training loop does NOT have unused parameters:

Check	Result
All backbone outputs flow to loss?	✅ Yes, via pfc() for all strategies
Branching skips modules?	❌ No, branching is input shape/indexing only
Optional modules exist?	Pooling head (use_head), but always used in training

What find_unused_parameters does

Forward → DDP traverses autograd graph → Marks unused params "ready"

If no params are unused, this traversal is pure overhead.

What static_graph=True does

Caches the parameter usage pattern after warmup, mitigating some overhead.
But if no params are unused, find_unused_parameters=False is still cleaner.

Testing Instructions

If this change causes errors:

# Error you might see:
RuntimeError: Expected to have finished reduction in the prior iteration...

# Debug with:
TORCH_DISTRIBUTED_DEBUG=INFO torchrun ... --max_steps 10

The debug output will show which parameters didn't receive gradients.

Rollback

If issues occur, simply change line to:

find_unused_parameters=True,  # Re-enabled due to [reason]

Files Changed

• training/train.py - DDP configuration with documentation