Updating TF to support 2 prng seeds.

Summary:
Updating PoplarCompiler to generate a replica identical prng seed during compilation. This seed gets used by the BaseVisitor to set the appropriate StochasticRoundingMethod for each instruction as it gets lowered. Seeds are then switched as we encounter instructions that require different methods. The seed state is stored in CompilerResources to keep it consistent during lowering, since multiple visitors will be used to lower the whole program.

Seed switching is behind the experimental_prng_stability flag and so disabled by default.

Test Plan: CI

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep, jakeh, samuelh

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep, samuelh

Subscribers: alfiee, samuelh

Maniphest Tasks: T41157

Differential Revision: https://phabricator.sourcevertex.net/D51984

Author: BabakkGraphcore

tensorflow/compiler/plugin/poplar/driver/compiler_resources.h

@@ -36,6 +36,7 @@ limitations under the License.
 #include "tensorflow/compiler/plugin/poplar/driver/compiler_information.h"
 #include "tensorflow/compiler/plugin/poplar/driver/config.pb.h"
 #include "tensorflow/compiler/plugin/poplar/driver/passes/convolution_classifier.h"
+#include "tensorflow/compiler/plugin/poplar/driver/prng_seed_state.h"
 #include "tensorflow/compiler/plugin/poplar/driver/tools/execution_counter_util.h"
 #include "tensorflow/compiler/plugin/poplar/driver/tools/generic_graph_caching.h"
 #include "tensorflow/compiler/plugin/poplar/driver/tools/mapping_helper.h"
@@ -186,6 +187,8 @@ struct CompilerResources {
   // The implementation of the progress bar.
   std::unique_ptr<ProgressBarBase> progress_bar;
 
+  PrngSeedState prng_seed_state;
+
   CompilerResources(
       HloModule* module, const CompilerInformation& information,
       const poplar::OptionFlags& conv_options,

tensorflow/compiler/plugin/poplar/driver/poplar_compiler.cc

@@ -23,17 +23,20 @@ limitations under the License.
 
 #include <algorithm>
 #include <fstream>
+#include <gcl/Collectives.hpp>
 #include <gcl/TileAllocation.hpp>
 #include <limits>
 #include <mutex>
 #include <popfloat/experimental/codelets.hpp>
 #include <poplar/CSRFunctions.hpp>
 #include <poplar/CodeletFileType.hpp>
 #include <poplar/CycleCount.hpp>
+#include <poplar/RandomSeed.hpp>
 #include <poplar/exceptions.hpp>
 #include <poplar/replication_factor.hpp>
 #include <poplin/codelets.hpp>
 #include <popnn/codelets.hpp>
+#include <popops/Cast.hpp>
 #include <popops/codelets.hpp>
 #include <poprand/RandomGen.hpp>
 #include <poprand/codelets.hpp>
@@ -537,7 +540,7 @@ HloPrintOptions GetPrintOptions() {
 }
 
 StatusOr<poplar::program::Program> InitializeSeed(
-    poplar::Graph& graph, int replication_factor,
+    poplar::Graph& graph, int replication_factor, CompilerResources& resources,
     const poplar::DebugContext& debug_context = {"__seed"}) {
   PoplarOpDefDebugInfo debug_info(debug_context, "InitializeSeed");
 
@@ -561,7 +564,29 @@ StatusOr<poplar::program::Program> InitializeSeed(
                         initializer.GetData(ShapeUtil::MakeShape(U32, {2})));
     TF_RETURN_IF_ERROR(SetInitialTensorValue(graph, seed, literal));
   }
-  poprand::setSeed(graph, seed, 0, seq, {debug_info, "set"});
+
+  if (replication_factor > 1 && resources.enable_experimental_prng_stability) {
+    // When running a replicated graph we need an additional seed, which is
+    // identical for all IPUs. This gets used to ensure determinism when
+    // performing stochastic rounding.
+    poplar::Tensor identical_seed =
+        // Using MEAN to preserve the seed value incase it's been explicitly set
+        // to be the same for each replica.
+        popops::cast(graph,
+                     gcl::allReduceCrossReplica(
+                         graph, popops::cast(graph, seed, poplar::FLOAT, seq),
+                         popops::CollectiveOperator::MEAN, seq,
+                         {debug_info, "allReduceSeed"}),
+                     poplar::UNSIGNED_INT, seq);
+
+    resources.prng_seed_state =
+        PrngSeedState::SetupSeeds(graph, identical_seed, seed, seq);
+  } else {
+    resources.prng_seed_state = PrngSeedState::SetupSeed(graph, seed, seq);
+  }
+
+  resources.prng_seed_state.ChangeStochasticRoundingMethod(
+      StochasticRoundingMethod_DifferingSeeds, seq);
 
   return seq;
 }
@@ -1679,6 +1704,11 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
     TF_RETURN_IF_ERROR(CreatePoplarGraphs(resources, module, poplar_executor));
     auto& main_graph = GetMasterGraph(resources);
 
+    // Set up the random seed.
+    TF_ASSIGN_OR_RETURN(
+        auto seed_setup,
+        InitializeSeed(main_graph, replication_factor, resources));
+
     EntryVisitor visitor(resources, entry);
     try {
       Tracepoint tracepoint("PoplarGraphConstruction");
@@ -1720,9 +1750,6 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
       main_program.add(poplar::program::Sync(poplar::SyncType::GLOBAL));
     }
 
-    // Set up the random seed.
-    TF_ASSIGN_OR_RETURN(auto seed_setup,
-                        InitializeSeed(main_graph, replication_factor));
     main_program.add(seed_setup);
 
     // Set up the floating point control register if required

tensorflow/compiler/plugin/poplar/driver/prng_seed_state.cc

@@ -45,6 +45,23 @@ std::unique_ptr<poputil::graphfn::TensorFunction> CreateChangeHwSeedsFn(
 }
 }  // namespace
 
