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: Achyuthan-S

deepspeed/checkpoint/affine.py

@@ -356,11 +356,16 @@ def _row_major_strides(shape):
     return tuple(strides)
 
 
-def replicated_map(shape, tp_degree):
+def replicated_map(shape, tp_degree, scale=1.0):
     """Every rank holds the whole parameter.
 
     One piece, named by every rank, which is what lets a converter read it from whichever
     rank is cheapest rather than from a designated owner.
+
+    ``scale`` belongs here rather than on a split, because the layouts that pre-divide a
+    value hold it whole on every rank: a row-parallel layer replicates its bias divided by
+    the world size so that summing the all-reduced outputs adds it exactly once. The weight
+    beside it is split and unscaled.
     """
     shape = tuple(shape)
     strides = _row_major_strides(shape)
@@ -371,7 +376,8 @@ def replicated_map(shape, tp_degree):
                     source_strides=strides,
                     dest_offset=0,
                     dest_strides=strides,
-                    locations=ranks)
+                    locations=ranks,
+                    scale=scale)
     ]
     return ParamAffineMap(logical_shape=shape,
                           shard_shapes={rank: shape
@@ -380,7 +386,7 @@ def replicated_map(shape, tp_degree):
                                           for rank in ranks})
 
 
-def contiguous_split_map(shape, per_rank_sizes, partition_dim, scale=1.0):
+def contiguous_split_map(shape, per_rank_sizes, partition_dim):
     """Each rank holds one contiguous block along ``partition_dim``.
 
     Covers row-parallel and column-parallel layers alike: they differ only in which axis
@@ -402,8 +408,7 @@ def contiguous_split_map(shape, per_rank_sizes, partition_dim, scale=1.0):
                         source_strides=source_strides,
                         dest_offset=0,
                         dest_strides=_row_major_strides(shard_shape),
-                        locations=[rank],
-                        scale=scale)
+                        locations=[rank])
         ]
         start += size
     return ParamAffineMap(logical_shape=shape, shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank)

tests/unit/checkpoint/test_affine_shard_map.py

@@ -567,3 +567,21 @@ def test_unscaled_optimizer_state_still_converts():
     shards = {rank: torch.ones(4, dtype=torch.float64) for rank in range(2)}
     for scale_power in (1, -1, -2):
         assert torch.equal(affine_map.rebuild(shards, scale_power), torch.ones(4, dtype=torch.float64))
+
+
+def test_replicated_map_carries_the_scale():
+    """The layouts that pre-divide a value replicate it whole, so the scale belongs here.
+
+    A row-parallel layer divides its bias by the world size and gives every rank the whole
+    thing, so summing the all-reduced outputs adds the bias once. The weight beside it is
+    split and unscaled, which is why the split constructors take no scale.
+    """
+    world_size = 4
+    full_bias = torch.randn(5, dtype=torch.float64)
+    affine_map = replicated_map((5, ), world_size, scale=1.0 / world_size)
+    affine_map.validate_coverage()
+
+    shards = {rank: affine_map.extract(full_bias, rank) for rank in range(world_size)}
+    assert torch.equal(shards[0], full_bias / world_size)
+    assert torch.equal(affine_map.rebuild(shards), full_bias)
+    assert affine_map.pieces_by_rank[0][0].locations == frozenset(range(world_size))