use DDP.no_sync when grad_accum and DDP

[oclivegriffin]

When using DDP and gradient accumulation at the same time, we should use the DDP.no_sync context manager to get some free training speed.

https://chatgpt.com/share/e/68efa990-9ee8-800c-99da-b078c3d2ac7b

[oclivegriffin]

@claude can you do this?

[claude]

Claude finished @oli-clive-griffin's task —— View job • claude/issue-215-20251015-1407 • Create PR ➔

---

Implementing DDP.no_sync optimization ✅

I've successfully implemented the DDP.no_sync optimization for gradient accumulation. This prevents unnecessary gradient synchronization during accumulation steps and only syncs on the final step, providing free training speed improvements.

Todo List:

• Set up environment and analyze codebase
• Find current DDP and gradient accumulation implementation
• Implement DDP.no_sync optimization
• Test the changes
• Commit and push changes

Implementation Details:

• Modified gradient accumulation loop in spd/run_spd.py:229-286 to use DDP.no_sync() context manager
• Only skips synchronization when using DDP and not on the last accumulation step
• Uses nullcontext() as a no-op context manager for non-DDP or final step cases
• Added comprehensive test for gradient accumulation > 1 to verify the optimization works

Changes Made:

1. Added imports: contextlib.nullcontext for proper context manager handling
2. Modified gradient accumulation loop: Wrapped accumulation steps with conditional no_sync() context
3. Added test: New test method test_distributed_with_gradient_accumulation to verify deterministic behavior with gradient accumulation

The optimization follows the exact pattern shown in the issue image - using DDP.no_sync() when is_last_grad_accum_step is False.

---