Back out "Reland "Make adding buffers more like adding parameters (py…

…torch#104069)" (pytorch#106224)" (pytorch#106743)

Summary:
Original commit changeset: 81319beb97f3

Original Phabricator Diff: D47961182

Test Plan: revert to maintain backward compat with legacy ads_dper3 production package. Read details in: S357822

Reviewed By: atuljangra

Differential Revision: D48131623

@diff-train-skip-merge
(D48131623 landed internally)
Pull Request resolved: pytorch#106743
Approved by: https://github.com/malfet

docs/source/nn.rst

@@ -22,7 +22,6 @@ These are the basic building blocks for graphs:
     :nosignatures:
     :template: classtemplate.rst
 
-    ~parameter.Buffer
     ~parameter.Parameter
     ~parameter.UninitializedParameter
     ~parameter.UninitializedBuffer

test/distributed/fsdp/test_fsdp_flatten_params.py

@@ -58,7 +58,7 @@ def _get_transformer(self, seed=0):
             dim_feedforward=128,
             dropout=0.1,
         )
-        module.dummy_buffer = nn.Buffer(torch.tensor(1.0))
+        module.register_buffer("dummy_buffer", torch.tensor(1.0))
 
         def get_input(device, dtype):
             torch.manual_seed(1)  # keep everything deterministic

test/distributed/fsdp/test_fsdp_misc.py

@@ -455,11 +455,11 @@ def init_nested_wrapped_module():
 
         # Check that `device_id` with `sync_module_states=True` works
         nested_wrapped_module = init_nested_wrapped_module()
-        nested_wrapped_module.buf = nn.Buffer(
-            torch.ones((2, 2), device="cpu") * self.rank
+        nested_wrapped_module.register_buffer(
+            "buf", torch.ones((2, 2), device="cpu") * self.rank
         )
-        nested_wrapped_module.module[0].buf = nn.Buffer(
-            torch.ones((3, 2), device="cpu") * self.rank
+        nested_wrapped_module.module[0].register_buffer(
+            "buf", torch.ones((3, 2), device="cpu") * self.rank
         )
         nested_wrapped_module = FSDP(
             nested_wrapped_module,
@@ -707,7 +707,7 @@ def __init__(self, rank):
                 torch.manual_seed(rank)
                 torch.cuda.manual_seed(rank)
                 self.lin = nn.Linear(10, 10, bias=False)
-                self.buffer = nn.Buffer(torch.ones(1) * rank)
+                self.register_buffer("buffer", torch.ones(1) * rank)
 
         m = MyModel(self.rank).cuda()
         _assert_module_states(

test/distributed/fsdp/test_fsdp_state_dict.py

@@ -87,7 +87,7 @@ def __init__(self, wrap_fsdp, register_buffers=False, ignore_inner=False):
         super().__init__()
         self.inner = Linear(*INNER_SHAPE)
         if register_buffers:
-            self.inner.buffer = nn.Buffer(torch.randn(BUFFER_SHAPE))
+            self.inner.register_buffer("buffer", torch.randn(BUFFER_SHAPE))
             self.inner.register_buffer(
                 "non_persistent_buffer", torch.randn(BUFFER_SHAPE), persistent=False
             )
@@ -97,7 +97,7 @@ def __init__(self, wrap_fsdp, register_buffers=False, ignore_inner=False):
             )
         self.outer = Linear(*OUTER_SHAPE)
         if register_buffers:
-            self.outer.buffer = nn.Buffer(torch.randn(BUFFER_SHAPE))
+            self.outer.register_buffer("buffer", torch.randn(BUFFER_SHAPE))
             self.outer.register_buffer(
                 "non_persistent_buffer", torch.randn(BUFFER_SHAPE), persistent=False
             )

test/distributed/fsdp/test_fsdp_unshard_params.py

@@ -423,7 +423,7 @@ def _test_named_parameters_and_buffers(self, prefix: str, recurse: bool):
             CUDAInitMode.CUDA_BEFORE,
             deterministic=True,
         )
-        model.buffer = nn.Buffer(torch.ones(1))
+        model.register_buffer("buffer", torch.ones(1))
         # Wrap the top-level with FSDP since `named_parameters()` and
         # `named_buffers` will contain FSDP prefixes if called on a non-FSDP
         # root module
@@ -436,7 +436,7 @@ def _test_named_parameters_and_buffers(self, prefix: str, recurse: bool):
             ),
             self.process_group,
         )
