Replication copies one variable value across an axis of a device mesh, so every device along that axis sees the same value. It differs from sharding, which splits a variable into separate chunks. section 18.2.2
Replication decisions apply independently to each variable dimension. In a two-by-four mesh, devices can support data parallelism along one axis while split model computation is replicated across the other. section 18.2.2
Use replication when mapping variables to a distributed Keras model:
- Create a LayoutMap associated with the device mesh. section 18.2.2
- Describe each variable using its path and a tuple with one entry per dimension. section 18.2.2
- Combine replicated model copies with sub-batches when using data parallelism. section 18.2.2
layout_map <- keras$distribution$LayoutMap(device_mesh)
layout_map["sequential/dense/kernel"] <- tuple(NULL, "model")
layout_map["sequential/dense/bias"] <- tuple("model")
layout_map["sequential/dense_1/kernel"] <- tuple(NULL, "model")
layout_map["sequential/dense_1/bias"] <- tuple("model")
Inspect a LayoutMap that assigns layouts to model variables by path. section 18.2.2
