Extend unmanaged gradient accumulation to ZeRO stage 3
Allow managed_gradient_accumulation=False for ZeRO stage 3 (non-offload).
Stage 3 partitions gradients into __param_id_to_grad_partition on every
backward (as in managed mode); step() only marks the boundary via
finalize_gradient_accumulation_boundary(). Since stage 3 is now permitted,
also reject ZeRO parameter offload (previously unreachable). ZeRO offload
(optimizer state and parameter) remains rejected pending the follow-up PR;
pipeline parallelism, DeepCompile, and Apex AMP stay unsupported.
Validated on a 2-GPU node (full -k Unmanaged suite, 29 passed) incl.
test_unmanaged_matches_managed[3] and test_unmanaged_varying_backward_count[3].
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