-        fsdp_model.buffer = nn.Buffer(torch.ones(1))
+        fsdp_model.register_buffer("buffer", torch.ones(1))
         with FSDP.summon_full_params(fsdp_model):
             for call in ["named_parameters", "named_buffers"]:
                 for (n1, p1), (n2, p2) in itertools.zip_longest(

test/distributed/optim/test_zero_redundancy_optimizer.py

@@ -855,7 +855,8 @@ def test_local_optimizer_parity(
                 torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM),
                 torch.nn.Linear(HIDDEN_DIM, OUTPUT_DIM),
             ).to(self.device)
-            model.test_buffer = torch.nn.Buffer(
+            model.register_buffer(
+                "test_buffer",
                 torch.ones((1), device=self.device) * self.rank,
             )
             # Define models/optimizers for DDP with ZeRO and DDP with local

test/distributed/test_data_parallel.py

@@ -34,8 +34,8 @@ def test_data_parallel_buffers_requiring_grad(self):
         class TestModule(nn.Module):
             def __init__(self, t):
                 super().__init__()
-                self.t_rg = nn.Buffer(t, t.requires_grad)
-                self.t_not_rg = nn.Buffer(t.clone().detach())
+                self.register_buffer('t_rg', t)
+                self.register_buffer('t_not_rg', t.clone().detach())
 
             def forward(self, x):
                 return x * self.t_rg + self.t_not_rg

test/distributed/test_dynamo_distributed.py

@@ -760,8 +760,9 @@ class DuplicateModule(nn.Module):
             def __init__(self) -> None:
                 super().__init__()
                 self._param = torch.randn((3,), device="cuda")
-                self._buf = torch.nn.Buffer(
-                    torch.randn((3,), requires_grad=False, device="cuda"))
+                self.register_buffer(
+                    "_buf", torch.randn((3,), requires_grad=False, device="cuda")
+                )
 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
                 # Use `_param` and `_buf` each twice in this compiled forward
@@ -788,8 +789,8 @@ def test_fsdp_dup_tensors_diff_source(self):
         class BufModule(nn.Module):
             def __init__(self) -> None:
                 super().__init__()
-                self._buf = nn.Buffer(
-                    torch.randn((3,), requires_grad=False, device="cuda")
+                self.register_buffer(
+                    "_buf", torch.randn((3,), requires_grad=False, device="cuda")
                 )
 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -801,7 +802,7 @@ def __init__(self) -> None:
                 self._param = nn.Parameter(torch.randn((1,), device="cuda"))
                 self._buf_module = BufModule()
                 # Share the buffer, meaning same tensor but different source
-                self._buf = self._buf_module._buf
+                self.register_buffer("_buf", self._buf_module._buf)
 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
                 # Use the same buffer tensor twice in the compiled forward,

test/dynamo/test_export.py

@@ -978,7 +978,7 @@ class MyBlock(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.weight = torch.nn.Parameter(torch.ones(1, 1))
-                self.buffer = torch.nn.Buffer(torch.ones(1, 1))
+                self.register_buffer("buffer", torch.ones(1, 1))
 
             def forward(self, x):
                 x = torch.nn.functional.linear(x, torch.randn(4, 4))
@@ -2670,7 +2670,7 @@ def test_not_functionalize(self):
         class Foo(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.buffer1 = torch.nn.Buffer(torch.ones(6, 2))
+                self.register_buffer("buffer1", torch.ones(6, 2))
 
             def forward(self, x):
                 x.add_(2)

test/dynamo/test_modules.py

@@ -1455,7 +1455,7 @@ def __init__(self):
         super().__init__()
         self.relu = torch.nn.ReLU()
         self.linear = torch.nn.Linear(10, 10)
-        self.buf0 = torch.nn.Buffer(torch.randn(10, 10))
+        self.register_buffer("buf0", torch.randn(10, 10))
 
     def forward(self, x):
         return self.relu(self.linear(x) + self.buf0)
@@ -1504,7 +1504,7 @@ class MockModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.linear = torch.nn.Linear(10, 10)
-                self.buf0 = torch.nn.Buffer(torch.randn(10, 10))
+                self.register_buffer("buf0", torch.randn(10, 10))
 
             def forward(self, x):
                 return self.r(torch.sin(x)) + self.buf0
@@ -1531,7 +1531,7 @@ class MockModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.linear = torch.nn.Linear(10, 10)
-                self.register_buffer("buf0", torch.nn.Buffer(torch.randn(10, 10)))
+                self.register_buffer("buf0", torch.randn(10, 10))
                 self.register_parameter(
                     name="param0", param=torch.nn.Parameter(torch.randn(10, 10))
                 )

test/dynamo/test_repros.py

@@ -1859,8 +1859,8 @@ def test_sort_out2(self):
         class MyModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.sorted = torch.nn.Buffer(torch.ones(4, 4))
-                self.indices = torch.nn.Buffer(torch.ones(4, 4, dtype=torch.long))
+                self.register_buffer("sorted", torch.ones(4, 4))
+                self.register_buffer("indices", torch.ones(4, 4, dtype=torch.long))
 
             def forward(self, x):
                 torch.sort(x, out=(self.sorted, self.indices))
@@ -1891,7 +1891,7 @@ def test_sigmoid_out2(self):
         class MyModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.base = torch.nn.Buffer(torch.ones(4, 4))
+                self.register_buffer("base", torch.ones(4, 4))
 
             def forward(self, x):
                 torch.sigmoid(x, out=self.base)
@@ -2174,8 +2174,8 @@ def test_named_buffers(self):
         class Foo(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.x = torch.nn.Buffer(torch.ones(3))
-                self.y = torch.nn.Buffer(torch.ones(3))
+                self.register_buffer("x", torch.ones(3))
+                self.register_buffer("y", torch.ones(3))
 
             def forward(self, inp):
                 res = 0

test/export/test_verifier.py

@@ -175,7 +175,8 @@ def test_aten_wrong_mem_format_buffer(self) -> None:
         class TestModel(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.a = torch.nn.Buffer(
+                self.register_buffer(
+                    "a",
                     torch.randn(1, 3, 100, 100).to(memory_format=torch.channels_last),
                 )
 

test/functorch/test_aotdispatch.py

@@ -1962,7 +1962,7 @@ def test_real_weights_in_symbolic_mode_with_inplace_ops(self):
         class M(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.buffer = torch.nn.Buffer(torch.ones(4, 5))
+                self.register_buffer("buffer", torch.ones(4, 5))
 
             def forward(self, x):
                 y = self.buffer.add_(3)
@@ -2157,7 +2157,7 @@ def test_aot_export_forward_mutation_no_buffer_mut_banned(self):
         class M(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.buffer1 = torch.nn.Buffer(torch.ones(6, 4))
+                self.register_buffer("buffer1", torch.ones(6, 4))
 
             def forward(self, x):
                 x.add_(4)
@@ -2170,7 +2170,7 @@ def test_aot_export_forward_mutation_multiple_mut_banned(self):
         class M(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.buffer1 = torch.nn.Buffer(torch.ones(6, 4))
+                self.register_buffer("buffer1", torch.ones(6, 4))
 
             def forward(self, x, y):
                 y.add_(4)

test/functorch/test_eager_transforms.py

@@ -3466,8 +3466,8 @@ def __init__(self):
                 super().__init__()
                 self.bias = nn.Parameter(torch.randn(3))
                 self.linear = nn.Linear(3, 3)
-                self.buffer = nn.Buffer(torch.randn(3))
-                self.buffer_tied = self.buffer
+                self.register_buffer('buffer', torch.randn(3))
+                self.register_buffer('buffer_tied', self.buffer)
 
             def forward(self, x):
                 x = self.linear(x)
@@ -3497,7 +3497,7 @@ class Foo(nn.Module):
             def __init__(self):
                 super().__init__()
                 self.linear = nn.Linear(3, 3)
-                self.buffer = nn.Buffer(torch.randn(3))
+                self.register_buffer('buffer', torch.randn(3))
 
             def forward(self, x):
                 x = self.linear(x)
@@ -3517,7 +3517,7 @@ class Foo(nn.Module):
             def __init__(self):
                 super().__init__()
                 self.linear = nn.Linear(3, 3)
-                self.buffer = nn.Buffer(torch.randn(3))
+                self.register_buffer('buffer', torch.randn(3))
 
             def forward(self, x):
                 x = self.linear(x)
@@ -3573,8 +3573,8 @@ def __init__(self):
                 self.linear = nn.Linear(3, 3)
                 self.weight = self.linear.weight
                 self.bias = self.linear.bias
-                self.buffer = nn.Buffer(torch.randn(3))
-                self.buffer_tied = self.buffer
+                self.register_buffer('buffer', torch.randn(3))
+                self.register_buffer('buffer_tied', self.buffer)
 
             def forward(self, x):
                 x = self.linear(x)

test/inductor/test_cuda_repro.py

@@ -558,7 +558,8 @@ def __init__(self):
                 start = math.log2(0.5)
                 end = math.log2(1 / (2**8))
 
-                self.scales = nn.Buffer(
+                self.register_buffer(
+                    "scales",
                     2
                     ** torch.arange(
                         start,

test/inductor/test_torchinductor.py

@@ -6275,7 +6275,9 @@ def fn(x, p1, p0):
         class Repro(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self._tensor_constant0 = nn.Buffer(torch.randn([], dtype=torch.float32))
+                self.register_buffer(
+                    "_tensor_constant0", torch.randn([], dtype=torch.float32)
+                )
 
             def forward(self, arg0_1, arg1_1):
                 convert_element_type = torch.ops.prims.convert_element_type.default(

test/jit/test_save_load.py

@@ -487,7 +487,6 @@ def __init__(self):
 
                 self.parameter_b = torch.nn.Parameter(torch.randn(4))
                 self.submodule_b = Submodule()
-                self.buffer_b = torch.nn.Buffer(torch.randn(4))
 
         m = TestModule()
         m_loaded = self.getExportImportCopy(torch.jit.script(m))
@@ -527,7 +526,7 @@ def __init__(self):
                 super().__init__()
                 self.foo = torch.nn.Linear(2, 3, device="meta")
                 self.bar = torch.nn.Linear(3, 4)
-                self.buffer = torch.nn.Buffer(torch.randn(4, device="meta"))
+                self.register_buffer("buffer", torch.randn(4, device="meta"))
 
             def forward(self, x):
                 x = self.foo(x)
@@ -1151,7 +1150,6 @@ def __init__(self):
 
                 self.parameter_b = torch.nn.Parameter(torch.randn(4))
                 self.submodule_b = Submodule()
-                self.buffer_b = torch.nn.Buffer(torch.randn(4))
 
         m = TestModule()
         m_loaded = self.getExportImportCopy(torch.jit.script(m))

test/nn/test_lazy_modules.py

@@ -5,7 +5,7 @@
 import torch
 import torch.nn as nn
 from torch.nn.parameter import UninitializedParameter, UninitializedBuffer
-from torch.nn import Buffer, Parameter
+from torch.nn import Parameter
 from torch.testing._internal.common_utils import TestCase, run_tests, suppress_warnings
 from torch.testing._internal.common_cuda import TEST_CUDA
 
@@ -47,29 +47,29 @@ def test_lazy_module_parameter(self):
     @suppress_warnings
     def test_lazy_module_buffer(self):
         module = LazyModule()
-        module.test_buffer = UninitializedBuffer()
+        module.register_buffer('test_buffer', UninitializedBuffer())
         self.assertTrue(module.has_uninitialized_params())
         state_dict = module.state_dict()
         self.assertIsInstance(state_dict['test_buffer'], UninitializedBuffer)
         new_module = LazyModule()
         # An error is raised when there is an attempt to replace an existing parameter
         # with an uninitialized one
-        new_module.test_buffer = Buffer(torch.ones(5, 5))
+        new_module.register_buffer('test_buffer', torch.ones(5, 5))
         with self.assertRaisesRegex(RuntimeError, 'shape of an uninitialized'):
             new_module.load_state_dict(state_dict)
         # Uninitialized parameters are overriden when the state dict to be loaded contains a valid one
         new_module = LazyModule()
-        new_module.test_buffer = Buffer(torch.ones(5, 5))
+        new_module.register_buffer('test_buffer', torch.ones(5, 5))
         module.load_state_dict(new_module.state_dict())
         self.assertEqual(module.test_buffer, torch.ones((5, 5)))
 
         # Uninitialized parameters are left unchanged
         module = LazyModule()
-        module.test_buffer = UninitializedBuffer()
+        module.register_buffer('test_buffer', UninitializedBuffer())
         self.assertTrue(module.has_uninitialized_params())
 
         new_module = LazyModule()
-        new_module.test_buffer = UninitializedBuffer()
+        new_module.register_buffer('test_buffer', UninitializedBuffer())
         module.load_state_dict(new_module.state_dict())
         module.load_state_dict(new_module.state_dict())
         self.assertTrue(module.has_uninitialized_params())
@@ -85,7 +85,7 @@ def test_lazy_module_jit_param(self):
     @suppress_warnings
     def test_lazy_module_jit_buffer(self):
         module = LazyModule()
-        module.test_buffer = UninitializedBuffer()
+        module.register_buffer('test_buffer', UninitializedBuffer())
         self.assertTrue(module.has_uninitialized_params())
         with self.assertRaisesRegex(RuntimeError, 'run a forward pass'):
             torch.jit.script(module)
@@ -101,7 +101,7 @@ def test_lazy_share_memory_param(self):
     @suppress_warnings
     def test_lazy_share_memory_buffer(self):
         module = LazyModule()
-        module.test_buffer = UninitializedBuffer()
+        module.register_buffer('test_buffer', UninitializedBuffer())
         self.assertTrue(module.has_uninitialized_params())
         with self.assertRaisesRegex(RuntimeError, 'share memory on an uninitialized'):
             module.share_memory()