add backward pass for spatial prallellism#993

Draft

mahf708 wants to merge 5 commits intoai2cm:mainfrom

E3SM-Project:spatial-parallel-training

Commits on Mar 10, 2026

Add test for backward step spatial parallelism
peterdschwartz
committed

Commits on Mar 19, 2026

add backward pass for spatial prallellism
mahf708
committed

Commits on Mar 20, 2026

Implement an autograd aware all_reduce op that doesn't double count in
peterdschwartz
committed
Merge branch 'sp-backward-step' into spatial-parallel-training
peterdschwartz
committed
add tests explicitly comparing gradients for backward step
peterdschwartz
committed