jax.experimental.mesh_utils.create_hybrid_device_mesh#

jax.experimental.mesh_utils.create_hybrid_device_mesh(mesh_shape, dcn_mesh_shape, devices=None, *, process_is_granule=False, should_sort_granules_by_key=True, allow_split_physical_axes=False)[源代码]#

为混合（例如，ICI 和 DCN）并行创建设备网格。

参数:

  • mesh_shape (Sequence[int]) – 更快/内部网络的逻辑网格的形状，按网络强度递增排序，例如 [副本，数据，模型]，其中模型具有最多的网络通信需求。

  • dcn_mesh_shape (Sequence[int]) – 较慢/外部网络的逻辑网格的形状，与 mesh_shape 的顺序相同。

  • devices (Sequence[Any] | None) – （可选）要为其构造网格的设备。默认为 jax.devices()。

  • process_is_granule (bool) – 如果为 True，此函数将进程视为较慢/外部网络的单位。 否则，它将在设备上查找 slice_index 属性，并将切片用作单位。 启用此选项旨在作为不设置 slice_index 的平台的后备方案。

  • should_sort_granules_by_key (bool) – 是否应按粒度键（切片或进程索引，取决于 process_is_granule）对设备粒度进行排序。

  • allow_split_physical_axes (bool) – 如果为 True，如有必要，我们将拆分物理轴以生成所需的设备网格。

引发:

ValueError – 如果 devices 所属的切片的数量不等于 dcn_mesh_shape 的乘积，或者如果属于任何单个切片的设备的数量不等于 mesh_shape 的乘积。

返回:

一个 JAX 设备的 np.ndarray，其形状为 mesh_shape * dcn_mesh_shape，可以将其馈送到 jax.sharding.Mesh 以进行混合并行处理。

返回类型:

np.ndarray