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>