+/*static*/ PrngSeedState PrngSeedState::SetupSeed(
+    poplar::Graph& graph, poplar::Tensor& seed,
+    poplar::program::Sequence& seq) {
+  const poplar::DebugContext& debug_context = {"__seed"};
+  PoplarOpDefDebugInfo debug_info(debug_context, "InitializeSeed");
+
+  poprand::setSeed(graph, seed, 0, seq, {debug_info, "set"});
+  auto differing_hw_seed = poplar::getHwSeeds(graph, seq, {debug_info, "get"});
+
+  // We want the behaviour to be consistent whether we're running
+  // with single or multiple seeds, so even when there's no replication
+  // we pretend that there's a separate identical seed and do everything
+  // else as normal.
+  return PrngSeedState(graph, StochasticRoundingMethod_DifferingSeeds,
+                       differing_hw_seed, differing_hw_seed);
+}
+
 /*static*/ PrngSeedState PrngSeedState::SetupSeeds(
     poplar::Graph& graph, poplar::Tensor& identical_seed,
     poplar::Tensor& differing_seed, poplar::program::Sequence& seq) {
@@ -57,7 +74,7 @@ std::unique_ptr<poputil::graphfn::TensorFunction> CreateChangeHwSeedsFn(
 
   poprand::setSeed(graph, differing_seed, 0, seq, {debug_info, "setDistinct"});
   auto differing_hw_seed =
-      poplar::getHwSeeds(graph, seq, {debug_info, "getIdenticalHw"});
+      poplar::getHwSeeds(graph, seq, {debug_info, "getDistinctHw"});
 
   // Speciifying DifferingSeeds since the last seed set was the differing one.
   return PrngSeedState(graph, StochasticRoundingMethod_DifferingSeeds,

tensorflow/compiler/plugin/poplar/driver/prng_seed_state.h

@@ -32,6 +32,10 @@ namespace poplarplugin {
 // flag.
 class PrngSeedState {
  public:
+  // Create a PrngSeedState from single or multiple seeds. This corresponds
+  // to whether we're running with replication (2 seeds) or not (1 seed).
+  static PrngSeedState SetupSeed(poplar::Graph& graph, poplar::Tensor& seed,
+                                 poplar::program::Sequence& seq);
   static PrngSeedState SetupSeeds(poplar::Graph& graph,
                                   poplar::Tensor& identical_seed,
                                   poplar::Tensor& differing_seed,

tensorflow/compiler/plugin/poplar/driver/visitors/visitor_arithmetic_expr.cc

@@ -46,7 +46,13 @@ ArithmeticExprVisitor::ArithmeticExprVisitor(
     const poplar::DebugNameAndId& debug_name_and_id)
     : BaseVisitor(res, debug_name_and_id),
       inputs_(std::move(inputs)),
-      caller_(caller) {}
+      caller_(caller) {
+  // We don need to change the seed as we visit instructions, since we will not
+  // be executing them individually but as fused unit via a single poplar call.
+  // Hence the seed we need to use will be set by the calling
+  // instruction/visitor.
+  allow_seed_changes_ = false;
+}
 
 StatusOr<std::unique_ptr<popops::expr::Expr>>
 ArithmeticExprVisitor::FindExpressionInput(const HloInstruction* inst) {

tensorflow/compiler/plugin/poplar/driver/visitors/visitor_base.cc

@@ -88,14 +88,32 @@ Status BaseVisitor::Preprocess(HloInstruction* inst) {
     stochastic_rounding_enabled_ = new_stochastic_rounding_enabled;
   }
 
+  if (allow_seed_changes_ && stochastic_rounding_enabled_) {
+    poplar::DebugNameAndId debug_name_and_id{"changeSRMethod"};
+    poplar::program::Sequence seq({}, debug_name_and_id);
+
+    const auto new_sr_method =
+        poplar_backend_config.stochastic_rounding_method();
+    if (resources_.prng_seed_state.ChangeStochasticRoundingMethod(
+            new_sr_method, seq, debug_name_and_id)) {
+      AddSequenceForInstruction(inst, seq);
+    }
+
+    VLOG(3) << "Using SR method "
+            << StochasticRoundingMethod_Name(
+                   resources_.prng_seed_state.GetStochasticRoundingMethod())
+            << " for instruction '" << inst->name() << "'";
+  }
+
   return Status::OK();
 }
 
 BaseVisitor::BaseVisitor(CompilerResources& resources,
                          const poplar::DebugNameAndId& debug_name_and_id)
     : resources_(resources),
       dnai_(debug_name_and_id),
-      execution_counters_(resources, debug_name_and_id) {
+      execution_counters_(resources, debug_name_and_id),
+      allow_seed_changes_(resources.enable_experimental_prng_stability) {
   stochastic_rounding_enabled_ =
       resources_.global_floating_point_behaviour.esr();
 

tensorflow/compiler/plugin/poplar/driver/visitors/visitor_base.h

@@ -207,6 +207,10 @@ class BaseVisitor : public DfsHloVisitor {
   // Scope execution counters.
   ExecutionCounters execution_counters_;
 
+  // Control whether changing seeds is allowed during instruction lowering,
+  // used to improve prng stability.
+  bool allow_seed_changes_ = false;
+
  private:
   Status CreateSequenceGroupedByInstruction(
       const HloInstruction* inst, const poplar::program::Sequence& seq);

tensorflow/python/ipu/tests/functional_ops_test.py

@@ -258,6 +258,7 @@ def body(x, labels):
     # pylint: disable=line-too-long
     ok = [
         '__seed/set/setMasterSeed',
+        '__seed/get/getSeeds',
         'matmul/dot*/Conv_1',
         'add_0/fusion/Op/Add',
         'Sigmoid/sigmoid/Nonlinearity',

tensorflow/python/ipu/tests/multi_conv_test.py

@@ -311,6 +311,7 @@ def inputs_fn():
         'ipu/Mean*/reduce*/Reduce',
         'ipu/add',
         '__seed/set/setMasterSeed',
+        '__seed/get/getSeeds',
         '[cC]opy',
     ]
     _compare_ipu_to_cpu(self,

tensorflow/python/ipu/tests/replicated_seed_control_test.py

@@ -18,9 +18,14 @@
 
 from tensorflow.python import ipu
 from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import test_util
+from tensorflow.python.framework import random_seed
+from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.compiler.plugin.poplar.ops import gen_popops_ops
 from tensorflow.compiler.plugin.poplar.tests import test_utils as tu
 from tensorflow.python.platform import googletest
 
@@ -170,6 +175,100 @@ def test_identical_replica_seeds(self):
       self.assertNotAllEqual(res1, res1_second)
       self.assertNotAllEqual(res2, res2_second)
 
+  @tu.test_uses_ipus(num_ipus=2)
+  @test_util.deprecated_graph_mode_only
+  def test_experimental_identical_seeds(self):
+    inp = array_ops.placeholder(np.float32, [10])
+
+    # This should produce the same result on each IPU
+    # as the stochastic rounding should be performed
+    # using the same seed.
+    with ipu.scopes.ipu_scope('/device:IPU:0'):
+      cast1 = math_ops.cast(inp, dtype=np.float16)
+      out0 = gen_popops_ops.ipu_all_gather(cast1, replication_factor=2)
+
+    # Configure the hardware
+    config = IPUConfig()
+    config.auto_select_ipus = [2]
+    config.experimental.enable_prng_stability = True
+    tu.add_hw_ci_connection_options(config)
+    # Enable stochastic rounding
+    config.floating_point_behaviour.esr = True
+    config.configure_ipu_system()
+
+    in_data = np.array([0.1] * 10)
+
+    with session_lib.Session() as sess:
+      res = sess.run(out0, {inp: in_data}).astype(np.float32)
+      # Compare the result of each IPU
+      self.assertAllEqual(res[0], res[1])
+
+  @tu.test_uses_ipus(num_ipus=4)
+  @test_util.deprecated_graph_mode_only
+  def test_experimental_identical_compute(self):
+    config = IPUConfig()
+    config.auto_select_ipus = [4]
+    config.experimental.enable_prng_stability = True
+    tu.add_hw_ci_connection_options(config)
+    # Enable stochastic rounding
+    config.floating_point_behaviour.esr = True
+    config.configure_ipu_system()
+
+    random_seed.set_seed(1234)
+    dataset = dataset_ops.Dataset.from_tensor_slices([[[0.1, 0.2, 0.3],
+                                                       [0.4, 0.5, 0.6]]])
+    dataset = dataset.repeat()
+
+    infeed = ipu.ipu_infeed_queue.IPUInfeedQueue(dataset)
+    outfeed = ipu.ipu_outfeed_queue.IPUOutfeedQueue()
+
+    noise = random_ops.random_uniform([2, 3], dtype=np.float32, seed=1)
+
+    def my_net():
+      # Perform some arbitrary simple compute that invokes SR.
+      def body(noise, infeed_value):
+        infeed_value = ipu.cross_replica_ops.assume_equal_across_replicas(
+            infeed_value)
+
+        noise_f16 = math_ops.cast(noise, dtype=np.float16)
+        infeed_val_f16 = math_ops.cast(infeed_value, dtype=np.float16)
+
+        result = infeed_val_f16 + noise_f16
+
+        const = constant_op.constant(0.01,
+                                     shape=result.shape,
+                                     dtype=np.float32)
+        const_f16 = math_ops.cast(const, dtype=np.float16)
+
+        result = result * const_f16
+        result_f32 = math_ops.cast(result, dtype=np.float32)
+
+        out = outfeed.enqueue(result_f32)
+        return (result_f32, out)
+
+      r = ipu.loops.repeat(10, body, [noise], infeed)
+      return r
+
+    with ipu.scopes.ipu_scope("/device:IPU:0"):
+      res = ipu.ipu_compiler.compile(my_net, inputs=[])
+
+    outqueue = outfeed.dequeue()
+    with session_lib.Session() as sess:
+      sess.run(infeed.initializer)
+      sess.run(res)
+      results = sess.run(outqueue)
+
+    # This contains a value per replica for each iteration of the loop. We want to test that
+    # for each iteration those values are all equal, since each replica is using the same data and
+    # we've enabled experimental prng stability.
+    self.assertEqual(len(results), 10)
+    for replica_results in results:
+      replicas_equal = [
+          replica_results[0].tolist() == result.tolist()
+          for result in replica_results
+      ]
+      self.assertTrue(all(replicas_equal))
+
 
 if __name__ == "__main__":
   googletest.main()