def _create_device_array(shape, device_type, host_id, local_device_ids=None):
"""Returns ID and device lists that can be used to create a mesh."""
num_global_devices = api.num_global_devices(device_type)
global_device_ids = np.arange(num_global_devices).reshape(shape)
local_device_list = api.local_devices(device_type)
# User can specify local_device_ids or use default list for multi host.
num_local_devices = len(local_device_list)
local_device_ids = [
x + host_id * num_local_devices for x in range(num_local_devices)
] if not local_device_ids else local_device_ids
return global_device_ids, local_device_ids, local_device_list
def create_distributed_mesh(mesh_dims: List[Tuple[str, int]],
mesh_name: str = '',
num_global_devices: Optional[int] = None,
num_clients: Optional[int] = None,
client_id: Optional[int] = None,
device_type: str = 'CPU') -> layout.Mesh:
"""Creates a single- or multi-client mesh.
For CPU and GPU meshes, users can choose to use fewer local devices than what
is available. If any argument is missing, it will be extracted from
environment variables. The default values for these environment variables
create a mesh using all devices (common for unit tests).
For TPU meshes, users should not specify any of the nullable arguments. The
DTensor runtime will set these arguments automatically, using all TPU cores
available in the entire cluster.
Args:
mesh_dims: A list of (dim_name, dim_size) tuples.
mesh_name: Name of the created mesh. Defaults to ''.
num_global_devices: Number of devices in the DTensor cluster. Defaults to
the corresponding environment variable.
num_clients: Number of clients in the DTensor cluster. Defaults to the
corresponding environment variable, DTENSOR_NUM_CLIENTS.
client_id: This client's ID. Defaults to the corresponding environment
variable, DTENSOR_CLIENT_ID.
device_type: Type of device to build the mesh for. Defaults to 'CPU'.
Returns:
A mesh created from specified or default arguments.
"""
dim_names, shape = zip(*mesh_dims)
if device_type.upper() in ['CPU', 'GPU']:
# For CPU and GPU meshes, user-specified args take precedence over env vars.
# This is particularly useful on single clients when users want to create
# meshes that use fewer logical devices than what's available.
if num_global_devices is None:
num_global_devices = api.num_global_devices(device_type)
if num_global_devices <= 0:
raise ValueError(
f'num_global_devices ({num_global_devices}) must be > 0')
if num_global_devices != np.prod(shape):
raise ValueError(
f'num_global_devices ({num_global_devices}) must be '
f'equal to total size of the mesh of shape {shape}')
if num_clients is None:
num_clients = api.num_clients()
if num_clients <= 0:
raise ValueError(f'num_clients ({num_clients}) must be > 0')
if _in_multi_client_mode is None and num_clients > 1:
raise ValueError(
'Invalid multi-client topology, run dtensor.initialize_multi_client() first'
)
if client_id is None:
client_id = api.client_id()
if client_id < 0:
raise ValueError(f'client_id ({client_id}) must be >= 0')
if client_id >= num_clients:
raise ValueError(
f'client_id ({client_id}) must be < {num_clients}')
if num_global_devices % num_clients != 0:
raise ValueError(
f'num_global_devices ({num_global_devices}) must be '
f'divisible by num_clients ({num_clients})')
num_local_devices = num_global_devices // num_clients
# It's allowed to create a CPU or GPU mesh using fewer logical devices than
# what's available. If so, just use the first N logical devices.
num_available_devices = api.num_local_devices(device_type)
if num_local_devices > num_available_devices:
raise ValueError(
f'Not enough devices; {num_local_devices} needed, '
f'only {num_available_devices} available')
local_devices = api.local_devices(device_type,
client_id)[:num_local_devices]
global_device_ids = np.arange(num_global_devices).reshape(shape)
flattened = np.ravel(global_device_ids).tolist()
start_idx = num_local_devices * client_id
local_device_ids = flattened[start_idx:start_idx + num_local_devices]
mesh = layout.Mesh(dim_names=dim_names,
global_device_ids=global_device_ids,
local_device_ids=local_device_ids,
local_devices=local_devices,
mesh_name=mesh_name)
_print_context(num_global_devices, num_clients, client_id, device_type,
mesh)
return mesh
if device_type.upper() == 'TPU':
# TPU meshes can only be configured through environment variables that
# reflect the actual TPU topology. Do not let users specify custom args.
if num_global_devices is not None:
raise ValueError(
f'Do not specify num_global_devices for {device_type.upper()} meshes. '
'It will be filled in automatically from environmental variables.'
'See api.py for the list of environmental variables for DTensor.'
)
if num_clients is not None:
raise ValueError(
f'Do not specify num_clients for {device_type.upper()} meshes. '
'It will be filled in automatically from environmental variables.'
'See api.py for the list of environmental variables for DTensor.'
)
if client_id is not None:
raise ValueError(
f'Do not specify client_id for {device_type.upper()} meshes. '
'It will be filled in automatically from environmental variables.'
'See api.py for the list of environmental variables for DTensor.'
)
mesh = tpu_util.create_tpu_mesh(dim_names, shape, mesh_name)
_print_context(api.num_global_devices(device_type), api.num_clients(),
api.client_id(), device_type, mesh)
return mesh
raise ValueError(f'Device type {device_type} is not CPU, GPU or TPU')
def dtensor_initialize_tpu_system(enable_coordination_service=False):
"""Initialize the TPU devices.
Args:
enable_coordination_service: If true, enable distributed coordination
service to make sure that workers know the devices on each other, a
prerequisite for data transfer through cross-worker rendezvous.
Raises:
RuntimeError: If running inside a tf.function.
NotFoundError: If no TPU devices found in eager mode.
"""
assert context.executing_eagerly()
in_multi_client_mode = api.job_name() != "localhost"
# Collective GRPC servers are only necessary in mutli-client setup.
# Single clients (e.g. Forge) can use local mode of collectives.
if in_multi_client_mode:
if api.jobs() is None:
raise ValueError(
"DTENSOR_JOBS environment variable is required when"
"using multi-client to properly set up communications between servers"
)
multi_client_util.initialize_multi_client_cluster(
job_name=api.job_name(),
dtensor_jobs=api.jobs(),
client_id=api.client_id(),
collective_leader=api.full_job_name(task_id=0),
enable_coordination_service=enable_coordination_service)
# Make sure the server change is fully propagated before attempting to run
# the core ID merging logic below.
context.ensure_initialized()
context.async_wait()
context.context()._clear_caches() # pylint: disable=protected-access
@function.defun
def _tpu_init_fn():
return gen_dtensor_ops.configure_and_initialize_global_tpu()
try:
with ops.device("/job:" + api.full_job_name() + "/device:TPU_SYSTEM:0"): # pylint: disable=protected-access
my_core_ids = _tpu_init_fn()
logging.info("TPU core IDs: %s", my_core_ids)
context.initialize_logical_devices()
# Configure virtual CPUs that is 1:1 mapped to TPU cores.
context.context().set_logical_cpu_devices(
len(api.local_devices(_TPU_DEVICE_TYPE)),
tf_device.DeviceSpec(
job=api.job_name(), replica=0, task=api.client_id()).to_string())
# `my_core_ids` contains the IDs of TPU cores attached to this host.
#
# To generate correct and efficient XLA AllReduce group assignment, we must
# merge these arrays from all hosts and broadcast the result back to all
# hosts, so all hosts can use these mappings in their MLIR passes.
#
# This is essentially doing what WaitForDistributedTpuOp and
# SetGlobalTPUArrayOp do, in our multi-client environment.
task_id = api.client_id()
num_tasks = api.num_clients()
num_devices = api.num_global_devices(_TPU_DEVICE_TYPE)
num_devices_per_task = int(num_devices / num_tasks)
# Create a one-time use mesh and layout just for merging core IDs.
mesh = layout_lib.Mesh([_MESH_DIM_X],
*_create_device_array((num_devices,),
_TPU_DEVICE_TYPE,
api.client_id()))
layout = layout_lib.Layout([_MESH_DIM_X, layout_lib.UNSHARDED], mesh)
device = dtensor_device.DTensorDevice(meshes=[mesh])
logging.info("TPU core locations: %s",
device.tpu_core_ids_to_locations(my_core_ids))
# At this point, we don't know which cores are attached to other hosts.
# The core ID mappings in the runtime haven't been set yet.
#
# The core ID merging AllReduce below is carefully written so it works
# without needing correct core mappings to be set in the runtime. We will
# use this AllReduce's result to set the core ID mappings, and all future
# user-initiated AllReduces will use the mappings.
#
# The runtime is hard-coded to ignore core ID mappings on this AllReduce.
all_core_ids = np.zeros([num_devices], dtype=np.int32)
for i in range(len(my_core_ids)):
all_core_ids[task_id * num_devices_per_task + i] = my_core_ids[i]
# Only one local device gets valid input: 8 local core IDs among
# (num_tasks - 1) * 8 zeros. The 8 core IDs are set using task ID as offset.
# The other 7 local devices get zero inputs. All devices on all host
# participate in one AllReduce, whose result will be core IDs arranged by
# task-device ordinals.
all_core_ids = constant_op.constant([all_core_ids])
zeros = array_ops.zeros_like(all_core_ids)
all_core_ids = [all_core_ids] + [zeros] * (num_devices_per_task - 1)
with ops.device(device.name):
all_core_ids = device.pack(all_core_ids, layout)
all_core_ids = math_ops.reduce_sum(all_core_ids, axis=[0])
unpacked_all_tpu_ids = device.unpack(all_core_ids)
all_core_ids = list(unpacked_all_tpu_ids[0].numpy())
logging.info("All TPU core IDs: %s", all_core_ids)
# Set the default core ID mappings in the runtime for legacy code and tests.
#
# Legacy code and tests create TPU meshes directly without using the
# `create_tpu_mesh` function below. Those meshes have global device IDs
# equal to TF task-device ordinals. The `all_core_ids` array happens to
# arrange core IDs by TF task-device ordinals. Using this array on those
# meshes guarantee correct although inefficient results.
device.set_tpu_core_ids("", all_core_ids)
# Remember enough global, immutable information to be able to build any ring
# we want prescribed by `create_tpu_mesh` in the future.
global _all_core_ids
_all_core_ids = all_core_ids
all_core_locations = device.tpu_core_ids_to_locations(all_core_ids)
all_core_locations = [
_CoreLocation(l[0], l[1], l[2], l[3]) for l in all_core_locations
]
global _all_core_locations
_all_core_locations = all_core_locations
logging.info("All TPU core locations: %s", all_core_locations)
tpu_topology = _create_tpu_topology(all_core_locations, num_tasks,
num_devices_per_task)
global _tpu_topology
_tpu_topology = tpu_topology
logging.vlog(1, "TPU Topology: %s, %s", tpu_topology.mesh_shape,
tpu_topology.device_coordinates)
global _dtensor_device
_dtensor_device = device
context.async_wait()
except errors.InvalidArgumentError as e:
raise errors.NotFoundError(
None, None, "Initialization failed, no valid TPUs found. " + str(e))
except errors.InternalError as e:
logging.error("Hit internal error during TPU system initialization. "
+ "It is likely hareware failure. \nPlease check the error "
+ "messages above to see whether that's the case. \nIf so, "
+ "consider to restart the job or try another machine.")
raise e
# Optionally exchange heartbeats between workers every minute.
if in_multi_client_mode and api.heartbeat_enabled():
logging.info(
"Starting DTensor heartbeat service exchanging signals every 10 minutes"
)
heartbeat.start(period=180)
# Clear out the eager context caches since the memory is invalid now.
logging.info("Clearing out eager caches")
context.context()._clear_caches() # pylint: disable=protected-access