wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/dtensor/python/BUILD
+341
View File
@@ -0,0 +1,341 @@
# DTensor Python API and libraries.
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
default_visibility = [
"//tensorflow/dtensor:dtensor-internal",
"//third_party/py/jax_tpu_embedding:__subpackages__",
]
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = default_visibility,
licenses = ["notice"],
)
# -----------------------------------------------------------------------------
# A list of all modules linked into TensorFlow.
py_library(
name = "dtensor",
srcs = ["__init__.py"],
strict_deps = True,
visibility = [
"//tensorflow/python:__pkg__",
"//tensorflow/python/tools/api/generator:__pkg__",
],
deps = [
":accelerator_util",
":api",
":config",
":d_checkpoint",
":d_variable",
":input_util",
":layout",
":mesh_util",
":save_restore",
":tpu_util",
],
)
# -----------------------------------------------------------------------------
# Implementations of the public API.
pytype_strict_library(
name = "api",
srcs = ["api.py"],
deps = [
":dtensor_device",
":gen_dtensor_ops",
":layout",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_library(
name = "config",
srcs = ["config.py"],
deps = [
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:device",
"//tensorflow/python/util:tf_export",
],
)
tf_gen_op_wrapper_py(
name = "gen_dtensor_ops",
out = "gen_dtensor_ops.py",
extra_py_deps = [
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
py_lib_rule = py_library,
deps = [
"//tensorflow/dtensor/cc:dtensor_ops",
"//tensorflow/dtensor/cc:dtensor_tpu_ops",
],
)
pytype_strict_library(
name = "layout",
srcs = ["layout.py"],
deps = [
"//tensorflow/dtensor/proto:layout_proto_py_pb2",
"//tensorflow/python:_pywrap_dtensor_device",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "d_random",
srcs = ["d_random.py"],
deps = [
":api",
":layout",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/ops:shape_util",
"//tensorflow/python/ops:stateless_random_ops_gen",
],
)
pytype_strict_library(
name = "d_variable",
srcs = ["d_variable.py"],
deps = [
":api",
":layout",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_library(
name = "d_checkpoint",
srcs = ["d_checkpoint.py"],
deps = [
":api",
":d_variable",
":gen_dtensor_ops",
":layout",
":save_restore",
"//tensorflow/core:protos_all_py",
"//tensorflow/dtensor/proto:layout_proto_py_pb2",
"//tensorflow/python/checkpoint",
"//tensorflow/python/checkpoint:checkpoint_options",
"//tensorflow/python/checkpoint:graph_view",
"//tensorflow/python/checkpoint:restore",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/trackable:base",
"//tensorflow/python/trackable:data_structures",
"//tensorflow/python/training:py_checkpoint_reader",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_library(
name = "numpy_util",
srcs = ["numpy_util.py"],
deps = [
":api",
":layout",
"//tensorflow/python/eager/polymorphic_function",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/types:core",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "save_restore",
srcs = ["save_restore.py"],
deps = [
":api",
":d_variable",
":gen_dtensor_ops",
":layout",
":mesh_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:io_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/util:tf_export",
],
)
# -----------------------------------------------------------------------------
# The DTensor runtime.
pytype_strict_library(
name = "dtensor_device",
srcs = ["dtensor_device.py"],
deps = [
":config",
":gen_dtensor_ops",
":layout",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:_pywrap_dtensor_device",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:core",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/util:_pywrap_utils",
"//third_party/py/numpy",
],
)
# -----------------------------------------------------------------------------
# Utilities.
pytype_strict_library(
name = "mesh_util",
srcs = ["mesh_util.py"],
visibility = default_visibility + [
"//tensorflow/dtensor:dtensor-users",
],
deps = [
":accelerator_util",
":api",
":config",
":layout",
":tpu_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:tfrt_utils",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
"@absl_py//absl/flags",
"@absl_py//absl/logging",
],
)
# TODO(b/245589661): Split accelerator_util to its module after
# The circular dependence is removed with dtensor_initialize_tpu_system.
pytype_strict_library(
name = "tpu_util",
srcs = [
"accelerator_util.py",
"tpu_util.py",
],
visibility = default_visibility + [
"//tensorflow/dtensor:dtensor-users",
],
deps = [
":config",
":dtensor_device",
":gen_dtensor_ops",
":layout",
"//tensorflow/core/protobuf:for_core_protos_py_proto",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:remote_utils",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/tpu:topology",
"//tensorflow/python/util:numpy_compat",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
"@absl_py//absl/logging",
],
)
pytype_strict_library(
name = "heartbeat",
srcs = ["heartbeat.py"],
deps = [
":api",
":config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:collective_ops",
"//tensorflow/python/platform:tf_logging",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "accelerator_util",
srcs = [],
deps = [
":tpu_util",
],
)
pytype_strict_library(
name = "input_util",
srcs = ["input_util.py"],
deps = [
":api",
":config",
":layout",
"//tensorflow/python/data/experimental/ops:data_service_ops",
"//tensorflow/python/data/experimental/ops:distribute",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/types:data",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
],
)
tensorflow/dtensor/python/__init__.py
+17
View File
@@ -0,0 +1,17 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DTensor Python API."""
# This file is left empty intentionally.
tensorflow/dtensor/python/accelerator_util.py
+298
View File
@@ -0,0 +1,298 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility for working with accelerator systems."""
from typing import List, Optional
from absl import logging
from tensorflow.core.protobuf import cluster_pb2
from tensorflow.core.protobuf import tensorflow_server_pb2
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import tpu_util
from tensorflow.python.eager import context
from tensorflow.python.framework import config as tf_config
from tensorflow.python.platform import remote_utils
from tensorflow.python.util.tf_export import tf_export
_INITIALIZED_ACCELERATOR_SYSTEM_TYPE = None
def is_initialized() -> bool:
"""Returns whether accelerator system has been initialized."""
return bool(_INITIALIZED_ACCELERATOR_SYSTEM_TYPE)
def set_initialized(value):
"""Sets if accelerator system has been initialized."""
global _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
_INITIALIZED_ACCELERATOR_SYSTEM_TYPE = value
def initialize_multi_client_cluster(job_name: str,
dtensor_jobs: List[str],
client_id: int,
collective_leader: str,
port: Optional[int] = None,
gpu_use_nccl_communication: bool = False,
enable_coordination_service: bool = True):
"""Initialize GRPC servers and collectives for multi-client DTensor setup.
This function can be used to initialize a multi-client cluster and enable
collective ops. GRPC servers are necessary in the multi-client mode, even
when the number of clientis is 1.
NOTE: this function must be called in an eager context.
Args:
job_name: The job name used by all clients in the DTensor cluster.
dtensor_jobs: A list of the DTensor client jobs participating in the
cluster. Must be strings of the form "hostname:port".
client_id: The ID of the DTensor client this function is being called in.
collective_leader: The job/task that will be used to run collectives.
port: The port this client's GRPC server will run on. If omitted, use the
port from dtensor_jobs for this client.
gpu_use_nccl_communication: if True, configure TensorFlow to use NCCL by
default.
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.
"""
assert context.executing_eagerly()
if not collective_leader.startswith("/job:"):
collective_leader = "/job:" + collective_leader
context.context().configure_collective_ops(
use_nccl_communication=gpu_use_nccl_communication,
collective_leader=collective_leader)
if enable_coordination_service:
context.context().configure_coordination_service(
service_type="standalone", service_leader=collective_leader)
config_proto = context.get_config()
# Construct server def from the host directly instead of relying on
# TF_CONFIG.
cluster_def = cluster_pb2.ClusterDef()
# Note that for bns addresses, we will currently rely on the sorted string
# of job name as the order of assigning task ids. This might be brittle once
# we have jobs across multiple cells.
cluster_def.job.add(name=job_name, tasks=dict(enumerate(dtensor_jobs)))
server_def = tensorflow_server_pb2.ServerDef(
cluster=cluster_def,
default_session_config=config_proto,
job_name=job_name,
task_index=client_id,
protocol=remote_utils.get_default_communication_protocol(),
port=port)
server_def.default_session_config.rpc_options.num_channels_per_target = 4
server_def.default_session_config.experimental.recv_buf_max_chunk = -1
logging.info("Enabling collectives with server_def: %s", server_def)
context.context().enable_collective_ops(server_def)
context.ensure_initialized()
@tf_export(
"experimental.dtensor.initialize_accelerator_system",
"experimental.dtensor.initialize_tpu_system",
"experimental.dtensor.initialize_multi_client",
v1=[])
def initialize_accelerator_system(
device_type: Optional[str] = None,
enable_coordination_service: Optional[bool] = True,
num_logical_cpu_devices: Optional[int] = None,
experimental_reset_context: Optional[bool] = False,
experimental_enable_megcore: Optional[bool] = False,
) -> str:
"""Initializes accelerators and communication fabrics for DTensor.
DTensor configures TensorFlow to run in the local mode or multi-client mode.
- In local mode, a mesh can only use devices attached to the current process.
- In multi-client mode, a mesh can span across devices from multiple clients.
If `DTENSOR_JOBS` is non-empty, DTensor configures TensorFlow to run in the
multi-client mode using the distributed runtime. In multi-client mode devices
on different clients can communicate with each other.
The following environment variables controls the behavior of this function.
- `DTENSOR_JOBS`: string, a comma separated list. Each item in the list is
of format `{hostname}:{port}`. If empty, DTensor runs in the local mode.
Examples of valid `DTENSOR_JOBS` values:
- 4 clients on localhost:
`localhost:10000,localhost:10001,localhost:10002,localhost:10003`
- 2 clients on host1, 2 clients on host2
`host1:10000,host1:10001,host2:10000,host2:10003`
If the hostnames are BNS addresses, the items must be sorted in
alphabetical order.
- `DTENSOR_CLIENT_ID`: integer, between `0` to `num_clients - 1`, to identify
the client id of the current process. The default value is `0`.
- `DTENSOR_JOB_NAME`: string, a string for the name of the TensorFlow job.
The job name controls the job name section of the TensorFlow DeviceSpecs,
e.g., `job:worker` in `/job:worker/replica:0/task:0/device:TPU:0` when
the job name is `worker`.
The default value is `localhost` in local mode, and
`worker` when in the multi-client mode. All DTensor clients within the
same multi-client cluster share the same job name.
Args:
device_type: Type of accelerator to use, can be CPU, GPU, or TPU. If None,
uses `tf.experimental.dtensor.preferred_device_type()`.
enable_coordination_service: If true, enable distributed coordination
service to make sure that workers know the devices on each other, when
there is more than 1 client.
num_logical_cpu_devices: the number of logical CPU devices per DTensor
client. Default to the current number of logical CPU
(`dtensor.num_local_devices("CPU")`),when `device_type` is CPU, otherwise
set automatially to match the number of local GPU/TPU devices.
experimental_reset_context: Reset the tensorflow context. Behaviors of
existing TensorFlow objects (e.g. Tensors) are undefined. Set this to True
as an escape hatch, if there is no clear way to refactor your code to call
initialize_accelerator_system() before calling TensorFlow APIs that
initialize the context.
experimental_enable_megcore: Optionally enable megcore in backend.
Returns:
device_type: the type of accelerator that was initialized.
"""
global _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
assert context.executing_eagerly()
if is_initialized():
raise ValueError(
"Accelerator system has already been initialized. "
"Call tf.experimental.dtensor.shutdown_accelerator_system() first.")
if experimental_reset_context:
if context.context()._initialized: # pylint: disable=protected-access
logging.warn(
"experimental_reset_context is True. "
"Resetting TensorFlow context. Existing TensorFlow objects "
"(e.g. Tensors and resources) are invalidated."
)
context.context().ensure_uninitialized()
if context.context()._initialized: # pylint: disable=protected-access
raise ValueError(
"TensorFlow has already been initialized. "
"tf.experimental.dtensor.initialize_accelerator_system() must be "
"called before TensorFlow is initialized.")
context.context()._clear_caches() # pylint: disable=protected-access
if device_type is None:
device_type = config.preferred_device_type()
device_type = device_type.upper()
if device_type not in {"CPU", "GPU", "TPU"}:
raise ValueError(f"Unknown device_type {device_type}. "
"Allowed values are CPU, GPU, or TPU")
if config.gpu_use_nccl_communication():
logical_gpu_count = config.num_local_devices("GPU")
physical_gpu_count = len(tf_config.list_physical_devices("GPU"))
if logical_gpu_count > physical_gpu_count:
raise ValueError(
"DTENSOR_GPU_USE_NCCL_COMMUNICATION is set for using NCCL. "
"NCCL Collectives require one to one mapping between logical and "
"physical GPUs. "
f"The number of logical GPU ({logical_gpu_count}) "
f"is more than the number of physical GPU ({physical_gpu_count})."
)
# Configure logical host CPU devices for accelerators.
if device_type in ("GPU", "TPU"):
num_local_devices = config.num_local_devices(device_type)
if num_logical_cpu_devices is None:
num_logical_cpu_devices = max(
config.num_local_devices("CPU"), num_local_devices
)
else:
if num_logical_cpu_devices < num_local_devices:
raise ValueError(
"If set, `num_logical_cpu_devices`"
f" (={num_logical_cpu_devices}) must be greater than or"
f" equal to the number of local {device_type} devices"
f" (={num_local_devices})"
)
if num_logical_cpu_devices is not None:
tf_config.set_logical_device_configuration(
tf_config.list_physical_devices("CPU")[0],
[context.LogicalDeviceConfiguration()]
* num_logical_cpu_devices,
)
if not config.is_local_mode():
initialize_multi_client_cluster(
job_name=config.job_name(),
dtensor_jobs=config.jobs(),
client_id=config.client_id(),
collective_leader=config.full_job_name(task_id=0),
gpu_use_nccl_communication=config.gpu_use_nccl_communication(),
enable_coordination_service=enable_coordination_service)
else:
if device_type == "GPU":
# Enables Nccl on local mode.
context.context( # pylint: disable=protected-access
)._collective_use_nccl_communication = config.gpu_use_nccl_communication(
)
if device_type == "TPU":
tpu_util.initialize_tpu_system(use_megacore=experimental_enable_megcore)
_INITIALIZED_ACCELERATOR_SYSTEM_TYPE = device_type
return device_type
@tf_export(
"experimental.dtensor.shutdown_accelerator_system",
"experimental.dtensor.shutdown_tpu_system",
v1=[])
def shutdown_accelerator_system() -> None:
"""Shuts down the accelerator system."""
global _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
try:
context.async_wait()
finally:
if not is_initialized():
raise ValueError(
"Accelerator system is not initialized. Call "
"tf.experimental.dtensor.initialize_accelerator_system first."
)
device_type = _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
if not config.is_local_mode():
raise ValueError(
"Shutting down accelerator system under multi-client mode is "
"not supported."
)
if device_type == "TPU":
tpu_util.shutdown_tpu_system()
# reset TF context to stop gRPC servers.
context._reset_context() # pylint: disable=protected-access
context.context()._clear_caches() # pylint: disable=protected-access
_INITIALIZED_ACCELERATOR_SYSTEM_TYPE = None
tensorflow/dtensor/python/api.py
+568
View File
@@ -0,0 +1,568 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Core DTensor Python API."""
import contextlib
import threading
from typing import Any, Callable, Optional, Sequence
from tensorflow.dtensor.python import dtensor_device
from tensorflow.dtensor.python import gen_dtensor_ops
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
_dtensor_singleton = None
_dtensor_singleton_lock = threading.Lock()
# -----------------------------------------------------------------------------
# Main methods to launch DTensor computations.
@tf_export("experimental.dtensor.call_with_layout", v1=[])
def call_with_layout(fn: Callable[...,
Any], layout: Optional[layout_lib.Layout],
*args, **kwargs) -> Any:
"""Calls a function in the DTensor device scope if `layout` is not None.
If `layout` is not None, `fn` consumes DTensor(s) as input and produces a
DTensor as output; a DTensor is a tf.Tensor with layout-related attributes.
If `layout` is None, `fn` consumes and produces regular tf.Tensors.
Args:
fn: A supported TF API function such as tf.zeros.
layout: Optional, the layout of the output DTensor.
*args: Arguments given to `fn`.
**kwargs: Keyword arguments given to `fn`.
Returns:
The return value of `fn` transformed to a DTensor if requested.
"""
if layout is not None:
if context.executing_eagerly():
with default_mesh(layout.mesh):
with _dtensor_device()._default_layout(layout): # pylint: disable=protected-access
return fn(*args, **kwargs)
else:
return relayout(fn(*args, **kwargs), layout)
return fn(*args, **kwargs)
@tf_export("experimental.dtensor.run_on", v1=[])
@deprecation.deprecated(None, "Use `dtensor.default_mesh` scope instead.")
@contextlib.contextmanager
def run_on(mesh: layout_lib.Mesh):
"""Runs enclosed functions in the DTensor device scope.
This function returns a scope. All the ops and tf.functions in this scope will
run on the DTensor device using the mesh provided.
This is useful for wrapping any tf.function that doesn't take a DTensor as
input but would like to produce DTensor as result. The scope will also make
sure all small constants be replicated as DTensor.
Args:
mesh: A Mesh instance to extract a default mesh from.
Yields:
A context in which all ops and tf.functions will run on the DTensor device.
"""
with default_mesh(mesh):
yield
@tf_export("experimental.dtensor.default_mesh", v1=[])
@contextlib.contextmanager
def default_mesh(mesh: layout_lib.Mesh):
"""Sets the default DTensor device mesh to use for enclosed functions.
This function returns a scope. All the ops and tf.functions in this scope will
default to this DTensor mesh if a mesh cannot be inferred from any of the
inputs
This is useful for wrapping any tf.function that doesn't take a DTensor as
input but would like to produce DTensor as result. The scope will also make
sure all small constants are replicated as DTensors.
Args:
mesh: A Mesh instance to extract a default mesh from.
Yields:
A context in which all ops and tf.functions will run on the given mesh.
"""
if not isinstance(mesh, layout_lib.Mesh):
raise ValueError(f"Expect `mesh` to be `Mesh`, got {type(mesh)}")
with _dtensor_device()._experimental_default_mesh(mesh): # pylint: disable=protected-access
with ops.device(device_name()):
yield
@tf_export("experimental.dtensor.get_default_mesh", v1=[])
def get_default_mesh() -> Optional[layout_lib.Mesh]:
"""Return the default mesh under the current dtensor device context.
In the case that dtensor device system is not initialized, this function
will return None.
Returns:
The current default mesh for the dtensor device context.
"""
if _dtensor_singleton is None:
return None
else:
return _dtensor_singleton._current_default_mesh # pylint: disable=protected-access
@tf_export("experimental.dtensor.device_name", v1=[])
def device_name() -> str:
"""Returns the singleton DTensor device's name.
This function can be used in the following way:
```python
import tensorflow as tf
with tf.device(dtensor.device_name()):
# ...
```
"""
return _dtensor_device().name
@tf_export("experimental.dtensor.is_dtensor", v1=[])
def is_dtensor(tensor) -> bool:
"""Check whether the input tensor is a DTensor.
In Python, a DTensor has the same type as a `tf.Tensor`. This method will
let you check and handle the tensor differently if a tf.Tensor is a DTensor.
Args:
tensor: an object to be checked.
Returns:
bool, True if the given tensor is a DTensor.
"""
return _dtensor_device().is_dtensor(tensor)
# -----------------------------------------------------------------------------
# Data transfer methods.
@tf_export("experimental.dtensor.copy_to_mesh", v1=[])
def copy_to_mesh(
tensor: Any,
layout: layout_lib.Layout,
source_layout: Optional[layout_lib.Layout] = None) -> tensor_lib.Tensor:
"""Copies a tf.Tensor onto the DTensor device with the given layout.
Copies a regular tf.Tensor onto the DTensor device. Use the mesh attached to
`layout` as target mesh. This method currently only supports replicated
layouts, or one-to-one copies for sharded layouts.
Args:
tensor: A regular tf.Tensor to be copied as a DTensor.
layout: Target layout (and mesh) for the result DTensor.
source_layout: Source layout of the tensor before copy. This argument
is deprecated.
Returns:
A DTensor on the DTensor device with the given layout.
"""
del source_layout
return relayout(tensor, layout)
@tf_export("experimental.dtensor.pack", v1=[])
def pack(tensors: Sequence[Any], layout: layout_lib.Layout) -> Any:
"""Packs `tf.Tensor` components into a DTensor.
Packing and unpacking are inverse operations:
```
* unpack(pack(tensors)) == tensors
* pack(unpack(dtensor)) == dtensor
```
1. For any DTensor on the mesh, `unpack` returns the raw components placed on
each underlying device.
2. Packing these raw components in the same order using `pack` returns a
DTensor which should be identical to the original DTensor--both the content
value and the layout.
**Shape, Rank, and Scalars**: The rank of the DTensor is the same as the
rank of its raw components, i.e., rank is preserved. This leads to a
consistent interpretation for packing scalar values into a DTensor. The only
valid layout for a scalar value is fully replicated, and the individual
components must be identical scalars.
Each input `tensors[i]` will be copied to `layout.mesh.local_device[i]`
if not already on the local device. Non-local components should not be passed
to `pack`; use `copy_to_mesh` and `relayout` to place tensors on all global
devices on a mesh.
It is the caller's responsibility to ensure that the underlying values
for `pack` adhere to the specified layout, and that only as many values are
specified as there are local devices. Pack does not move data between clients.
See examples below for more detail about layouts.
For example, assume we have a mesh `[X(2), Y(3)]`, which has in total 6
underlying devices. Futuremore, assume that the device location mapping is
the following:
```
device_ID | location X, Y
0 0, 0
1 0, 1
2 0, 2
3 1, 0
4 1, 1
5 1, 2
```
1. For 1-D vector DTensor with shape `[128]` with layout `[mesh.X]` and value
as `range(128)`, the raw components will have shape `[64]` each, and the
raw components will be:
```
device_ID | raw component
0 range(0, 64)
1 range(0, 64)
2 range(0, 64)
3 range(64, 128)
4 range(64, 128)
5 range(64, 128)
```
This also means for a 1-D DTensor with shape `[2]` and layout `[mesh.X]`,
the raw components have shape `[1]` rather than the shape for scalar values
`[]`.
2. For 2-D vector DTensor with shape `[2, 3]` with layout `[mesh.X, mesh.Y]`
and value as `range(6)`, this is basically a fully-sharded DTensor.
From global view, the content looks like
```
[
[0.0, 1.0, 2.0],
[3.0, 4.0, 5.0],
]
```
The raw components will have shape `[1, 1]` each, and have the following
content:
```
device_ID | raw component
0 [[0.0]]
1 [[1.0]]
2 [[2.0]]
3 [[3.0]]
4 [[4.0]]
5 [[5.0]]
```
3. For a scalar value `123.0` DTensor, it can only have one legitimate layout
`[]` (no dimension, but fully replicated).
The raw components will have shape `[]` each, and have the following
content:
```
device_ID | raw component
0 123.0
1 123.0
2 123.0
3 123.0
4 123.0
5 123.0
```
Again, caller of `pack` is expected to provide 6 identical value raw
components with scalar shapes.
4. For 3-D vector DTensor with shape `[2, 2, 3]` with layout
`[X, unsharded, unsharded]` and value as `range(12)`,
From global view, the content looks like:
```
[
[
[0.0, 1.0, 2.0],
[3.0, 4.0, 5.0],
],
[
[6.0, 7.0, 8.0],
[9.0, 10., 11.],
],
]
```
The raw components will have shape `[1, 2, 3]` each, and have the following
content:
```
device_ID | raw component
0 range(6).reshape([1, 2, 3])
1 range(6).reshape([1, 2, 3])
2 range(6).reshape([1, 2, 3])
3 range(6, 12).reshape([1, 2, 3])
4 range(6, 12).reshape([1, 2, 3])
5 range(6, 12).reshape([1, 2, 3])
```
Args:
tensors: The list of local tensor components to pack into a DTensor.
layout: The layout of the DTensor to be created.
Returns:
A DTensor created from the individual component tensors.
Raises:
RuntimeError: When `pack` is not called eagerly.
"""
return _dtensor_device().pack(tensors, layout)
@tf_export("experimental.dtensor.unpack", v1=[])
def unpack(tensor: Any) -> Sequence[Any]:
"""Unpacks a DTensor into `tf.Tensor` components.
Packing and unpacking are inverse operations:
```
* unpack(pack(tensors)) == tensors
* pack(unpack(dtensor)) == dtensor
```
1. For any DTensor on the mesh, `unpack` returns the raw components placed on
each underlying device.
2. Packing these raw components in the same order using `pack` returns a
DTensor which should be identical to the original DTensor--both the content
value and the layout.
See the documentation for `pack` for more information about how packing and
unpacking works.
Args:
tensor: The DTensor to unpack.
Returns:
The individual component tensors of the DTensor. This will include only the
client-local components, i.e. the components placed on the local devices.
Raises:
RuntimeError: When `unpack` is not called eagerly.
"""
return _dtensor_device().unpack(tensor)
# -----------------------------------------------------------------------------
# Layout-related methods.
@tf_export("experimental.dtensor.fetch_layout", v1=[])
def fetch_layout(tensor: tensor_lib.Tensor) -> layout_lib.Layout:
"""Fetches the layout of a DTensor.
Args:
tensor: The DTensor whose layout is to be fetched.
Returns:
The `Layout` of this DTensor.
Raises:
RuntimeError: When not called eagerly.
"""
return _dtensor_device().fetch_layout(tensor)
@tf_export("experimental.dtensor.check_layout", v1=[])
def check_layout(tensor: tensor_lib.Tensor, layout: layout_lib.Layout) -> None:
"""Asserts that the layout of the DTensor is `layout`.
Args:
tensor: A DTensor whose layout is to be checked.
layout: The `Layout` to compare against.
Raises:
ValueError: If the layout of `tensor` does not match the supplied `layout`.
"""
if fetch_layout(tensor) != layout:
raise ValueError("Layout of tensor: " + str(fetch_layout(tensor)) +
", did not match expected layout: " + str(layout))
@tf_export("experimental.dtensor.relayout", v1=[])
def relayout(
tensor: tensor_lib.Tensor,
layout: layout_lib.Layout,
name: Optional[str] = None,
) -> tensor_lib.Tensor:
"""Changes the layout of `tensor`.
Changes the layout of `tensor` to `layout`. This is used to fine-tune the
behavior of ops following/connected to `tensor`, such as choosing one SPMD
expansion pattern over another. This works by forward propagating `layout`
to connected TensorFlow computation graphs during layout propagation.
Currently, only converting layouts from replicated to sharded or sharded to
replicated per mesh dimension is supported. That is, "x, y" -> "unsharded, y"
is supported, while "x, y" -> "z, y" is not supported.
We also support a special "match" sharding spec, which instructs the relayout
to act as an identity operation with respect to any sharding on these
mesh dimensions.
Relayout is internally lowered to a set of Split and/or AllToAll ops. When
tensor layouts are converted from replicated to sharded, the cost is
comparatively low because we only insert Split ops and no cross-device
communication is needed. However, when tensor layouts are converted from
sharded to replicated, cross-device communication may occur, causing potential
performance impact.
Args:
tensor: A DTensor to specify a new layout for.
layout: A Layout object specifying a new sharding spec.
name: name of the Op.
Returns:
A DTensor output from the Relayout op.
"""
layout_str = layout.to_string()
with default_mesh(layout.mesh):
return gen_dtensor_ops.relayout(tensor, layout_str, name=name)
@tf_export("experimental.dtensor.relayout_like", v1=[])
def relayout_like(
tensor: tensor_lib.Tensor,
layout_tensor: tensor_lib.Tensor,
name: Optional[str] = None,
) -> tensor_lib.Tensor:
"""Changes the layout of `tensor` to the same as `layout_tensor`.
`relayout_like` is often used inside a `tf.function`, to ensure a tensor is
placed to the same mesh and with the same layout as another tensor.
The backward gradient of a `relayout` is a `relayout_like` operation, to
ensure the backward tensor has the same layout as the forward input tensor:
```
@ops.RegisterGradient("Relayout")
def _relayout_gradient(op, grad):
return relayout_like(grad, layout_input=op.inputs[0])
```
Here is another illustrative example:
```
@tf.function
def func(x):
z = tf.ones(x.shape)
z = dtensor.relayout_like(z, x)
return x + z
with dtensor.default_mesh(cpu_mesh):
x = tf.ones((4, 4))
with dtensor.default_mesh(gpu_mesh):
y = func(x)
# y would be on the cpu mesh, following the mesh of x.
```
Args:
tensor: A DTensor to specify a new layout for.
layout_tensor: A Tensor object whose layout will be used for the layout of
result. The shape and type of layout_tensor are irrelevant.
name: name of the Op.
Returns:
A DTensor output from the RelayoutLike op.
"""
return gen_dtensor_ops.relayout_like(
input=tensor, layout_input=layout_tensor, name=name
)
@tf_export("experimental.dtensor._reset_dtensor_device", v1=[])
def reset_dtensor_device(is_async: bool) -> None:
"""Resets the Eager execution device for DTensor.
This function is only intended for testing and diagnostics.
Args:
is_async: If True, the device uses async execution.
"""
global _dtensor_singleton
device = dtensor_device.DTensorDevice(meshes=[], is_async=is_async)
_dtensor_singleton = device
def _dtensor_device() -> dtensor_device.DTensorDevice:
with _dtensor_singleton_lock:
if _dtensor_singleton is None:
reset_dtensor_device(is_async=True)
return _dtensor_singleton
def _reset() -> None:
global _dtensor_singleton
with _dtensor_singleton_lock:
if _dtensor_singleton is not None:
_dtensor_singleton.clear_tpu_core_ids()
_dtensor_singleton = None
# ----------------------------------------------------------------------------
# Gradients
@ops.RegisterGradient("Relayout")
def _relayout_gradient(op, grad):
grad = gen_dtensor_ops.relayout_like(grad, layout_input=op.inputs[0])
return grad
@ops.RegisterGradient("RelayoutLike")
def _relayout_grad_gradient(op, grad):
# Gradient of RelayoutGrad is relayout to the original Relayout's output.
grad = gen_dtensor_ops.relayout_like(grad, layout_input=op.inputs[0])
# Return None for forward_input's partial gradient since it is not connected
# to the target's gradient.
return grad, None
@ops.RegisterGradient("CopyToMesh")
def _copy_to_mesh_gradient(op, grad):
grad = gen_dtensor_ops.copy_to_mesh_grad(
grad,
forward_input=op.inputs[0],
)
return grad
@ops.RegisterGradient("CopyToMeshGrad")
def _copy_to_mesh_grad_gradient(op, grad):
grad = gen_dtensor_ops.copy_to_mesh_grad(
grad,
forward_input=op.inputs[0],
)
return grad, None
tensorflow/dtensor/python/config.py
+222
View File
@@ -0,0 +1,222 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DTensor Configuration API."""
import os
from typing import List, Optional, Union
from tensorflow.python.eager import context
from tensorflow.python.framework import config as tf_config
from tensorflow.python.framework import device as tf_device
from tensorflow.python.util.tf_export import tf_export
_DT_CLIENT_ID = "DTENSOR_CLIENT_ID"
# DTENSOR_NUM_CLIENTS is removed, but some DTensor users still use this symbol.
_DT_NUM_CLIENTS = "DTENSOR_NUM_CLIENTS"
_DT_JOB_NAME = "DTENSOR_JOB_NAME"
_DT_JOBS = "DTENSOR_JOBS"
_DT_HEARTBEAT_ENABLED = "DTENSOR_ENABLE_HEARTBEAT"
# All functions in this file can be used before calling
# `tf.experimental.dtensor.initialize_accelerator_system`.
# -----------------------------------------------------------------------------
# Distributed training-related methods.
#
# Most users should use DTensor utility methods to create a mesh. The methods
# here are only for advanced users who want to fully customize their meshes.
# Note that local_devices and num_local_devices return the actual number of
# locally attached devices. The others are set through environment variables.
@tf_export("experimental.dtensor.local_devices", v1=[])
def local_devices(
device_type: str, for_client_id: Optional[int] = None
) -> List[tf_device.DeviceSpec]:
"""Returns a list of device specs configured on this client."""
if device_type.upper() not in ["CPU", "GPU", "TPU"]:
raise ValueError(f"Device type {device_type} is not CPU, GPU, or TPU.")
if for_client_id is None:
for_client_id = client_id()
# Return fully qualified device specs, sorted by increasing device index.
return [
tf_device.DeviceSpec( # pylint: disable=g-complex-comprehension
job=job_name(),
replica=0, # replica is deprecated and mostly hard-coded now.
task=for_client_id,
device_type=device_type,
device_index=i,
)
for i in range(num_local_devices(device_type))
]
@tf_export("experimental.dtensor.num_local_devices", v1=[])
def num_local_devices(device_type: str) -> int:
"""Returns the number of devices of device_type configured on this client."""
# Reads from config because CPU and GPU can use logical devices.
if device_type.upper() in ["CPU", "GPU"]:
context_config = context.get_config()
return context_config.device_count[device_type.upper()]
return len(tf_config.list_physical_devices(device_type))
@tf_export("experimental.dtensor.num_global_devices", v1=[])
def num_global_devices(device_type: str) -> int:
"""Returns the number of devices of device_type in this DTensor cluster."""
return num_local_devices(device_type) * num_clients()
@tf_export("experimental.dtensor.client_id", v1=[])
def client_id() -> int:
"""Returns this client's ID."""
# If missing, assume running with a single client with client_id of 0.
client_id_value = int(os.environ.get(_DT_CLIENT_ID, "0"))
if client_id_value < 0:
raise ValueError(
f"Environment variable {_DT_CLIENT_ID} "
f"must be >= 0, got {client_id_value}. "
)
if client_id_value >= num_clients():
raise ValueError(
f"Environment variable {_DT_CLIENT_ID} "
f"must be < {num_clients()}, got {client_id_value}"
)
return client_id_value
@tf_export("experimental.dtensor.num_clients", v1=[])
def num_clients() -> int:
"""Returns the number of clients in this DTensor cluster."""
if is_local_mode():
return 1
return len(jobs())
@tf_export("experimental.dtensor.job_name", v1=[])
def job_name() -> str:
"""Returns the job name used by all clients in this DTensor cluster."""
# If missing, assumes the program runs locally and use localhost as job name
# per TensorFlow convention.
return os.environ.get(
_DT_JOB_NAME, "localhost" if num_clients() == 1 else "worker"
)
@tf_export("experimental.dtensor.full_job_name", v1=[])
def full_job_name(task_id: Optional[int] = None) -> str:
"""Returns the fully qualified TF job name for this or another task."""
# If task_id is None, use this client's ID, which is equal to its task ID.
if task_id is None:
task_id = client_id()
# In local runs and unit tests, there should be exactly one client running
# on one TF task.
if num_clients() == 1 and task_id != 0:
raise ValueError(f"Unexpected task ID {task_id} in local runs")
return f"{job_name()}/replica:0/task:{task_id}"
def _bns_task_id(job: str) -> Union[int, str]:
"""Tries to extract an integer task ID from a job name.
For example, for `job` = '/.../tpu_worker/0:port_name', return 0.
Args:
job: A job name to extract task ID from.
Returns:
The task ID on success, or the original job name on failure.
"""
maybe_task_id = job.rsplit("/")[-1].rsplit(":")[0]
try:
return int(maybe_task_id)
except ValueError:
return job
@tf_export("experimental.dtensor.jobs", v1=[])
def jobs() -> List[str]:
"""Returns a list of job names of all clients in this DTensor cluster."""
d_jobs = os.environ.get(_DT_JOBS)
if d_jobs is None:
return []
d_jobs_list = d_jobs.split(",")
# Validate ordering for BNS style job names.
# For definition of BNS, refer to https://research.google/pubs/pub43438/.
if any([name.startswith("/bns/") for name in d_jobs_list]):
if d_jobs_list != sorted(d_jobs_list, key=_bns_task_id):
raise ValueError(
f"Unexpected DTENSOR_JOBS content {d_jobs}. Sort entries "
"in DTENSOR_JOBS because cluster construction relies on "
"the order."
)
return d_jobs_list
@tf_export("experimental.dtensor.heartbeat_enabled", v1=[])
def heartbeat_enabled() -> bool:
"""Returns true if DTensor heartbeat service is enabled."""
return os.environ.get(_DT_HEARTBEAT_ENABLED, "true").lower() in ("true", "1")
def is_local_mode() -> bool:
"""Returns true if DTensor shall run in local mode."""
return not jobs()
def is_tpu_present() -> bool:
"""Returns true if TPU devices are present."""
# Check if TPU is present from initialized context.
# TPU_SYSTEM is a device that indicates TPUs are present.
tpu_system_devices = tf_config.list_physical_devices("TPU_SYSTEM")
return bool(tpu_system_devices)
def is_gpu_present() -> bool:
"""Returns true if TPU devices are present."""
return bool(tf_config.list_physical_devices("GPU"))
@tf_export("experimental.dtensor.preferred_device_type", v1=[])
def preferred_device_type() -> str:
"""Returns the preferred device type for the accelerators.
The returned device type is determined by checking the first present device
type from all supported device types in the order of 'TPU', 'GPU', 'CPU'.
"""
if is_tpu_present():
return "TPU"
elif is_gpu_present():
return "GPU"
return "CPU"
def use_multi_device_mode() -> bool:
"""Return True if environment indicates multi-device mode is enabled."""
return os.environ.get("DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION", "0") != "0"
def gpu_use_nccl_communication() -> bool:
"""Return True if environment indicates NCCL shall be used for GPU."""
return os.environ.get("DTENSOR_GPU_USE_NCCL_COMMUNICATION", "0") != "0"
tensorflow/dtensor/python/d_checkpoint.py
+463
View File
@@ -0,0 +1,463 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DTensor Checkpoint.
Note that this module contains deprecated functionality, and the DTensor related
checkpoint has been integrated with tf.train.Checkpoint. It can be used out of
the box to save and restore dtensors.
"""
from typing import Dict, List, Optional
import weakref
from tensorflow.core.protobuf import trackable_object_graph_pb2
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import gen_dtensor_ops
from tensorflow.dtensor.python import layout
from tensorflow.dtensor.python import save_restore
from tensorflow.python.checkpoint import checkpoint as util
from tensorflow.python.checkpoint import checkpoint_options
from tensorflow.python.checkpoint import graph_view as graph_view_lib
from tensorflow.python.checkpoint import restore as restore_lib
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.trackable import base
from tensorflow.python.trackable import data_structures
from tensorflow.python.training import py_checkpoint_reader
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.util import deprecation
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export
class _DSaver: # pylint: disable=protected-access
"""A single device saver that places tensors on DTensor Device."""
def __init__(self, mesh: layout.Mesh,
saveable_objects: List[saveable_object.SaveableObject]):
self._saveable_objects = saveable_objects
self._mesh = mesh
def save(
self,
file_prefix: str,
options: Optional[checkpoint_options.CheckpointOptions] = None
) -> Optional[ops.Operation]:
"""Saves the saveable objects to a checkpoint with `file_prefix`.
Also query the generated shards from the distributed DTensor SaveV2 ops and
do a MergeV2 on those. Each op here is backed by a global_barrier to avoid
racing from multiple clients.
Args:
file_prefix: A string or scalar string Tensor containing the prefix to
save under.
options: Optional `CheckpointOptions` object. This is unused in DTensor.
Returns:
An `Operation`, or None when executing eagerly.
"""
if options is not None and options.experimental_io_device is not None:
raise ValueError(
"Specified experimental_io_device in DTensor checkpoint is not supported."
)
del options
tensor_names = []
tensors = []
tensor_slices = []
for saveable in self._saveable_objects:
for spec in saveable.specs:
tensor = spec.tensor
# A tensor value of `None` indicates that this SaveableObject gets
# recorded in the object graph, but that no value is saved in the
# checkpoint.
if tensor is not None:
if api.device_name() != spec.device:
# Some small tensors are placed on CPU0 from save manager and
# broadcasted to DTensor mesh, e,g., SaveCounter.
tensor = api.pack([tensor] *
self._mesh.host_mesh().num_local_devices(),
layout.Layout.replicated(
self._mesh.host_mesh(),
rank=tensor.shape.rank))
tensor_names.append(spec.name)
tensors.append(tensor)
tensor_slices.append(spec.slice_spec)
return save_restore.sharded_save(self._mesh, file_prefix, tensor_names,
tensor_slices, tensors)
def restore(
self,
file_prefix: str,
options: Optional[checkpoint_options.CheckpointOptions] = None
) -> Dict[str, ops.Operation]:
"""Restore the saveable objects from a checkpoint with `file_prefix`.
Args:
file_prefix: A string or scalar string Tensor containing the prefix for
files to read from.
options: Optional `CheckpointOptions` object. This is unused in DTensor.
Returns:
A dictionary mapping from SaveableObject names to restore operations.
"""
if options is not None and options.experimental_io_device is not None:
raise ValueError(
"Specified experimental_io_device in DTensor checkpoint is not "
"supported.")
del options
restore_specs = []
tensor_structure = []
for saveable in self._saveable_objects:
saveable_tensor_structure = []
tensor_structure.append(saveable_tensor_structure)
# DTensor change 1 : Gather shapes and layout from original saveable
# specs.
# Note that this relies on the fact that the variables are already
# initialized -- which isn't the behavior we want eventually.
# TODO(b/159035705): Handle the variable initialization in restore.
for spec in saveable.specs:
saveable_tensor_structure.append(spec.name)
if isinstance(spec, d_variable.DSaveSpec):
restore_specs.append((spec.name, spec.slice_spec, spec.dtype,
spec.layout, spec.global_shape))
# Fall back to replicated layouts for non-DTensor saves that constructs
# normal SaveSpec.
elif isinstance(spec, saveable_object.SaveSpec):
restore_specs.append(
(spec.name, spec.slice_spec, spec.dtype,
layout.Layout.replicated(self._mesh.host_mesh(),
spec.tensor.shape.rank).to_string(),
spec.tensor.shape.as_list()))
tensor_names, tensor_slices, tensor_dtypes, layouts, global_shapes = zip(
*restore_specs)
with ops.device(api.device_name()):
# DTensor change 2 : Run on customized DTensor RestoreV2 op rather than
# stock TF io_ops.RestoreV2.
restored_tensors = gen_dtensor_ops.d_tensor_restore_v2(
prefix=file_prefix,
tensor_names=tensor_names,
shape_and_slices=tensor_slices,
input_shapes=global_shapes,
input_layouts=layouts,
dtypes=tensor_dtypes)
structured_restored_tensors = nest.pack_sequence_as(tensor_structure,
restored_tensors)
restore_ops = {}
for saveable, restored_tensors in zip(self._saveable_objects,
structured_restored_tensors):
restore_ops[saveable.name] = saveable.restore(
restored_tensors, restored_shapes=None)
return restore_ops
class _DCheckpointRestoreCoordinator(util._CheckpointRestoreCoordinator): # pylint: disable=protected-access
"""Holds the status of an object-based checkpoint load."""
def __init__(self, mesh: layout.Mesh, **kwargs):
super().__init__(**kwargs)
self._mesh = mesh
def restore_saveables(self,
tensor_saveables: Dict[str,
saveable_object.SaveableObject],
python_positions: List[restore_lib.CheckpointPosition],
registered_savers: Optional[Dict[str, Dict[
str, base.Trackable]]] = None,
reader: py_checkpoint_reader.NewCheckpointReader = None
) -> Optional[List[ops.Operation]]:
"""Run or build restore operations for SaveableObjects.
Args:
tensor_saveables: `SaveableObject`s which correspond to Tensors.
python_positions: `CheckpointPosition`s which correspond to `PythonState`
Trackables bound to the checkpoint.
registered_savers: a dict mapping saver names-> object name -> Trackable.
This argument is not implemented for DTensorCheckpoint.
reader: A CheckpointReader. Creates one lazily if None.
Returns:
When graph building, a list of restore operations, either cached or newly
created, to restore `tensor_saveables`.
"""
del registered_savers
restore_ops = []
# Eagerly run restorations for Python state.
if python_positions:
# Lazily create the NewCheckpointReader, since this requires file access
# and we may not have any Python saveables.
if reader is None:
reader = py_checkpoint_reader.NewCheckpointReader(self.save_path_string)
for position in python_positions:
key = position.object_proto.attributes[0].checkpoint_key
position.trackable.deserialize(reader.get_tensor(key))
# If we have new SaveableObjects, extract and cache restore ops.
if tensor_saveables:
validated_saveables = saveable_object_util.validate_and_slice_inputs(
tensor_saveables)
validated_names = set(saveable.name for saveable in validated_saveables)
if set(tensor_saveables.keys()) != validated_names:
raise AssertionError(
("Saveable keys changed when validating. Got back %s, was "
"expecting %s") % (tensor_saveables.keys(), validated_names))
# DTensor change: Use _DSaver that does restore on DTensor with
# customized DTensorRestoreV2 op.
new_restore_ops = _DSaver(self._mesh, validated_saveables).restore(
self.save_path_tensor, self.options)
if not context.executing_eagerly():
for name, restore_op in sorted(new_restore_ops.items()):
restore_ops.append(restore_op)
assert name not in self.restore_ops_by_name
self.restore_ops_by_name[name] = restore_op
return restore_ops
class DTrackableSaver(util.TrackableSaver):
"""A DTensor trackable saver that uses _SingleDeviceSaver."""
def __init__(self, mesh: layout.Mesh, graph_view):
super(DTrackableSaver, self).__init__(graph_view)
self._mesh = mesh
def _gather_saveables(self, object_graph_tensor=None):
# Since the base Checkpoint class does not return SaveableObjects, re-use
# the saveables cache or generate new Saveables.
(serialized_tensors, feed_additions, registered_savers,
graph_proto) = self._gather_serialized_tensors(object_graph_tensor)
saveables_dict = self._saveables_cache
if saveables_dict is None:
# Get and remove object graph tensor from `serialized_tensors`, because
# the function `serialized_tensors_to_saveable_cache` isn't equipped
# to handle it.
object_graph_tensor = serialized_tensors.pop(
None)[base.OBJECT_GRAPH_PROTO_KEY]
saveables_dict = (
saveable_object_util.serialized_tensors_to_saveable_cache(
serialized_tensors))
named_saveable_objects = []
for saveable_by_name in saveables_dict.values():
for saveables in saveable_by_name.values():
named_saveable_objects.extend(saveables)
named_saveable_objects.append(
base.NoRestoreSaveable(
tensor=object_graph_tensor,
name=base.OBJECT_GRAPH_PROTO_KEY))
return (named_saveable_objects, graph_proto, feed_additions,
registered_savers)
def _save_cached_when_graph_building(self,
file_prefix,
object_graph_tensor,
options,
update_ckpt_state=False):
"""Create or retrieve save ops, overrides parents's private method.
Args:
file_prefix: The prefix for saved checkpoint files.
object_graph_tensor: A `Tensor` to which the current object graph will be
fed.
options: `CheckpointOptions` object.
update_ckpt_state: Optional bool flag. Indiciate whether the internal
checkpoint state needs to be updated. This is used for async checkpoint,
which DTrackableSaver currently does not support.
TODO(chienchunh): Implement async checkpoint for DTrackableSaver.
Returns:
A two-element tuple with a filename tensor and a feed_dict of tensors to
feed when running it (if graph building). The feed dict contains the
current object graph and any Python state to be saved in the
checkpoint. When executing eagerly only the first argument is meaningful.
"""
(named_saveable_objects, graph_proto, feed_additions,
unused_registered_savers) = self._gather_saveables(
object_graph_tensor=object_graph_tensor)
if (self._last_save_object_graph != graph_proto
# When executing eagerly, we need to re-create SaveableObjects each time
# save() is called so they pick up new Tensors passed to their
# constructors. That means the Saver needs to be copied with a new
# var_list.
or context.executing_eagerly() or ops.inside_function()):
# This is needed to avoid MultiDeviceSaver creating unnecessary MergeV2
# ops in DTensor. It is an issue when saving TPU Variables on host CPU
# mesh given our limited expressiveness in API and hard-coded logic in
# broadcasting -- for a small constant Tensor with no extra information,
# we place it on the first registered mesh(A.K.A. default mesh).
saver = _DSaver(self._mesh, named_saveable_objects)
save_op = saver.save(file_prefix, options=options)
with ops.device("/cpu:0"):
with ops.control_dependencies([save_op]):
self._cached_save_operation = array_ops.identity(file_prefix)
self._last_save_object_graph = graph_proto
return self._cached_save_operation, feed_additions
# TODO(b/180466245): Use proper mesh placement semantic.
def restore(self, save_path, options=None):
"""Restore a training checkpoint with host mesh placement."""
options = options or checkpoint_options.CheckpointOptions()
if save_path is None:
return util.InitializationOnlyStatus(self._graph_view, ops.uid())
reader = py_checkpoint_reader.NewCheckpointReader(save_path)
graph_building = not context.executing_eagerly()
if graph_building:
dtype_map = None
else:
dtype_map = reader.get_variable_to_dtype_map()
try:
object_graph_string = reader.get_tensor(base.OBJECT_GRAPH_PROTO_KEY)
except errors_impl.NotFoundError:
# The object graph proto does not exist in this checkpoint. Try the
# name-based compatibility mode.
restore_coordinator = util._NameBasedRestoreCoordinator( # pylint: disable=protected-access
save_path=save_path,
dtype_map=dtype_map)
if not graph_building:
for existing_trackable in self._graph_view.list_objects():
# pylint: disable=protected-access
existing_trackable._maybe_initialize_trackable()
existing_trackable._name_based_restores.add(restore_coordinator)
existing_trackable._name_based_attribute_restore(restore_coordinator)
# pylint: enable=protected-access
return util.NameBasedSaverStatus(
restore_coordinator, graph_view=self._graph_view)
if graph_building:
if self._file_prefix_placeholder is None:
# DTensor change: provide a hint for mesh broadcasting to put the input
# onto the host mesh.
self._file_prefix_placeholder = api.pack(
[constant_op.constant("model")] * self._mesh.num_local_devices(),
layout.Layout.replicated(self._mesh.host_mesh(), rank=0))
file_prefix_tensor = self._file_prefix_placeholder
file_prefix_feed_dict = {self._file_prefix_placeholder: save_path}
else:
# DTensor change: provide a hint for mesh broadcasting to put the input
# onto the host mesh.
file_prefix_tensor = api.pack(
[constant_op.constant(save_path)] * self._mesh.num_local_devices(),
layout.Layout.replicated(self._mesh.host_mesh(), rank=0))
file_prefix_feed_dict = None
object_graph_proto = (trackable_object_graph_pb2.TrackableObjectGraph())
object_graph_proto.ParseFromString(object_graph_string)
# DTensor Change: Hook the proper DSaver in restore.
checkpoint = _DCheckpointRestoreCoordinator(
mesh=self._mesh,
object_graph_proto=object_graph_proto,
save_path=save_path,
save_path_tensor=file_prefix_tensor,
reader=reader,
restore_op_cache=self._restore_op_cache,
graph_view=self._graph_view,
options=options,
saveables_cache=self._saveables_cache)
restore_lib.CheckpointPosition(
checkpoint=checkpoint, proto_id=0).restore(self._graph_view.root)
# Attached dependencies are not attached to the root, so should be restored
# separately.
if self._graph_view.attached_dependencies:
for ref in self._graph_view.attached_dependencies:
if ref.name == "root":
# Root dependency is automatically added to attached dependencies --
# this can be ignored since it maps back to the root object.
continue
proto_id = None
# Find proto ID of attached dependency (if it is in the proto).
for proto_ref in object_graph_proto.nodes[0].children:
if proto_ref.local_name == ref.name:
proto_id = proto_ref.node_id
break
if proto_id in checkpoint.object_by_proto_id:
# Object has already been restored. This can happen when there's an
# indirect connection from the attached object to the root.
continue
restore_lib.CheckpointPosition(
checkpoint=checkpoint, proto_id=proto_id).restore(ref.ref)
load_status = util.CheckpointLoadStatus(
checkpoint,
graph_view=self._graph_view,
feed_dict=file_prefix_feed_dict,
options=options)
return load_status
@deprecation.deprecated(
date=None,
instructions="Please use tf.train.Checkpoint instead of DTensorCheckpoint. "
"DTensor is integrated with tf.train.Checkpoint and it can be "
"used out of the box to save and restore dtensors.")
@tf_export("experimental.dtensor.DTensorCheckpoint", v1=[])
class DTensorCheckpoint(util.Checkpoint):
"""Manages saving/restoring trackable values to disk, for DTensor."""
def __init__(self, mesh: layout.Mesh, root=None, **kwargs):
super(DTensorCheckpoint, self).__init__(root=root, **kwargs)
self._mesh = mesh
saver_root = self
attached_dependencies = None
self._save_counter = None # Created lazily for restore-on-create.
self._save_assign_op = None
if root:
util._assert_trackable(root, "root")
saver_root = root
attached_dependencies = []
# All keyword arguments (including root itself) are set as children
# of root.
kwargs["root"] = root
root._maybe_initialize_trackable()
self._save_counter = data_structures.NoDependency(
root._lookup_dependency("save_counter"))
self._root = data_structures.NoDependency(root)
for k, v in sorted(kwargs.items(), key=lambda item: item[0]):
setattr(self, k, v)
# Call getattr instead of directly using v because setattr converts
# v to a Trackable data structure when v is a list/dict/tuple.
converted_v = getattr(self, k)
util._assert_trackable(converted_v, k)
if root:
# Make sure that root doesn't already have dependencies with these names
attached_dependencies = attached_dependencies or []
child = root._lookup_dependency(k)
if child is None:
attached_dependencies.append(base.TrackableReference(k, converted_v))
elif child != converted_v:
raise ValueError(
"Cannot create a Checkpoint with keyword argument {name} if "
"root.{name} already exists.".format(name=k))
# DTensor Change:
# Override the parents saver with DTrackableSaver with _SingleDeviceSaver.
self._saver = DTrackableSaver(
mesh,
graph_view_lib.ObjectGraphView(
weakref.ref(saver_root),
attached_dependencies=attached_dependencies))
tensorflow/dtensor/python/d_random.py
+331
View File
@@ -0,0 +1,331 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DTensor helpers for random generators."""
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import gen_stateless_random_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import shape_util
# ------------------------------------------------------------------------------
# stateless rngs
# ------------------------------------------------------------------------------
# TODO(b/171746536): switch all rng ops to official versions once supported.
def _old_tf_random_stateless_normal(
shape,
seed,
mean=0.0,
stddev=1.0,
dtype=dtypes.float32,
name=None,
layout=None,
):
"""DTensor stateless normal implementation that takes an layout."""
with ops.name_scope(
name, "stateless_random_normal", [shape, seed, mean, stddev]
) as name:
seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed")
shape = shape_util.shape_tensor(shape)
mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
rnd = api.call_with_layout(
gen_stateless_random_ops.stateless_random_normal,
layout,
shape,
seed,
dtype,
)
result = math_ops.add(rnd * stddev, mean, name=name)
shape_util.maybe_set_static_shape(result, shape)
return result
def _old_tf_random_stateless_uniform(
shape,
seed,
minval=0,
maxval=None,
dtype=dtypes.float32,
name=None,
layout=None,
):
"""DTensor stateless uniform implementation that takes an layout."""
dtype = dtypes.as_dtype(dtype)
accepted_dtypes = (
dtypes.float16,
dtypes.bfloat16,
dtypes.float32,
dtypes.float64,
dtypes.int32,
dtypes.int64,
dtypes.uint32,
dtypes.uint64,
)
if dtype not in accepted_dtypes:
raise ValueError(
f"Argument `dtype` got invalid value {dtype}. Accepted dtypes are "
f"{accepted_dtypes}."
)
if dtype.is_integer:
if (minval is None) != (maxval is None):
raise ValueError(
f"For integer `dtype` argument {dtype}, argument `minval` and "
f"`maxval` must be both None or not None. Got `minval`={minval} and "
f"`maxval`={maxval}."
)
if minval is not None and dtype in (dtypes.uint32, dtypes.uint64):
raise ValueError(
f"Argument `dtype` got invalid value {dtype} when argument `minval` "
"is not None. Please don't use unsigned integers in this case."
)
shape = shape_util.shape_tensor(shape)
with ops.name_scope(
name, "stateless_random_uniform", [shape, seed, minval, maxval]
) as name:
seed = ops.convert_to_tensor(seed, dtype_hint=dtypes.int32, name="seed")
if dtype.is_integer and minval is None and maxval is None:
result = api.call_with_layout(
gen_stateless_random_ops.stateless_random_uniform_full_int,
layout,
shape,
seed=seed,
dtype=dtype,
name=name,
)
else:
if not dtype.is_integer and maxval is None:
maxval = 1
val_range = ops.convert_to_tensor(
maxval - minval, dtype=dtype, name="range"
)
minval = ops.convert_to_tensor(minval, dtype=dtype, name="min")
if dtype.is_integer:
result = api.call_with_layout(
gen_stateless_random_ops.stateless_random_uniform_int,
layout,
shape,
seed=seed,
minval=minval,
maxval=maxval,
)
else:
rnd = api.call_with_layout(
gen_stateless_random_ops.stateless_random_uniform,
layout,
shape,
seed=seed,
dtype=dtype,
)
result = math_ops.add(rnd * val_range, minval, name=name)
shape_util.maybe_set_static_shape(result, shape)
return result
def _old_tf_stateless_truncated_normal(
shape,
seed,
mean=0.0,
stddev=1.0,
dtype=dtypes.float32,
name=None,
layout=None,
):
"""DTensor stateless truncated normal implementation that takes an layout."""
with ops.name_scope(
name, "stateless_truncated_normal", [shape, seed, mean, stddev]
) as name:
seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed")
shape = shape_util.shape_tensor(shape)
mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
rnd = api.call_with_layout(
gen_stateless_random_ops.stateless_truncated_normal,
layout,
shape,
seed,
dtype,
)
result = math_ops.add(rnd * stddev, mean, name=name)
shape_util.maybe_set_static_shape(result, shape)
return result
def stateless_random_normal(
shape,
seed,
mean=0.0,
stddev=1.0,
dtype=dtypes.float32,
name=None,
layout=None,
):
"""DTensor stateless RNG."""
if not context.executing_eagerly():
layout = None
return _old_tf_random_stateless_normal(
shape,
seed=seed,
mean=mean,
stddev=stddev,
dtype=dtype,
name=name,
layout=layout,
)
def stateless_random_uniform(
shape,
seed,
minval=0,
maxval=None,
dtype=dtypes.float32,
name=None,
layout=None,
):
"""DTensor stateless random uniform."""
if not context.executing_eagerly():
layout = None
return _old_tf_random_stateless_uniform(
shape,
seed=seed,
minval=minval,
maxval=maxval,
dtype=dtype,
name=name,
layout=layout,
)
def stateless_truncated_normal(
shape,
seed,
mean=0.0,
stddev=1.0,
dtype=dtypes.float32,
name=None,
layout=None,
):
"""DTensor stateless RNG."""
if not context.executing_eagerly():
layout = None
return _old_tf_stateless_truncated_normal(
shape,
seed=seed,
mean=mean,
stddev=stddev,
dtype=dtype,
name=name,
layout=layout,
)
def stateless_split(seed, num=2, mesh=None):
seed = ops.convert_to_tensor(seed)
layout = None
if mesh:
layout = layout_lib.Layout.replicated(mesh, rank=2)
return stateless_random_uniform(
shape=[num, 2],
seed=seed,
dtype=seed.dtype,
minval=None,
maxval=None,
layout=layout,
)
# ------------------------------------------------------------------------------
# stateless dropout.
# ------------------------------------------------------------------------------
def _get_noise_shape(x, noise_shape):
"""Noisve shape util copied from tf nn_ops."""
# If noise_shape is none return immediately.
if noise_shape is None:
return array_ops.shape(x)
try:
# Best effort to figure out the intended shape.
# If not possible, let the op to handle it.
# In eager mode exception will show up.
noise_shape_ = tensor_shape.as_shape(noise_shape)
except (TypeError, ValueError):
return noise_shape
if x.shape.dims is not None and len(x.shape.dims) == len(noise_shape_.dims):
new_dims = []
for i, dim in enumerate(x.shape.dims):
if noise_shape_.dims[i].value is None and dim.value is not None:
new_dims.append(dim.value)
else:
new_dims.append(noise_shape_.dims[i].value)
return tensor_shape.TensorShape(new_dims)
return noise_shape
# TODO(b/171213877, b/169909066): Fix layout prop in function case for the rng
# Op used. The layout prop should be able to propagate the layout from input
# tensor `x` to the tf.mul and then back propagate the layout to the
# `random_tensor`.
def dropout(x, rate, noise_shape=None, seed=None, name=None):
"""DTensor replacement for dropout."""
if not isinstance(rate, float):
raise ValueError("rate should be float for dropout.")
if seed is None:
raise ValueError("seed must be specified for DTensor dropout. Got: None")
with ops.name_scope(name, "dropout", [x]):
x_dtype = x.dtype
keep_prob = 1 - rate
scale = 1 / keep_prob
scale = ops.convert_to_tensor(scale, dtype=x_dtype)
ret = gen_math_ops.mul(x, scale)
noise_shape = _get_noise_shape(x, noise_shape)
# stateless_random_uniform requires a shape [2] seed.
seed = [seed, 0]
if context.executing_eagerly():
layout = api.fetch_layout(x)
else:
layout = None
random_tensor = _old_tf_random_stateless_uniform(
noise_shape, seed=seed, minval=0, maxval=1, dtype=x_dtype, layout=layout
)
keep_mask = random_tensor >= rate
ret = gen_math_ops.mul(ret, gen_math_ops.cast(keep_mask, x_dtype))
if not context.executing_eagerly():
ret.set_shape(x.get_shape())
return ret
# TODO(b/195413777): error out for stateful dropout.
tensorflow/dtensor/python/d_variable.py
+260
View File
@@ -0,0 +1,260 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DTensor variable and saveable."""
import functools
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.trackable import base as trackable
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.util.tf_export import tf_export
class DSaveSpec(saveable_object.SaveSpec):
"""DTensor SaveSpec that additionaly captures global_shape and layout."""
def __init__(self,
tensor,
slice_spec,
name,
global_shape,
layout,
dtype=None,
device=None):
super().__init__(
tensor=tensor,
slice_spec=slice_spec,
name=name,
dtype=dtype,
device=device)
self.global_shape = global_shape
self.layout = layout
class _DVariableSaveable(saveable_object.SaveableObject):
"""Class for defining how to save/restore DTensor variable."""
def __init__(self, dvariable, name):
with ops.device(dvariable.device):
original_layout = api.fetch_layout(dvariable)
# Record original layout to allow restore.
self._original_layout = original_layout
self._dvariable = dvariable
def pack(tensors, layout):
with ops.device(dvariable.device):
return api.pack(tensors, layout)
host_layout = layout_lib.Layout(original_layout.sharding_specs,
original_layout.mesh.host_mesh())
def get_host_dtensor():
# Copy to host mesh if needed.
if original_layout.mesh.device_type().upper() != 'CPU':
# Prefer pack and unpack in eager mode because it supports sharded
# layouts.
if context.executing_eagerly():
host_dtensor = api.pack(
api.unpack(dvariable.read_value()), host_layout)
else:
host_dtensor = api.copy_to_mesh(dvariable.read_value(), host_layout)
else:
host_dtensor = dvariable.read_value()
return (math_ops.cast(host_dtensor, dtypes.bfloat16)
if self.should_cast(host_dtensor) else host_dtensor)
num_local_devices = original_layout.mesh.num_local_devices()
super(_DVariableSaveable, self).__init__(
None,
[
DSaveSpec(
tensor=get_host_dtensor,
slice_spec=pack([''] * num_local_devices,
layout_lib.Layout.replicated(
original_layout.mesh.host_mesh(), rank=0)),
name=pack([name] * num_local_devices,
layout_lib.Layout.replicated(
original_layout.mesh.host_mesh(), rank=0)),
global_shape=dvariable.shape,
# Layout is attached as attribute, no need to put it as a
# Tensor on DTensorDevice.
layout=host_layout.to_string(),
dtype=dtypes.bfloat16
if self.should_cast(dvariable) else dvariable.dtype,
device=dvariable.device)
],
name)
def should_cast(self, v):
"""Returns True if v has float32 dtype and is intructed to save as bf16.
Args:
v : The variable that determines whether to cast.
Returns:
True if current savable DVariable is instructed to save as bfloat16 and
the variable has dtype float32.
"""
return self._dvariable.save_as_bf16 and v.dtype == dtypes.float32
def restore(self, restored_tensors, restored_shapes):
"""Restores the same value into all variables."""
tensor, = restored_tensors
@def_function.function
def _restore(t):
with ops.device(self._dvariable.device):
return api.copy_to_mesh(t, self._original_layout)
# This assign establishes connections from restored tensor and tensors
# being restored to -- so that restore in SPMD can backtrack the DVariable
# and its layout, given that we're using tf.function style restore.
# Note that the restored dvaraible is on CPU no matter what as the restoreV2
# op must run on CPU.
# TODO(b/159035705): Allow restore for Tensor objects as well?
# Restore the dvariable back to original layout.
if self._original_layout.mesh.device_type().upper() != 'CPU':
tensor = _restore(tensor)
return self._dvariable.assign(
math_ops.cast(tensor, dtype=self._dvariable.dtype) if self._dvariable
.save_as_bf16 else tensor)
@tf_export('experimental.dtensor.DVariable', v1=[])
class DVariable(resource_variable_ops.ResourceVariable):
"""A replacement for tf.Variable which follows initial value placement.
The class also handles restore/save operations in DTensor. Note that,
DVariable may fall back to normal tf.Variable at this moment if
`initial_value` is not a DTensor.
"""
def __init__(self, initial_value, *args, dtype=None, **kwargs):
"""Overrides tf.Variable to fix VarHandleOp placements."""
# Variables by default use the current device scope for placement. This
# wrapper has them follow the initial value's placement instead (which will
# be the DTensor device if the initial value has a layout).
# Pop layout from kwargs since keras make_variable may pass a 'layout'
# keyword argument. We need to pop it because we are passing kwargs to
# super class constructor.
layout = kwargs.pop('layout', None)
shape = kwargs.get('shape', None)
if callable(initial_value):
unwrapped = initial_value
if issubclass(type(initial_value), functools.partial):
unwrapped = initial_value.func
# If wrapped is a CheckpointInitialValueCallable, this means that
# we are creating a Variable during a checkpoint restore.
# Thus the restore will happen now through this callable
# and we will create the DVariable with the restored dtensor.
if issubclass(type(unwrapped), trackable.CheckpointInitialValueCallable):
if not shape or not layout:
raise ValueError('Expected shape and layout to be not None.')
# CheckpointInitialValueCallable will call an eager tf.RestoreV2,
# which does not have any shape information or layout information
# attached. Thus we will do two things to have them correctly specified:
#
# The default layout scope allows us to correctly specify the output
# layout of the tf.RestoreV2 that will be called
#
# Passing shard_info with the correct shape allows the tf.RestoreV2
# ShapeInference to extract the shape.
initial_value = api.call_with_layout(
initial_value,
layout,
shard_info=trackable.ShardInfo(
shape=shape, offset=[0] * len(shape)))
else:
initial_value = initial_value()
# When the initial value came from a Checkpoint restoration, fetch tensor.
if isinstance(initial_value, trackable.CheckpointInitialValue):
initial_value = initial_value.wrapped_value
initial_value = ops.convert_to_tensor(initial_value, dtype=dtype)
variable_device = initial_value.device
self._save_as_bf16 = False
# TODO(b/159035705): The following code enables variable creation inside
# a tf.function. However, it requires a global dtensor device.
# if not variable_device and not tf.executing_eagerly():
# try:
# initial_value.op.get_attr("_layout")
# except ValueError:
# pass
# else:
# # The initial value is a DTensor, but because the DTensor device is
# # only active during eager execution at the moment we need to
# # translate that into a placement for the eager VarHandleOp.
# variable_device = _dtensor_device().name
with ops.device(variable_device):
# If initial tensor assigned to DVariable is DTensor, record the layout of
# the resource so that this can be queried.
if context.executing_eagerly():
if api.is_dtensor(initial_value):
value_layout = api.fetch_layout(initial_value)
if layout is not None and layout != value_layout:
raise errors_impl.InvalidArgumentError(
None,
None,
'Conflicting layout are provided for initial '
f'value layout ({value_layout}) and variable ({layout}).',
)
layout = value_layout
elif layout is not None:
initial_value = api.relayout(initial_value, layout)
else:
raise errors_impl.InvalidArgumentError(
None,
None,
'Neither layout nor DTensor initial value are provided.',
)
self.layout = layout
with api.default_mesh(layout.mesh):
super(DVariable, self).__init__(
initial_value, *args, dtype=dtype, **kwargs
)
else:
# FIXME(175928457): Record value layout in graph mode.
if layout is not None:
initial_value = api.relayout(initial_value, layout)
super(DVariable, self).__init__(
initial_value, *args, dtype=dtype, **kwargs)
@property
def save_as_bf16(self):
return self._save_as_bf16
@save_as_bf16.setter
def save_as_bf16(self, save_as_bf16):
"""Enables saving float32 as bfloat16."""
self._save_as_bf16 = save_as_bf16 and self.dtype == dtypes.float32
def _gather_saveables_for_checkpoint(self):
return {
trackable.VARIABLE_VALUE_KEY:
functools.partial(_DVariableSaveable, self)
}
tensorflow/dtensor/python/dtensor_device.py
+437
View File
@@ -0,0 +1,437 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Propagates information about tensor layouts across operations."""
import contextlib
import logging
import threading
from typing import Any, List, Sequence, Set
import numpy as np
from tensorflow.core.framework import attr_value_pb2
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.python import _pywrap_dtensor_device
from tensorflow.python.eager import context
from tensorflow.python.eager import core
from tensorflow.python.framework import device as tf_device
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_util
from tensorflow.python.util import _pywrap_utils
# TODO(allenl): Allow something other than "CUSTOM" so we don't need device
# numbering hacks to avoid collisions between parallel devices and dtensor
# devices.
_next_device_number = 0
_next_device_number_lock = threading.Lock()
class DTensorDevice(object):
"""Wraps a custom device which attempts to propagate tensor layouts."""
def __init__(self,
meshes: List[layout_lib.Mesh],
is_async=True,
in_flight_nodes_limit=8):
"""Create a new DTensorDevice which executes ops on `underlying_device`.
Args:
meshes: A list of `Mesh` objects indicating groups of devices to execute
on. These may also be registered lazily.
is_async: Indicates whether DTensor operations on this client will return
immediately (with "non-ready" handles) or block until executed. This is
on by default and is exposed as an option for ease of debugging.
in_flight_nodes_limit: Indicates the limit of in-flight nodes before
enqueueing of async operations to DTensorDevice is blocked. This limit
is per mesh. 0 for no limits from DTensor. Default is 8.
"""
if any(not isinstance(mesh, layout_lib.Mesh) for mesh in meshes):
raise TypeError(
"Expected a flat list of Mesh objects, got {}".format(meshes))
global _next_device_number
ctx = context.context()
with _next_device_number_lock:
self.name = "{}/device:CUSTOM:{}".format(ctx.host_address_space(),
_next_device_number)
_next_device_number += 1
device, device_info = _pywrap_dtensor_device.Allocate(
self.name, is_async, in_flight_nodes_limit
)
context.register_custom_device(device, self.name, device_info)
self._device_info = device_info
self._current_output_layout = None
self._current_default_mesh = None
self._meshes = set()
self._mesh_lock = threading.Lock()
for mesh in meshes:
self._register_mesh(mesh)
def _create_host_array(self, shape, host_id):
"""Returns ID and device lists that can be used to create a host mesh."""
num_global_devices = np.prod(shape)
global_device_ids = np.arange(num_global_devices).reshape(shape)
local_device_list = [
tf_device.DeviceSpec(
job=config.full_job_name(), device_type="CPU", device_index=0)
]
num_local_devices = len(local_device_list)
local_device_ids = [
x + host_id * num_local_devices for x in range(num_local_devices)
]
return global_device_ids, local_device_ids, local_device_list
def _register_mesh(self, mesh: layout_lib.Mesh):
"""Idempotently register `mesh` with the dtensor device."""
with self._mesh_lock:
if mesh not in self._meshes:
_pywrap_dtensor_device.AddMesh(
self._device_info, mesh.to_string(), False
)
self._meshes.add(mesh)
if mesh.device_type().upper() == "TPU":
logging.info(
"Registering virtual 1:1 mapped host mesh %s for mesh %s",
mesh.host_mesh().to_string(), mesh.to_string())
_pywrap_dtensor_device.AddMesh(
self._device_info, mesh.host_mesh().to_string(), True
)
self._meshes.add(mesh.host_mesh())
@property
def meshes(self) -> Set[layout_lib.Mesh]:
return self._meshes
def pack(self, tensors: Sequence[Any], layout: layout_lib.Layout) -> Any:
"""Packs tensors into a DTensor handle on this DTensor device.
Packing and unpacking are inverse operations:
```
* unpack(pack(tensors)) == tensors
* pack(unpack(dtensor)) == dtensor
```
Refer to `dtensor.pack` for more information.
Args:
tensors: The list of tensors to pack into a DTensor.
layout: The layout of the DTensor to be created.
Returns:
A DTensor created from the individual component tensors.
Raises:
RuntimeError: When not called eagerly.
"""
if not context.executing_eagerly():
raise RuntimeError("`pack` must be called eagerly.")
self._register_mesh(layout.mesh)
with ops.device(self.name):
if all(isinstance(t, sparse_tensor.SparseTensor) for t in tensors):
if not all(t.shape == tensors[0].shape for t in tensors):
raise TypeError("All input SparseTensors to Pack must be same shape.")
is_sparse = True
tensors = [t.indices for t in tensors] + [t.values for t in tensors] + [
ops.convert_to_tensor(t.shape, dtype=dtypes.int64) for t in tensors
]
elif any(isinstance(t, sparse_tensor.SparseTensor) for t in tensors):
raise TypeError("Cannot Pack SparseTensors with Tensors.")
else:
is_sparse = False
try:
return _pywrap_dtensor_device.Pack(
context.context()._handle, # pylint: disable=protected-access
tensors,
layout.to_string(),
self._device_info,
is_sparse)
except core._NotOkStatusException as e: # pylint: disable=protected-access
raise core._status_to_exception(e) from None # pylint: disable=protected-access
def unpack(self, dtensor: Any) -> Sequence[Any]:
"""Unpacks a DTensor handle on this DTensor device.
Packing and unpacking are inverse operations:
```
* unpack(pack(tensors)) == tensors
* pack(unpack(dtensor)) == dtensor
```
Refer to `dtensor.unpack` for more information.
Args:
dtensor: The DTensor to unpack.
Returns:
The raw underlying tensor components of the DTensor.
Raises:
RuntimeError: When not called eagerly.
"""
if not context.executing_eagerly():
raise RuntimeError("`unpack` must be called eagerly.")
try:
tensors = _pywrap_dtensor_device.Unpack(
context.context()._handle, # pylint: disable=protected-access
dtensor,
self._device_info)
except core._NotOkStatusException as e: # pylint: disable=protected-access
raise core._status_to_exception(e) from None # pylint: disable=protected-access
is_sparse = _pywrap_dtensor_device.IsSparseDTensor(
context.context()._handle, # pylint: disable=protected-access.
dtensor,
self._device_info)
if is_sparse:
result = []
for i in range(len(tensors) // 3):
result.append(
sparse_tensor.SparseTensor(tensors[i],
tensors[i + len(tensors) // 3],
tensors[i + 2 * len(tensors) // 3]))
return result
else:
return tensors
def fetch_layout(self, dtensor: Any) -> layout_lib.Layout:
"""Fetches the layout of the DTensor.
Args:
dtensor: The DTensor whose layout is to be fetched.
Returns:
The `Layout` of this DTensor.
Raises:
RuntimeError: When not called eagerly.
"""
if not context.executing_eagerly():
raise RuntimeError("`fetch_layout` must be called eagerly.")
if _pywrap_utils.IsVariable(dtensor):
dtensor = dtensor.read_value()
try:
layout_string = _pywrap_dtensor_device.FetchLayout(
context.context()._handle, # pylint: disable=protected-access
dtensor,
self._device_info)
except core._NotOkStatusException as e: # pylint: disable=protected-access
raise core._status_to_exception(e) from None # pylint: disable=protected-access
if layout_string is None:
return None
return layout_lib.Layout.from_string(layout_string)
def is_dtensor(self, tensor: Any) -> bool:
"""Check whether the input tensor is a DTensor.
In Python, a DTensor has the same type as a `tf.Tensor`. This method will
let you check and handle the tensor differently if a tf.Tensor is a DTensor.
Args:
tensor: an object to be checked.
Returns:
bool, True if the given tensor is a DTensor.
Raises:
RuntimeError: When not called eagerly.
"""
if not context.executing_eagerly():
raise RuntimeError("`is_dtensor` must be called eagerly.")
if not tensor_util.is_tensor(tensor):
return False
if _pywrap_utils.IsVariable(tensor):
tensor = tensor._handle # pylint: disable=protected-access
return _pywrap_dtensor_device.IsDTensor(
context.context()._handle, # pylint: disable=protected-access
tensor,
self._device_info,
)
def set_tpu_core_ids(self, mesh_name, tpu_core_ids):
"""Sets the singleton global device ID-to-physical core ID map.
Args:
mesh_name: The name of a mesh. If empty, set the default mapping.
tpu_core_ids: TPU core IDs sorted by TF task/device ordinal.
"""
_pywrap_dtensor_device.SetTPUCoreIDs(self._device_info, mesh_name,
tpu_core_ids)
def clear_tpu_core_ids(self):
_pywrap_dtensor_device.ClearTPUCoreIDs(self._device_info)
def tpu_core_ids_to_locations(self, tpu_core_ids):
"""Translates TPU core IDs to TPU core locations.
Args:
tpu_core_ids: A list of TPU core IDs. Each one is an unsigned integer.
Returns:
A list of corresponding TPU core locations.
"""
return _pywrap_dtensor_device.TPUCoreIDsToLocations(
context.context()._handle, # pylint: disable=protected-access
self._device_info,
tpu_core_ids)
def tpu_core_locations_to_ids(self, tpu_core_locations):
"""Translates TPU core locations to TPU core IDs.
Args:
tpu_core_locations: A list of TPU core locations. Each one is a list of
four unsigned integers, [x, y, z, core].
Returns:
A list of corresponding TPU core IDs.
"""
return _pywrap_dtensor_device.TPUCoreLocationsToIDs(
context.context()._handle, # pylint: disable=protected-access
self._device_info,
tpu_core_locations)
def _get_stats(self):
"""Returns the number of cache hit and miss for function compilation.
Returns:
A dictionary.
'miss': number of cache misses;
'hit': number of cache hits; and
'size': size of cache;
miss count.
"""
return _pywrap_dtensor_device.GetStats(
context.context()._handle, # pylint: disable=protected-access,
self._device_info,
)
def set_iterator_element_layouts(self, iterator_resource_dtensor,
layouts: List[layout_lib.Layout]):
"""Sets the element layouts on an iterator resource tensor.
Args:
iterator_resource_dtensor: a DTensor created by packing the individiual
iterator resource tensors.
layouts: the flattened list of layouts to be applied to the elements
emitted by the iterator resource DTensor.
"""
_pywrap_dtensor_device.SetIteratorElementLayouts(
context.context()._handle, # pylint: disable=protected-access
iterator_resource_dtensor,
[layout.to_string() for layout in layouts],
self._device_info)
@contextlib.contextmanager
def _experimental_default_mesh(self, mesh: layout_lib.Mesh):
"""Sets a default mesh for all ops in the scope.
Note: This is an internal helper method, which is not user facing api.
Useful for requesting a specific mesh for ops which would have no inferred
layout, e.g. tf.zeros.
Args:
mesh: A Mesh to be used for ops without Mesh.
Yields:
Nothing.
"""
previous_default = self._current_default_mesh
self._register_mesh(mesh)
_pywrap_dtensor_device.ExperimentalSetDefaultMesh(
self._device_info,
mesh.to_string().encode("utf-8"))
self._current_default_mesh = mesh
yield
_pywrap_dtensor_device.ExperimentalClearDefaultMesh(self._device_info)
if previous_default:
_pywrap_dtensor_device.ExperimentalSetDefaultMesh(
self._device_info,
previous_default.to_string().encode("utf-8"))
self._current_default_mesh = previous_default
@contextlib.contextmanager
def _default_layout(self, layout: layout_lib.Layout):
"""Sets a default output layout for all ops in the scope.
Note: This is an internal helper method, which is not user facing api.
Useful for requesting a specific layout for ops which would have no inferred
layout, e.g. tf.zeros.
Caveats:
- Currently only affects the first output of an op. For Op with multiple
outputs, this does not support yet.
- All Ops in the scope will be attached with the same layout. This might not
be valid as the rank is different. The current suggestion is: Try to wrap
the raw op wheneven possible.
Args:
layout: A Layout for the outputs of all operations in this scope.
Yields:
Nothing.
"""
previous_default = None
previous_graph_size = None
graph = None
self._register_mesh(layout.mesh)
try:
previous_default = self._current_output_layout
self._current_output_layout = layout.to_string().encode("utf-8")
_pywrap_dtensor_device.ExperimentalSetDefaultLayout(
self._device_info, self._current_output_layout)
if context.executing_eagerly():
with ops.device(self.name):
yield
else:
# Custom devices currently don't affect graph building, so we need a
# separate way to indicate layouts.
#
# TODO(allenl): Remove this case once the DTensor device is active
# during tracing.
graph = ops.get_default_graph()
previous_graph_size = len(graph.get_operations())
yield
finally:
if graph is not None:
# Tag operations added under this scope
for operation in graph.get_operations()[previous_graph_size:]:
# Set layout directly on the Op itself.
operation._set_attr( # pylint: disable=protected-access
"_layout",
attr_value_pb2.AttrValue(
list=attr_value_pb2.AttrValue.ListValue(
s=[self._current_output_layout])))
operation._set_attr( # pylint: disable=protected-access
"_mesh",
attr_value_pb2.AttrValue(
s=layout.mesh.to_string().encode("utf-8")))
self._current_output_layout = previous_default
if self._current_output_layout is None:
_pywrap_dtensor_device.ExperimentalClearDefaultLayout(self._device_info)
else:
_pywrap_dtensor_device.ExperimentalSetDefaultLayout(
self._device_info, self._current_output_layout.decode("utf-8"))
tensorflow/dtensor/python/heartbeat.py
+178
View File
@@ -0,0 +1,178 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""A heartbeat service periodically pinging all workers.
In normal cases, all workers will exchange the same randomly generated number
until normal program termination. If any worker stops or restarts, other workers
will detect that and crash themselves.
In this module, logging.fatal is used to guarantee a worker crash no matter how
the functions below are called, in a thread or not.
"""
import atexit
import threading
import time
import numpy as np
from tensorflow.dtensor.python import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import device as tf_device
from tensorflow.python.framework import ops
from tensorflow.python.ops.collective_ops import all_reduce
from tensorflow.python.platform import tf_logging as logging
# More than these many consecutive failures will cause a crash.
_CONSECUTIVE_FAILURES_LIMIT = 3
_failure_count = 0
_heartbeat_timer = None
def _heartbeat(
period: int, # in seconds
timer: threading.Event,
token: int,
num_tasks: int,
task_id: int,
device: tf_device.DeviceSpec,
):
"""Periodically sends and receives a heartbeat signal."""
logging.info('Starting a heartbeat thread')
global _failure_count
while True:
# `timer.wait` blocks until one of two things happens.
# It returns True if the timer is explicitly set at process exit, and we
# should gracefully end this heartbeat thread.
# Otherwise, it returns False when `period` has elapsed, meaning it's time
# for the next heartbeat exchange.
# See https://docs.python.org/3/library/threading.html#threading.Event.wait.
if timer.wait(period):
logging.info('Exiting the heartbeat thread normally')
return
# Every worker fills in one element of the signal with `token`.
signal = np.zeros([num_tasks], dtype=np.int32)
signal[task_id] = token
logging.vlog(2, 'Sending heartbeat signal %s', signal)
try:
with ops.device(device):
# Always use 0 for group and instance keys to reduce unnecessary
# collective hangs and simplify failure analysis. This also avoid
# collision with normal collectives.
signal = all_reduce(
constant_op.constant(signal),
group_size=num_tasks,
group_key=0,
instance_key=0,
timeout=max(period - 10, 2)).numpy()
except Exception as e: # pylint: disable=broad-except
_failure_count += 1
if _failure_count < _CONSECUTIVE_FAILURES_LIMIT:
logging.warning('Heartbeat failure %d, %d more until limit: %s',
_failure_count,
_CONSECUTIVE_FAILURES_LIMIT - _failure_count, e)
else:
logging.fatal('Heartbeat failure %d, limit of %d reached: %s',
_failure_count, _CONSECUTIVE_FAILURES_LIMIT, e)
logging.vlog(2, 'Received heartbeat signal %s', signal)
# Out of sync workers will cause this, crash immediately.
if not np.all(signal == token):
logging.fatal('Unexpected heartbeat signal received: %s', signal)
# Any success resets the failure counter.
_failure_count = 0
def start(period: int) -> threading.Event:
"""Starts a persistent thread exchanging heartbeats between workers.
Args:
period: Heartbeat interval in seconds. Heartbeat timeout is set to the
larger of `period` - 10 and 2s.
Returns:
A threading.Event object. Users can choose to call its set() method to shut
down the heartbeat service gracefully. This isn't necessary in most cases,
because the heartbeat service automatically shuts down at successful program
exit through atexit handlers. But in situations when atexit handlers are not
invoked, such as when multiprocessing processes exit in tests, users can
manually request a shutdown.
"""
global _heartbeat_timer
if _heartbeat_timer is not None:
logging.warning('A heartbeat thread is already running, skipping this one.')
return _heartbeat_timer
task_id = config.client_id()
num_tasks = config.num_clients()
# Worker 0 generates a random token. All other workers receive that token.
if task_id == 0:
token = np.random.randint(0, pow(2, 16) - 1) # reserve the other 16 bits
signal = np.full([num_tasks], token, dtype=np.int32)
else:
signal = np.zeros([num_tasks], dtype=np.int32)
logging.info('Initial heartbeat signal: %s', signal)
device = tf_device.DeviceSpec(
job=config.job_name(),
replica=0,
task=task_id,
device_type='CPU',
device_index=0)
# Always use 0 for group and instance keys to reduce unnecessary
# collective hangs and simplify failure analysis. This also avoid
# collision with normal collectives.
with ops.device(device):
signal = all_reduce(
constant_op.constant(signal),
group_size=num_tasks,
group_key=0,
instance_key=0,
timeout=max(period - 10, 2)).numpy()
logging.info('Merged heartbeat signal %s', signal)
# The merged signal should have equal elements. If not, some worker(s) may be
# out of sync, and we should terminate all workers.
if task_id == 0:
if not np.all(signal == token):
logging.fatal('Merged heartbeat signal has value != %d', token)
else:
if len(set(signal)) != 1:
logging.fatal('Merged heartbeat signal has unequal elements')
token = signal[0]
# On normal main process exit, set the timer to stop the heartbeat thread.
_heartbeat_timer = threading.Event()
def stop_heartbeat():
logging.info('Stopping the heartbeat thread')
_heartbeat_timer.set()
# Give the threads some time to clean up.
time.sleep(max(period // 10, 2))
atexit.register(stop_heartbeat)
# Start the persistent heartbeat thread.
thread = threading.Thread(
target=_heartbeat,
args=[period, _heartbeat_timer, token, num_tasks, task_id, device],
daemon=True)
thread.start()
return _heartbeat_timer
tensorflow/dtensor/python/input_util.py
+666
View File
@@ -0,0 +1,666 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""APIs to deal with input datasets efficiently in DTensor.
When using tf.data with DTensor, the `DTensorDataset` API can be used to
efficiently handle loading the input data and correctly packing it to the
corresponding devices. This API is intended to work with unbatched data and can
be used for both data and model parallel setups.
Example usage:
>>> # 1-D mesh with 4 devices
>>> mesh = dtensor.Mesh(dim_names=['batch'], ...)
>>> layout = dtensor.Layout.batch_sharded(mesh, 'batch', rank=1)
>>> dataset = tf.data.Dataset.range(256)
>>> d_dataset = dtensor.DTensorDataset(
... dataset=dataset,
... global_batch_size=16,
... mesh=mesh,
... layouts=layout,
... batch_dim='batch')
>>> d_iter = iter(d_dataset)
>>> # Each batch is a length 16 tensor sharded across 4 devices
>>> batch_0_dtensor = next(d_iter)
>>> batch_0_dtensor
<tf.Tensor: shape=(16,),
dtype=int64,
value={"CPU:0": [ 0 1 2 4],
"CPU:1": [ 5 6 7 8],
"CPU:2": [ 9 10 11 12],
"CPU:3": [13 14 15 16]}>
>>> batch_1_dtensor = next(d_iter)
>>> batch_1_dtensor
<tf.Tensor: shape=(16,),
dtype=int64,
value={"CPU:0": [17 18 19 20],
"CPU:1": [21 22 23 24],
"CPU:2": [25 26 27 28],
"CPU:3": [29 30 31 32]}>
For multi-client setups, `DTensorDataset` interacts with tf.data service to
correctly distribute the dataset among the participating clients. DTensor works
with tf.data service in co-located mode where each worker is running alongside
the DTensor client (the Tensorflow Python process). The `TFDataServiceConfig`
dataclass can be filled with information about the tf.data service cluster, and
passed to `DTensorDataset` to enable distribution.
"""
import dataclasses
import operator
from typing import Any, List, Optional, Sequence, Tuple
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.python.data.experimental.ops import data_service_ops
from tensorflow.python.data.experimental.ops import distribute
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.types import data as data_types
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export
@dataclasses.dataclass
class TFDataServiceConfig:
"""Specifies the tf.data service configuration to use.
Attributes:
dispatcher_address: a string specifying the address of the tf.data service
dispatcher server.
job_name: a non-empty string identifying the shared job that will be created
on tf.data service to process this dataset.
"""
dispatcher_address: str
job_name: str
# TODO(b/223275517): Add support for get_next_as_optional().
class _DTensorIterator(iterator_ops.OwnedIterator):
"""An iterator for a tf.data.Dataset distributed using DTensor.
DTensorIterator encapsulates multiple underlying dataset iterators. It handles
retrieving the tensors to be placed on each underlying device and then uses
the 'pack' operation to create and return a DTensor. Thus users need only
interact with a single DTensorIterator to automatically distribute dataset
tensors onto devices.
"""
def __init__(
self,
dtensor_components: Tuple[tensor.Tensor],
global_element_spec: tensor_spec.TensorSpec,
layouts: Any):
"""Initializes a distributed iterator for DTensor datasets.
This iterator encapsulates tf.data iterators for the underlying devices, and
treats it as a packed DTensor of iterator resource tensors.
Args:
dtensor_components: a tuple containing the underlying iterator resources
packed into a DTensor. This is expected to be a tuple with a single
element.
global_element_spec: the underlying dataset's element spec from a global
view.
layouts: a structure of DTensor layouts to be applied to the elements
returned by the underlying iterators. This can be a single layout or
(possibly nested) tuples or dictionaries of layouts, and the structure
must match the structure of the iterator elements.
"""
# dtensor_components is expected to be a single-element tuple.
[self._iterator_resource_dtensor] = dtensor_components
self._global_element_spec = global_element_spec
self._layouts = layouts
self._layouts_str = nest.map_structure(
lambda layout: layout.to_string(), layouts)
super().__init__(
components=dtensor_components, element_spec=global_element_spec)
def __next__(self):
try:
# IteratorGetNext will return a DTensor on the host, so move it to the
# device mesh. If the dataset layouts are on the host mesh itself, this
# is handled by DTensor as a no-op.
host_elem = self._next_internal()
context.async_wait()
device_elem = nest.map_structure(
api.copy_to_mesh, host_elem, self._layouts)
context.async_wait()
return device_elem
except errors.OutOfRangeError as e:
# Match TF2 eager executor behavior by raising StopIteration when iterator
# is out of range.
if context.executing_eagerly():
raise StopIteration from e
else:
raise e
@property
def _type_spec(self):
return _DTensorIteratorSpec(self._global_element_spec, self._layouts_str)
class _DTensorIteratorSpec(iterator_ops.IteratorSpec):
"""Type specification for `_DTensorIterator`."""
__slots__ = ['_global_element_spec', '_layouts_str']
def __init__(
self, global_element_spec: tensor_spec.TensorSpec, layouts_str: Any):
super().__init__(global_element_spec)
self._global_element_spec = global_element_spec
self._layouts_str = layouts_str
@property
def value_type(self):
return _DTensorIterator
def _serialize(self):
return (self._global_element_spec, self._layouts_str)
@property
def _component_specs(self):
return (tensor_spec.TensorSpec([], dtypes.resource),)
def _to_components(self, value):
return (value._iterator_resource_dtensor,) # pylint: disable=protected-access
def _from_components(self, components):
layouts = nest.map_structure(
layout_lib.Layout.from_string, self._layouts_str)
return _DTensorIterator(
dtensor_components=components,
global_element_spec=self._global_element_spec,
layouts=layouts)
@classmethod
def from_value(cls, value):
return cls(value._global_element_spec, value._layouts_str) # pylint: disable=protected-access
def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
flattened_elem_spec: Sequence[tensor_spec.TensorSpec],
dataset_already_batched: bool):
"""Checks that the dataset's layouts and element specs are compatible.
Args:
flattened_layouts: the flattened list of layouts used to distribute the
dataset.
flattened_elem_spec: the flattened list of element specs used in the
dataset's components.
dataset_already_batched: whether the dataset to be validated is already
batched.
Raises:
ValueError: if the dataset's inputs are incompatible.
"""
if not flattened_elem_spec:
raise ValueError(
'Expected input element spec of at least one element, was empty.')
first_elem_shape = flattened_elem_spec[0].shape
for layout, elem_spec in zip(flattened_layouts, flattened_elem_spec):
if elem_spec.shape.rank is None:
raise ValueError(
'Dataset element shape must have a valid rank, got spec %s.' %
elem_spec)
# Check that layout's rank matches the element's rank. If dataset is not yet
# batched, then the layout's rank must be one greater than the element's
# rank.
expected_rank = elem_spec.shape.rank
if not dataset_already_batched:
expected_rank += 1
if layout.rank != expected_rank:
raise ValueError(
('Expected layout with rank %d for element spec %s, got layout %s. '
'Check that the dataset is not batched before passing to '
'DTensorDataset.') %
(expected_rank, elem_spec, layout.sharding_specs))
if dataset_already_batched:
# Check that the batch dimension size of all dataset elements match.
batch_dim_size = first_elem_shape.as_list()[0]
if batch_dim_size is None:
raise ValueError(
('Size of batch dimension of element spec %s is None. Ensure '
'drop_remainder=True when batching the dataset.') % elem_spec)
if elem_spec.shape.as_list()[0] != batch_dim_size:
raise ValueError(
('Size of batch dimension of element spec %s does not match '
'expected size %d.') % (elem_spec, batch_dim_size))
def _shard_counts(layout: layout_lib.Layout,
batch_dim: Optional[str] = None) -> List[int]:
"""Computes a list of the number of shards in each dimension of the layout.
The shard counts are used to slice each dataset element. The batch dimension's
count is overridden to 1 since we only consider how many shards to make
locally (within each local replica). Sharding across clients is handled by
either tf.data.Dataset's shard transformation (in the single-client case) or
tf.data service's distribute function (in the multi-client case).
Args:
layout: the layout to compute the shard counts for.
batch_dim: the name of the batch dimension of the layout, if present.
Returns:
A list of shard counts, one element per dimension of the layout.
"""
shard_counts = []
for spec in layout.sharding_specs:
if spec in (batch_dim, layout_lib.UNSHARDED):
shard_counts.append(1)
else:
shard_counts.append(layout.mesh.dim_size(spec))
return shard_counts
def _index_matrix(layout: layout_lib.Layout,
elem_spec: tensor_spec.TensorSpec) -> tensor.Tensor:
"""Computes a utility matrix to derive device-based slice offsets.
This function builds a matrix of shape `[mesh.rank, layout.rank]` for each
dataset element. This matrix can be used to slice the DTensor components
returned by the iterator according to the local device that component is to be
placed on. This can be done by multiplying the device offsets of shape
`[1, mesh.rank]` with this index matrix to get a `[1, layout.rank]` shape
tensor containing the slice offsets.
Note: the index on the batch dim is always 0 since sharding on the batch
dimension is handled by either tf.data.Dataset's shard transformation (in the
single-client case) or tf.data service's distribute function (in the
multi-client case). If there is no sharding on the batch dimension (or any
other dimension), the slice index remains 0.
Args:
layout: the layout of the dataset element.
elem_spec: the spec of the dataset element.
Returns:
The index matrix as a tensor.
"""
matrix = []
for dim in layout.mesh.dim_names:
row = [0]
for layout_idx, spec in enumerate(layout.sharding_specs[1:]):
if spec == layout_lib.UNSHARDED or spec != dim:
row.append(0)
else:
row.append(elem_spec.shape[layout_idx] // layout.mesh.dim_size(dim))
matrix.append(row)
return constant_op.constant(matrix, dtype=dtypes.int32)
def _pack_iterator_resource_dtensor(
datasets: List[Tuple[int, data_types.DatasetV2]],
layouts: Any,
mesh: layout_lib.Mesh,
num_local_devices_per_replica: int):
"""Creates a DTensor iterator resource for the per-replica datasets.
Given a list of replica ID to tf.data.Dataset mappings, this function creates
iterators for each device and then packs the underlying iterator resource
tensors into a single DTensor. This resource tensor is used by the
IteratorGetNext op to retrieve the next element in the dataset.
Args:
datasets: a list of tuples of each unique local replica ID to the dataset
object whose elements will be placed on the devices corresponding to that
replica.
layouts: a structure of DTensor layouts to be applied to the elements
returned by the underlying iterators. This can be a single layout or
(possibly nested) tuples or dictionaries of layouts, and the structure
must match the structure of the iterator elements.
mesh: the DTensor mesh to place the iterator batches on.
num_local_devices_per_replica: the number of devices in each data-parallel
replica.
Returns:
A DTensor of the underlying iterator resource tensors.
"""
host_mesh_devices = mesh.host_mesh().local_devices()
device_idx = 0
iterators = []
for _, dataset in datasets:
for idx in range(num_local_devices_per_replica):
with ops.device_v2(host_mesh_devices[device_idx]):
device_dataset = dataset.shard(
num_shards=num_local_devices_per_replica, index=idx)
iterators.append(iter(device_dataset))
device_idx += 1
if device_idx != len(host_mesh_devices):
raise ValueError(
'The `datasets` argument does not have the correct number of'
f' underlying datasets, found {device_idx} but expected'
f' {len(host_mesh_devices)}.')
host_layouts = nest.map_structure(
lambda l: layout_lib.Layout(l.sharding_specs, mesh.host_mesh()), layouts)
# Pack the iterator resource tensors into a replicated 0-dimensional DTensor
# and set the element layouts.
iterator_resources = [it._iterator_resource for it in iterators] # pylint: disable=protected-access
d_iterator_resource = api.pack(
iterator_resources,
layout_lib.Layout.replicated(mesh=mesh.host_mesh(), rank=0))
api._dtensor_device().set_iterator_element_layouts( # pylint: disable=protected-access
d_iterator_resource, nest.flatten(host_layouts))
return d_iterator_resource
@tf_export('experimental.dtensor.DTensorDataset', v1=[])
class DTensorDataset(dataset_ops.UnaryUnchangedStructureDataset):
"""A dataset of DTensors.
DTensorDataset encapsulates a `tf.data.Dataset` whose elements are
automatically packed and returned as DTensors based on a given mesh and
layouts.
"""
def __init__(self,
dataset: data_types.DatasetV2,
*,
mesh: layout_lib.Mesh,
layouts: Any,
global_batch_size: int,
dataset_already_batched: bool = False,
batch_dim: Optional[str] = None,
prefetch: Optional[int] = None,
tf_data_service_config: Optional[TFDataServiceConfig] = None):
"""Creates a DTensorDataset.
DTensorDataset automatically handles distribution of the dataset elements to
each client's devices. It can be used to create an iterator that returns
DTensors of the input data on each iteration.
DTensorDataset works best with unbatched datasets. It takes the mesh and the
provided layouts to automatically calculate how to batch the input locally
for each replica.
If the provided dataset is already batched according to the per-replica
batch size, then `dataset_already_batched` must be set and DTensorDataset
will check that the batch size is consistent with the intended
`global_batch_size` using the layout information. Each replica receives a
separate slice of the global batch, thus the per-replica batch size can be
computed as the global batch size divided by the number of model replicas.
For a DTensor mesh, the number of replicas is equal to the size of the
mesh's batch dimension.
Note: `tf.experimental.dtensor.DTensorDataset` instances do *not* implement
the full interface of `tf.data.Dataset`. It only supports two usages we will
mention below: iteration and `element_spec`. We don't support any other APIs
to transform or inspect the dataset.
TODO(b/223275517): add support for input datasets that are already batched
to the global batch size.
Args:
dataset: a `tf.data.Dataset` object.
mesh: the DTensor mesh to place the dataset batches on.
layouts: a structure of DTensor layouts to be applied to the input dataset
values. This can be a single layout or (possibly nested) tuples or
dictionaries of layouts, and the structure must match the structure of
the dataset. Either all or none of the layouts should be sharded on the
batch dimension; having only a subset of layouts batch sharded will not
work and raises a ValueError.
global_batch_size: the desired global batch size.
dataset_already_batched: must be set only if the dataset is already
batched to the per-replica batch size. The batched dataset must have
`drop_remainder=True` set since DTensor requires static shapes for
slicing the input tensors.
batch_dim: the mesh dimension on which the input's batch dimension is
sharded. Set to None if the input layouts do not shard on the batch
dimension.
prefetch: number of batches to prefetch using Dataset.prefetch.
tf_data_service_config: if operating in multi-client mode, this config
specifies the tf.data service configuration to use.
Raises:
ValueError: on any of the following situations,
1. if the structures and ranks of layouts and the dataset do not match.
2. if the shapes in the dataset's spec are not fully defined.
3. if batch_dim is specified and all layouts are not batch-sharded.
4. if per_replica_batch_size is specified for an already batched Dataset
but it does not match the expected per-replica size based on the
provided mesh.
TypeError: if type of structures of layouts and the dataset do not match.
"""
super().__init__(dataset, dataset_ops.to_variant(dataset))
# TODO(b/271162918): fix multi-client use case.
if tf_data_service_config is not None:
raise NotImplementedError(
'Multi-client DTensorDataset is currently not supported.'
' Check b/271162918.')
self._mesh = mesh
self._layouts = layouts
self._batch_dim = batch_dim
self._prefetch = prefetch
self._tf_data_service_config = tf_data_service_config
nest.assert_same_structure(dataset.element_spec, layouts)
flattened_layouts = nest.flatten(layouts)
flattened_elem_spec = nest.flatten(dataset.element_spec)
if batch_dim:
self.num_global_replicas = mesh.dim_size(batch_dim)
self._local_replica_ids = list(
dict.fromkeys(
[loc[batch_dim] for loc in mesh.local_device_locations()]))
for layout in flattened_layouts:
if batch_dim != layout.sharding_specs[0]:
raise ValueError(
('batch_dim %s was specified but at least one layout did not '
'contain it: %s') % (batch_dim, layout))
else:
# Only one replica since there is no sharding on the batch dimension.
self.num_global_replicas = 1
self._local_replica_ids = [0]
# Validate layout and element spec compatibility, and raise ValueError if
# invalid.
_validate_input(
flattened_layouts,
flattened_elem_spec,
dataset_already_batched=dataset_already_batched)
expected_batch_size = global_batch_size // self.num_global_replicas
if not dataset_already_batched:
self._batched_dataset = dataset.batch(
expected_batch_size, drop_remainder=True)
else:
per_replica_batch_size = flattened_elem_spec[0].shape.as_list()[0]
if per_replica_batch_size != expected_batch_size:
raise ValueError(
('per_replica_batch_size does not matched expected size based on '
'the mesh, got %d but expected %d.') %
(per_replica_batch_size, expected_batch_size))
self._batched_dataset = dataset
# Construct a global element spec of the dataset.
flattened_global_elem_spec = []
batch_tensor_shape = tensor_shape.as_shape([global_batch_size])
for elem_spec in nest.flatten(self._batched_dataset.element_spec):
new_elem_spec = tensor_spec.TensorSpec(
shape=operator.concat(batch_tensor_shape, elem_spec.shape[1:]),
dtype=elem_spec.dtype,
name=elem_spec.name)
flattened_global_elem_spec.append(new_elem_spec)
self._global_element_spec = nest.pack_sequence_as(
dataset.element_spec, flattened_global_elem_spec)
num_global_devices_per_replica = config.num_global_devices(
mesh.device_type()) // self.num_global_replicas
self._num_local_replicas = len(self._local_replica_ids)
self._num_local_devices_per_replica = mesh.num_local_devices(
) // self._num_local_replicas
# The number of clients each replica is split over.
self._num_clients_per_replica = (
num_global_devices_per_replica // self._num_local_devices_per_replica)
# In the case where a replica is split across multiple clients, an offset
# needs to be added to the index used by the partitioning logic such that
# the local devices on that client can be correctly matched to slices of the
# input tensor(s). If replicas are wholly contained within a client, then
# this offset is always 0.
self._partition_offset = (config.client_id() % self._num_clients_per_replica
) * self._num_local_devices_per_replica
# Helper data structures used in partitioning the dataset tensors.
self._all_shard_counts = [
_shard_counts(layout, batch_dim) for layout in flattened_layouts
]
self._index_matrices = [
_index_matrix(layout, elem_spec)
for layout, elem_spec in zip(flattened_layouts, flattened_elem_spec)
]
def __iter__(self):
datasets: List[Tuple[int, data_types.DatasetV2]] = []
# Start with the batched the dataset.
local_dataset = self._batched_dataset
if self._batch_dim is not None:
if self._num_clients_per_replica > 1:
# If a replica is split over multiple clients then each batch needs to
# be repeated before distribution as many times as there are clients
# corresponding to that replica.
local_dataset = self._repeat_batch(local_dataset,
self._num_clients_per_replica)
sharding_policy = data_service_ops.ShardingPolicy.DATA
else:
# Replicas are unique to each client, so FILE based sharding can be used
# which is more performant since each worker does not need to read the
# entire dataset.
sharding_policy = data_service_ops.ShardingPolicy.FILE
else:
# No batch dimension sharding specified so disable dataset sharding during
# the distribute step.
sharding_policy = data_service_ops.ShardingPolicy.OFF
# Apply distribution here (if specified) so all remaining transformations
# are executed locally.
if self._tf_data_service_config is not None:
local_dataset = local_dataset.apply(
data_service_ops.distribute(
processing_mode=sharding_policy,
service=self._tf_data_service_config.dispatcher_address,
job_name=f'{self._tf_data_service_config.job_name}_{config.client_id()}',
target_workers='LOCAL'))
for local_replica_idx, replica_id in enumerate(self._local_replica_ids):
# Select the shard for the corresponding replica.
dataset = distribute._AutoShardDataset(
local_dataset,
num_workers=self._num_local_replicas,
index=local_replica_idx,
num_replicas=self.num_global_replicas)
# Repeat each batch for each local device in the replica.
dataset = self._repeat_batch(dataset, self._num_local_devices_per_replica)
# Slice each shard further for all non-batch dim shards. If there is no
# non-batch dim sharding, this slice is essentially a no-op.
dataset = self._partition(dataset)
# Apply prefetch as the last step. Since each batch is repeated, the
# number of elements to prefetch has to be scaled by the same size.
if self._prefetch is not None:
dataset = dataset.prefetch(
self._prefetch * self._num_local_devices_per_replica)
datasets.append((replica_id, dataset))
# Convert the datasets into iterators placed on the host.
d_iterator_resource = _pack_iterator_resource_dtensor(
datasets=datasets,
layouts=self._layouts,
mesh=self._mesh,
num_local_devices_per_replica=self._num_local_devices_per_replica)
return _DTensorIterator(
dtensor_components=(d_iterator_resource,),
global_element_spec=self._global_element_spec,
layouts=self._layouts)
def _repeat_batch(self, dataset, repeats):
if repeats == 1:
# Remove this shortcut if tf.data can optimize this away.
return dataset
def repeat(*x):
return dataset_ops.DatasetV2.from_tensors(x).repeat(repeats)
return dataset.flat_map(repeat)
def _partition(self, dataset):
"""Slices each dataset element on any sharded non-batch dimension."""
if self._num_local_devices_per_replica == 1 and self._partition_offset == 0:
# Remove this shortcut if tf.data can optimize this away.
return dataset
# TODO(b/223275517): decouple from self and make testable.
def slice_batch(index, batch):
flattened_batch = nest.flatten(batch)
flattened_output = []
norm_index = math_ops.cast(
index % self._num_local_devices_per_replica, dtype=dtypes.int32)
norm_index += self._partition_offset
coords = self._mesh.coords(norm_index)
coords = array_ops.reshape(coords, (1, -1))
for element, shard_counts, idx_matrix in zip(flattened_batch,
self._all_shard_counts,
self._index_matrices):
indexes = math_ops.matmul(coords, idx_matrix)
start = array_ops.reshape(indexes, (-1,))
size = array_ops.shape_v2(
element, out_type=dtypes.int32) // shard_counts
flattened_output.append(
array_ops.slice(element, begin=start, size=size))
return nest.pack_sequence_as(batch, flattened_output)
enumerated_dataset = dataset.enumerate()
partitioned_dataset = enumerated_dataset.map(slice_batch)
return partitioned_dataset
@property
def element_spec(self):
return self._global_element_spec
tensorflow/dtensor/python/layout.py
+550
View File
@@ -0,0 +1,550 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Python definitions for `Mesh` and `Layout`."""
import collections
import functools
import itertools
from typing import List, Dict, Optional, Union
import numpy as np
from tensorflow.dtensor.proto import layout_pb2
from tensorflow.python import _pywrap_dtensor_device
from tensorflow.python.framework import device as tf_device
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.util.tf_export import tf_export
# UNSHARDED indicates a tensor dimension is not sharded over any mesh dimension.
UNSHARDED = 'unsharded'
MATCH = 'match'
USE_XLA_SPMD = False
tf_export(
'experimental.dtensor.UNSHARDED',
v1=[]).export_constant(__name__, 'UNSHARDED')
tf_export(
'experimental.dtensor.MATCH', v1=[]).export_constant(__name__, 'MATCH')
MeshDimension = collections.namedtuple('MeshDimension', ['name', 'size'])
def _compute_mesh_strides(shape: List[int]) -> List[int]:
strides = [1]
for idx, dim_size in enumerate(reversed(shape[1:])):
strides.append(strides[idx] * dim_size)
strides.reverse()
return strides
@tf_export('experimental.dtensor.Mesh', v1=[])
class Mesh(_pywrap_dtensor_device.Mesh):
"""Represents a Mesh configuration over a certain list of Mesh Dimensions.
A mesh consists of named dimensions with sizes, which describe how a set of
devices are arranged. Defining tensor layouts in terms of mesh dimensions
allows us to efficiently determine the communication required when computing
an operation with tensors of different layouts.
A mesh provides information not only about the placement of the tensors but
also the topology of the underlying devices. For example, we can group 8 TPUs
as a 1-D array for data parallelism or a `2x4` grid for (2-way) data
parallelism and (4-way) model parallelism.
Refer to [DTensor Concepts](https://www.tensorflow.org/guide/dtensor_overview)
for in depth discussion and examples.
Note: the utilities `dtensor.create_mesh` and
`dtensor.create_distributed_mesh` provide a simpler API to create meshes for
single- or multi-client use cases.
"""
def __init__(
self,
dim_names: List[str],
global_device_ids: np.ndarray,
local_device_ids: List[int],
local_devices: List[Union[tf_device.DeviceSpec, str]],
mesh_name: str = '',
global_devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
use_xla_spmd: bool = USE_XLA_SPMD,
):
"""Builds a Mesh.
The `dim_names` and `global_device_ids` arguments describe the dimension
names and shape for the mesh.
For example,
```python
dim_names = ('x', 'y'),
global_device_ids = [[0, 1],
[2, 3],
[4, 5]]
```
defines a 2D mesh of shape 3x2. A reduction over the 'x' dimension will
reduce across columns (0, 2, 4) and (1, 3, 5), and a reduction over the 'y'
dimension reduces across rows.
Note: the utilities `dtensor.create_mesh` and
`dtensor.create_distributed_mesh` provide a simpler API to create meshes for
single- or multi-client use cases.
Args:
dim_names: A list of strings indicating dimension names.
global_device_ids: An ndarray of global device IDs is used to compose
DeviceSpecs describing the mesh. The shape of this array determines the
size of each mesh dimension. Values in this array should increment
sequentially from 0. This argument is the same for every DTensor client.
local_device_ids: A list of local device IDs equal to a subset of values
in global_device_ids. They indicate the position of local devices in the
global mesh. Different DTensor clients must contain distinct
local_device_ids contents. All local_device_ids from all DTensor clients
must cover every element in global_device_ids.
local_devices: The list of devices hosted locally. The elements correspond
1:1 to those of local_device_ids.
mesh_name: The name of the mesh. Currently, this is rarely used, and is
mostly used to indicate whether it is a CPU, GPU, or TPU-based mesh.
global_devices (optional): The list of global devices. Set when multiple
device meshes are in use.
use_xla_spmd (optional): Boolean when True, will use XLA SPMD instead of
DTensor SPMD.
"""
# Check if input args are valid.
if not isinstance(global_device_ids, np.ndarray):
raise ValueError('Variable global_device_ids must be an ndarray.')
if global_device_ids.size == 0:
raise ValueError('Variable global_device_ids must be non-empty.')
flat_global_device_ids = global_device_ids.flatten()
# global_device_ids are expected to be consecutive numbers.
# LINT.IfChange
distance = flat_global_device_ids[0]
if any(
(gid - i != distance) for i, gid in enumerate(flat_global_device_ids)):
raise ValueError('global_device_ids must sequentially increase: %s' %
global_device_ids)
# LINT.ThenChange(//tensorflow/dtensor/cc/dtensor_device.cc)
# TODO(b/242201545): This class is only for args type transformation for
# exported C++ Mesh class after the unification is complete. Any other
# logics should reside in the C++ layer, including validation checks, shall
# go to C++.
if len(dim_names) != global_device_ids.ndim:
raise ValueError(
'Number of mesh dimensions does not match number of dimension names.')
if not isinstance(local_device_ids, list):
raise ValueError('Variable local_device_ids must be a list of integers.')
if not isinstance(local_devices, list):
raise ValueError('Variable local_devices must be a list of DeviceSpecs.')
if global_devices and not isinstance(global_devices, list):
raise ValueError('Variable global_devices must be a list of DeviceSpecs.')
if not local_devices and not global_devices:
raise ValueError('Empty list of devices not allowed.')
# Transform args format for C++ Mesh constructor
global_device_ids_flatten = global_device_ids.flatten()
global_device_ids_shape = global_device_ids.shape
def to_str(d) -> str:
if isinstance(d, tf_device.DeviceSpec):
return d.to_string()
return d
def to_spec(d) -> tf_device.DeviceSpec:
if not isinstance(d, tf_device.DeviceSpec):
return tf_device.DeviceSpec.from_string(d)
return d
local_devices_str = [to_str(d) for d in local_devices]
local_devices_spec = [to_spec(d) for d in local_devices]
if not global_devices:
global_devices = []
global_devices_str = [to_str(d) for d in global_devices]
global_devices_spec = [to_spec(d) for d in global_devices]
local_devices_set = set(local_devices_spec)
local_device_only_contains_host_cpu = (
len(local_devices_set) == 1 and
list(local_devices_set)[0].device_type == 'CPU')
if not local_device_only_contains_host_cpu and len(local_devices) != len(
local_devices_set):
raise ValueError('Duplicate devices found in mesh specification %s.' %
[d for d in local_devices if local_devices.count(d) > 1])
if len(local_device_ids) != len(local_devices):
raise ValueError(
'Variable local_device_ids does not have same size as local_devices.')
if len(local_device_ids) > len(np.ravel(global_device_ids)):
raise ValueError('Cannot have more local than gobal device IDs.')
device_types = set([device.device_type for device in local_devices_spec])
if not device_types:
device_types = set([device.device_type for device in global_devices_spec])
if None in device_types:
raise ValueError('device_type is required')
if len(device_types) > 1:
raise ValueError('Devices containing multiple device_types : %s' %
device_types)
device_type = device_types.pop()
if use_xla_spmd and device_type != 'TPU':
raise ValueError('XLA SPMD is not currently not supported for %s mesh.' %
device_type)
super().__init__(
mesh_name,
dim_names,
global_device_ids_shape,
global_device_ids_flatten,
global_devices_str,
local_device_ids,
local_devices_str,
use_xla_spmd,
)
@classmethod
def _new_object(cls, *args, **kwargs):
# Need to explicitly invoke the base class __init__ because
# Mesh.__init__ overrode it with a different signature.
self = _pywrap_dtensor_device.Mesh.__new__(cls)
super().__init__(self, *args, **kwargs)
return self
def global_device_ids(self) -> np.ndarray:
"""Returns a global device list as an array."""
return np.array(super().global_device_ids(), dtype=np.int64).reshape(
self.shape()
)
def __getitem__(self, dim_name: str) -> MeshDimension:
return MeshDimension(name=dim_name, size=self.dim_size(dim_name))
def __hash__(self):
return hash(self.as_proto().SerializeToString(deterministic=True))
def __repr__(self) -> str:
return f'Mesh.from_string({self.to_string()})'
# TODO(panzf): change to pybind11 pickle implementation in the last step
def __reduce__(self):
return Mesh.from_string, (self.to_string(),)
# TODO(b/242201545): implement this in Mesh C++ class
def coords(self, device_idx: int) -> tensor.Tensor:
"""Converts the device index into a tensor of mesh coordinates."""
strides = ops.convert_to_tensor(self.strides)
shape = ops.convert_to_tensor(self.shape())
return (device_idx // strides) % shape
@classmethod
def from_proto(cls, proto: layout_pb2.MeshProto) -> 'Mesh':
"""Construct a mesh instance from input `proto`."""
return cls._new_object(mesh_proto=proto)
@classmethod
def from_string(cls, mesh_str: str) -> 'Mesh':
return cls._new_object(mesh_str=mesh_str)
@classmethod
def from_device(cls, device: str) -> 'Mesh':
"""Constructs a single device mesh from a device string."""
return cls._new_object(single_device=device)
@classmethod
def _from_mesh(cls, mesh: _pywrap_dtensor_device.Mesh):
"""Creates a copy from an existing pywrap mesh object."""
return cls._new_object(mesh=mesh)
@functools.cached_property
def _host_mesh(self) -> 'Mesh':
return Mesh._from_mesh(super().host_mesh())
def host_mesh(self) -> 'Mesh':
"""Returns a host mesh."""
# TODO(b/242201545): Find a way to get the super class to return correct
# typed objects.
return self._host_mesh
# TODO(b/242201545): implement this in Mesh C++ class
def local_device_locations(self) -> List[Dict[str, int]]:
"""Returns a list of local device locations.
A device location is a dictionary from dimension names to indices on those
dimensions.
"""
mapping = self.unravel_index()
return [mapping[device_id] for device_id in self.local_device_ids()]
# TODO(b/242201545): implement this in Mesh C++ class
@property
def strides(self) -> List[int]:
"""Returns the strides tensor array for this mesh.
If the mesh shape is `[a, b, c, d]`, then the strides array can be computed
as `[b*c*d, c*d, d, 1]`. This array can be useful in computing local device
offsets given a device ID. Using the same example, the device coordinates of
the mesh can be computed as:
```
[(device_id / (b*c*d)) % a,
(device_id / (c*d)) % b,
(device_id / (d)) % c,
(device_id) % d]
```
This is the same as `(device_id // mesh.strides) % mesh.shape`.
Returns:
The mesh strides as an integer tensor.
"""
return _compute_mesh_strides(self.shape())
# TODO(b/242201545): implement this in Mesh C++ class
def unravel_index(self):
"""Returns a dictionary from device ID to {dim_name: dim_index}.
For example, for a 3x2 mesh, return this:
```
{ 0: {'x': 0, 'y', 0},
1: {'x': 0, 'y', 1},
2: {'x': 1, 'y', 0},
3: {'x': 1, 'y', 1},
4: {'x': 2, 'y', 0},
5: {'x': 2, 'y', 1} }
```
"""
idx_ranges = [range(self.dim_size(dim_name)) for dim_name in self.dim_names]
mesh_pos = itertools.product(*idx_ranges)
mapping = {}
for device_id, device_pos in enumerate(mesh_pos):
device_loc = {}
for dim_name, dim_index in zip(self.dim_names, device_pos):
device_loc[dim_name] = dim_index
mapping[device_id] = device_loc
return mapping
LayoutType = _pywrap_dtensor_device.LayoutType
# TODO(hthu): Consider making this class immutable.
@tf_export('experimental.dtensor.Layout', v1=[])
class Layout(_pywrap_dtensor_device.Layout):
"""Represents the layout information of a DTensor.
A layout describes how a distributed tensor is partitioned across a mesh (and
thus across devices). For each axis of the tensor, the corresponding
sharding spec indicates which dimension of the mesh it is sharded over. A
special sharding spec `UNSHARDED` indicates that axis is replicated on
all the devices of that mesh.
Refer to [DTensor Concepts](https://www.tensorflow.org/guide/dtensor_overview)
for in depth discussion and examples.
For example, let's consider a 1-D mesh:
```
Mesh(["TPU:0", "TPU:1", "TPU:2", "TPU:3", "TPU:4", "TPU:5"], [("x", 6)])
```
This mesh arranges 6 TPU devices into a 1-D array. `Layout([UNSHARDED], mesh)`
is a layout for rank-1 tensor which is replicated on the 6 devices.
For another example, let's consider a 2-D mesh:
```
Mesh(["TPU:0", "TPU:1", "TPU:2", "TPU:3", "TPU:4", "TPU:5"],
[("x", 3), ("y", 2)])
```
This mesh arranges 6 TPU devices into a `3x2` 2-D array.
`Layout(["x", UNSHARDED], mesh)` is a layout for rank-2 tensor whose first
axis is sharded on mesh dimension "x" and the second axis is replicated. If we
place `np.arange(6).reshape((3, 2))` using this layout, the individual
components tensors would look like:
```
Device | Component
TPU:0 [[0, 1]]
TPU:1 [[0, 1]]
TPU:2 [[2, 3]]
TPU:3 [[2, 3]]
TPU:4 [[4, 5]]
TPU:5 [[4, 5]]
```
"""
def __init__(self, sharding_specs: List[str], mesh: Mesh):
"""Builds a Layout from a list of dimension names and a Mesh.
Args:
sharding_specs: List of sharding specifications, each corresponding to a
tensor axis. Each specification (dim_sharding) can either be a mesh
dimension or the special value UNSHARDED.
mesh: A mesh configuration for the Tensor.
Returns:
A valid Layout built with given layout & mesh.
"""
# Validate mesh
if not isinstance(mesh, Mesh):
raise ValueError('mesh is not a valid Mesh object.')
# Validate sharding spec
for _, dim_sharding in enumerate(sharding_specs):
# If special value no need to check for uniqueness, just skip.
if dim_sharding == UNSHARDED or dim_sharding == MATCH:
continue
# Check dim_sharding is unique.
if sharding_specs.count(dim_sharding) > 1:
raise ValueError(
('Mesh dimension {mesh_dim} was repeated in sharding ' +
'specification {sharding_specs}. Mesh dimensions must be unique ' +
'in a layout.').format(
mesh_dim=dim_sharding, sharding_specs=sharding_specs))
# Check dim_sharding is mesh dimension.
if dim_sharding not in mesh:
raise ValueError(
('{dim_sharding}: A dimension sharding must either be a ' +
'valid mesh dimension or UNSHARDED.').format(
dim_sharding=dim_sharding))
super().__init__(
type=LayoutType.STATIC, sharding_specs=sharding_specs, mesh=mesh
)
@classmethod
def _new_object(cls, *args, **kwargs):
# Need to explicitly invoke the base class __init__ because
# Layout.__init__ overrode it with a different signature.
self = _pywrap_dtensor_device.Layout.__new__(cls)
super().__init__(self, *args, **kwargs)
return self
def __repr__(self) -> str:
return f'Layout.from_string({self.to_string()})'
def __hash__(self):
return hash(self.as_proto().SerializeToString(deterministic=True))
# TODO(panzf): change to pybind11 pickle implementation in the last step
def __reduce__(self):
return Layout.from_string, (self.to_string(),)
@property
def mesh(self):
return Mesh._from_mesh(mesh=super().mesh) # pylint: disable=protected-access
@property
def shape(self):
return self.mesh.shape()
@classmethod
def batch_sharded(
cls, mesh: Mesh, batch_dim: str, rank: int, axis: int = 0
) -> 'Layout':
"""Returns a layout sharded on batch dimension."""
return cls._new_object(
# Watchout for the different ordering.
mesh=mesh,
rank=rank,
batch_dim=batch_dim,
axis=axis,
)
# TODO(b/242201545): Move this to C++ / find the corresponding function there.
def delete(self, dims: List[int]) -> 'Layout':
"""Returns the layout with the give dimensions deleted."""
if not isinstance(dims, list):
dims = [dims]
new_specs = [
spec for i, spec in enumerate(self.sharding_specs) if i not in dims
]
return Layout(new_specs, self.mesh)
@classmethod
def from_proto(cls, layout_proto: layout_pb2.LayoutProto) -> 'Layout':
"""Creates an instance from a LayoutProto."""
return cls._new_object(layout_proto=layout_proto)
@classmethod
def from_string(cls, layout_str: str) -> 'Layout':
"""Creates an instance from a human-readable string."""
return cls._new_object(layout_str=layout_str)
def to_parted(self) -> 'Layout':
"""Returns a "parted" layout from a static layout.
A parted layout contains axes that are treated as independent by most of
SPMD expanders.
FIXME(b/285905569): The exact semantics is still being investigated.
"""
return Layout._new_object(layout=super().to_parted())
@classmethod
def inner_sharded(cls, mesh: Mesh, inner_dim: str, rank: int) -> 'Layout':
"""Returns a layout sharded on inner dimension."""
return cls.batch_sharded(mesh, inner_dim, rank, axis=rank - 1)
@classmethod
def from_single_device_mesh(cls, mesh: Mesh) -> 'Layout':
"""Constructs a single device layout from a single device mesh."""
return cls._new_object(mesh=mesh)
@classmethod
def from_device(cls, device: str) -> 'Layout':
"""Constructs a single device layout from a single device mesh."""
return cls.from_single_device_mesh(Mesh.from_device(device))
# TODO(b/242201545): Move this to C++ / find the corresponding function there.
def offset_to_shard(self):
"""Mapping from offset in a flattened list to shard index."""
unravel_index = self.mesh.unravel_index()
locations = [None] * self.mesh.size
for offset, mesh_loc in unravel_index.items():
loc = []
for dim_sharding in self.sharding_specs:
if dim_sharding == UNSHARDED:
loc.append(0)
else:
loc.append(mesh_loc[dim_sharding])
locations[offset] = tuple(loc)
return locations
# TODO(b/242201545): Move this to C++ / find the corresponding function there.
def offset_tuple_to_global_index(self, offset_tuple):
"""Mapping from offset to index in global tensor."""
index = 0
for i, o in enumerate(offset_tuple):
m = 1
for x in range(i + 1, self.rank):
m = m * self.num_shards(x)
index = index + m * o
return index
@classmethod
def replicated(cls, mesh: Mesh, rank: int) -> 'Layout':
"""Returns a replicated layout of rank `rank`."""
return cls._new_object(mesh=mesh, rank=rank)
tensorflow/dtensor/python/mesh_util.py
+310
View File
@@ -0,0 +1,310 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities to help with mesh creation."""
from typing import Dict, List, Optional, Tuple, Union
from absl import logging
import numpy as np
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import layout
from tensorflow.dtensor.python import tpu_util
from tensorflow.python.eager import context
from tensorflow.python.framework import device as tf_device
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.util.tf_export import tf_export
def _print_context(num_global_devices: int, num_clients: int, client_id: int,
device_type: str, mesh: layout.Mesh) -> None:
logging.info('This is client %d of %d clients', client_id, num_clients)
logging.info('Number of global %s devices: %d', device_type.upper(),
num_global_devices)
# pylint: disable=protected-access
logging.info('Global device IDs: %s', mesh.global_device_ids())
logging.info('Local device IDs: %s', mesh.local_device_ids())
logging.info('Local devices: %s', mesh.local_devices())
# pylint: enable=protected-access
def _make_device_specs(
devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
device_type: Optional[str] = None
) -> Tuple[List[tf_device.DeviceSpec], str]:
"""Makes device specs for all local devices or from a provided list."""
if devices is None:
if device_type is None:
device_type = 'CPU'
devices = config.local_devices(device_type)
else:
if isinstance(devices[0], str):
devices = [tf_device.DeviceSpec.from_string(d) for d in devices]
if device_type is None:
device_type = devices[0].device_type
if device_type.upper() != devices[0].device_type.upper():
raise ValueError(
f'Conflicting devices {str(devices)} and device_type {device_type}'
)
return devices, device_type
@tf_export('experimental.dtensor.create_mesh', v1=[])
def create_mesh(
mesh_dims: Optional[Union[List[Tuple[str, int]], Dict[str, int]]] = None,
mesh_name: str = '',
devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
device_type: Optional[str] = None,
use_xla_spmd: bool = layout.USE_XLA_SPMD,
) -> layout.Mesh:
"""Creates a single-client mesh.
If both `mesh_dims` and `devices` are specified, they must match each otehr.
As a special case, when all arguments are missing, this creates a 1D CPU mesh
with an empty name, assigning all available devices to that dimension.
Args:
mesh_dims: A dict of dim_name: dim_size, or a list of (dim_name, dim_size)
tuples. Defaults to a single batch-parallel dimension called 'x' usin all
devices. As a special case, a single-element mesh_dims whose dim_size is
-1 also uses all devices. e.g. `{'x' : 4, 'y' : 1}` or `[('x', 4), ('y',
1)]`.
mesh_name: Name of the created mesh. Defaults to ''.
devices: String representations of devices to use. This is the device part
of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available logical devices.
device_type: If `devices` is missing, the type of devices to use. Defaults
to 'CPU'.
use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.
Returns:
A single-client mesh created from specified or default arguments.
"""
device_specs, device_type = _make_device_specs(devices, device_type)
local_spec = tf_device.DeviceSpec(job=config.job_name(), replica=0, task=0)
device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
if isinstance(mesh_dims, dict):
mesh_dims = list(mesh_dims.items())
if mesh_dims is None:
mesh_dims = [('x', len(device_specs))]
elif len(mesh_dims) == 1 and mesh_dims[0][1] == -1:
# Replace -1 dim_size in a 1D mesh will the number of all devices.
mesh_dims[0] = (mesh_dims[0][0], len(device_specs))
dim_names = [d[0] for d in mesh_dims]
shape = [d[1] for d in mesh_dims]
if np.prod(shape) != len(device_specs):
raise ValueError(f'length of devices ({len(device_specs)}) must be '
f'equal to total size of the mesh of shape {shape}')
global_device_ids = np.arange(len(device_specs)).reshape(shape)
local_device_ids = np.ravel(global_device_ids).tolist()
mesh = layout.Mesh(
dim_names=dim_names,
global_device_ids=global_device_ids,
local_device_ids=local_device_ids,
local_devices=device_specs,
mesh_name=mesh_name,
use_xla_spmd=use_xla_spmd)
_print_context(
num_global_devices=len(device_specs),
num_clients=1,
client_id=0,
device_type=device_type,
mesh=mesh)
return mesh
@tf_export('experimental.dtensor.create_distributed_mesh', v1=[])
def create_distributed_mesh(
mesh_dims: Union[List[Tuple[str, int]], Dict[str, int]],
mesh_name: str = '',
local_devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
device_type: Optional[str] = None,
use_xla_spmd: bool = layout.USE_XLA_SPMD,
) -> layout.Mesh:
"""Creates a distributed mesh.
This is similar to `create_mesh`, but with a different set of arguments to
create a mesh that spans evenly across a multi-client DTensor cluster.
For CPU and GPU meshes, users can choose to use fewer local devices than what
is available `local_devices`.
For TPU, only meshes that uses all TPU cores is supported by the DTensor
runtime.
Args:
mesh_dims: A dict of dim_name: dim_size, or a list of (dim_name, dim_size)
tuples. e.g. `{'x' : 4, 'y' : 1}` or `[('x', 4), ('y', 1)]`.
mesh_name: Name of the created mesh. Defaults to ''.
local_devices: String representations of devices to use. This is the device
part of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available local
logical devices.
device_type: Type of device to build the mesh for. Defaults to 'CPU'.
Supported values are 'CPU', 'GPU', 'TPU'.6
use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.
Returns:
A mesh that spans evenly across all DTensor clients in the cluster.
"""
if isinstance(mesh_dims, dict):
mesh_dims = list(mesh_dims.items())
dim_names, shape = zip(*mesh_dims)
if not accelerator_util.is_initialized():
raise ValueError('Accelerators are uninitialized, please run '
'dtensor.initialize_accelerator_system() first.')
if device_type and device_type.upper() == 'TPU':
# TODO(b/185940495): Allow multi-mesh and partial on TPU.
# TPU meshes can only be configured through environment variables that
# reflect the actual TPU topology. Do not let users specify custom args.
if local_devices is not None:
raise ValueError(
f'Do not specify devices for {device_type.upper()} meshes. '
f'Using a partial list of devices for {device_type.upper()} '
f'is not supported.')
device_specs, device_type = _make_device_specs(local_devices, device_type)
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.
local_spec = tf_device.DeviceSpec(
job=config.job_name(), replica=0, task=config.client_id())
device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
# Assumes identical number of local devices per client.
num_global_devices = len(device_specs) * config.num_clients()
if np.prod(shape) != num_global_devices:
raise ValueError(
f'Global number of devices '
f'({len(device_specs)} per client * {config.num_clients()} clients '
f'= {num_global_devices}) must be '
f'equal to total size of the mesh of shape {shape}')
global_device_ids = np.arange(num_global_devices).reshape(shape)
flattened = np.ravel(global_device_ids).tolist()
start_idx = len(device_specs) * config.client_id()
local_device_ids = flattened[start_idx:start_idx + len(device_specs)]
mesh = layout.Mesh(
dim_names=dim_names,
global_device_ids=global_device_ids,
local_device_ids=local_device_ids,
local_devices=device_specs,
mesh_name=mesh_name,
use_xla_spmd=use_xla_spmd)
_print_context(num_global_devices, config.num_clients(), config.client_id(),
device_type, mesh)
return mesh
if device_type.upper() == 'TPU':
mesh = tpu_util.create_tpu_mesh(
mesh_dim_names=dim_names,
mesh_shape=shape,
mesh_name=mesh_name,
use_xla_spmd=use_xla_spmd)
_print_context(
config.num_global_devices(device_type), config.num_clients(),
config.client_id(), device_type, mesh)
return mesh
raise ValueError(f'Device type {device_type} is not CPU, GPU or TPU')
_BARRIER_DICT = {}
@tf_export('experimental.dtensor.barrier', v1=[])
def barrier(mesh: layout.Mesh,
barrier_name: Optional[str] = None,
timeout_in_ms: Optional[int] = None):
"""Runs a barrier on the mesh.
Upon returning from the barrier, all operations run before the barrier
would have completed across all clients. Currently we allocate a fully
sharded tensor with mesh shape and run an all_reduce on it.
Example:
A barrier can be used before application exit to ensure completion of pending
ops.
```python
x = [1, 2, 3]
x = dtensor.relayout(x, dtensor.Layout.batch_sharded(mesh, 'batch', 1))
dtensor.barrier(mesh)
# At this point all devices on all clients in the mesh have completed
# operations before the barrier. Therefore it is OK to tear down the clients.
sys.exit()
```
Args:
mesh: The mesh to run the barrier on.
barrier_name: The name of the barrier. Mainly used for logging purpose.
timeout_in_ms: The timeout of the barrier in ms. If omitted, blocks
indefinitely till the barrier is reached from all clients.
"""
if barrier_name is None:
barrier_name = '(barrier)'
logging.info('entering barrier before op: %s', barrier_name)
# Make sure all ops are consumed before running the sync.
context.async_wait()
# Reduction on a fully sharded tensor requires all devices to participate
# and serves as a barrier on the mesh.
component = array_ops.reshape(1.0, [1] * len(mesh.shape()))
ones = api.pack([component] * mesh.num_local_devices(),
layout.Layout(mesh.dim_names, mesh))
mesh_size = math_ops.reduce_sum(ones)
if mesh_size != mesh.size:
raise ValueError(
'Global barrier produced wrong mesh size : {0} while mesh has actual'
'size : {1}'.format(mesh_size, mesh.size))
# TODO(hthu): This isn't strictly needed but might cause confusing behaviors
# from users. Consider dropping this if there is a `big` performance hit.
context.async_wait()
if context.context().coordination_service:
if timeout_in_ms is None:
timeout_in_ms = 24 * 60 * 60 * 1000 # 24 hours to stand in for infinite.
num_calls = _BARRIER_DICT.setdefault(barrier_name, 0)
_BARRIER_DICT[barrier_name] = num_calls + 1
barrier_id = f'{barrier_name}:{num_calls}'
context.context().wait_at_barrier(barrier_id, timeout_in_ms)
logging.info('finished running barrier across all clients after '
'op: %s', barrier_name)
tensorflow/dtensor/python/numpy_util.py
+130
View File
@@ -0,0 +1,130 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities to convert data buffers to/from DTensor tensors."""
from typing import List
import numpy as np
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.types.core import Tensor, TensorLike # pylint: disable=g-multiple-import
# FIXME(b/262894693): Functions in this file are buggy.
# They do not distinguish between the client-local data and the global view.
def _split(value, splits, axis=0, split_fn=np.split, stack_fn=np.stack):
"""Split `value` into a sharded nparray/tf tensor based on the number of splits.
"""
# During graph tracing a dimension can be dynamic (`None`), so only run the
# divisibility check when the size is statically known. A raw modulo on a
# `None` dimension would otherwise raise a confusing TypeError.
dim_size = value.shape[axis]
if hasattr(dim_size, "value"):
dim_size = dim_size.value
if dim_size is not None and dim_size % splits[0] != 0:
raise ValueError(
f"Tensor shape along dimension {axis} ({dim_size}) is not evenly "
f"divisible by the number of splits ({splits[0]}) for that dimension."
)
children = split_fn(value, splits[0], axis=axis)
if len(splits) > 1:
splits = splits[1:]
children = [_split(child, splits, axis + 1) for child in children]
return stack_fn(children)
def to_numpy(tensor: TensorLike) -> np.ndarray:
"""Copy `input` DTensor to an equivalent local numpy array."""
layout = api.fetch_layout(tensor)
if layout.mesh.is_remote():
return np.array([None])
unpacked = [tensor.numpy() for tensor in api.unpack(tensor)]
return unpacked_to_numpy(unpacked, layout)
def unpacked_to_numpy(unpacked: List[TensorLike],
layout: layout_lib.Layout) -> np.ndarray:
"""Heals local Tensor components to a numpy array."""
if len(unpacked) != len(layout.offset_to_shard()):
raise ValueError('Wrong number of component Tensors.')
unravelled = np.ndarray([layout.num_shards(i) for i in range(layout.rank)],
dtype=object)
for offset, loc in enumerate(layout.offset_to_shard()):
unravelled[loc] = unpacked[offset]
concat_tensor = np.block(unravelled.tolist())
# np.block can introduce empty initial dimensions, peel these off until
# the output matches the rank of the input tensors.
while concat_tensor.ndim > unpacked[0].ndim:
concat_tensor = np.squeeze(concat_tensor, axis=0)
return concat_tensor
# TODO(feyu): rename to slice.
def unpack(t: TensorLike,
layout: layout_lib.Layout,
split_fn=np.split,
stack_fn=np.stack) -> List[TensorLike]:
"""Slice `t` into a flattened list of tensors suitable for `pack`."""
if not layout.rank:
return [t] * layout.mesh.size
sharded_tensor = _split(
t, [layout.num_shards(i) for i in range(layout.rank)],
split_fn=split_fn,
stack_fn=stack_fn)
flattened = [np.ndarray([])] * layout.mesh.size
for offset, shard in enumerate(layout.offset_to_shard()):
flattened[offset] = sharded_tensor[tuple(shard)]
return flattened
def pack_numpy(value: np.ndarray,
layout: layout_lib.Layout,
make_sparse: bool = False) -> Tensor:
assert value is not None
unpacked = unpack(value, layout)
if make_sparse:
return api.pack([sparse_ops.from_dense(t) for t in unpacked], layout)
return api.pack(unpacked, layout)
def pack_tf_tensor(value: Tensor, layout: layout_lib.Layout) -> Tensor:
if value is None:
raise ValueError('pack requires values to be passed in')
unpacked = unpack(
value, layout, split_fn=array_ops.split, stack_fn=array_ops_stack.stack)
return api.pack(unpacked, layout)
@polymorphic_function.function
def stateless_random_uniform(shape, seed, layout):
"""Creates uniform random tensor with the given layout."""
return api.relayout(
stateless_random_ops.stateless_random_uniform(shape=shape, seed=seed),
layout=layout,
)
tensorflow/dtensor/python/save_restore.py
+222
View File
@@ -0,0 +1,222 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Contains functionaility for Checkpoint/SavedModel in DTensor."""
import collections
from typing import Dict, List, Union
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import gen_dtensor_ops
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import mesh_util
from tensorflow.python.eager import context
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import variables as tf_variables
from tensorflow.python.util.tf_export import tf_export
@tf_export('experimental.dtensor.sharded_save', v1=[])
def sharded_save(
mesh: layout_lib.Mesh,
file_prefix: Union[str, tensor_lib.Tensor],
tensor_names: Union[List[str], tensor_lib.Tensor],
shape_and_slices: Union[List[str], tensor_lib.Tensor],
tensors: List[Union[tensor_lib.Tensor, tf_variables.Variable]],
):
"""Saves given named tensor slices in a sharded, multi-client safe fashion.
The method makes sure the checkpoint directory state is correct in a sharded
mutli-client saving. Namely, we place a barrier after SaveV2 to make sure
every client has done writing the files. And another one after
MergeV2Checkpoints to make sure all Metadata is properly merged.
Upon existing, the checkpoint is completed and the all directory operations
are done.
Args:
mesh: The Mesh that contains the Tensors to save.
file_prefix: The prefix of checkpoint.
tensor_names: a list of tensor names used in save op.
shape_and_slices: a list of shape and slice specification used in save op.
The only supported value is "" as we don't support distributed saving with
slices yet.
tensors: a list of tensors used in save op. The order should match
tensor_names.
Returns:
A MergeV2Checkpoints op that merged all Metadata.
"""
with ops.device(api.device_name()):
io_ops.save_v2(file_prefix, tensor_names, shape_and_slices, tensors)
# Make sure all clients have written the files
mesh_util.barrier(mesh.host_mesh(), 'SaveV2') # pylint: disable=protected-access
with api.default_mesh(mesh.host_mesh()):
merge_op = io_ops.MergeV2Checkpoints(
checkpoint_prefixes=[file_prefix],
destination_prefix=file_prefix,
delete_old_dirs=True)
# Make sure first device in first host has finished merge.
mesh_util.barrier(mesh.host_mesh(), 'MergeV2Checkpoints')
return merge_op
@tf_export('experimental.dtensor.enable_save_as_bf16', v1=[])
def enable_save_as_bf16(variables: List[tf_variables.Variable]):
"""Allows float32 DVariables to be checkpointed and restored as bfloat16.
The method only affects the DVariable part inside the model and leaves
non-DTensor Variables/Tensors untouched.
Args:
variables: A list of tf.Variable to be enabled with bfloat16 save/restore.
Only has effect on DTensor Variables as they go through d_variables with
DTensor Specific logis.
"""
for v in variables:
if isinstance(v, d_variable.DVariable):
v.save_as_bf16 = True
@tf_export('experimental.dtensor.name_based_restore', v1=[])
def name_based_restore(
mesh: layout_lib.Mesh,
checkpoint_prefix: str,
name_tensor_dict: Dict[
str, Union[tensor_lib.Tensor, tf_variables.Variable]],
):
"""Restores from checkpoint_prefix to name based DTensors.
It is required to have already-initialized DTensor variables that have same
shape/dtype for the tensors being restored.
Also, we currently only support a named based restore on a single mesh.
Args:
mesh: The single mesh that all Tensors would be restored to.
checkpoint_prefix : The prefix of checkpoint to be restored.
name_tensor_dict: A ordered dictionary of tensor_names to a DTensor. The
DTensor shape/dtype must match the tensors being saved/restored for now.
Returns:
A dictionary of name to its restored DTensor value.
"""
if not context.executing_eagerly():
raise ValueError('name based restore must run eagerly.')
ordered_name_tensor_dict = name_tensor_dict
if not isinstance(name_tensor_dict, collections.OrderedDict):
ordered_name_tensor_dict = collections.OrderedDict(name_tensor_dict)
# Make sure that all tensors are on CPU mesh for now.
# This might not be a hard limitation in the future.
for name, tensor in ordered_name_tensor_dict.items():
try:
if api.fetch_layout(tensor).mesh.device_type().upper() != 'CPU':
raise ValueError(
'Restoring a non CPU Tensor is not supported currently. Offending '
'tensor name : {tensor_name}'.format(tensor_name=name))
except errors_impl.OpError as op_error:
raise ValueError(
'Saving/Restoring tensor must be a DTensor') from op_error
# Now that we have all tensors on CPU mesh, do a DTensorRestoreV2.
checkpoint_prefix = api.pack(
[checkpoint_prefix] * mesh.num_local_devices(),
layout_lib.Layout.replicated(mesh.host_mesh(), rank=0))
# Explicitly pack to mesh to avoid implicit small constant extraction, which
# does not work larger restores that has lots of names.
tensor_names = api.pack(
[list(ordered_name_tensor_dict.keys())] * mesh.num_local_devices(),
layout_lib.Layout.replicated(mesh.host_mesh(), rank=1))
shape_and_slices = api.pack(
[[''] * len(ordered_name_tensor_dict)] * mesh.num_local_devices(),
layout_lib.Layout.replicated(mesh.host_mesh(), rank=1))
# A list of TensorShape representing all shapes for the input tensors.
input_shapes = [tensor.shape for tensor in ordered_name_tensor_dict.values()]
input_layouts = [
api.fetch_layout(tensor).to_string()
for tensor in ordered_name_tensor_dict.values()
]
with ops.device(api.device_name()):
restored_cpu_tensors = gen_dtensor_ops.d_tensor_restore_v2(
prefix=checkpoint_prefix,
tensor_names=tensor_names,
shape_and_slices=shape_and_slices,
input_shapes=input_shapes,
input_layouts=input_layouts,
dtypes=[tensor.dtype for tensor in ordered_name_tensor_dict.values()],
)
return collections.OrderedDict(
zip(ordered_name_tensor_dict.keys(), restored_cpu_tensors)
)
@tf_export('experimental.dtensor.name_based_save', v1=[])
def name_based_save(
mesh: layout_lib.Mesh,
checkpoint_prefix: Union[str, tensor_lib.Tensor],
name_tensor_dict: Dict[
str, Union[tensor_lib.Tensor, tf_variables.Variable]],
):
"""Saves name based Tensor into a Checkpoint.
The function prepares the input dictionary to the format of a `sharded_save`,
so that it can take advantage of DTensor SPMD based distributed save.
Same as restore, the function only supports saving on the single mesh.
Args:
mesh: The single mesh that all Tensors would be restored to.
checkpoint_prefix : The prefix of checkpoint to be restored.
name_tensor_dict: A ordered dictionary of tensor_names to a DTensor. The
DTensor shape/dtype must match the tensors being saved/restored for now.
"""
if not context.executing_eagerly():
raise ValueError('name based save must run eagerly.')
ordered_name_tensor_dict = name_tensor_dict
if not isinstance(name_tensor_dict, collections.OrderedDict):
ordered_name_tensor_dict = collections.OrderedDict(name_tensor_dict)
# Current _dtensor_device() in api.py is the correct way of specifying
# DTensor device singletons. The API itself will be eventually be moved to
# a public API and provides global singleton in DTensor context.
# For now, we just use the current `internal` API and aim at migrating in
# one shot later.
# TODO(hthu): Provide _dtensor_device() singleton as a public API.
# pylint: disable=protected-access
checkpoint_prefix = api.pack([checkpoint_prefix] * mesh.num_local_devices(),
layout_lib.Layout.replicated(
mesh.host_mesh(), rank=0))
tensor_names = api.pack(
[list(ordered_name_tensor_dict.keys())] * mesh.num_local_devices(),
layout_lib.Layout.replicated(mesh.host_mesh(), rank=1))
sharded_save(
mesh,
file_prefix=checkpoint_prefix,
tensor_names=tensor_names,
shape_and_slices=[''] * len(ordered_name_tensor_dict),
tensors=list(ordered_name_tensor_dict.values()))
tensorflow/dtensor/python/tests/BUILD
+1041
View File
File diff suppressed because it is too large Load Diff
tensorflow/dtensor/python/tests/api_test.py
+305
View File
@@ -0,0 +1,305 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for the internal DTensor Python API."""
from absl.testing import parameterized
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_random
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.platform import test
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
class APITest(test_util.DTensorBaseTest):
def setUp(self):
super(APITest, self).setUp()
global_ids = test_util.create_device_ids_array((2, 2))
local_device_ids = np.ravel(global_ids).tolist()
mesh_dict = {
'CPU': Mesh(
[_MESH_DIM_X, _MESH_DIM_Y],
global_ids,
local_device_ids,
test_util.create_device_list((2, 2), 'CPU'),
)
}
self.mesh = self.configTestMesh(mesh_dict)
self.layouts_1d = [
Layout.replicated(self.mesh, rank=1),
Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=1),
Layout.batch_sharded(self.mesh, _MESH_DIM_Y, rank=1),
]
self.layouts_2d = [
Layout.replicated(self.mesh, rank=2),
Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=2),
Layout.inner_sharded(self.mesh, _MESH_DIM_X, rank=2),
Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh),
]
def testV2API(self):
layout = Layout.replicated(self.mesh, rank=1)
zero_tensor = array_ops.zeros([10], layout=layout)
zero_like_tensor = array_ops.zeros_like_v2(zero_tensor, layout=layout)
self.assertAllEqual(zero_like_tensor.numpy(), zero_tensor.numpy())
ones_tensor = array_ops.ones([10], layout=layout)
ones_like_tensor = array_ops.ones_like_v2(zero_tensor, layout=layout)
self.assertAllEqual(ones_like_tensor.numpy(), ones_tensor.numpy())
def testStatelessRandom(self):
# test dtype default float32 random
result = stateless_random_ops.stateless_random_uniform(
[10],
seed=constant_op.constant([0, 0], dtype=dtypes.int64),
minval=0.0,
maxval=10.0,
)
self.assertEqual([10], result.shape)
# test dtype default int32 minval maxval are both None
result = stateless_random_ops.stateless_random_uniform(
[10],
seed=constant_op.constant([1, 2], dtype=dtypes.int64),
dtype=dtypes.int32,
minval=None,
maxval=None,
)
self.assertEqual([10], result.shape)
# test maxval is None or not given
result = stateless_random_ops.stateless_random_uniform(
[10],
seed=constant_op.constant([1, 2], dtype=dtypes.int64),
maxval=12,
dtype=dtypes.int32,
)
self.assertEqual([10], result.shape)
self.assertAllInRange(result, 0, 12)
def testStatelessRandomNormal(self):
# test dtype default float32 random
result = stateless_random_ops.stateless_random_normal(
[10], seed=constant_op.constant([0, 0], dtype=dtypes.int32)
)
self.assertEqual([10], result.shape)
# test dtype double
result = stateless_random_ops.stateless_random_normal(
[10],
seed=constant_op.constant([1, 2], dtype=dtypes.int32),
dtype=dtypes.double,
)
self.assertEqual([10], result.shape)
# test mean and stddev
result = stateless_random_ops.stateless_random_normal(
[10],
seed=constant_op.constant([1, 2], dtype=dtypes.int32),
mean=0,
stddev=0,
)
self.assertEqual([10], result.shape)
self.assertAllInRange(result, 0, 0)
# test dtensor version of each, check layouts
layout = Layout.replicated(self.mesh, rank=1)
# test dtype default float 32 random
result = d_random.stateless_random_normal(
[10],
seed=constant_op.constant([0, 0], dtype=dtypes.int32),
layout=layout,
)
self.assertEqual([10], result.shape)
self.assertEqual(layout, api.fetch_layout(result))
# test dtype double
result = d_random.stateless_random_normal(
[10],
seed=constant_op.constant([1, 2], dtype=dtypes.int32),
dtype=dtypes.double,
layout=layout,
)
self.assertEqual([10], result.shape)
self.assertEqual(layout, api.fetch_layout(result))
# test mean and stddev
result = d_random.stateless_random_normal(
[10],
seed=constant_op.constant([1, 2], dtype=dtypes.int32),
mean=0,
stddev=0,
layout=layout,
)
self.assertEqual([10], result.shape)
self.assertAllInRange(result, 0, 0)
self.assertEqual(layout, api.fetch_layout(result))
@parameterized.named_parameters(*set(
test_util.product((('_labels_unsharded', 0), ('_labels_batch', 1),
('_labels_inner', 2), ('_labels_both', 3)),
(('_logits_unsharded', 0), ('_logits_batch', 1),
('_logits_inner', 2), ('_logits_both', 3)))))
def testSoftmaxCrossentropyWithLogits(self, labels_layout, logits_layout):
expected_layout = Layout.replicated(self.mesh, rank=1)
if (labels_layout == 1 or labels_layout == 3 or logits_layout == 1 or
logits_layout == 3):
expected_layout = Layout.inner_sharded(self.mesh, _MESH_DIM_X, rank=1)
labels_layout = self.layouts_2d[labels_layout]
logits_layout = self.layouts_2d[logits_layout]
labels_numpy = np.random.uniform(size=[6, 4])
logits_numpy = np.random.uniform(size=[6, 4])
labels = constant_op.constant(labels_numpy, dtype=dtypes.float32)
logits = constant_op.constant(logits_numpy, dtype=dtypes.float32)
# Should we test against the built in version or the patched version?
expected = nn_ops.softmax_cross_entropy_with_logits_v2(
labels=labels, logits=logits
)
labels = numpy_util.pack_numpy(labels, labels_layout)
logits = numpy_util.pack_numpy(logits, logits_layout)
dtensor_result = nn_ops.softmax_cross_entropy_with_logits_v2(
labels=labels, logits=logits
)
self.assertDTensorEqual(expected, expected_layout, dtensor_result)
@parameterized.named_parameters(*set(
test_util.product((('_labels_unsharded', 0), ('_labels_batch_x', 1),
('_labels_batch_y', 2)),
(('_logits_unsharded', 0), ('_logits_batch', 1),
('_logits_inner', 2), ('_logits_both', 3)))))
def testSparseSoftmaxCrossentropyWithLogits(self, labels_layout,
logits_layout):
expected_layout = Layout.replicated(self.mesh, rank=1)
if labels_layout == 1 or logits_layout == 1 or logits_layout == 3:
expected_layout = Layout.inner_sharded(self.mesh, _MESH_DIM_X, rank=1)
elif labels_layout == 2:
expected_layout = Layout.inner_sharded(self.mesh, _MESH_DIM_Y, rank=1)
labels_layout = self.layouts_1d[labels_layout]
logits_layout = self.layouts_2d[logits_layout]
labels_numpy = np.random.randint(size=[6], low=0, high=4)
logits_numpy = np.random.uniform(size=[6, 4])
labels = constant_op.constant(labels_numpy, dtype=dtypes.int64)
logits = constant_op.constant(logits_numpy, dtype=dtypes.float32)
# Should we test against the built in version or the patched version?
expected = nn_ops.sparse_softmax_cross_entropy_with_logits_v2(
labels=labels, logits=logits
)
labels = numpy_util.pack_numpy(labels, labels_layout)
logits = numpy_util.pack_numpy(logits, logits_layout)
dtensor_result = nn_ops.sparse_softmax_cross_entropy_with_logits_v2(
labels=labels, logits=logits
)
self.assertDTensorEqual(expected, expected_layout, dtensor_result)
def test_dropout_raises_on_none_seed(self):
with api.default_mesh(self.mesh):
with self.assertRaisesRegex(ValueError, 'seed must be specified'):
_ = d_random.dropout(
array_ops.ones([2, 2], dtype=dtypes.float32), rate=0.5, seed=None
)
def test_default_mesh(self):
@polymorphic_function.function
def func(a):
return a + 3.0
with api.default_mesh(self.mesh):
a = array_ops.zeros(shape=())
result = func(a)
self.assertEqual(result, 3.0)
self.assertEqual(api.fetch_layout(result).mesh, self.mesh)
self.assertTrue(api.fetch_layout(result).is_fully_replicated())
self.assertEqual(result.device, api.device_name())
# Also make sure it works as wrapper
@api.default_mesh(self.mesh)
def func2():
b = array_ops.ones(shape=())
return func(b)
result = func2()
self.assertEqual(result, 4.0)
self.assertEqual(api.fetch_layout(result).mesh, self.mesh)
self.assertTrue(api.fetch_layout(result).is_fully_replicated())
self.assertEqual(result.device, api.device_name())
with self.assertRaisesRegex(ValueError, 'Expect `mesh` to be `Mesh`'):
with api.default_mesh(None):
pass
def test_default_mesh_with_constant(self):
@polymorphic_function.function
def func():
return constant_op.constant([3, 4])
with api.default_mesh(self.mesh):
result = func()
self.assertAllEqual(result, [3, 4])
self.assertEqual(api.fetch_layout(result).mesh, self.mesh)
self.assertTrue(api.fetch_layout(result).is_fully_replicated())
self.assertEqual(result.device, api.device_name())
def test_error_no_default_mesh(self):
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
'No default mesh has been registered to DTensor',
):
with ops.device_v2(api.device_name()):
_ = constant_op.constant(3.0)
def test_get_default_mesh(self):
self.assertIsNone(api.get_default_mesh())
with api.default_mesh(self.mesh):
self.assertEqual(api.get_default_mesh(), self.mesh)
with api.default_mesh(self.mesh.host_mesh()):
self.assertEqual(api.get_default_mesh(), self.mesh.host_mesh())
self.assertEqual(api.get_default_mesh(), self.mesh)
self.assertIsNone(api.get_default_mesh())
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/array_ops_test.py
+180
View File
@@ -0,0 +1,180 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for DTensor support in array_ops."""
import numpy as np
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import combinations
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
Layout = layout_lib.Layout
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
_MESH_DIMS = [_MESH_DIM_X, _MESH_DIM_Y]
class ArrayOpsTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
global_ids = test_util.create_device_ids_array((2, 1))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = { # pylint: disable=g-complex-comprehension
device: layout_lib.Mesh(
_MESH_DIMS,
global_ids,
local_ids,
test_util.create_device_list((2, 1), device),
)
for device in ('CPU', 'GPU', 'TPU')
}
self.mesh = self.configTestMesh(mesh_dict)
@combinations.generate(
combinations.combine(is_graph=[False, True], size=[32, 4096])
)
def testTwoFills(self, is_graph, size):
layout_x = Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=1)
layout_y = Layout.batch_sharded(self.mesh, _MESH_DIM_Y, rank=1)
def fn():
return (
array_ops.fill([size], 0.0, layout=layout_x),
array_ops.fill([size], 0.0, layout=layout_y),
)
if is_graph:
fn = polymorphic_function.function(fn)
with api.default_mesh(self.mesh):
dtensor_x, dtensor_y = fn()
tensor = array_ops.zeros([size], layout=None)
self.assertDTensorEqual(tensor, layout_x, dtensor_x)
self.assertDTensorEqual(tensor, layout_y, dtensor_y)
@combinations.generate(
combinations.combine(
is_graph=[False, True],
size=[32, 4096],
nullary_op=[array_ops.zeros, array_ops.ones],
)
)
def testNullaryOp(self, is_graph, size, nullary_op):
layout_y = Layout.batch_sharded(self.mesh, _MESH_DIM_Y, rank=1)
tensor = nullary_op([size], layout=None)
def fn():
return nullary_op([size], layout=layout_y)
if is_graph:
fn = polymorphic_function.function(fn)
with api.default_mesh(self.mesh):
dtensor = fn()
self.assertDTensorEqual(tensor, layout_y, dtensor)
@combinations.generate(
combinations.combine(
is_graph=[False, True],
size=[32, 4096],
nullary_op=[array_ops.zeros_like_v2, array_ops.ones_like_v2],
)
)
def testNullaryLikeOpWithLayout(self, is_graph, size, nullary_op):
layout_x = Layout.batch_sharded(self.mesh, batch_dim=_MESH_DIM_X, rank=1)
layout_y = Layout.batch_sharded(self.mesh, batch_dim=_MESH_DIM_Y, rank=1)
tensor = array_ops.zeros([size], layout=None)
tensor_like = nullary_op(tensor, layout=None)
dtensor = array_ops.zeros([size], layout=layout_x)
self.assertDTensorEqual(tensor, layout_x, dtensor)
def fn(layout):
return nullary_op(dtensor, layout=layout)
if is_graph:
fn = polymorphic_function.function(fn)
with api.default_mesh(self.mesh):
dtensor_like = fn(layout_y)
self.assertDTensorEqual(tensor_like, layout_y, dtensor_like)
@combinations.generate(
combinations.combine(
is_graph=[True],
size=[32, 4096],
nullary_op=[array_ops.zeros_like_v2, array_ops.ones_like_v2],
)
)
def testNullaryLikeOpWithoutLayoutEager(self, is_graph, size, nullary_op):
layout_x = Layout.batch_sharded(self.mesh, batch_dim=_MESH_DIM_X, rank=1)
layout_replicated = Layout.replicated(self.mesh, rank=1)
tensor = array_ops.zeros([size], layout=None)
tensor_like = nullary_op(tensor, layout=None)
dtensor = array_ops.zeros([size], layout=layout_x)
self.assertDTensorEqual(tensor, layout_x, dtensor)
def fn(layout):
return nullary_op(dtensor, layout=layout)
if is_graph:
fn = polymorphic_function.function(fn)
with api.default_mesh(self.mesh):
dtensor_like = fn(None)
self.assertDTensorEqual(tensor_like, layout_replicated, dtensor_like)
@combinations.generate(
combinations.combine(
is_graph=[False],
size=[32, 4096],
nullary_op=[array_ops.zeros_like_v2, array_ops.ones_like_v2],
)
)
def testNullaryLikeOpWithoutLayoutGraph(self, is_graph, size, nullary_op):
layout_x = Layout.batch_sharded(self.mesh, batch_dim=_MESH_DIM_X, rank=1)
tensor = array_ops.zeros([size], layout=None)
tensor_like = nullary_op(tensor, layout=None)
dtensor = array_ops.zeros([size], layout=layout_x)
self.assertDTensorEqual(tensor, layout_x, dtensor)
def fn(layout):
return nullary_op(dtensor, layout=layout)
if is_graph:
fn = polymorphic_function.function(fn)
with api.default_mesh(self.mesh):
dtensor_like = fn(None)
self.assertDTensorEqual(tensor_like, layout_x, dtensor_like)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/batchparallel_spmd_test.py
+660
View File
@@ -0,0 +1,660 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for batchparallel_spmd."""
import itertools
from absl.testing import parameterized
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.dtensor.python.tests import test_util_ops
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import gen_image_ops
from tensorflow.python.ops import gen_linalg_ops
from tensorflow.python.ops import nn_impl
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import test
# pylint: enable=g-direct-tensorflow-import
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
class DTensorBatchParallelSPMDTest(test_util.DTensorBaseTest):
def setUp(self):
super(DTensorBatchParallelSPMDTest, self).setUp()
self.skipForDeviceType(['TPU'],
'all tests require 8 TPU cores.',
unless_device_count_equals_to=8)
# Builds a 8x2 mesh.
self._mesh_dim_b = 'b'
self._mesh_dim_x = 'x'
self._dims = [self._mesh_dim_b, self._mesh_dim_x]
global_ids = test_util.create_device_ids_array((4, 2))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: Mesh(
self._dims,
global_ids,
local_ids,
test_util.create_device_list((4, 2), device),
)
for device in ('CPU', 'GPU', 'TPU')
}
self.mesh = self.configTestMesh(mesh_dict)
context.ensure_initialized()
# Creates a bunch of common layouts used by tests later.
# 4-d
self.replicated_layout_4d = Layout.replicated(self.mesh, rank=4)
self.batch_layout_4d = Layout.batch_sharded(
self.mesh, self._mesh_dim_b, rank=4)
# 5-d
self.replicated_layout_5d = Layout.replicated(self.mesh, rank=5)
self.batch_layout_5d = Layout.batch_sharded(
self.mesh, self._mesh_dim_b, rank=5)
@parameterized.named_parameters(('NoBatchDim', 0), ('SingleBatchDim', 1),
('TwoBatchDim', 2))
def testCholesky(self, num_batch_dim):
# Input needs to be symmetric and positive definite.
x = constant_op.constant(
[[1, 1, 1, 1], [1, 5, 5, 5], [1, 5, 14, 14], [1, 5, 14, 17]],
dtype=dtypes.float32,
)
for _ in range(num_batch_dim):
x = array_ops.expand_dims_v2(x, 0)
s = [4] + [1 for _ in range(array_ops.rank(x) - 1)]
x = gen_array_ops.tile(x, s)
expected_result = gen_linalg_ops.cholesky(x)
if num_batch_dim == 0:
layout_spec = []
elif num_batch_dim == 1:
layout_spec = [self._mesh_dim_b]
elif num_batch_dim == 2:
layout_spec = [self._mesh_dim_b, self._mesh_dim_x]
layout = Layout(layout_spec + ['unsharded'] * 2, self.mesh)
x = numpy_util.pack_numpy(x, layout)
got = gen_linalg_ops.cholesky(input=x)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(
test_util.product(
[('NoBatchDim', 0), ('SingleBatchDim', 1), ('TwoBatchDim', 2)],
test_util_ops.FFT_OPS,
)
)
def testFFT(self, num_batch_dim, fft_op, num_nonbatch_dim):
shape = [4 for i in range(num_batch_dim + num_nonbatch_dim)]
np.random.seed(123)
x = constant_op.constant(
np.random.normal(0.0, 1.0, np.prod(shape)).reshape(shape),
dtype=dtypes.complex64,
)
expected_result = fft_op(input=x)
if num_batch_dim == 0:
layout_spec = []
elif num_batch_dim == 1:
layout_spec = [self._mesh_dim_b]
elif num_batch_dim == 2:
layout_spec = [self._mesh_dim_b, self._mesh_dim_x]
layout = Layout(layout_spec + ['unsharded'] * num_nonbatch_dim, self.mesh)
x = numpy_util.pack_numpy(x, layout)
got = fft_op(input=x)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(
test_util.product(
[('NoBatchDim', 0), ('SingleBatchDim', 1), ('TwoBatchDim', 2)],
test_util_ops.RFFT_OPS,
)
)
def testRFFT(self, num_batch_dim, rfft_op, num_nonbatch_dim, dtype):
self.skipForDeviceType(['GPU'], 'RFFT has numerical issues on GPU')
shape = [4 for i in range(num_batch_dim + num_nonbatch_dim)]
np.random.seed(123)
x = constant_op.constant(
np.random.normal(0.0, 1.0, np.prod(shape)).reshape(shape), dtype=dtype
)
expected_result = rfft_op(input=x, fft_length=[2] * num_nonbatch_dim)
if num_batch_dim == 0:
layout_spec = []
elif num_batch_dim == 1:
layout_spec = [self._mesh_dim_b]
elif num_batch_dim == 2:
layout_spec = [self._mesh_dim_b, self._mesh_dim_x]
layout = Layout(layout_spec + ['unsharded'] * num_nonbatch_dim, self.mesh)
x = numpy_util.pack_numpy(x, layout)
got = rfft_op(input=x, fft_length=[2] * num_nonbatch_dim)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(
test_util.product(
[('Replicated', 'replicated'), ('Sharded', 'batch')],
[
(
'SamePadding',
'SAME',
),
(
'ValidPadding',
'VALID',
),
],
test_util_ops.BATCH_PARALLEL_2D_WINDOW_OPS,
)
)
def test2DWindowOp(self, layout_spec, padding, op):
np.random.seed(123)
row_window_size = 3
col_window_size = 4
window_size = [1, row_window_size, col_window_size, 1]
stride_size = [1, row_window_size - 1, col_window_size - 1, 1]
num_rows = (row_window_size - 1) * 5 + 1
num_cols = (col_window_size - 1) * 7 + 1
x_in = np.random.normal(0.0, 1.0, 8 * num_rows * num_cols * 3).reshape(
[8, num_rows, num_cols, 3])
inputs = constant_op.constant(x_in, dtype=dtypes.float32)
expected_result = op(inputs, window_size, stride_size, padding)
if layout_spec == 'replicated':
layout = self.replicated_layout_4d
else:
layout = self.batch_layout_4d
x = numpy_util.pack_numpy(inputs, layout)
got = op(x, window_size, stride_size, padding)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(
test_util.product(
[('Replicated', 'replicated'), ('BatchSharded', 'batch')],
[
(
'SamePadding',
'SAME',
),
(
'ValidPadding',
'VALID',
),
],
test_util_ops.BATCH_PARALLEL_3D_WINDOW_OPS,
)
)
def test3DWindowOp(self, layout_spec, padding, op):
np.random.seed(123)
dep_window_size = 2
row_window_size = 3
col_window_size = 4
window_size = [1, dep_window_size, row_window_size, col_window_size, 1]
stride_size = [
1, dep_window_size - 1, row_window_size - 1, col_window_size - 1, 1
]
num_deps = 3
num_rows = (row_window_size - 1) * 5 + 1
num_cols = (col_window_size - 1) * 7 + 1
x_in = np.random.normal(0.0, 1.0, 8 * num_deps * num_rows * num_cols *
3).reshape([8, num_deps, num_rows, num_cols, 3])
inputs = constant_op.constant(x_in, dtype=dtypes.float32)
expected_result = op(inputs, window_size, stride_size, padding)
if layout_spec == 'replicated':
layout = self.replicated_layout_5d
else:
layout = self.batch_layout_5d
x = numpy_util.pack_numpy(inputs, layout)
got = op(x, window_size, stride_size, padding)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(test_util_ops.PADDINGS)
def testDepthwiseConv2dNative(self, padding):
np.random.seed(123)
x_in = np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1])
kernel_in = np.array([
[[[2, 0.1]], [[3, 0.2]]],
[[[0, 0.3]], [[1, 0.4]]],
])
inputs = constant_op.constant(x_in, dtype=dtypes.float32)
kernel = constant_op.constant(kernel_in, dtype=dtypes.float32)
expected_result = nn_impl.depthwise_conv2d_v2(
inputs, kernel, strides=[1, 1, 1, 1], padding=padding
)
layout = self.batch_layout_4d
x = numpy_util.pack_numpy(inputs, layout)
kernel = numpy_util.pack_numpy(kernel, self.replicated_layout_4d)
got = nn_impl.depthwise_conv2d_v2(
x, kernel, strides=[1, 1, 1, 1], padding=padding
)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(('Sharded', 'sharded'),
('Replicated', 'replicated'))
def testResizeBilinear(self, shard_spec):
np.random.seed(123)
images = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1]),
dtype=dtypes.float32,
)
expected_result = gen_image_ops.resize_bilinear(
images=images,
size=[3, 3],
align_corners=False,
half_pixel_centers=False,
name=None,
)
if shard_spec == 'sharded':
layout = self.batch_layout_4d
else:
layout = self.replicated_layout_4d
images = numpy_util.pack_numpy(images, layout)
got = gen_image_ops.resize_bilinear(
images=images,
size=[3, 3],
align_corners=False,
half_pixel_centers=False,
name=None,
)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(('Sharded', 'sharded'),
('Replicated', 'replicated'))
def testResizeNearestNeighbor(self, shard_spec):
np.random.seed(123)
images = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1]),
dtype=dtypes.float32,
)
expected_result = gen_image_ops.resize_nearest_neighbor(
images=images,
size=[3, 3],
align_corners=False,
half_pixel_centers=False,
name=None,
)
if shard_spec == 'sharded':
layout = self.batch_layout_4d
else:
layout = self.replicated_layout_4d
images = numpy_util.pack_numpy(images, layout)
got = gen_image_ops.resize_nearest_neighbor(
images=images,
size=[3, 3],
align_corners=False,
half_pixel_centers=False,
name=None,
)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(('Sharded', 'sharded'),
('Replicated', 'replicated'))
def testAdjustContrastv2(self, shard_spec):
np.random.seed(123)
images = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9 * 3).reshape([8, 9, 9, 3]),
dtype=dtypes.float32,
)
expected_result = gen_image_ops.adjust_contrastv2(
images=images, contrast_factor=0.5
)
if shard_spec == 'sharded':
layout = self.batch_layout_4d
else:
layout = self.replicated_layout_4d
images = numpy_util.pack_numpy(images, layout)
got = gen_image_ops.adjust_contrastv2(images=images, contrast_factor=0.5)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.named_parameters(('Sharded', 'sharded'),
('Replicated', 'replicated'))
def testAdjustSaturation(self, shard_spec):
np.random.seed(123)
images = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9 * 3).reshape([8, 9, 9, 3]),
dtype=dtypes.float32,
)
expected_result = gen_image_ops.adjust_saturation(images=images, scale=0.5)
if shard_spec == 'sharded':
layout = self.batch_layout_4d
else:
layout = self.replicated_layout_4d
images = numpy_util.pack_numpy(images, layout)
got = gen_image_ops.adjust_saturation(images=images, scale=0.5)
self.assertDTensorEqual(expected_result, layout, got)
@parameterized.parameters(
itertools.permutations(['sharded', 'replicated'], 2))
def testResizeBilinearGradBatchSharded(self, spec1, spec2):
np.random.seed(123)
images = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1]),
dtype=dtypes.float32,
)
grads = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1]),
dtype=dtypes.float32,
)
expected_result = gen_image_ops.resize_bilinear_grad(
grads=grads,
original_image=images,
align_corners=False,
half_pixel_centers=False,
name=None,
)
specs = [spec1, spec2]
layouts = [
self.batch_layout_4d if spec == 'sharded' else self.replicated_layout_4d
for spec in specs
]
# Test images is replicated, grads is batch sharded
images = numpy_util.pack_numpy(images, layouts[0])
grads = numpy_util.pack_numpy(grads, layouts[1])
got = gen_image_ops.resize_bilinear_grad(
grads=grads,
original_image=images,
align_corners=False,
half_pixel_centers=False,
name=None,
)
self.assertDTensorEqual(expected_result, self.batch_layout_4d, got)
def testResizeBilinearGradReplicated(self):
np.random.seed(123)
images = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1]),
dtype=dtypes.float32,
)
grads = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1]),
dtype=dtypes.float32,
)
expected_result = gen_image_ops.resize_bilinear_grad(
grads=grads,
original_image=images,
align_corners=False,
half_pixel_centers=False,
name=None,
)
images = numpy_util.pack_numpy(images, self.replicated_layout_4d)
grads = numpy_util.pack_numpy(grads, self.replicated_layout_4d)
got = gen_image_ops.resize_bilinear_grad(
grads=grads,
original_image=images,
align_corners=False,
half_pixel_centers=False,
name=None,
)
self.assertDTensorEqual(expected_result, self.replicated_layout_4d, got)
@parameterized.named_parameters(
test_util.product([('Replicated', 'replicated'), ('Sharded', 'batch')], [(
'SamePadding',
'SAME',
), (
'ValidPadding',
'VALID',
)]))
def testMaxPool3DGrad(self, shard_spec, padding):
np.random.seed(123)
dep_window_size = 2
row_window_size = 3
col_window_size = 4
window_size = [1, dep_window_size, row_window_size, col_window_size, 1]
stride_size = [
1, dep_window_size - 1, row_window_size - 1, col_window_size - 1, 1
]
num_deps = 3
num_rows = (row_window_size - 1) * 5 + 1
num_cols = (col_window_size - 1) * 7 + 1
x_in = np.random.normal(0.0, 1.0, 8 * num_deps * num_rows * num_cols *
3).reshape([8, num_deps, num_rows, num_cols, 3])
inputs = constant_op.constant(x_in, dtype=dtypes.float32)
with backprop.GradientTape() as tape:
tape.watch([inputs])
expected_result = nn_ops.max_pool3d(
inputs, window_size, stride_size, padding
)
expected_grad = tape.gradient(expected_result, [inputs])
layout = (
self.batch_layout_5d
if shard_spec == 'sharded'
else self.replicated_layout_5d
)
inputs = numpy_util.pack_numpy(inputs, layout)
with ops.device_v2(api.device_name()):
with backprop.GradientTape() as tape:
tape.watch([inputs])
dtensor_result = nn_ops.max_pool3d(
inputs, window_size, stride_size, padding
)
dtensor_grad = tape.gradient(dtensor_result, [inputs])
self.assertDTensorEqual(expected_grad[0], layout, dtensor_grad[0])
@parameterized.named_parameters(
test_util.product([('Replicated', 'replicated'), ('Sharded', 'batch')], [(
'SamePadding',
'SAME',
), (
'ValidPadding',
'VALID',
)]))
def testMaxPool3DGradGrad(self, shard_spec, padding):
np.random.seed(123)
dep_window_size = 2
row_window_size = 3
col_window_size = 4
window_size = [1, dep_window_size, row_window_size, col_window_size, 1]
stride_size = [
1, dep_window_size - 1, row_window_size - 1, col_window_size - 1, 1
]
num_deps = 3
num_rows = (row_window_size - 1) * 5 + 1
num_cols = (col_window_size - 1) * 7 + 1
x_in = np.random.normal(0.0, 1.0, 8 * num_deps * num_rows * num_cols *
3).reshape([8, num_deps, num_rows, num_cols, 3])
inputs = constant_op.constant(x_in, dtype=dtypes.float32)
with backprop.GradientTape() as outer_tape:
with backprop.GradientTape() as inner_tape:
outer_tape.watch([inputs])
inner_tape.watch([inputs])
expected_result = nn_ops.max_pool3d(
inputs, window_size, stride_size, padding
)
expected_first_grad = inner_tape.gradient(expected_result, [inputs])
expected_second_grad = outer_tape.gradient(expected_first_grad, [inputs])
if shard_spec == 'sharded':
layout = self.batch_layout_5d
else:
layout = self.replicated_layout_5d
inputs = numpy_util.pack_numpy(inputs, layout)
@polymorphic_function.function()
def compute_gradients(inputs):
with backprop.GradientTape() as outer_tape:
with backprop.GradientTape() as inner_tape:
outer_tape.watch([inputs])
inner_tape.watch([inputs])
dtensor_result = nn_ops.max_pool3d(
inputs, window_size, stride_size, padding
)
dtensor_first_grad = inner_tape.gradient(dtensor_result, [inputs])
dtensor_second_grad = outer_tape.gradient(dtensor_first_grad[0], [inputs])
return dtensor_first_grad, dtensor_second_grad
dtensor_first_grad, dtensor_second_grad = compute_gradients(inputs)
self.assertDTensorEqual(expected_first_grad[0], layout,
dtensor_first_grad[0])
self.assertDTensorEqual(expected_second_grad[0], layout,
dtensor_second_grad[0])
@parameterized.named_parameters(
test_util.product([('Replicated', 'replicated'), ('Sharded', 'batch')], [(
'SamePadding',
'SAME',
), (
'ValidPadding',
'VALID',
)]))
def testMaxPoolGradGrad(self, shard_spec, padding):
np.random.seed(123)
row_window_size = 3
col_window_size = 4
window_size = [1, row_window_size, col_window_size, 1]
stride_size = [1, row_window_size - 1, col_window_size - 1, 1]
num_rows = (row_window_size - 1) * 5 + 1
num_cols = (col_window_size - 1) * 7 + 1
x_in = np.random.normal(0.0, 1.0, 8 * num_rows * num_cols * 3).reshape(
[8, num_rows, num_cols, 3])
inputs = constant_op.constant(x_in, dtype=dtypes.float32)
with backprop.GradientTape() as outer_tape:
with backprop.GradientTape() as inner_tape:
outer_tape.watch([inputs])
inner_tape.watch([inputs])
expected_result = nn_ops.max_pool_v2(
inputs, window_size, stride_size, padding
)
expected_first_grad = inner_tape.gradient(expected_result, [inputs])
expected_second_grad = outer_tape.gradient(expected_first_grad, [inputs])
if shard_spec == 'sharded':
layout = self.batch_layout_4d
else:
layout = self.replicated_layout_4d
inputs = numpy_util.pack_numpy(inputs, layout)
@polymorphic_function.function()
def compute_gradients(inputs):
with backprop.GradientTape() as outer_tape:
with backprop.GradientTape() as inner_tape:
outer_tape.watch([inputs])
inner_tape.watch([inputs])
dtensor_result = nn_ops.max_pool_v2(
inputs, window_size, stride_size, padding
)
dtensor_first_grad = inner_tape.gradient(dtensor_result, [inputs])
dtensor_second_grad = outer_tape.gradient(dtensor_first_grad[0], [inputs])
return dtensor_first_grad, dtensor_second_grad
dtensor_first_grad, dtensor_second_grad = compute_gradients(inputs)
self.assertDTensorEqual(expected_first_grad[0], layout,
dtensor_first_grad[0])
self.assertDTensorEqual(expected_second_grad[0], layout,
dtensor_second_grad[0])
@parameterized.named_parameters(('Sharded', 'sharded'),
('Replicated', 'replicated'))
def testResizeNearestNeighborGrad(self, shard_spec):
np.random.seed(123)
grads = constant_op.constant(
np.random.normal(0.0, 1.0, 8 * 9 * 9).reshape([8, 9, 9, 1]),
dtype=dtypes.float32,
)
expected_result = gen_image_ops.resize_nearest_neighbor_grad(
grads=grads,
size=[3, 3],
align_corners=False,
half_pixel_centers=False,
name=None,
)
if shard_spec == 'sharded':
layout = self.batch_layout_4d
else:
layout = self.replicated_layout_4d
grads = numpy_util.pack_numpy(grads, layout)
got = gen_image_ops.resize_nearest_neighbor_grad(
grads=grads,
size=[3, 3],
align_corners=False,
half_pixel_centers=False,
name=None,
)
self.assertDTensorEqual(expected_result, layout, got)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/cache_test.py
+330
View File
@@ -0,0 +1,330 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests DTensor device cache for compiled function computation."""
import gc
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import combinations
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_stateless_random_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
# Convenient constants to use for tests.
_BATCH_DIM = "batch"
_MESH_DIM_X = "x"
# Shorter notation.
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
def diff_dicts(dict1, dict2):
keys = set(dict1.keys()) | set(dict2.keys())
return {key: dict1.get(key, 0) - dict2.get(key, 0) for key in keys}
class DTensorDeviceCacheTest(test_util.DTensorBaseTest):
def setUp(self):
super(DTensorDeviceCacheTest, self).setUp()
device_ids = test_util.create_device_ids_array((2,))
local_device_ids = np.ravel(device_ids).tolist()
mesh_dict = {
device: Mesh(
[_BATCH_DIM],
device_ids,
local_device_ids,
test_util.create_device_list((2,), device),
)
for device in ("CPU", "GPU", "TPU")
}
self.mesh = self.configTestMesh(mesh_dict)
def testBasic(self):
@polymorphic_function.function
def func0(a):
return a + 1
@polymorphic_function.function
def func1(a):
return a + 2
c0 = api.copy_to_mesh(
constant_op.constant(1.0), Layout.replicated(self.mesh, rank=0)
)
c1 = api.copy_to_mesh(
constant_op.constant([2.0, 3.0]), Layout.replicated(self.mesh, rank=1)
)
c2 = api.copy_to_mesh(
constant_op.constant([4.0]), Layout.replicated(self.mesh, rank=1)
)
c3 = api.copy_to_mesh(
constant_op.constant(1, dtype=dtypes.int32),
Layout.replicated(self.mesh, rank=0),
)
# c0 and c1 have different layouts. c1 and c2 have different shapes.
# c0 and c3 have different dtypes.
self.assertAllEqual(func0(c0), 2.0)
self.assertAllEqual(func0(c1), [3.0, 4.0])
self.assertAllEqual(func0(c2), [5.0])
self.assertAllEqual(func0(c3), 2)
# func0 and func1 have different names.
self.assertAllEqual(func1(c0), 3.0)
def testFunctionInputConstantFoldingCacheHits(self):
@polymorphic_function.function
def add(a, b):
return a + b
c0 = api.copy_to_mesh(
constant_op.constant(17.0), Layout.replicated(self.mesh, rank=0)
)
c1 = api.copy_to_mesh(
constant_op.constant(21.0), Layout.replicated(self.mesh, rank=0)
)
stats1 = api._dtensor_device()._get_stats()
self.assertAllEqual(add(c0, c1), 38.0)
self.assertAllEqual(add(c0, c1), 38.0)
# First call should miss and second should hit.
stats2 = api._dtensor_device()._get_stats()
diff = {key: stats2[key] - stats1[key] for key in stats1.keys()}
self.assertEqual(diff["function_manager.miss"], 1)
self.assertEqual(diff["function_manager.hit"], 1)
def testFunctionInputConstantFoldingCacheMiss(self):
@polymorphic_function.function
def add(a, b):
return a + b
c0 = api.copy_to_mesh(
constant_op.constant(17.0), Layout.replicated(self.mesh, rank=0)
)
c1 = api.copy_to_mesh(
constant_op.constant(21.0), Layout.replicated(self.mesh, rank=0)
)
c2 = api.copy_to_mesh(
constant_op.constant(0.0), Layout.replicated(self.mesh, rank=0)
)
stats1 = api._dtensor_device()._get_stats()
# First call should log a cache miss.
self.assertAllEqual(add(c0, c1), 38.0)
# Second call should also log a cache miss since second constant changed.
self.assertAllEqual(add(c0, c2), 17.0)
# Third call should not log a cache miss since the same input as the prev.
self.assertAllEqual(add(c0, c2), 17.0)
# Fourth call should log a cache miss since first input changed.
self.assertAllEqual(add(c1, c2), 21.0)
stats2 = api._dtensor_device()._get_stats()
diff = {key: stats2[key] - stats1[key] for key in stats1.keys()}
self.assertEqual(diff["function_manager.miss"], 3)
self.assertEqual(diff["function_manager.hit"], 1)
def testCacheWithRNG(self):
with api._dtensor_device()._default_layout(
Layout.replicated(self.mesh, rank=1)):
v0 = gen_stateless_random_ops.stateless_random_normal(
shape=[1], seed=[1, 2]
)
with api._dtensor_device()._default_layout(
Layout.replicated(self.mesh, rank=1)):
v1 = gen_stateless_random_ops.stateless_random_normal(
shape=[1], seed=[1, 2]
)
v2 = gen_stateless_random_ops.stateless_random_normal(
shape=[2], seed=[1, 2]
)
v3 = gen_stateless_random_ops.stateless_random_normal(
shape=[1], seed=[3, 4]
)
# v0 and v1 have same layouts.
self.assertAllEqual(v0, v1)
api.check_layout(v0, Layout.replicated(self.mesh, rank=1))
api.check_layout(v1, Layout.replicated(self.mesh, rank=1))
# v1 and v2 have different shapes.
self.assertNotEqual(v1.shape, v2.shape)
# v1 and v3 have different seeds.
self.assertNotEqual(v1.numpy(), v3.numpy())
def testCacheWithVariable(self):
c0 = api.copy_to_mesh(
constant_op.constant(1.0), Layout.replicated(self.mesh, rank=0)
)
c1 = api.copy_to_mesh(
constant_op.constant([2.0, 3.0]), Layout.replicated(self.mesh, rank=1)
)
a = constant_op.constant([4.0])
b = constant_op.constant([5.0])
c2 = api.pack(
[a, b], layout=Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=1)
)
v0 = d_variable.DVariable(c0)
v1 = d_variable.DVariable(c1)
v2 = d_variable.DVariable(c2)
self.assertAllEqual(v0.read_value(), 1.0)
self.assertAllEqual(v1.read_value(), [2.0, 3.0])
unpacked_tensor = api.unpack(v2.read_value())
self.assertAllClose([4.0], unpacked_tensor[0])
self.assertAllClose([5.0], unpacked_tensor[1])
@combinations.generate(
combinations.combine(size=[16, 40], same_value=[True, False])
)
def testManyFunctions(self, size, same_value):
r = range(100)
values = [np.reshape(r[i : i + size], (4, size // 4)) for i in range(10)]
c_layout = Layout.replicated(self.mesh, rank=2)
values = [constant_op.constant(v, dtype=dtypes.float32) for v in values]
c0 = [api.copy_to_mesh(v, c_layout) for v in values]
c0 = [c0[0 if same_value else i] for i in range(10)]
e0 = [values[0 if same_value else i] for i in range(10)]
stats1 = api._dtensor_device()._get_stats()
for i in range(10):
# Use a special to ensure no conflicts with otherwise used names.
@polymorphic_function.function
def fn_31415926(c):
return math_ops.reduce_sum(c)
self.assertAllEqual(fn_31415926(c0[i]).numpy(), np.sum(e0[i]))
del fn_31415926
gc.collect()
stats2 = api._dtensor_device()._get_stats()
diff = diff_dicts(stats2, stats1)
self.assertEqual(diff["function_manager.size"], 0)
self.assertEqual(diff["kernel_cache.size"], 0)
self.assertEqual(diff["device_cache.size"], 0)
@combinations.generate(
combinations.combine(size=[16, 40], same_value=[True, False])
)
def testManyEagerOps(self, size, same_value):
if self.mesh.device_type() != "TPU":
# For the CPU/GPU mesh, we have a shortcut that doesn't go through the
# MLIR, but run the eager op locally and broadcast to all the devices.
expected_cache_diff = 0
expected_kernel_cache = 0
expected_device_cache = 0
expected_eager_pure_hit = 10
else:
# TODO(b/287529295): Remove this branch after the TPU issue is fixed.
expected_device_cache = 0
expected_eager_pure_hit = 0
if same_value:
expected_cache_diff = 1
expected_kernel_cache = 2
else:
if size >= 20:
expected_cache_diff = 1
expected_kernel_cache = 2
else:
expected_cache_diff = 2
expected_kernel_cache = 4
r = range(100)
c_layout = Layout.replicated(self.mesh, rank=2)
values = [np.reshape(r[i : i + size], (4, size // 4)) for i in range(10)]
values = [constant_op.constant(v, dtype=dtypes.float32) for v in values]
c0 = [api.copy_to_mesh(v, c_layout) for v in values]
c0 = [c0[0 if same_value else i] for i in range(10)]
e0 = [values[0 if same_value else i] for i in range(10)]
stats1 = api._dtensor_device()._get_stats()
for i in range(10):
self.assertAllEqual(array_ops.identity(c0[i]).numpy(), e0[i])
gc.collect()
stats2 = api._dtensor_device()._get_stats()
diff = diff_dicts(stats2, stats1)
if same_value:
self.assertEqual(diff["function_manager.size"], expected_cache_diff)
self.assertEqual(
diff["eager_pure_optimization.hit"], expected_eager_pure_hit
)
# TFRT doesn't use eager cache.
if not test_util.is_tfrt_enabled():
self.assertEqual(diff["kernel_cache.size"], expected_kernel_cache)
self.assertEqual(diff["device_cache.size"], expected_device_cache)
else:
# FIXME(feyu): Update these when the leaks are fixed.
if size >= 20:
self.assertEqual(diff["function_manager.size"], expected_cache_diff)
self.assertEqual(
diff["eager_pure_optimization.hit"], expected_eager_pure_hit
)
# TFRT doesn't use eager cache.
if not test_util.is_tfrt_enabled():
self.assertEqual(diff["kernel_cache.size"], expected_kernel_cache)
self.assertEqual(diff["device_cache.size"], expected_device_cache)
else:
self.assertEqual(diff["function_manager.size"], expected_cache_diff)
self.assertEqual(
diff["eager_pure_optimization.hit"], expected_eager_pure_hit
)
# TFRT doesn't use eager cache.
if not test_util.is_tfrt_enabled():
self.assertEqual(diff["kernel_cache.size"], expected_kernel_cache)
self.assertEqual(diff["device_cache.size"], expected_device_cache)
def testManyEagerOpsVaryInput(self):
c_layout = Layout.replicated(self.mesh, rank=10)
c0 = constant_op.constant(
[[[[[[[[[[0, 1, 2, 3], [4, 5, 6, 7]]]]]]]]]], dtype=dtypes.float32
)
e0 = c0.numpy()
c0 = api.copy_to_mesh(c0, c_layout)
for ax in range(10):
self.assertAllEqual(
math_ops.reduce_sum(c0, axis=ax).numpy(), np.sum(e0, axis=ax)
)
if __name__ == "__main__":
test.main()
tensorflow/dtensor/python/tests/collective_test.py
+574
View File
@@ -0,0 +1,574 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for DTensor collectives."""
import os
from absl.testing import parameterized
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import dtensor_device
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
# pylint: enable=g-direct-tensorflow-import
Layout = layout_lib.Layout
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
_MESH_DIMS = [_MESH_DIM_X, _MESH_DIM_Y]
class CollectiveTest(test_util.DTensorBaseTest):
def setUp(self):
super(CollectiveTest, self).setUp()
global_ids = test_util.create_device_ids_array((2, 1))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout_lib.Mesh(_MESH_DIMS, global_ids, local_ids,
test_util.create_device_list((2, 1), device))
for device in ('CPU', 'GPU', 'TPU')
}
self.mesh = self.configTestMesh(mesh_dict)
self.fully_replicated_layout_2d = Layout.replicated(self.mesh, rank=2)
self.first_dimension_sharded_layout_2d = Layout.batch_sharded(
self.mesh, _MESH_DIM_X, 2)
self.scalar_layout = Layout.replicated(self.mesh, rank=0)
def testReduceOnBfloat16(self):
self.skipForDeviceType(['GPU'],
'GPUs do not support bfloat16 collective reduce')
self.skipForDeviceType(['TPU'],
'This test only needs to run on 2 cores.',
unless_device_count_equals_to=2)
a = constant_op.constant(
np.array([[1, 2, 3, 4], [5.0, 6.0, 7.0, 8.0]]), dtype=dtypes.bfloat16)
expected_result = math_ops.reduce_sum(a)
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
dtensor_result = math_ops.reduce_sum(sharded_a)
self.assertDTensorEqual(expected_result, self.scalar_layout, dtensor_result)
def testReduceOnInt32(self):
a = constant_op.constant(
np.array([[1, 2, 3, 4], [5, 6, 7, 8]]), dtype=dtypes.int32)
expected_result = math_ops.reduce_sum(a)
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
dtensor_result = math_ops.reduce_sum(sharded_a)
self.assertDTensorEqual(expected_result, self.scalar_layout, dtensor_result)
def testReduceOnInt8(self):
a = constant_op.constant(
np.array([[1, 2, 3, 4], [5, 6, 7, 8]]), dtype=dtypes.int8
)
expected_result = math_ops.reduce_sum(a)
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
dtensor_result = math_ops.reduce_sum(sharded_a)
self.assertDTensorEqual(expected_result, self.scalar_layout, dtensor_result)
def testTwoReducesWithAssign(self):
# FIXME(b/238384852): The purpose of this test is to validate the control
# dependency added by DTensor.
# However, as we have no way of testing the per-device graph
# produced by the DTensor SPMD expansion, we have to use manual logging
# to verify if the results are correct.
# For example, this test is used to check cl/459358521.
# Uses dtypes.float32 to avoid int32 problem with VarHandleOp on GPUs.
a = constant_op.constant(
np.array([[1, 2, 3, 4], [5, 6, 7, 8]]), dtype=dtypes.float32)
b = constant_op.constant(
np.array([[11, 12, 13, 4], [15, 16, 17, 18]]), dtype=dtypes.float32)
expected_result = math_ops.reduce_sum(a) * math_ops.reduce_sum(b)
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
sharded_b = api.relayout(b, self.first_dimension_sharded_layout_2d)
sharded_v = d_variable.DVariable(sharded_b)
@polymorphic_function.function
def func(a, b):
a1 = math_ops.reduce_sum(a, name='reducea')
sharded_v.assign(a)
b1 = math_ops.reduce_sum(b, name='reduceb')
return a1 * b1
with api.default_mesh(self.mesh):
dtensor_result = func(sharded_a, sharded_b)
self.assertDTensorEqual(expected_result, self.scalar_layout, dtensor_result)
@parameterized.named_parameters(
('_all', math_ops.reduce_all),
('_any', math_ops.reduce_any),
)
def testReduceOnBool(self, reduction):
# TODO(b/193531363): Track the work to support int32 reduce.
self.skipForDeviceType(['GPU'],
'GPUs do not support int32 collective reduce')
self.skipForDeviceType(['TPU'],
'This test only needs to run on 2 cores.',
unless_device_count_equals_to=2)
a = constant_op.constant(
np.array([[True, False, False, True], [False, False, False, True]]),
dtype=dtypes.bool)
expected_result = reduction(a)
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
dtensor_result = reduction(sharded_a)
self.assertDTensorEqual(expected_result, self.scalar_layout, dtensor_result)
def testAllToAllOnBool(self):
# TODO(b/193531363): Track the work to support int32 reduce.
self.skipForDeviceType(['GPU'],
'GPUs do not support int32 collective reduce')
self.skipForDeviceType(['TPU'],
'This test only needs to run on 2 cores.',
unless_device_count_equals_to=2)
a = constant_op.constant(
np.array([[True, False, False, True], [False, False, False, True]]),
dtype=dtypes.bool)
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
unsharded_a = api.relayout(sharded_a, self.fully_replicated_layout_2d)
self.assertDTensorEqual(a, self.fully_replicated_layout_2d, unsharded_a)
def testAllToAllOnInt32(self):
# TODO(b/193471732): Tracking the work to do int32 all-concat.
#
# Currently, the test will fail with segfault.
self.skipForDeviceType(['GPU'],
'GPUs do not support int32 StridedSliceXXX Ops')
self.skipForDeviceType(['TPU'],
'This test only needs to run on 2 cores.',
unless_device_count_equals_to=2)
a = constant_op.constant(np.array([[1, 2], [3, 4]]), dtype=dtypes.int32)
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
unsharded_a = api.relayout(sharded_a, self.fully_replicated_layout_2d)
self.assertDTensorEqual(a, self.fully_replicated_layout_2d, unsharded_a)
def testCollectiveOpsOnComplex64(self):
# This functions tests for AllScatter, AllGather, and AllReduce.
a = constant_op.constant(
np.array([[1, 2 + 2j], [3 + 1j, 4 + 5j]]), dtype=dtypes.complex64
)
# Tests AllScatter
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
# Tests AllGather / AllReduce
unsharded_a = api.relayout(sharded_a, self.fully_replicated_layout_2d)
self.assertDTensorEqual(a, self.fully_replicated_layout_2d, unsharded_a)
def testCollectiveOpsOnComplex128(self):
# This function tests for AllScattering, AllReduce, and AllToAll.
self.skipForDeviceType(['TPU'], 'TPU does not support comolex128')
expected_layout = Layout.inner_sharded(self.mesh, 'x', rank=2)
initial_layout = Layout.batch_sharded(self.mesh, 'x', rank=2)
a = constant_op.constant(
np.array([[1, 2 + 2j], [3 + 1j, 4 + 5j]]), dtype=dtypes.complex128)
# Tests AllScatter
sharded_a_initial = api.relayout(a, initial_layout)
# Tests AllToAll / AllReduce
sharded_a = api.relayout(sharded_a_initial, expected_layout)
api.check_layout(sharded_a, expected_layout)
def testNoOpAllToAll(self):
self.skipForDeviceType(['TPU'],
'This test only needs to run on 2 cores.',
unless_device_count_equals_to=2)
a = constant_op.constant(
np.array([[1., 2., 3., 4.], [5.0, 6.0, 7.0, 8.0]]),
dtype=dtypes.float32)
expected_result = a
sharded_a = api.relayout(a, self.first_dimension_sharded_layout_2d)
dtensor_result = api.relayout(sharded_a, self.fully_replicated_layout_2d)
self.assertDTensorEqual(expected_result, self.fully_replicated_layout_2d,
dtensor_result)
# Regression test for b/184401449.
def testDeviceIdTensorOnSplitHost(self):
if not test_util.is_tpu_present():
self.skipTest('This test only runs on TPUs.')
self.skipForDeviceType(['TPU'],
'This test requires 8 TPU cores.',
unless_device_count_equals_to=8)
global_ids = test_util.create_device_ids_array((2, 4))
local_ids = [0, 1, 4, 5, 2, 3, 6, 7] # not sequentially increasing
mesh = layout_lib.Mesh(_MESH_DIMS, global_ids, local_ids,
test_util.create_device_list((2, 4), 'TPU'),
'tpu_mesh')
# This works because on 2x2, global device IDs are equal to physical TPU
# core IDs: both are range(8). So local device IDs happen to be usable here.
# TODO(b/180046115): Add a device.get_tpu_core_ids method and translate
# device IDs to core IDs before setting the list here.
device = dtensor_device.DTensorDevice(meshes=[mesh])
device.set_tpu_core_ids('tpu_mesh', local_ids)
layout_x = Layout.batch_sharded(mesh, _MESH_DIM_X, 2)
layout_y = Layout.batch_sharded(mesh, _MESH_DIM_Y, 2)
# Create a 2x4 batch-sharded d-tensor, with batch IDs in its first column
# and zeros in other columns.
replica_ids = constant_op.constant(
np.array([[0, 0, 0, 0], [1, 0, 0, 0]]), dtype=dtypes.int32
)
replica_ids = api.relayout(replica_ids, layout_x)
# Create a 4x4 y-sharded d-tensor filled with ones.
ones = constant_op.constant(
np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]),
dtype=dtypes.int32,
)
ones = api.relayout(ones, layout_y)
# If `a` has a layout of [x, unsharded], and `b` has a layout of
# [y, unsharded], the matmul will slice `a` to [x, y], do a local matmul,
# and all-reduce to produce a final result with a layout of [x, unsharded].
#
# Because `a` only has non-zero values in its first column, local devices
# must be given a correct device ID tensor (as the first argument of every
# function) to produce correct `begin` values for slicing `a`.
#
# Although this function only contains a single op, running it in op-by-op
# mode doesn't produce the intended effect because the output of
# math_ops.matmul would have a layout of [y, unsharded] instead of
# [x, unsharded].
@polymorphic_function.function
def func(a, b):
return math_ops.matmul(a, b)
dtensor_result = func(replica_ids, ones)
# The correct result is a 2x4 batch-sharded d-tensor, with rows filled with
# batch IDs.
expected_result = [
constant_op.constant(
[loc[_MESH_DIM_X]] * 4, dtype=dtypes.int32, shape=[1, 4])
for loc in mesh.local_device_locations()
]
self.assertEqual(api.fetch_layout(dtensor_result), layout_x)
dtensor_result = [t.numpy() for t in api.unpack(dtensor_result)]
self.assertAllEqual(expected_result, dtensor_result)
def testDifferentShapesBetweenCalls(self):
self.skipForTfrt(
'b/269333905, TFRT cpu fails due to step_id not propagated.'
)
self.skipForDeviceType(
['TPU'],
'Known failure under TPU for legalization requires a static shape.',
)
# The error only happens across the batch, where the value of
# tf.unique are differnet.
def produce_data(inputs, label):
inputs = api.relayout(
inputs, Layout.batch_sharded(self.mesh, _MESH_DIM_X, 1)
)
label = api.relayout(
label, Layout.batch_sharded(self.mesh, _MESH_DIM_X, 1)
)
return inputs, label
@polymorphic_function.function
def train_fn(inputs, label):
inputs, indices = array_ops.unique(inputs)
return math_ops.unsorted_segment_sum(label, indices, len(inputs))
input1, label1 = produce_data([6, 0, 6, 0], [1, 2, 3, 4])
input2, label2 = produce_data([2, 1, 2, 0], [1, 2, 3, 4])
result1 = train_fn(input1, label1)
result2 = train_fn(input2, label2)
self.assertAllEqual(
result1.numpy(),
[
4,
6,
],
)
self.assertAllEqual(
result2.numpy(),
[
4,
2,
4,
],
)
class CollectiveTestWithCustomMesh(test_util.DTensorBaseTest):
# Create two independent global AllReduce ops that should get combined.
def testGlobalAllReduceCombiner(self):
self.skipForDeviceType(['TPU'],
'This test requires 8 TPU cores.',
unless_device_count_equals_to=8)
# Create and use an eight-device mesh just for this test.
global_ids = test_util.create_device_ids_array((8,))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout_lib.Mesh([_MESH_DIM_X], global_ids, local_ids,
test_util.create_device_list((8,), device))
for device in ('CPU', 'GPU', 'TPU')
}
mesh = self.configTestMesh(mesh_dict)
fully_replicated_layout_1d = Layout.replicated(mesh, rank=1)
first_dimension_sharded_layout_2d = Layout.batch_sharded(
mesh, _MESH_DIM_X, 2)
@polymorphic_function.function
def func(a, b):
a = math_ops.reduce_sum(a, axis=[0])
b = math_ops.reduce_sum(b, axis=[0])
# Do something with the results before adding them, to make sure the MLIR
# pass can handle dependent ops sandwiched between two all-reduce ops.
return gen_math_ops.square(a) + gen_math_ops.square(b)
row = constant_op.constant(np.array([[1., 2.0]]), dtype=dtypes.float32)
a = array_ops.repeat(row, repeats=[8], axis=0)
b = gen_array_ops.reverse_v2(a, axis=[1])
expected_result = func(a, b)
a = api.relayout(a, first_dimension_sharded_layout_2d)
b = api.relayout(b, first_dimension_sharded_layout_2d)
dtensor_result = func(a, b)
self.assertDTensorEqual(expected_result, fully_replicated_layout_1d,
dtensor_result)
# Create two independent global AllReduce ops that should get combined.
# Create two independent global AllReduce ops with different reductions, that
# should not get combined
def testGlobalAllReduceCombinerDifferentReduce(self):
self.skipForDeviceType(['TPU'],
'This test requires 8 TPU cores.',
unless_device_count_equals_to=8)
# Create and use an eight-device mesh just for this test.
global_ids = test_util.create_device_ids_array((8,))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout_lib.Mesh([_MESH_DIM_X], global_ids, local_ids,
test_util.create_device_list((8,), device))
for device in ('CPU', 'GPU', 'TPU')
}
mesh = self.configTestMesh(mesh_dict)
fully_replicated_layout_1d = Layout.replicated(mesh, rank=1)
first_dimension_sharded_layout_2d = Layout.batch_sharded(
mesh, _MESH_DIM_X, 2)
@polymorphic_function.function
def func(a, b):
a = math_ops.reduce_sum(a, axis=[0])
b = math_ops.reduce_mean(b, axis=[0])
# Do something with the results before adding them, to make sure the MLIR
# pass can handle dependent ops sandwiched between two all-reduce ops.
return gen_math_ops.square(a) + gen_math_ops.square(b)
row = constant_op.constant(np.array([[1., 2.0]]), dtype=dtypes.float32)
a = array_ops.repeat(row, repeats=[8], axis=0)
b = gen_array_ops.reverse_v2(a, axis=[1])
expected_result = func(a, b)
a = api.relayout(a, first_dimension_sharded_layout_2d)
b = api.relayout(b, first_dimension_sharded_layout_2d)
dtensor_result = func(a, b)
self.assertDTensorEqual(expected_result, fully_replicated_layout_1d,
dtensor_result)
# Create two independent subgroup AllReduce ops that should get combined.
def testSubgroupAllReduceCombiner(self):
self.skipForDeviceType(['TPU'],
'This test requires 8 TPU cores.',
unless_device_count_equals_to=8)
# Create and use an eight-device mesh just for this test.
global_ids = test_util.create_device_ids_array((4, 2))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout_lib.Mesh(_MESH_DIMS, global_ids, local_ids,
test_util.create_device_list((4, 2), device))
for device in ('CPU', 'GPU', 'TPU')
}
mesh = self.configTestMesh(mesh_dict)
fully_sharded_layout_2d = Layout(_MESH_DIMS, mesh)
@polymorphic_function.function
def func(a, b):
a = math_ops.reduce_sum(a, axis=[0])
b = math_ops.reduce_sum(b, axis=[0])
# Do something with the results before adding them, to make sure the MLIR
# pass can handle dependent ops sandwiched between two all-reduce ops.
return gen_math_ops.square(a) + gen_math_ops.square(b)
row = constant_op.constant(np.array([[1., 2.0]]), dtype=dtypes.float32)
a = array_ops.repeat(row, repeats=[8], axis=0)
b = gen_array_ops.reverse_v2(a, axis=[1])
expected_result = func(a, b)
a = api.relayout(a, fully_sharded_layout_2d)
b = api.relayout(b, fully_sharded_layout_2d)
dtensor_result = func(a, b)
self.assertDTensorEqual(expected_result, Layout([_MESH_DIM_Y], mesh),
dtensor_result)
# TODO(b/188605096): also add a MixedPrecisionReduceScatter test
def testMixedPrecisionAllReduce(self):
has_enable_dtensor_mixed_precision_reduce = (
'DTENSOR_ENABLE_MIXED_PRECISION_REDUCE' in os.environ
)
has_dtensor_reduce_in_bfloat16_max_group_size = (
'DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE' in os.environ
)
if has_dtensor_reduce_in_bfloat16_max_group_size:
old_dtensor_reduce_in_bfloat16_max_group_size = os.environ[
'DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE']
os.environ['DTENSOR_ENABLE_MIXED_PRECISION_REDUCE'] = ''
os.environ['DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE'] = '4'
self.skipForDeviceType(['GPU'],
'GPUs do not support bfloat16 reduce')
self.skipForDeviceType(['TPU'],
'This test requires 8 TPU cores.',
unless_device_count_equals_to=8)
# Create and use an 8-device mesh just for this test. Mixed-precision
# AllReduce will be in effect since the reduction will be across 8 devices
# which is larger than the max group size flag value of 4.
global_ids = test_util.create_device_ids_array((8,))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout_lib.Mesh([_MESH_DIM_X], global_ids, local_ids,
test_util.create_device_list((8,), device))
for device in ('CPU', 'GPU', 'TPU')
}
mesh = self.configTestMesh(mesh_dict)
replicated_layout_1d = Layout.replicated(mesh, rank=1)
first_dim_sharded_layout_1d = Layout.batch_sharded(
mesh, _MESH_DIM_X, rank=2)
@polymorphic_function.function
def func(x):
return math_ops.reduce_sum(x, axis=0)
# Reduce across 8 devices.
inp = constant_op.constant(
np.arange(48.).reshape((8, 6)), dtype=dtypes.bfloat16)
expected_result = np.sum(inp, axis=0)
inp_dtensor = api.relayout(inp, first_dim_sharded_layout_1d)
dtensor_result = func(inp_dtensor)
self.assertDTensorEqual(
expected_result, replicated_layout_1d, dtensor_result)
if not has_enable_dtensor_mixed_precision_reduce:
del os.environ['DTENSOR_ENABLE_MIXED_PRECISION_REDUCE']
if has_dtensor_reduce_in_bfloat16_max_group_size:
os.environ['DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE'] = (
old_dtensor_reduce_in_bfloat16_max_group_size
)
else:
del os.environ['DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE']
# Create two independent AllReduce ops with indirect dependency, that should
# not get combined.
def testAllReduceCombinerWithIndirectDependency(self):
# The purpose of this test is to validate the depdency check in AllReduce
# AllReduce combiner (dtensor_allreduce_combine_optimization). Specifically,
# the side effects from indirect dependency.
self.skipForDeviceType(['TPU'],
'This test requires 8 TPU cores.',
unless_device_count_equals_to=8)
# Create and use an eight-device mesh just for this test.
global_ids = test_util.create_device_ids_array((8,))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout_lib.Mesh([_MESH_DIM_X], global_ids, local_ids,
test_util.create_device_list((8,), device))
for device in ('CPU', 'GPU', 'TPU')
}
mesh = self.configTestMesh(mesh_dict)
first_dim_sharded_layout_1d = Layout.batch_sharded(
mesh, _MESH_DIM_X, rank=1)
init_value = constant_op.constant(np.ones(32), dtype=dtypes.float32)
init_value = api.relayout(init_value, first_dim_sharded_layout_1d)
# 2nd reduction indrectly depend on the 1st reduction
@polymorphic_function.function
def func(v):
a = math_ops.reduce_sum(v)
v.assign_add(v+a)
b = math_ops.reduce_sum(v)
return b
v = d_variable.DVariable(init_value)
dtensor_result = func(v)
# Replicate the scenario above without using dtensor
expected_result = constant_op.constant(np.ones(32), dtype=dtypes.float32)
expected_result += expected_result + math_ops.reduce_sum(expected_result)
expected_result = math_ops.reduce_sum(expected_result)
self.assertDTensorEqual(
expected_result, Layout.replicated(mesh=mesh, rank=0), dtensor_result
)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/config_test.py
+92
View File
@@ -0,0 +1,92 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for the open source DTensor Python API."""
import os
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager import context
from tensorflow.python.framework import device as tf_device
from tensorflow.python.platform import test as tf_test
# pylint: enable=g-direct-tensorflow-import
class ConfigTest(tf_test.TestCase):
def setUp(self):
super().setUp()
test_util.reset_logical_devices('CPU', 2)
if test_util.is_gpu_present():
test_util.reset_logical_devices('GPU', 2)
def tearDown(self):
os.environ.pop(config._DT_JOBS, [])
super().tearDown()
def test_env_vars(self):
self.assertEqual(config.client_id(), 0)
self.assertEqual(config.num_clients(), 1)
self.assertEqual(config.job_name(), 'localhost')
self.assertEqual(config.full_job_name(), 'localhost/replica:0/task:0')
self.assertEqual(config.jobs(), [])
def test_list_devices(self):
device_type = config.preferred_device_type()
local_devices = [
tf_device.DeviceSpec.from_string(
f'/job:localhost/replica:0/task:0/device:{device_type}:0'),
tf_device.DeviceSpec.from_string(
f'/job:localhost/replica:0/task:0/device:{device_type}:1'),
]
self.assertEqual(config.local_devices(device_type), local_devices)
self.assertEqual(config.num_local_devices(device_type), 2)
self.assertEqual(config.num_global_devices(device_type), 2)
# The eager context should not be initialized by any of the calls
self.assertFalse(context.context()._initialized) # pylint: disable=protected-access
def test_sort_jobs_with_bns_names(self):
# bns names must be sorted in the bns order.
dtensor_jobs = [
'/bns/localhost/{task_id}'.format(task_id=i) for i in range(16)
]
os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
self.assertListEqual(dtensor_jobs, config.jobs())
dtensor_jobs = [
'/bns/localhost/{task_id}:8888'.format(task_id=i) for i in range(16)
]
os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
self.assertListEqual(dtensor_jobs, config.jobs())
dtensor_jobs = [
'/bns/localhost/{task_id}'.format(task_id=100 - i) for i in range(16)
]
os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
with self.assertRaisesRegex(ValueError, 'Unexpected DTENSOR_JOBS'):
config.jobs()
def test_jobs_with_ip_port(self):
# The ip port format is not a bns address, and not required to sorted.
dtensor_jobs = ['localhost:{port}'.format(port=16 - i) for i in range(16)]
os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
self.assertListEqual(dtensor_jobs, config.jobs())
if __name__ == '__main__':
tf_test.main()
tensorflow/dtensor/python/tests/conv_test.py
+350
View File
@@ -0,0 +1,350 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for executing ops needed to implement image model."""
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import special_math_ops
from tensorflow.python.platform import test
UNSHARDED = layout_lib.UNSHARDED
Mesh = layout_lib.Mesh
Layout = layout_lib.Layout
BATCH_DIM = 'batch'
DEPTH_DIM = 'depth'
HEIGHT_DIM = 'height'
WIDTH_DIM = 'width'
BATCH_SIZE = 4
DEPTH = 8
HEIGHT = 12
WIDTH = 12
CHANNEL_IN = 1
CHANNEL_OUT = 3
class ConvOpTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
global_ids = test_util.create_device_ids_array((2, 2, 2))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {}
for device in ('CPU', 'GPU', 'TPU'):
mesh_dict[device] = Mesh(
[BATCH_DIM, HEIGHT_DIM, WIDTH_DIM],
global_ids,
local_ids,
test_util.create_device_list((2, 2, 2), device),
)
self.mesh = self.configTestMesh(mesh_dict)
self.replicated_2d = Layout.replicated(self.mesh, 2)
self.batch_sharded_2d = Layout.batch_sharded(self.mesh, BATCH_DIM, 2)
@parameterized.named_parameters(
test_util.product(
*[
[
(
'Conv2D',
nn_ops.conv2d_v2,
(BATCH_SIZE, HEIGHT, WIDTH, CHANNEL_IN),
(2, 2, CHANNEL_IN, CHANNEL_OUT),
'bhwc,xy->by',
[1, 2, 1, 1],
),
(
'Conv3D',
nn_ops.conv3d_v2,
(BATCH_SIZE, DEPTH, HEIGHT, WIDTH, CHANNEL_IN),
(2, 2, 2, CHANNEL_IN, CHANNEL_OUT),
'bdhwc,xy->by',
[1, 1, 2, 1, 1],
),
],
[
('Eager', True),
('Graph', False),
],
[
('ReplicatedInput', 'replicated'),
('BatchShardedInput', 'batch_sharded'),
],
[
('ValidPadding', 'VALID'),
('SamePadding', 'SAME'),
],
]
)
)
def testConvFollowedByEinsum(self, conv_op, input_size, kernel_size,
einsum_eq, strides, eager_mode, input_sharding,
padding):
x_in = constant_op.constant(
np.random.random(size=input_size), dtype=dtypes.float32
)
kernel_in = constant_op.constant(
np.random.random(size=kernel_size), dtype=dtypes.float32
)
weight = constant_op.constant(
np.random.random(size=(2, 2)), dtype=dtypes.float32
)
def conv_fn(inputs, img_kernel, layer_weights):
output = conv_op(inputs, img_kernel, strides=strides, padding=padding)
output = special_math_ops.einsum(einsum_eq, output, layer_weights)
return output
if not eager_mode:
conv_fn = polymorphic_function.function(conv_fn)
golden_result = conv_fn(x_in, kernel_in, weight)
if input_sharding == 'replicated':
input_layout = Layout.replicated(self.mesh, len(input_size))
output_layout = self.replicated_2d
elif input_sharding == 'batch_sharded':
input_layout = Layout.batch_sharded(self.mesh, BATCH_DIM, len(input_size))
output_layout = self.batch_sharded_2d
kernel_layout = Layout.replicated(self.mesh, len(kernel_size))
d_x_in = numpy_util.pack_numpy(x_in, input_layout)
d_kernel_in = numpy_util.pack_numpy(kernel_in, kernel_layout)
d_weight = numpy_util.pack_numpy(weight, self.replicated_2d)
d_result = conv_fn(d_x_in, d_kernel_in, d_weight)
self.assertDTensorEqual(golden_result, output_layout, d_result)
@parameterized.named_parameters(
test_util.product(
*[
[
(
'Conv2D',
nn_ops.conv2d_v2,
(BATCH_SIZE, HEIGHT, WIDTH, CHANNEL_IN),
(2, 2, CHANNEL_IN, CHANNEL_OUT),
'bhwc,xy->by',
[1, 1, 1, 1],
),
(
'Conv3D',
nn_ops.conv3d_v2,
(BATCH_SIZE, DEPTH, HEIGHT, WIDTH, CHANNEL_IN),
(2, 2, 2, CHANNEL_IN, CHANNEL_OUT),
'bdhwc,xy->by',
[1, 1, 1, 1, 1],
),
],
[
('ReplicatedInput', 'replicated'),
('BatchShardedInput', 'batch_sharded'),
],
[
('ValidPadding', 'VALID'),
('SamePadding', 'SAME'),
],
]
)
)
def testConvFollowedByEinsumWithGradient(self, conv_op, input_size,
kernel_size, einsum_eq, strides,
input_sharding, padding):
x_in = constant_op.constant(
np.random.random(size=input_size), dtype=dtypes.float32
)
kernel_in = constant_op.constant(
np.random.random(size=kernel_size), dtype=dtypes.float32
)
weight = constant_op.constant(
np.random.random(size=(2, 2)), dtype=dtypes.float32
)
@polymorphic_function.function
def conv_fn(inputs, img_kernel, layer_weights):
with backprop.GradientTape() as tape:
tape.watch([inputs, img_kernel, layer_weights])
output = conv_op(inputs, img_kernel, strides=strides, padding=padding)
output = special_math_ops.einsum(einsum_eq, output, layer_weights)
inputs_grad, kernel_grad, weight_grad = tape.gradient(
output, [inputs, img_kernel, layer_weights])
return output, inputs_grad, kernel_grad, weight_grad
result, inputs_grad, kernel_grad, weight_grad = conv_fn(
x_in, kernel_in, weight)
if input_sharding == 'replicated':
input_layout = Layout.replicated(self.mesh, len(input_size))
output_layout = self.replicated_2d
elif input_sharding == 'batch_sharded':
input_layout = Layout.batch_sharded(self.mesh, BATCH_DIM, len(input_size))
output_layout = self.batch_sharded_2d
kernel_layout = Layout.replicated(self.mesh, len(kernel_size))
d_x_in = numpy_util.pack_numpy(x_in, input_layout)
d_kernel_in = numpy_util.pack_numpy(kernel_in, kernel_layout)
d_weight = numpy_util.pack_numpy(weight, self.replicated_2d)
d_result, d_inputs_grad, d_kernel_grad, d_weight_grad = conv_fn(
d_x_in, d_kernel_in, d_weight)
self.assertDTensorEqual(result, output_layout, d_result)
# TODO(b/208700444): layout of input grads should match layout of input.
self.assertDTensorEqual(
inputs_grad,
Layout.replicated(self.mesh, len(input_size)),
d_inputs_grad,
)
self.assertDTensorEqual(kernel_grad, kernel_layout, d_kernel_grad)
self.assertDTensorEqual(weight_grad, self.replicated_2d, d_weight_grad)
SPATIALLY_PARTITIONED_CONV_TEST_CASES = [
[
('Case1', (BATCH_SIZE, 8, 16, CHANNEL_IN), (3, 5, CHANNEL_IN,
CHANNEL_OUT)),
('Case2', (BATCH_SIZE, 8, 128, CHANNEL_IN), (3, 9, CHANNEL_IN,
CHANNEL_OUT)),
],
[
('ValidPadding', 'VALID'),
('SamePadding', 'SAME'),
],
[
('Batch_1d_2x4', [BATCH_DIM, UNSHARDED, WIDTH_DIM, UNSHARDED], (2, 4)),
('2d_2x4', [UNSHARDED, HEIGHT_DIM, WIDTH_DIM, UNSHARDED], (2, 4)),
('Batch_2d_2x2x2', [BATCH_DIM, HEIGHT_DIM, WIDTH_DIM,
UNSHARDED], (2, 2, 2)),
],
]
class SpatiallyPartitionedConvOpTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
# TODO(b/261485237): Enable CPU testing once CollectivePermute is supported
# on CPU's.
if not test_util.is_tpu_present():
self.skipTest('This test only runs on TPUs.')
def _create_mesh(self, mesh_dims, topology):
global_ids = test_util.create_device_ids_array(topology)
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {}
for device in ('CPU', 'GPU', 'TPU'):
mesh_dict[device] = Mesh(
mesh_dims,
global_ids,
local_ids,
test_util.create_device_list(topology, device),
)
return self.configTestMesh(mesh_dict)
@parameterized.named_parameters(
test_util.product(*SPATIALLY_PARTITIONED_CONV_TEST_CASES))
def testConv(self, input_shape, kernel_shape, padding, sharding_specs,
topology):
mesh_dims = [spec for spec in sharding_specs if spec != UNSHARDED]
mesh = self._create_mesh(mesh_dims, topology)
x_in = constant_op.constant(
np.random.random(size=input_shape), dtype=dtypes.float32
)
kernel_in = constant_op.constant(
np.random.random(size=kernel_shape), dtype=dtypes.float32
)
expected_output = nn_ops.conv2d_v2(
x_in, kernel_in, strides=[1, 1, 1, 1], padding=padding
)
input_layout = Layout(sharding_specs, mesh)
kernel_layout = Layout.replicated(mesh, 4)
d_x_in = numpy_util.pack_numpy(x_in, input_layout)
d_kernel_in = numpy_util.pack_numpy(kernel_in, kernel_layout)
d_output = nn_ops.conv2d_v2(
d_x_in, d_kernel_in, strides=[1, 1, 1, 1], padding=padding
)
self.assertDTensorEqual(expected_output, input_layout, d_output)
@parameterized.named_parameters(
test_util.product(*SPATIALLY_PARTITIONED_CONV_TEST_CASES))
def testConvWithGradient(self, input_shape, kernel_shape, padding,
sharding_specs, topology):
# TODO(b/208700444): add support for SPMD expansion of spatially partitioned
# conv backprop.
self.skipTest(
'b/208700444: Spatially partitioned conv backprop not implemented.')
mesh_dims = [spec for spec in sharding_specs if spec != UNSHARDED]
mesh = self._create_mesh(mesh_dims, topology)
x_in = constant_op.constant(
np.random.random(size=input_shape), dtype=dtypes.float32
)
kernel_in = constant_op.constant(
np.random.random(size=kernel_shape), dtype=dtypes.float32
)
@polymorphic_function.function
def conv_fn(inputs, img_kernel, padding):
with backprop.GradientTape() as tape:
tape.watch([inputs, img_kernel])
output = nn_ops.conv2d_v2(
inputs, img_kernel, strides=[1, 1, 1, 1], padding=padding
)
inputs_grad, kernel_grad = tape.gradient(output, [inputs, img_kernel])
return output, inputs_grad, kernel_grad
expected_output, expected_inputs_grad, expected_kernel_grad = conv_fn(
x_in, kernel_in, padding)
input_layout = Layout(sharding_specs, mesh)
kernel_layout = Layout.replicated(mesh, 4)
d_x_in = numpy_util.pack_numpy(x_in, input_layout)
d_kernel_in = numpy_util.pack_numpy(kernel_in, kernel_layout)
d_output, d_inputs_grad, d_kernel_grad = conv_fn(d_x_in, d_kernel_in,
padding)
self.assertDTensorEqual(expected_output, input_layout, d_output)
self.assertDTensorEqual(expected_inputs_grad, input_layout, d_inputs_grad)
self.assertDTensorEqual(expected_kernel_grad, kernel_layout, d_kernel_grad)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/device_test.py
+798
View File
@@ -0,0 +1,798 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for DTensorDevice in python."""
import os
import threading
from unittest import mock
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import dtensor_device
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_collective_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
# Convenient constants to use for tests.
_BATCH_DIM = "batch"
_MESH_DIM_X = "x"
# Shorter notation
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
UNSHARDED = layout_lib.UNSHARDED
class DTensorDeviceTest(test_util.DTensorBaseTest, parameterized.TestCase):
def setUp(self):
super(DTensorDeviceTest, self).setUp()
device_ids = test_util.create_device_ids_array((2,))
local_device_ids = np.ravel(device_ids).tolist()
mesh_dict = { # pylint: disable=g-complex-comprehension
device: Mesh(
[_BATCH_DIM],
device_ids,
local_device_ids,
test_util.create_device_list((2,), device),
)
for device in ("CPU", "GPU", "TPU")
}
self.mesh = self.configTestMesh(mesh_dict)
def testInvalidLayout(self):
a = api.copy_to_mesh(
constant_op.constant([1.0]), Layout.replicated(self.mesh, rank=1)
)
b = array_ops.identity(a)
with self.assertRaises(ValueError):
api.check_layout(b, Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=1))
@parameterized.parameters(True, False)
def testAsyncOption(self, is_async):
try:
# There isn't a great way to test whether something actually executed
# synchronously; this test just exercises the option.
api.reset_dtensor_device(is_async=is_async)
with api._dtensor_device()._experimental_default_mesh(self.mesh):
with ops.device_v2(api.device_name()):
a = api.copy_to_mesh(
constant_op.constant([1.0]), Layout.replicated(self.mesh, rank=1)
)
b = array_ops.identity(a)
self.assertEqual([1.0], b.numpy())
finally:
api._reset() # pylint: disable=protected-access
def testBasicTypeBasedDispatch(self):
# Tests for b = Op(a).
a = constant_op.constant([1.0, 2.0, 3.0, 4.0])
a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=1))
# __getitem__
b = a[2:-2]
api.check_layout(b, Layout.replicated(self.mesh, rank=1))
c = a * 2
api.check_layout(b, Layout.replicated(self.mesh, rank=1))
self.assertAllEqual(a.numpy(), [1., 2., 3., 4.])
self.assertAllEqual(c.numpy(), [2., 4., 6., 8.])
@mock.patch.dict(os.environ, {"DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION": "1"})
def testMultiDeviceExpansion(self):
# Tests for b = Op(a).
a = constant_op.constant([1.0, 2.0, 3.0, 4.0])
a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=1))
# __getitem__
b = a[2:-2]
api.check_layout(b, Layout.replicated(self.mesh, rank=1))
c = a * 2
api.check_layout(b, Layout.replicated(self.mesh, rank=1))
self.assertAllEqual(a.numpy(), [1.0, 2.0, 3.0, 4.0])
self.assertAllEqual(c.numpy(), [2.0, 4.0, 6.0, 8.0])
def testNoImplicitCopyOnForLargeIntegerTensors(self):
a = array_ops.ones([10, 10], dtype=dtypes.int32)
a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=2))
big = array_ops.ones([10, 10], dtype=dtypes.int32)
small = array_ops.ones([10], dtype=dtypes.int32)
with self.assertRaises(errors_impl.UnimplementedError):
a + big # pylint:disable=pointless-statement
a + small # pylint:disable=pointless-statement
def testConcurrentExecute(self):
results = {}
def func(thread_id):
@polymorphic_function.function
def update_variable(initial_value, num_round):
y = math_ops.multiply(initial_value, num_round)
return math_ops.add(initial_value, y)
for n in range(10):
with api._dtensor_device()._experimental_default_mesh(self.mesh):
x = stateless_random_ops.stateless_random_uniform(
[10], seed=(1, 2), minval=0, maxval=255
)
y = api.copy_to_mesh(x, Layout.replicated(self.mesh, rank=1))
y = update_variable(y, n + 1)
results[thread_id] = y
threads = {}
for a in range(10):
t = threading.Thread(target=func, args=(a,))
threads[a] = t
t.start()
for thrad_id, thread in threads.items():
thread.join()
self.assertIsNotNone(results[thrad_id])
def testNoImplicitCopyOnForScalarVariableOnNonCPUMesh(self):
self.skipForTfrt("b/235088250")
self.skipForDeviceType(["CPU"], "CPU mesh implicit copy is allowed.")
init_value = api.call_with_layout(
array_ops.ones, shape=(1), layout=Layout.replicated(self.mesh, rank=1)
)
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
r"Using a non-DTensor variable with DTensor is only supported for..*\n"
r".*Shape: \[1\].*\n"
".*device_test.py.*",
):
api.copy_to_mesh(
variables.Variable(init_value), Layout.replicated(self.mesh, rank=1)
)
@parameterized.named_parameters(
test_util.product(
[
("Int32", dtypes.int32),
("Float32", dtypes.float32),
("Int64", dtypes.int64),
("Float64", dtypes.float64),
],
[
(
"Scalar",
[],
),
(
"RankOne",
[1],
),
("RankTwo", [2, 2]),
],
)
)
def testImplicitCopyVariableOnCPUMesh(self, dtype, shape):
self.skipForTfrt("b/235088250")
self.skipForDeviceType(
["GPU", "TPU"], "Variable implicit copy is only allowed for CPU mesh.")
variable = variables.Variable(array_ops.ones(shape=shape, dtype=dtype))
new_value = array_ops.zeros(shape=shape, dtype=dtype)
@polymorphic_function.function
def assign_function(v, new_value):
return v.assign(new_value)
layout = Layout.replicated(self.mesh, rank=len(shape))
# Run explicitly on the dtensor device with a default mesh since
# we do not have any registered mesh to broadcast the inputs to.
with api.default_mesh(self.mesh):
assign_function(variable, api.pack([new_value] * self.mesh.size, layout))
read_value = variable.read_value()
self.assertDTensorEqual(new_value, layout, read_value)
def testNumpyCallWithReplicatedInput(self):
a = constant_op.constant([1.0, 2.0, 3.0, 4.0])
a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=1))
b = a.numpy()
self.assertAllEqual(b, [1., 2., 3., 4.])
def testTensorIteration(self):
a = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=1))
iterator = iter(a)
self.assertAllClose(1., next(iterator))
def testCopyToMeshWithSameLayout(self):
a = constant_op.constant([1.0, 2.0, 3.0, 4.0])
a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=1))
a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=1))
api.check_layout(a, Layout.replicated(self.mesh, rank=1))
def testSetDefaultLayoutEager(self):
tensor = constant_op.constant([[1.0], [1.0]])
tensor = api.copy_to_mesh(tensor, Layout.replicated(self.mesh, rank=2))
with api._dtensor_device()._default_layout(
Layout.replicated(self.mesh, rank=1)):
tensor = array_ops.reshape(tensor, [-1])
api.check_layout(tensor, Layout.replicated(self.mesh, rank=1))
self.assertAllClose([1., 1.], tensor.numpy())
def testSetDefaultLayoutFunction(self):
@polymorphic_function.function
def func():
tensor = constant_op.constant([[1.0], [1.0]])
return array_ops.reshape(tensor, [-1]), array_ops.reshape(tensor, [-1])
with api._dtensor_device()._default_layout(
Layout.batch_sharded(self.mesh, batch_dim=_BATCH_DIM, rank=1)
):
tensor1, tensor2 = func()
api.check_layout(
tensor1, Layout.batch_sharded(self.mesh, batch_dim=_BATCH_DIM, rank=1)
)
api.check_layout(tensor2, Layout.replicated(self.mesh, rank=1))
tensor1 = api.relayout(tensor1, Layout.replicated(self.mesh, rank=1))
self.assertAllClose([1.0, 1.0], tensor1.numpy())
self.assertAllClose([1.0, 1.0], tensor2.numpy())
@parameterized.named_parameters(
# pylint: disable=unnecessary-lambda
# Needed for the DVariable monkey patch to work.
("Variable", lambda x: d_variable.DVariable(x)),
# pylint: enable=unnecessary-lambda
("Tensor", lambda x: x),
)
def testStringRepresentation(self, transform):
replicated = api.copy_to_mesh(
constant_op.constant(8.0), Layout.replicated(self.mesh, rank=0)
)
replicated = transform(replicated)
replicated_str = str(replicated)
self.assertIn("8", replicated_str)
self.assertIn("layout", replicated_str)
sharded = api.pack(
[constant_op.constant([8.0]), constant_op.constant([9.0])],
layout=Layout([_BATCH_DIM], self.mesh),
)
sharded = transform(sharded)
sharded_str = str(sharded)
self.assertIn("8", sharded_str)
self.assertIn("9", sharded_str)
self.assertIn("layout", sharded_str)
@parameterized.named_parameters(("Async", True), ("Sync", False))
def testCancellation(self, is_async):
self.skipForTfrt("b/181368626: support cancellation in tfrt.")
self.skipForDeviceType(["TPU"], "b/195552283: Fix cancellation on TPU.")
device = dtensor_device.DTensorDevice(meshes=[self.mesh], is_async=is_async)
@polymorphic_function.function
def f(x):
# Integer division by 0 on one device, which returns a bad status.
x = math_ops.cast(gen_math_ops.div(x=x, y=x), dtypes.float32)
# A reduction requiring a collective, which would normally deadlock with
# one of its participants missing.
return math_ops.reduce_sum(x, axis=0)
a = constant_op.constant([[1, 2]])
b = constant_op.constant([[0, 1]])
x = device.pack([a, b], layout=Layout([_BATCH_DIM, UNSHARDED], self.mesh))
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError, "Integer division by zero"
):
y = f(x)
y.numpy()
z = array_ops.identity(x)
self.assertAllClose([[[1, 2]], [[0, 1]]], device.unpack(z))
def testCopyToMeshShapeFn(self):
@polymorphic_function.function
def f():
c = constant_op.constant([1.0, 2.0])
on_mesh = api.copy_to_mesh(c, Layout.replicated(self.mesh, rank=1))
return on_mesh
output, = f.get_concrete_function().outputs
self.assertEqual([2], output.shape)
def testUnpackInvalidInput(self):
# Test for b/255629824
with self.assertRaisesRegex(TypeError, "Expecting a Tensor"):
api.unpack(
**{
"tensor": [[
41.8684053521925,
731.610023060566,
356.0701500440248,
9.62928117100512,
185.0041559439026,
225.87663065861508,
450.2403652750002,
268.7273627027147,
]]
}
)
def testIsDTensorInvalidInput(self):
# Test for b/272381211
self.assertFalse(api.fetch_layout(**{"tensor": -1024}))
def testFetchLayoutInvalidInput(self):
# Test for b/272381211
self.assertIsNone(api.fetch_layout(**{"tensor": -1024}))
def testFetchLayoutForDVariablesReturnsCorrectLayout(self):
layout = Layout.replicated(self.mesh, 2)
with api._dtensor_device()._experimental_default_mesh(self.mesh):
dvariable = d_variable.DVariable(
api.call_with_layout(
array_ops.ones, shape=[2, 3], dtype=dtypes.float32, layout=layout
)
)
self.assertEqual(layout, api.fetch_layout(dvariable))
def testFetchLayoutForDTensorReturnsCorrectLayout(self):
layout = Layout.replicated(self.mesh, 2)
tensor = api.call_with_layout(
array_ops.ones, shape=[2, 3], dtype=dtypes.float32, layout=layout
)
self.assertEqual(layout, api.fetch_layout(tensor))
def testFetchLayoutForRegularTensorsThrowsError(self):
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
"FetchLayout expects a tensor placed on the layout device.",
):
api.fetch_layout(constant_op.constant([2, 3]))
def testFetchLayoutNotEagerlyRaisesRuntimeError(self):
@polymorphic_function.function
def f(dtensor_input):
api.fetch_layout(dtensor_input)
with self.assertRaisesRegex(RuntimeError,
"`fetch_layout` must be called eagerly."):
f(
api.copy_to_mesh(
constant_op.constant(1.0), Layout.replicated(self.mesh, rank=0)
)
)
def testIsDTensor(self):
normal_tensor = array_ops.zeros(shape=[10, 10])
self.assertFalse(api.is_dtensor(normal_tensor))
layout = Layout.replicated(self.mesh, rank=1)
d_tensor = api.call_with_layout(array_ops.zeros, layout=layout, shape=[10])
self.assertTrue(api.is_dtensor(d_tensor))
var = d_variable.DVariable(d_tensor)
self.assertTrue(api.is_dtensor(var))
self.assertFalse(api.is_dtensor([0, 1]))
self.assertFalse(api.is_dtensor({False: True}))
self.assertFalse(api.is_dtensor(1))
class C:
pass
self.assertFalse(api.is_dtensor(C()))
def testIsDTensorNotEagerlyRaisesRuntimeError(self):
@polymorphic_function.function
def f(dtensor_input):
api.is_dtensor(dtensor_input)
with self.assertRaisesRegex(
RuntimeError, "`is_dtensor` must be called eagerly."):
f(
api.copy_to_mesh(
constant_op.constant(1.0), Layout.replicated(self.mesh, 0)
)
)
def testSingleDeviceMesh(self):
# FIXME(b/274647196): Add a mesh_util API that takes CPU:0.
cpu0_mesh = Mesh.from_device("/job:localhost/replica:0/task:0/device:CPU:0")
with api.default_mesh(cpu0_mesh):
a = constant_op.constant(1.0)
self.assertFalse(api.is_dtensor(a))
self.assertIn("CPU:0", a.device)
with api.default_mesh(cpu0_mesh):
b = array_ops.ones(shape=(3, 3))
self.assertTrue(api.is_dtensor(b))
self.assertEqual(api.fetch_layout(b).mesh, cpu0_mesh)
def testUnsupportedOpReplicatedInput(self):
with api.default_mesh(self.mesh):
t = array_ops.ones(shape=(8, 3))
a = gen_collective_ops.collective_reduce_v2(
t,
group_size=1,
group_key=1030,
instance_key=1,
merge_op="Add",
final_op="Id",
ordering_token=[],
)
self.assertFalse(api.is_dtensor(a))
self.assertAllClose(a.numpy(), t.numpy())
def testUnsupportedOpShardedInput(self):
with api.default_mesh(self.mesh):
t = array_ops.ones(shape=(8, 3))
t = api.relayout(
t, Layout.batch_sharded(mesh=self.mesh, batch_dim=_BATCH_DIM, rank=2)
)
with self.assertRaisesRegex(
errors_impl.UnimplementedError, "not supported"
):
# This is an Op that we don't have a SPMD expander.
gen_collective_ops.collective_reduce_v2(
t,
group_size=1,
group_key=1030,
instance_key=1,
merge_op="Add",
final_op="Id",
ordering_token=[],
)
class DTensorPackUnpackOnOneDMeshTest(test_util.DTensorBaseTest):
def setUp(self):
super(DTensorPackUnpackOnOneDMeshTest, self).setUp()
global_ids = test_util.create_device_ids_array((2,))
local_device_ids = np.ravel(global_ids).tolist()
mesh_dict = { # pylint: disable=g-complex-comprehension
device: Mesh(
[_BATCH_DIM],
global_ids,
local_device_ids,
test_util.create_device_list((2,), device),
)
for device in ("CPU", "GPU", "TPU")
}
self.mesh = self.configTestMesh(mesh_dict)
def testUnpack(self):
with api.default_mesh(self.mesh):
v = constant_op.constant(1.0)
v = api.copy_to_mesh(v, Layout.replicated(self.mesh, rank=0))
self.assertAllClose([1.0, 1.0], api.unpack(v))
def testUnpackVariables(self):
v0 = d_variable.DVariable(
api.call_with_layout(
array_ops.ones,
shape=[2, 3],
dtype=dtypes.float32,
layout=Layout.replicated(self.mesh, 2),
)
)
with self.assertRaisesRegex(TypeError, "Expecting a Tensor"):
api._dtensor_device().unpack(v0)
def testUnpackingRegularTensorRaisesInvalidArgumentError(self):
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
"DTensorUnpack expects a tensor placed on the DTensor device",
):
api._dtensor_device().unpack(constant_op.constant([1.0, 2.0]))
def testUnpackingNotEagerlyRaisesRuntimeError(self):
@polymorphic_function.function
def f(dtensor_input):
api._dtensor_device().unpack(dtensor_input)
with self.assertRaisesRegex(
RuntimeError, "`unpack` must be called eagerly."):
f(
api.copy_to_mesh(
constant_op.constant(1.0), Layout.replicated(self.mesh, rank=0)
)
)
def testPack(self):
a = constant_op.constant([1.0, 2.0])
b = constant_op.constant([3.0, 4.0])
with ops.device_v2(api.device_name()):
packed_tensor = api.pack(
[a, b], layout=Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=1)
)
api.check_layout(
packed_tensor, Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=1)
)
self.assertAllEqual([
4,
], packed_tensor.shape)
unpacked_tensor = api.unpack(packed_tensor)
self.assertAllClose([1., 2.], unpacked_tensor[0])
self.assertAllClose([3., 4.], unpacked_tensor[1])
def testPackingNotEagerlyRaisesRuntimeError(self):
@polymorphic_function.function
def f(a):
api.pack([a, a], layout=Layout.replicated(self.mesh, rank=1))
with self.assertRaisesRegex(RuntimeError, "`pack` must be called eagerly."):
f(constant_op.constant([1.0]))
def testPackingVariablesRaisesError(self):
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError, "Variable input is not supported."
):
api._dtensor_device().pack(
[
variables.Variable(array_ops.ones([2, 3])),
variables.Variable(array_ops.ones([2, 3])),
],
Layout.replicated(self.mesh, rank=2),
)
def testPackDevice(self):
a = constant_op.constant([1.0, 2.0])
b = constant_op.constant([3.0, 4.0])
with ops.device_v2(api.device_name()):
packed_tensor = api.pack(
[a, b], layout=Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=1)
)
unpacked_tensor = api.unpack(packed_tensor)
self.assertAllEqual(self.mesh.local_devices(),
[t.device for t in unpacked_tensor])
def testPackScalar(self):
a = constant_op.constant(1.0)
with ops.device_v2(api.device_name()):
packed_layout = Layout([], self.mesh)
packed_tensor = api.pack([a, a], layout=packed_layout)
api.check_layout(packed_tensor, packed_layout)
self.assertAllEqual([], packed_tensor.shape)
unpacked_tensor = api.unpack(packed_tensor)
self.assertAllClose([a, a], unpacked_tensor)
def testPackHigherRankValue(self):
# Pack a rank 3 matrix into a 1d mesh.
a = constant_op.constant(
[[[1, 2, 3], [4, 5, 6]], [[2, 3, 4], [5, 6, 7]]]
) # 2x2x3
b = constant_op.constant(
[[[3, 2, 1], [6, 5, 4]], [[4, 3, 2], [7, 6, 5]]]
) # 2x2x3
pack_layout = Layout([_BATCH_DIM, UNSHARDED, UNSHARDED], self.mesh)
with ops.device_v2(api.device_name()):
# pack to 4x2x3
packed_tensor = api.pack([a, b], layout=pack_layout)
api.check_layout(packed_tensor, pack_layout)
self.assertAllEqual([4, 2, 3], packed_tensor.shape)
class DTensorPackUnpackOnTwoDMeshTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
self.skipForDeviceType(["TPU"],
"all tests require 8 TPU cores.",
unless_device_count_equals_to=8)
global_ids = test_util.create_device_ids_array((2, 4))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = { # pylint: disable=g-complex-comprehension
device: Mesh(
[_BATCH_DIM, _MESH_DIM_X],
global_ids,
local_ids,
test_util.create_device_list((2, 4), device),
)
for device in ("CPU", "GPU", "TPU")
}
self.mesh = self.configTestMesh(mesh_dict)
def testPackWithScalars(self):
a = constant_op.constant(1.23)
with ops.device_v2(api.device_name()):
packed_layout = Layout([], self.mesh)
packed_tensor = api.pack([a, a, a, a, a, a, a, a], layout=packed_layout)
api.check_layout(packed_tensor, packed_layout)
self.assertAllEqual([], packed_tensor.shape)
unpacked_tensor = api.unpack(packed_tensor)
self.assertAllClose([a, a, a, a, a, a, a, a], unpacked_tensor)
def testPackWithScalarsWithInvalidRank(self):
a = constant_op.constant(1.23)
with ops.device_v2(api.device_name()):
invalid_packed_layout = Layout([UNSHARDED], self.mesh)
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
"Packed layout should have the same rank",
):
_ = api.pack([a, a, a, a, a, a, a, a], layout=invalid_packed_layout)
def testPackWithBatchSharding(self):
a = constant_op.constant([[1.0], [2.0]])
b = constant_op.constant([[3.0], [4.0]])
with ops.device_v2(api.device_name()):
packed_tensor = api.pack(
[a, a, a, a, b, b, b, b],
layout=Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2),
)
api.check_layout(
packed_tensor, Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
)
self.assertAllEqual([4, 1], packed_tensor.shape)
unpacked_tensor = api.unpack(packed_tensor)
self.assertAllClose([a, a, a, a, b, b, b, b], unpacked_tensor)
def testPackWithFullReplicated(self):
a = constant_op.constant([[1.0], [2.0]])
with ops.device_v2(api.device_name()):
packed_layout = Layout([UNSHARDED, UNSHARDED], self.mesh)
packed_tensor = api.pack([a, a, a, a, a, a, a, a], layout=packed_layout)
api.check_layout(packed_tensor, packed_layout)
self.assertAllEqual([2, 1], packed_tensor.shape)
unpacked_tensor = api.unpack(packed_tensor)
self.assertAllClose([a, a, a, a, a, a, a, a], unpacked_tensor)
def testFillUsesSpecifiedLayout(self):
with api.default_mesh(self.mesh):
# TODO(allenl): Figure out why the embedded constant (triggered for
# dtypes.int32) gets a sharded layout by default.
dims = constant_op.constant([4, 1], dtype=dtypes.int64)
value = constant_op.constant(1.0)
with api._dtensor_device()._default_layout(
Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)):
filled = array_ops.fill(value=value, dims=dims)
api.check_layout(
filled, Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
)
unpacked_tensor = api.unpack(filled)
self.assertAllClose(8 * [[[1.], [1.]]], unpacked_tensor)
self.assertEqual([4, 1], filled.shape)
class DTensorSparse(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
self.skipForDeviceType(["TPU"],
"all tests require 8 TPU cores.",
unless_device_count_equals_to=8)
global_ids = test_util.create_device_ids_array((2, 4))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = { # pylint: disable=g-complex-comprehension
device: Mesh(
[_BATCH_DIM, _MESH_DIM_X],
global_ids,
local_ids,
test_util.create_device_list((2, 4), device),
)
for device in ("CPU", "GPU", "TPU")
}
self.mesh = self.configTestMesh(mesh_dict)
@parameterized.named_parameters(
test_util.product([("Replicated", "replicated"), ("Sharded", "batch")], [(
"RankTwo",
[2, 4],
), (
"RankThree",
[2, 2, 3],
)]))
def testPackUnpackReturnsCorrectValuesAndDevices(self, sharding, shape):
a = sparse_ops.from_dense(
stateless_random_ops.stateless_random_uniform(shape, seed=[0, 1])
)
b = sparse_ops.from_dense(
stateless_random_ops.stateless_random_uniform(shape, seed=[0, 1])
)
if sharding == "replicated":
layout = Layout(["unsharded"] * len(shape), self.mesh)
input_tensor = 8 * [a]
expected = 8 * [sparse_ops.sparse_tensor_to_dense(a)]
expected_shape = shape
else:
layout = Layout([_BATCH_DIM] + ["unsharded"] * (len(shape) - 1),
self.mesh)
input_tensor = 4 * [a] + 4 * [b]
expected = 4 * [sparse_ops.sparse_tensor_to_dense(a)] + 4 * [
sparse_ops.sparse_tensor_to_dense(b)
]
expected_shape = [shape[0] * 2] + shape[1:]
with ops.device_v2(api._dtensor_device().name):
packed_tensor = api.pack(input_tensor, layout)
api.check_layout(packed_tensor, layout)
unpacked_tensor = api.unpack(packed_tensor)
got = [sparse_ops.sparse_tensor_to_dense(t) for t in unpacked_tensor]
# Check shape of packed tensor.
self.assertAllEqual(expected_shape, packed_tensor.shape)
# Check values.
self.assertAllClose(expected, got)
# Check devices.
self.assertAllEqual(self.mesh.local_devices(),
[t.indices.device for t in unpacked_tensor])
self.assertAllEqual(self.mesh.local_devices(),
[t.values.device for t in unpacked_tensor])
def testPackingMixedTensorTypesRaisesTypeError(self):
tensor = stateless_random_ops.stateless_random_uniform([2, 4], seed=[0, 1])
sparse_tensor = sparse_ops.from_dense(tensor)
with ops.device_v2(api.device_name()):
with self.assertRaisesRegex(TypeError,
"Cannot Pack SparseTensors with Tensors."):
api.pack(
4 * [tensor] + 4 * [sparse_tensor],
Layout.replicated(self.mesh, rank=2),
)
def testPackingTensorsWithDifferentShapesRaisesTypeError(self):
a = sparse_ops.from_dense(
stateless_random_ops.stateless_random_uniform([2, 2], seed=[0, 1])
)
b = sparse_ops.from_dense(
stateless_random_ops.stateless_random_uniform([4, 4], seed=[0, 1])
)
with ops.device_v2(api.device_name()):
with self.assertRaisesRegex(
TypeError, "All input SparseTensors to Pack must be same shape."):
api.pack(4 * [a] + 4 * [b], Layout.replicated(self.mesh, rank=2))
def testPackingSparseTensorsReturnsCorrectLayout(self):
layout = Layout.replicated(self.mesh, 2)
a = sparse_ops.from_dense(
stateless_random_ops.stateless_random_uniform([16, 16], seed=[0, 1])
)
with ops.device_v2(api.device_name()):
api.check_layout(api.pack(8 * [a], layout), layout)
if __name__ == "__main__":
test.main()
tensorflow/dtensor/python/tests/input_util_test.py
+660
View File
@@ -0,0 +1,660 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for DTensor input pipeline utilities."""
import contextlib
import threading
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import input_util
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import mesh_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test as tf_test
MESH_DIM_BATCH = 'batch'
MESH_DIM_HEIGHT = 'height'
MESH_DIM_WIDTH = 'width'
MESH_SIZE_BATCH = 4
MESH_SIZE_HEIGHT = 2
MESH_SIZE_WIDTH = 2
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
UNSHARDED = layout_lib.UNSHARDED
class DTensorDatasetTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
self._num_devices = MESH_SIZE_BATCH * MESH_SIZE_HEIGHT * MESH_SIZE_WIDTH
self.mesh = mesh_util.create_mesh(
devices=['CPU:%d' % i for i in range(self._num_devices)],
mesh_dims=[(MESH_DIM_BATCH, MESH_SIZE_BATCH),
(MESH_DIM_HEIGHT, MESH_SIZE_HEIGHT),
(MESH_DIM_WIDTH, MESH_SIZE_WIDTH)])
self.mesh = self.configTestMesh({'CPU': self.mesh})
self.images = self._images([8, 8, 3])
self.labels = self._labels([1])
def _images(self, shape):
return stateless_random_ops.stateless_random_uniform(
shape, seed=(1, 2), minval=0, maxval=255)
def _labels(self, shape):
return stateless_random_ops.stateless_random_uniform(
shape, seed=(1, 2), minval=0, maxval=10, dtype=dtypes.float32)
def testIterableFailsWithUnknownShapeDatasetSpec(self):
def gen():
yield constant_op.constant([1, 2], dtype=dtypes.int32)
dataset = dataset_ops.DatasetV2.from_generator(
gen,
output_signature=tensor_spec.TensorSpec(
tensor_shape.TensorShape(None), dtype=dtypes.int32))
with self.assertRaisesRegex(
ValueError, 'Dataset element shape must have a valid rank'):
input_util.DTensorDataset(
dataset=dataset,
global_batch_size=8,
mesh=self.mesh,
layouts=Layout.replicated(self.mesh, rank=2))
def testIterMismatchedLayoutFails(self):
dataset = dataset_ops.DatasetV2.from_tensors(self.images).repeat()
# Mismatched rank-3 layout for rank-4 input (after batching)
images_layout = Layout(
[MESH_DIM_BATCH, MESH_DIM_HEIGHT, MESH_DIM_WIDTH], self.mesh)
with self.assertRaisesRegex(ValueError, 'Expected layout with rank 4'):
_ = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=32,
mesh=self.mesh,
layouts=images_layout,
batch_dim=MESH_DIM_BATCH)
@parameterized.named_parameters(('Eager', False), ('Graph', True))
def testRangeIteration(self, is_graph):
batch_size = 8
num_batches = 4
dataset = dataset_ops.DatasetV2.from_tensor_slices(
self._images([batch_size * num_batches, 8, 8, 3]))
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=self.mesh,
layouts=images_layout,
batch_dim=MESH_DIM_BATCH)
def train(iterator, steps):
iters = 1
output = next(iterator)
for _ in math_ops.range(steps - 1):
output += next(iterator)
iters += 1
if not is_graph:
mesh_util.barrier(self.mesh)
return output, iters
train_fn = polymorphic_function.function(train) if is_graph else train
exception = errors_impl.OutOfRangeError if is_graph else StopIteration
iterator = iter(dataset.batch(batch_size, drop_remainder=True))
output, iters = train_fn(iterator, num_batches)
d_iterator = iter(d_dataset)
d_output, d_iters = train_fn(d_iterator, num_batches)
mesh_util.barrier(self.mesh)
# Try one more iteration which will raise an exception since the iterator is
# exhausted.
with self.assertRaises(exception):
if is_graph:
# FIXME(b/285884302): This flakily raises error
# "Cannot add 'while_cond' function, because a different function"
# Since num_batches is changed to 1, it retriggers SPMD expansion.
# Recreating polymorphic function to avoid running into the error.
train_fn = polymorphic_function.function(train)
train_fn(d_iterator, 1)
# In the graph case, we need to wait for the executor to finish all async
# calls after invoking the tf.function to ensure any pending error is
# raised.
mesh_util.barrier(self.mesh)
self.assertEqual(iters, d_iters)
self.assertDTensorEqual(output, images_layout, d_output)
@parameterized.named_parameters(('Eager', False), ('Graph', True))
def testForInIteration(self, is_graph):
batch_size = 8
num_batches = 4
dataset = dataset_ops.DatasetV2.from_tensor_slices(
self._images([batch_size * num_batches, 8, 8, 3]))
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=self.mesh,
layouts=images_layout,
batch_dim=MESH_DIM_BATCH)
def train(iterator):
iters = 1
output = next(iterator)
for img in iterator:
output += img
iters += 1
if not is_graph:
mesh_util.barrier(self.mesh)
return output, iters
train_fn = polymorphic_function.function(train) if is_graph else train
iterator = iter(dataset.batch(batch_size, drop_remainder=True))
output, iters = train_fn(iterator)
d_iterator = iter(d_dataset)
d_output, d_iters = train_fn(d_iterator)
self.assertEqual(iters, d_iters)
self.assertDTensorEqual(output, images_layout, d_output)
@parameterized.named_parameters(('Eager', False), ('Graph', True))
def testIterSingleInput(self, is_graph):
dataset = dataset_ops.DatasetV2.from_tensors(self.images).repeat()
batch_size = 32
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=self.mesh,
layouts=images_layout,
batch_dim=MESH_DIM_BATCH)
self.assertEqual(d_dataset.element_spec.shape, [batch_size, 8, 8, 3])
def train(iterator):
it = next(iterator)
return it
train_fn = polymorphic_function.function(train) if is_graph else train
d_iterator = iter(d_dataset)
self.assertEqual(d_iterator.element_spec.shape, [batch_size, 8, 8, 3])
d_images = train_fn(d_iterator)
mesh_util.barrier(self.mesh)
expected = next(iter(dataset.batch(batch_size, drop_remainder=True)))
mesh_util.barrier(self.mesh)
self.assertDTensorEqual(expected, images_layout, d_images)
@parameterized.named_parameters(('Eager', False), ('Graph', True))
def testIterTupleInputs(self, is_graph):
dataset = dataset_ops.DatasetV2.from_tensors(
(self.images, self.labels)
).repeat()
batch_size = 32
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
labels_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=2)
layouts = (images_layout, labels_layout)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=self.mesh,
layouts=layouts,
batch_dim=MESH_DIM_BATCH)
def train(iterator):
return next(iterator)
train_fn = polymorphic_function.function(train) if is_graph else train
d_iterator = iter(d_dataset)
d_images, d_labels = train_fn(d_iterator)
expected_images, expected_labels = next(
iter(dataset.batch(batch_size, drop_remainder=True)))
self.assertDTensorEqual(expected_images, images_layout, d_images)
self.assertDTensorEqual(expected_labels, labels_layout, d_labels)
@parameterized.named_parameters(('Eager', False), ('Graph', True))
def testIterDictInputs(self, is_graph):
dataset = dataset_ops.DatasetV2.from_tensors({
'images': self.images,
'labels': self.labels,
}).repeat()
batch_size = 32
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
labels_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=2)
layouts = {'images': images_layout, 'labels': labels_layout}
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=self.mesh,
layouts=layouts,
batch_dim=MESH_DIM_BATCH)
def train(iterator):
return next(iterator)
train_fn = polymorphic_function.function(train) if is_graph else train
d_iterator = iter(d_dataset)
d_element = train_fn(d_iterator)
expected = next(iter(dataset.batch(batch_size, drop_remainder=True)))
self.assertDTensorEqual(expected['images'], images_layout,
d_element['images'])
self.assertDTensorEqual(expected['labels'], labels_layout,
d_element['labels'])
@parameterized.named_parameters(('Eager', False), ('Graph', True))
def testIterOnBatchedDataset(self, is_graph):
dataset = dataset_ops.DatasetV2.from_tensors({
'images': self.images,
'labels': self.labels,
}).repeat()
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
labels_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=2)
layouts = {'images': images_layout, 'labels': labels_layout}
global_batch_size = 32
per_replica_batch_size = global_batch_size // MESH_SIZE_BATCH
batched_dataset = dataset.batch(per_replica_batch_size, drop_remainder=True)
d_dataset = input_util.DTensorDataset(
dataset=batched_dataset,
global_batch_size=global_batch_size,
dataset_already_batched=True,
mesh=self.mesh,
layouts=layouts,
batch_dim=MESH_DIM_BATCH)
def train(iterator):
return next(iterator)
train_fn = polymorphic_function.function(train) if is_graph else train
d_iterator = iter(d_dataset)
d_element = train_fn(d_iterator)
expected = next(iter(dataset.batch(global_batch_size, drop_remainder=True)))
self.assertDTensorEqual(expected['images'], images_layout,
d_element['images'])
self.assertDTensorEqual(expected['labels'], labels_layout,
d_element['labels'])
def testIterOnBatchedDatasetFailsOnIncorrectBatchSize(self):
dataset = dataset_ops.DatasetV2.from_tensors({
'images': self.images,
'labels': self.labels,
}).repeat()
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
labels_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=2)
layouts = {'images': images_layout, 'labels': labels_layout}
global_batch_size = 32
per_replica_batch_size = 16 # correct value would be: 32 // 4 = 8
batched_dataset = dataset.batch(
per_replica_batch_size, drop_remainder=True)
with self.assertRaisesRegex(
ValueError,
('per_replica_batch_size does not matched expected size based on the '
'mesh, got 16 but expected 8.')):
_ = input_util.DTensorDataset(
dataset=batched_dataset,
global_batch_size=global_batch_size,
dataset_already_batched=True,
mesh=self.mesh,
layouts=layouts,
batch_dim=MESH_DIM_BATCH)
def testIterOnBatchedDatasetFailsNoDropLastBatch(self):
dataset = dataset_ops.DatasetV2.from_tensors({
'images': self.images,
'labels': self.labels,
}).repeat()
images_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=4)
labels_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=2)
layouts = {'images': images_layout, 'labels': labels_layout}
global_batch_size = 32
per_replica_batch_size = global_batch_size // MESH_SIZE_BATCH
batched_dataset = dataset.batch(
per_replica_batch_size, drop_remainder=False)
with self.assertRaisesRegex(
ValueError, 'Ensure drop_remainder=True when batching the dataset.'):
_ = input_util.DTensorDataset(
dataset=batched_dataset,
global_batch_size=global_batch_size,
dataset_already_batched=True,
mesh=self.mesh,
layouts=layouts,
batch_dim=MESH_DIM_BATCH)
@parameterized.named_parameters(('Disabled', False), ('Enabled', True))
def testIterPrefetch(self, prefetch):
condition = threading.Condition()
counter = variables.Variable(0)
def count(x):
counter.assign_add(1)
return x
num_batches = 8
batch_size = 4
total_elems = num_batches * batch_size
prefetch_buffer_size = 2 if prefetch else 0
inputs = np.arange(total_elems)
dataset = dataset_ops.DatasetV2.from_tensor_slices(inputs)
dataset = dataset.map(count)
inputs_layout = Layout.batch_sharded(
self.mesh, batch_dim=MESH_DIM_BATCH, rank=1)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=self.mesh,
layouts=inputs_layout,
batch_dim=MESH_DIM_BATCH,
prefetch=prefetch_buffer_size if prefetch else None)
# Check nothing was prefetched before iterators were created.
self.assertEqual(counter.numpy(), 0)
# Check nothing was prefetched before the first iteration.
d_iterator = iter(d_dataset)
self.assertEqual(counter.numpy(), 0)
# The number of elements that are expected to be fetched in each iteration.
multiple = batch_size * (self.mesh.size // MESH_SIZE_BATCH)
# Check the number of elements fetched upon for each batch.
for i in range(num_batches):
elem = next(d_iterator)
with condition:
count = min((i + prefetch_buffer_size) * multiple,
num_batches * multiple)
result = condition.wait_for(lambda: counter.numpy() >= count, timeout=5)
self.assertTrue(result)
start_idx, end_idx = i * batch_size, (i + 1) * batch_size
self.assertDTensorEqual(inputs[start_idx:end_idx], inputs_layout, elem)
@parameterized.product(
(
dict(
images_sharding=[UNSHARDED, UNSHARDED, UNSHARDED, UNSHARDED],
labels_sharding=[UNSHARDED, UNSHARDED],
),
dict(
images_sharding=[MESH_DIM_BATCH, UNSHARDED, UNSHARDED, UNSHARDED],
labels_sharding=[MESH_DIM_BATCH, UNSHARDED],
),
dict(
images_sharding=[
UNSHARDED,
MESH_DIM_HEIGHT,
MESH_DIM_WIDTH,
UNSHARDED,
],
labels_sharding=[UNSHARDED, UNSHARDED],
),
dict(
images_sharding=[
UNSHARDED,
MESH_DIM_WIDTH,
MESH_DIM_HEIGHT,
UNSHARDED,
],
labels_sharding=[UNSHARDED, UNSHARDED],
),
dict(
images_sharding=[
MESH_DIM_BATCH,
MESH_DIM_HEIGHT,
MESH_DIM_WIDTH,
UNSHARDED,
],
labels_sharding=[MESH_DIM_BATCH, UNSHARDED],
),
dict(
images_sharding=[
MESH_DIM_BATCH,
MESH_DIM_WIDTH,
MESH_DIM_HEIGHT,
UNSHARDED,
],
labels_sharding=[MESH_DIM_BATCH, UNSHARDED],
),
),
is_graph=[False, True],
through_dtensor=[False, True],
)
def testIterWithLayouts(
self, images_sharding, labels_sharding, is_graph, through_dtensor
):
if through_dtensor:
scope = api.default_mesh(self.mesh)
else:
scope = contextlib.nullcontext()
with scope:
batch_size = 32
dataset = dataset_ops.DatasetV2.from_tensors(
(self.images, self.labels)
).repeat()
batched_dataset = dataset.batch(batch_size, drop_remainder=True)
images_layout = Layout(images_sharding, self.mesh)
labels_layout = Layout(labels_sharding, self.mesh)
layouts = (images_layout, labels_layout)
batch_dim = None
if MESH_DIM_BATCH in images_sharding or MESH_DIM_BATCH in labels_sharding:
batch_dim = MESH_DIM_BATCH
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=self.mesh,
layouts=layouts,
batch_dim=batch_dim,
)
def train(iterator):
return next(iterator)
train_fn = polymorphic_function.function(train) if is_graph else train
d_iterator = iter(d_dataset)
d_images, d_labels = train_fn(d_iterator)
iterator = iter(batched_dataset)
images, labels = train_fn(iterator)
self.assertDTensorEqual(images, images_layout, d_images)
self.assertDTensorEqual(labels, labels_layout, d_labels)
def testMixedLayoutsFails(self):
dataset = dataset_ops.DatasetV2.from_tensors(
(self.images, self.labels)
).repeat()
images_layout = Layout(
[UNSHARDED, MESH_DIM_HEIGHT, MESH_DIM_WIDTH, UNSHARDED], self.mesh)
labels_layout = Layout([MESH_DIM_BATCH, UNSHARDED], self.mesh)
layouts = (images_layout, labels_layout)
with self.assertRaisesRegex(ValueError, (
f'batch_dim {MESH_DIM_BATCH} was specified but at least one layout did '
'not contain it')):
input_util.DTensorDataset(
dataset=dataset,
global_batch_size=32,
mesh=self.mesh,
layouts=layouts,
batch_dim=MESH_DIM_BATCH)
class InputUtilHelpersTest(test_util.DTensorBaseTest):
@parameterized.parameters(
{
'mesh_dims': [(MESH_DIM_BATCH, 8)],
'layout_specs': [UNSHARDED],
'batch_dim': None,
'counts': [1],
}, {
'mesh_dims': [(MESH_DIM_BATCH, 8)],
'layout_specs': [MESH_DIM_BATCH],
'batch_dim': None,
'counts': [8],
}, {
'mesh_dims': [(MESH_DIM_BATCH, 8)],
'layout_specs': [MESH_DIM_BATCH],
'batch_dim': MESH_DIM_BATCH,
'counts': [1],
}, {
'mesh_dims': [(MESH_DIM_BATCH, 2),
(MESH_DIM_HEIGHT, 4),
(MESH_DIM_WIDTH, 2)],
'layout_specs': [UNSHARDED, MESH_DIM_HEIGHT],
'batch_dim': None,
'counts': [1, 4],
}, {
'mesh_dims': [(MESH_DIM_BATCH, 2),
(MESH_DIM_HEIGHT, 4),
(MESH_DIM_WIDTH, 2)],
'layout_specs': [MESH_DIM_BATCH, MESH_DIM_WIDTH, MESH_DIM_HEIGHT],
'batch_dim': None,
'counts': [2, 2, 4],
}, {
'mesh_dims': [(MESH_DIM_BATCH, 2),
(MESH_DIM_HEIGHT, 4),
(MESH_DIM_WIDTH, 2)],
'layout_specs': [MESH_DIM_BATCH, MESH_DIM_WIDTH, MESH_DIM_HEIGHT],
'batch_dim': MESH_DIM_BATCH,
'counts': [1, 2, 4],
})
def testShardCounts(self, mesh_dims, layout_specs, batch_dim, counts):
num_devices = np.prod([size for _, size in mesh_dims])
mesh = mesh_util.create_mesh(
mesh_dims=mesh_dims, devices=['CPU:%d' % i for i in range(num_devices)])
layout = Layout(layout_specs, mesh)
self.assertEqual(input_util._shard_counts(layout, batch_dim), counts)
class DTensorIteratorSpecTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
mesh = mesh_util.create_mesh(
devices=['CPU:%d' % i for i in range(8)],
mesh_dims=[(MESH_DIM_BATCH, 8)])
self.mesh = self.configTestMesh({'CPU': mesh})
self.images = stateless_random_ops.stateless_random_uniform(
[8, 8, 3], seed=(1, 2), minval=0, maxval=255)
def testToTensorList(self):
dataset = dataset_ops.DatasetV2.from_tensors(self.images).repeat(8)
images_layout = Layout.replicated(self.mesh, rank=4)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=4,
mesh=self.mesh,
layouts=images_layout)
d_iterator = iter(d_dataset)
spec = input_util._DTensorIteratorSpec(
global_element_spec=d_iterator._global_element_spec,
layouts_str=d_iterator._layouts_str)
value = d_iterator
tensor_list = spec._to_tensor_list(value)
self.assertListEqual(tensor_list, [d_iterator._iterator_resource_dtensor])
def testFromTensorList(self):
dataset = dataset_ops.DatasetV2.from_tensors(self.images).repeat(8)
images_layout = Layout.replicated(self.mesh, rank=4)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=4,
mesh=self.mesh,
layouts=images_layout)
d_iterator = iter(d_dataset)
spec = input_util._DTensorIteratorSpec(
global_element_spec=d_iterator._global_element_spec,
layouts_str=d_iterator._layouts_str)
tensor_list = [d_iterator._iterator_resource_dtensor]
value = spec._from_tensor_list(tensor_list)
self.assertIsInstance(value, input_util._DTensorIterator)
self.assertIs(value._global_element_spec, d_iterator._global_element_spec)
self.assertEqual(value._layouts, d_iterator._layouts)
if __name__ == '__main__':
tf_test.main()
tensorflow/dtensor/python/tests/layout_propagation_test.py
+391
View File
@@ -0,0 +1,391 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for layout propagation."""
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import layout
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
UNSHARDED = layout.UNSHARDED
# Convenient constants to use for tests.
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
class LayoutPropagationV2Test(test_util.DTensorBaseTest):
def setUp(self):
super(LayoutPropagationV2Test, self).setUp()
global_ids = test_util.create_device_ids_array((1, 2))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = { # pylint: disable=g-complex-comprehension
device: layout.Mesh(
[_MESH_DIM_X, _MESH_DIM_Y],
global_ids,
local_ids,
test_util.create_device_list((1, 2), device),
)
for device in ('CPU', 'GPU', 'TPU')
}
self.mesh = self.configTestMesh(mesh_dict)
# 1D Layouts
self.unsharded_layout = layout.Layout.replicated(self.mesh, rank=1)
self.x_layout = layout.Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=1)
self.y_layout = layout.Layout.batch_sharded(self.mesh, _MESH_DIM_Y, rank=1)
# 2D Layouts
self.unsharded_unsharded_layout = layout.Layout.replicated(
self.mesh, rank=2
)
self.x_unsharded_layout = layout.Layout.batch_sharded(
self.mesh, _MESH_DIM_X, rank=2
)
self.unsharded_x_layout = layout.Layout.inner_sharded(
self.mesh, _MESH_DIM_X, rank=2
)
self.unsharded_y_layout = layout.Layout.inner_sharded(
self.mesh, _MESH_DIM_Y, rank=2
)
def test_layout_prop_v2_with_const_tf_function(self):
a = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
b = constant_op.constant([[10.0, 20.0], [30.0, 40.0]])
golden_result = math_ops.add(a, b)
c = api.copy_to_mesh(a, self.unsharded_unsharded_layout)
@polymorphic_function.function
def add_function():
d = constant_op.constant([[10.0, 20.0], [30.0, 40.0]])
return math_ops.add(c, d)
dtensor_result = add_function()
self.assertDTensorEqual(golden_result, self.unsharded_unsharded_layout,
dtensor_result)
def test_layout_prop_v2_while(self):
a = constant_op.constant([0, 1, 2, 1], dtype=dtypes.float32)
num_iterations = 10
@polymorphic_function.function
def function_with_while(t):
for _ in math_ops.range(num_iterations):
random_number = stateless_random_ops.stateless_random_normal(
shape=[4], seed=[1, 2], dtype=dtypes.float32
)
t = t + random_number
return t
golden_result = function_with_while(a)
a = numpy_util.pack_numpy(a, self.unsharded_layout)
dtensor_result = function_with_while(a)
self.assertDTensorEqual(golden_result, self.unsharded_layout,
dtensor_result)
@parameterized.named_parameters(
dict(testcase_name='unsharded', sharded_layout=0, use_split=False),
dict(testcase_name='x_sharded', sharded_layout=1, use_split=False),
dict(testcase_name='unsharded_split', sharded_layout=0, use_split=True),
dict(testcase_name='x_sharded_split', sharded_layout=1, use_split=True))
def test_while_microbatch(self, sharded_layout, use_split):
layouts = [self.unsharded_unsharded_layout, self.x_unsharded_layout]
sharded_layout = layouts[sharded_layout]
np.random.seed(0)
random_initial_value = np.random.uniform(size=4 * 4).reshape([4, 4])
if use_split:
random_batch = np.random.uniform(size=12 * 4).reshape([12, 4])
else:
random_batch = np.random.uniform(size=4 * 4).reshape([4, 4])
golden_variable = variables.Variable(random_initial_value)
@polymorphic_function.function
def update_weights(batch, variable):
accum_grads = array_ops.zeros_like_v2(variable)
for i in math_ops.range(3):
if use_split:
reshaped = array_ops.reshape(batch, [4, 3, 4])
mini_batch = array_ops.gather_v2(reshaped, i, axis=1)
else:
mini_batch = batch
with backprop.GradientTape() as tape:
logits = variable * variable + mini_batch
loss = math_ops.reduce_sum(logits * logits)
accum_grads += tape.gradient(loss, variable)
new_variable = variable + accum_grads
variable.assign(new_variable)
return accum_grads
golden_accum = update_weights(
constant_op.constant(random_batch), golden_variable
)
random_batch = numpy_util.pack_numpy(random_batch, sharded_layout)
random_initial_value = numpy_util.pack_numpy(random_initial_value,
sharded_layout)
dtensor_variable = d_variable.DVariable(random_initial_value)
dtensor_accum = update_weights(random_batch, dtensor_variable)
self.assertDTensorEqual(golden_accum, sharded_layout, dtensor_accum)
@parameterized.named_parameters(
dict(testcase_name='unsharded_split', sharded_layout=0),
dict(testcase_name='x_sharded_split', sharded_layout=1))
def test_while_microbatch_with_reused_gradient_accumulator(
self, sharded_layout):
layouts = [
self.unsharded_unsharded_layout, self.x_unsharded_layout,
self.unsharded_x_layout
]
sharded_layout_0 = layouts[sharded_layout]
sharded_layout_1 = layouts[2]
np.random.seed(0)
random_initial_value_1 = np.random.uniform(size=4 * 4).reshape(
[4, 4]).astype(np.float32)
random_initial_value_2 = np.random.uniform(size=4 * 4).reshape(
[4, 4]).astype(np.float32)
random_batch = np.random.uniform(size=12 * 4).reshape([12, 4
]).astype(np.float32)
golden_variable_1 = variables.Variable(random_initial_value_1)
golden_variable_2 = variables.Variable(random_initial_value_2)
@polymorphic_function.function
def update_weights(batch, variable1, variable2):
accum_grads = array_ops.zeros_like_v2(variable1)
accum_grads_2 = array_ops.zeros_like_v2(variable2)
for i in math_ops.range(3):
reshaped = array_ops.reshape(batch, [4, 3, 4])
mini_batch = array_ops.gather_v2(reshaped, i, axis=1)
with backprop.GradientTape() as tape:
logits_1 = variable1 * variable1 + mini_batch
logits_2 = variable2 * variable2 + mini_batch
loss_1 = math_ops.reduce_sum(logits_1 * logits_1)
loss_2 = math_ops.reduce_sum(logits_2 * logits_2)
loss = loss_1 + loss_2
grads = tape.gradient(loss, [variable1, variable2])
accum_grads += grads[0]
accum_grads_2 += grads[1]
new_variable = variable1 + accum_grads
new_variable_2 = variable2 + accum_grads_2
variable1.assign(new_variable)
variable2.assign(new_variable_2)
return accum_grads, accum_grads_2
golden_accum, golden_accum_2 = update_weights(
constant_op.constant(random_batch), golden_variable_1, golden_variable_2
)
random_batch = numpy_util.pack_numpy(random_batch, sharded_layout_0)
random_initial_value_1 = numpy_util.pack_numpy(random_initial_value_1,
sharded_layout_0)
dtensor_variable_1 = d_variable.DVariable(random_initial_value_1)
random_initial_value_2 = numpy_util.pack_numpy(random_initial_value_2,
sharded_layout_1)
dtensor_variable_2 = d_variable.DVariable(random_initial_value_2)
dtensor_accum, dtensor_accum_2 = update_weights(random_batch,
dtensor_variable_1,
dtensor_variable_2)
self.assertDTensorEqual(golden_accum, sharded_layout_0, dtensor_accum)
self.assertDTensorEqual(golden_accum_2, sharded_layout_1, dtensor_accum_2)
def test_layout_like(self):
a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.float32)
a = numpy_util.pack_numpy(a, self.unsharded_layout)
b = constant_op.constant([0, 1, 2, 3], dtype=dtypes.int64)
b = numpy_util.pack_numpy(b, self.x_layout)
@polymorphic_function.function
def layout_like_function(i, t):
i = math_ops.sqrt(i)
return api.relayout_like(i, t)
# Triggers relayout_like with different dtype.
dtensor_result = layout_like_function(a, b)
api.check_layout(dtensor_result, self.x_layout)
def test_layout_prop_v2_if(self):
a = constant_op.constant([0, 1, 2, 1], dtype=dtypes.float32)
a = numpy_util.pack_numpy(a, self.unsharded_layout)
@polymorphic_function.function
def function_with_if(t):
if math_ops.equal(math_ops.reduce_sum(t), 0):
t = math_ops.sqrt(t)
return api.relayout(t, self.x_layout)
else:
return array_ops.zeros_like_v2(t)
dtensor_result = function_with_if(a)
api.check_layout(dtensor_result, self.x_layout)
def test_layout_prop_v2_if_with_different_layouts_for_branches(self):
unsharded_unsharded = layout.Layout.replicated(self.mesh, rank=2)
unsharded_y = layout.Layout.inner_sharded(self.mesh, _MESH_DIM_Y, rank=2)
x_unsharded = layout.Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=2)
a = np.random.uniform(size=16).reshape([4, 4])
a = numpy_util.pack_numpy(a, unsharded_unsharded)
@polymorphic_function.function
def function_with_if(t):
if math_ops.equal(math_ops.reduce_sum(t), 0):
t = math_ops.sqrt(t)
return api.relayout(t, unsharded_y)
else:
t = array_ops.zeros_like_v2(a)
return api.relayout(t, x_unsharded)
dtensor_result = function_with_if(a)
x_y_sharded = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)
api.check_layout(dtensor_result, x_y_sharded)
def test_partial_relayout_in_function(self):
sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)
a = np.random.uniform(size=16).reshape([4, 4])
a = numpy_util.pack_numpy(a, sharded_layout)
replicated_layout = layout.Layout(
[layout.MATCH, layout.UNSHARDED], mesh=self.mesh
)
@polymorphic_function.function
def func_with_relayout(t):
out = math_ops.cast(t, dtypes.float32)
out = math_ops.sqrt(out)
return api.relayout(out, replicated_layout)
out = func_with_relayout(a)
expected_layout = layout.Layout([_MESH_DIM_X, layout.UNSHARDED], self.mesh)
api.check_layout(out, expected_layout)
def test_partial_relayout_in_eager(self):
sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)
a = np.random.uniform(size=16).reshape([4, 4])
a = numpy_util.pack_numpy(a, sharded_layout)
replicated_layout = layout.Layout(
[layout.MATCH, layout.UNSHARDED], mesh=self.mesh
)
a = math_ops.cast(a, dtypes.float32)
a = math_ops.sqrt(a)
out = api.relayout(a, replicated_layout)
expected_layout = layout.Layout([_MESH_DIM_X, layout.UNSHARDED], self.mesh)
api.check_layout(out, expected_layout)
@parameterized.named_parameters(
dict(testcase_name='unsharded_unsharded', sharded_layout=0),
dict(testcase_name='x_unsharded', sharded_layout=1),
dict(testcase_name='unsharded_y', sharded_layout=2),
)
def test_relayout_equivalent_layouts(self, sharded_layout):
layouts = [
self.unsharded_unsharded_layout, self.x_unsharded_layout,
self.unsharded_y_layout]
expected_layout = layouts[sharded_layout]
inputs = constant_op.constant([[0, 1], [2, 1]], dtype=dtypes.float32)
sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)
inputs = numpy_util.pack_numpy(inputs, sharded_layout)
output = api.relayout(inputs, expected_layout)
api.check_layout(output, expected_layout)
def test_strided_slice_grad(self):
np.random.seed(0)
random_initial_value = np.random.uniform(size=4).reshape([4])
random_initial_value = numpy_util.pack_numpy(random_initial_value,
self.unsharded_layout)
@polymorphic_function.function
def fn_with_strided_slice(t):
a = array_ops.strided_slice(t, [1], [2], shrink_axis_mask=1)
return math_ops.sqrt(a)
random_variable = d_variable.DVariable(random_initial_value)
with backprop.GradientTape() as tape:
output = fn_with_strided_slice(random_variable)
grads = tape.gradient(output, [random_variable])
self.assertTrue(api.fetch_layout(grads[0]).is_fully_replicated())
def test_layout_prop_v2_infinite_loop(self):
x_unsharded = layout.Layout([_MESH_DIM_X, layout.UNSHARDED], self.mesh)
unsharded_x = layout.Layout([layout.UNSHARDED, _MESH_DIM_X], self.mesh)
@polymorphic_function.function
def func(input_a, input_b):
out = array_ops.identity(
math_ops.matmul(input_a, array_ops.identity(input_b))
)
return api.relayout(out, unsharded_x)
result = func(
api.call_with_layout(
array_ops.ones, shape=(16, 16), layout=x_unsharded
),
api.call_with_layout(
array_ops.ones, shape=(16, 16), layout=unsharded_x
),
)
api.check_layout(result, unsharded_x)
def test_layout_prop_parted_layout(self):
value = np.array([1.0, 2.0, 3.0, 4.0])
expected = nn_ops.softmax_v2(gen_nn_ops.relu(value))
@polymorphic_function.function
def func(value):
return nn_ops.softmax_v2(gen_nn_ops.relu(value))
parted_layout = self.x_layout.to_parted()
result = func(api.relayout(value, parted_layout))
# Verifies the parted layout can be propagated through a chain of ops to the
# final output.
self.assertDTensorEqual(expected, parted_layout, result)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/layout_test.py
+638
View File
@@ -0,0 +1,638 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for layout.py."""
import copy
import itertools
import pickle
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import layout
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import combinations
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.platform import test
UNSHARDED = layout.UNSHARDED
# Convenient constants to use for tests.
_MESH_DIM_BATCH = 'batch'
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
_MESH_2D_STRING = (
'|batch=2,x=2|0,1,2,3|0,1,2,3|'
'/job:localhost/replica:0/task:0/device:TPU:0,'
'/job:localhost/replica:0/task:0/device:TPU:1,'
'/job:localhost/replica:0/task:0/device:TPU:2,'
'/job:localhost/replica:0/task:0/device:TPU:3'
)
_2D_GLOBAL_IDS = test_util.create_device_ids_array((2, 2))
_2D_MESH = layout.Mesh([_MESH_DIM_BATCH, _MESH_DIM_X], _2D_GLOBAL_IDS,
np.ravel(_2D_GLOBAL_IDS).tolist(),
test_util.create_device_list((2, 2), 'TPU'))
_2D_X_Y_MESH = layout.Mesh([_MESH_DIM_X, _MESH_DIM_Y], _2D_GLOBAL_IDS,
np.ravel(_2D_GLOBAL_IDS).tolist(),
test_util.create_device_list((2, 2), 'CPU'))
_SINGLE_DEVICE_MESH = layout.Mesh.from_device(
'job:localhost/replica:0/task:0/TPU:0'
)
class MeshTest(test_util.DTensorBaseTest, parameterized.TestCase):
def test_mesh_single_device(self):
self.assertTrue(_SINGLE_DEVICE_MESH.is_single_device())
def test_mesh_single_device_to_string(self):
roundtrip = layout.Mesh.from_string(_SINGLE_DEVICE_MESH.to_string())
self.assertTrue(roundtrip.is_single_device())
self.assertEqual(roundtrip.single_device, _SINGLE_DEVICE_MESH.single_device)
def test_mesh_single_device_to_proto(self):
roundtrip = layout.Mesh.from_proto(_SINGLE_DEVICE_MESH.as_proto())
self.assertTrue(roundtrip.is_single_device())
self.assertEqual(roundtrip.single_device, _SINGLE_DEVICE_MESH.single_device)
def test_mesh_reciprocal_mock_string_and_object(self):
generated_mesh_from_string = layout.Mesh.from_string(_MESH_2D_STRING)
self.assertProtoEquals(_2D_MESH.as_proto(),
generated_mesh_from_string.as_proto())
def test_mesh_reciprocal_string_rep(self):
new_mesh_str = layout.Mesh.from_string(_MESH_2D_STRING).to_string()
self.assertEqual(_MESH_2D_STRING, new_mesh_str)
def test_mesh_repr(self):
device_ids = test_util.create_device_ids_array((4, 2))
mesh = layout.Mesh([_MESH_DIM_BATCH, _MESH_DIM_X], device_ids,
np.ravel(device_ids).tolist(),
test_util.create_device_list((4, 2), 'CPU'))
self.assertIn('batch=4,x=2', repr(mesh))
self.assertIn('CPU:0', repr(mesh))
self.assertIn('CPU:7', repr(mesh))
def test_mesh_contains_dim(self):
self.assertTrue(_2D_MESH.contains_dim('batch'))
self.assertTrue(_2D_MESH.contains_dim('x'))
self.assertFalse(_2D_MESH.contains_dim('y'))
def test_mesh_contains(self):
self.assertIn('batch', _2D_MESH)
self.assertIn('x', _2D_MESH)
self.assertNotIn('y', _2D_MESH)
def test_mesh_dim_names_property(self):
self.assertSequenceEqual(_2D_MESH.dim_names, ['batch', 'x'])
def test_mesh_size_property(self):
self.assertEqual(_2D_MESH.size, 4)
def test_mesh_device_type(self):
self.assertEqual(_2D_MESH.device_type(), 'TPU')
self.assertEqual(_2D_X_Y_MESH.device_type(), 'CPU')
def test_mesh_num_local_devices(self):
self.assertEqual(_2D_MESH.num_local_devices(), 4)
def test_mesh_min_global_device_id(self):
self.assertEqual(_2D_MESH.min_global_device_id(), 0)
def test_mesh_is_remote(self):
self.assertFalse(_2D_MESH.is_remote())
def test_mesh_local_device_ids(self):
self.assertSequenceEqual(_2D_MESH.local_device_ids(), [0, 1, 2, 3])
def test_mesh_local_devices(self):
self.assertSequenceEqual(_2D_MESH.local_devices(), [
'/job:localhost/replica:0/task:0/device:TPU:0',
'/job:localhost/replica:0/task:0/device:TPU:1',
'/job:localhost/replica:0/task:0/device:TPU:2',
'/job:localhost/replica:0/task:0/device:TPU:3'
])
def test_mesh_shape(self):
self.assertSequenceEqual(_2D_MESH.shape(), [2, 2])
def test_mesh_pickle(self):
pickled = pickle.dumps(_2D_MESH)
unpickled = pickle.loads(pickled)
self.assertEqual(_2D_MESH, unpickled)
def test_mesh_pickle_w_modification_after_init(self):
mesh = copy.copy(_2D_MESH)
mesh._name = 'fake_name'
pickled = pickle.dumps(mesh)
unpickled = pickle.loads(pickled)
self.assertEqual(mesh, unpickled)
def test_mesh_dims(self):
device_ids = test_util.create_device_ids_array((4, 2))
mesh = layout.Mesh(
[_MESH_DIM_BATCH, _MESH_DIM_X],
device_ids,
np.ravel(device_ids).tolist(),
test_util.create_device_list((4, 2), 'CPU'))
self.assertIn(_MESH_DIM_BATCH, mesh)
self.assertIn(_MESH_DIM_X, mesh)
self.assertNotIn(_MESH_DIM_Y, mesh)
self.assertEqual(mesh[_MESH_DIM_BATCH].name, _MESH_DIM_BATCH)
self.assertEqual(mesh[_MESH_DIM_BATCH].size, 4)
self.assertEqual(mesh[_MESH_DIM_X].name, _MESH_DIM_X)
self.assertEqual(mesh[_MESH_DIM_X].size, 2)
@parameterized.parameters(
{
'mesh_dims': ('a',),
'mesh_shape': (8,),
'strides': [1],
}, {
'mesh_dims': ('a', 'b', 'c'),
'mesh_shape': (2, 4, 2),
'strides': [8, 2, 1],
}, {
'mesh_dims': ('a', 'b', 'c', 'd'),
'mesh_shape': (8, 16, 2, 4),
'strides': [128, 8, 4, 1],
})
def test_mesh_strides(self, mesh_dims, mesh_shape, strides):
device_ids = test_util.create_device_ids_array(mesh_shape)
mesh = layout.Mesh(
dim_names=list(mesh_dims),
global_device_ids=device_ids,
local_device_ids=np.ravel(device_ids).tolist(),
local_devices=test_util.create_device_list(mesh_shape, 'CPU'))
self.assertEqual(mesh.strides, strides)
def test_mesh_coords(self):
mesh_shape = (2, 4, 2)
device_ids = test_util.create_device_ids_array(mesh_shape)
mesh = layout.Mesh(
dim_names=['a', 'b', 'c'],
global_device_ids=device_ids,
local_device_ids=np.ravel(device_ids).tolist(),
local_devices=test_util.create_device_list(mesh_shape, 'CPU'))
coords = itertools.product(range(2), range(4), range(2))
# Repeat coords to overflow idx, mesh coords should still be correct.
coords = itertools.chain.from_iterable(
itertools.repeat(tuple(coords), times=3))
for idx, (a, b, c) in enumerate(coords):
self.assertAllEqual(mesh.coords(idx), [a, b, c])
@parameterized.named_parameters(('use_xla_spmd', True),
('do_not_use_xla_spmd', False))
def test_mesh_use_xla_spmd_tpu_mesh(self, use_xla_spmd):
mesh_shape = (2, 4, 2)
device_ids = test_util.create_device_ids_array(mesh_shape)
mesh = layout.Mesh(
dim_names=['a', 'b', 'c'],
global_device_ids=device_ids,
local_device_ids=np.ravel(device_ids).tolist(),
local_devices=test_util.create_device_list(mesh_shape, 'TPU'),
use_xla_spmd=use_xla_spmd)
self.assertEqual(use_xla_spmd, mesh.use_xla_spmd())
@parameterized.named_parameters(('gpu', 'GPU'), ('cpu', 'CPU'))
def test_mesh_use_xla_spmd_for_non_tpu_mesh_raises_error(self, mesh_type):
mesh_shape = (2, 4, 2)
device_ids = test_util.create_device_ids_array(mesh_shape)
with self.assertRaisesRegex(ValueError,
'XLA SPMD is not currently not supported for'):
layout.Mesh(
dim_names=['a', 'b', 'c'],
global_device_ids=device_ids,
local_device_ids=np.ravel(device_ids).tolist(),
local_devices=test_util.create_device_list(mesh_shape, mesh_type),
use_xla_spmd=True)
@parameterized.named_parameters(
('use_xla_spmd', True), ('do_not_use_xla_spmd', False)
)
def test_mesh_as_proto_use_xla_spmd(self, use_xla_spmd):
mesh_shape = (2, 4, 2)
device_ids = test_util.create_device_ids_array(mesh_shape)
mesh = layout.Mesh(
dim_names=['a', 'b', 'c'],
global_device_ids=device_ids,
local_device_ids=np.ravel(device_ids).tolist(),
local_devices=test_util.create_device_list(mesh_shape, 'TPU'),
use_xla_spmd=use_xla_spmd,
)
mesh_proto = mesh.as_proto()
self.assertEqual(mesh_proto.use_xla_spmd, mesh.use_xla_spmd())
def test_mesh_from_string_with_use_xla_spmd(self):
mesh_str_without_global_device_ids = (
'|batch=2|0,1|0,1|/job:localhost/replica:0/task:0'
'/device:TPU:0,/job:localhost/replica:0/task:0/device:TPU:1|use_xla_spmd'
)
mesh = layout.Mesh.from_string(mesh_str_without_global_device_ids)
self.assertTrue(mesh.use_xla_spmd())
def test_mesh_from_string_with_use_xla_spmd_and_global_devices(self):
mesh_str_with_global_device_ids = (
'|batch=2|0,1|0|/job:localhost/replica:0/task:0'
'/device:TPU:0|/job:localhost/replica:0/task:0/device:TPU:0,'
'/job:localhost/replica:0/task:0/device:TPU:1|use_xla_spmd'
)
mesh = layout.Mesh.from_string(mesh_str_with_global_device_ids)
self.assertTrue(mesh.use_xla_spmd())
def test_non_unique_device_type(self):
a = test_util.create_device_array((2,), 'CPU')
b = test_util.create_device_array((2,), 'TPU')
c = np.vstack([a, b])
global_ids = test_util.create_device_ids_array(c.shape)
local_ids = np.ravel(global_ids).tolist()
with self.assertRaisesRegex(ValueError,
'Devices containing multiple device_types'):
layout.Mesh([_MESH_DIM_BATCH, _MESH_DIM_X], global_ids, local_ids,
np.ravel(c).tolist())
def test_duplicated_devices(self):
a = test_util.create_device_array((2,), 'CPU')
b = test_util.create_device_array((2,), 'CPU')
c = np.vstack([a, b])
global_ids = test_util.create_device_ids_array((2, 2))
local_ids = global_ids.flatten().tolist()
with self.assertRaisesRegex(
ValueError, 'Duplicate devices found in mesh specification'):
layout.Mesh([_MESH_DIM_BATCH, _MESH_DIM_X], global_ids, local_ids,
np.ravel(c).tolist())
def test_inconsecutive_device_ids(self):
a = test_util.create_device_array((2,), 'CPU')
global_ids = test_util.create_device_ids_array((2))
global_ids = np.flip(global_ids)
local_ids = global_ids.flatten().tolist()
with self.assertRaisesRegex(ValueError,
'global_device_ids must sequentially increase'):
layout.Mesh([_MESH_DIM_BATCH], global_ids, local_ids,
np.ravel(a).tolist())
def test_host_mesh(self):
device_ids = test_util.create_device_ids_array((4, 2))
mesh = layout.Mesh(
[_MESH_DIM_BATCH, _MESH_DIM_X],
device_ids,
np.ravel(device_ids).tolist(),
test_util.create_device_list((4, 2), 'GPU'),
mesh_name='name_not_preserved',
)
expected = layout.Mesh(
[_MESH_DIM_BATCH, _MESH_DIM_X],
device_ids,
np.ravel(device_ids).tolist(),
test_util.create_device_list((4, 2), 'CPU'),
)
host_mesh = mesh.host_mesh()
self.assertEqual(host_mesh, expected)
class LayoutTest(test_util.DTensorBaseTest, parameterized.TestCase):
def test_empty_sharding_spec_different_from_single_unsharded(self):
layout_str_single_unsharded = (
'sharding_specs:unsharded, mesh:' + _MESH_2D_STRING
)
layout_str_empty_sharding_spec = 'sharding_specs: mesh:' + _MESH_2D_STRING
self.assertNotEqual(
layout.Layout.from_string(layout_str_single_unsharded).to_string(),
layout.Layout.from_string(layout_str_empty_sharding_spec).to_string(),
)
@parameterized.named_parameters(
dict(
testcase_name='sharded_batch_and_x',
test_layout_str='sharding_specs:batch,x, mesh:' + _MESH_2D_STRING,
),
dict(
testcase_name='unsharded_explicit',
test_layout_str='sharding_specs:'
+ UNSHARDED
+ ','
+ UNSHARDED
+ ','
+ ' mesh:'
+ _MESH_2D_STRING,
),
)
def test_layout_reciprocal_string_rep(self, test_layout_str):
new_layout_str = layout.Layout.from_string(test_layout_str).to_string()
self.assertEqual(test_layout_str, new_layout_str)
def test_layout_pickle(self):
replicated = layout.Layout.replicated(_2D_MESH, rank=3)
pickled = pickle.dumps(replicated)
unpickled = pickle.loads(pickled)
self.assertEqual(replicated, unpickled)
def test_layout_repr(self):
tensor_layout = layout.Layout.batch_sharded(
_2D_MESH, _MESH_DIM_BATCH, rank=2)
self.assertIn('batch,unsharded', repr(tensor_layout))
def test_throws_for_non_mesh(self):
with self.assertRaisesRegex(ValueError, 'mesh is not a valid Mesh object'):
layout.Layout([_MESH_DIM_BATCH, _MESH_DIM_X], 'string_mesh')
def test_throws_for_repeated_dimension(self):
with self.assertRaisesRegex(ValueError, 'Mesh dimensions must be unique.'):
layout.Layout([_MESH_DIM_BATCH, _MESH_DIM_BATCH], _2D_MESH)
def test_throws_for_invalid_sharding_spec(self):
with self.assertRaisesRegex(
ValueError,
'A dimension sharding must either be a valid mesh dimension or ' +
'UNSHARDED.'):
layout.Layout(['WRONG_SHARDING_SPEC', 'UNSHARDED'], _2D_MESH)
def test_data_parallel_layout(self):
tensor_layout = layout.Layout.batch_sharded(
_2D_MESH, _MESH_DIM_BATCH, rank=2)
self.assertEqual(
tensor_layout.num_shards(0), _2D_MESH.dim_size(_MESH_DIM_BATCH))
self.assertEqual(tensor_layout.num_shards(1), 1)
def test_global_shape_from_local_shape(self):
tensor_layout = layout.Layout(
[_MESH_DIM_BATCH, _MESH_DIM_X, layout.UNSHARDED],
mesh=_2D_MESH,
)
self.assertEqual(
tensor_layout.global_shape_from_local_shape(
tensor_shape.TensorShape((1, 3, 5))
),
(2, 6, 5),
)
def test_local_shape_from_global_shape(self):
tensor_layout = layout.Layout(
[_MESH_DIM_BATCH, _MESH_DIM_X, layout.UNSHARDED],
mesh=_2D_MESH,
)
self.assertEqual(
tensor_layout.local_shape_from_global_shape(
tensor_shape.TensorShape((2, 6, 5))
),
(1, 3, 5),
)
def test_single_device_layout(self):
tensor_layout = layout.Layout.from_single_device_mesh(_SINGLE_DEVICE_MESH)
tensor_layout2 = layout.Layout.from_device(
_SINGLE_DEVICE_MESH.single_device
)
self.assertTrue(tensor_layout.is_single_device())
self.assertEqual(tensor_layout.mesh, _SINGLE_DEVICE_MESH)
self.assertEqual(tensor_layout, tensor_layout2)
def test_single_device_layout_from_string(self):
tensor_layout = layout.Layout.from_single_device_mesh(_SINGLE_DEVICE_MESH)
roundtrip = layout.Layout.from_string(tensor_layout.to_string())
self.assertEqual(roundtrip, tensor_layout)
def test_single_device_layout_from_proto(self):
tensor_layout = layout.Layout.from_single_device_mesh(_SINGLE_DEVICE_MESH)
roundtrip = layout.Layout.from_proto(tensor_layout.as_proto())
self.assertEqual(roundtrip, tensor_layout)
def test_parted_layout(self):
tensor_layout = layout.Layout.batch_sharded(
_2D_MESH, _MESH_DIM_BATCH, rank=2
)
parted_layout = tensor_layout.to_parted()
self.assertEqual(parted_layout.type, layout.LayoutType.PARTED)
class RelayoutTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
global_ids = test_util.create_device_ids_array((2, 2))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = { # pylint: disable=g-complex-comprehension
device: layout.Mesh(
[_MESH_DIM_X, _MESH_DIM_Y],
global_ids,
local_ids,
test_util.create_device_list((2, 2), device),
)
for device in ('CPU', 'GPU', 'TPU')
}
self.mesh = self.configTestMesh(mesh_dict)
# 1D Layouts
self.x_layout = layout.Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=1)
self.y_layout = layout.Layout.batch_sharded(self.mesh, _MESH_DIM_Y, rank=1)
# 2D Layouts
self.unsharded_unsharded_layout = layout.Layout.replicated(
self.mesh, rank=2
)
self.x_unsharded_layout = layout.Layout.batch_sharded(
self.mesh, _MESH_DIM_X, rank=2
)
self.unsharded_x_layout = layout.Layout.inner_sharded(
self.mesh, _MESH_DIM_X, rank=2
)
@combinations.generate(
combinations.combine(is_graph=[False, True], is_replicated=[False, True])
)
def test_relayout(self, is_graph, is_replicated):
inp = stateless_random_ops.stateless_random_uniform([4, 4], seed=[0, 1])
if is_replicated:
to_layout = self.unsharded_unsharded_layout
else:
to_layout = self.x_unsharded_layout
def do_relayout():
return api.relayout(inp, to_layout)
if is_graph:
relayout_fn = polymorphic_function.function(do_relayout)
self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
"No OpKernel was registered to support Op 'Relayout'",
relayout_fn,
)
else:
self.assertDTensorEqual(inp, to_layout, do_relayout())
@combinations.generate(combinations.combine(is_graph=[False, True]))
def test_relayout_to_parted(self, is_graph):
data = np.array([1, 2, 3, 4.0], dtype='f4')
inp = api.relayout(data, self.y_layout)
def do_relayout():
return api.relayout(inp, self.y_layout.to_parted())
if is_graph:
do_relayout = polymorphic_function.function(do_relayout)
with api.default_mesh(self.mesh):
result = do_relayout()
self.assertDTensorEqual(data, self.y_layout.to_parted(), result)
@combinations.generate(combinations.combine(is_graph=[False, True]))
def test_relayout_like_simple(self, is_graph):
data = np.array([1, 2, 3, 4.0], dtype='f4')
inp = api.relayout(data, self.y_layout)
inp_layout = api.relayout(data, self.x_layout)
def do_relayout():
return api.relayout_like(inp, inp_layout)
if is_graph:
do_relayout = polymorphic_function.function(do_relayout)
with api.default_mesh(self.mesh):
result = do_relayout()
self.assertDTensorEqual(data, self.x_layout, result)
def test_relayout_like_init_scope(self):
data = np.array([1, 2, 3, 4.0], dtype='f4')
inp = api.relayout(data, self.y_layout)
inp_layout = api.relayout(data, self.x_layout)
@polymorphic_function.function
def do_relayout(x):
with ops.init_scope():
return api.relayout_like(inp, x)
with api.default_mesh(self.mesh):
with self.assertRaisesRegex(
TypeError, 'is out of scope and cannot be used here'
):
result = do_relayout(inp_layout)
result.numpy()
def test_nested_relayout_gradient_preserves_layout(self):
# Test that nesting gradient tapes with relayouts preserves the layout of
# the original DTensor input. The second-order gradient should have a layout
# equivalent to the original input, even if the inner gradient tape
# relayouts the DTensor to a different layout.
@polymorphic_function.function
def inner(x):
with backprop.GradientTape() as tape:
tape.watch(x)
t = x * 1.0
t = api.relayout(t, self.unsharded_x_layout)
cube = t * t * t
grad = tape.gradient(cube, x)
return grad
@polymorphic_function.function
def outer(x):
with backprop.GradientTape() as tape:
tape.watch(x)
t = api.relayout(x, self.x_unsharded_layout)
grad = inner(t)
out = grad + t
out_grad = tape.gradient(out, x)
return out_grad
a = stateless_random_ops.stateless_random_uniform([8, 8], seed=[0, 1])
a_dt = api.relayout(a, self.unsharded_unsharded_layout)
with ops.device_v2(api.device_name()):
inner_grad = inner(a_dt)
outer_grad = outer(a_dt)
self.assertDTensorEqual(
3 * a * a, self.unsharded_unsharded_layout, inner_grad
)
self.assertDTensorEqual(
6 * a + 1, self.unsharded_unsharded_layout, outer_grad
)
def test_wus_using_relayout(self):
sharded_layout = layout.Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=2)
w = stateless_random_ops.stateless_random_uniform(
[4, 4], seed=[0, 1], dtype=dtypes.float32
)
sharded_w = api.relayout(w, sharded_layout)
replicated_layout = layout.Layout(
[layout.UNSHARDED, layout.UNSHARDED], mesh=self.mesh
)
@polymorphic_function.function
def func_with_relayout(t):
with backprop.GradientTape() as tape:
tape.watch(t)
t = t + t
out = api.relayout(t, replicated_layout)
loss = math_ops.reduce_sum(out)
grad = tape.gradient(loss, t)
t = t - grad
return t
func_with_relayout(sharded_w)
@combinations.generate(
combinations.combine(size=[16, 4096], is_graph=[False, True])
)
def test_call_with_layout(self, size, is_graph):
layout_x = layout.Layout.batch_sharded(
self.mesh, batch_dim=_MESH_DIM_X, rank=1
)
layout_y = layout.Layout.batch_sharded(
self.mesh, batch_dim=_MESH_DIM_Y, rank=1
)
expected = array_ops.zeros(shape=[size])
def func():
tensor_x = api.call_with_layout(array_ops.zeros, layout_x, shape=[size])
tensor_y = api.call_with_layout(array_ops.zeros, layout_y, shape=[size])
return tensor_x, tensor_y
if is_graph:
func = polymorphic_function.function(func)
with api.default_mesh(self.mesh):
tensor_x, tensor_y = func()
self.assertDTensorEqual(expected, layout_x, tensor_x)
self.assertDTensorEqual(expected, layout_y, tensor_y)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/mesh_util_test.py
+364
View File
@@ -0,0 +1,364 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for mesh_util."""
import os
from unittest import mock
from absl.testing import parameterized
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import mesh_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager import context
from tensorflow.python.framework import config as tf_config
from tensorflow.python.framework import device as tf_device
from tensorflow.python.platform import test
class MeshUtilTest(test_util.DTensorBaseTest):
"""Tests for mesh_util that do not require accelerator initialization."""
def test_mesh_creation(self):
self.skipForDeviceType(
['TPU'], reason='Test is intended for CPUs and GPUs.'
)
mesh = mesh_util.create_mesh()
num_devices = len(test_util.list_local_logical_devices(mesh.device_type()))
self.assertEqual(mesh.num_local_devices(), num_devices)
self.assertEqual(mesh.size, num_devices)
def test_mesh_dict_creation(self):
self.skipForDeviceType(
['TPU'], reason='Test is intended for CPUs and GPUs.'
)
num_devices = len(test_util.list_local_logical_devices('CPU'))
mesh = mesh_util.create_mesh({'x': num_devices, 'y': 1}, device_type='CPU')
num_devices = len(test_util.list_local_logical_devices(mesh.device_type()))
self.assertEqual(mesh.num_local_devices(), num_devices)
self.assertEqual(mesh.dim_names, ['x', 'y'])
self.assertEqual(mesh.size, num_devices)
def test_tpu_mesh_creation(self):
self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
mesh = mesh_util.create_mesh(mesh_name='1d_mesh', device_type='TPU')
num_devices = len(test_util.list_local_logical_devices('TPU'))
self.assertEqual(mesh.num_local_devices(), num_devices)
self.assertEqual(mesh.size, num_devices)
@parameterized.named_parameters(('use_xla_spmd', True),
('do_not_use_xla_spmd', False))
def test_tpu_2d_mesh_creation(self, use_xla_spmd):
self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
self.skipForDeviceType(['TPU'],
reason='Test requires exactly 2 cores',
unless_device_count_equals_to=2)
devices = test_util.list_local_logical_devices('TPU')
self.assertLen(devices, 2)
mesh = mesh_util.create_mesh([('x', 2), ('y', 1)],
device_type='TPU',
use_xla_spmd=use_xla_spmd)
self.assertEqual(mesh.num_local_devices(), 2)
self.assertEqual(mesh.size, 2)
self.assertAllEqual(mesh.dim_names, ['x', 'y'])
self.assertEqual(mesh.use_xla_spmd(), use_xla_spmd)
def test_tpu_2d_mesh_creation_with_devices(self):
self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
self.skipForDeviceType(['TPU'],
reason='Test requires at least 2 cores',
unless_device_count_equals_to=2)
devices = test_util.list_local_logical_devices('TPU')
self.assertLen(devices, 2)
mesh = mesh_util.create_mesh([('x', 2), ('y', 1)],
devices=['/device:tpu:0', '/device:tpu:1'])
self.assertEqual(mesh.num_local_devices(), 2)
self.assertEqual(mesh.size, 2)
self.assertAllEqual(mesh.dim_names, ['x', 'y'])
def test_tpu_2d_mesh_creation_with_device_specs(self):
self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
self.skipForDeviceType(['TPU'],
reason='Test requires at least 2 cores',
unless_device_count_equals_to=2)
devices = test_util.list_local_logical_devices('TPU')
self.assertLen(devices, 2)
mesh = mesh_util.create_mesh(
[('x', 2), ('y', 1)],
devices=[
tf_device.DeviceSpec.from_string('/tpu:0'),
tf_device.DeviceSpec.from_string('/tpu:1'),
],
)
self.assertEqual(mesh.num_local_devices(), 2)
self.assertEqual(mesh.size, 2)
self.assertAllEqual(mesh.dim_names, ['x', 'y'])
def test_single_client_mesh_creation(self):
self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
num_devices = len(test_util.list_local_logical_devices('CPU'))
with mock.patch.object(accelerator_util,
'is_initialized') as is_initialized:
is_initialized.return_value = True
mesh = mesh_util.create_distributed_mesh(
mesh_name='single_client_1d_mesh', mesh_dims=[('x', num_devices)])
self.assertEqual(mesh.num_local_devices(), num_devices)
self.assertEqual(mesh.size, num_devices)
def test_single_client_mesh_dict_creation(self):
self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
num_devices = len(test_util.list_local_logical_devices('CPU'))
with mock.patch.object(
accelerator_util, 'is_initialized'
) as is_initialized:
is_initialized.return_value = True
mesh = mesh_util.create_distributed_mesh(
mesh_name='single_client_1d_mesh',
mesh_dims={'x': num_devices, 'y': 1},
)
self.assertEqual(mesh.num_local_devices(), num_devices)
self.assertEqual(mesh.dim_names, ['x', 'y'])
self.assertEqual(mesh.size, num_devices)
def test_single_client_mesh_with_local_devices(self):
self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
with mock.patch.object(accelerator_util,
'is_initialized') as is_initialized:
is_initialized.return_value = True
mesh = mesh_util.create_distributed_mesh(
mesh_name='single_client_1d_mesh',
mesh_dims=[('x', 1)],
local_devices=['CPU:0'])
self.assertEqual(mesh.num_local_devices(), 1)
self.assertEqual(mesh.size, 1)
def test_create_distributed_mesh_requires_initialize(self):
self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
with mock.patch.object(accelerator_util,
'is_initialized') as is_initialized:
is_initialized.return_value = False
with self.assertRaisesRegex(ValueError, 'Accelerators are uninitialized'):
_ = mesh_util.create_distributed_mesh(
mesh_name='single_client_1d_mesh',
mesh_dims=[('x', 1)],
local_devices=['CPU:0'])
def test_single_client_mesh_creation_wrong_shape(self):
self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
num_devices = len(test_util.list_local_logical_devices('CPU'))
with mock.patch.object(accelerator_util,
'is_initialized') as is_initialized:
is_initialized.return_value = True
with self.assertRaisesRegex(ValueError,
'must be equal to total size of the mesh'):
mesh_util.create_distributed_mesh(
mesh_name='single_client_1d_mesh',
mesh_dims=[('x', num_devices * 2)])
def test_single_client_mesh_creation_using_fewer_devices(self):
self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
test_util.reset_logical_devices('CPU', 4)
with mock.patch.object(accelerator_util,
'is_initialized') as is_initialized:
is_initialized.return_value = True
mesh = mesh_util.create_distributed_mesh(
mesh_name='single_client_1d_mesh',
mesh_dims=[('x', 2)],
local_devices=['CPU:0', 'CPU:1'])
self.assertEqual(mesh.num_local_devices(), 2)
self.assertEqual(mesh.size, 2)
mesh = mesh_util.create_distributed_mesh(
mesh_name='single_client_1d_mesh',
mesh_dims=[('x', 2)],
local_devices=['CPU:0', 'CPU:1'])
self.assertEqual(mesh.num_local_devices(), 2)
self.assertEqual(mesh.size, 2)
def test_single_client_mesh_creation_with_xla_spmd_raises_error(self):
self.skipForDeviceType(['TPU'],
reason='Test is intended for non TPU devices')
test_util.reset_logical_devices('CPU', 4)
with mock.patch.object(accelerator_util,
'is_initialized') as is_initialized:
is_initialized.return_value = True
with self.assertRaisesRegex(
ValueError, 'XLA SPMD is not currently not supported for'):
mesh_util.create_distributed_mesh(
mesh_name='single_client_mesh',
mesh_dims=[('x', 2)],
local_devices=['CPU:0', 'CPU:1'],
use_xla_spmd=True)
@mock.patch.object(config, 'num_clients')
@mock.patch.object(accelerator_util, 'is_initialized')
def test_multi_client_mesh_creation(self, num_clients, is_initialized):
self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
with mock.patch.object(accelerator_util,
'is_initialized') as is_initialized:
with mock.patch.object(config, 'num_clients') as num_clients:
num_clients.return_value = 2
is_initialized.return_value = True
test_util.reset_context()
cpus = tf_config.list_physical_devices('CPU')
tf_config.set_logical_device_configuration(
cpus[0], [context.LogicalDeviceConfiguration()] * 4
)
with mock.patch.object(config, 'client_id', return_value=0):
mesh_1 = mesh_util.create_distributed_mesh(
mesh_name='multi_client_1d_mesh_1',
mesh_dims=[('x', 4)],
local_devices=['CPU:0', 'CPU:1'])
self.assertEqual(mesh_1.num_local_devices(), 2)
self.assertEqual(mesh_1.size, 4)
with mock.patch.object(config, 'client_id', return_value=1):
mesh_2 = mesh_util.create_distributed_mesh(
mesh_name='multi_client_1d_mesh_2',
mesh_dims=[('x', 4)],
local_devices=['CPU:2', 'CPU:3'])
self.assertEqual(mesh_2.num_local_devices(), 2)
self.assertEqual(mesh_2.size, 4)
class InitializedMeshUtilTest(test_util.DTensorBaseTest):
"""Tests for mesh_util that require accelerator initialization."""
def setUp(self):
super().setUp()
device_type = config.preferred_device_type()
accelerator_util.initialize_accelerator_system(device_type)
def tearDown(self):
super().tearDown()
context._reset_context() # pylint: disable=protected-access
def test_is_initialized(self):
self.assertTrue(accelerator_util.is_initialized())
def test_initialize_accelerator_system(self):
accelerator_util.shutdown_accelerator_system()
device_type = accelerator_util.initialize_accelerator_system('CPU')
self.assertEqual(device_type, 'CPU')
# Default uses preferred_device_type.
accelerator_util.shutdown_accelerator_system()
device_type = accelerator_util.initialize_accelerator_system()
self.assertEqual(device_type, config.preferred_device_type())
@mock.patch.dict(os.environ, {'DTENSOR_GPU_USE_NCCL_COMMUNICATION': '1'})
def test_initialize_error_vgpu_with_nccl(self):
self.skipForDeviceType(['CPU', 'TPU'], reason='Test is intended for GPUs')
accelerator_util.shutdown_accelerator_system()
num_physical_devices = config.num_local_devices('GPU')
test_util.reset_logical_devices('GPU', 2 * num_physical_devices)
with self.assertRaisesRegex(ValueError,
'DTENSOR_GPU_USE_NCCL_COMMUNICATION'):
accelerator_util.initialize_accelerator_system('GPU')
@mock.patch.dict(os.environ, {'DTENSOR_GPU_USE_NCCL_COMMUNICATION': '1'})
def test_initialize_with_nccl(self):
self.skipForDeviceType(['CPU', 'TPU'], reason='Test is intended for GPUs')
accelerator_util.shutdown_accelerator_system()
accelerator_util.initialize_accelerator_system('GPU')
num_devices = len(test_util.list_local_logical_devices('GPU'))
mesh = mesh_util.create_mesh([('dim', num_devices)], device_type='GPU')
# The following shall run, but there is no clear way to check if it uses
# Collectives backed by NCCL.
mesh_util.barrier(mesh)
def test_initialize_after_tensorflow(self):
accelerator_util.shutdown_accelerator_system()
context.ensure_initialized()
with self.assertRaisesRegex(ValueError,
'TensorFlow has already been initialized'):
accelerator_util.initialize_accelerator_system('CPU')
def test_initialize_after_tensorflow_with_reset(self):
accelerator_util.shutdown_accelerator_system()
test_util.reset_logical_devices('CPU', 32)
context.ensure_initialized()
with self.assertLogs(level='WARNING') as log:
accelerator_util.initialize_accelerator_system(
'CPU', experimental_reset_context=True
)
self.assertIn('experimental_reset_context', log[0][0].message)
# Preserves the original logical device setting.
self.assertLen(test_util.list_local_logical_devices('CPU'), 32)
@parameterized.parameters(
dict(
device_type='CPU', skip_for=[]
), # We can create CPU meshes on TPU and GPU platforms!
dict(device_type='GPU', skip_for=['CPU', 'TPU']),
dict(device_type='TPU', skip_for=['CPU', 'GPU']),
)
def test_initialize_with_manual_logical_cpu_devices(
self, device_type: str, skip_for: list[str]
):
self.skipForDeviceType(
skip_for,
reason=f'Test is not intended for {skip_for}',
)
accelerator_util.shutdown_accelerator_system()
test_util.reset_logical_devices('CPU', 1)
accelerator_util.initialize_accelerator_system(
device_type, num_logical_cpu_devices=32
)
self.assertLen(test_util.list_local_logical_devices('CPU'), 32)
def test_shutdown_accelerator_system(self):
self.assertTrue(accelerator_util.is_initialized())
accelerator_util.shutdown_accelerator_system()
self.assertFalse(accelerator_util.is_initialized())
with self.assertRaisesRegex(ValueError, 'not initialized'):
accelerator_util.shutdown_accelerator_system()
def test_distributed_tpu_mesh_creation(self):
self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs')
self.skipForDeviceType(['TPU'],
reason='Test requires exactly 8 cores',
unless_device_count_equals_to=8)
num_devices = len(test_util.list_local_logical_devices('TPU'))
mesh = mesh_util.create_distributed_mesh(
mesh_name='distributed_1d_mesh',
mesh_dims=[('x', num_devices)],
device_type='TPU')
self.assertEqual(mesh.num_local_devices(), 8)
self.assertEqual(mesh.size, 8)
def test_mesh_barrier(self):
device_type = config.preferred_device_type()
num_devices = len(test_util.list_local_logical_devices(device_type))
mesh = mesh_util.create_mesh([('dim', num_devices)],
device_type=device_type)
# FIXME(b/235416015): To really test this we'll need a new eager async
# API. The following shall run, but the barrier semantics is not tested.
mesh_util.barrier(mesh, 'Name')
mesh_util.barrier(mesh)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/mnist_test.py
+197
View File
@@ -0,0 +1,197 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DTensor MNIST test."""
from absl.testing import parameterized
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import input_util
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import backprop
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
_BATCH_DIM = 'batch'
_DEVICE_IDS = test_util.create_device_ids_array((2,))
_ONE_D_MESH = layout_lib.Mesh(
[_BATCH_DIM],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2,), 'CPU'),
)
_ONE_D_TPU_MESH = layout_lib.Mesh(
[_BATCH_DIM],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2,), 'TPU'),
)
_BATCH_SIZE = 1024
_STEPS = 5
_LR = 1e-3
_ATOL = 1 # absolute error becomes large as gradients approach zero.
_RTOL = 1e-3
Layout = layout_lib.Layout
def mnist_fake_dataset():
imgs = []
labels = []
for i in range(_STEPS * _BATCH_SIZE):
img = stateless_random_ops.stateless_random_uniform(
shape=(28, 28, 1),
seed=[1, i],
minval=0,
maxval=256,
dtype=dtypes.float32,
)
imgs.append(img)
label = stateless_random_ops.stateless_random_uniform(
shape=(1,), seed=[2, i], minval=0, maxval=10, dtype=dtypes.int64
)
labels.append(label)
return dataset_ops.DatasetV2.from_tensor_slices(
(array_ops_stack.stack(imgs), array_ops_stack.stack(labels))
)
def _run_step(inputs, w, b, k):
with backprop.GradientTape() as g:
g.watch([w, b])
logits = nn_ops.conv2d_v2(inputs, k, strides=[1, 1, 1, 1], padding='SAME')
logits = array_ops.reshape(logits, [logits.shape[0], -1])
logits = math_ops.matmul(logits, w)
logits = logits + b
loss = math_ops.reduce_sum(logits, axis=[0, 1])
gw, gb = g.gradient(loss, [w, b])
for v, v_grad in zip([w, b], [gw, gb]):
v.assign_sub(_LR * v_grad)
return gw, gb, loss
class DTensorMNISTTest(test_util.DTensorBaseTest):
def setUp(self):
super(DTensorMNISTTest, self).setUp()
global_ids = test_util.create_device_ids_array((2,))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout_lib.Mesh(
[_BATCH_DIM],
global_ids,
local_ids,
test_util.create_device_list((2,), device),
)
for device in ['TPU', 'GPU', 'CPU']
}
self.mesh = self.configTestMesh(mesh_dict)
def init_var(self, mesh):
# Initialize TF randon normal variables(without using DTensor).
w_initializer = stateless_random_ops.stateless_random_normal(
shape=[28 * 28, 10], seed=[0, 1]
)
b_initializer = stateless_random_ops.stateless_random_normal(
shape=[10], seed=[1, 2]
)
# A filter with 3x3 shape, 1 input channel and 1 output channel.
k_initializer = stateless_random_ops.stateless_random_normal(
[3, 3, 1, 1], seed=[2, 3]
)
n_w = variables.Variable(w_initializer)
n_b = variables.Variable(b_initializer)
n_k = variables.Variable(k_initializer)
# Initialize DTensor variables.
w_initializer_on_mesh = api.copy_to_mesh(
w_initializer, Layout.replicated(mesh, 2)
)
b_initializer_on_mesh = api.copy_to_mesh(
b_initializer, Layout.replicated(mesh, rank=1)
)
k_initializer_on_mesh = api.copy_to_mesh(
k_initializer, Layout.replicated(mesh, rank=4)
)
w = d_variable.DVariable(w_initializer_on_mesh)
b = d_variable.DVariable(b_initializer_on_mesh)
k = d_variable.DVariable(k_initializer_on_mesh)
return (n_w, n_b, n_k), (w, b, k)
@parameterized.named_parameters(('Eager', False), ('Function', True))
def testMnist(self, on_function):
mnist_dataset = mnist_fake_dataset()
(n_w, n_b, n_k), (w, b, k) = self.init_var(self.mesh)
n_dataset = mnist_dataset.batch(_BATCH_SIZE, drop_remainder=True)
n_iter = iter(n_dataset)
input_layout = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=4)
label_layout = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
dtensor_dataset = input_util.DTensorDataset(
dataset=mnist_dataset,
global_batch_size=_BATCH_SIZE,
mesh=self.mesh,
layouts=(input_layout, label_layout),
batch_dim=_BATCH_DIM,
)
dtensor_iter = iter(dtensor_dataset)
step_fn = (
polymorphic_function.function(_run_step) if on_function else _run_step
)
# Training loop.
for _ in range(_STEPS):
# Normal run without DTensor.
n_input, _ = next(n_iter)
g_nw, g_nb, n_loss = step_fn(n_input, n_w, n_b, n_k)
# DTensor Run
dtensor_input, _ = next(dtensor_iter)
with ops.device_v2(api.device_name()):
gw, gb, loss = step_fn(dtensor_input, w, b, k)
loss_unpack = api.unpack(loss)
self.assertAllEqual(loss_unpack[0], loss_unpack[1])
self.assertAllClose(n_loss, loss, atol=_ATOL, rtol=_RTOL)
self.assertAllClose(g_nw, gw, atol=_ATOL, rtol=_RTOL)
self.assertAllClose(g_nb, gb, atol=_ATOL, rtol=_RTOL)
self.assertAllClose(n_w, w, atol=_ATOL, rtol=_RTOL)
self.assertAllClose(n_b, b, atol=_ATOL, rtol=_RTOL)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/multi_client_input_util_test.py
+548
View File
@@ -0,0 +1,548 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Multi-client tests for input_util."""
import os
from typing import Any, List, Mapping, Optional, Tuple
from absl import logging
from absl.testing import parameterized
import numpy as np
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import input_util
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import mesh_util
from tensorflow.dtensor.python.tests import multi_client_test_util
from tensorflow.dtensor.python.tests import test_backend_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.data.experimental.service import server_lib
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import readers
from tensorflow.python.eager import context
from tensorflow.python.framework import config as tf_config
from tensorflow.python.framework import device_spec
from tensorflow.python.framework import dtypes
from tensorflow.python.lib.io import tf_record
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import gen_parsing_ops
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import parsing_config
from tensorflow.python.ops import parsing_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.platform import test
mp_context = test_backend_util.get_mp_context()
# Multi-client test constants.
JOB_NAME = 'worker'
TF_DATA_SERVICE_JOB_NAME = 'dtensor_tf_data'
NUM_CLIENTS = 4
NUM_DEVICES_PER_CLIENT = 4
# Mesh constants.
MESH_DIM_BATCH = 'batch'
MESH_DIM_HEIGHT = 'height'
MESH_DIM_WIDTH = 'width'
# Data constants.
IMG_HEIGHT = 8
IMG_WIDTH = 8
IMG_CHANNELS = 3
UNSHARDED = layout_lib.UNSHARDED
Mesh = layout_lib.Mesh
Layout = layout_lib.Layout
def redirect_output(file_name):
# Redirect stderr/stdout to undeclared outputs on sponge.
artifact_dir = os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', '')
if artifact_dir:
with open(os.path.join(artifact_dir, file_name), 'wb') as fp:
os.dup2(fp.fileno(), 1)
os.dup2(fp.fileno(), 2)
def create_dispatcher(test_name, worker_addresses, port, pipe=None):
dispatcher = server_lib.DispatchServer(
config=server_lib.DispatcherConfig(
port=port, protocol='grpc', worker_addresses=worker_addresses
)
)
dispatcher.start()
if pipe is None:
# Dispatcher is not within subprocess, so do not block.
return dispatcher, dispatcher._address
else:
redirect_output(f'test-{test_name}-dispatcher.log')
pipe.send(dispatcher._address)
signal = pipe.recv() # blocks until a 'stop' signal is received
if signal == 'stop':
dispatcher._stop()
pipe.send('stopped')
else:
raise ValueError('Got unknown signal %s' % signal)
def create_worker(test_name, dispatcher_address, port=None, pipe=None):
worker = server_lib.WorkerServer(
config=server_lib.WorkerConfig(
port=port, dispatcher_address=dispatcher_address, protocol='grpc'
)
)
worker.start()
if pipe is None:
# Worker is not within subprocess, so do not block.
return worker, worker._address
else:
redirect_output(f'test-{test_name}-worker.log')
pipe.send(worker._address)
signal = pipe.recv() # blocks until a 'stop' signal is received
if signal == 'stop':
worker._stop()
pipe.send('stopped')
else:
raise ValueError('Got unknown signal %s' % signal)
class TFDataServiceCluster:
"""tf.data service cluster with dispatcher and workers as subprocesses.
To run the cluster in co-located mode, set `num_workers` to 0 and create the
tf.data service workers manually in each client process.
"""
def __init__(self,
test_name,
num_workers,
worker_ports=None,
worker_addresses=None):
self._test_name = test_name
self._num_workers = num_workers
self._start_dispatcher(worker_addresses)
self._start_workers(worker_ports)
def _start_dispatcher(self, worker_addresses, port=0):
self._pipe_to_dispatcher, dispatcher_pipe = mp_context.Pipe(True)
logging.info(
'Starting remote dispatcher on port %d with worker addresses: %s', port,
worker_addresses)
self._dispatcher_process = mp_context.Process(
target=create_dispatcher,
args=(self._test_name, worker_addresses, port, dispatcher_pipe),
)
self._dispatcher_process.start()
self._dispatcher_address = self._pipe_to_dispatcher.recv()
def dispatcher_address(self):
return self._dispatcher_address
def _start_workers(self, worker_ports=None):
self._workers = []
self._worker_addresses = []
self._worker_pipes = []
for idx in range(self._num_workers):
port = worker_ports[idx] if worker_ports else None
self._start_worker(port)
def _start_worker(self, port=None):
pipe_to_worker, worker_pipe = mp_context.Pipe(True)
logging.info(
'Starting remote worker on port %d with dispatcher address: %s', port,
self._dispatcher_address)
worker_process = mp_context.Process(
target=create_worker,
args=(self._test_name, self._dispatcher_address, port, worker_pipe),
)
worker_process.start()
worker_address = self._pipe_to_worker.recv()
self._workers.append(worker_process)
self._worker_addresses.append(worker_address)
self._worker_pipes.append(pipe_to_worker)
def worker_addresses(self):
return self._worker_addresses
def stop(self):
# Segfault logs may still be printed because clean exit of child processes
# is not always possible. This will not affect the outcome of the test.
logging.info('Will try to stop TFDataServiceCluster!')
for idx in range(self._num_workers):
address = self._worker_addresses[idx]
pipe_to_worker = self._worker_pipes[idx]
logging.info('Stopping worker %s...', address)
pipe_to_worker.send('stop')
if pipe_to_worker.poll(2):
if pipe_to_worker.recv() == 'stopped':
logging.info('Successfully stopped worker %s', address)
self._workers[idx].terminate()
logging.info('Stopping dispatcher...')
self._pipe_to_dispatcher.send('stop')
if self._pipe_to_dispatcher.poll(2):
if self._pipe_to_dispatcher.recv() == 'stopped':
logging.info('Successfully stopped dispatcher')
self._dispatcher_process.terminate()
def setup_local_devices(num_devices):
physical_cpus = tf_config.list_physical_devices('CPU')
tf_config.set_logical_device_configuration(
physical_cpus[0],
[context.LogicalDeviceConfiguration() for _ in range(num_devices)],
)
def setup_client(client_id: int, test_name: str, env: Mapping[str, str],
num_local_devices: int):
"""Set up a DTensor client for use in multi-client tests.
Args:
client_id: the index of the client.
test_name: the name of the test under which this client is running, used To
identify the log file artifact containing the test output.
env: a dictionary of environment variables to update.
num_local_devices: number of local devices to set up.
"""
# Redirect client's stderr/stdout to undeclared outputs on sponge.
redirect_output(f'test-{test_name}-process-{client_id}.log')
# Update any specified environment variables.
for var, val in env.items():
os.environ[var] = val
# Set up local devices.
setup_local_devices(num_local_devices)
# Set up DTensor cluster and enable collectives.
accelerator_util.initialize_accelerator_system()
def run_client(
client_id: int,
test_name: str,
env: Mapping[str, str],
num_local_devices: int,
dispatcher_address: str,
worker_port: int,
batch_size: int,
dataset_paths: List[str],
mesh: Mesh,
batch_dim: Optional[str],
layouts: Tuple[Layout, Layout],
) -> List[Tuple[Any, Any]]:
# Co-located tf.data service mode. It is important to hold the worker object
# until the end otherwise it will get garbage collected.
worker, worker_address = create_worker( # pylint: disable=unused-variable
test_name, dispatcher_address, port=worker_port)
logging.info(
'tf.data service worker running at %s',
worker_address,
)
setup_client(client_id, test_name, env, num_local_devices)
def decode_fn(record_bytes):
decoded = parsing_ops.parse_single_example_v2(
serialized=record_bytes,
features={
'idx': parsing_config.FixedLenFeature([], dtype=dtypes.int64),
'elem': parsing_config.FixedLenFeature([], dtype=dtypes.string),
},
)
parsed_elem = gen_parsing_ops.parse_tensor(decoded['elem'], dtypes.int32)
elem = check_ops.ensure_shape(
parsed_elem, [IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS]
)
return decoded['idx'], elem
dataset = dataset_ops.DatasetV2.from_tensor_slices(dataset_paths)
dataset = dataset.interleave(readers.TFRecordDatasetV2)
dataset = dataset.map(decode_fn)
tf_data_service_config = input_util.TFDataServiceConfig(
dispatcher_address=dispatcher_address, job_name=TF_DATA_SERVICE_JOB_NAME
)
d_dataset = input_util.DTensorDataset(
dataset=dataset,
global_batch_size=batch_size,
mesh=mesh,
layouts=layouts,
batch_dim=batch_dim,
tf_data_service_config=tf_data_service_config,
)
# Subprocesses cannot return a sharded DTensor as it triggers a copy and
# copying non-replicated DTensors is not supported. So instead we unpack it
# and return the component tensors.
ret = []
for batch_idx, elem in d_dataset:
n_batch_idx = api.unpack(batch_idx)
n_elem = api.unpack(elem)
ret.append((n_batch_idx, n_elem))
return ret
class MultiClientDTensorDatasetTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
logging.info('Check per client log in Test artifacts.')
self.server_ports = [
multi_client_test_util.pick_unused_port() for _ in range(NUM_CLIENTS)
]
self.worker_ports = [
multi_client_test_util.pick_unused_port() for _ in range(NUM_CLIENTS)
]
worker_addresses = [f'localhost:{port}' for port in self.worker_ports]
self.cluster = TFDataServiceCluster(
test_name=self._testMethodName,
num_workers=0, # Co-located mode.
worker_addresses=worker_addresses)
def tearDown(self):
super().tearDown()
self.cluster.stop()
def write_dataset(self, dataset, num_files, num_elems):
"""Writes a dataset_ops.DatasetV2 to multiple files."""
dataset_paths = []
dataset_iter = iter(dataset)
for file_idx in range(num_files):
dataset_path = os.path.join(self.get_temp_dir(),
f'dataset-{file_idx}.tfrecords')
dataset_paths.append(dataset_path)
with tf_record.TFRecordWriter(dataset_path) as writer:
for _ in range(num_elems // num_files):
idx, elem = next(dataset_iter)
elem_bytes = example_pb2.Example(
features=feature_pb2.Features(
feature={
'idx': feature_pb2.Feature(
int64_list=feature_pb2.Int64List(value=[idx])
),
'elem': feature_pb2.Feature(
bytes_list=feature_pb2.BytesList(
value=[io_ops.serialize_tensor(elem).numpy()]
)
),
}
)
).SerializeToString()
writer.write(elem_bytes)
return dataset_paths
@parameterized.product(
(
{
# batch=4 x height=2 x width=2
# 1 replica per client.
'mesh_dims': [(MESH_DIM_BATCH, 4),
(MESH_DIM_HEIGHT, 2),
(MESH_DIM_WIDTH, 2)],
}, {
# batch=4 x height=2 x width=2 (transposed)
# 1 replica per client with reordered local partitions.
'mesh_dims': [(MESH_DIM_BATCH, 4),
(MESH_DIM_WIDTH, 2),
(MESH_DIM_HEIGHT, 2)],
}, {
# batch=8 x height=2 x width=1
# 2 replicas per client.
'mesh_dims': [(MESH_DIM_BATCH, 8),
(MESH_DIM_HEIGHT, 2),
(MESH_DIM_WIDTH, 1)],
}, {
# batch=8 x height=2 x width=1 (transposed)
# 2 replicas per client with reordered partitions.
'mesh_dims': [(MESH_DIM_BATCH, 8),
(MESH_DIM_WIDTH, 1),
(MESH_DIM_HEIGHT, 2)],
}, {
# batch=2 x height=4 x width=2
# 1 replica split over 2 clients.
'mesh_dims': [(MESH_DIM_BATCH, 2),
(MESH_DIM_HEIGHT, 4),
(MESH_DIM_WIDTH, 2)],
}, {
# batch=2 x height=4 x width=2 (transposed)
# 1 replica split over 2 clients with reordered partitions.
'mesh_dims': [(MESH_DIM_BATCH, 2),
(MESH_DIM_WIDTH, 2),
(MESH_DIM_HEIGHT, 4)],
},
),
(
{
# Replicated
'idx_sharding': [UNSHARDED],
'images_sharding': [UNSHARDED, UNSHARDED, UNSHARDED, UNSHARDED],
}, {
# Batch sharded
'idx_sharding': [MESH_DIM_BATCH],
'images_sharding':
[MESH_DIM_BATCH, UNSHARDED, UNSHARDED, UNSHARDED],
}, {
# Spatially sharded
'idx_sharding': [UNSHARDED],
'images_sharding':
[UNSHARDED, MESH_DIM_HEIGHT, MESH_DIM_WIDTH, UNSHARDED],
}, {
# Batch and spatially sharded
'idx_sharding': [MESH_DIM_BATCH],
'images_sharding':
[MESH_DIM_BATCH, MESH_DIM_HEIGHT, MESH_DIM_WIDTH, UNSHARDED],
}
))
def testMultiClientIter(self, mesh_dims, idx_sharding, images_sharding):
num_batches = 4
batch_size = 16
num_elems = num_batches * batch_size
images = stateless_random_ops.stateless_random_uniform(
[num_elems, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS],
seed=(1, 2),
minval=0,
maxval=255,
dtype=dtypes.int32,
)
dataset = dataset_ops.DatasetV2.from_tensor_slices(images)
# Enumerate the dataset elements to make it easier to identify the batches
# returned by the DTensorDataset.
dataset = dataset.enumerate()
# Store a mapping of index to dataset elements which can be looked up later
# to identify the batches returned by the DTensorDataset.
all_elems = {idx.numpy(): elem for idx, elem in dataset}
# Write the dataset and shard it among multiple files.
dataset_paths = self.write_dataset(
dataset, num_files=8, num_elems=num_elems)
# Construct args for starmap.
args = []
mesh_dim_names, mesh_dim_sizes = zip(*mesh_dims)
global_device_ids = test_util.create_device_ids_array(mesh_dim_sizes)
device_ids_split = np.split(np.ravel(global_device_ids), NUM_CLIENTS)
dtensor_jobs = [
f'localhost:{self.server_ports[i]}' for i in range(NUM_CLIENTS)
]
for client_id in range(NUM_CLIENTS):
# Manually specify DTensor environment variables since we are in a test
# environment.
env = {
config._DT_CLIENT_ID: str(client_id),
config._DT_JOB_NAME: str(JOB_NAME),
config._DT_JOBS: ','.join(dtensor_jobs)
}
local_device_ids = device_ids_split[client_id].tolist()
local_devices = [
device_spec.DeviceSpecV2( # pylint: disable=g-complex-comprehension
job=JOB_NAME,
replica=0,
task=client_id,
device_type='CPU',
device_index=i,
)
for i in range(len(local_device_ids))
]
mesh = Mesh(
dim_names=mesh_dim_names,
global_device_ids=global_device_ids,
local_device_ids=local_device_ids,
local_devices=local_devices,
)
idx_layout = Layout(idx_sharding, mesh)
images_layout = Layout(images_sharding, mesh)
batch_dim = MESH_DIM_BATCH if MESH_DIM_BATCH in images_sharding else None
args.append((client_id, self._testMethodName, env, NUM_DEVICES_PER_CLIENT,
self.cluster.dispatcher_address(),
self.worker_ports[client_id], batch_size, dataset_paths,
mesh, batch_dim, (idx_layout, images_layout)))
def get_results():
# Run the DTensor client processes and get the DTensor dataset components.
with mp_context.Pool(NUM_CLIENTS) as pool:
results = pool.starmap(run_client, args)
pool.close()
pool.join()
return results
# TODO(b/271162918): fix multi-client use case.
with self.assertRaises(NotImplementedError):
results = get_results()
return
# pylint: disable=unreachable
# Create a mesh on the main test process. The tensor components returned
# from each DTensor client subprocess will be packed onto this mesh to
# verify correctness.
test_mesh = mesh_util.create_mesh(
mesh_dims=mesh_dims,
devices=[
'CPU:%d' % i for i in range(NUM_CLIENTS * NUM_DEVICES_PER_CLIENT)
])
test_mesh = self.configTestMesh({'CPU': test_mesh})
idx_test_layout = Layout(idx_sharding, test_mesh)
images_test_layout = Layout(images_sharding, test_mesh)
for batch_elems in zip(*results):
# Collect the tensor components returned from each client.
idx_components = []
images_components = []
for client_id in range(NUM_CLIENTS):
local_idx, local_images = batch_elems[client_id]
idx_components.extend(local_idx)
images_components.extend(local_images)
# Pack the dataset elements into a DTensor on the test mesh.
d_idx = api.pack(idx_components, idx_test_layout)
d_images = api.pack(images_components, images_test_layout)
# Get the batch of elements from the original dataset using the element
# indices.
batch_stack = []
for elem_idx in d_idx:
batch_stack.append(all_elems.pop(elem_idx.numpy()))
batch = array_ops_stack.stack(batch_stack)
self.assertDTensorEqual(batch, images_test_layout, d_images)
self.assertEmpty(
all_elems, 'Not all batches were returned by DTensorDataset.')
if __name__ == '__main__':
test_backend_util.handle_test_main(test.main)
tensorflow/dtensor/python/tests/multi_client_test.py
+242
View File
@@ -0,0 +1,242 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for DTensor multi-client setup."""
import os
from absl import flags
import numpy as np
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import api as d_api
from tensorflow.dtensor.python import config as d_config
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import layout as d_layout
from tensorflow.dtensor.python import mesh_util
from tensorflow.dtensor.python.tests import multi_client_test_util
from tensorflow.dtensor.python.tests import test_backend_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import config
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test as tf_test
_MODEL_DIM_SIZE = flags.DEFINE_integer('model_dim_size', 4,
'Size of the model dimension.')
_BATCH_DIM = 'batch'
_MODEL_DIM = 'model'
_BATCH_SIZE = 8
_STEPS = 5
_LR = 1e-3
@polymorphic_function.function
def _run_step(inputs, w, b, k):
with backprop.GradientTape() as g:
g.watch([w, b])
logits = nn_ops.conv2d_v2(inputs, k, strides=[1, 1, 1, 1], padding='SAME')
logits = array_ops.reshape(logits, [logits.shape[0], -1])
logits = math_ops.matmul(logits, w)
logits = logits + b
loss = math_ops.reduce_sum(logits, axis=[0, 1])
gw, gb = g.gradient(loss, [w, b])
for v, v_grad in zip([w, b], [gw, gb]):
v.assign_sub(_LR * v_grad)
return gw, gb, loss
def init_var(mesh):
w_initializer = stateless_random_ops.stateless_random_normal(
[28 * 28, 16], seed=[0, 1]
)
b_initializer = stateless_random_ops.stateless_random_normal(
[16], seed=[0, 2]
)
k_initializer = stateless_random_ops.stateless_random_normal(
[3, 3, 1, 1], seed=[0, 3]
)
n_w = variables.Variable(w_initializer)
n_b = variables.Variable(b_initializer)
n_k = variables.Variable(k_initializer)
# Initialize DTensor variables.
w_initializer_on_mesh = d_api.copy_to_mesh(
w_initializer, d_layout.Layout.replicated(mesh, rank=2)
)
b_initializer_on_mesh = d_api.copy_to_mesh(
b_initializer, d_layout.Layout.replicated(mesh, rank=1)
)
k_initializer_on_mesh = d_api.copy_to_mesh(
k_initializer, d_layout.Layout.replicated(mesh, rank=4)
)
w = d_variable.DVariable(
d_api.relayout(
w_initializer_on_mesh,
d_layout.Layout(['unsharded', _MODEL_DIM], mesh),
)
)
b = d_variable.DVariable(
d_api.relayout(b_initializer_on_mesh, d_layout.Layout([_MODEL_DIM], mesh))
)
k = d_variable.DVariable(k_initializer_on_mesh)
return (n_w, n_b, n_k), (w, b, k)
class DTensorMNISTTest(tf_test.TestCase):
def setUp(self):
super(DTensorMNISTTest, self).setUp()
if config.list_physical_devices('GPU'):
device_type = 'GPU'
elif test_util.is_tpu_present():
device_type = 'TPU'
else:
device_type = 'CPU'
local_devices = d_config.local_devices(device_type)
num_devices = len(local_devices)
global_device_ids = test_util.create_device_ids_array((
d_config.num_clients() * num_devices // _MODEL_DIM_SIZE.value,
_MODEL_DIM_SIZE.value,
))
device_ids_list = np.ravel(global_device_ids).tolist()
index = d_config.client_id() * num_devices
local_device_ids = device_ids_list[index:(index + num_devices)]
self.mesh = d_layout.Mesh(
[_BATCH_DIM, _MODEL_DIM],
global_device_ids=global_device_ids,
local_device_ids=local_device_ids,
local_devices=local_devices,
)
def tearDown(self):
# A barrier prevents a client from disconnecting prematurely in tests.
mesh_util.barrier(self.mesh)
super().tearDown()
def test_mnist(self):
def train():
input_layout = d_layout.Layout.batch_sharded(
self.mesh, _BATCH_DIM, rank=4
)
(n_w, n_b, n_k), (w, b, k) = init_var(self.mesh)
for i in range(_STEPS):
data = stateless_random_ops.stateless_random_normal(
[_BATCH_SIZE, 28, 28, 1], seed=[0, i]
)
g_nw, g_nb, n_loss = _run_step(data.numpy(), n_w, n_b, n_k)
input_image = d_api.copy_to_mesh(
data, layout=d_layout.Layout.replicated(self.mesh, rank=4)
)
input_image = d_api.relayout(input_image, layout=input_layout)
with ops.device_v2(self.mesh.local_devices()[0]):
gw, gb, loss = _run_step(input_image, w, b, k)
gw = d_api.relayout(gw, d_layout.Layout.replicated(self.mesh, rank=2))
w = d_api.relayout(w, d_layout.Layout.replicated(self.mesh, rank=2))
gb = d_api.relayout(gb, d_layout.Layout.replicated(self.mesh, rank=1))
b = d_api.relayout(b, d_layout.Layout.replicated(self.mesh, rank=1))
return (n_loss, g_nw, g_nb, n_w, n_b), (loss, gw, gb, w, b)
(n_loss, g_nw, g_nb, n_w, n_b), (loss, gw, gb, w, b) = train()
self.assertAllClose(n_loss, loss, atol=5e-4)
self.assertAllClose(g_nw, gw, atol=1e-5)
self.assertAllClose(g_nb, gb, atol=1e-5)
self.assertAllClose(n_w, w, atol=1e-5)
self.assertAllClose(n_b, b, atol=1e-5)
def test_copy_to_mesh(self):
layout = d_layout.Layout([_BATCH_DIM, 'unsharded'], self.mesh)
host_layout = d_layout.Layout(layout.sharding_specs, self.mesh.host_mesh())
x = d_api.pack(
[array_ops.ones((2, 2), dtype=dtypes.float32)]
* len(self.mesh.local_devices()),
layout,
)
@polymorphic_function.function
def d2h(x):
return d_api.copy_to_mesh(x, host_layout)
@polymorphic_function.function
def h2d(x):
return d_api.copy_to_mesh(x, layout)
y = d2h(x)
ys = d_api.unpack(y)
for i in ys:
self.assertAllClose(i, array_ops.ones((2, 2)), atol=1e-5)
z = h2d(y)
zs = d_api.unpack(z)
for i in zs:
self.assertAllClose(i, array_ops.ones((2, 2)), atol=1e-5)
def client_config_function(config_params):
num_clients = config_params['num_clients']
dtensor_client_id = config_params['client_id']
dtensor_jobs = config_params['worker_jobs']
num_devices = config_params['num_devices']
# pylint: disable=protected-access
if num_clients != 0:
os.environ[d_config._DT_CLIENT_ID] = f'{dtensor_client_id}'
os.environ[d_config._DT_JOB_NAME] = 'worker'
os.environ[d_config._DT_JOBS] = ','.join(dtensor_jobs)
# pylint: enable=protected-access
if config.list_physical_devices('GPU'):
device_type = 'GPU'
elif test_util.is_tpu_present():
device_type = 'TPU'
else:
device_type = 'CPU'
# reset_logical_devices
test_util.reset_context()
if device_type != 'TPU':
# Configure virtual devices. This does not initialize the TensorFlow
# context.
test_util.reset_logical_devices(device_type, num_devices)
accelerator_util.initialize_accelerator_system(
device_type, enable_coordination_service=True)
# Validates the correct number of devices are created.
logical_devices = test_util.list_local_logical_devices(device_type)
assert len(logical_devices) == num_devices, (
logical_devices,
f'Test is mis-configured: expecting {num_devices} logical_devices.')
if __name__ == '__main__':
test_backend_util.handle_test_main(
multi_client_test_util.multi_client_main, client_config_function)
tensorflow/dtensor/python/tests/multi_client_test_util.py
+144
View File
@@ -0,0 +1,144 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for multi-client setup."""
import os
import sys
from absl import flags
import portpicker
from tensorflow.dtensor.python.tests import test_backend_util
from tensorflow.python.platform import test as tf_test
_NUM_LOCAL_DEVICES = flags.DEFINE_integer(
'num_local_devices', 4,
'Number of local devices. 4 is the only allowed value for TPU.')
_NUM_CLIENTS = flags.DEFINE_integer(
'num_clients', 2,
'Number of clients. 0 for local mode. 2 is the only allowed value for TPU.')
def pick_unused_port():
"""Helper function to return an unused port."""
return portpicker.pick_unused_port()
def multi_client_main(client_config_function):
"""Creates a Flock of TensorFlow Processes on localhost."""
flags.FLAGS(sys.argv, known_only=True)
num_clients = _NUM_CLIENTS.value
num_process = num_clients or 1
num_local_devices = _NUM_LOCAL_DEVICES.value
# No GPU visible to the flock controller.
os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ['HIP_VISIBLE_DEVICES'] = ''
# Python multiprocess module in OSS.
mp_context = test_backend_util.get_mp_context()
print('Check per client log in Test artifacts.', flush=True)
# Inverts the order of ports intentionally to rule out ordering bugs.
server_ports = sorted(
[pick_unused_port() for _ in range(num_process)], reverse=True
)
additional_ports = sorted([pick_unused_port() for _ in range(num_process)])
# Starts processes
procs = []
for client_idx in range(num_process):
proc = mp_context.Process(
target=run_client,
args=(client_idx, num_clients, server_ports, additional_ports,
num_local_devices, client_config_function),
name=f'Client-{client_idx}',
)
proc.start()
procs.append(proc)
# Joins processes
exitcode = 0
for proc in procs:
proc.join()
if proc.exitcode != 0:
exitcode = proc.exitcode
sys.exit(exitcode)
def run_client(idx, num_clients, server_ports, additional_ports,
num_local_devices, client_config_function):
"""Runs test.main() from a DTensor Client process on localhost.
This function runs in a separate process so that the eager context is
properly separated, which resembles real world multi-client setup.
Virtual devices are configured before test.main() is called.
Each client is configured to only have access to the physical GPU device
corresponding to its client id via CUDA_VISIBLE_DEVICES/HIP_VISIBLE_DEVICES.
Each client is configured to only have access to some TPU cores
corresponding to its client id via flags.
The clients redirect stdout and stderr to files under Test Artifacts.
Args:
idx: integer task number represents the client's id from global picture.
num_clients: total number of clients.
server_ports: A list of ports that is allocated and to be used to construct
GRPC server. server_ports[idx] will be the GRPC server on the
corresponding client.
additional_ports: A list of ports that is allocated and to be used to
construct the backends.
num_local_devices: Number of devices per client.
client_config_function: A function, for each of the client to config the
local environment variables, etc. Note that the function will be called
with a dict of extra params, eg:
{'num_clients': 2
'client_id': 0,
'worker_jobs': ['localhost:port1', 'localhost:port2'],
'num_devices': 4,
}
"""
test_backend_util.slice_host_devices_for_multiworker(
num_clients, idx, additional_ports
)
artifact_dir = os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', '')
# Redirect extra client's stderr/stdout to undeclared outputs on sponge.
if artifact_dir:
with open(
os.path.join(artifact_dir, f'test-client-process-{idx}.log'),
'wb') as fp:
os.dup2(fp.fileno(), 1)
os.dup2(fp.fileno(), 2)
# Set up cluster and enable collectives.
worker_jobs = [f'localhost:{port:06d}' for port in server_ports]
client_config_func_param = {
'num_clients': num_clients,
'client_id': idx,
'worker_jobs': worker_jobs,
'num_devices': num_local_devices,
}
client_config_function(client_config_func_param)
# The following function call never returns.
tf_test.main()
tensorflow/dtensor/python/tests/multi_mesh_test.py
+808
View File
@@ -0,0 +1,808 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for multiple meshes in DTensor."""
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
UNSHARDED = layout_lib.UNSHARDED
# Makes a 1D mesh with dimension X(2).
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
_DEVICES_IDS = test_util.create_device_ids_array((2,))
_ONE_D_CPU_MESH = Mesh(
[_MESH_DIM_X],
_DEVICES_IDS,
np.ravel(_DEVICES_IDS).tolist(),
test_util.create_device_list((2,), 'CPU'),
)
_ONE_D_GPU_MESH = Mesh(
[_MESH_DIM_X],
_DEVICES_IDS,
np.ravel(_DEVICES_IDS).tolist(),
test_util.create_device_list((2,), 'GPU'),
)
_ONE_D_TPU_MESH = Mesh(
[_MESH_DIM_X],
_DEVICES_IDS,
np.ravel(_DEVICES_IDS).tolist(),
test_util.create_device_list((2,), 'TPU'),
)
_ONE_D_CPU_MESH_Y = Mesh(
[_MESH_DIM_Y],
_DEVICES_IDS,
np.ravel(_DEVICES_IDS).tolist(),
test_util.create_device_list((2,), 'CPU'),
)
_OTHER_CPU_MESH = Mesh(
[_MESH_DIM_X],
_DEVICES_IDS[:1],
np.ravel(_DEVICES_IDS[:1]).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
class MultiMeshTest(test_util.DTensorBaseTest):
def setUp(self):
super(MultiMeshTest, self).setUp()
self.first_mesh = _ONE_D_CPU_MESH
if test_util.is_tpu_present():
self.second_mesh = _ONE_D_TPU_MESH
elif test_util.is_gpu_present():
self.second_mesh = _ONE_D_GPU_MESH
else:
self.second_mesh = _ONE_D_CPU_MESH_Y
device_type = config.preferred_device_type()
if device_type != 'TPU':
test_util.reset_logical_devices(device_type, 2)
accelerator_util.initialize_accelerator_system(device_type)
def testBasicCopyToMesh(self):
target_layout = Layout.replicated(self.first_mesh, rank=1)
numpy_value = np.zeros([3], dtype=np.int32)
dtensor_copy_from_numpy = api.copy_to_mesh(numpy_value, target_layout)
self.assertDTensorEqual(numpy_value, target_layout, dtensor_copy_from_numpy)
numpy_value = np.ones([3], dtype=np.int32)
src_mesh = api.copy_to_mesh(
numpy_value, Layout.replicated(self.second_mesh, rank=1)
)
dtensor_from_another_mesh = api.copy_to_mesh(src_mesh, target_layout)
self.assertDTensorEqual(
numpy_value, target_layout, dtensor_from_another_mesh
)
@parameterized.named_parameters(
dict(testcase_name='Graph', is_eager=False),
dict(testcase_name='Eager', is_eager=True),
)
def testCopyToMeshOneToOneSharded(self, is_eager):
if not test_util.is_tpu_present():
self.skipForDeviceType(
['CPU'], 'Need at least one device mesh for this test.'
)
replicated_layout = Layout.replicated(self.first_mesh, rank=1)
first_layout = Layout([_MESH_DIM_X], self.first_mesh)
second_layout = Layout([_MESH_DIM_X], self.second_mesh)
numpy_value = np.zeros([8], dtype=np.int32)
dt_value = api.copy_to_mesh(numpy_value, replicated_layout)
self.assertDTensorEqual(numpy_value, replicated_layout, dt_value)
def fn(val):
dt_source = api.relayout(val, first_layout)
dt_target = api.copy_to_mesh(dt_source, second_layout)
return dt_source, dt_target
if not is_eager:
fn = polymorphic_function.function(fn)
dt_source, dt_target = fn(dt_value)
self.assertDTensorEqual(numpy_value, first_layout, dt_source)
self.assertDTensorEqual(numpy_value, second_layout, dt_target)
@parameterized.named_parameters(
dict(testcase_name='Graph', is_eager=False),
dict(testcase_name='Eager', is_eager=True),
)
def testCopyToMeshToShardedLayout(self, is_eager):
target_layout = Layout([_MESH_DIM_X], self.first_mesh)
a_np = array_ops.zeros([8], dtype=dtypes.int32)
def fn(val):
return api.copy_to_mesh(val, target_layout)
if not is_eager:
fn = polymorphic_function.function(fn)
with api.default_mesh(self.first_mesh):
dt_value = fn(a_np)
self.assertDTensorEqual(a_np, target_layout, dt_value)
def testNestedDefaultMesh(self):
@polymorphic_function.function
def func(a):
return a + 3.0
with api.default_mesh(self.first_mesh):
with api.default_mesh(self.second_mesh):
with api.default_mesh(self.first_mesh):
result = func(array_ops.ones(shape=()))
self.assertEqual(api.fetch_layout(result).mesh, self.first_mesh)
result = func(array_ops.ones(shape=()))
self.assertEqual(api.fetch_layout(result).mesh, self.second_mesh)
result = func(array_ops.ones(shape=()))
self.assertEqual(api.fetch_layout(result).mesh, self.first_mesh)
def testImplicitCopyToCPUMeshForStrings(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
host_device_id = test_util.create_device_ids_array((1,))
host_cpu_mesh = Mesh(
[_MESH_DIM_X],
host_device_id,
np.ravel(host_device_id).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
replicated_layout_on_cpu = Layout.replicated(host_cpu_mesh, rank=0)
replicated_layout_on_tpu = Layout.replicated(self.second_mesh, rank=0)
string_tensor = constant_op.constant('hello')
@polymorphic_function.function
def f(tensor, dtensor_a, dtensor_b):
# Return an identity op for all three inputs so that linter does not
# complain about unused variables.
return tensor, dtensor_a, dtensor_b
cpu_dtensor = api.copy_to_mesh(
constant_op.constant(1), replicated_layout_on_cpu
)
tpu_dtensor = api.copy_to_mesh(
constant_op.constant(1), replicated_layout_on_tpu
)
string_dtensor, _, _ = f(string_tensor, cpu_dtensor, tpu_dtensor)
# Regular string tensor should be implicitly copied onto the CPU mesh,
# not the TPU mesh.
self.assertEqual(api.fetch_layout(string_dtensor), replicated_layout_on_cpu)
def testMultiMeshBroadcast(self):
first_mesh_a = api.copy_to_mesh(
np.zeros([3], dtype=np.int32),
Layout.replicated(self.first_mesh, rank=1),
)
second_mesh_a = api.copy_to_mesh(
np.ones([3], dtype=np.int32),
Layout.replicated(self.second_mesh, rank=1),
)
self.assertDTensorEqual(
np.asarray([1, 1, 1], dtype=np.int32),
Layout.replicated(self.first_mesh, rank=1), first_mesh_a + 1)
# Run an add with small constant - the constant should be broadcasted
# onto the second mesh rather than the first.
self.assertDTensorEqual(
np.asarray([2, 2, 2], dtype=np.int32),
Layout.replicated(self.second_mesh, rank=1), second_mesh_a + 1)
def testMultiMeshAdd(self):
a = constant_op.constant(1, dtype=dtypes.int32)
b = constant_op.constant(2, dtype=dtypes.int32)
with ops.device_v2(api.device_name()):
first_mesh_a = api.copy_to_mesh(
a, Layout.replicated(self.first_mesh, rank=0)
)
first_mesh_b = api.copy_to_mesh(
b, Layout.replicated(self.first_mesh, rank=0)
)
# Copy-to-mesh doesn't work with multi-mesh as we always broadcast to
# default mesh.
# TODO(hthu): Use copy-to-mesh after the generic Relayout CL is in.
second_mesh_a = api.copy_to_mesh(
np.ones([3], dtype=np.int32),
Layout.replicated(self.second_mesh, rank=1),
)
second_mesh_b = api.copy_to_mesh(
np.zeros([3], dtype=np.int32),
Layout.replicated(self.second_mesh, rank=1),
)
first_mesh_result = first_mesh_a + first_mesh_b
second_mesh_result = second_mesh_a + second_mesh_b
self.assertDTensorEqual(
np.asarray(3, dtype=np.int32),
Layout.replicated(self.first_mesh, rank=0), first_mesh_result)
self.assertDTensorEqual(
np.ones([3], dtype=np.int32),
Layout.replicated(self.second_mesh, rank=1), second_mesh_result)
def testMultiMeshFunc(self):
a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
with ops.device_v2(api.device_name()):
first_mesh_a = api.copy_to_mesh(
a, Layout.replicated(self.first_mesh, rank=1)
)
second_mesh_a = api.copy_to_mesh(
np.ones([4], dtype=np.int32),
Layout.replicated(self.second_mesh, rank=1),
)
with self.assertRaises(errors_impl.UnknownError):
# fails mesh propagation as it requires all inputs to be on the same
# mesh.
# pylint: disable=pointless-statement
first_mesh_a + second_mesh_a
# pylint: enable=pointless-statement
def testMultiMeshInSideFunctionLayoutV2(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
replicated_layout_on_tpu = Layout.replicated(self.second_mesh, rank=1)
host_device_id = test_util.create_device_ids_array((1,))
host_cpu_mesh = Mesh(
[_MESH_DIM_X],
host_device_id,
np.ravel(host_device_id).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
replicated_layout_on_cpu = Layout.replicated(host_cpu_mesh, rank=0)
a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
def func(t):
t = math_ops.cast(t, dtypes.float32)
t = math_ops.reduce_sum(t)
return math_ops.sqrt(t)
golden_result = func(a)
a = api.copy_to_mesh(a, replicated_layout_on_tpu)
@polymorphic_function.function
def cpu_func(t):
return math_ops.sqrt(t)
@polymorphic_function.function
def tpu_func(t):
t = math_ops.cast(t, dtypes.float32)
t = math_ops.reduce_sum(t)
cpu_tensor = api.copy_to_mesh(t, replicated_layout_on_cpu)
return cpu_func(cpu_tensor)
with ops.device_v2(api.device_name()):
output = tpu_func(a)
self.assertDTensorEqual(golden_result, replicated_layout_on_cpu, output)
def testMultiMeshCancellation(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
host_device_id = test_util.create_device_ids_array((1,))
host_cpu_mesh = Mesh(
[_MESH_DIM_X],
host_device_id,
np.ravel(host_device_id).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
replicated_layout_on_cpu = Layout([UNSHARDED], host_cpu_mesh)
replicated_layout_on_tpu = Layout([UNSHARDED], self.second_mesh)
@polymorphic_function.function
def cpu_func(x):
# Integer division by 0, which returns a bad status.
x = math_ops.cast(gen_math_ops.div(x=x, y=x), dtypes.float32)
return math_ops.cast(x, dtypes.float32)
@polymorphic_function.function
def tpu_func(cpu_tensor):
cpu_result = cpu_func(cpu_tensor)
tpu_tensor = api.copy_to_mesh(cpu_result, replicated_layout_on_tpu)
# A reduction on the TPU mesh which must be cancelled in response to the
# CPU mesh's failure.
return math_ops.reduce_sum(tpu_tensor)
cpu_tensor = api.copy_to_mesh(
constant_op.constant([0, 1]), replicated_layout_on_cpu
)
with self.assertRaisesRegex(Exception, 'Integer division by zero'):
# Ensure any errors are raised at end of scope.
with context.async_scope():
with ops.device_v2(api.device_name()):
tpu_func(cpu_tensor)
def testMultiMeshCPUToTPUTransfer(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
multiple_host_device_id = test_util.create_device_ids_array((2,))
host_multi_cpu_mesh = Mesh(
[_MESH_DIM_X],
multiple_host_device_id,
np.ravel(multiple_host_device_id).tolist(),
test_util.create_device_list((2,), 'CPU'),
)
replicated_layout_on_cpu = Layout.replicated(host_multi_cpu_mesh, rank=1)
sharded_layout_on_tpu_r1 = Layout([_MESH_DIM_X], self.second_mesh)
replicated_layout_on_tpu_r1 = Layout.replicated(self.second_mesh, rank=1)
a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
a = api.copy_to_mesh(a, replicated_layout_on_cpu)
@polymorphic_function.function
def tpu_func(t):
return api.relayout(t, sharded_layout_on_tpu_r1)
@polymorphic_function.function
def cpu_func(t):
t = math_ops.cast(t, dtypes.float32)
tpu_tensor = api.copy_to_mesh(t, replicated_layout_on_tpu_r1)
return tpu_func(tpu_tensor)
with ops.device_v2(api.device_name()):
output = cpu_func(a)
api.check_layout(output, sharded_layout_on_tpu_r1)
def testMultiMeshUnsupportedTypes(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
host_device_id = test_util.create_device_ids_array((1,))
host_cpu_mesh = Mesh(
[_MESH_DIM_X],
host_device_id,
np.ravel(host_device_id).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
replicated_layout_on_cpu = Layout.replicated(host_cpu_mesh, rank=1)
replicated_layout_on_tpu_r1 = Layout.replicated(self.second_mesh, rank=1)
s = constant_op.constant(['a', 'b', 'c'], dtype=dtypes.string)
s = api.copy_to_mesh(s, replicated_layout_on_cpu)
@polymorphic_function.function
def tpu_func(t):
return array_ops.identity(t)
@polymorphic_function.function
def cpu_func(t):
t = array_ops.identity(t)
tpu_tensor = api.copy_to_mesh(t, replicated_layout_on_tpu_r1)
return tpu_func(tpu_tensor)
with self.assertRaises(errors_impl.UnknownError) as ex:
with ops.device_v2(api.device_name()):
_ = str(cpu_func(s))
self.assertIn('unsupported output type', ex.exception.message)
def testMultiMeshCPUToCPUTransfer(self):
send_device_id = test_util.create_device_ids_array((1,))
send_cpu_mesh = Mesh(
[_MESH_DIM_X],
send_device_id,
np.ravel(send_device_id).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
recv_cpu_mesh = Mesh.from_string(
'|x=1|0|0|/job:localhost/replica:0/task:0/device:CPU:1'
)
replicated_layout_on_cpu_send = Layout.replicated(send_cpu_mesh, rank=1)
replicated_layout_on_cpu_recv = Layout.replicated(recv_cpu_mesh, rank=1)
replicated_layout_on_cpu_r0 = Layout.replicated(recv_cpu_mesh, rank=0)
def func(t):
t = math_ops.cast(t, dtypes.float32)
t = math_ops.reduce_sum(t)
return math_ops.sqrt(t)
@polymorphic_function.function
def cpu_recv_func(t):
t = math_ops.reduce_sum(t)
t = math_ops.sqrt(t)
return t
@polymorphic_function.function
def cpu_send_func(t):
t = math_ops.cast(t, dtypes.float32)
cpu_recv_tensor = api.copy_to_mesh(t, replicated_layout_on_cpu_recv)
t = cpu_recv_func(cpu_recv_tensor)
return t
a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
golden_result = func(a)
a = api.copy_to_mesh(a, replicated_layout_on_cpu_send)
with ops.device_v2(api.device_name()):
output = cpu_send_func(a)
self.assertDTensorEqual(golden_result, replicated_layout_on_cpu_r0,
output)
def testMultiMeshCPUTest(self):
device_ids = test_util.create_device_ids_array((2,))
cpu_mesh_a = Mesh(
['x'],
device_ids,
np.ravel(device_ids).tolist(),
test_util.create_device_list((2,), 'CPU'),
)
cpu_mesh_b = Mesh(
['y'],
device_ids,
np.ravel(device_ids).tolist(),
test_util.create_device_list((2,), 'CPU'),
)
replicated_layout_on_a = Layout.replicated(cpu_mesh_a, rank=1)
replicated_layout_on_b = Layout.replicated(cpu_mesh_b, rank=1)
x = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
y = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
a = api.copy_to_mesh(x, replicated_layout_on_a)
b = api.copy_to_mesh(y, replicated_layout_on_b)
@polymorphic_function.function
def func2(t1, t2):
t1 = math_ops.cast(t1, dtypes.float32)
t1 = t1 * t1
t2 = math_ops.cast(t2, dtypes.float32)
t2 = math_ops.sqrt(t2)
return t1, t2
with ops.device_v2(api.device_name()):
output1, output2 = func2(a, b)
api.check_layout(output1, replicated_layout_on_a)
api.check_layout(output2, replicated_layout_on_b)
def testFunctionWithMultiMeshInputOutputs(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
host_device_id = test_util.create_device_ids_array((1,))
host_cpu_mesh = Mesh(
[_MESH_DIM_X],
host_device_id,
np.ravel(host_device_id).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
replicated_layout_on_cpu = Layout.replicated(host_cpu_mesh, rank=1)
replicated_layout_on_cpu_r0 = Layout.replicated(host_cpu_mesh, rank=0)
replicated_layout_on_tpu_r0 = Layout.replicated(self.second_mesh, rank=0)
replicated_layout_on_tpu = Layout.replicated(self.second_mesh, rank=1)
a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
b = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
def golden_func(t1, t2):
t1 = math_ops.cast(t1, dtypes.float32)
t1 = t1 * t1
t2 = math_ops.cast(t2, dtypes.float32)
t2 = math_ops.reduce_sum(t2)
out1 = gen_math_ops.neg(t2)
t2 = t2 + t1
out0 = math_ops.sqrt(t2)
return out0, out1
golden_result0, golden_result1 = golden_func(a, b)
cpu_dtensor = api.copy_to_mesh(a, replicated_layout_on_cpu)
tpu_dtensor = api.copy_to_mesh(b, replicated_layout_on_tpu)
@polymorphic_function.function
def cpu_func(t1, t2):
t2 = t2 + t1
return math_ops.sqrt(t2)
@polymorphic_function.function
def func(tpu_input, cpu_input):
cpu_input = math_ops.cast(cpu_input, dtypes.float32)
cpu_input = cpu_input * cpu_input
tpu_input = math_ops.cast(tpu_input, dtypes.float32)
tpu_input = math_ops.reduce_sum(tpu_input)
tpu_output = gen_math_ops.neg(tpu_input)
cpu_tensor = api.copy_to_mesh(tpu_input, replicated_layout_on_cpu_r0)
cpu_output = cpu_func(cpu_tensor, cpu_input)
return cpu_output, tpu_output
with ops.device_v2(api.device_name()):
output0, output1 = func(tpu_dtensor, cpu_dtensor)
self.assertDTensorEqual(golden_result0, replicated_layout_on_cpu, output0)
self.assertDTensorEqual(golden_result1, replicated_layout_on_tpu_r0,
output1)
def testMultiMeshWithResourceOps(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
host_device_id = test_util.create_device_ids_array((1,))
host_cpu_mesh = Mesh(
[_MESH_DIM_X],
host_device_id,
np.ravel(host_device_id).tolist(),
test_util.create_device_list((1,), 'CPU'),
)
replicated_layout_on_cpu = Layout.replicated(host_cpu_mesh, rank=0)
replicated_layout_on_tpu = Layout.replicated(self.second_mesh, rank=1)
a = constant_op.constant(
[1, 2, 3, 4], dtype=dtypes.int64
) # NOTE(b/274627284): Variable of int32 type on GPU doesn't work.
def func(t):
t = math_ops.cast(t, dtypes.float32)
t = math_ops.reduce_sum(t)
return math_ops.sqrt(t)
golden_result = func(a)
@polymorphic_function.function
def cpu_func(t):
return math_ops.sqrt(t)
@polymorphic_function.function
def tpu_func(t):
t = math_ops.cast(t, dtypes.float32)
t = math_ops.reduce_sum(t)
cpu_tensor = api.copy_to_mesh(t, replicated_layout_on_cpu)
return cpu_func(cpu_tensor)
with ops.device_v2(api.device_name()):
v = api.copy_to_mesh(a, replicated_layout_on_tpu)
w = d_variable.DVariable(v)
output = tpu_func(w)
self.assertDTensorEqual(golden_result, replicated_layout_on_cpu, output)
@parameterized.named_parameters(
('_host_to_dev_sharded_i32', True, True, dtypes.int32),
('_dev_to_host_sharded_i32', False, True, dtypes.int32),
('_host_to_dev_replicated_i32', True, False, dtypes.int32),
('_dev_to_host_replicated_i32', False, False, dtypes.int32),
('_host_to_dev_sharded_bf16', True, True, dtypes.bfloat16),
('_dev_to_host_sharded_bf16', False, True, dtypes.bfloat16),
('_host_to_dev_replicated_bf16', True, False, dtypes.bfloat16),
('_dev_to_host_replicated_bf16', False, False, dtypes.bfloat16),
('_host_to_dev_sharded_f32', True, True, dtypes.float32),
('_dev_to_host_sharded_f32', False, True, dtypes.float32),
('_host_to_dev_replicated_f32', True, False, dtypes.float32),
('_dev_to_host_replicated_f32', False, False, dtypes.float32),
('_host_to_dev_sharded_f64', True, True, dtypes.float64),
('_dev_to_host_sharded_f64', False, True, dtypes.float64),
('_host_to_dev_replicated_f64', True, False, dtypes.float64),
('_dev_to_host_replicated_f64', False, False, dtypes.float64),
)
def testMultiMeshHostDeviceTransfer(self, host_to_dev, sharded, dtype):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
def run_copy_to_mesh(data, src_layout, dst_layout):
@polymorphic_function.function
def func(x):
return api.copy_to_mesh(x, dst_layout)
if src_layout.is_fully_replicated():
src_data = api.copy_to_mesh(data, src_layout)
else:
src_data = api.copy_to_mesh(
data, Layout.replicated(src_layout.mesh, rank=len(data.shape))
)
src_data = api.relayout(src_data, src_layout)
dst_data = func(src_data)
return (src_data, dst_data)
dev_mesh = self.first_mesh
cpu_mesh = self.second_mesh
if host_to_dev:
src_mesh, dst_mesh = cpu_mesh, dev_mesh
else:
src_mesh, dst_mesh = dev_mesh, cpu_mesh
if sharded:
src_layout = Layout.batch_sharded(src_mesh, src_mesh.dim_names[0], rank=2)
dst_layout = Layout.batch_sharded(dst_mesh, dst_mesh.dim_names[0], rank=2)
else:
src_layout = Layout.replicated(src_mesh, rank=2)
dst_layout = Layout.replicated(dst_mesh, rank=2)
data = array_ops.ones([8, 8], dtype=dtype)
src, dst = run_copy_to_mesh(data, src_layout, dst_layout)
self.assertDTensorEqual(data, src_layout, src)
self.assertDTensorEqual(data, dst_layout, dst)
@parameterized.named_parameters(('_host_to_tpu', True),
('_tpu_to_host', False))
def testMultiMeshWithHostMesh(self, host_to_tpu):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
sharded_layout_on_tpu = Layout([_MESH_DIM_X], self.second_mesh)
host_layout = Layout(sharded_layout_on_tpu.sharding_specs,
sharded_layout_on_tpu.mesh.host_mesh())
if host_to_tpu:
source_layout = host_layout
target_layout = sharded_layout_on_tpu
else:
source_layout = sharded_layout_on_tpu
target_layout = host_layout
numpy_a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
# TODO(b/193443769): switch to a single copy_to_mesh when this is supported.
replicated_layout = Layout.replicated(source_layout.mesh,
source_layout.rank)
a = api.copy_to_mesh(numpy_a, replicated_layout)
a = api.relayout(a, source_layout)
@polymorphic_function.function
def func(t):
target_tensor = api.copy_to_mesh(t, target_layout)
return array_ops.identity(target_tensor)
with ops.device_v2(api.device_name()):
dtensor_output = func(a)
self.assertDTensorEqual(numpy_a, target_layout, dtensor_output)
def testMultiMeshBackward(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
replicated_layout_on_tpu = Layout.replicated(self.second_mesh, rank=1)
host_layout = Layout.replicated(self.second_mesh.host_mesh(), rank=1)
source_layout = host_layout
target_layout = replicated_layout_on_tpu
@polymorphic_function.function
def func(x):
with backprop.GradientTape() as tape:
tape.watch(x)
x = x * 4.0
t = api.copy_to_mesh(x, target_layout)
sqrt = math_ops.sqrt(t)
sqrt_grad = tape.gradient(sqrt, x)
return sqrt_grad
@polymorphic_function.function
def second(x):
with backprop.GradientTape() as tape:
tape.watch(x)
sqrt_grad = func(x)
sqrt_grad_grad = tape.gradient(sqrt_grad, x)
return sqrt_grad_grad
numpy_a = constant_op.constant([1, 4, 16, 64], dtype=dtypes.float32)
a = api.copy_to_mesh(numpy_a, source_layout)
with ops.device_v2(api.device_name()):
a_grad = func(a)
self.assertDTensorEqual(0.5 * 0.5 * (1 / numpy_a)**0.5, host_layout, a_grad)
with ops.device_v2(api.device_name()):
a_grad_grad = second(a)
self.assertDTensorEqual(-0.5 * 0.5 * 0.5 * (1 / numpy_a)**1.5, host_layout,
a_grad_grad)
def testMultiMeshMultipleCopyToMesh(self):
self.skipForDeviceType(
['CPU'],
'Skipping test as only CPU mesh is available for multi-meshtest.',
)
sharded_layout_on_tpu = Layout([_MESH_DIM_X], self.second_mesh)
host_layout = Layout(
sharded_layout_on_tpu.sharding_specs,
sharded_layout_on_tpu.mesh.host_mesh(),
)
source_layout = host_layout
target_layout = sharded_layout_on_tpu
numpy_a = constant_op.constant([1, 2, 3, 4], dtype=dtypes.int32)
numpy_b = constant_op.constant([2, 2, 3, 4], dtype=dtypes.int32)
# TODO(b/193443769): switch to a single copy_to_mesh when this is supported.
replicated_layout = Layout.replicated(
source_layout.mesh, source_layout.rank
)
a = api.copy_to_mesh(numpy_a, replicated_layout)
b = api.copy_to_mesh(numpy_b, replicated_layout)
a = api.relayout(a, source_layout)
b = api.relayout(b, source_layout)
@polymorphic_function.function
def func(a, b):
a = api.copy_to_mesh(a, target_layout)
b = api.copy_to_mesh(b, target_layout)
return array_ops.identity(a), array_ops.identity(b)
with ops.device_v2(api.device_name()):
dtensor_a, dtensor_b = func(a, b)
self.assertDTensorEqual(numpy_a, target_layout, dtensor_a)
self.assertDTensorEqual(numpy_b, target_layout, dtensor_b)
def testDVariableDefaultMesh(self):
other_layout = Layout.replicated(_OTHER_CPU_MESH, rank=0)
first_layout = Layout.replicated(_ONE_D_CPU_MESH, rank=0)
_ = api.copy_to_mesh(1.0, other_layout)
init_value = api.copy_to_mesh(1.0, first_layout)
_ = d_variable.DVariable(init_value)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/numerics_test.py
+125
View File
@@ -0,0 +1,125 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for numerics in DTensor Ops."""
import os
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.platform import test
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
UNSHARDED = layout_lib.UNSHARDED
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
_MESH_DIMS = [_MESH_DIM_X, _MESH_DIM_Y]
class NumericTest(test_util.DTensorBaseTest):
def setUp(self):
super(NumericTest, self).setUp()
self.skipForDeviceType(['TPU'],
'all tests require 8 TPU cores.',
unless_device_count_equals_to=8)
test_util.reset_logical_devices('CPU', 8)
accelerator_util.initialize_accelerator_system()
self.stateless_random_seed = [0, 1]
def _create_mesh(self, topology, device):
device_ids = test_util.create_device_ids_array(topology)
return Mesh(
_MESH_DIMS,
device_ids,
np.ravel(device_ids).tolist(),
test_util.create_device_list(topology, device),
)
# Tests AllReduce numerics with and without mixed precision reduce enabled,
# based on go/dtensor-numerics.
@parameterized.named_parameters(('_without_mixed_precision_reduce', False),
('_with_mixed_precision_reduce', True))
def test_all_reduce(self, enable_mixed_precision_reduce):
if enable_mixed_precision_reduce:
os.environ['DTENSOR_ENABLE_MIXED_PRECISION_REDUCE'] = ''
# Override group size since we are testing on smaller mesh.
os.environ['DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE'] = '4'
else:
if 'DTENSOR_ENABLE_MIXED_PRECISION_REDUCE' in os.environ:
del os.environ['DTENSOR_ENABLE_MIXED_PRECISION_REDUCE']
@polymorphic_function.function
def _compute_reduction(inp):
return math_ops.reduce_sum(inp, axis=[2])
input_tensor = stateless_random_ops.stateless_random_uniform(
shape=(8, 8, 8, 64),
seed=self.stateless_random_seed,
minval=-5.0,
maxval=5.0,
dtype=dtypes.bfloat16,
)
expected = _compute_reduction(input_tensor)
# Compute reduction on 8x1, since dim 2 is unsharded AllReduce will not be
# needed.
mesh_8x1 = self._create_mesh((8, 1), 'TPU')
input_8x1 = numpy_util.pack_numpy(
input_tensor,
Layout([_MESH_DIM_X, UNSHARDED, UNSHARDED, UNSHARDED], mesh_8x1),
)
result_8x1 = _compute_reduction(input_8x1)
result_8x1_np = numpy_util.to_numpy(result_8x1)
# Compute reduction on 1x8, AllReduce will be needed since dim 2 is sharded.
mesh_1x8 = self._create_mesh((1, 8), 'TPU')
input_1x8 = numpy_util.pack_numpy(
input_tensor,
Layout([_MESH_DIM_X, UNSHARDED, _MESH_DIM_Y, UNSHARDED], mesh_1x8),
)
result_1x8 = _compute_reduction(input_1x8)
result_1x8_np = numpy_util.to_numpy(result_1x8)
self.assertEqual(result_8x1.dtype, dtypes.bfloat16)
self.assertEqual(result_1x8.dtype, dtypes.bfloat16)
# Mixed precision does not apply since AllReduce was not used, result will
# always be close to the expected value.
self.assertAllClose(result_8x1_np, expected, atol=1e-5, rtol=1e-5)
# AllReduce was needed, so result will be more accurate if mixed precision
# is enabled.
if enable_mixed_precision_reduce:
self.assertAllClose(result_1x8_np, expected, atol=1e-5, rtol=1e-5)
else:
self.assertNotAllClose(result_1x8_np, expected, atol=1e-5, rtol=1e-5)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/numpy_util_test.py
+136
View File
@@ -0,0 +1,136 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for the buffer and DTensor conversion helpers."""
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import layout
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.platform import test as tf_test
# pylint: enable=g-direct-tensorflow-import
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
class NumpyUtilTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
global_ids = test_util.create_device_ids_array((2, 2))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: layout.Mesh([_MESH_DIM_X, _MESH_DIM_Y], global_ids, local_ids,
test_util.create_device_list((2, 2), device))
for device in ('CPU', 'GPU', 'TPU')
}
self.mesh = self.configTestMesh(mesh_dict)
def test_tensor_from_replicated(self):
tensors = [np.arange(4) for i in range(self.mesh.size)]
replicated_layout = layout.Layout([layout.UNSHARDED, layout.UNSHARDED],
mesh=self.mesh)
self.assertAllClose(
np.arange(4), numpy_util.unpacked_to_numpy(tensors, replicated_layout))
def test_tensor_x_sharded(self):
t00 = np.arange(8).reshape(2, 4)
t01 = np.arange(8).reshape(2, 4)
t10 = np.arange(8, 16).reshape(2, 4)
t11 = np.arange(8, 16).reshape(2, 4)
tensors = [t00, t01, t10, t11]
sharded_on_x = layout.Layout([_MESH_DIM_X, layout.UNSHARDED],
mesh=self.mesh)
self.assertAllClose(
np.arange(16).reshape(4, 4),
numpy_util.unpacked_to_numpy(tensors, sharded_on_x))
def test_tensor_y_sharded(self):
# [[0,1], [4,5], [8,9], [12,13]]
t00 = np.arange(16).reshape(4, 4)[:, :-2]
# [[2,3], [6,7], [10,11], [14,15]]
t01 = np.arange(16).reshape(4, 4)[:, 2:4]
t10 = np.arange(16).reshape(4, 4)[:, :-2]
t11 = np.arange(16).reshape(4, 4)[:, 2:4]
tensors = [t00, t01, t10, t11]
sharded_on_y = layout.Layout([layout.UNSHARDED, _MESH_DIM_Y],
mesh=self.mesh)
self.assertAllClose(
numpy_util.unpacked_to_numpy(tensors, sharded_on_y),
np.arange(16).reshape(4, 4))
def test_tensor_x_sharded_on_mesh_y(self):
t00 = np.arange(8).reshape(2, 4)
t01 = np.arange(8, 16).reshape(2, 4)
t10 = np.arange(8).reshape(2, 4)
t11 = np.arange(8, 16).reshape(2, 4)
tensors = [t00, t01, t10, t11]
sharded_on_y = layout.Layout([_MESH_DIM_Y, layout.UNSHARDED],
mesh=self.mesh)
self.assertAllClose(
numpy_util.unpacked_to_numpy(tensors, sharded_on_y),
np.arange(16).reshape(4, 4))
def test_tensor_y_sharded_on_mesh_x(self):
# [[0,1], [4,5], [8,9], [12,13]]
t00 = np.arange(16).reshape(4, 4)[:, :-2]
t01 = np.arange(16).reshape(4, 4)[:, :-2]
# [[2,3], [6,7], [10,11], [14,15]]
t10 = np.arange(16).reshape(4, 4)[:, 2:4]
t11 = np.arange(16).reshape(4, 4)[:, 2:4]
tensors = [t00, t01, t10, t11]
sharded_on_x = layout.Layout([layout.UNSHARDED, _MESH_DIM_X],
mesh=self.mesh)
self.assertAllClose(
numpy_util.unpacked_to_numpy(tensors, sharded_on_x),
np.arange(16).reshape(4, 4))
def test_tensor_x_y_sharded_x_y(self):
t00 = np.array([[0, 1], [4, 5]])
t01 = np.array([[2, 3], [6, 7]])
t10 = np.array([[8, 9], [12, 13]])
t11 = np.array([[10, 11], [14, 15]])
tensors = [t00, t01, t10, t11]
sharded_on_x_y = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
self.assertAllClose(
numpy_util.unpacked_to_numpy(tensors, sharded_on_x_y),
np.arange(16).reshape(4, 4))
def test_tensor_x_y_sharded_y_x(self):
t00 = np.array([[0, 1], [4, 5]])
t01 = np.array([[8, 9], [12, 13]])
t10 = np.array([[2, 3], [6, 7]])
t11 = np.array([[10, 11], [14, 15]])
tensors = [t00, t01, t10, t11]
sharded_on_y_x = layout.Layout([_MESH_DIM_Y, _MESH_DIM_X], mesh=self.mesh)
self.assertAllClose(
numpy_util.unpacked_to_numpy(tensors, sharded_on_y_x),
np.arange(16).reshape(4, 4))
def test_unpack_uneven_split_raises(self):
value = np.arange(5)
layout = layout_lib.Layout.batch_sharded(self.mesh, batch_dim='x', rank=1)
with self.assertRaisesRegex(ValueError, 'not evenly divisible'):
numpy_util.unpack(value, layout)
if __name__ == '__main__':
tf_test.main()
tensorflow/dtensor/python/tests/rng_test.py
+665
View File
@@ -0,0 +1,665 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from absl.testing import parameterized
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.dtensor.python.tests import test_util_ops
from tensorflow.python.distribute import tpu_strategy
from tensorflow.python.distribute.cluster_resolver.tpu import tpu_cluster_resolver
from tensorflow.python.eager import remote
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_bitwise_ops
from tensorflow.python.ops import gen_stateful_random_ops
from tensorflow.python.ops import gen_stateless_random_ops_v2
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.tpu import device_assignment as device_assignment_lib
# pylint: enable=g-direct-tensorflow-import
# Makes a 2-D mesh with dimensions as, X(2) and Y(4).
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
_MESH_DIMS = [_MESH_DIM_X, _MESH_DIM_Y]
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
# Create a random local IDs to make tests more challenging.
_LOCAL_IDS = [7, 3, 1, 4, 2, 0, 6, 5]
# The row and col indices for each local id, e.g., 7 is (row=1, col=3)
_ROW_INDEX = [i / 4 for i in _LOCAL_IDS]
_COL_INDEX = [i % 4 for i in _LOCAL_IDS]
# The index of local id for the row head.
#
# For example, local id 7 is on row 1, the head is local id 4, whose index in
# _LOCAL_IDS is 3, i.e., _LOCAL_IDS[3] == 4
_ROW_0_HEAD = 3
_ROW_1_HEAD = 5
_ROW_HEAD = [3, 5, 5, 3, 5, 5, 3, 3]
# The index of local id for the col head. Similar to row id before.
_COL_0_HEAD = 5
_COL_1_HEAD = 2
_COL_2_HEAD = 4
_COL_3_HEAD = 1
_COL_HEAD = [1, 1, 2, 5, 4, 5, 4, 2]
_tpu_strategy = None
def _call_op(op, seed, shape, dtype, key, counter, alg, minval, maxval,
op_version):
if op_version == 'V1':
return op(shape=shape, seed=seed, dtype=dtype)
elif op_version == 'V2':
return op(shape=shape, key=key, counter=counter, alg=alg, dtype=dtype)
elif op_version == 'V2_RANGE':
return op(
shape=shape,
key=key,
counter=counter,
alg=alg,
minval=minval,
maxval=maxval)
else:
raise ValueError('op_version argument was invalid.')
def _call_dtensor_op(op, seed, shape, dtype, key, counter, alg, minval, maxval,
op_version, mesh):
if op_version == 'V1':
return op(shape=shape, seed=seed, dtype=dtype)
shape = numpy_util.pack_numpy(
constant_op.constant(shape), Layout.replicated(mesh, 1)
)
key = numpy_util.pack_numpy(key, Layout.replicated(mesh, 1))
counter = numpy_util.pack_numpy(counter, Layout.replicated(mesh, 1))
if op_version == 'V2':
return op(shape=shape, key=key, counter=counter, alg=alg, dtype=dtype)
elif op_version == 'V2_RANGE':
return op(
shape=shape,
key=key,
counter=counter,
alg=alg,
minval=minval,
maxval=maxval)
else:
raise ValueError('op_version argument was invalid.')
def get_tpu_strategy():
"""Returns a single-core TPUStrategy."""
global _tpu_strategy
if _tpu_strategy is not None:
return _tpu_strategy
resolver = tpu_cluster_resolver.TPUClusterResolver(tpu='')
remote.connect_to_cluster(resolver)
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
device_assignment = device_assignment_lib.DeviceAssignment.build(
topology, num_replicas=1
)
strategy = tpu_strategy.TPUStrategyV2(
resolver, experimental_device_assignment=device_assignment
)
_tpu_strategy = strategy
return strategy
def rng_op_spmd(op,
device_id,
seed,
shape,
dtype,
key,
counter,
alg,
minval,
maxval,
op_version,
device_index_fn,
full_replicated=False,
is_tpu=False):
if not is_tpu:
return rng_op_spmd_fn(
op,
device_id,
seed,
shape,
dtype,
key,
counter,
alg,
minval,
maxval,
op_version,
device_index_fn,
full_replicated=full_replicated)
# As of 2021-April, TPU eager and multi-device function produce different
# stateless rng results compared with bridge compiled function. As DTensor
# uses bridge to lower TPU function by default, we need to create a
# TPUStrategy for single core and invoke `run` on it.
@polymorphic_function.function
def tpu_fn(device_id, seed):
return rng_op_spmd_fn(
op,
device_id,
seed,
shape,
dtype,
key,
counter,
alg,
minval,
maxval,
op_version,
device_index_fn,
full_replicated=full_replicated)
return get_tpu_strategy().run(tpu_fn, args=(device_id, seed))
def rng_op_spmd_fn(op,
device_id,
seed,
shape,
dtype,
key,
counter,
alg,
minval,
maxval,
op_version,
device_index_fn,
full_replicated=False):
if full_replicated:
# TODO(bfontain,xiejw): Consider to make this consistent with non-replicated
# case. Seems very confusing.
new_seed, new_key = seed, key
else:
# Runs on TF2 non-DTensor pure eager. This code should align the same
# logic in RandomOpSPMDExpander.
x_cord = device_id // 4
y_cord = device_id % 4
device_index = device_index_fn(x_cord, y_cord)
device_id_seed = device_index * 65536 + 65521
new_seed = gen_bitwise_ops.bitwise_xor(seed, device_id_seed)
new_key = gen_bitwise_ops.bitwise_xor(
key, math_ops.cast(device_id_seed, dtype=dtypes.uint64)
)
return _call_op(
op=op,
seed=new_seed,
shape=shape,
dtype=dtype,
key=new_key,
counter=counter,
alg=alg,
minval=minval,
maxval=maxval,
op_version=op_version)
class DTensorRNGTest(test_util.DTensorBaseTest):
def setUp(self):
super(DTensorRNGTest, self).setUp()
global_ids = test_util.create_device_ids_array((2, 4))
local_ids = _LOCAL_IDS
mesh_dict = {
device: Mesh(
[_MESH_DIM_X, _MESH_DIM_Y],
global_ids,
local_ids,
test_util.create_device_list((2, 4), device),
)
for device in ('CPU', 'GPU', 'TPU')
}
self.mesh = self.configTestMesh(mesh_dict)
# Creates a bunch of common layouts used by tests later.
self.replicated_layout_2d = Layout.replicated(self.mesh, rank=2)
self.shardings = {
'batch': Layout.batch_sharded,
'inner': Layout.inner_sharded
}
# Creates a bunch of parameters for rng V2 ops
self.key = constant_op.constant([123], dtype=dtypes.uint64)
self.counter = constant_op.constant([1, 1], dtype=dtypes.uint64)
self.alg = 1
self.minval = 1
self.maxval = 100
@parameterized.named_parameters(test_util_ops.RANDOM_OPS)
def testStatelessRNGWithFullyReplicated(self, op, dtype, op_version):
layout = self.replicated_layout_2d
shape = [16, 16]
seed = [123, 321]
with ops.device_v2(api.device_name()):
with api._dtensor_device()._default_layout(layout):
b = _call_dtensor_op(
op=op,
seed=seed,
shape=shape,
dtype=dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
mesh=self.mesh)
api.check_layout(b, layout)
self.assertListEqual(shape, list(b.shape))
b = [tensor.numpy() for tensor in api.unpack(b)]
for i in range(self.mesh.num_local_devices() - 1):
self.assertAllEqual(b[i], b[i + 1])
@parameterized.named_parameters(test_util_ops.RANDOM_OPS)
def testStatelessRNGWithFullyReplicatedComparingWithNonDTensor(
self, op, dtype, op_version):
layout = self.replicated_layout_2d
shape = [16, 16]
seed = [123, 321]
with ops.device_v2(api.device_name()):
with api._dtensor_device()._default_layout(layout):
b = _call_dtensor_op(
op=op,
seed=seed,
shape=shape,
dtype=dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
mesh=self.mesh)
api.check_layout(b, layout)
self.assertListEqual(shape, list(b.shape))
b = [tensor.numpy() for tensor in api.unpack(b)]
local_shape = shape
for index, device_id in enumerate(_LOCAL_IDS):
self.assertAllEqual(
b[index],
rng_op_spmd(
op,
device_id,
seed,
local_shape,
dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
device_index_fn=None, # not needed
full_replicated=True,
is_tpu=self.mesh.device_type().upper() == 'TPU'))
@parameterized.named_parameters(
test_util_ops.expand_test_config(
test_util_ops.RANDOM_OPS,
[
{
'dim': _MESH_DIM_X,
'shard_type': 'batch',
},
{
'dim': _MESH_DIM_Y,
'shard_type': 'batch',
},
{
'dim': _MESH_DIM_X,
'shard_type': 'inner',
},
{'dim': _MESH_DIM_Y, 'shard_type': 'inner'},
],
)
)
def testStatelessRNGOpsWithSingleDimensionSharded(self, op, dtype, op_version,
dim, shard_type):
shape = [128, 128]
seed = [123, 321]
sharding = self.shardings[shard_type]
layout = sharding(self.mesh, dim, rank=2)
# Raw rng Ops do not have inputs, so we need to place the Op DTensor device
# explicitly.
with ops.device_v2(api.device_name()):
with api._dtensor_device()._default_layout(layout):
b = _call_dtensor_op(
op=op,
seed=seed,
shape=shape,
dtype=dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
mesh=self.mesh)
api.check_layout(b, layout)
b = [tensor.numpy() for tensor in api.unpack(b)]
if dim == _MESH_DIM_X:
if shard_type == 'batch':
self.assertAllEqual(b[0].shape, [64, 128])
else:
assert shard_type == 'inner'
self.assertAllEqual(b[0].shape, [128, 64])
# first check that each component is same as the row header.
for i in range(self.mesh.num_local_devices()):
self.assertAllEqual(b[i], b[_ROW_HEAD[i]])
# then check the row header are NOT identital.
self.assertNotAllEqual(b[_ROW_0_HEAD], b[_ROW_1_HEAD])
elif dim == _MESH_DIM_Y:
if shard_type == 'batch':
self.assertAllEqual(b[0].shape, [32, 128])
else:
assert shard_type == 'inner'
self.assertAllEqual(b[0].shape, [128, 32])
# first check elements in same columns are identical
for i in range(self.mesh.num_local_devices()):
self.assertAllEqual(b[i], b[_COL_HEAD[i]])
col_heads = [_COL_0_HEAD, _COL_1_HEAD, _COL_2_HEAD, _COL_3_HEAD]
# then check the column header are not identital (mutually)
for i in range(self.mesh.num_local_devices() - 1):
for j in range(self.mesh.num_local_devices()):
if i == j:
continue
if i in col_heads and j in col_heads:
self.assertNotAllEqual(b[i], b[j])
else:
self.fail('should not reach here.')
@parameterized.named_parameters(
test_util_ops.expand_test_config(
test_util_ops.RANDOM_OPS,
[
{
'dim': _MESH_DIM_X,
'shard_type': 'batch',
},
{
'dim': _MESH_DIM_Y,
'shard_type': 'batch',
},
{
'dim': _MESH_DIM_X,
'shard_type': 'inner',
},
{'dim': _MESH_DIM_Y, 'shard_type': 'inner'},
],
)
)
def testStatelessRNGOpsWithSingleDimensionShardedComparingWithNonDTensor(
self, op, dtype, op_version, dim, shard_type):
shape = [128, 128]
seed = [123, 321]
sharding = self.shardings[shard_type]
layout = sharding(self.mesh, dim, rank=2)
# Raw rng Ops do not have inputs, so we need to place the Op DTensor device
# explicitly.
with ops.device_v2(api.device_name()):
with api._dtensor_device()._default_layout(layout):
b = _call_dtensor_op(
op=op,
seed=seed,
shape=shape,
dtype=dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
mesh=self.mesh)
api.check_layout(b, layout)
b = [tensor.numpy() for tensor in api.unpack(b)]
if dim == _MESH_DIM_X:
if shard_type == 'batch':
local_shape = [64, 128]
else:
local_shape = [128, 64]
def device_index_fn(x_cord, y_cord):
# See todo of device_index_fn in 2d sharding case.
del y_cord
return x_cord
for index, device_id in enumerate(_LOCAL_IDS):
self.assertAllEqual(
b[index],
rng_op_spmd(
op,
device_id,
seed,
local_shape,
dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
device_index_fn=device_index_fn,
is_tpu=self.mesh.device_type().upper() == 'TPU'))
elif dim == _MESH_DIM_Y:
if shard_type == 'batch':
local_shape = [32, 128]
else:
local_shape = [128, 32]
def device_index_fn(x_cord, y_cord):
# See todo of device_index_fn in 2d sharding case. note this case is
# particulary interesting as 2*y_cord is more natual.
del x_cord
return y_cord
for index, device_id in enumerate(_LOCAL_IDS):
self.assertAllEqual(
b[index],
rng_op_spmd(
op,
device_id,
seed,
local_shape,
dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
device_index_fn=device_index_fn,
is_tpu=self.mesh.device_type().upper() == 'TPU'))
else:
self.fail('should not reach here.')
@parameterized.named_parameters(test_util_ops.RANDOM_OPS)
def testStatelessRNGOpsWith2DSharding(self, op, dtype, op_version):
shape = [128, 128]
seed = [123, 321]
layout = Layout([_MESH_DIM_Y, _MESH_DIM_X], self.mesh)
# Raw rng Ops do not have inputs, so we need to place the Op DTensor device
# explicitly.
with ops.device_v2(api.device_name()):
with api._dtensor_device()._default_layout(layout):
b = _call_dtensor_op(
op=op,
seed=seed,
shape=shape,
dtype=dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
mesh=self.mesh)
api.check_layout(b, layout)
b = [tensor.numpy() for tensor in api.unpack(b)]
# check all raw components are not identital (mutually)
for i in range(self.mesh.num_local_devices() - 1):
for j in range(self.mesh.num_local_devices()):
if i == j:
continue
self.assertNotAllEqual(b[i], b[j])
@parameterized.named_parameters(test_util_ops.RANDOM_OPS)
def testStatelessRNGOpsWith2DShardingComparingWithNonDTensor(
self, op, dtype, op_version):
shape = [128, 128]
seed = [123, 321]
layout = Layout([_MESH_DIM_Y, _MESH_DIM_X], self.mesh)
local_shape = [128 // 4, 128 // 2]
# Raw rng Ops do not have inputs, so we need to place the Op DTensor device
# explicitly.
with ops.device_v2(api.device_name()):
with api._dtensor_device()._default_layout(layout):
b = _call_dtensor_op(
op=op,
seed=seed,
shape=shape,
dtype=dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
mesh=self.mesh)
api.check_layout(b, layout)
b = [tensor.numpy() for tensor in api.unpack(b)]
def device_index_fn(x_cord, y_cord):
# TODO(bfontain,xiejw): Currently, the device index is x+2y. But it is
# more natual to use 4x+y for a mesh<x=2, y=4>. Consider to change this
# once all correctness tests are done.
return x_cord + 2 * y_cord
for index, device_id in enumerate(_LOCAL_IDS):
self.assertAllEqual(
b[index],
rng_op_spmd(
op,
device_id,
seed,
local_shape,
dtype,
key=self.key,
counter=self.counter,
alg=self.alg,
minval=self.minval,
maxval=self.maxval,
op_version=op_version,
device_index_fn=device_index_fn,
is_tpu=self.mesh.device_type().upper() == 'TPU'))
def testRNGReadAndSkip(self):
replicated_layout = Layout.replicated(self.mesh, 1)
a = constant_op.constant([1, 2, 3], dtype=dtypes.int64)
v = variables.Variable(a)
expected = gen_stateful_random_ops.rng_read_and_skip(
resource=v.handle,
alg=1,
delta=constant_op.constant(1, dtype=dtypes.uint64),
)
a = numpy_util.pack_numpy(a, replicated_layout)
v = d_variable.DVariable(a)
got = gen_stateful_random_ops.rng_read_and_skip(
resource=v.handle,
alg=1,
delta=constant_op.constant(1, dtype=dtypes.uint64),
)
self.assertDTensorEqual(expected, replicated_layout, got)
def testStatelessRandomGetKeyCounter(self):
seed = constant_op.constant([7, 17], dtypes.int32)
# TPU computation result is different from CPU computation.
# We force it to run on the TPU using tpu_strategy for TPU mesh
# so that we compare equal values.
@polymorphic_function.function
def tpu_fn():
return gen_stateless_random_ops_v2.stateless_random_get_key_counter(
seed=seed
)
if self.mesh.device_type().upper() == 'TPU':
expected = get_tpu_strategy().run(tpu_fn)
else:
expected = gen_stateless_random_ops_v2.stateless_random_get_key_counter(
seed=seed
)
replicated_1d_layout = Layout.replicated(self.mesh, 1)
seed = numpy_util.pack_numpy(seed, replicated_1d_layout)
got = gen_stateless_random_ops_v2.stateless_random_get_key_counter(
seed=seed
)
self.assertDTensorEqual(expected[0], replicated_1d_layout, got[0])
self.assertDTensorEqual(expected[1], replicated_1d_layout, got[1])
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/save_restore_v2_test.py
+337
View File
@@ -0,0 +1,337 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import gc
from absl.testing import parameterized
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import checkpoint_management
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.platform import test
Mesh = layout_lib.Mesh
Layout = layout_lib.Layout
UNSHARDED = layout_lib.UNSHARDED
# Makes a 2D mesh with dimension X(2) and dimension Y(4).
_MESH_DIM_X = 'x'
_MESH_DIM_Y = 'y'
_DEVICE_IDS = test_util.create_device_ids_array((2, 4))
_TWO_D_CPU_MESH = Mesh(
[_MESH_DIM_X, _MESH_DIM_Y],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2, 4), 'CPU'),
)
_TWO_D_TPU_MESH = Mesh(
[_MESH_DIM_X, _MESH_DIM_Y],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2, 4), 'TPU'),
)
_TWO_D_GPU_MESH = Mesh(
[_MESH_DIM_X, _MESH_DIM_Y],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2, 4), 'GPU'),
)
class DTensorSaveRestoreV2Test(test_util.DTensorBaseTest):
def setUp(self):
super(DTensorSaveRestoreV2Test, self).setUp()
self.skipForDeviceType(['TPU'],
'all tests require 8 TPU cores.',
unless_device_count_equals_to=8)
mesh_dict = {
'CPU': _TWO_D_CPU_MESH,
'GPU': _TWO_D_GPU_MESH,
'TPU': _TWO_D_TPU_MESH,
}
self.mesh = self.configTestMesh(mesh_dict)
self.skipForTfrt(
'b/235088250, DTensorCheckpointingV2 requires upcasting TF scalar '
'variables to replicated DTensor scalar variables, which is not '
'supported in TFRT.')
@parameterized.named_parameters(
('x_unsharded', [_MESH_DIM_X, UNSHARDED]),
('unsharded_x', [UNSHARDED, _MESH_DIM_X]),
('x_y', [_MESH_DIM_X, _MESH_DIM_Y]),
('unsharded_unsharded', [UNSHARDED, UNSHARDED]),
)
def test_checkpoint_simple(self, shard_spec):
tensor_a = stateless_random_ops.stateless_random_uniform(
shape=[4, 8], seed=[0, 1]
)
tensor_b = stateless_random_ops.stateless_random_uniform(
shape=[2, 4], seed=[0, 1]
)
layout = Layout(shard_spec, self.mesh)
dvariable_a = d_variable.DVariable(numpy_util.pack_numpy(tensor_a, layout))
dvariable_b = d_variable.DVariable(numpy_util.pack_numpy(tensor_b, layout))
# Record a checkpoint with two dvariables.
ckpt = checkpoint.Checkpoint(a=dvariable_a, b=dvariable_b)
saved_path = ckpt.save(self.get_temp_dir())
# Zero out the values of the DVariables so that we can restore
# and check that the values are restored to the initial random values.
dvariable_a.assign(
numpy_util.pack_numpy(
array_ops.zeros([4, 8], dtype=dtypes.float32), layout
)
)
dvariable_b.assign(
numpy_util.pack_numpy(
array_ops.zeros([2, 4], dtype=dtypes.float32), layout
)
)
ckpt.restore(saved_path)
self.assertDTensorEqual(tensor_a, layout, dvariable_a.read_value())
self.assertDTensorEqual(tensor_b, layout, dvariable_b.read_value())
@parameterized.named_parameters(
('x_unsharded', [_MESH_DIM_X, UNSHARDED]),
('unsharded_x', [UNSHARDED, _MESH_DIM_X]),
('x_y', [_MESH_DIM_X, _MESH_DIM_Y]),
('unsharded_unsharded', [UNSHARDED, UNSHARDED]),
)
def test_checkpoint_write(self, shard_spec):
tensor_a = stateless_random_ops.stateless_random_uniform(
shape=[4, 8], seed=[0, 1]
)
tensor_b = stateless_random_ops.stateless_random_uniform(
shape=[2, 4], seed=[0, 1]
)
layout = Layout(shard_spec, self.mesh)
dvariable_a = d_variable.DVariable(numpy_util.pack_numpy(tensor_a, layout))
dvariable_b = d_variable.DVariable(numpy_util.pack_numpy(tensor_b, layout))
ckpt = checkpoint.Checkpoint(a=dvariable_a, b=dvariable_b)
saved_path = ckpt.write(self.get_temp_dir())
dvariable_a.assign(
numpy_util.pack_numpy(
array_ops.zeros([4, 8], dtype=dtypes.float32), layout
)
)
dvariable_b.assign(
numpy_util.pack_numpy(
array_ops.zeros([2, 4], dtype=dtypes.float32), layout
)
)
ckpt.restore(saved_path)
self.assertDTensorEqual(tensor_a, layout, dvariable_a.read_value())
self.assertDTensorEqual(tensor_b, layout, dvariable_b.read_value())
@parameterized.named_parameters(
('x_unsharded', [_MESH_DIM_X, UNSHARDED]),
('unsharded_x', [UNSHARDED, _MESH_DIM_X]),
('x_y', [_MESH_DIM_X, _MESH_DIM_Y]),
('unsharded_unsharded', [UNSHARDED, UNSHARDED]),
)
def test_checkpoint_manager(self, shard_spec):
tensor_a = stateless_random_ops.stateless_random_uniform(
shape=[8, 16], seed=[0, 1]
)
tensor_b = stateless_random_ops.stateless_random_uniform(
shape=[4, 4], seed=[0, 1]
)
layout = Layout(shard_spec, self.mesh)
dvariable_a = d_variable.DVariable(numpy_util.pack_numpy(tensor_a, layout))
dvariable_b = d_variable.DVariable(numpy_util.pack_numpy(tensor_b, layout))
# Record a checkpoint with two dvariables.
ckpt = checkpoint.Checkpoint(a=dvariable_a, b=dvariable_b)
checkpoint_manager = checkpoint_management.CheckpointManager(
ckpt, self.get_temp_dir(), max_to_keep=None
)
saved_path = checkpoint_manager.save()
# Zero out the values of the DVariables so that we can restore
# and check that the values are restored to the initial random values.
dvariable_a.assign(
numpy_util.pack_numpy(
array_ops.zeros([8, 16], dtype=dtypes.float32), layout
)
)
dvariable_b.assign(
numpy_util.pack_numpy(
array_ops.zeros([4, 4], dtype=dtypes.float32), layout
)
)
ckpt.restore(saved_path)
self.assertDTensorEqual(tensor_a, layout, dvariable_a.read_value())
self.assertDTensorEqual(tensor_b, layout, dvariable_b.read_value())
@parameterized.named_parameters(
('x_unsharded', [_MESH_DIM_X, UNSHARDED]),
('unsharded_x', [UNSHARDED, _MESH_DIM_X]),
('x_y', [_MESH_DIM_X, _MESH_DIM_Y]),
('unsharded_unsharded', [UNSHARDED, UNSHARDED]),
)
def test_checkpoint_restore_with_different_layout(self, shard_spec):
tensor_a = stateless_random_ops.stateless_random_uniform(
shape=[4, 8], seed=[0, 1]
)
tensor_b = stateless_random_ops.stateless_random_uniform(
shape=[2, 4], seed=[0, 1]
)
layout = Layout(shard_spec, self.mesh)
dvariable_a = d_variable.DVariable(numpy_util.pack_numpy(tensor_a, layout))
dvariable_b = d_variable.DVariable(numpy_util.pack_numpy(tensor_b, layout))
# Record a checkpoint with two dvariables.
checkpoint_1 = checkpoint.Checkpoint(a=dvariable_a, b=dvariable_b)
saved_path = checkpoint_1.save(self.get_temp_dir())
new_layout = Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)
# Create new Dvariables, zero'd out with different layouts
# from the layouts we saved the tensors.
dvariable_a = d_variable.DVariable(
numpy_util.pack_numpy(
array_ops.zeros([4, 8], dtype=dtypes.float32), new_layout
)
)
dvariable_b = d_variable.DVariable(
numpy_util.pack_numpy(
array_ops.zeros([2, 4], dtype=dtypes.float32), new_layout
)
)
checkpoint_2 = checkpoint.Checkpoint(a=dvariable_a, b=dvariable_b)
checkpoint_2.restore(saved_path)
self.assertDTensorEqual(tensor_a, new_layout, dvariable_a.read_value())
self.assertDTensorEqual(tensor_b, new_layout, dvariable_b.read_value())
@parameterized.named_parameters(
('x_unsharded', [_MESH_DIM_X, UNSHARDED]),
('unsharded_x', [UNSHARDED, _MESH_DIM_X]),
)
def test_checkpoint_in_a_train_loop(self, shard_dims):
# This test is a parallel test with save_restore_test's
# DTensorSaveRestoreTest.test_checkpoint
class M(module.Module):
# Pass in both replicated and sharded for better coverage.
def __init__(self, replicated_value, sharded_value):
# This is actually a DVariable.
self.r = d_variable.DVariable(replicated_value)
self.s = d_variable.DVariable(sharded_value)
def __call__(self, x):
return math_ops.reduce_sum(x + self.r) + math_ops.reduce_sum(x + self.s)
directory = self.get_temp_dir()
sharded_np = np.arange(8).reshape((2, 4)).astype(np.float32)
replicated_np = np.arange(16).reshape((8, 2)).astype(np.float32)
replicated_layout = Layout.replicated(self.mesh, rank=2)
one_d_sharded_layout = Layout(shard_dims, self.mesh)
replicated_value = api.copy_to_mesh(replicated_np, replicated_layout)
replicated_zeros = api.copy_to_mesh(
np.zeros((8, 2)).astype(np.float32), replicated_layout
)
sharded_value = numpy_util.pack_numpy(sharded_np, one_d_sharded_layout)
sharded_zeros = numpy_util.pack_numpy(
np.zeros((2, 4)).astype(np.float32), one_d_sharded_layout)
# Training loop that just increments the model's variable every "epoch"
# to test checkpointing.
for epoch in range(5):
m = M(replicated_value, sharded_value)
ckpt = checkpoint.Checkpoint(model=m)
manager = checkpoint_management.CheckpointManager(
ckpt, directory=directory, max_to_keep=None
)
ckpt.restore(manager.latest_checkpoint)
# Ensure that the variable is created
m(api.copy_to_mesh(1.0, Layout.replicated(self.mesh, rank=0)))
self.assertDTensorEqual(epoch + replicated_np, replicated_layout, m.r)
self.assertDTensorEqual(epoch + sharded_np, one_d_sharded_layout, m.s)
m.s.assign_add(
numpy_util.pack_numpy(
np.ones((2, 4), dtype=np.float32), one_d_sharded_layout))
m.r.assign_add(
api.copy_to_mesh(
constant_op.constant(np.ones((8, 2), dtype=np.float32)),
replicated_layout,
)
)
checkpoint_number = epoch + 1
stats1 = api._dtensor_device()._get_stats()
manager.save(checkpoint_number=checkpoint_number)
gc.collect()
stats2 = api._dtensor_device()._get_stats()
keys = set(stats2.keys())
keys.update(stats1.keys())
diff = {k: stats2.get(k, 0) - stats1.get(k, 0) for k in keys}
diff = {k: v for k, v in diff.items() if v != 0}
m.s.assign(sharded_zeros)
m.r.assign(replicated_zeros)
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/sparse_test.py
+141
View File
@@ -0,0 +1,141 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from absl.testing import parameterized
import numpy as np
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
# Convenient constants to use for tests.
_BATCH_DIM = "batch"
_MESH_DIM_X = "x"
# Shorter notation
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
class DTensorSPMDTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
self.skipForDeviceType(["GPU", "TPU"],
"SparseTensors only supported on CPU.")
global_ids = test_util.create_device_ids_array((2, 2))
local_ids = np.ravel(global_ids).tolist()
mesh_dict = {
device: Mesh(
[_BATCH_DIM, _MESH_DIM_X],
global_ids,
local_ids,
test_util.create_device_list((2, 2), device),
)
for device in ("CPU", "GPU", "TPU")
}
self.mesh = self.configTestMesh(mesh_dict)
@parameterized.parameters(
[dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64]
)
def testIdentityOpWithSparseTensorInputSimple(self, dtype):
inputs = array_ops.ones([6, 4], dtype=dtype)
layout = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
@polymorphic_function.function
def f(x):
return array_ops.identity(x)
self.assertDTensorEqual(
inputs, layout,
f(numpy_util.pack_numpy(inputs, layout, make_sparse=True)))
@parameterized.product(
dtype=[dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64],
is_sparse_a=[True, False],
is_sparse_b=[True, False],
)
def testIdentityOpWithSparseTensorInputComplex(self, dtype, is_sparse_a,
is_sparse_b):
inputs_a = array_ops.ones([2, 1], dtype=dtype)
inputs_b = array_ops.ones([32, 16], dtype=dtype)
layout_a = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
layout_b = Layout.replicated(self.mesh, rank=2)
@polymorphic_function.function
def f(x, y):
return array_ops.identity(x), array_ops.identity(y)
got_a, got_b = f(
numpy_util.pack_numpy(inputs_a, layout_a, make_sparse=is_sparse_a),
numpy_util.pack_numpy(inputs_b, layout_b, make_sparse=is_sparse_b))
self.assertDTensorEqual(inputs_a, layout_a, got_a)
self.assertDTensorEqual(inputs_b, layout_b, got_b)
def testMultipleIdentityOpFromOneSparseTensor(self):
inputs_a = array_ops.ones([2, 1])
layout_a = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
@polymorphic_function.function
def f(x):
return array_ops.identity(x), array_ops.identity(x)
got_a, got_b = f(
numpy_util.pack_numpy(inputs_a, layout_a, make_sparse=True))
self.assertDTensorEqual(inputs_a, layout_a, got_a)
self.assertDTensorEqual(inputs_a, layout_a, got_b)
@parameterized.product(
is_sparse_a=[True, False],
is_sparse_b=[True, False],
shard_type=["Replicated", "Sharded"])
def testSparseTensorDenseMatMul(self, is_sparse_a, is_sparse_b, shard_type):
inputs_a = array_ops.ones([16, 16])
inputs_b = array_ops.ones([16, 16])
if shard_type == "Replicated":
layout_a = Layout.replicated(self.mesh, rank=2)
layout_b = Layout.replicated(self.mesh, rank=2)
else:
layout_a = Layout([_MESH_DIM_X, _BATCH_DIM], self.mesh)
layout_b = Layout(["unsharded", _MESH_DIM_X], self.mesh)
expected = math_ops.matmul(inputs_a, inputs_b)
@polymorphic_function.function
def f(x, y):
return math_ops.matmul(x, y)
got = f(
numpy_util.pack_numpy(inputs_a, layout_a, make_sparse=is_sparse_a),
numpy_util.pack_numpy(inputs_b, layout_b, make_sparse=is_sparse_b))
self.assertDTensorEqual(expected, Layout.replicated(self.mesh, rank=2), got)
if __name__ == "__main__":
test.main()
tensorflow/dtensor/python/tests/spmd_test.py
+5086
View File
File diff suppressed because it is too large Load Diff
tensorflow/dtensor/python/tests/test_backend_name.py
+38
View File
@@ -0,0 +1,38 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""List of test backend names."""
import enum
import os
# LINT.IfChange(backend_name)
class DTensorTestUtilBackend(enum.Enum):
"""DTensor backend the test is being run on."""
UNSPECIFIED = 'unspecified'
CPU = 'cpu'
GPU = 'gpu'
GPU_2DEVS_BACKEND = '2gpus'
TPU = 'tpu'
TPU_STREAM_EXECUTOR = 'tpu_se'
TPU_V3_DONUT_BACKEND = 'tpu_v3_2x2'
TPU_V4_DONUT_BACKEND = 'tpu_v4_2x2'
DTENSOR_TEST_UTIL_BACKEND = DTensorTestUtilBackend(
os.getenv('DTENSOR_TEST_UTIL_BACKEND', default='unspecified')
)
# LINT.ThenChange()
tensorflow/dtensor/python/tests/test_backend_util.py
+69
View File
@@ -0,0 +1,69 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility to set up DTensor backend in tests."""
# LINT.IfChange
import multiprocessing
import os
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python.tests.test_backend_name import DTENSOR_TEST_UTIL_BACKEND
from tensorflow.python.platform import test as tf_test
class DTensorTestBackendConfigurator:
"""Configurate test backends."""
def __init__(self, test_case: tf_test.TestCase):
self._test_case = test_case
# TODO(b/260771689): Refactor common backend set up logic to here.
def tearDown(self):
# Only need to explicitly shuts down TPU system in TFRT since in current
# runtime, the shutdown is done in initialization process.
if accelerator_util.is_initialized():
accelerator_util.shutdown_accelerator_system()
def slice_host_devices_for_multiworker(num_clients, client_id, ports):
"""Configure the current process to only use a slice of devices."""
if num_clients == 0:
# All GPUs are visible to the client.
del os.environ['CUDA_VISIBLE_DEVICES']
del os.environ['HIP_VISIBLE_DEVICES']
else:
# Make the client_id-th GPU visible to the client.
os.environ['CUDA_VISIBLE_DEVICES'] = f'{client_id}'
os.environ['HIP_VISIBLE_DEVICES'] = f'{client_id}'
# Make the client_id-th (4x) TPU cores visible to the client.
os.environ['CLOUD_TPU_TASK_ID'] = f'{client_id}'
if 'tpu' in DTENSOR_TEST_UTIL_BACKEND.value:
del ports # Unused due to lack of implementation.
# We need to find out if there is a way to slice a CloudTPU host to
# multiple workers.
raise NotImplementedError(
'OSS multi-client tests of TPU is not supported.'
)
def get_mp_context():
return multiprocessing.get_context('forkserver')
def handle_test_main(main, *args, **kwargs):
main(*args, **kwargs)
# LINT.ThenChange(test_backend_util.py)
tensorflow/dtensor/python/tests/test_util.py
+419
View File
@@ -0,0 +1,419 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility methods for DTensor testing."""
import collections
import copy
import itertools
import json
import os
import typing
from absl import flags
from absl.testing import parameterized
import numpy as np
# pylint: disable=g-direct-tensorflow-import
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python.config import is_gpu_present # pylint: disable=unused-import
from tensorflow.dtensor.python.config import is_tpu_present # pylint: disable=unused-import
from tensorflow.dtensor.python.config import preferred_device_type # pylint: disable=unused-import
from tensorflow.dtensor.python.config import use_multi_device_mode # pylint: disable=unused-import
from tensorflow.dtensor.python.tests.test_backend_name import DTENSOR_TEST_UTIL_BACKEND
from tensorflow.dtensor.python.tests.test_backend_name import DTensorTestUtilBackend
from tensorflow.dtensor.python.tests.test_backend_util import DTensorTestBackendConfigurator
from tensorflow.python.compat import v2_compat
from tensorflow.python.eager import context
from tensorflow.python.framework import config as tf_config
from tensorflow.python.framework import device as tf_device
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.platform import test as tf_test
from tensorflow.python.util import numpy_compat
# pylint: enable=g-direct-tensorflow-import
# DTensor only runs with TF V2.
v2_compat.enable_v2_behavior()
DEFAULT_TOL = 1e-5
_DEFAULT_GPU_MEMORY_LIMIT = 1024 # 1G
def get_use_xla_spmd(device_type):
"""Returns True when device_type is TPU and environment variable is set.
Args:
device_type: A str representing the type of device on the mesh.
Returns:
bool: True when device_type is TPU and environment variable is set.
"""
return device_type == 'TPU' and '0' != os.environ.get(
'DTENSOR_TEST_USE_XLA_SPMD', '0'
)
def create_device_ids_array(shape):
device_count = np.prod(shape)
return np.arange(device_count).reshape(shape)
def create_device_array(shape, device_type):
device_count = np.prod(shape)
return numpy_compat.np_asarray([
tf_device.DeviceSpec( # pylint: disable=g-complex-comprehension
job='localhost/replica:0/task:0',
device_type=device_type,
device_index=i,
)
for i in range(device_count)
]).reshape(shape)
def create_device_list(shape, device_type):
devices = create_device_array(shape, device_type)
return np.ravel(devices).tolist()
def reset_context():
context._reset_context() # pylint: disable=protected-access
def reset_logical_devices(device_type, count):
"""Resets logical devices for CPU/GPU.
Logical devices can only be instantiated once on a particular context. For
now, context re-use is triggering some function duplication errors, so we
reset the context on each call.
Args:
device_type: The device_type to reset.
count: numbers of virtual device to reset to.
"""
reset_context()
devices = tf_config.list_physical_devices(device_type)
if device_type.upper() not in ('CPU', 'GPU'):
raise ValueError(
'resetting logical device for non-supported device type : %s'
% device_type
)
if count < len(devices):
devices = devices[:count]
tf_config.set_visible_devices(devices, device_type=device_type.upper())
for i, device in enumerate(devices):
n = (i + 1) * count // len(devices) - i * count // len(devices)
assert n > 0 # guaranteed if count >= len(devices)
configs = []
for ordinal in range(n):
if device_type.upper() == 'GPU':
dev_config = context.LogicalDeviceConfiguration(
memory_limit=_DEFAULT_GPU_MEMORY_LIMIT,
experimental_device_ordinal=ordinal,
)
else:
dev_config = context.LogicalDeviceConfiguration()
configs.append(dev_config)
tf_config.set_logical_device_configuration(device, configs)
def list_local_logical_devices(device_type):
"""Returns a list of local logial devices."""
# When coordinator service is enabled, list_logical_devices returns
# a global list.
devices = tf_config.list_logical_devices(device_type)
def is_local(device):
spec = tf_device.DeviceSpec.from_string(device.name)
if spec.job is None or spec.job == 'localhost':
return True
elif spec.job == config.job_name() and spec.task == config.client_id():
return True
return False
return [d for d in devices if is_local(d)]
def is_tfrt_enabled():
return context.is_tfrt_enabled()
FLAGS = flags.FLAGS
class DTensorBaseTest(tf_test.TestCase, parameterized.TestCase):
"""Provides comparison helper for dtensor vs local results."""
@classmethod
def setUpClass(cls):
super(DTensorBaseTest, cls).setUpClass()
def setUp(self):
super().setUp()
self._backend_configurator = DTensorTestBackendConfigurator(self)
def tearDown(self):
# Make sure all async ops finish.
try:
context.async_wait()
finally:
# TODO(hthu): Remove the reset once we fixed the CopyToMesh with
# DefaultMesh placement issue.
reset_dtensor()
self._backend_configurator.tearDown()
super().tearDown()
@staticmethod
def configTestMesh( # pylint: disable=invalid-name
device_type_mesh_map: typing.Dict[typing.Text, layout_lib.Mesh],
) -> layout_lib.Mesh:
"""Configs corresponding mesh given test context.
If runs on a CPU mesh, set virtual device on CPU.
If runs on a GPU mesh, sets virtual device on GPU with proper memory limits.
if runs on a TPU mesh, initializes TPU system.
Args:
device_type_mesh_map: A dictionary containing device_type -> mesh mapping.
Returns:
A properly configured mesh for use in test.
"""
reset_context()
def get_mesh(device_type):
mesh = device_type_mesh_map.get(device_type, None)
if mesh is None:
raise ValueError(
'Requires a %s mesh to run test on %s.' % (device_type, device_type)
)
return mesh
mesh = None
if is_tpu_present():
mesh = get_mesh('TPU')
reset_context()
accelerator_util.initialize_accelerator_system('TPU')
elif tf_config.list_physical_devices('GPU'):
mesh = get_mesh('GPU')
reset_logical_devices('GPU', np.prod(mesh.shape()))
accelerator_util.initialize_accelerator_system('GPU')
else:
mesh = get_mesh('CPU')
reset_logical_devices('CPU', np.prod(mesh.shape()))
accelerator_util.initialize_accelerator_system('CPU')
return mesh
def skipForDeviceType( # pylint: disable=invalid-name
self,
device_type: typing.List[str],
reason: str,
unless_device_count_equals_to=None,
):
"""Skip the test for the specific device_type.
Args:
device_type: list of device types, one of "CPU", "GPU", or "TPU".
reason: string that describe the reason for skipping the test.
unless_device_count_equals_to: Optional int. This parameter only works if
device_type is "TPU". If set, the test will be skipped unless the number
of TPUs equals to the specified count.
"""
physical_device_types = set(
[d.device_type for d in tf_config.list_physical_devices()]
)
for device in device_type:
if device == 'TPU' and is_tpu_present():
if unless_device_count_equals_to is None:
self.skipTest(reason)
elif (
len(list_local_logical_devices(device))
!= unless_device_count_equals_to
):
self.skipTest(reason)
if (
device == 'CPU'
and len(physical_device_types) == 1
and 'CPU' in physical_device_types
):
# Make sure we skip when only `CPU` is present.
self.skipTest(reason)
if device == 'GPU' and 'GPU' in physical_device_types:
self.skipTest(reason)
def skipForTfrt(self, reason: str): # pylint: disable=invalid-name
if is_tfrt_enabled():
self.skipTest(reason)
def skipTest(self, reason): # pylint: disable=invalid-name
# skipTest() may be called in super().setUp()
if hasattr(self, '_backend_configurator'):
self._backend_configurator.tearDown()
super().skipTest(reason)
def assertDTensorEqual(
self, # pylint: disable=invalid-name
expected_result,
expected_layout,
result_dtensor,
tol=DEFAULT_TOL,
):
"""Asserts DTensor is of the particular value."""
if issubclass(
type(result_dtensor), resource_variable_ops.BaseResourceVariable
):
result_dtensor = result_dtensor.value()
if expected_layout is not None:
# This, the assertEqual, is a pure proto raw bytes comparison. To make it
# human-readable, use the `to_string` api for Layout for the dedicated msg
# field.
#
# Futhurmore, as the mesh part is very long and usually identical. Try to
# cut them as well, to make it easier to read.
expected_str = expected_layout.to_string()
got_str = api.fetch_layout(result_dtensor).to_string()
index_for_mesh = expected_str.find('mesh:')
if (
index_for_mesh != -1
and got_str.find(expected_str[index_for_mesh:]) != -1
):
# the mesh part is same. cut them so it is more readable.
expected_str = expected_str[:index_for_mesh]
got_str = got_str[: got_str.find('mesh:')]
self.assertEqual(
api.fetch_layout(result_dtensor),
expected_layout,
msg=(
'=======\nexpected layout is\n {}\n\nwhile got layout is\n {}\n'
.format(expected_str, got_str)
),
)
layout = api.fetch_layout(result_dtensor)
unpacked = [t.numpy() for t in api.unpack(result_dtensor)]
# Check global shape.
self.assertAllEqual(expected_result.shape, result_dtensor.shape)
result_dtensor = numpy_util.to_numpy(result_dtensor)
# Check dtype.
# Note: This check needs be after result_dtensor is converted
# into numpy, due to failure with Numpy version 1.18.5.
self.assertEqual(
expected_result.dtype, result_dtensor.dtype, result_dtensor
)
# Check value on concatenated result DTensor.
self.assertAllClose(expected_result, result_dtensor, atol=tol, rtol=tol)
# In addition to check the 'concatenated' DTensor, we also check all
# "replicated" parts are same.
#
# The algorithm is simple:
# 1. For a mesh with topology (x,y,z,p), and a DTensor with layout ('',z,x).
# 2. Create some data structures:
# - create a mapping from device id (called offset below) to mesh
# location. For the mesh above, loc {x:1,y:2,z:2,p:0} means the device
# is located at that coordinates in the 4-D mesh.
# - create a mapping from mesh location to device id.
# 3. Find all replicated mesh dimension names, i.e., 'y' and `p` in the
# example above.
# 4. Iterate over all unpacked components, translate the offset (device id)
# to mesh location, called (x',y',z',p').
# - For `y`, which is replicated dim in the mesh, check all unpacked
# components at (x',*,z',p') are same as the component at (x',0,z',p').
# - For `p`, which is also replicated dim in the mesh, check all unpacked
# components at (x',y',z',*) are same as the component at (x',y',z',0).
def hash_key(loc):
"""Hash key for Python dict."""
# Python dict is unhashable. Creates a sorted dict and dumps as json str.
d = collections.OrderedDict(sorted(loc.items(), key=lambda x: x[0]))
return json.dumps(d)
offset_to_mesh_loc_dict = layout.mesh.unravel_index()
mesh_loc_to_offset_dict = {}
for offset, loc in offset_to_mesh_loc_dict.items():
mesh_loc_to_offset_dict[hash_key(loc)] = offset
# pylint: disable=protected-access
replicated_dims = [
x for x in layout.mesh.dim_names if x not in layout.sharding_specs
]
# pylint: enable=protected-access
for offset, tensor in enumerate(unpacked):
mesh_loc = offset_to_mesh_loc_dict[offset]
for dim_sharding in replicated_dims:
if mesh_loc[dim_sharding] != 0:
mesh_loc = copy.deepcopy(mesh_loc) # deepcopy as we will mutate
mesh_loc[dim_sharding] = 0
offset = mesh_loc_to_offset_dict[hash_key(mesh_loc)]
# tol is be as low as possible as they should match "exactly". so, we
# ignore the `tol` passed by caller and choose the default one.
self.assertAllClose(tensor, unpacked[offset])
def product(*lists):
"""Makes a product of names parameters list."""
# Each element lists should be a tuple of tuples of the form
# (("test1", ...), ("test2", ...), ...).
# Function returns the product of the lists with the labels concatenated.
return [ # pylint: disable=g-complex-comprehension
(''.join(p[0] for p in elt), *sum((p[1:] for p in elt), ()))
for elt in itertools.product(*lists)
]
def reset_dtensor():
"""Resets the singleton DTensor Device.
This behavior is not generally exposed and only meant to be used in tests.
"""
api._reset() # pylint: disable=protected-access
__all__ = [
'DEFAULT_TOL',
'DTensorTestUtilBackend',
'DTENSOR_TEST_UTIL_BACKEND',
'create_device_ids_array',
'create_device_array',
'create_device_list',
'reset_context',
'reset_logical_devices',
'list_local_logical_devices',
'is_tfrt_enabled',
'FLAGS',
'DTensorBaseTest',
'product',
'reset_dtensor',
'is_tpu_present',
'is_gpu_present',
'use_multi_device_mode',
]
tensorflow/dtensor/python/tests/test_util_ops.py
+658
View File
@@ -0,0 +1,658 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility methods for DTensor testing."""
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import clip_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import gen_bitwise_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import gen_spectral_ops
from tensorflow.python.ops import gen_stateless_random_ops
from tensorflow.python.ops import gen_stateless_random_ops_v2
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import special_math_ops
def expand_test_config(op_list, test_configs):
"""Returns a list of test case args that covers ops and test_configs.
The list is a Cartesian product between op_list and test_configs.
Args:
op_list: A list of dicts, with items keyed by 'testcase_name' and 'op'.
Available lists are defined later in this module.
test_configs: A list of dicts, additional kwargs to be appended for each
test parameters.
Returns:
test_configurations: a list of test parameters that covers all
provided ops in op_list and args in test_configs.
"""
test_configurations = []
for op_info in op_list:
test_index = 0
for added_test_config in test_configs:
test_config = op_info.copy()
test_config.update(added_test_config)
test_config['testcase_name'] = op_info['testcase_name'] + '_' + str(
test_index)
test_index += 1
test_configurations.append(test_config)
return test_configurations
# Disable pyformat for this block to force compact style.
# pyformat: disable
#
# Disable g-long-lambda to make the unit test suits compact (avoid def new func)
# pylint: disable=g-long-lambda
UNARY_OPS = [
{
'testcase_name': 'Identity',
'op': array_ops.identity
},
{
'testcase_name': 'ZerosLike',
'op': array_ops.zeros_like_v2
},
{
'testcase_name': 'Abs',
'op': math_ops.abs
},
{
'testcase_name': 'Negative',
'op': gen_math_ops.neg
},
{
'testcase_name': 'Cast',
'op': lambda x: math_ops.cast(x, dtypes.int32)
},
{
'testcase_name': 'ErfOp',
'op': gen_math_ops.erf
},
{
'testcase_name': 'Softmax',
'op': nn_ops.softmax_v2
},
{
'testcase_name': 'LogSoftmax',
'op': nn_ops.log_softmax_v2
},
{
'testcase_name': 'StopGradient',
'op': array_ops.stop_gradient
},
{
'testcase_name': 'Exp',
'op': math_ops.exp
},
{
'testcase_name': 'Sqrt',
'op': math_ops.sqrt
},
{
'testcase_name': 'Rsqrt',
'op': math_ops.rsqrt
},
{
'testcase_name': 'Reciprocal',
'op': gen_math_ops.reciprocal
},
{
'testcase_name': 'Relu',
'op': gen_nn_ops.relu
},
{
'testcase_name': 'Square',
'op': gen_math_ops.square
},
{
'testcase_name': 'Tanh',
'op': gen_math_ops.tanh
},
{
'testcase_name': 'Cos',
'op': gen_math_ops.cos
},
{
'testcase_name': 'Sigmoid',
'op': math_ops.sigmoid
},
{
'testcase_name': 'Acos',
'op': math_ops.acos
},
{
'testcase_name': 'Acosh',
'op': gen_math_ops.acosh
},
{
'testcase_name': 'Angle',
'op': math_ops.angle
},
{
'testcase_name': 'Asin',
'op': gen_math_ops.asin
},
{
'testcase_name': 'Asinh',
'op': gen_math_ops.asinh
},
{
'testcase_name': 'Atan',
'op': gen_math_ops.atan
},
{
'testcase_name': 'Bessel0e',
'op': special_math_ops.bessel_i0e
},
{
'testcase_name': 'Bessel1e',
'op': special_math_ops.bessel_i1e
},
{
'testcase_name': 'Bitcast',
'op': lambda x: gen_array_ops.bitcast(x, type=dtypes.int32)
},
{
'testcase_name': 'Ceil',
'op': math_ops.ceil
},
{
'testcase_name': 'CheckNumbers',
'op': (lambda x: gen_array_ops.check_numerics(x, message='bug'))
},
{
'testcase_name': 'ClipByValue',
'op': (lambda x: clip_ops.clip_by_value(x, 1.5, 2.5))
},
{
'testcase_name': 'Conj',
'op': math_ops.conj
},
{
'testcase_name': 'Cosh',
'op': gen_math_ops.cosh
},
{
'testcase_name': 'Digamma',
'op': gen_math_ops.digamma
},
{
'testcase_name':
'ComplexAbs',
'op':
lambda x: gen_math_ops.complex_abs(
x=math_ops.cast(x, dtypes.complex64), Tout=float, name='raw')
},
{
'testcase_name': 'Sign',
'op': math_ops.sign
},
{
'testcase_name': 'Elu',
'op': gen_nn_ops.elu
},
{
'testcase_name': 'Erfc',
'op': gen_math_ops.erfc
},
{
'testcase_name': 'Expm1',
'op': gen_math_ops.expm1
},
{
'testcase_name': 'Floor',
'op': math_ops.floor
},
{
'testcase_name': 'Imag',
'op': math_ops.imag
},
{
'testcase_name': 'Inv',
'op': (lambda x: gen_math_ops.inv(x=x, name='Inv'))
},
{
'testcase_name': 'IsInf',
'op': gen_math_ops.is_inf
},
{
'testcase_name': 'IsNan',
'op': gen_math_ops.is_nan
},
{
'testcase_name': 'LeakyRelu',
'op': (lambda x: nn_ops.leaky_relu((x - 2), alpha=0.3)),
},
{
'testcase_name': 'Lgamma',
'op': gen_math_ops.lgamma,
},
{
'testcase_name': 'Log1p',
'op': gen_math_ops.log1p,
},
{
'testcase_name': 'Ndtri',
'op': (lambda x: math_ops.ndtri(x / 100)),
},
{
'testcase_name': 'Selu',
'op': gen_nn_ops.selu,
},
{
'testcase_name': 'Sin',
'op': gen_math_ops.sin,
},
{
'testcase_name': 'Sinh',
'op': gen_math_ops.sinh,
},
{
'testcase_name': 'Softplus',
'op': math_ops.softplus,
},
{
'testcase_name': 'Softsign',
'op': gen_nn_ops.softsign,
},
{
'testcase_name': 'Tan',
'op': gen_math_ops.tan,
},
{
'testcase_name': 'Round',
'op': math_ops.round,
},
{
'testcase_name': 'Rint',
'op': gen_math_ops.rint,
},
{
'testcase_name': 'Relu6',
'op': nn_ops.relu6,
},
{
'testcase_name': 'Real',
'op': math_ops.real,
},
{
'testcase_name': 'PreventGradient',
'op': lambda x: gen_array_ops.prevent_gradient(input=x),
},
]
BINARY_ANY_TYPE_OPS_WITH_BROADCASTING_SUPPORT = [
{
'testcase_name': 'Add',
'op': math_ops.add
},
{
'testcase_name': 'Subtract',
'op': math_ops.subtract
},
{
'testcase_name': 'Multiply',
'op': math_ops.multiply
},
{
'testcase_name': 'Maximum',
'op': gen_math_ops.maximum
},
{
'testcase_name': 'Minimum',
'op': gen_math_ops.minimum
},
{
'testcase_name': 'Squared_Difference',
'op': gen_math_ops.squared_difference
},
{
'testcase_name': 'GreaterEqual',
'op': gen_math_ops.greater_equal
},
{
'testcase_name': 'Equal',
'op': math_ops.equal
},
{
'testcase_name': 'NotEqual',
'op': math_ops.not_equal
},
{
'testcase_name': 'LessEqual',
'op': gen_math_ops.less_equal
},
{
'testcase_name': 'Less',
'op': gen_math_ops.less
},
{
'testcase_name': 'Pow',
'op': math_ops.pow
},
]
BINARY_FLOAT_OPS_WITH_BROADCASTING_SUPPORT = [
{
'testcase_name': 'Real_Divide',
'op': math_ops.divide
},
{
'testcase_name': 'DivNoNan',
'op': math_ops.div_no_nan
},
] + BINARY_ANY_TYPE_OPS_WITH_BROADCASTING_SUPPORT
BINARY_INT_OPS_WITH_BROADCASTING_SUPPORT = [
{
'testcase_name': 'LeftShift',
'op': gen_bitwise_ops.left_shift
},
{
'testcase_name': 'RightShift',
'op': gen_bitwise_ops.right_shift
},
{
'testcase_name': 'BitwiseOr',
'op': gen_bitwise_ops.bitwise_or
},
{
'testcase_name': 'BitwiseAnd',
'op': gen_bitwise_ops.bitwise_and
},
{
'testcase_name': 'BitwiseXor',
'op': gen_bitwise_ops.bitwise_xor
},
{
'testcase_name': 'TruncateDiv',
'op': gen_math_ops.truncate_div
},
{
'testcase_name': 'TruncateMod',
'op': gen_math_ops.truncate_mod
},
] + BINARY_ANY_TYPE_OPS_WITH_BROADCASTING_SUPPORT
BINARY_BOOL_OPS = [{
'testcase_name': 'LogicalOr',
'op': gen_math_ops.logical_or
}]
BINARY_FLOAT_OPS = [
{
'testcase_name': 'RsqrtGrad',
'op': lambda y, dy: gen_math_ops.rsqrt_grad(y=y, dy=dy)
},
{
'testcase_name': 'SqrtGrad',
'op': lambda y, dy: gen_math_ops.sqrt_grad(y=y, dy=dy)
},
{
'testcase_name': 'Atan2',
'op': gen_math_ops.atan2
},
{
'testcase_name': 'Betainc',
'op': lambda a, b: gen_math_ops.betainc(a, b, 1.0)
},
{
'testcase_name': 'Complex',
'op': math_ops.complex
},
{
'testcase_name':
'EluGrad',
'op': (lambda x, y: gen_nn_ops.elu_grad(
gradients=x, outputs=y, name='op_elugrad'))
},
{
'testcase_name': 'Igamma',
'op': gen_math_ops.igamma
},
{
'testcase_name':
'IgammaGradA',
'op': (lambda a, x: gen_math_ops.igamma_grad_a(
a=a, x=x, name='IgammaGradA'))
},
{
'testcase_name':
'LeakyReluGrad',
'op':
(lambda x, y: gen_nn_ops.leaky_relu_grad(gradients=x, features=y)),
},
{
'testcase_name': 'MulNoNan',
'op': (lambda x, y: gen_math_ops.mul_no_nan(x=x, y=y)),
},
{
'testcase_name': 'NextAfter',
'op': gen_math_ops.next_after,
},
{
'testcase_name': 'PolyGamma',
'op': gen_math_ops.polygamma,
},
{
'testcase_name': 'SeluGrad',
'op': (lambda x, y: gen_nn_ops.selu_grad(gradients=x, outputs=y)),
},
{
'testcase_name': 'Relu6Grad',
'op': (lambda x, y: gen_nn_ops.relu6_grad(gradients=x, features=y)),
},
{
'testcase_name': 'ReciprocalGrad',
'op': (lambda x, y: gen_math_ops.reciprocal_grad(y=x, dy=y)),
},
{
'testcase_name': 'Xdivy',
'op': math_ops.xdivy,
},
{
'testcase_name': 'Xlog1py',
'op': math_ops.xlog1py,
},
{
'testcase_name': 'Xlogy',
'op': gen_math_ops.xlogy,
},
{
'testcase_name': 'Zeta',
'op': gen_math_ops.zeta,
},
] + BINARY_FLOAT_OPS_WITH_BROADCASTING_SUPPORT
BINARY_INT_OPS = [] + BINARY_INT_OPS_WITH_BROADCASTING_SUPPORT
REDUCTION_OPS = [
{
'testcase_name': 'Sum',
'op': math_ops.reduce_sum
},
{
'testcase_name': 'Mean',
'op': math_ops.reduce_mean
},
{
'testcase_name': 'Prod',
'op': math_ops.reduce_prod
},
{
'testcase_name': 'Max',
'op': math_ops.reduce_max
},
{
'testcase_name': 'Min',
'op': math_ops.reduce_min
},
]
# TODO(b/171746536): added v2 rng ops here once supported.
RANDOM_OPS = [{
'testcase_name': 'StatelessNorm',
'op': gen_stateless_random_ops.stateless_random_normal,
'dtype': dtypes.float32,
'op_version': 'V1'
}, {
'testcase_name': 'StatelessTruncatedNorm',
'op': gen_stateless_random_ops.stateless_truncated_normal,
'dtype': dtypes.float32,
'op_version': 'V1'
}, {
'testcase_name': 'StatelessUniform',
'op': gen_stateless_random_ops.stateless_random_uniform,
'dtype': dtypes.float32,
'op_version': 'V1'
}, {
'testcase_name': 'StatelessUniformFullInt',
'op': gen_stateless_random_ops.stateless_random_uniform_full_int,
'dtype': dtypes.int32,
'op_version': 'V1'
}, {
'testcase_name': 'StatelessRandomUniformFullIntV2',
'op': gen_stateless_random_ops_v2.stateless_random_uniform_full_int_v2,
'dtype': dtypes.int32,
'op_version': 'V2'
}, {
'testcase_name': 'StatelessRandomNormalV2',
'op': gen_stateless_random_ops_v2.stateless_random_normal_v2,
'dtype': dtypes.float32,
'op_version': 'V2'
}, {
'testcase_name': 'StatelessTruncatedNormalV2',
'op': gen_stateless_random_ops_v2.stateless_truncated_normal_v2,
'dtype': dtypes.float32,
'op_version': 'V2'
}, {
'testcase_name': 'StatelessRandomUniformV2',
'op': gen_stateless_random_ops_v2.stateless_random_uniform_v2,
'dtype': dtypes.float32,
'op_version': 'V2'
}, {
'testcase_name': 'StatelessRandomUniformIntV2',
'op': gen_stateless_random_ops_v2.stateless_random_uniform_int_v2,
'dtype': dtypes.int32,
'op_version': 'V2_RANGE'
}]
# op(inputs()) is expected to return an NxM tensor (N, M both even) with a
# flexible output sharding, depending on the context `op` runs in.
EXPANSION_OPS = [
dict(
testcase_name='TileFrom1x1Array',
inputs=lambda: (constant_op.constant([[1.]]), [4, 4]),
op=gen_array_ops.tile),
dict(
testcase_name='TileFrom2x2Array',
inputs=lambda: (constant_op.constant([[1., 2.], [3., 4.]]), [2, 4]),
op=gen_array_ops.tile),
dict(
testcase_name='Fill',
inputs=lambda: ([2, 4], constant_op.constant(1.)),
op=array_ops.fill),
]
BATCH_PARALLEL_2D_WINDOW_OPS = [(
'AvgPool',
nn_ops.avg_pool_v2,
)]
BATCH_PARALLEL_3D_WINDOW_OPS = [(
'MaxPool3D',
nn_ops.max_pool3d,
), (
'AvgPool3D',
nn_ops.avg_pool3d,
)]
FFT_OPS = [(
'FFT',
gen_spectral_ops.fft,
1,
), (
'FFT2D',
gen_spectral_ops.fft2d,
2,
), (
'FFT3D',
gen_spectral_ops.fft3d,
3,
), (
'IFFT',
gen_spectral_ops.ifft,
1,
), (
'IFFT2D',
gen_spectral_ops.ifft2d,
2,
), (
'IFFT3D',
gen_spectral_ops.ifft3d,
3,
)]
RFFT_OPS = [(
'IRFFT',
gen_spectral_ops.irfft,
1,
dtypes.complex64,
), (
'IRFFT2D',
gen_spectral_ops.irfft2d,
2,
dtypes.complex64,
), (
'IRFFT3D',
gen_spectral_ops.irfft3d,
3,
dtypes.complex64,
), (
'RFFT',
gen_spectral_ops.rfft,
1,
dtypes.float32,
), (
'RFFT2D',
gen_spectral_ops.rfft2d,
2,
dtypes.float32,
), (
'RFFT3D',
gen_spectral_ops.rfft3d,
3,
dtypes.float32,
)]
PADDINGS = [
{
'testcase_name': 'SamePadding',
'padding': 'SAME'
},
{
'testcase_name': 'ValidPadding',
'padding': 'VALID'
},
]
# pylint: enable=g-long-lambda
# pyformat: enable
tensorflow/dtensor/python/tests/tpu_device_assignment_test.py
+889
View File
@@ -0,0 +1,889 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for TPU device assignment."""
from tensorflow.dtensor.python import accelerator_util
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python import numpy_util
from tensorflow.dtensor.python import tpu_util
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
Layout = layout_lib.Layout
Mesh = layout_lib.Mesh
class DeviceAssignmentTest(test_util.DTensorBaseTest):
def setUp(self):
super().setUp()
accelerator_util.initialize_accelerator_system('TPU')
def tearDown(self):
accelerator_util.shutdown_accelerator_system()
super().tearDown()
def _build_all_reduce_ring(self, core_locations):
permutation = tpu_util._build_all_reduce_ring(core_locations)
return [core_locations[element] for element in permutation]
# Picture of chips:
# 0 -- 1
# | |
# 3 -- 2
def testBuildAllReduceRing4Replicas(self):
core_locations = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 0),
]
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 0),
]
result = self._build_all_reduce_ring(core_locations)
self.assertAllEqual(result, expected)
# Picture of chips with core0/core1 assignments:
# 0/1 -- 2/3
# | |
# 6/7 -- 4/5
def testBuildAllReduceRing8ReplicasUsingTwoCores(self):
core_locations = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 1),
]
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
]
result = self._build_all_reduce_ring(core_locations)
self.assertAllEqual(result, expected)
# Picture of chips:
# 0 -- 1 -- 2 -- 3
# | |
# 15 6 -- 5 -- 4
# | |
# 14 7 -- 8 -- 9
# | |
# 13-- 12-- 11-- 10
def testBuildAllReduceRing32Replicas(self):
core_locations = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
]
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
]
result = self._build_all_reduce_ring(core_locations)
self.assertAllEqual(result, expected)
# Picture of chips:
# 7 -- 0 6 -- 5
# | |
# 2 -- 1 3 -- 4
def testBuildAllReduceRing3D(self):
core_locations = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 0, 1, 0),
tpu_util._CoreLocation(0, 0, 1, 1),
tpu_util._CoreLocation(0, 1, 1, 0),
tpu_util._CoreLocation(0, 1, 1, 1),
tpu_util._CoreLocation(1, 0, 1, 0),
tpu_util._CoreLocation(1, 0, 1, 1),
tpu_util._CoreLocation(1, 1, 1, 0),
tpu_util._CoreLocation(1, 1, 1, 1),
]
expected = [
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(0, 1, 1, 1),
tpu_util._CoreLocation(0, 1, 1, 0),
tpu_util._CoreLocation(1, 1, 1, 1),
tpu_util._CoreLocation(1, 1, 1, 0),
tpu_util._CoreLocation(1, 0, 1, 1),
tpu_util._CoreLocation(1, 0, 1, 0),
tpu_util._CoreLocation(0, 0, 1, 0),
tpu_util._CoreLocation(0, 0, 1, 1),
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
]
result = self._build_all_reduce_ring(core_locations)
self.assertAllEqual(result, expected)
# Picture of chips:
# 31-- 0 -- 1 -- 2 30--29--28--27
# | |
# 14 5 -- 4 -- 3 15 24--25--26
# | | | |
# 13 6 -- 7 -- 8 16 23--22--21
# | | | |
# 12-- 11-- 10-- 9 17--18--19--20
def testBuildAllReduceRing3DLarge(self):
core_locations = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
tpu_util._CoreLocation(0, 0, 1, 0),
tpu_util._CoreLocation(0, 0, 1, 1),
tpu_util._CoreLocation(1, 0, 1, 0),
tpu_util._CoreLocation(1, 0, 1, 1),
tpu_util._CoreLocation(2, 0, 1, 0),
tpu_util._CoreLocation(2, 0, 1, 1),
tpu_util._CoreLocation(3, 0, 1, 0),
tpu_util._CoreLocation(3, 0, 1, 1),
tpu_util._CoreLocation(0, 1, 1, 0),
tpu_util._CoreLocation(0, 1, 1, 1),
tpu_util._CoreLocation(1, 1, 1, 0),
tpu_util._CoreLocation(1, 1, 1, 1),
tpu_util._CoreLocation(2, 1, 1, 0),
tpu_util._CoreLocation(2, 1, 1, 1),
tpu_util._CoreLocation(3, 1, 1, 0),
tpu_util._CoreLocation(3, 1, 1, 1),
tpu_util._CoreLocation(0, 2, 1, 0),
tpu_util._CoreLocation(0, 2, 1, 1),
tpu_util._CoreLocation(1, 2, 1, 0),
tpu_util._CoreLocation(1, 2, 1, 1),
tpu_util._CoreLocation(2, 2, 1, 0),
tpu_util._CoreLocation(2, 2, 1, 1),
tpu_util._CoreLocation(3, 2, 1, 0),
tpu_util._CoreLocation(3, 2, 1, 1),
tpu_util._CoreLocation(0, 3, 1, 0),
tpu_util._CoreLocation(0, 3, 1, 1),
tpu_util._CoreLocation(1, 3, 1, 0),
tpu_util._CoreLocation(1, 3, 1, 1),
tpu_util._CoreLocation(2, 3, 1, 0),
tpu_util._CoreLocation(2, 3, 1, 1),
tpu_util._CoreLocation(3, 3, 1, 0),
tpu_util._CoreLocation(3, 3, 1, 1),
]
expected = [
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(0, 1, 1, 1),
tpu_util._CoreLocation(0, 1, 1, 0),
tpu_util._CoreLocation(0, 2, 1, 1),
tpu_util._CoreLocation(0, 2, 1, 0),
tpu_util._CoreLocation(0, 3, 1, 1),
tpu_util._CoreLocation(0, 3, 1, 0),
tpu_util._CoreLocation(1, 3, 1, 1),
tpu_util._CoreLocation(1, 3, 1, 0),
tpu_util._CoreLocation(2, 3, 1, 1),
tpu_util._CoreLocation(2, 3, 1, 0),
tpu_util._CoreLocation(3, 3, 1, 1),
tpu_util._CoreLocation(3, 3, 1, 0),
tpu_util._CoreLocation(3, 2, 1, 1),
tpu_util._CoreLocation(3, 2, 1, 0),
tpu_util._CoreLocation(2, 2, 1, 1),
tpu_util._CoreLocation(2, 2, 1, 0),
tpu_util._CoreLocation(1, 2, 1, 1),
tpu_util._CoreLocation(1, 2, 1, 0),
tpu_util._CoreLocation(1, 1, 1, 1),
tpu_util._CoreLocation(1, 1, 1, 0),
tpu_util._CoreLocation(2, 1, 1, 1),
tpu_util._CoreLocation(2, 1, 1, 0),
tpu_util._CoreLocation(3, 1, 1, 1),
tpu_util._CoreLocation(3, 1, 1, 0),
tpu_util._CoreLocation(3, 0, 1, 1),
tpu_util._CoreLocation(3, 0, 1, 0),
tpu_util._CoreLocation(2, 0, 1, 1),
tpu_util._CoreLocation(2, 0, 1, 0),
tpu_util._CoreLocation(1, 0, 1, 1),
tpu_util._CoreLocation(1, 0, 1, 0),
tpu_util._CoreLocation(0, 0, 1, 0),
tpu_util._CoreLocation(0, 0, 1, 1),
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
]
result = self._build_all_reduce_ring(core_locations)
self.assertAllEqual(result, expected)
# Picture of chips:
# 0 -- 1 4 -- 5
# | | | |
# 3 -- 2 7 -- 6
#
# 12-- 13 8 -- 9
# | | | |
# 15-- 14 11-- 10
def testBuildOrthogonalAllReduceRings(self):
core_locations = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
]
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
]
result = tpu_util._build_orthogonal_rings(
core_locations, ring_size=8, rotate_ring_across_rings=False)
self.assertAllEqual(result, expected)
# Picture of chips:
# 0 -- 1 12 -- 13
# | | | |
# 3 -- 2 15 -- 14
#
# 4 -- 5 8 -- 9
# | | | |
# 7 -- 6 11-- 10
def testBuildOrthogonalRotatedAllReduceRings(self):
core_locations = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
]
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
]
result = tpu_util._build_orthogonal_rings(
core_locations, ring_size=8, rotate_ring_across_rings=True)
self.assertAllEqual(result, expected)
# Create a 4x8 mesh on a 4x4 DF slice, disallowing splitting hosts.
def testCreateDFMeshNoSplittingHosts(self):
result = tpu_util._enumerate_core_locations(
[4, 4, 1, 2], [4, 4, 1, 2], ['core', 'y', 'z', 'x'],
can_split_host_across_rings=False,
ring_size=8)
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
]
self.assertAllEqual(result, expected)
# Create a 4x8 mesh on a 4x4 DF slice with at most 2, 2, 1, 2 devices from
# each dimension, disallowing splitting hosts.
def testCreateDFMeshWithRingBoundsNoSplittingHosts(self):
result = tpu_util._enumerate_core_locations(
[4, 4, 1, 2], [2, 2, 1, 2], ['core', 'x', 'y', 'z'],
can_split_host_across_rings=False,
ring_size=8)
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
]
self.assertAllEqual(result, expected)
# Create a 4x8 mesh on a 4x4 DF slice, allowing splitting hosts.
def testCreateDFMeshSplittingHosts(self):
result = tpu_util._enumerate_core_locations(
[4, 4, 1, 2], [4, 4, 1, 2], ['core', 'y', 'z', 'x'],
can_split_host_across_rings=True,
ring_size=8)
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
]
self.assertAllEqual(result, expected)
# Create a 2x64 mesh on a 4x4x4 PF slice, allowing splitting hosts.
def testCreateMeshPFSplittingHosts(self):
result = tpu_util._enumerate_core_locations(
[4, 4, 4, 2], [4, 4, 4, 2], ['core', 'x', 'y', 'z'],
can_split_host_across_rings=True,
ring_size=64)
expected = [
tpu_util._CoreLocation(0, 0, 0, 0),
tpu_util._CoreLocation(0, 0, 0, 1),
tpu_util._CoreLocation(1, 0, 0, 0),
tpu_util._CoreLocation(1, 0, 0, 1),
tpu_util._CoreLocation(2, 0, 0, 0),
tpu_util._CoreLocation(2, 0, 0, 1),
tpu_util._CoreLocation(3, 0, 0, 0),
tpu_util._CoreLocation(3, 0, 0, 1),
tpu_util._CoreLocation(0, 1, 0, 0),
tpu_util._CoreLocation(0, 1, 0, 1),
tpu_util._CoreLocation(1, 1, 0, 0),
tpu_util._CoreLocation(1, 1, 0, 1),
tpu_util._CoreLocation(2, 1, 0, 0),
tpu_util._CoreLocation(2, 1, 0, 1),
tpu_util._CoreLocation(3, 1, 0, 0),
tpu_util._CoreLocation(3, 1, 0, 1),
tpu_util._CoreLocation(0, 2, 0, 0),
tpu_util._CoreLocation(0, 2, 0, 1),
tpu_util._CoreLocation(1, 2, 0, 0),
tpu_util._CoreLocation(1, 2, 0, 1),
tpu_util._CoreLocation(2, 2, 0, 0),
tpu_util._CoreLocation(2, 2, 0, 1),
tpu_util._CoreLocation(3, 2, 0, 0),
tpu_util._CoreLocation(3, 2, 0, 1),
tpu_util._CoreLocation(0, 3, 0, 0),
tpu_util._CoreLocation(0, 3, 0, 1),
tpu_util._CoreLocation(1, 3, 0, 0),
tpu_util._CoreLocation(1, 3, 0, 1),
tpu_util._CoreLocation(2, 3, 0, 0),
tpu_util._CoreLocation(2, 3, 0, 1),
tpu_util._CoreLocation(3, 3, 0, 0),
tpu_util._CoreLocation(3, 3, 0, 1),
tpu_util._CoreLocation(0, 0, 1, 0),
tpu_util._CoreLocation(0, 0, 1, 1),
tpu_util._CoreLocation(1, 0, 1, 0),
tpu_util._CoreLocation(1, 0, 1, 1),
tpu_util._CoreLocation(2, 0, 1, 0),
tpu_util._CoreLocation(2, 0, 1, 1),
tpu_util._CoreLocation(3, 0, 1, 0),
tpu_util._CoreLocation(3, 0, 1, 1),
tpu_util._CoreLocation(0, 1, 1, 0),
tpu_util._CoreLocation(0, 1, 1, 1),
tpu_util._CoreLocation(1, 1, 1, 0),
tpu_util._CoreLocation(1, 1, 1, 1),
tpu_util._CoreLocation(2, 1, 1, 0),
tpu_util._CoreLocation(2, 1, 1, 1),
tpu_util._CoreLocation(3, 1, 1, 0),
tpu_util._CoreLocation(3, 1, 1, 1),
tpu_util._CoreLocation(0, 2, 1, 0),
tpu_util._CoreLocation(0, 2, 1, 1),
tpu_util._CoreLocation(1, 2, 1, 0),
tpu_util._CoreLocation(1, 2, 1, 1),
tpu_util._CoreLocation(2, 2, 1, 0),
tpu_util._CoreLocation(2, 2, 1, 1),
tpu_util._CoreLocation(3, 2, 1, 0),
tpu_util._CoreLocation(3, 2, 1, 1),
tpu_util._CoreLocation(0, 3, 1, 0),
tpu_util._CoreLocation(0, 3, 1, 1),
tpu_util._CoreLocation(1, 3, 1, 0),
tpu_util._CoreLocation(1, 3, 1, 1),
tpu_util._CoreLocation(2, 3, 1, 0),
tpu_util._CoreLocation(2, 3, 1, 1),
tpu_util._CoreLocation(3, 3, 1, 0),
tpu_util._CoreLocation(3, 3, 1, 1),
tpu_util._CoreLocation(0, 0, 2, 0),
tpu_util._CoreLocation(0, 0, 2, 1),
tpu_util._CoreLocation(1, 0, 2, 0),
tpu_util._CoreLocation(1, 0, 2, 1),
tpu_util._CoreLocation(2, 0, 2, 0),
tpu_util._CoreLocation(2, 0, 2, 1),
tpu_util._CoreLocation(3, 0, 2, 0),
tpu_util._CoreLocation(3, 0, 2, 1),
tpu_util._CoreLocation(0, 1, 2, 0),
tpu_util._CoreLocation(0, 1, 2, 1),
tpu_util._CoreLocation(1, 1, 2, 0),
tpu_util._CoreLocation(1, 1, 2, 1),
tpu_util._CoreLocation(2, 1, 2, 0),
tpu_util._CoreLocation(2, 1, 2, 1),
tpu_util._CoreLocation(3, 1, 2, 0),
tpu_util._CoreLocation(3, 1, 2, 1),
tpu_util._CoreLocation(0, 2, 2, 0),
tpu_util._CoreLocation(0, 2, 2, 1),
tpu_util._CoreLocation(1, 2, 2, 0),
tpu_util._CoreLocation(1, 2, 2, 1),
tpu_util._CoreLocation(2, 2, 2, 0),
tpu_util._CoreLocation(2, 2, 2, 1),
tpu_util._CoreLocation(3, 2, 2, 0),
tpu_util._CoreLocation(3, 2, 2, 1),
tpu_util._CoreLocation(0, 3, 2, 0),
tpu_util._CoreLocation(0, 3, 2, 1),
tpu_util._CoreLocation(1, 3, 2, 0),
tpu_util._CoreLocation(1, 3, 2, 1),
tpu_util._CoreLocation(2, 3, 2, 0),
tpu_util._CoreLocation(2, 3, 2, 1),
tpu_util._CoreLocation(3, 3, 2, 0),
tpu_util._CoreLocation(3, 3, 2, 1),
tpu_util._CoreLocation(0, 0, 3, 0),
tpu_util._CoreLocation(0, 0, 3, 1),
tpu_util._CoreLocation(1, 0, 3, 0),
tpu_util._CoreLocation(1, 0, 3, 1),
tpu_util._CoreLocation(2, 0, 3, 0),
tpu_util._CoreLocation(2, 0, 3, 1),
tpu_util._CoreLocation(3, 0, 3, 0),
tpu_util._CoreLocation(3, 0, 3, 1),
tpu_util._CoreLocation(0, 1, 3, 0),
tpu_util._CoreLocation(0, 1, 3, 1),
tpu_util._CoreLocation(1, 1, 3, 0),
tpu_util._CoreLocation(1, 1, 3, 1),
tpu_util._CoreLocation(2, 1, 3, 0),
tpu_util._CoreLocation(2, 1, 3, 1),
tpu_util._CoreLocation(3, 1, 3, 0),
tpu_util._CoreLocation(3, 1, 3, 1),
tpu_util._CoreLocation(0, 2, 3, 0),
tpu_util._CoreLocation(0, 2, 3, 1),
tpu_util._CoreLocation(1, 2, 3, 0),
tpu_util._CoreLocation(1, 2, 3, 1),
tpu_util._CoreLocation(2, 2, 3, 0),
tpu_util._CoreLocation(2, 2, 3, 1),
tpu_util._CoreLocation(3, 2, 3, 0),
tpu_util._CoreLocation(3, 2, 3, 1),
tpu_util._CoreLocation(0, 3, 3, 0),
tpu_util._CoreLocation(0, 3, 3, 1),
tpu_util._CoreLocation(1, 3, 3, 0),
tpu_util._CoreLocation(1, 3, 3, 1),
tpu_util._CoreLocation(2, 3, 3, 0),
tpu_util._CoreLocation(2, 3, 3, 1),
tpu_util._CoreLocation(3, 3, 3, 0),
tpu_util._CoreLocation(3, 3, 3, 1),
]
self.assertAllEqual(result, expected)
def testCreateMeshNoSplittingHostsUnfulfillable(self):
with self.assertRaises(ValueError):
tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_unfulfillable_without_splitting_hosts',
can_split_host_across_rings=False)
def testCreateMeshWithDefaultOptions(self):
mesh = tpu_util.create_tpu_mesh(['x'], [2], 'mesh_with_default_options')
self.assertAllEqual(mesh.shape(), [2])
self.assertEqual(mesh.num_local_devices(), 2)
def testCreateMeshWithWrongShape(self):
with self.assertRaises(ValueError):
tpu_util.create_tpu_mesh(['x'], [1], 'mesh_with_wrong_shape')
# Build rings for the batch dimension.
def testCreateMeshWithPositiveRingDims(self):
mesh = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_with_positive_ring_dims',
ring_dims=1)
self.assertAllEqual(mesh.shape(), [2, 1])
self.assertEqual(mesh.num_local_devices(), 2)
# Build rings for all non-batch dimensions.
def testCreateMeshWithNegativeRingDims(self):
mesh = tpu_util.create_tpu_mesh(['x', 'y', 'z'], [1, 2, 1],
'mesh_with_negative_ring_dims',
ring_dims=-2)
self.assertAllEqual(mesh.shape(), [1, 2, 1])
self.assertEqual(mesh.num_local_devices(), 2)
# Build single-core rings.
def testCreateMeshWithZeroRingDims(self):
mesh = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_with_zero_ring_dims',
ring_dims=0)
self.assertAllEqual(mesh.shape(), [2, 1])
self.assertEqual(mesh.num_local_devices(), 2)
def testCreateMeshWithCustomAxes(self):
mesh = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_with_custom_axes',
ring_axes=['x', 'z', 'y', 'core'])
self.assertAllEqual(mesh.shape(), [2, 1])
self.assertEqual(mesh.num_local_devices(), 2)
# More cores (2 cores) on the first axis (core) than ring size (1).
def testCreateMeshWithDividedAxis(self):
mesh = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_with_divided_axis',
ring_dims=-1,
ring_axes=['core', 'z', 'y', 'x'])
self.assertAllEqual(mesh.shape(), [2, 1])
self.assertEqual(mesh.num_local_devices(), 2)
# Both meshes should produce the same result despite different `ring_dim`.
def testCreateMultipleMeshes(self):
a = constant_op.constant([[0, 1], [2, 3]], dtype=dtypes.int32)
b_expected = math_ops.reduce_sum(a)
mesh_1 = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1], 'mesh_1', ring_dims=1)
a_1 = numpy_util.pack_numpy(a, Layout(['x', 'y'], mesh_1))
b_1 = math_ops.reduce_sum(a_1)
self.assertDTensorEqual(b_expected, Layout.replicated(mesh_1, rank=0), b_1)
mesh_2 = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_2',
ring_dims=-1)
a_2 = numpy_util.pack_numpy(a, Layout(['x', 'y'], mesh_2))
b_2 = math_ops.reduce_sum(a_2)
self.assertDTensorEqual(b_expected, Layout.replicated(mesh_2, rank=0), b_2)
def testCreateMeshWithEmptyName(self):
tpu_util.create_tpu_mesh(['x'], [2], '')
def testCreateMeshWithExistingName(self):
tpu_util.create_tpu_mesh(['x'], [2], 'mesh_with_existing_name')
with self.assertRaises(ValueError):
tpu_util.create_tpu_mesh(['x'], [2], 'mesh_with_existing_name')
def testGetDeviceIDs(self):
mesh = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_to_get_device_ids')
self.assertAllEqual(tpu_util.get_device_ids(mesh), [0, 1])
def testGetDeviceLocations(self):
mesh = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
'mesh_to_get_device_locations')
self.assertAllEqual(
tpu_util.get_device_locations(mesh), [{
'x': 0,
'y': 0
}, {
'x': 1,
'y': 0
}])
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tests/variable_test.py
+226
View File
@@ -0,0 +1,226 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for DTensor support of Variables."""
import numpy as np
from tensorflow.dtensor.python import api
from tensorflow.dtensor.python import d_variable
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.dtensor.python.tests import test_util
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
# Makes a 1D mesh with dimension X(2).
_MESH_DIM_X = 'x'
_DEVICE_IDS = test_util.create_device_ids_array((2,))
_ONE_D_CPU_MESH = layout_lib.Mesh(
[_MESH_DIM_X],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2,), 'CPU'),
)
_ONE_D_TPU_MESH = layout_lib.Mesh(
[_MESH_DIM_X],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2,), 'TPU'),
)
_ONE_D_GPU_MESH = layout_lib.Mesh(
[_MESH_DIM_X],
_DEVICE_IDS,
np.ravel(_DEVICE_IDS).tolist(),
test_util.create_device_list((2,), 'GPU'),
)
Layout = layout_lib.Layout
UNSHARDED = layout_lib.UNSHARDED
DVariable = d_variable.DVariable
class Var(object):
def __init__(self):
self.v = None
class DTensorVariableTest(test_util.DTensorBaseTest):
def setUp(self):
super(DTensorVariableTest, self).setUp()
mesh_dict = {
'CPU': _ONE_D_CPU_MESH,
'GPU': _ONE_D_GPU_MESH,
'TPU': _ONE_D_TPU_MESH,
}
self.mesh = self.configTestMesh(mesh_dict)
self._replicated_layout = Layout([UNSHARDED, UNSHARDED], self.mesh)
self._one_d_replicated_layout = Layout([UNSHARDED], self.mesh)
self._scalar_replicated_layout = Layout([], self.mesh)
self._one_d_shard_layout = Layout([_MESH_DIM_X], self.mesh)
self._first_d_shard_layout = Layout([_MESH_DIM_X, UNSHARDED], self.mesh)
def testNonDtensorVariable(self):
non_dtensor_variable = variables.Variable(1.0)
with ops.device_v2(api.device_name()):
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
'No default mesh has been registered to DTensor',
):
non_dtensor_variable.read_value()
def testDVariableNoLayout(self):
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
'Neither layout nor DTensor initial value are provided.',
):
DVariable(1.0)
def testDVariableConflictingLayout(self):
a = api.relayout([1, 2, 3, 4], self._one_d_replicated_layout)
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError, 'Conflicting layout are provided'
):
DVariable(a, layout=self._one_d_shard_layout)
def testVariable(self):
with ops.device_v2(api.device_name()):
initial_value = api.relayout([1.0], self._one_d_replicated_layout)
v = variables.Variable(initial_value)
v = api.relayout(v, self._one_d_replicated_layout)
api.check_layout(v, self._one_d_replicated_layout)
def testVariableWithInitialValue(self):
a = constant_op.constant([1.0])
a = api.relayout(a, self._one_d_replicated_layout)
with ops.device_v2(api.device_name()):
v = variables.Variable(initial_value=a)
api.check_layout(v, self._one_d_replicated_layout)
to_add = api.relayout([1.0], self._one_d_replicated_layout)
v = v.assign_add(to_add)
api.check_layout(v, self._one_d_replicated_layout)
def testVarAssignmentOpByOp(self):
v = constant_op.constant(1.0)
v = api.relayout(v, Layout.replicated(self.mesh, rank=0))
w = d_variable.DVariable(v)
api.check_layout(w, Layout.replicated(self.mesh, rank=0))
self.assertEqual(w.numpy(), 1.0)
w.assign_add(v)
api.check_layout(w, Layout.replicated(self.mesh, rank=0))
self.assertEqual(w.numpy(), 2.0)
def testVarInitOutsideTfFunction(self):
v = constant_op.constant(1.0)
v = api.relayout(v, Layout.replicated(self.mesh, rank=0))
w = d_variable.DVariable(v)
@polymorphic_function.function()
def assign_var(x):
w.assign(x * 2)
return w + x
out = assign_var(constant_op.constant(1.0))
api.check_layout(w, Layout.replicated(self.mesh, rank=0))
self.assertEqual(w.numpy(), 2.0)
self.assertEqual(out.numpy(), 3.0)
def testDVariableInitFromValues(self):
# Python value 1 will be converted as dtypes.int32 without a dtype.
non_dtensor_variable = variables.Variable(1, dtype=dtypes.int64)
with ops.device_v2(api.device_name()):
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
'No default mesh has been registered to DTensor',
):
non_dtensor_variable.read_value()
dtensor_variable = DVariable(
api.relayout(
constant_op.constant(1, dtype=dtypes.int64),
Layout.replicated(self.mesh, rank=0),
)
)
with ops.device_v2(api.device_name()):
dtensor_variable = DVariable(
api.relayout(
constant_op.constant(1, dtype=dtypes.int64),
Layout.replicated(self.mesh, rank=0),
)
)
self.assertEqual(dtensor_variable.numpy(), 1)
def testCreateVarInsideFunctionWithInitScope(self):
var = Var()
@polymorphic_function.function
def assign_add():
with ops.init_scope():
if var.v is None:
c = constant_op.constant(1.0)
c = api.relayout(c, Layout.replicated(self.mesh, rank=0))
var.v = variables.Variable(c)
var.v.assign_add(1.0)
with api._dtensor_device()._default_layout(
Layout.replicated(self.mesh, rank=0)):
assign_add()
output = var.v.read_value()
api.check_layout(output, Layout.replicated(self.mesh, rank=0))
self.assertAllEqual(output, 2.)
def testBufferAliasingOnDF(self):
# TODO(b/239471086): re-enable once b/239471086 is fixed
self.skipTest('Disabled due to b/239471086')
self.skipForDeviceType(['GPU', 'CPU'], 'Test only applies to DF TPU')
@polymorphic_function.function
def add_var(v):
new_v = array_ops_stack.stack([v, v])
v.assign(math_ops.reduce_sum(new_v, axis=0))
# Note that this only works with DF. When updated to PF, we need to
# adjust tensor size accordingly.
#
# Without aliasing, the returned tensor is 3.5G * 4 = 14G, plus the
# reserved 500MB and 3.5G arguments, this exceeds the 16G memory.
# With aliasing, the 3.5G arguments will be aliased so it lowers the
# memory pressure to 18-3.5 = 14.5G, which barely fits the memory.
return v, new_v
# Around 3.5G tensor.
v = DVariable(
initial_value=api.relayout(
array_ops.ones((7, 512, 1024, 256), dtype=dtypes.float32),
Layout.replicated(self.mesh, rank=4),
)
)
add_var(v)
self.assertEqual(api.fetch_layout(v), Layout.replicated(self.mesh, rank=4))
if __name__ == '__main__':
test.main()
tensorflow/dtensor/python/tpu_util.py
+879
View File
@@ -0,0 +1,879 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""TPU-specific utilities for DTensor."""
import functools
import time
from typing import Dict, List, Optional
import numpy as np
from tensorflow.dtensor.python import config
from tensorflow.dtensor.python import dtensor_device
from tensorflow.dtensor.python import gen_dtensor_ops
from tensorflow.dtensor.python import layout as layout_lib
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.tpu import topology
from tensorflow.python.util import numpy_compat
from tensorflow.python.util.tf_export import tf_export
_MESH_DIM_X = "x"
_TPU_DEVICE_TYPE = "TPU"
# A dedicated, hidden device used to make C++ API calls.
_dtensor_device = None
# `_topology._mesh_shape` contains the TPU hardware slice size.
# `_topology.device_coordinates` maps TF task-device ordinals to TPU core IDs.
_tpu_topology = None
# Cache core ID <-> location mappings so we need not make repeated C++ calls.
# Both are indexed by TF task-device ordinals.
_all_core_ids = None
_all_core_locations = None
class _CoreLocation:
"""Represents a TPU core's location in the mesh."""
def __init__(self, x: int = 0, y: int = 0, z: int = 0, core: int = 0):
self.x = x
self.y = y
self.z = z
self.core = core
def __eq__(self, other):
if not isinstance(other, _CoreLocation):
return False
return (
self.x == other.x
and self.y == other.y
and self.z == other.z
and self.core == other.core
)
def __ne__(self, other):
if not isinstance(other, _CoreLocation):
return True
return not self == other
def __hash__(self):
return hash((self.x, self.y, self.z, self.core))
def __repr__(self):
return (
f"{type(self).__name__}(x={self.x}, y={self.y}, z={self.z},"
f" core={self.core})"
)
def to_list(self):
return [self.x, self.y, self.z, self.core]
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 = config.num_global_devices(device_type)
global_device_ids = np.arange(num_global_devices).reshape(shape)
local_device_list = config.local_devices(device_type)
# User can specify local_device_ids or use default list for multi host.
num_local_devices = len(local_device_list)
if not local_device_ids:
local_device_ids = [
x + host_id * num_local_devices for x in range(num_local_devices) # pytype: disable=unsupported-operands
]
return global_device_ids, local_device_ids, local_device_list
def _create_tpu_topology(
core_locations: List[_CoreLocation],
num_tasks: int,
num_devices_per_task: int,
) -> topology.Topology:
"""Returns a Topology object build from a _CoreLocation list.
Args:
core_locations: A list of _CoreLocation objects sorted first by TF task ID
and then by per-task device ordinals.
num_tasks: The number of TF tasks in the cluster.
num_devices_per_task: The number of TPU devices local to each task.
"""
assert min([l.x for l in core_locations]) == 0
assert min([l.y for l in core_locations]) == 0
assert min([l.z for l in core_locations]) == 0
assert min([l.core for l in core_locations]) == 0
x_max = max([l.x for l in core_locations])
y_max = max([l.y for l in core_locations])
z_max = max([l.z for l in core_locations])
core_max = max([l.core for l in core_locations])
mesh_shape = [x_max + 1, y_max + 1, z_max + 1, core_max + 1]
device_coordinates = [[l.x, l.y, l.z, l.core] for l in core_locations]
device_coordinates = numpy_compat.np_asarray(device_coordinates).reshape(
num_tasks, num_devices_per_task, 4
)
return topology.Topology(
mesh_shape=mesh_shape, device_coordinates=device_coordinates
)
def shutdown_tpu_system():
"""Shuts down the TPU system."""
@def_function.function
def _shutdown_tpu_system():
return gen_dtensor_ops.shutdown_tpu_system()
success = _shutdown_tpu_system() if context.is_tfrt_enabled() else True
if success:
logging.info("TPU system shut down.")
else:
logging.warning("TPU system fails to shut down.")
def tpu_system_init_helper(
task_id,
num_tasks,
num_devices,
use_tfrt_host_runtime=True,
use_megacore=False,
):
"""A helper function to initialize multi-client tpu system."""
@def_function.function
def _tpu_init_fn():
return gen_dtensor_ops.configure_and_initialize_global_tpu(
use_tfrt_host_runtime=use_tfrt_host_runtime
)
@def_function.function
def _set_global_tpu_array_fn(topology_proto):
gen_dtensor_ops.d_tensor_set_global_tpu_array(topology_proto)
with ops.device("/job:" + config.full_job_name() + "/device:TPU_SYSTEM:0"): # pylint: disable=protected-access
my_core_ids = _tpu_init_fn()
if use_megacore:
logging.info("Using TPU megacore")
my_core_ids = my_core_ids * 2
logging.info("TPU core IDs: %s", my_core_ids)
# `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.
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, config.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 a valid input. To give an example, assume we have
# 2 tasks and each of them has 8 local devices, then `all_core_ids` in task 0
# will have 8 tensors, where 1 of them may have its value as
# [0,1,2,3,4,5,6,7,0,0,0,0,0,0,0,0] and the other tensors are all zeros. For
# task 1, the case may be one with [0,0,0,0,0,0,0,0,8,9,10,11,12,13,14,15]
# and other 7 are all zeros.
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)
# All devices on all hosts participate in one AllReduce, whose result will be
# core IDs arranged by task-device ordinals. For the above example, the result
# will be [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15].
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
)
_set_global_tpu_array_fn(tpu_topology.serialized())
return tpu_topology, device
def initialize_tpu_system(use_megacore=False):
"""Initializes the TPU system."""
# 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
use_tfrt_host_runtime = context.context().use_tfrt
logging.info("Using TFRT host runtime is set to %s", use_tfrt_host_runtime)
try:
task_id = config.client_id()
num_tasks = config.num_clients()
num_devices = config.num_global_devices(_TPU_DEVICE_TYPE)
tpu_topology, device = tpu_system_init_helper(
task_id,
num_tasks,
num_devices,
use_tfrt_host_runtime=use_tfrt_host_runtime,
use_megacore=use_megacore,
)
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)
) from e
except errors.InternalError as e:
logging.error(
"Hit internal error during TPU system initialization. "
+ "It is likely hardware 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
# 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
def _enumerate_cores(
bounds: List[int],
ring_bounds: List[int],
ring_sizes: List[int],
host_bounds: List[int],
host_sizes: List[int],
) -> List[List[int]]:
"""Enumerates cores within `bounds` from fatest to slowest varying axes.
Args:
bounds: Upper bounds of axes, from fastest to slowest varying.
ring_bounds: Upper bounds of ring size per axis in the same axis order.
ring_sizes: Number consecutive cores in the ring built so far, cumulatively.
host_bounds: Number of axis values per host in the same axis order.
host_sizes: Number consecutive cores on one host, cumulatively.
Returns:
Cores represented as a list of 4 integers in the same axis order.
"""
if not bounds:
return [[]]
# Recursively enumerate cores under all but the slowest varying axis.
partials = _enumerate_cores(
bounds[:-1],
ring_bounds[:-1],
ring_sizes[:-1],
host_bounds[:-1],
host_sizes[:-1],
)
# Append the slowest varying axis to the end of all partial results.
# From ring_i|j to host_i|j to core_i|j, use progressively smaller or equal
# iteration groupings until every one of the bounds[-1] * len(partials)
# combinations is iterated on.
# Despite the six levels of nested loops below, the total time complexity for
# this invocation is O(N), where N is the number of cores in the topology.
results = []
for ring_i in range(0, bounds[-1], ring_bounds[-1]):
for ring_j in range(0, len(partials), ring_sizes[-1]):
for host_i in range(ring_i, ring_i + ring_bounds[-1], host_bounds[-1]):
for host_j in range(ring_j, ring_j + ring_sizes[-1], host_sizes[-1]):
for i in range(host_i, host_i + host_bounds[-1]):
for j in range(host_j, host_j + host_sizes[-1]):
results.append(partials[j] + [i])
return results
def _enumerate_core_locations(
bounds: List[int],
ring_bounds: List[int],
axes: List[str],
can_split_host_across_rings: bool,
ring_size: int,
) -> List[_CoreLocation]:
"""Enumerates all possible core locations under the axis iteration order.
Args:
bounds: A list of 4 positive integers, upper bound values for x, y, z, core.
ring_bounds: A list of 4 positive integers, upper bound values for ring size
in x, y, z, core axes.
axes: A permutation of ["x", "y", "z", "core"], the axis iteration order.
can_split_host_across_rings: If true, devices attached to the same host may
get assigned to different rings.
ring_size: Number of devices in a ring, only for argument validation.
Returns:
A list of all CoreLocation objects defined in a TPU slice of shape `bounds`,
sorted by axis iteration order specified by `axes`.
For example, given bounds=[2, 2, 1, 2] and axes=["core", "z", "y", "x"],
return 8 core locations expressed in (x, y, z, core) format but iterated in
core -> z -> y -> x order (fatest to slowest varying):
[_CoreLocation(0, 0, 0, 0),
_CoreLocation(0, 0, 0, 1),
_CoreLocation(0, 1, 0, 0),
_CoreLocation(0, 1, 0, 1),
_CoreLocation(1, 0, 0, 0),
_CoreLocation(1, 0, 0, 1),
_CoreLocation(1, 1, 0, 0),
_CoreLocation(1, 1, 0, 1)]
Raises:
ValueError: If ring_size cannot be fulfilled without splitting hosts.
"""
num_cores_per_chip = bounds[3]
if num_cores_per_chip != 1 and num_cores_per_chip != 2:
raise ValueError("Unsupported TPU slice size: %s" % bounds)
# Translate `axes` from string to integer format.
axes = [{"x": 0, "y": 1, "z": 2, "core": 3}[axis] for axis in axes]
# Reorder bounds from fastest to slowest varying axes.
bounds = [bounds[i] for i in axes]
# Set and validate host_bounds.
if can_split_host_across_rings:
# If we can split hosts, shrink every host to effectively contain 1 device.
host_bounds = [1, 1, 1, 1]
elif np.prod(bounds) <= 2:
# We must be running on 1x1 or 1x1x1 Forge.
host_bounds = [[1, 1, 1, num_cores_per_chip][i] for i in axes]
else:
# Other cases including 2x2 Forge and Borg must use a full donut.
host_bounds = [[2, 2, 1, num_cores_per_chip][i] for i in axes]
# host_sizes is the cumulative products of host_bounts.
host_sizes = [1]
for host_bound in host_bounds:
host_sizes.append(host_sizes[-1] * host_bound)
host_size = host_sizes.pop()
# When can_split_host_across_rings is false, a ring must contain at least as
# many devices as a host has.
if ring_size < host_size:
assert not can_split_host_across_rings
raise ValueError(
"Rings too small for can_split_host_across_rings = False: %d"
% ring_size
)
# Reorder ring_bounds and validate it's element-wise >= host_bounds.
ring_bounds = [ring_bounds[i] for i in axes]
if ring_bounds < host_bounds:
raise ValueError(
"ring_bounds %s should be >= host_bounds %s"
% (ring_bounds, host_bounds)
)
ring_sizes = [1]
# ring_sizes is the cumulative products of ring_bounds.
for ring_bound in ring_bounds:
ring_sizes.append(ring_sizes[-1] * ring_bound)
ring_sizes.pop()
# Enumerate cores in the given iteration order. Each core is represented as a
# list of int, which are offsets from fatest to slowest varying axes.
cores = _enumerate_cores(
bounds, ring_bounds, ring_sizes, host_bounds, host_sizes
)
# Reorder offsets of each core back to the x, y, z, core order.
core_locations = []
for core in cores:
core = [core[axes.index(i)] for i in range(4)]
core_locations.append(_CoreLocation(core[0], core[1], core[2], core[3]))
return core_locations
def _build_all_reduce_ring(
core_locations: List[_CoreLocation], rotate: bool = False
) -> List[int]:
"""Reorders a list of TPU cores to optimize for AllReduce performance.
This is ported from the C++ tensorflow::BuildAllReduceRing function,
mixed with some logic from TF TPU's device_assignment._ring_3d.
Args:
core_locations: A list of core locations expressed as [x, y, z, core].
rotate: If true, scan the cores in a column-major order. False by default.
Returns:
A permutation of the input list such that neighbors in the sequence are
nearby in the TPU topology.
"""
permutation = list(range(len(core_locations)))
if not permutation:
return permutation
logging.vlog(2, "Core locations in: %s", core_locations)
first_column = min([l.x for l in core_locations])
first_row = min([l.y for l in core_locations])
same_z = len(set([l.z for l in core_locations])) == 1
logging.vlog(2, "first_column: %d", first_column)
logging.vlog(2, "first_row: %d", first_row)
logging.vlog(2, "same_z: %s", same_z)
def _cmp_2d(ia: int, ib: int) -> int:
if not rotate:
a = core_locations[ia]
b = core_locations[ib]
# Order the first column last in the sequence, except for the first row.
a_first = a.x == first_column and a.y != first_row
b_first = b.x == first_column and b.y != first_row
if a_first != b_first:
return -1 if b_first else 1
# Order rows in increasing order, unless in the first column.
if a.y != b.y:
return b.y - a.y if a_first else a.y - b.y
# Order even rows left to right, odd rows right to left.
if a.x != b.x:
return a.x - b.x if a.y % 2 == 0 else b.x - a.x
# Order cores in increasing order.
return a.core - b.core
else:
a = core_locations[ia]
b = core_locations[ib]
# Order the first row last in the sequence, except for the first column.
a_first = a.y == first_row and a.x != first_column
b_first = b.y == first_row and b.x != first_column
if a_first != b_first:
return -1 if b_first else 1
# Order columns in increasing order, unless in the first row.
if a.x != b.x:
return b.x - a.x if a_first else a.x - b.x
# Order even columns top down, odd columns bottom up.
if a.y != b.y:
return a.y - b.y if a.x % 2 == 0 else b.y - a.y
# Order cores in increasing order.
return a.core - b.core
def _cmp_3d(ia: int, ib: int) -> int:
a = core_locations[ia]
b = core_locations[ib]
a_corner = a.x == first_column and a.y == first_row
b_corner = b.x == first_column and b.y == first_row
# If both are in the corner, order in reverse z then core order.
if a_corner and b_corner:
return b.z - a.z if a.z != b.z else a.core - b.core
# Corner cores always go after non-corner cores.
if a_corner != b_corner:
return -1 if b_corner else 1
# Both non-corner cores are on the same z-plane. Reverse odd z-planes.
if a.z == b.z:
return _cmp_2d(ia, ib) if a.z % 2 == 0 else -_cmp_2d(ia, ib)
# Both non-corner cores are on different z-planes. Smaller z goes first.
return a.z - b.z
# If all cores are on the same z-plane, order as usual. Otherwise, order
# neighbor z-planes in opposite orders. Stack all z-planes along the z axis
# and connect them in one corner.
if same_z:
permutation.sort(key=functools.cmp_to_key(_cmp_2d))
else:
permutation.sort(key=functools.cmp_to_key(_cmp_3d))
logging.vlog(2, "Permutation out: %s", permutation)
return permutation
def _build_orthogonal_rings(
core_locations: List[_CoreLocation],
ring_size: int,
rotate_ring_across_rings: bool,
) -> List[_CoreLocation]:
"""Build two all-reduce rings orthogonal to each other.
One ring includes every `ring_size` consecutive core locations. It is usually
applied to the model-parallel dimension of a mesh to achieve best 1D
all-reduce performance. The other ring includes core locations separated by
a stride of `ring_size`. It is usually applied to the data-parallel dimension
of a mesh to get predictable strided all-reduce performance.
Args:
core_locations: A list of core locations expressed as [x, y, z, core].
ring_size: The number of core locations in the consecutive ring.
rotate_ring_across_rings: Build column-major secondary rings.
Returns:
A permutation of the input list forming the described rings.
"""
# Build a ring for the first `ring_size` cores, and apply that permutation to
# every group of `ring_size` cores.
num_cores = len(core_locations)
permutation = _build_all_reduce_ring(core_locations[:ring_size])
for r in range(0, num_cores, ring_size):
core_locations[r : r + ring_size] = [
core_locations[r + permutation[i]] for i in range(ring_size)
]
logging.vlog(1, "Permutated core locations: %s", core_locations)
# Build a "ring" for the collection of devices consisting of the 0th device
# from every group, and apply that permutation to every i-th device group.
# This is achieved by transposing the list and back.
transposed = []
for i in range(ring_size):
transposed += [
core_locations[g + i] for g in range(0, num_cores, ring_size)
]
num_rings = int(num_cores / ring_size)
permutation = _build_all_reduce_ring(
transposed[:num_rings], rotate=rotate_ring_across_rings
)
for r in range(0, num_cores, num_rings):
transposed[r : r + num_rings] = [
transposed[r + permutation[i]] for i in range(num_rings)
]
untransposed = []
for i in range(num_rings):
untransposed += [transposed[g + i] for g in range(0, num_cores, num_rings)]
logging.vlog(1, "Stride-permutated core locations: %s", untransposed)
return untransposed
@tf_export("experimental.dtensor.create_tpu_mesh", v1=[])
def create_tpu_mesh(
mesh_dim_names: List[str],
mesh_shape: List[int],
mesh_name: str,
ring_dims: Optional[int] = None,
ring_axes: Optional[List[str]] = None,
ring_bounds: Optional[List[int]] = None,
can_split_host_across_rings: bool = True,
build_ring_across_rings: bool = False,
rotate_ring_across_rings: bool = False,
use_xla_spmd: bool = layout_lib.USE_XLA_SPMD,
) -> layout_lib.Mesh:
"""Returns a distributed TPU mesh optimized for AllReduce ring reductions.
Only as many as leading axes specified by `ring_axes` as necessary will be
used to build rings, as long as the subslice formed by these axes have enough
cores to contain a ring of the required size. The leftover axes in `ring_axes`
won't affect results.
This function always uses all TPU devices, and offers more customization than
`tf.experimental.dtensor.create_distributed_mesh`.
Args:
mesh_dim_names: List of mesh dimension names.
mesh_shape: Shape of the mesh.
mesh_name: A unique name for the mesh. If empty, internally generate one.
ring_dims: Optional; The number of leading (ring_dims > 0) or trailing
(ring_dims < 0) mesh dimensions to build rings for. If unspecified, build
rings for all but the first dimension.
ring_axes: Optional; A permutation of ["x", "y", "z", "core"], specifying
the order of TPU topology axes to build rings in. If unspecified, default
to ["core", "x", "y", "z"].
ring_bounds: Optional; The maximum number of devices on each axis, in the x,
y, z, core order. If unspecified, default to physical topology limits.
can_split_host_across_rings: Optional; If true, devices attached to the same
host (i.e., DTensor client) may get assigned to different rings. Setting
it to false may cause some combinations of arguments to be infeasible; see
DeviceAssignmentTest.testCreateMesh[No]SplittingHosts* for examples.
build_ring_across_rings: Optional; If true, also build a data-parallel ring
across model-parallel rings. This ring could be strided.
rotate_ring_across_rings: Optional; If true, build the data-parallel ring in
column-major instead of row-major order.
use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.
"""
logging.info("Building a TPU mesh %s of shape %s", mesh_name, mesh_shape)
logging.info("Requested ring_dims: %s", ring_dims)
logging.info("Requested ring_axes: %s", ring_axes)
logging.info("Requested ring_bounds: %s", ring_bounds)
logging.info(
"Requested can_split_host_across_rings: %s", can_split_host_across_rings
)
if not mesh_name:
mesh_name = "mesh_%f" % time.time()
logging.info("Requested mesh_name: %s", mesh_name)
# By default, build rings for all but the first (usually batch) dimension.
if ring_dims is None:
ring_dims = 1 - len(mesh_shape)
elif ring_dims < -len(mesh_shape) or ring_dims > len(mesh_shape):
raise ValueError("Invalid ring_dims value: %d" % ring_dims)
logging.info("Actual ring_dims: %s", ring_dims)
# By default, vary axes in the core -> x -> y -> z order.
if ring_axes is None:
ring_axes = ["core", "x", "y", "z"]
elif len(ring_axes) != 4:
raise ValueError("Expected 4 elements in ring_axes, got %s" % ring_axes)
elif sorted(ring_axes) != ["core", "x", "y", "z"]:
raise ValueError("Invalid ring_axes value: %s" % ring_axes)
logging.info("Actual ring_axes: %s", ring_axes)
# Validate ring_bounds values.
if _tpu_topology is None:
raise ValueError(
"Invalid TPU topology, run dtensor.initialize_tpu_system() first"
)
topology_shape = list(_tpu_topology.mesh_shape)
if ring_bounds is None:
ring_bounds = topology_shape
elif len(ring_bounds) != 4:
raise ValueError("Expected 4 elements in ring_bounds, got %s" % ring_bounds)
elif ring_bounds > topology_shape:
raise ValueError(
"ring_bounds %s should be <= topology sizes %s"
% (ring_bounds, topology_shape)
)
logging.info("Actual ring_bounds: %s", ring_bounds)
# Compute ring_size, the number of cores in a ring.
if ring_dims > 0:
ring_size = np.prod(mesh_shape[:ring_dims])
elif ring_dims < 0:
ring_size = np.prod(mesh_shape[ring_dims:])
else:
ring_size = 1 # single-core rings
logging.info("Actual ring_size: %d", ring_size)
# Rearrange all cores according to the axis iteration order.
global_core_locations = _enumerate_core_locations(
topology_shape,
ring_bounds,
ring_axes,
can_split_host_across_rings,
ring_size,
)
logging.vlog(1, "Enumerated core locations: %s", global_core_locations)
num_cores = len(global_core_locations)
# The mesh to be created must use all TPU cores in the system.
mesh_size = np.prod(mesh_shape)
if mesh_size != num_cores:
raise ValueError(
"Invalid mesh size: mesh shape %s cannot 1:1 map to %d TPU cores"
% (mesh_shape, num_cores)
)
# Build a ring for the `ring_size` dimension and, if required, a strided ring
# for the orthogonal dimension.
if build_ring_across_rings:
global_core_locations = _build_orthogonal_rings(
global_core_locations, ring_size, rotate_ring_across_rings
)
else:
permutation = _build_all_reduce_ring(global_core_locations[:ring_size])
for r in range(0, num_cores, ring_size):
global_core_locations[r : r + ring_size] = [
global_core_locations[r + permutation[i]] for i in range(ring_size)
]
logging.vlog(1, "Permutated core locations: %s", global_core_locations)
# For this point on, change from List[CoreLocation] to List[List[int]] for
# easier interaction with the C++ API.
global_core_locations = [l.to_list() for l in global_core_locations]
if _dtensor_device is None:
raise ValueError(
"Invalid system device, "
"run dtensor.initialize_accelerator_system() first"
)
global_core_ids = _dtensor_device.tpu_core_locations_to_ids(
global_core_locations
)
# Store a per-mesh mapping in the runtime.
_dtensor_device.set_tpu_core_ids(mesh_name, global_core_ids)
# Create the mesh by manually specifying local_device_ids.
local_core_locations = _tpu_topology.device_coordinates[config.client_id()]
indexes = [
global_core_locations.index(list(local_core_location))
for local_core_location in local_core_locations
]
global_device_ids, local_device_ids, local_device_list = _create_device_array(
mesh_shape, _TPU_DEVICE_TYPE, None, local_device_ids=indexes
)
return layout_lib.Mesh(
mesh_dim_names,
global_device_ids,
local_device_ids,
local_device_list,
mesh_name,
use_xla_spmd=use_xla_spmd,
)
def get_device_ids(
mesh: layout_lib.Mesh, client_id: Optional[int] = None
) -> List[int]:
"""Returns the device IDs of all TPU cores local to the given client.
A device ID is a non-negative integer that uniquely identifies a device in the
mesh. For example, for a 2x2 mesh ('x', 'y'), this function returns a
permutation of [0, 1, 2, 3].
Note that device IDs and device locations are equivalent. The former is a
linearization of the latter along mesh dimensions.
Args:
mesh: A TPU mesh.
client_id: Optional; A DTensor client ID. If empty, query this client.
"""
if mesh.device_type() != _TPU_DEVICE_TYPE:
raise ValueError("The mesh must be a TPU mesh")
if client_id is None or client_id == config.client_id():
return mesh.local_device_ids()
# It's not clear we should ever allow a client to query other clients for
# their device IDs.
raise NotImplementedError(
"Looking up other clients' device IDs is not supported"
)
def get_device_locations(
mesh: layout_lib.Mesh, client_id: Optional[int] = None
) -> List[Dict[str, int]]:
"""Returns the device locations of all TPU cores local to the given client.
A device location is a dictionary from dimension names to indices on those
dimensions. For example, for a 2x2 mesh ('x', 'y'), this function returns a
permutation of this list:
[{'x': 0, 'y': 0},
{'x': 0, 'y': 1},
{'x': 1, 'y': 0},
{'x': 1, 'y': 1}].
Note that device IDs and device locations are equivalent. The former is a
linearization of the latter along mesh dimensions.
Args:
mesh: A TPU mesh.
client_id: Optional; A DTensor client ID. If empty, query this client.
"""
if mesh.device_type() != _TPU_DEVICE_TYPE:
raise ValueError("The mesh must be a TPU mesh")
if client_id is None or client_id == config.client_id():
return mesh.local_device_locations()
# It's not clear we should ever allow a client to query other clients for
# their device locations.
raise NotImplementedError(
"Looking up other clients' device locations is not supported"
)
# TODO(b/245589661): Remove dtensor_initialize_tpu_system() and
# dtensor_shutdown_tpu_system() after users stopped using them.
def dtensor_initialize_tpu_system(enable_coordination_service=False):
"""Deprecated way to initialize the TPU system."""
from . import accelerator_util # pylint: disable=g-import-not-at-top
accelerator_util.initialize_accelerator_system(
"TPU", enable_coordination_service=enable_coordination_service
)
def dtensor_shutdown_tpu_system():
"""Deprecated way to shutodwn the TPU system."""
from . import accelerator_util # pylint: disable=g-import-not-at-top
accelerator_util.shutdown_accelerator_system()