Support ZeRO overlap_comm in unmanaged mode for stage 2
Stage 2 reduces/partitions gradients on every backward, and the epilogue
synchronizes the reduction stream before finalizing accumulated partitions,
so overlap_comm is safe. Relax the guard to allow overlap_comm with stage 2
while still rejecting it for stage 0/1 (where reduction is deferred to step()).
Add a stage-2 equivalence test (unmanaged vs managed with overlap_comm) and
update docs to describe the stage-2 support.
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
