DeepSpeed
bce0adf0 - Carry the affine scale on the replicated map, not the split (#8477)

Commit
1 day ago
Carry the affine scale on the replicated map, not the split (#8477) Follow-up to #8385. The scale move answering @delock's question was pushed after the merge queue had already snapshotted the branch, so it did not land with the rest. The layouts that pre-divide a value hold it whole on every rank: Yuan's o_proj and the last conv layer both replicate the bias divided by the world size, so the all-reduced sum adds it exactly once. The weight beside them is what gets split, and it is unscaled — so no in-tree layout scales a split, and the split constructors no longer take the argument. Adds a test covering the replicated case. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Author
Parents
Loading