Add support for torch.Generator type in TorchScript (pytorch#110413)

- Add support for `torch.Generator` type in TorchScript
- Add `generator` args to all `torch.nn.init` functions that call `uniform_` or `normal_`
- Add support for `torch.Generator` in LTC's TorchScript backend (CC: @wconstab)

CC: @eellison @davidberard98 @GlebKazantaev @behzad-a
Pull Request resolved: pytorch#110413
Approved by: https://github.com/wconstab, https://github.com/albanD, https://github.com/glebk-cerebras, https://github.com/davidberard98

Author: antoniojkim

aten/src/ATen/core/ivalue.cpp

@@ -644,6 +644,13 @@ std::ostream& IValue::repr(
       c10::printQuotedString(out, device_stream.str());
       return out << ")";
     }
+    case IValue::Tag::Generator: {
+      auto generator = v.toGenerator();
+      out << "torch.Generator(device=";
+      c10::printQuotedString(out, generator.device().str());
+      out << ", seed=" << generator.current_seed() << ")";
+      return out;
+    }
     case IValue::Tag::GenericDict:
       return printMaybeAnnotatedDict(out, v, formatter);
     case IValue::Tag::Enum: {
@@ -956,6 +963,7 @@ IValue IValue::deepcopy(
     case IValue::Tag::SymBool:
     case IValue::Tag::Bool:
     case IValue::Tag::Device:
+    case IValue::Tag::Generator:
     case IValue::Tag::Uninitialized: {
       copy = *this;
     } break;

aten/src/ATen/core/type_factory.cpp

@@ -28,6 +28,7 @@ namespace c10 {
   _(complex, ComplexType)           \
   _(str, StringType)                \
   _(Device, DeviceObjType)          \
+  _(Generator, GeneratorType)       \
   _(Stream, StreamObjType)          \
   _(number, NumberType)             \
   _(None, NoneType)                 \

aten/src/ATen/native/ts_native_functions.yaml

@@ -168,6 +168,9 @@ full_codegen:
   - slice_scatter
   - diagonal_scatter
   - as_strided_scatter
+  # random ops
+  - normal_functional
+  - uniform
 ir_gen:
   - selu
 supported:
@@ -177,7 +180,6 @@ supported:
   - empty.memory_format
   - empty_strided
   - fill_.Scalar
-  - normal_
   - max_pool3d_with_indices
   - max_pool3d_with_indices_backward
   - _to_copy

build_variables.bzl

@@ -447,7 +447,6 @@ lazy_tensor_ts_sources = [
     "torch/csrc/lazy/ts_backend/dynamic_ir.cpp",
     "torch/csrc/lazy/ts_backend/config.cpp",
     "torch/csrc/lazy/ts_backend/ops/device_data.cpp",
-    "torch/csrc/lazy/ts_backend/ops/random_ops.cpp",
     "torch/csrc/lazy/ts_backend/ops/generic.cpp",
     "torch/csrc/lazy/ts_backend/tensor_aten_ops.cpp",
     "torch/csrc/lazy/ts_backend/ts_autograd_functions.cpp",

docs/source/jit_unsupported.rst

@@ -88,4 +88,3 @@ we suggest using :meth:`torch.jit.trace`.
   * :class:`torch.nn.AdaptiveLogSoftmaxWithLoss`
   * :class:`torch.autograd.Function`
   * :class:`torch.autograd.enable_grad`
-  * :class:`torch.Generator`

test/jit/test_generator.py

@@ -0,0 +1,195 @@
+# Owner(s): ["oncall: jit"]
+
+import io
+import math
+import unittest
+
+import torch
+from torch.nn import init
+from torch.testing._internal.common_utils import skipIfLegacyJitExecutor
+from torch.testing._internal.jit_utils import JitTestCase
+
+
+if __name__ == "__main__":
+    raise RuntimeError(
+        "This test file is not meant to be run directly, use:\n\n"
+        "\tpython test/test_jit.py TESTNAME\n\n"
+        "instead."
+    )
+
+
+class TestGenerator(JitTestCase):
+    # torch.jit.trace does not properly capture the generator manual seed
+    # and thus is non deterministic even if the generator is manually seeded
+    @skipIfLegacyJitExecutor("legacy JIT executor does not support Generator type")
+    @unittest.expectedFailure
+    def test_trace(self):
+        def f():
+            generator = torch.Generator()
+            generator.seed()
+            generator.manual_seed(2023)
+            generator.initial_seed()
+            tensor = torch.empty(2, 2)
+            tensor.uniform_(0, 1, generator=generator)
+            return tensor
+
+        traced_f = torch.jit.trace(f, ())
+
+        # Run this 3 times to ensure that the generator is being manually seeded
+        # each time the traced function is run
+        for i in range(3):
+            torch.manual_seed(1)
+
+            eager_tensor = f()
+
+            # Change the seed of the default generator to
+            # check that we're using the generator from the
+            # trace
+            torch.manual_seed(2)
+            traced_tensor = traced_f()
+
+            self.assertEqual(eager_tensor, traced_tensor)
+
+    def test_script(self):
+        def f():
+            generator = torch.Generator()
+            generator.seed()
+            generator.manual_seed(2023)
+            generator.initial_seed()
+            tensor = torch.empty(2, 2)
+            tensor.normal_(-1.0, 1.0, generator=generator)
+            return tensor
+
+        script_f = torch.jit.script(f, ())
+
+        # Run this 3 times to ensure that the generator is being manually seeded
+        # each time the traced function is run
+        for i in range(3):
+            torch.manual_seed(1)
+
+            eager_tensor = f()
+
+            # Change the seed of the default generator to
+            # check that we're using the generator from the
+            # trace
+            torch.manual_seed(2)
+
+            script_tensor = script_f()
+
+            self.assertEqual(eager_tensor, script_tensor)
+
+    def test_default_generator(self):
+        def f():
+            # check that calling manual seed for the default generator works
+            torch.manual_seed(2023)
+            tensor = torch.empty(2, 2)
+            tensor.normal_(-1.0, 1.0)
+            return tensor
+
+        torch.manual_seed(1)
+
+        eager_tensor = f()
+
+        torch.manual_seed(2)
+
+        script_f = torch.jit.script(f, ())
+        script_tensor = script_f()
+
+        self.assertEqual(eager_tensor, script_tensor)
+
+    def test_generator_arg(self):
+        def f(generator: torch.Generator):
+            tensor = torch.empty(2, 2)
+            tensor.normal_(-1.0, 1.0, generator=generator)
+            return tensor
+
+        generator = torch.Generator()
+        generator.manual_seed(2023)
+
+        script_f = torch.jit.script(f, (generator,))
+
+        for i in range(3):
+            generator = torch.Generator()
+            generator.manual_seed(2023 + i)
+
+            torch.manual_seed(1 + i)
+
+            eager_tensor = f(generator)
+
+            generator = torch.Generator()
+            generator.manual_seed(2023 + i)
+
+            torch.manual_seed(1 + i)
+
+            script_tensor = script_f(generator)
+
+            self.assertEqual(eager_tensor, script_tensor)
+
+    def test_save_load(self):
+        class Foo(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.foo = torch.nn.Linear(2, 2, bias=False)
+                self.bar = torch.nn.Linear(2, 2, bias=False)
+
+                self.reset_parameters()
+
+            def reset_linear(self, module, generator):
+                init.kaiming_uniform_(
+                    module.weight, a=math.sqrt(5), generator=generator
+                )
+
+            def reset_parameters(self):
+                generator = torch.Generator()
+                generator.manual_seed(1)
+                self.reset_linear(self.foo, generator)
+
+                generator = torch.Generator()
+                generator.manual_seed(2)
+                self.reset_linear(self.bar, generator)
+
+            def forward(self, x):
+                x = self.foo(x)
+                x = self.bar(x)
+
+                generator = torch.Generator()
+                generator.manual_seed(3)
+                r = torch.empty_like(x)
+                r.normal_(0.0, 1.0, generator=generator)
+
+                return x, r
+
+        eager_foo = Foo()
+
+        script_module = torch.jit.script(Foo())
+        saved_module = io.BytesIO()
+        torch.jit.save(script_module, saved_module)
+        saved_module.seek(0)
+
+        loaded_module = torch.jit.load(saved_module)
+
+        self.assertEqual(eager_foo.foo.weight, loaded_module.foo.weight)
+        self.assertEqual(eager_foo.bar.weight, loaded_module.bar.weight)
+
+        try:
+            # Run this 3 times so make sure that the generator seed is being set
+            # every time forward is called
+            for i in range(3):
+                x = torch.ones(2, 2)
+                out1, r1 = eager_foo(x)
+                out2, r2 = loaded_module(x)
+
+                try:
+                    self.assertEqual(out1, out2)
+                except:  # noqa: B001, E722
+                    print(f"Iteration {i}:\n{out1=}\n{out2=}")
+                    raise
+
+                try:
+                    self.assertEqual(r1, r2)
+                except:  # noqa: B001, E722
+                    print(f"Iteration {i}:\n{r1=}\n{r2=}")
+                    raise
+        except:  # noqa: B001, E722
+            print(loaded_module.forward.code)
+            raise

test/lazy/test_generator.py

@@ -0,0 +1,103 @@
+# Owner(s): ["oncall: jit"]
+
+import torch
+import torch._lazy.metrics as metrics
+import torch._lazy.ts_backend
+from torch.testing._internal.common_utils import run_tests, skipIfTorchDynamo, TestCase
+
+torch._lazy.ts_backend.init()
+
+
+class LazyGeneratorTest(TestCase):
+    def test_generator(self):
+        """
+        Test that generators are being inserted into the TorchScript
+        graph by setting different seeds before each call to
+        generate_tensor but the resulting tensor is the same
+        """
+
+        def generate_tensor():
+            g1 = torch.Generator()
+            g1.manual_seed(2023)
+            t1 = torch.tensor(1.0)
+            t1.uniform_(generator=g1)
+
+            g2 = torch.Generator()
+            g2.manual_seed(2024)
+            t2 = torch.tensor(1.0)
+            t2.normal_(generator=g2)
+
+            return t1, t2
+
+        torch.manual_seed(1)
+
+        with torch.device("cpu"):
+            cpu_t1, cpu_t2 = generate_tensor()
+
+        torch.manual_seed(2)
+
+        with torch.device("lazy"):
+            lazy_t1, lazy_t2 = generate_tensor()
+
+        torch._lazy.mark_step()
+
+        assert torch.allclose(
+            cpu_t1, lazy_t1.to("cpu")
+        ), f"Expected {cpu_t1}, got {lazy_t1.to('cpu')}"
+        assert torch.allclose(
+            cpu_t2, lazy_t2.to("cpu")
+        ), f"Expected {cpu_t2}, got {lazy_t2.to('cpu')}"
+
+    @skipIfTorchDynamo("Torch Dynamo does not support torch.Generator type")
+    def test_generator_causes_multiple_compiles(self):
+        """
+        Test that inserting generators with different seed caused recompile
+        """
+
+        def generate_tensor(seed):
+            t = torch.tensor(1.0)
+            g = torch.Generator()
+            g.manual_seed(seed)
+            t.uniform_(-1, 1, generator=g)
+            return t
+
+        metrics.reset()
+
+        with torch.device("lazy"):
+            t = generate_tensor(1)
+            torch._lazy.mark_step()
+
+            uncached_compile = metrics.counter_value("UncachedCompile")
+            assert (
+                uncached_compile == 1
+            ), f"Expected 1 uncached compiles, got {uncached_compile}"
+
+            t = generate_tensor(2)
+            torch._lazy.mark_step()
+
+            uncached_compile = metrics.counter_value("UncachedCompile")
+            assert (
+                uncached_compile == 2
+            ), f"Expected 2 uncached compiles, got {uncached_compile}"
+
+            t = generate_tensor(1)
+            torch._lazy.mark_step()
+
+            uncached_compile = metrics.counter_value("UncachedCompile")
+            assert (
+                uncached_compile == 2
+            ), f"Expected 2 uncached compiles, got {uncached_compile}"
+            cached_compile = metrics.counter_value("CachedCompile")
+            assert (
+                cached_compile == 1
+            ), f"Expected 1 cached compile, got {cached_compile}"
+
+        metrics.reset()
+
+        latest_graph = torch._C._lazy_ts_backend._get_latest_computation_graph()
+        assert 'torch.Generator(device="cpu", seed=1)' in latest_graph
+        assert "aten::uniform" in latest_graph
+
+
+if __name__ == "__main__":
+    run_tests()

test/lazy/test_ts_opinfo.py

@@ -231,13 +231,17 @@ def assert_allclose_rec(t):
 
         samples = op.sample_inputs("lazy", dtype, requires_grad=False)
         for sample in samples:
+            # Need to run mark step so that all random ops are computed in the right order
+            torch._lazy.mark_step()
+
             args = [sample.input] + list(sample.args)
             kwargs = sample.kwargs
             copy_args = clone_to_device(args, test_device)
 
             r_exp = op(*copy_args, **kwargs)
             r_actual = op(*args, **kwargs)
 
+            torch._lazy.mark_step()
             assert_allclose_rec((r_actual, r_exp))
 
     @ops([op for op in op_db if op.name in LAZY_OPS_LIST and op.name not in SKIP_RUNTIME_ERROR_LIST | SKIP_INCORRECT_RESULTS_LIST], allowed_dtypes=(torch.float,))  # noqa: B950
@@ -263,6 +267,9 @@ def assert_allclose_rec(t):
 
         samples = op.sample_inputs("lazy", dtype, requires_grad=False)
         for sample in samples:
+            # Need to run mark step so that all random ops are computed in the right order
+            torch._lazy.mark_step()
+
             args = [sample.input] + list(sample.args)
             kwargs = sample.kwargs
             copy_args = clone_to_device(args, test_device)