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/python/checkpoint/sharding/BUILD
+104
View File
@@ -0,0 +1,104 @@
# Description:
# Utilities for sharding object-based checkpoints.
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
pytype_strict_library(
name = "sharding_policies",
srcs = ["sharding_policies.py"],
deps = [
":sharding_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:base",
"//tensorflow/python/util:tf_export",
"@absl_py//absl/logging",
],
)
tf_py_strict_test(
name = "sharding_policies_test",
srcs = ["sharding_policies_test.py"],
tags = [
"notap",
"oss_excluded",
], # See b/455621783.
deps = [
":sharding_policies",
":sharding_util",
"//tensorflow/python/checkpoint",
"//tensorflow/python/checkpoint:checkpoint_options",
"//tensorflow/python/checkpoint:graph_view",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:remote",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/module",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
"@absl_py//absl/logging",
],
)
pytype_strict_library(
name = "sharding_util",
srcs = ["sharding_util.py"],
deps = [
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "sharding_util_test",
srcs = ["sharding_util_test.py"],
deps = [
":sharding_policies",
":sharding_util",
"//tensorflow/python/checkpoint",
"//tensorflow/python/checkpoint:graph_view",
"//tensorflow/python/eager:remote",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
],
)
tensorflow/python/checkpoint/sharding/sharding_policies.py
+353
View File
@@ -0,0 +1,353 @@
# 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.
# ==============================================================================
"""Checkpoint policies that determine how tensors are split into shards."""
import math
import operator
from typing import MutableSequence, Sequence
from absl import logging
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.eager import context
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variables
from tensorflow.python.trackable import base
from tensorflow.python.util import tf_export
@tf_export.tf_export("train.experimental.ShardByTaskPolicy")
class ShardByTaskPolicy(sharding_util.ShardingCallback):
"""Policy that splits tensors into shards based on their device spec task."""
@property
def description(self) -> str:
return "Split tensors into shards based on their device spec task."
def __call__(
self,
shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
"""Callback to split tensors into shards based on their device spec task.
Args:
shardable_tensors: A list of ShardableTensors.
Returns:
List of shard dicts containing tensors.
[ {checkpoint key: {slice_spec: tensor} } ]
"""
tensors_by_task = {}
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
(tensors_by_task
.setdefault(checkpoint_key, {})[slice_spec]) = tensor
return [tensors_by_task]
_OffsetAndShape = tuple[Sequence[int], Sequence[int]]
@tf_export.tf_export("train.experimental.MaxShardSizePolicy")
class MaxShardSizePolicy(sharding_util.ShardingCallback):
"""Policy that splits tensors into shards with a max shard size.
Shards may exceed the max shard size if they contain 1. a single scalar/string
tensor that could not be sliced and exceeds the max shard size or 2. the
checkpoint object graph, whose size cannot be calculated when saving.
"""
class MaxShardSizePartitioner():
"""Partition tensors into shards with a max shard size."""
max_shard_size: int
_large_scalars: MutableSequence[sharding_util.Shard]
_tensors_by_shard: MutableSequence[sharding_util.Shard]
_shard_size_remaining: int
_checkpoint_key: str
_dtype: dtypes.DType
_device: device_lib.DeviceSpec
_root_tensor: tensor_lib.Tensor
_slice_spec: variables.Variable.SaveSliceInfo
_full_shape: tensor_shape.TensorShape
_root_shape: tensor_shape.TensorShape
_root_offset: Sequence[int]
_dtype_size: int
_working_tensor_offset: MutableSequence[float]
_working_tensor_shape: tensor_shape.TensorShape
def _get_next_partition(self) -> tuple[int, float]:
"""Gets tensor partition with size closest to shard_size_remaining.
Returns:
A tuple containing the axis and size of the next partition.
"""
rank = self._working_tensor_shape.rank
if rank is None or rank == 0:
return 0, math.inf
num_elems = self._working_tensor_shape.num_elements()
def num_partitions(axis: int) -> float:
axis_len = self._working_tensor_shape.dims[axis].value
slice_elems = num_elems // axis_len
bytes_per_slice = slice_elems * self._dtype_size
slices_per_shard = self._shard_size_remaining // bytes_per_slice
if slices_per_shard == 0:
return math.inf
return math.ceil(axis_len / slices_per_shard)
# Find axis with minimum partitions. (axis with maximum partition size)
# (max partition size is as close as possible to the shard_size_remaining)
min_parts = num_partitions(0)
min_axis = 0
for axis in range(1, rank):
parts_along_axis = num_partitions(axis)
part_size = num_elems * self._dtype_size / parts_along_axis
if (parts_along_axis < min_parts and
part_size <= self._shard_size_remaining):
min_axis, min_parts = axis, int(parts_along_axis)
return (min_axis,
math.ceil(int(self._working_tensor_shape[min_axis]) / min_parts))
def _add_partition(self, part_axis: int, part_size: float):
"""Adds the tensor partition to the shard, if possible.
Args:
part_axis: The axis of the partition.
part_size: The size of the partition.
Raises:
RuntimeError: When the slice size is larger than the remaining shard
size.
"""
# Add what we can to the current shard.
relative_offset = list(
map(operator.sub, self._working_tensor_offset, self._root_offset))
slice_shape = list(map(operator.sub, self._root_shape, relative_offset))
slice_shape[part_axis] = part_size
slice_size_in_bytes = int(math.prod(slice_shape)) * self._dtype_size
with ops.device(self._device):
tensor_slice = array_ops.slice(
self._root_tensor, begin=relative_offset, size=slice_shape)
slice_spec = variables.Variable.SaveSliceInfo(
full_name=self._checkpoint_key,
full_shape=self._full_shape,
var_offset=self._working_tensor_offset,
var_shape=slice_shape).spec.strip()
if slice_size_in_bytes > self.max_shard_size:
logging.warning("Tensor %s's minimum slice %s has size %s bytes and "
"cannot be partitioned into a shard of max shard size "
"%s bytes. It will be added as an individual shard "
"that exceeds the max shard size.",
self._checkpoint_key, slice_spec, slice_size_in_bytes,
self.max_shard_size)
self._large_scalars.append(
{self._checkpoint_key: {slice_spec: tensor_slice}})
elif slice_size_in_bytes > self._shard_size_remaining:
raise RuntimeError(
f"Slice size ({slice_size_in_bytes} bytes) is larger than the "
f"remaining shard size ({self._shard_size_remaining} bytes). This "
"should have been caught in MaxShardSizePolicy._add_partition().")
else:
(self._tensors_by_shard[-1]
.setdefault(self._checkpoint_key, {})[slice_spec]) = tensor_slice
self._shard_size_remaining -= slice_size_in_bytes
if self._shard_size_remaining == 0:
self._tensors_by_shard.append({})
self._shard_size_remaining = self.max_shard_size
# Get remaining portion of tensor to add to the next shard(s).
self._working_tensor_offset[part_axis] += part_size
relative_offset[part_axis] += part_size
self._working_tensor_shape = tensor_shape.TensorShape(list(
map(operator.sub, self._root_shape, relative_offset)))
def get_shards(
self,
max_shard_size: int,
shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
"""Callback to split tensors into shards with a max shard size.
Args:
max_shard_size: The maximum size of a shard file in bytes.
shardable_tensors: A list of ShardableTensors.
Returns:
List of shard dicts containing tensors.
[ {checkpoint key: {slice_spec: tensor} } ]
"""
self.max_shard_size = max_shard_size
self._tensors_by_shard = [{}]
self._large_scalars = []
string_size_warning_printed = False
self._shard_size_remaining = self.max_shard_size
for shardable_tensor in shardable_tensors:
self._checkpoint_key = shardable_tensor.checkpoint_key
self._dtype = shardable_tensor.dtype
self._device = shardable_tensor.device
self._root_tensor = shardable_tensor.tensor
self._slice_spec = shardable_tensor.slice_spec
# If the tensor has already been sliced, make sure to keep track of its
# parent tensor's shape & offset. These will be used when creating slice
# specs later.
if self._slice_spec:
save_slice_info = variables.Variable.SaveSliceInfo.from_spec(
self._slice_spec)
self._full_shape = tensor_shape.TensorShape(
save_slice_info.full_shape)
self._root_shape = tensor_shape.TensorShape(save_slice_info.var_shape)
self._root_offset = save_slice_info.var_offset
else:
self._full_shape = self._root_shape = shardable_tensor.shape
self._root_offset = [0] * self._root_shape.rank
self._dtype_size = dtypes.as_dtype(self._dtype).size
total_size = self._root_shape.num_elements() * self._dtype_size # bytes
# Calculate string tensor sizes.
if self._checkpoint_key == base.OBJECT_GRAPH_PROTO_KEY:
# In graph mode, the object graph is populated using feed_additions
# when the session is run. So, we can't calculate the size here.
# Fortunately, the serialized object graph string will never be that
# big, so we just place it in the current shard without worrying about
# its size.
total_size = self._dtype_size = 0
elif self._dtype == dtypes.variant:
# Can't determine a variant's type, so can't calculate its size.
total_size = self._dtype_size = 0
elif (self._dtype == dtypes.string
and not context.executing_eagerly()
and ops.get_default_session() is None):
# TODO(b/326287351): Get string tensor size in tf.function.
total_size = self._dtype_size = 0
if not string_size_warning_printed:
logging.warning("The checkpoint sharding policy is being executed "
"in a tf.function. The size of the string/variant "
"constant cannot be obtained.")
string_size_warning_printed = True
elif self._dtype == dtypes.string:
with ops.device(self._device):
if not context.executing_eagerly():
self._root_tensor = ops.get_default_session().run(
self._root_tensor)
if self._root_shape.rank is None or self._root_shape.rank == 0:
sizes = [string_ops.string_length(self._root_tensor,
unit="BYTE")]
else:
sizes = [string_ops.string_length(elem, unit="BYTE")
for elem in self._root_tensor]
if context.executing_eagerly():
sizes = [size.numpy() for size in sizes]
else:
sizes = ops.get_default_session().run(sizes)
total_size = sum(sizes)
self._dtype_size = max(sizes)
if self._root_shape.rank is None or self._root_shape.rank == 0:
if total_size > self.max_shard_size:
logging.warning(
"Tensor %s is a %s scalar of size %s bytes and cannot be "
"partitioned into a shard of max shard size %s bytes. It will "
"be added as an individual shard that exceeds the max shard "
"size.", self._checkpoint_key, self._dtype, total_size,
self.max_shard_size)
self._large_scalars.append(
{self._checkpoint_key: {self._slice_spec: self._root_tensor}})
else:
if total_size > self._shard_size_remaining:
self._tensors_by_shard.append({})
self._shard_size_remaining = self.max_shard_size
(self._tensors_by_shard[-1]
.setdefault(self._checkpoint_key, {})
[self._slice_spec]) = self._root_tensor
self._shard_size_remaining -= total_size
continue
# Partition tensor and add partitions to shards.
self._working_tensor_offset = self._root_offset[:]
self._working_tensor_shape = self._root_shape
working_tensor_size = total_size
while working_tensor_size > self._shard_size_remaining:
(part_axis, part_size) = self._get_next_partition()
if part_size == 0:
# Tensor partition couldn't fit in remaining shard space. Try again
# with the next full shard.
self._tensors_by_shard.append({})
self._shard_size_remaining = self.max_shard_size
else:
self._add_partition(part_axis, part_size)
working_tensor_size = (
int(math.prod(self._working_tensor_shape)) * self._dtype_size)
if self._working_tensor_shape.num_elements() > 0:
if self._working_tensor_offset and self._working_tensor_shape:
with ops.device(self._device):
working_tensor = array_ops.slice(
self._root_tensor,
begin=list(map(
operator.sub,
self._working_tensor_offset, self._root_offset)),
size=self._working_tensor_shape.as_list())
else:
working_tensor = self._root_tensor
remaining_tensor_slice_spec = variables.Variable.SaveSliceInfo(
full_name=self._checkpoint_key,
full_shape=self._full_shape,
var_offset=self._working_tensor_offset,
var_shape=self._working_tensor_shape).spec.strip()
(self._tensors_by_shard[-1]
.setdefault(self._checkpoint_key, {})
[remaining_tensor_slice_spec]) = working_tensor
self._shard_size_remaining -= working_tensor_size
shards = []
if self._tensors_by_shard[0]:
shards.extend(self._tensors_by_shard)
shards.extend(self._large_scalars)
return shards
def __init__(self, max_shard_size: int):
self.max_shard_size = max_shard_size
@property
def description(self) -> str:
return "Split tensors into shards with a max shard size."
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
return self.MaxShardSizePartitioner().get_shards(
self.max_shard_size, shardable_tensors)
tensorflow/python/checkpoint/sharding/sharding_policies_test.py
+810
View File
@@ -0,0 +1,810 @@
# 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 checkpoint sharding policies."""
import random
import re
import string
from absl import logging
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import checkpoint_options
from tensorflow.python.checkpoint import graph_view
from tensorflow.python.checkpoint.sharding import sharding_policies
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.eager import def_function
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.module import module
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.training import server_lib
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.training.saving import saveable_object_util
class ShardingPoliciesTest(test.TestCase):
def _get_shardable_tensors_by_task(self, root):
serialized_tensors, _, _, _ = (
checkpoint.TrackableSaver(graph_view.ObjectGraphView(root))
._gather_serialized_tensors(None))
shardable_tensors_by_task = {}
for obj, tensor_dict in serialized_tensors.items():
# Divide tensor_dict by device.
for checkpoint_key, tensor_slice_dict in tensor_dict.items():
if not isinstance(tensor_slice_dict, dict):
# Make sure that maybe_tensor is structured as {slice_spec -> tensor}.
tensor_slice_dict = {"": tensor_slice_dict}
for slice_spec, tensor_save_spec in tensor_slice_dict.items():
if not isinstance(tensor_save_spec, saveable_object.SaveSpec):
tensor_save_spec = saveable_object.SaveSpec(
tensor=tensor_save_spec,
slice_spec=slice_spec,
name=checkpoint_key,
dtype=tensor_save_spec.dtype,
device=tensor_save_spec.device)
save_spec_tensor = tensor_save_spec.tensor
device = (device_lib.DeviceSpec.from_string(tensor_save_spec.device)
if isinstance(tensor_save_spec.device, str)
else tensor_save_spec.device)
task = device_lib.DeviceSpec.from_string(
saveable_object_util.set_cpu0(device.to_string()))
shardable_tensors_by_task.setdefault(task, []).append(
sharding_util.ShardableTensor(
_tensor_save_spec=tensor_save_spec,
tensor=save_spec_tensor,
dtype=tensor_save_spec.dtype,
device=device,
name=tensor_save_spec.name,
shape=save_spec_tensor.shape,
slice_spec=slice_spec,
checkpoint_key=checkpoint_key,
trackable=obj))
return shardable_tensors_by_task.values()
def test_ShardByTaskPolicy(self):
servers = [server_lib.Server.create_local_server() for _ in range(3)]
cluster_spec = server_lib.ClusterSpec({
"worker": [s.target[len("grpc://"):] for s in servers]})
remote.connect_to_cluster(cluster_spec)
root = module.Module()
with ops.device("/job:worker/task:0/cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
with ops.device("/job:worker/task:2/cpu:0"):
v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
root.v0 = v0
root.v1 = v1
root.v2 = v2
shardable_tensors = self._get_shardable_tensors_by_task(root)
callback = sharding_policies.ShardByTaskPolicy()
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertAllEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE"},
{"v1/.ATTRIBUTES/VARIABLE_VALUE"},
{"v2/.ATTRIBUTES/VARIABLE_VALUE"},
{"_CHECKPOINTABLE_OBJECT_GRAPH"}
])
self.assertEqual(
self.evaluate(shards[0]["v0/.ATTRIBUTES/VARIABLE_VALUE"][""]),
v0.numpy())
self.assertEqual(
self.evaluate(shards[1]["v1/.ATTRIBUTES/VARIABLE_VALUE"][""]),
v1.numpy())
self.assertEqual(
self.evaluate(shards[2]["v2/.ATTRIBUTES/VARIABLE_VALUE"][""]),
v2.numpy())
def test_CheckpointOption_ShardByTaskPolicy(self):
servers = [server_lib.Server.create_local_server() for _ in range(3)]
cluster_spec = server_lib.ClusterSpec({
"worker": [s.target[len("grpc://"):] for s in servers]})
remote.connect_to_cluster(cluster_spec)
root = module.Module()
with ops.device("/job:worker/task:0/cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
self.evaluate(v0.initializer)
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
self.evaluate(v1.initializer)
with ops.device("/job:worker/task:2/cpu:0"):
v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
self.evaluate(v2.initializer)
root.v0 = v0
root.v1 = v1
root.v2 = v2
tmp_dir = self.create_tempdir("ckpt")
ckpt = checkpoint.Checkpoint(root)
save_path = ckpt.save(
tmp_dir, options=checkpoint_options.CheckpointOptions(
experimental_sharding_callback=(
sharding_policies.ShardByTaskPolicy())))
self.assertLen(gfile.Glob(save_path + ".data*"), 4)
ckpt.restore(save_path)
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_1D(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([0.0, 1.0, 2.0, 3.0],
name="v0",
dtype=dtypes.float32)
v1 = resource_variable_ops.ResourceVariable([[4],
[5],
[6],
[7]],
name="v1",
dtype=dtypes.int32)
self.evaluate(v0.initializer)
self.evaluate(v1.initializer)
root.v0 = v0
root.v1 = v1
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
v1_name = "v1/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v0_name,
full_shape=tensor_shape.TensorShape(dims=[4]),
var_offset=var_offset,
var_shape=var_shape)
class V1SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v1_name,
full_shape=tensor_shape.TensorShape(dims=[4, 1]),
var_offset=var_offset,
var_shape=var_shape)
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the v0 & v1 tensors with different max shard sizes.
# max_shard_size: 4 bytes
# Each element of v0/v1 is a 32 bit/4 byte value, so each variable should be
# split into 4 shards.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=4)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[0][v0_name][slice_spec]), 0.0)
slice_spec = V0SaveSliceInfo(var_offset=[1], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[1][v0_name][slice_spec]), 1.0)
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[2][v0_name][slice_spec]), 2.0)
slice_spec = V0SaveSliceInfo(var_offset=[3], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[3][v0_name][slice_spec]), 3.0)
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[4][v1_name][slice_spec]), [4])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[5][v1_name][slice_spec]), [5])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[6][v1_name][slice_spec]), [6])
slice_spec = V1SaveSliceInfo(var_offset=[3, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[7][v1_name][slice_spec]), [7])
# max_shard_size: 8 bytes
# v0/v1 haven't changed, so they should now be split into 2 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0])
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]), [[4], [5]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[6], [7]])
# max_shard_size: 10 bytes
# 10 bytes is an uneven boundary for 4 byte elements. v0/v1 should be split
# into 2 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=10)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0])
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]), [[4], [5]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[6], [7]])
# max_shard_size: 16 bytes
# 16 bytes the exact size of each variable, so they should get 1 shard each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=16)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[4]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0, 2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[4, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v1_name][slice_spec]), [[4], [5], [6], [7]])
# max_shard_size: 18 bytes
# 18 bytes slightly larger than the size of each variable, but not large
# enough to fit another 4 byte element, so they should get 1 shard each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=18)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[4]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0, 2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[4, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v1_name][slice_spec]), [[4], [5], [6], [7]])
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_2D(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([[0, 1],
[2, 3],
[4, 5]],
name="v0")
v1 = resource_variable_ops.ResourceVariable([[[6.0], [7.0]],
[[8.0], [9.0]],
[[10.0], [11.0]]], name="v1")
self.evaluate(v0.initializer)
self.evaluate(v1.initializer)
root.v0 = v0
root.v1 = v1
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
v1_name = "v1/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v0_name,
full_shape=tensor_shape.TensorShape(dims=[3, 2]),
var_offset=var_offset,
var_shape=var_shape)
class V1SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v1_name,
full_shape=tensor_shape.TensorShape(dims=[3, 2, 1]),
var_offset=var_offset,
var_shape=var_shape)
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the v0 & v1 tensors with different max shard sizes.
# max_shard_size: 8 bytes
# Each element of v0/v1 is a 32 bit/4 byte value, so each variable should be
# split into 3 shards.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[2, 3]])
slice_spec = V0SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[2][v0_name][slice_spec]), [[4, 5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[[6.0], [7.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[4][v1_name][slice_spec]), [[[8.0], [9.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[5][v1_name][slice_spec]), [[[10.0], [11.0]]])
# max_shard_size: 10 bytes
# 10 bytes is an uneven boundary for 4 byte elements. v0/v1 should be split
# into 3 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=10)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[2, 3]])
slice_spec = V0SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[2][v0_name][slice_spec]), [[4, 5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[[6.0], [7.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[4][v1_name][slice_spec]), [[[8.0], [9.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[5][v1_name][slice_spec]), [[[10.0], [11.0]]])
# max_shard_size: 12 bytes
# 12 bytes is enough to fit 3 elements per variable in each shard.
# v0/v1 should be split into 2 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=12)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[3, 1]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0], [2], [4]])
slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[3, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[1], [3], [5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[3, 1, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]),
[[[6.0]], [[8.0]], [[10.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[0, 1, 0], var_shape=[3, 1, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]),
[[[7.0]], [[9.0]], [[11.0]]])
# max_shard_size: 16 bytes
# Each variable should be split into 1.5 shards. The middle shard will
# contain elements from both variables.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=16)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE", "v1/.ATTRIBUTES/VARIABLE_VALUE"},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[2, 2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1], [2, 3]])
slice_spec = V0SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[4, 5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v1_name][slice_spec]), [[[6.0], [7.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0, 0], var_shape=[2, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]),
[[[8.0], [9.0]], [[10.0], [11.0]]])
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_Strings(self):
v_strings = [
"".join(random.choices(string.ascii_uppercase + string.digits, k=10))
for _ in range(4)]
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(v_strings, name="v0",
dtype=dtypes.string)
self.evaluate(v0.initializer)
root.v0 = v0
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v0_name,
full_shape=tensor_shape.TensorShape(dims=[4]),
var_offset=var_offset,
var_shape=var_shape)
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the v0 & v1 tensors with different max shard sizes.
# max_shard_size: 10 bytes
# Each string in v0 is 10 bytes, so there should be 1 string per shard.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=10)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [v_strings[0]])
slice_spec = V0SaveSliceInfo(var_offset=[1], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [v_strings[1]])
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v0_name][slice_spec]), [v_strings[2]])
slice_spec = V0SaveSliceInfo(var_offset=[3], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v0_name][slice_spec]), [v_strings[3]])
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_LargeScalar(self):
v_string = "".join(random.choices(
string.ascii_uppercase + string.digits, k=10)).encode("utf-8")
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(
v_string, name="v0", dtype=dtypes.string)
self.evaluate(v0.initializer)
root.v0 = v0
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
shardable_tensors = self._get_shardable_tensors_by_task(root)
# max_shard_size: 8 bytes
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"_CHECKPOINTABLE_OBJECT_GRAPH",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",}
])
tensor_val = (self.evaluate(shards[1][v0_name][""])
if ops.context.executing_eagerly()
else shards[1][v0_name][""])
self.assertEqual(tensor_val, v_string)
@test_util.run_in_graph_and_eager_modes
def test_CheckpointOption_MaxShardSizePolicy(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([[0, 1],
[2, 3],
[4, 5]],
name="v0")
v1 = resource_variable_ops.ResourceVariable([[[6.0], [7.0]],
[[8.0], [9.0]],
[[10.0], [11.0]]], name="v1")
v2 = resource_variable_ops.ResourceVariable("test_string", name="v1")
self.evaluate(v0.initializer)
self.evaluate(v1.initializer)
self.evaluate(v2.initializer)
root.v0 = v0
root.v1 = v1
root.v2 = v2
tmp_dir = self.create_tempdir("ckpt")
ckpt = checkpoint.Checkpoint(root)
save_path = ckpt.save(
tmp_dir, options=checkpoint_options.CheckpointOptions(
experimental_sharding_callback=(
sharding_policies.MaxShardSizePolicy(max_shard_size=10))))
# 8 files = 3 shards for v0, 3 for v1, 1 for v2, and 1 for the object graph
self.assertLen(gfile.Glob(save_path + ".data*"), 8)
ckpt.restore(save_path)
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_PreSlicedTensor(self):
root = module.Module()
sliced_v0_name = "sliced_v0/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=sliced_v0_name,
full_shape=tensor_shape.TensorShape(dims=[2, 5]),
var_offset=var_offset,
var_shape=var_shape)
v0_slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[2, 3])
class ResourceVariableWithSliceSpec(resource_variable_ops.ResourceVariable):
def _serialize_to_tensors(self):
ckpt_key, tensor = list(super()._serialize_to_tensors().items())[0]
return {ckpt_key: {v0_slice_spec.spec: tensor}}
with ops.device("cpu:0"):
# full_v0 = [[0.0, 1.0, 2.0, 3.0, 4.0],
# [5.0, 6.0, 7.0, 8.0, 9.0]]
sliced_v0 = ResourceVariableWithSliceSpec([[1.0, 2.0, 3.0],
[6.0, 7.0, 8.0]],
name="sliced_v0",
dtype=dtypes.float32)
sliced_v0._set_save_slice_info(v0_slice_spec)
self.evaluate(sliced_v0.initializer)
root.sliced_v0 = sliced_v0
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the pre-sliced v0 tensor with different max shard sizes.
# max_shard_size: 8 bytes
# Each element of v0 is a 32 bit/4 byte value, so v0 should be split into 3
# shards containing 2 elements each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",
"_CHECKPOINTABLE_OBJECT_GRAPH",},
])
slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[0][sliced_v0_name][slice_spec]), [[1.0], [6.0]])
slice_spec = V0SaveSliceInfo(var_offset=[0, 2], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][sliced_v0_name][slice_spec]), [[2.0, 3.0]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 2], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[2][sliced_v0_name][slice_spec]), [[7.0, 8.0]])
# max_shard_size: 12 bytes
# Each element of v0 is a 32 bit/4 byte value, so v0 should be split into 2
# shards containing 3 elements each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=12)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",
"_CHECKPOINTABLE_OBJECT_GRAPH",},
])
slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[1, 3]).spec
self.assertAllEqual(
self.evaluate(shards[0][sliced_v0_name][slice_spec]), [[1.0, 2.0, 3.0]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 1], var_shape=[1, 3]).spec
self.assertAllEqual(
self.evaluate(shards[1][sliced_v0_name][slice_spec]), [[6.0, 7.0, 8.0]])
def test_MaxShardSizePolicy_TFFunction(self):
v_string = "".join(random.choices(
string.ascii_uppercase + string.digits, k=10)).encode("utf-8")
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(
v_string, name="v0", dtype=dtypes.string)
self.evaluate(v0.initializer)
root.v0 = v0
shardable_tensors = self._get_shardable_tensors_by_task(root)
@def_function.function
def wrapped_policy(shardable_tensors):
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=4)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
return shards
# TODO(b/326287351): Get string tensor size in tf.function.
# This test case should be changed when the bug is fixed/warning removed.
with self.assertLogs(level="WARNING") as log_output:
log_level = logging.get_verbosity()
logging.set_verbosity(logging.WARNING)
try:
wrapped_policy(shardable_tensors)
finally:
logging.set_verbosity(log_level)
output = log_output[0][0].message
self.assertTrue(
re.search("sharding policy is being executed in a tf.function", output))
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()
tensorflow/python/checkpoint/sharding/sharding_util.py
+288
View File
@@ -0,0 +1,288 @@
# 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.
# ==============================================================================
"""Data structures and utilities for checkpoint sharding."""
import abc
import dataclasses
import inspect
from typing import Hashable, MutableMapping, Sequence
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import variables
from tensorflow.python.trackable import base
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.util import tf_export
TensorSlices = MutableMapping[tensor_spec.TensorSpec, tensor_lib.Tensor]
# A mapping from a checkpoint key (full tensor name) to the corresponding tensor
# slices of the full tensor. It represents the collection of tensors stored in a
# checkpoint shard data file.
Shard = MutableMapping[str, TensorSlices]
@tf_export.tf_export("train.experimental.ShardableTensor")
@dataclasses.dataclass(frozen=True)
class ShardableTensor:
"""Tensor wrapper containing data necessary for sharding.
The tensor representation used as inputs to pre-made and custom
`tf.train.experiemental.ShardingCallback`s, which can be specified using the
`experimental_sharding_callback` option in `tf.train.CheckpointOptions`.
"""
_tensor_save_spec: saveable_object.SaveSpec
tensor: tensor_lib.Tensor
dtype: dtypes.DType
device: device_lib.DeviceSpec
name: str
shape: tensor_shape.TensorShape
slice_spec: variables.Variable.SaveSliceInfo
checkpoint_key: str
trackable: base.Trackable
def __hash__(self) -> int:
return hash((self.name, self.dtype, str(self.device), self.checkpoint_key))
def __repr__(self) -> str:
return (f"\n{self.__class__.__name__}:\n"
f" _tensor_save_spec={self._tensor_save_spec!r}\n"
f" tensor={self.tensor!r}\n"
f" dtype={self.dtype!r}\n"
f" device={self.device!r}\n"
f" name={self.name!r}\n"
f" shape={self.shape!r}\n"
f" slice_spec={self.slice_spec!r}\n"
f" checkpoint_key={self.checkpoint_key!r}\n"
f" trackable={self.trackable!r}")
@tf_export.tf_export("train.experimental.ShardingCallback")
class ShardingCallback(abc.ABC):
"""Checkpoint sharding callback function, along with a text description.
A callback function wrapper that will be executed to determine how tensors
will be split into shards when the saver writes the checkpoint shards to disk.
The callback takes a list of `tf.train.experimental.ShardableTensor`s as input
(as well as any kwargs defined by the `tf.train.experimental.ShardingCallback`
subclass), and organizes the input tensors into different shards. Tensors are
first organized by device task (see `tf.DeviceSpec`), then the callback will
be called for each collection of tensors.
There are a few restrictions to keep in mind when creating a custom callback:
- Tensors must not be removed from the checkpoint.
- Tensors must not be reshaped.
- Tensor dtypes must not change.
- Tensors within a shard must belong to the same task.
Validation checks will be performed after the callback function is executed to
ensure these restrictions aren't violated.
Here's an example of a simple custom callback:
```
# Place all tensors in a single shard.
class AllInOnePolicy(tf.train.experimental.ShardingCallback):
@property
def description(self):
return "Place all tensors in a single shard."
def __call__(self, shardable_tensors):
tensors = {}
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor_save_spec.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
tensors.set_default(checkpoint_key, {})[slice_spec] = tensor
return [tensors]
ckpt.save(
"path",
options=tf.train.CheckpointOptions(
experimental_sharding_callback=AllInOnePolicy()))
```
The `description` attribute is used to identify the callback and to aid
debugging during saving and restoration.
To take in kwargs, simply define the constructor and pass them in:
```
class ParameterPolicy(tf.train.experimental.ShardingCallback):
def __init__(self, custom_param):
self.custom_param = custom_param
...
ckpt.save(
"path",
options=tf.train.CheckpointOptions(
experimental_sharding_callback=ParameterPolicy(custom_param=...)))
```
"""
@property
@abc.abstractmethod
def description(self) -> str:
"""Returns a text description of the sharding policy."""
pass
@abc.abstractmethod
def __call__(
self, shardable_tensors: Sequence[ShardableTensor]
) -> Sequence[Shard]:
"""Returns a list of shards for the given shardable tensors."""
pass
def __hash__(self) -> int:
hash_val = hash(self.description)
# vars() only includes user-defined attributes.
for attr_name, attr_val in vars(self).items():
if not (inspect.ismethod(attr_val) or inspect.isfunction(attr_val)):
hash_val ^= hash(attr_name)
if isinstance(attr_val, Hashable):
hash_val ^= hash(attr_val)
return hash_val
def validate_shards(
shards: Sequence[Shard],
shardable_tensors: Sequence[ShardableTensor],
callback_description: str
) -> None:
"""Validates shards generated by the sharding_callback."""
unseen_tensor_dict = {}
for shardable_tensor in shardable_tensors:
unseen_tensor_dict.setdefault(
shardable_tensor.checkpoint_key, {}
)[shardable_tensor.slice_spec] = shardable_tensor.tensor
seen_tensor_set = set()
for shard_tensors in shards:
task_tensor = None
for checkpoint_key, tensor_slice_dict in shard_tensors.items():
for slice_spec, shard_tensor in tensor_slice_dict.items():
slice_spec = slice_spec.strip()
# Validate uniqueness.
if (checkpoint_key, slice_spec) in seen_tensor_set:
raise RuntimeError(
"After executing the checkpoint sharding callback, multiple "
"tensors with the same checkpoint key and slice spec were "
"found:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n")
# Validate no added tensors.
if checkpoint_key not in unseen_tensor_dict:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"not originally in the object graph was found in the "
"checkpoint shards:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n")
# Validate no shape change.
target_shape = unseen_tensor_dict[checkpoint_key][slice_spec].shape
if shard_tensor.shape != target_shape:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"was found with an altered shape:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" original tensor_shape: {target_shape}\n"
f" new tensor_shape: {shard_tensor.shape}\n")
# Validate no dtype change.
target_dtype = unseen_tensor_dict[checkpoint_key][slice_spec].dtype
if shard_tensor.dtype != target_dtype:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"was found with an altered dtype:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" original tensor_dtype: {target_dtype}\n"
f" new tensor_dtype: {shard_tensor.dtype}\n")
# Validate no task change.
target_task = device_lib.DeviceSpec.from_string(
unseen_tensor_dict[checkpoint_key][slice_spec].device).task
shard_tensor_task = device_lib.DeviceSpec.from_string(
shard_tensor.device).task
if shard_tensor_task != target_task:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"was found with an altered task:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" original tensor_task: {target_task}\n"
f" new tensor_task: {shard_tensor_task}\n")
# Validate tensors in shard have the same task.
if task_tensor is None:
task_tensor = ShardableTensor(
_tensor_save_spec=None,
tensor=None,
dtype=None,
device=shard_tensor.device,
name=None,
shape=None,
slice_spec=slice_spec,
checkpoint_key=checkpoint_key,
trackable=None)
else:
task1 = device_lib.DeviceSpec.from_string(task_tensor.device).task
task2 = device_lib.DeviceSpec.from_string(shard_tensor.device).task
if task1 is not None and task2 is not None and task1 != task2:
raise RuntimeError(
"After executing the checkpoint sharding callback, tensors "
"with different tasks were found in the same shard:\n"
f" callback_description: {callback_description}\n"
" tensor #1:"
f" checkpoint_key: {task_tensor.checkpoint_key}\n"
f" slice_spec: {task_tensor.slice_spec}\n"
f" task: {task1}\n"
" tensor #2:"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" task: {task2}\n")
del unseen_tensor_dict[checkpoint_key][slice_spec]
if not unseen_tensor_dict[checkpoint_key]:
del unseen_tensor_dict[checkpoint_key]
seen_tensor_set.add((checkpoint_key, slice_spec))
# validate no tensor removal
if unseen_tensor_dict:
tensors_info = ""
for ckpt_key, slice_spec in unseen_tensor_dict.items():
tensors_info += " tensor:\n"
tensors_info += f" checkpoint_key: {ckpt_key}\n"
tensors_info += f" slice_spec: {slice_spec}\n"
raise RuntimeError(
"After executing the checkpoint sharding callback, tensors in the "
"object graph were not found in the checkpoint shards:\n"
f" callback_description: {callback_description}\n"
f"{tensors_info}")
tensorflow/python/checkpoint/sharding/sharding_util_test.py
+430
View File
@@ -0,0 +1,430 @@
# 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 checkpoint sharding structures and utilities."""
from typing import Sequence
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import graph_view
from tensorflow.python.checkpoint.sharding import sharding_policies
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
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 resource_variable_ops
from tensorflow.python.training import server_lib
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.training.saving import saveable_object_util
class ShardingUtilTest(test.TestCase):
def _get_shardable_tensors_by_task(self, root):
serialized_tensors, _, _, _ = (
checkpoint.TrackableSaver(graph_view.ObjectGraphView(root))
._gather_serialized_tensors(None))
shardable_tensors_by_task = {}
for obj, tensor_dict in serialized_tensors.items():
for checkpoint_key, tensor_slice_dict in tensor_dict.items():
if not isinstance(tensor_slice_dict, dict):
# Make sure that maybe_tensor is structured as {slice_spec -> tensor}.
tensor_slice_dict = {"": tensor_slice_dict}
for slice_spec, tensor_save_spec in tensor_slice_dict.items():
if not isinstance(tensor_save_spec, saveable_object.SaveSpec):
tensor_save_spec = saveable_object.SaveSpec(
tensor=tensor_save_spec,
slice_spec=slice_spec,
name=checkpoint_key,
dtype=tensor_save_spec.dtype,
device=tensor_save_spec.device)
save_spec_tensor = tensor_save_spec.tensor
device = (device_lib.DeviceSpec.from_string(tensor_save_spec.device)
if isinstance(tensor_save_spec.device, str)
else tensor_save_spec.device)
task = device_lib.DeviceSpec.from_string(
saveable_object_util.set_cpu0(device.to_string()))
shardable_tensors_by_task.setdefault(task, []).append(
sharding_util.ShardableTensor(
_tensor_save_spec=tensor_save_spec,
tensor=save_spec_tensor,
dtype=tensor_save_spec.dtype,
device=device,
name=tensor_save_spec.name,
shape=save_spec_tensor.shape,
slice_spec=slice_spec.strip(),
checkpoint_key=checkpoint_key,
trackable=obj))
return shardable_tensors_by_task.values()
def test_hash_ShardingCallback(self):
class BlankCallback(sharding_util.ShardingCallback):
@property
def description(self):
return ""
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
pass
self.assertEqual(hash(BlankCallback()), hash(BlankCallback()))
class ValueCallback(sharding_util.ShardingCallback):
def __init__(self, val):
self.val = val
@property
def description(self):
return "value callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
pass
self.assertEqual(hash(ValueCallback(1)), hash(ValueCallback(1)))
self.assertNotEqual(hash(ValueCallback(1)), hash(ValueCallback(2)))
def test_validate_shards_correct(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
with ops.device("cpu:1"):
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
with ops.device("cpu:2"):
v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
root.v0 = v0
root.v1 = v1
root.v2 = v2
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = sharding_policies.ShardByTaskPolicy()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
self.assertEqual(
[list(shard.keys()) for shard in shards],
[[
"v0/.ATTRIBUTES/VARIABLE_VALUE",
"v1/.ATTRIBUTES/VARIABLE_VALUE",
"v2/.ATTRIBUTES/VARIABLE_VALUE",
"_CHECKPOINTABLE_OBJECT_GRAPH"
]])
self.assertEqual(
shards[0]["v0/.ATTRIBUTES/VARIABLE_VALUE"][""].numpy(),
v0.numpy())
self.assertEqual(
shards[0]["v1/.ATTRIBUTES/VARIABLE_VALUE"][""].numpy(),
v1.numpy())
self.assertEqual(
shards[0]["v2/.ATTRIBUTES/VARIABLE_VALUE"][""].numpy(),
v2.numpy())
def test_validate_shards_duplicate_tensor(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
with ops.device("cpu:1"):
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
root.v0 = v0
root.v1 = v1
class DuplicateTensorCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "duplicate tensor callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
tensor = shardable_tensors[0].tensor
checkpoint_key = shardable_tensors[0].checkpoint_key
slice_spec = shardable_tensors[0].slice_spec
shards = [
{checkpoint_key: {slice_spec: tensor}},
{checkpoint_key: {slice_spec: tensor}}
]
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = DuplicateTensorCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"multiple tensors with the same checkpoint "
"key and slice spec were found"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_added_tensor(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
root.v0 = v0
class AddedTensorCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "added tensor callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
checkpoint_key = "ADDED_TENSOR_ABC123"
slice_spec = ""
tensor = tensor_lib.Tensor()
return [{checkpoint_key: {slice_spec: tensor}}]
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = AddedTensorCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor not originally in the object graph"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_shape_change(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([[0.0, 1.0]], name="v0")
root.v0 = v0
class ShapeChangeCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "shape change callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shards = []
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
if checkpoint_key == "v0/.ATTRIBUTES/VARIABLE_VALUE":
tensor = array_ops.transpose(tensor)
shards.append({checkpoint_key: {slice_spec: tensor}})
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = ShapeChangeCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor was found with an altered shape"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_dtype_change(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
root.v0 = v0
class DtypeChangeCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "dtype change callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shards = []
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
if checkpoint_key == "v0/.ATTRIBUTES/VARIABLE_VALUE":
tensor = math_ops.cast(tensor, dtype=dtypes.int32)
shards.append({checkpoint_key: {slice_spec: tensor}})
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = DtypeChangeCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor was found with an altered dtype"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_task_change(self):
servers = [server_lib.Server.create_local_server() for _ in range(2)]
cluster_spec = server_lib.ClusterSpec({
"worker": [s.target[len("grpc://"):] for s in servers]})
remote.connect_to_cluster(cluster_spec)
root = module.Module()
with ops.device("/job:worker/task:0/cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(0.0, name="v1")
root.v0 = v0
root.v1 = v1
class TaskChangeCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "task change callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shards = []
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
if checkpoint_key == "v0/.ATTRIBUTES/VARIABLE_VALUE":
with ops.device("/job:worker/task:1/cpu:0"):
tensor = array_ops.identity(tensor)
shards.append({checkpoint_key: {slice_spec: tensor}})
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = TaskChangeCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor was found with an altered task"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_different_tasks(self):
servers = [server_lib.Server.create_local_server() for _ in range(2)]
cluster_spec = server_lib.ClusterSpec({
"worker": [s.target[len("grpc://"):] for s in servers]})
remote.connect_to_cluster(cluster_spec)
root = module.Module()
with ops.device("/job:worker/task:0/cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(0.0, name="v1")
root.v0 = v0
root.v1 = v1
class DifferentTasksCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "different tasks callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shard = {}
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
shard.setdefault(checkpoint_key, {})[slice_spec] = tensor
return [shard]
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = DifferentTasksCallback()
shards = sharding_callback(shardable_tensors_flat)
with self.assertRaisesRegex(RuntimeError,
"tensors with different tasks were found"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_tensor_removal(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
root.v0 = v0
class TensorRemovalCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "tensor removal callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
return []
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = TensorRemovalCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"tensors in the object graph were not found"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
if __name__ == "__main__":
test.main()