tf.keras.distribution.ModelParallel

tf.keras.distribution.ModelParallel(
    device_mesh, layout_map, batch_dim_name=None
)

devices = list_devices()    # Assume there are 8 devices.

# Create a mesh with 2 devices for data parallelism and 4 devices for
# model parallelism.
device_mesh = DeviceMesh(shape=(2, 4), axis_names=('batch', 'model'),
                         devices=devices)
# Create a layout map that shard the `Dense` layer and `Conv2D`
# layer variables on the last dimension.
# Based on the `device_mesh`, this means the variables
# will be split across 4 devices. Any other variable that doesn't
# match any key in the layout map will be fully replicated.
layout_map = LayoutMap(device_mesh)
layout_map['dense.*kernel'] = (None, 'model')
layout_map['dense.*bias'] = ('model',)
layout_map['conv2d.*kernel'] = (None, None, None, 'model')
layout_map['conv2d.*bias'] = ('model',)

distribution = ModelParallel(device_mesh=device_mesh,
                             layout_map=layout_map,
                             batch_dim_name='batch')
# Set the global distribution, or via `with distribution.scope():`
set_distribution(distribution)

model = model_creation()
model.compile()
model.fit(data)

# With only the shape change for the device mesh, the variables will be
# sharded across 8 devices instead of 4, which further reduces the memory
# footprint of variables on each of the device.
device_mesh = DeviceMesh(shape=(1, 8), axis_names=('batch', 'model'),
                         devices=devices)

model = create_model()
for v in model.variables:
    print(v.path)

distribute_dataset(
    dataset
)

get_data_layout(
    data_shape
)

get_tensor_layout(
    path
)

get_variable_layout(
    variable
)

@contextlib.contextmanager
scope()