Sharding splits a variable or data object into distinct chunks. In multi-device training, each device receives a different chunk instead of a full copy. This matters when a model is too large to fit on one device. section 18.2.2 section 18.2.1
Split a variable across four devices
Suppose a variable has shape (32, 64). Split its last dimension into four chunks of shape (32, 16), and assign one chunk to each device along a device mesh’s "model" axis. The first dimension remains replicated. A LayoutMap expresses this choice with tuple(NULL, "model") for a Dense layer’s kernel; its bias uses tuple("model"). Set the ModelParallel distribution configuration before creating and fitting the model. 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")
These LayoutMap entries show how kernel and bias variable paths receive different layouts: the kernel’s last dimension and the bias’s only dimension follow the "model" mesh axis. section 18.2.2
A split can apply at different levels:
- Data parallelism divides a batch into sub-batches for parallel processing. Model parallelism instead divides one model among devices when a full copy will not fit. section 18.2.1
- Horizontal model partitioning assigns different layers to different devices. Passing a large layer output between devices can create a communication bottleneck. section 18.2.1
- Vertical model partitioning divides a layer’s variables across devices, computes each device’s part, and merges the outputs. It is the suggested fit for large models here. section 18.2.1
A file shard is a different, storage-level use of the term: a text dataset may arrive as separate files. For example, a pretraining dataset has 50 text shards of about 75 MB each. File shards do not, by themselves, specify how model variables are placed on devices. section 16.2 section 18.2.2
