Replication copies a variable across an axis of a device mesh so every device along that axis sees the same value. Unlike sharding, it does not divide the variable into chunks. section 18.2.2
A replicated variable
For a variable shaped (32, 64), replication copies the complete value across a mesh axis. Sharding instead creates four chunks shaped (32, 16), with each device receiving a different chunk. section 18.2.2
Use these practical rules when assigning layouts:
- Choose replication or sharding independently for each variable dimension. section 18.2.2
- Describe the choice with LayoutMap, which maps a variable path to one layout entry per dimension. section 18.2.2
- Use replication alongside data parallelism or model parallelism when different device axes serve different roles. section 18.2.2
layout_map = keras.distribution.LayoutMap(device_mesh) layout_map["sequential/dense/kernel"] = (None, "model") layout_map["sequential/dense/bias"] = ("model",) layout_map["sequential/dense_1/kernel"] = (None, "model") layout_map["sequential/dense_1/bias"] = ("model",)
Inspect a concrete LayoutMap that assigns mesh dimensions to model kernels and biases. section 18.2.2
