Extend unmanaged gradient accumulation to ZeRO stage 2
Finalize averaged_gradients at step() after per-backward reduce/partition.
Still rejects ZeRO stage 3, ZeRO offload, and pipeline parallelism.
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
