chore: import upstream snapshot with attribution

test/collective/fleet/test_zero_bubble_utils.py

@@ -0,0 +1,180 @@
+#   Copyright (c) 2025 PaddlePaddle 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.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.distributed.fleet.meta_parallel import zero_bubble_utils
+
+
+class TestZuroBubble(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(42)
+
+    def test_weight_grad_store(self):
+        def fake_func():
+            return
+
+        zero_bubble_utils.WeightGradStore.put(fake_func)
+        np.testing.assert_equal(
+            zero_bubble_utils.WeightGradStore.funcs_queue.empty(), True
+        )
+        zero_bubble_utils.WeightGradStore.flush()
+        np.testing.assert_equal(
+            zero_bubble_utils.WeightGradStore.funcs_queue.empty(), False
+        )
+        zero_bubble_utils.WeightGradStore.pop()
+        np.testing.assert_equal(
+            zero_bubble_utils.WeightGradStore.funcs_queue.empty(), True
+        )
+
+    def test_zero_bubble_utils(self):
+        zero_bubble_utils.WeightGradStore.enabled = False
+
+        paddle.seed(42)
+        input = paddle.randn([2, 4096, 2048])
+        input.stop_gradient = False
+        splitbw_linear = zero_bubble_utils.SplitBWLinear(
+            2048, 2048, bias_attr=True
+        )
+        o = splitbw_linear(input)
+        o.mean().backward()
+
+        paddle.seed(42)
+        ref_input = paddle.randn([2, 4096, 2048])
+        ref_input.stop_gradient = False
+        ref_linear = paddle.nn.Linear(2048, 2048, bias_attr=True)
+        o_ref = ref_linear(ref_input)
+        o_ref.mean().backward()
+
+        np.testing.assert_equal(o._md5sum(), o_ref._md5sum())
+        np.testing.assert_equal(input.grad._md5sum(), ref_input.grad._md5sum())
+        np.testing.assert_equal(
+            splitbw_linear.weight.grad._md5sum(),
+            ref_linear.weight.grad._md5sum(),
+        )
+        np.testing.assert_equal(
+            splitbw_linear.bias.grad._md5sum(), ref_linear.bias.grad._md5sum()
+        )
+
+        zero_bubble_utils.WeightGradStore.enabled = True
+
+        paddle.seed(42)
+        input = paddle.randn([2, 4096, 2048])
+        input.stop_gradient = False
+        splitbw_linear = zero_bubble_utils.SplitBWLinear(
+            2048, 2048, bias_attr=True
+        )
+        o = splitbw_linear(input)
+        o.mean().backward()
+        np.testing.assert_equal(splitbw_linear.weight.grad, None)
+        zero_bubble_utils.WeightGradStore.flush()
+        zero_bubble_utils.WeightGradStore.pop()
+        np.testing.assert_equal(
+            splitbw_linear.weight.grad._md5sum(),
+            ref_linear.weight.grad._md5sum(),
+        )
+
+    def test_zero_bubble_utils_no_bias(self):
+        zero_bubble_utils.WeightGradStore.enabled = True
+
+        paddle.seed(42)
+        ref_input = paddle.randn([2, 4096, 2048])
+        ref_input.stop_gradient = False
+        ref_linear = paddle.nn.Linear(2048, 2048, bias_attr=False)
+        o_ref = ref_linear(ref_input)
+        o_ref.mean().backward()
+
+        paddle.seed(42)
+        input = paddle.randn([2, 4096, 2048])
+        input.stop_gradient = False
+        splitbw_linear = zero_bubble_utils.SplitBWLinear(
+            2048, 2048, bias_attr=False
+        )
+        o = splitbw_linear(input)
+        o.mean().backward()
+
+        np.testing.assert_equal(splitbw_linear.weight.grad, None)
+
+        zero_bubble_utils.WeightGradStore.flush()
+        zero_bubble_utils.WeightGradStore.pop()
+
+        np.testing.assert_equal(o._md5sum(), o_ref._md5sum())
+        np.testing.assert_equal(input.grad._md5sum(), ref_input.grad._md5sum())
+        np.testing.assert_equal(
+            splitbw_linear.weight.grad._md5sum(),
+            ref_linear.weight.grad._md5sum(),
+        )
+
+    def test_zero_bubble_with_main_grad(self):
+        def _update_main_grad_hook(param):
+            @paddle.autograd.no_grad()
+            def param_hook(tmp_grad):
+                if tmp_grad is not None and tmp_grad._is_initialized():
+                    if param.main_grad is None:
+                        param.main_grad = (
+                            paddle.base.framework.core.eager.Tensor(
+                                value=tmp_grad.cast(paddle.float32).value(),
+                                place=tmp_grad.place,
+                                name="main_grad@" + param.name,
+                            )
+                        )
+                    else:
+                        # 梯度累加
+                        param.main_grad.add_(tmp_grad)
+
+                    tmp_grad._clear_data()
+
+            return param_hook
+
+        zero_bubble_utils.WeightGradStore.enabled = False
+
+        paddle.seed(42)
+        ref_input = paddle.randn([2, 4096, 2048])
+        ref_input.stop_gradient = False
+        ref_linear = paddle.nn.Linear(2048, 2048, bias_attr=False)
+
+        paddle.seed(42)
+        input = paddle.randn([2, 4096, 2048])
+        input.stop_gradient = False
+        splitbw_linear = zero_bubble_utils.SplitBWLinear(
+            2048, 2048, bias_attr=False
+        )
+
+        for param in splitbw_linear.parameters():
+            if not hasattr(param, "main_grad"):
+                param.main_grad = None
+                param._register_grad_hook(_update_main_grad_hook(param))
+        for param in ref_linear.parameters():
+            if not hasattr(param, "main_grad"):
+                param.main_grad = None
+                param._register_grad_hook(_update_main_grad_hook(param))
+
+        o_ref = ref_linear(ref_input)
+        o_ref.mean().backward()
+        o = splitbw_linear(input)
+        o.mean().backward()
+
+        np.testing.assert_equal(o._md5sum(), o_ref._md5sum())
+        np.testing.assert_equal(input.grad._md5sum(), ref_input.grad._md5sum())
+        np.testing.assert_equal(
+            splitbw_linear.weight.main_grad._md5sum(),
+            ref_linear.weight.main_grad._md5sum(),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()