# Copyright 2018 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 distributed values library.""" import itertools import uuid from absl.testing import parameterized from tensorflow.python.checkpoint import checkpoint as trackable_utils from tensorflow.python.checkpoint import checkpoint_management as ckpt_manager from tensorflow.python.distribute import collective_all_reduce_strategy from tensorflow.python.distribute import combinations from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import strategy_combinations from tensorflow.python.distribute import strategy_test_lib from tensorflow.python.distribute import test_util from tensorflow.python.distribute import values from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver from tensorflow.python.eager import context from tensorflow.python.eager import def_function from tensorflow.python.eager import test from tensorflow.python.framework import config from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import indexed_slices 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 random_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variable_v1 from tensorflow.python.ops import variables as variables_lib from tensorflow.python.util import variable_utils def strategy_and_run_tf_function_combinations(): # Test the combination of different strategies and whether a tf.function # is passed into strategy.run.""" # TODO(b/197981388): re-enable MWMS test # return combinations.combine( # distribution=[ # strategy_combinations.mirrored_strategy_with_gpu_and_cpu, # ], # mode=["graph", "eager"], # experimental_run_tf_function=[True, False], # use_var_policy=[True, False]) + return combinations.combine( distribution=[ strategy_combinations.tpu_strategy, strategy_combinations.tpu_strategy_packed_var, ], mode=["graph", "eager"], experimental_run_tf_function=[True], use_var_policy=[True, False]) def strategy_with_var_policy(): return combinations.combine( distribution=[ strategy_combinations.mirrored_strategy_with_gpu_and_cpu, # TODO(b/197981388): re-enable MWMS test # strategy_combinations.multi_worker_mirrored_2x1_cpu, # strategy_combinations.multi_worker_mirrored_2x1_gpu, strategy_combinations.tpu_strategy, strategy_combinations.tpu_strategy_packed_var, ], mode=["graph", "eager"], use_var_policy=[True, False]) class OnWriteVariableSync(test.TestCase, parameterized.TestCase): @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssign(self, distribution, experimental_run_tf_function): def assign(fn, v, update_value, cross_replica): update_fn = lambda: getattr(v, fn)(update_value) if cross_replica: return update_fn() else: if experimental_run_tf_function: update_fn = def_function.function(update_fn) return test_util.gather(distribution, distribution.run(update_fn)) updates = [("assign", 1.), ("assign_add", 1.), ("assign_sub", -1.)] aggregations = [ variables_lib.VariableAggregation.NONE, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] options = list( x for x in itertools.product(updates, aggregations, [True, False])) for update, aggregation, cross_replica in options: # assign in replica context with SUM does not make sense cause you can # just do value * num replicas error is 1. is not a distributed value and # is unsupported for aggregation SUM if (not cross_replica and aggregation == variables_lib.VariableAggregation.SUM): continue with distribution.scope(): v = variable_v1.VariableV1( 0., aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) fn, update_value = update self.evaluate(assign(fn, v, update_value, cross_replica)) for component in v._values: self.assertAllEqual(self.evaluate(component.read_value()), self.evaluate(array_ops.ones_like(component))) @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssignOnWriteVar(self, distribution, experimental_run_tf_function): with distribution.scope(): v_to_assign = variable_v1.VariableV1( 2., aggregation=variables_lib.VariableAggregation.MEAN) v_to_assign_sub = variable_v1.VariableV1( -2., aggregation=variables_lib.VariableAggregation.MEAN) def assign(fn, v, update_value, cross_replica): update_fn = lambda: getattr(v, fn)(update_value) if cross_replica: return update_fn() else: if experimental_run_tf_function: update_fn = def_function.function(update_fn) return test_util.gather(distribution, distribution.run(update_fn)) updates = [("assign", v_to_assign), ("assign_add", v_to_assign), ("assign_sub", v_to_assign_sub)] aggregations = [ variables_lib.VariableAggregation.NONE, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] options = list( x for x in itertools.product(updates, aggregations, [True, False])) for update, aggregation, cross_replica in options: # assign in replica context with SUM does not make sense cause you can # just do value * num replicas error is 1. is not a distributed value and # is unsupported for aggregation SUM if aggregation == variables_lib.VariableAggregation.SUM: continue with distribution.scope(): v = variable_v1.VariableV1( 0., aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) fn, update_value = update self.evaluate(assign(fn, v, update_value, cross_replica)) for component in v._values: self.assertAllEqual(2.0, self.evaluate(component.read_value())) @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssignPerReplicaVal(self, distribution, experimental_run_tf_function): if strategy_test_lib.is_tpu_strategy(distribution): self.skipTest("Assigning PerReplica values is not supported. See" " sponge/80ba41f8-4220-4516-98ce-bbad48f9f11a.") with distribution.scope(): per_replica_value = values.PerReplica( [constant_op.constant(2.0), constant_op.constant(2.0)]) per_replica_sub_value = values.PerReplica( [constant_op.constant(-2.0), constant_op.constant(-2.0)]) def assign(fn, v, update_value, cross_replica): update_fn = lambda: getattr(v, fn)(update_value) if cross_replica: return update_fn() else: if experimental_run_tf_function: update_fn = def_function.function(update_fn) return test_util.gather(distribution, distribution.run(update_fn)) updates = [("assign", per_replica_value), ("assign_add", per_replica_value), ("assign_sub", per_replica_sub_value)] # We don't support assigning PerReplica valus to vars in replica context # with aggregation=NONE. aggregations = [ variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] options = list( x for x in itertools.product(updates, aggregations, [True, False])) for update, aggregation, cross_replica in options: # assign in replica context with SUM does not make sense cause you can # just do value * num replicas error is 1. is not a distributed value and # is unsupported for aggregation SUM if cross_replica: # We don't support assigning PerReplica values to MirroredVariables in # cross replica context continue with distribution.scope(): v = variable_v1.VariableV1( 0., aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) fn, update_value = update self.evaluate(assign(fn, v, update_value, cross_replica)) if aggregation == variables_lib.VariableAggregation.SUM: expected = 4.0 else: expected = 2.0 for component in v._values: self.assertAllEqual(expected, self.evaluate(component.read_value())) @combinations.generate(strategy_with_var_policy()) def testValueInReplicaContext(self, distribution): with distribution.scope(): v = variables_lib.Variable( 1., aggregation=variables_lib.VariableAggregation.MEAN) self.evaluate(variables_lib.global_variables_initializer()) @def_function.function def f(): with ops.control_dependencies([v.assign_add(1.)]): return v.value() results = self.evaluate( test_util.gather(distribution, distribution.run(f))) for value in results: self.assertEqual(2., value) @combinations.generate(strategy_with_var_policy()) def testValueInReplicaContextAssignDirectValue(self, distribution, use_var_policy): with distribution.scope(): v = variables_lib.Variable( 1., aggregation=variables_lib.VariableAggregation.MEAN) self.evaluate(variables_lib.global_variables_initializer()) @def_function.function def f(): with ops.control_dependencies([v.assign_add(1.)]): return v.value() results = self.evaluate( test_util.gather(distribution, distribution.run(f))) for value in results: self.assertEqual(2., value) @combinations.generate(strategy_and_run_tf_function_combinations()) def testReadValueInReplicaContext(self, distribution, experimental_run_tf_function): aggregations = [ variables_lib.VariableAggregation.NONE, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] for aggregation in aggregations: with distribution.scope(): v = variable_v1.VariableV1( 0., aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) if experimental_run_tf_function: read_var_fn = def_function.function(v.read_value) else: read_var_fn = v.read_value results = self.evaluate( test_util.gather(distribution, distribution.run(read_var_fn))) for component, value in zip(v._values, results): self.assertAllEqual(self.evaluate(component.read_value()), value) @combinations.generate(strategy_and_run_tf_function_combinations()) def testReadValueInCrossReplicaContext(self, distribution, experimental_run_tf_function): aggregations = [ variables_lib.VariableAggregation.NONE, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] for aggregation in aggregations: with distribution.scope(): v = variable_v1.VariableV1( 2., aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) if experimental_run_tf_function: read_var_fn = def_function.function(v.read_value) else: read_var_fn = v.read_value results = read_var_fn() for component in v._values: self.assertEqual(self.evaluate(component.read_value()), self.evaluate(results)) @combinations.generate(strategy_with_var_policy()) def testAssignOutOfScope(self, distribution): with distribution.scope(): mirrored = variables_lib.Variable(1.) self.evaluate(mirrored.assign(3.)) self.assertEqual(self.evaluate(mirrored.read_value()), 3.) for component in mirrored.values: self.assertEqual(self.evaluate(component.read_value()), 3.) @combinations.generate(strategy_with_var_policy()) def testInitializedToSameValueInsideEagerRun(self, distribution): if not context.executing_eagerly(): self.skipTest("eager only test") if isinstance(distribution.extended, collective_all_reduce_strategy.CollectiveAllReduceExtended): self.skipTest("Test for more than 1 device per worker only.") v = [None] @def_function.function def step(): def f(): if v[0] is None: v[0] = variables_lib.Variable(random_ops.random_normal([])) distribution.run(f) context.set_global_seed(None) step() vals = self.evaluate(v[0].values) self.assertAllEqual(vals[0], vals[1]) @combinations.generate(strategy_with_var_policy()) def testAggregationOnlyFirstReplica(self, distribution): if isinstance(distribution.extended, collective_all_reduce_strategy.CollectiveAllReduceExtended): self.skipTest("b/212945803") with distribution.scope(): v = variable_v1.VariableV1( 15., synchronization=variables_lib.VariableSynchronization.ON_WRITE, aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA) self.evaluate(variables_lib.global_variables_initializer()) @def_function.function def assign(): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group return v.assign(math_ops.cast(replica_id, dtypes.float32)) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(assign))) # The per-replica values should always match the first replicas value. self.assertAllEqual( array_ops.zeros(distribution.num_replicas_in_sync, dtypes.float32), per_replica_results) @combinations.generate(strategy_with_var_policy()) def testInitScope(self, distribution): if not context.executing_eagerly(): self.skipTest("eager only") class C(object): pass obj = C() obj.w = None obj.v = None @def_function.function def assign(): with ops.init_scope(): if obj.w is None: obj.w = variables_lib.Variable( 0., aggregation=variables_lib.VariableAggregation.MEAN) obj.v = variables_lib.Variable( obj.w.read_value(), aggregation=variables_lib.VariableAggregation.MEAN) self.evaluate(variables_lib.global_variables_initializer()) return obj.v.assign_add(2.) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(assign))) self.assertAllEqual([2., 2.], per_replica_results) @combinations.generate(strategy_with_var_policy()) def testOperatorOverride(self, distribution): if not context.executing_eagerly() and isinstance( distribution.extended, collective_all_reduce_strategy.CollectiveAllReduceExtended): self.skipTest("b/212954197") with distribution.scope(): v = variable_v1.VariableV1( 1, aggregation=variables_lib.VariableAggregation.SUM) self.evaluate(variables_lib.global_variables_initializer()) self.assertEqual(2, self.evaluate(v + 1)) @def_function.function def add(): return v + 1 per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(add))) self.assertAllEqual([2, 2], per_replica_results) @combinations.generate( combinations.combine( strategy=[ strategy_combinations.mirrored_strategy_with_gpu_and_cpu, strategy_combinations.tpu_strategy, strategy_combinations.tpu_strategy_packed_var, strategy_combinations.multi_worker_mirrored_2x1_cpu, strategy_combinations.multi_worker_mirrored_2x1_gpu, ], mode=["eager"], use_var_policy=[True, False])) def testSaveAndRestoreOnWrite(self, strategy): aggregation = [ variable_scope.VariableAggregation.NONE, variable_scope.VariableAggregation.ONLY_FIRST_REPLICA, variable_scope.VariableAggregation.SUM, variable_scope.VariableAggregation.MEAN ] for agg in aggregation: v_normal_restore = variables_lib.Variable(1.0) v_normal_save = variables_lib.Variable(3.0) with strategy.scope(): v_on_write = variables_lib.Variable(2.0, aggregation=agg) # Save ONWRITE Restore ONWRITE # Save ckpt = trackable_utils.Checkpoint(var=v_on_write) manager = ckpt_manager.CheckpointManager( ckpt, "/tmp/ckpt_" + str(uuid.uuid4()), max_to_keep=None) manager.save() # Restore ckpt.restore(manager.latest_checkpoint) self.assertEqual(2.0, self.evaluate(v_on_write._values[0])) self.assertEqual(2.0, self.evaluate(v_on_write.read_value())) # Save Mirrored Restore Normal # We've already saved Mirrored, so we only need to restore normal ckpt_normal = trackable_utils.Checkpoint(var=v_normal_restore) ckpt_normal.restore(manager.latest_checkpoint) self.assertEqual(2.0, self.evaluate(v_on_write._values[0])) self.assertEqual(2.0, self.evaluate(v_normal_restore.read_value())) # Save Normal Restore Mirrored # Save ckpt = trackable_utils.Checkpoint(var=v_normal_save) manager_2 = ckpt_manager.CheckpointManager( ckpt, "/tmp/ckptckpt_" + str(uuid.uuid4()), max_to_keep=None) manager_2.save() # Restore ckpt_on_write = trackable_utils.Checkpoint(var=v_on_write) ckpt_on_write.restore(manager_2.latest_checkpoint) self.assertEqual(3.0, self.evaluate(v_on_write._values[0])) self.assertEqual(3.0, self.evaluate(v_on_write.read_value())) ms_combination = combinations.combine( distribution=[strategy_combinations.mirrored_strategy_with_gpu_and_cpu], mode=["graph", "eager"]) tpu_combination = combinations.combine( distribution=[strategy_combinations.tpu_strategy_packed_var], mode=["graph", "eager"]) class OnWriteVariableSyncScatterTests(test.TestCase, parameterized.TestCase): @combinations.generate(ms_combination) def testScatterSub(self, distribution): with distribution.scope(): v = variables_lib.Variable( [0., 0., 0.], aggregation=variables_lib.VariableAggregation.MEAN) self.evaluate(v.initializer) @def_function.function def scatter_sub(): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group value = indexed_slices.IndexedSlices( values=array_ops_stack.stack([ math_ops.cast(replica_id, dtypes.float32), math_ops.cast(replica_id + 1, dtypes.float32) ]), indices=array_ops_stack.stack([replica_id, replica_id + 1]), dense_shape=(3,)) return v.scatter_sub(value) per_replica_results = self.evaluate( distribution.experimental_local_results( distribution.run(scatter_sub))) self.assertAllEqual([[0., -1., -1.], [0., -1., -1.]], per_replica_results) @combinations.generate(ms_combination) def testScatterAdd(self, distribution): with distribution.scope(): v = variables_lib.Variable( [0, 0, 0], aggregation=variables_lib.VariableAggregation.SUM) self.evaluate(v.initializer) @def_function.function def scatter_add(): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group value = indexed_slices.IndexedSlices( values=array_ops_stack.stack([replica_id, replica_id + 1]), indices=array_ops_stack.stack([replica_id, replica_id + 1]), dense_shape=(3,)) return v.scatter_add(value) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(scatter_add))) self.assertAllEqual([[0, 2, 2], [0, 2, 2]], per_replica_results) @combinations.generate(ms_combination) def testScatterDiv(self, distribution): with distribution.scope(): v = variables_lib.Variable( [1, 6, 1], aggregation=variables_lib.VariableAggregation.SUM) self.evaluate(v.initializer) @def_function.function def scatter_div(): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group value = indexed_slices.IndexedSlices( values=array_ops.reshape(replica_id + 2, [1]), indices=array_ops.reshape(replica_id, [1]), dense_shape=(3,)) return v.scatter_div(value) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(scatter_div))) self.assertAllEqual([[0, 2, 1], [0, 2, 1]], per_replica_results) @combinations.generate(ms_combination) def testScatterMul(self, distribution): with distribution.scope(): v = variables_lib.Variable( [2., 1., 1.], aggregation=variables_lib.VariableAggregation.MEAN) self.evaluate(v.initializer) @def_function.function def scatter_mul(): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group value = indexed_slices.IndexedSlices( values=array_ops.reshape( math_ops.cast(replica_id + 2, dtypes.float32), [1]), indices=array_ops.reshape(replica_id, [1]), dense_shape=(3,)) return v.scatter_mul(value) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(scatter_mul))) self.assertAllClose([[2., 1.5, 1.], [2., 1.5, 1.]], per_replica_results) @combinations.generate(ms_combination) def testScatterMin(self, distribution): with distribution.scope(): v1 = variables_lib.Variable( [0, 2, 0], aggregation=variables_lib.VariableAggregation.SUM) v2 = variables_lib.Variable( [0, 2, 0], aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA) self.evaluate(variables_lib.global_variables_initializer()) @def_function.function def scatter_min(v): value = indexed_slices.IndexedSlices( values=array_ops.identity([1]), indices=array_ops.identity([1]), dense_shape=(3,)) return v.scatter_min(value) with self.assertRaisesRegex(NotImplementedError, "scatter_min.*"): self.evaluate( test_util.gather(distribution, distribution.run(scatter_min, args=(v1,)))) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(scatter_min, args=(v2,)))) self.assertAllClose([[0, 1, 0], [0, 1, 0]], per_replica_results) @combinations.generate(ms_combination) def testScatterMax(self, distribution): with distribution.scope(): v1 = variables_lib.Variable( [0, 0, 0], aggregation=variables_lib.VariableAggregation.SUM) v2 = variables_lib.Variable( [0, 0, 0], aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA) self.evaluate(variables_lib.global_variables_initializer()) @def_function.function def scatter_max(v): value = indexed_slices.IndexedSlices( values=array_ops.identity([1]), indices=array_ops.identity([0]), dense_shape=(3,)) return v.scatter_max(value) with self.assertRaisesRegex(NotImplementedError, "scatter_max.*"): self.evaluate( test_util.gather(distribution, distribution.run(scatter_max, args=(v1,)))) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(scatter_max, args=(v2,)))) self.assertAllClose([[1, 0, 0], [1, 0, 0]], per_replica_results) @combinations.generate(ms_combination) def testScatterUpdate(self, distribution): with distribution.scope(): v1 = variables_lib.Variable( [0, 0, 0], aggregation=variables_lib.VariableAggregation.SUM) v2 = variables_lib.Variable( [0, 0, 0], aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA) self.evaluate(variables_lib.global_variables_initializer()) @def_function.function def scatter_update(v): value = indexed_slices.IndexedSlices( values=array_ops.identity([3]), indices=array_ops.identity([1]), dense_shape=(3,)) return v.scatter_update(value) with self.assertRaisesRegex(NotImplementedError, "scatter_update.*"): self.evaluate( test_util.gather(distribution, distribution.run(scatter_update, args=(v1,)))) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(scatter_update, args=(v2,)))) self.assertAllClose([[0, 3, 0], [0, 3, 0]], per_replica_results) @combinations.generate(ms_combination + tpu_combination) def testScatterOpsWithNoneAggregation(self, distribution): config.disable_mlir_bridge() def assert_close(v, op, delta, expect): scatter_op = getattr(v, op) @def_function.function def scatter_xxx(): return scatter_op(delta) per_replica_results = self.evaluate( variable_utils.convert_variables_to_tensors( distribution.experimental_local_results( distribution.run(scatter_xxx)))) self.assertAllClose([expect, expect], per_replica_results) with distribution.scope(): v = variables_lib.Variable( [4.], aggregation=variables_lib.VariableAggregation.NONE) self.evaluate(variables_lib.global_variables_initializer()) delta = indexed_slices.IndexedSlices( values=array_ops.identity([2.]), indices=array_ops.identity([0]), dense_shape=(1,)) assert_close(v, "scatter_sub", delta, [2.]) assert_close(v, "scatter_add", delta, [4.]) assert_close(v, "scatter_max", delta, [4.]) assert_close(v, "scatter_min", delta, [2.]) assert_close(v, "scatter_mul", delta, [4.]) assert_close(v, "scatter_div", delta, [2.]) assert_close(v, "scatter_update", delta, [2.]) @combinations.generate(ms_combination + tpu_combination) def testScatterOpsInCrossReplicaContext(self, distribution): with distribution.scope(): v1 = variables_lib.Variable( [1, 1, 1], aggregation=variables_lib.VariableAggregation.SUM) v2 = variables_lib.Variable([1, 1, 1]) self.evaluate(variables_lib.global_variables_initializer()) value = indexed_slices.IndexedSlices( values=array_ops.identity([2]), indices=array_ops.identity([0]), dense_shape=(3,)) with distribution.scope(): self.evaluate(v1.scatter_add(value)) self.assertAllEqual([3, 1, 1], self.evaluate(v1.read_value())) self.evaluate(v2.scatter_min(value)) self.assertAllEqual([1, 1, 1], self.evaluate(v2.read_value())) class OnReadVariableSyncTest(test.TestCase, parameterized.TestCase): @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssign(self, distribution, experimental_run_tf_function): def assign(fn, v, update_value, cross_replica): update_fn = lambda: getattr(v, fn)(update_value) if cross_replica: return update_fn() else: if experimental_run_tf_function: update_fn = def_function.function(update_fn) return test_util.gather(distribution, distribution.run(update_fn)) updates = [("assign", 1.), ("assign_add", 1.), ("assign_sub", -1.)] aggregations = [ variables_lib.VariableAggregation.NONE, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] options = list( x for x in itertools.product(updates, aggregations, [True, False])) for update, aggregation, cross_replica in options: # VariableAggregation.SUM in cross-replica mode is tested below, # VariableAggregation.NONE in cross-replica mode is not supported. if cross_replica and aggregation in [ variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.NONE, ]: continue with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) fn, update_value = update self.evaluate(assign(fn, v, update_value, cross_replica)) for component in v._values: self.assertAllEqual(self.evaluate(component.read_value()), self.evaluate(array_ops.ones_like(component))) @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssignOnReadVar(self, distribution, experimental_run_tf_function): with distribution.scope(): v_to_assign = variable_v1.VariableV1( 2., aggregation=variables_lib.VariableAggregation.MEAN) v_to_assign_sub = variable_v1.VariableV1( -2., aggregation=variables_lib.VariableAggregation.MEAN) def assign(fn, v, update_value, cross_replica): update_fn = lambda: getattr(v, fn)(update_value) if cross_replica: return update_fn() else: if experimental_run_tf_function: update_fn = def_function.function(update_fn) return test_util.gather(distribution, distribution.run(update_fn)) updates = [("assign", v_to_assign), ("assign_add", v_to_assign), ("assign_sub", v_to_assign_sub)] expected_cross_replica = { variables_lib.VariableAggregation.SUM: 1.0, variables_lib.VariableAggregation.MEAN: 2.0, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA: 2.0 } expected_replica = { variables_lib.VariableAggregation.SUM: 2.0, variables_lib.VariableAggregation.MEAN: 2.0, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA: 2.0 } # aggregation=NONE is not supported for OnReadVariables. aggregations = [ variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] options = list( x for x in itertools.product(updates, aggregations, [True, False])) for update, aggregation, cross_replica in options: # assign in replica context with SUM does not make sense cause you can # just do value * num replicas error is 1. is not a distributed value and # is unsupported for aggregation SUM if aggregation == variables_lib.VariableAggregation.SUM: continue with distribution.scope(): v = variable_v1.VariableV1( 0., aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) fn, update_value = update self.evaluate(assign(fn, v, update_value, cross_replica)) if cross_replica: for component in v._values: self.assertAllEqual(expected_cross_replica.get(aggregation), self.evaluate(component.read_value())) else: for component in v._values: self.assertAllEqual(expected_replica.get(aggregation), self.evaluate(component.read_value())) @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssignPerReplicaVal(self, distribution, experimental_run_tf_function): if strategy_test_lib.is_tpu_strategy(distribution): self.skipTest("Assigning PerReplica values is not supported. See" " sponge/80ba41f8-4220-4516-98ce-bbad48f9f11a.") self.skipTest( "We don't support assigning PerReplica values in cross " "replica context or replica context. see error in " "sponge/2b2e54c1-eda6-4534-82e1-c73b1dcd517f." ) with distribution.scope(): per_replica_value = values.PerReplica( [constant_op.constant(2.0), constant_op.constant(2.0)]) def assign(fn, v, update_value, cross_replica): update_fn = lambda: getattr(v, fn)(update_value) if cross_replica: return update_fn() else: if experimental_run_tf_function: update_fn = def_function.function(update_fn) return test_util.gather(distribution, distribution.run(update_fn)) updates = [("assign", per_replica_value)] # We don't support assigning PerReplica valus to vars in replica context # with aggregation=NONE. aggregations = [ variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] options = list( x for x in itertools.product(updates, aggregations, [True, False])) for update, aggregation, cross_replica in options: # assign in replica context with SUM does not make sense cause you can # just do value * num replicas error is 1. is not a distributed value and # is unsupported for aggregation SUM with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) fn, update_value = update # with self.assertRaisesRegex(ValueError, "Attempt to convert a value "): self.evaluate(assign(fn, v, update_value, cross_replica)) if aggregation == variables_lib.VariableAggregation.SUM: expected = 4.0 else: expected = 2.0 for component in v._values: self.assertAllEqual(expected, self.evaluate(component.read_value())) @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssignDtypeConversion(self, distribution, experimental_run_tf_function): def assign(fn, v, update_value, cross_replica): update_fn = lambda: getattr(v, fn)(update_value) if cross_replica: return update_fn() else: if experimental_run_tf_function: update_fn = def_function.function(update_fn) return test_util.gather(distribution, distribution.run(update_fn)) updates = [("assign", 1), ("assign_add", 1), ("assign_sub", -1)] aggregations = [ variables_lib.VariableAggregation.NONE, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] options = list( x for x in itertools.product(updates, aggregations, [True, False])) for update, aggregation, cross_replica in options: # VariableAggregation.SUM in cross-replica mode is tested below, # VariableAggregation.NONE in cross-replica mode is not supported. if cross_replica and aggregation in [ variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.NONE, ]: continue with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) fn, update_value = update self.evaluate(assign(fn, v, update_value, cross_replica)) for component in v._values: self.assertAllEqual(self.evaluate(component.read_value()), self.evaluate(array_ops.ones_like(component))) @combinations.generate(strategy_with_var_policy()) def testAssignWithAggregationSum(self, distribution): with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=variables_lib.VariableAggregation.SUM) self.evaluate(variables_lib.global_variables_initializer()) self.evaluate(v.assign(1. * distribution.num_replicas_in_sync)) for component in v._values: self.assertAllEqual(self.evaluate(component.read_value()), self.evaluate(array_ops.ones_like(component))) @combinations.generate(strategy_with_var_policy()) def testAssignAddSubWithAggregationSum(self, distribution): with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=variables_lib.VariableAggregation.SUM) self.evaluate(variables_lib.global_variables_initializer()) with self.assertRaisesRegex( ValueError, "SyncOnReadVariable does not support "): self.evaluate(v.assign_add(1.)) with self.assertRaisesRegex( ValueError, "SyncOnReadVariable does not support "): self.evaluate(v.assign_sub(1.)) @combinations.generate(strategy_and_run_tf_function_combinations()) def testReadValueInReplicaContext(self, distribution, experimental_run_tf_function): aggregations = [ variables_lib.VariableAggregation.NONE, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] for aggregation in aggregations: with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) if experimental_run_tf_function: read_var_fn = def_function.function(v.read_value) else: read_var_fn = v.read_value results = self.evaluate( test_util.gather(distribution, distribution.run(read_var_fn))) for component, value in zip(v._values, results): self.assertAllEqual(self.evaluate(component.read_value()), value) @combinations.generate(strategy_and_run_tf_function_combinations()) def testReadValueInCrossReplicaContext(self, distribution, experimental_run_tf_function): aggregations = [ variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] for aggregation in aggregations: if strategy_test_lib.is_tpu_strategy(distribution): resolver = tpu_cluster_resolver.TPUClusterResolver("") tpu_cluster_resolver.initialize_tpu_system(resolver) with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) def assign(v=v): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group return v.assign(math_ops.cast(replica_id, dtypes.float32)) if experimental_run_tf_function: assign = def_function.function(assign) self.evaluate(test_util.gather(distribution, distribution.run(assign))) num_replicas = distribution.num_replicas_in_sync sum_of_replica_values = num_replicas * (num_replicas - 1) / 2. if aggregation == variables_lib.VariableAggregation.SUM: expected = sum_of_replica_values elif aggregation == variables_lib.VariableAggregation.MEAN: expected = sum_of_replica_values / num_replicas else: expected = 0 self.assertEqual(expected, self.evaluate(v.read_value()), aggregation) self.assertEqual(expected, self.evaluate(v.value()), aggregation) self.assertEqual(expected, self.evaluate(v), aggregation) self.assertEqual(expected, self.evaluate(array_ops.identity(v)), aggregation) @combinations.generate(strategy_and_run_tf_function_combinations()) def testAllReduce(self, distribution, experimental_run_tf_function): with distribution.scope(): v = variable_v1.VariableV1( 2., synchronization=variables_lib.VariableSynchronization.ON_WRITE, aggregation=variables_lib.VariableAggregation.MEAN) self.evaluate(variables_lib.global_variables_initializer()) def all_reduce(): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group return ctx.all_reduce("SUM", v) + math_ops.cast(replica_id, dtypes.float32) if experimental_run_tf_function: all_reduce = def_function.function(all_reduce) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(all_reduce))) expected_result = [] for i in range(distribution.num_replicas_in_sync): expected_result.append(2.0 * distribution.num_replicas_in_sync + 1.0 * i) self.assertAllEqual(per_replica_results, tuple(expected_result)) @combinations.generate(strategy_and_run_tf_function_combinations()) def testAssignPerReplicaBeforeRead(self, distribution, experimental_run_tf_function): aggregations = [ variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ] for aggregation in aggregations: with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(variables_lib.global_variables_initializer()) def assign(var=v): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group return var.assign(math_ops.cast(replica_id, dtypes.float32)) if experimental_run_tf_function: assign = def_function.function(assign) per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(assign))) expected_result = [] for i in range(distribution.num_replicas_in_sync): expected_result.append(1.0 * i) self.assertAllEqual(per_replica_results, tuple(expected_result)) @combinations.generate(strategy_with_var_policy()) def testReadValueWithAggregationNoneInCrossReplicaContext(self, distribution): with distribution.scope(): v = variable_v1.VariableV1( 0., synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=variables_lib.VariableAggregation.NONE) self.evaluate(variables_lib.global_variables_initializer()) with self.assertRaisesRegex( ValueError, "Could not convert from .* VariableAggregation\\.NONE"): self.evaluate(v.read_value()) @combinations.generate(strategy_with_var_policy()) def testInitializedToSameValueInsideEagerRun(self, distribution): if not context.executing_eagerly(): self.skipTest("eager only") if isinstance(distribution.extended, collective_all_reduce_strategy.CollectiveAllReduceExtended): self.skipTest("Test for more than 1 device per worker only.") v = [None] @def_function.function def step(): def f(): if v[0] is None: v[0] = variables_lib.Variable( random_ops.random_normal([]), synchronization=variables_lib.VariableSynchronization.ON_READ) distribution.run(f) context.set_global_seed(None) step() vals = self.evaluate(v[0].values) self.assertAllEqual(vals[0], vals[1]) @combinations.generate(strategy_with_var_policy()) def testOperatorOverride(self, distribution): with distribution.scope(): v = variable_v1.VariableV1( 0.0, synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=variables_lib.VariableAggregation.MEAN) self.evaluate(variables_lib.global_variables_initializer()) @def_function.function def assign(): ctx = distribute_lib.get_replica_context() replica_id = ctx.replica_id_in_sync_group return v.assign(math_ops.cast(replica_id, dtypes.float32)) # Assign different replicas with different values. self.evaluate(test_util.gather(distribution, distribution.run(assign))) self.assertEqual(1.5, self.evaluate(v + 1)) @def_function.function def add(): return v + 1 per_replica_results = self.evaluate( test_util.gather(distribution, distribution.run(add))) self.assertAllEqual([1, 2], per_replica_results) @combinations.generate( combinations.combine( strategy=[ strategy_combinations.mirrored_strategy_with_gpu_and_cpu, strategy_combinations.tpu_strategy, strategy_combinations.tpu_strategy_packed_var, strategy_combinations.multi_worker_mirrored_2x1_cpu, strategy_combinations.multi_worker_mirrored_2x1_gpu, ], mode=["eager"], use_var_policy=[True, False])) def testSaveAndRestoreOnRead(self, strategy): aggregation = [variable_scope.VariableAggregation.SUM, variable_scope.VariableAggregation.MEAN] for agg in aggregation: v_normal_restore = variables_lib.Variable(1.0) v_normal_save = variables_lib.Variable(2.0) with strategy.scope(): v_on_read = variables_lib.Variable( 1.0, synchronization=variable_scope.VariableSynchronization.ON_READ, aggregation=agg) @def_function.function def assign_fn(): cluster_resolver = strategy.cluster_resolver replica_ctx = distribute_lib.get_replica_context() if ((cluster_resolver and cluster_resolver.task_type == "worker") or math_ops.equal(replica_ctx.replica_id_in_sync_group, constant_op.constant(1))): v_on_read.assign(3.) # pylint:disable=cell-var-from-loop else: v_on_read.assign(4.) # pylint:disable=cell-var-from-loop strategy.run(assign_fn) # Save ONREAD, restore ONREAD # Saves v[0] + v[1] = 7 for SUM and 3.5 for MEAN. ckpt = trackable_utils.Checkpoint(var=v_on_read) manager = ckpt_manager.CheckpointManager( ckpt, "/tmp/ckpt_" + str(uuid.uuid4()), max_to_keep=None) manager.save() # Restores a value of 7/2 = 3.5 for SUM and 3.5 for MEAN. ckpt.restore(manager.latest_checkpoint) self.assertEqual(3.5, self.evaluate(v_on_read._values[0])) # Save ONREAD, restore normal ckpt_normal = trackable_utils.Checkpoint(var=v_normal_restore) ckpt_normal.restore(manager.latest_checkpoint) if agg == variable_scope.VariableAggregation.SUM: self.assertEqual(7.0, self.evaluate(v_normal_restore.read_value())) else: self.assertEqual(3.5, self.evaluate(v_normal_restore.read_value())) # Save normal, restore ONREAD ckpt = trackable_utils.Checkpoint(var=v_normal_save) manager = ckpt_manager.CheckpointManager( ckpt, "/tmp/ckpt_" + str(uuid.uuid4()), max_to_keep=None) manager.save() # Restores a value of 2/2 = 1.0 for SUM and 2.0 for MEAN. ckpt_on_read = trackable_utils.Checkpoint(var=v_on_read) ckpt_on_read.restore(manager.latest_checkpoint) if agg == variable_scope.VariableAggregation.SUM: self.assertEqual(1.0, self.evaluate(v_on_read._values[0])) else: self.assertEqual(2.0, self.evaluate(v_on_read._values[0])) @combinations.generate( combinations.combine( distribution=[ strategy_combinations.mirrored_strategy_with_gpu_and_cpu, strategy_combinations.multi_worker_mirrored_2x1_cpu, strategy_combinations.multi_worker_mirrored_2x1_gpu, ], aggregation=[ variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ], mode=["graph", "eager"], use_var_policy=[True, False])) class SyncOnReadScatterReplicaTest(test.TestCase, parameterized.TestCase): def testScatterSub(self, distribution, aggregation): with distribution.scope(): v = variables_lib.Variable( [1., 1., 1.], synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(v.initializer) delta = values.PerReplica([ indexed_slices.IndexedSlices( values=[[0.], [1.]], indices=[0, 1], dense_shape=(3,)), indexed_slices.IndexedSlices( values=[[1.], [2.]], indices=[1, 2], dense_shape=(3,)), ]) with self.assertRaises(NotImplementedError): self.evaluate(distribution.run(v.scatter_sub, args=(delta,))) def testScatterAdd(self, distribution, aggregation): with distribution.scope(): v = variables_lib.Variable( [1., 1., 1.], synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(v.initializer) delta = values.PerReplica([ indexed_slices.IndexedSlices( values=[[0.], [1.]], indices=[0, 1], dense_shape=(3,)), indexed_slices.IndexedSlices( values=[[1.], [2.]], indices=[1, 2], dense_shape=(3,)), ]) with self.assertRaises(NotImplementedError): self.evaluate(distribution.run(v.scatter_add, args=(delta,))) def testScatterDiv(self, distribution, aggregation): with distribution.scope(): v = variables_lib.Variable( [2., 6., 1.], synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(v.initializer) delta = values.PerReplica([ indexed_slices.IndexedSlices( values=[[2.], [2.]], indices=[0, 1], dense_shape=(3,)), indexed_slices.IndexedSlices( values=[[3.], [3.]], indices=[1, 2], dense_shape=(3,)), ]) with self.assertRaises(NotImplementedError): self.evaluate(distribution.run(v.scatter_div, args=(delta,))) def testScatterMul(self, distribution, aggregation): with distribution.scope(): v = variables_lib.Variable( [2., 1., 1.], synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(v.initializer) delta = values.PerReplica([ indexed_slices.IndexedSlices( values=[[2.], [3.]], indices=[0, 1], dense_shape=(3,)), indexed_slices.IndexedSlices( values=[[4.], [5.]], indices=[1, 2], dense_shape=(3,)), ]) with self.assertRaises(NotImplementedError): self.evaluate(distribution.run(v.scatter_mul, args=(delta,))) def testScatterMin(self, distribution, aggregation): with distribution.scope(): v = variables_lib.Variable( [3., 4., 5.], synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(v.initializer) delta = values.PerReplica([ indexed_slices.IndexedSlices( values=[[1.], [8.]], indices=[0, 1], dense_shape=(3,)), indexed_slices.IndexedSlices( values=[[9.], [2.]], indices=[1, 2], dense_shape=(3,)), ]) with self.assertRaises(NotImplementedError): self.evaluate(distribution.run(v.scatter_min, args=(delta,))) def testScatterMax(self, distribution, aggregation): with distribution.scope(): v = variables_lib.Variable( [3., 4., 5.], synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(v.initializer) delta = values.PerReplica([ indexed_slices.IndexedSlices( values=[[1.], [8.]], indices=[0, 1], dense_shape=(3,)), indexed_slices.IndexedSlices( values=[[9.], [2.]], indices=[1, 2], dense_shape=(3,)), ]) with self.assertRaises(NotImplementedError): self.evaluate(distribution.run(v.scatter_max, args=(delta,))) def testScatterUpdate(self, distribution, aggregation): with distribution.scope(): v = variables_lib.Variable( [0., 0., 0.], synchronization=variables_lib.VariableSynchronization.ON_READ, aggregation=aggregation) self.evaluate(v.initializer) delta = values.PerReplica([ indexed_slices.IndexedSlices( values=[[1.], [2.]], indices=[0, 1], dense_shape=(3,)), indexed_slices.IndexedSlices( values=[[3.], [4.]], indices=[1, 2], dense_shape=(3,)), ]) with self.assertRaises(NotImplementedError): self.evaluate(distribution.run(v.scatter_min, args=(delta,))) if __name__ == "__main__": test_util.main()