wehub-resource-sync
261 lines
10 KiB
Python
261 lines
10 KiB
Python
# 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)
|
|
}
|