chore: import upstream snapshot with attribution

python/paddle/distributed/fleet/meta_parallel/zero_bubble_utils.py

@@ -0,0 +1,133 @@
+# 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.
+
+# The file has been adapted from DeepSeek DualPipe project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DualPipe/blob/main/LICENSE
+
+
+import queue
+from functools import partial
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.autograd import PyLayer
+
+
+class WeightGradStore:
+    enabled = False
+    cache = []
+    funcs_queue = queue.Queue()
+
+    @classmethod
+    def put(cls, func) -> None:
+        cls.cache.append(func)
+
+    @classmethod
+    def flush(cls) -> None:
+        cls.funcs_queue.put(cls.cache)
+        cls.cache = []
+
+    @classmethod
+    def pop(cls) -> None:
+        assert not cls.funcs_queue.empty(), "Pop empty queue."
+        funcs = cls.funcs_queue.get()
+        for func in funcs:
+            func()
+
+    @classmethod
+    def clear(cls) -> None:
+        cls.cache = []
+        cls.funcs_queue = queue.Queue()
+
+
+class EventStore:
+    event = None
+
+    @classmethod
+    def set(cls, event) -> None:
+        cls.event = event
+
+
+def fold_init_dims(tensor):
+    # NOTE(zhangyuqin1998): Reshape a rank-3 tensor from P x M x N to (P * M) x N,
+    # to keep weight_grad in a correct rank. See phi::FoldInitDims.
+    if tensor.ndim == 3:
+        tensor = paddle.reshape(tensor, [-1, tensor.shape[-1]])
+    return tensor
+
+
+def grad_weight_fn(input, weight, out_grad, inplace_update_grad=True):
+    if weight.stop_gradient:
+        return
+    with paddle.no_grad():
+        weight_grad = paddle.matmul(
+            x=fold_init_dims(input),
+            y=fold_init_dims(out_grad),
+            transpose_x=True,
+            transpose_y=False,
+        )
+
+        if hasattr(weight, "main_grad"):
+            if weight.main_grad is None:
+                weight.main_grad = paddle.base.framework.core.eager.Tensor(
+                    value=weight_grad.cast(paddle.float32).value(),
+                    place=weight_grad.place,
+                    name="main_grad@" + weight.name,
+                )
+            else:
+                weight.main_grad.add_(weight_grad)
+            weight_grad._clear_data()
+        else:
+            if weight.grad is None:
+                weight.grad = paddle.zeros_like(weight, dtype=weight.dtype)
+            weight.grad = paddle.add(weight.grad, weight_grad)
+
+
+class SplitBWMatmul(PyLayer):
+    @staticmethod
+    def forward(ctx, input, weight, bias):
+        ctx.save_for_backward(input, weight, bias)
+        out = F.linear(x=input, weight=weight, bias=bias)
+        return out
+
+    @staticmethod
+    def backward(ctx, out_grad):
+        input, weight, bias = ctx.saved_tensor()
+
+        if WeightGradStore.enabled:
+            WeightGradStore.put(
+                partial(grad_weight_fn, input, weight, out_grad)
+            )
+        else:
+            grad_weight_fn(input, weight, out_grad)
+
+        input_grad = None
+        if not input.stop_gradient:
+            input_grad = paddle.matmul(
+                x=out_grad, y=weight, transpose_x=False, transpose_y=True
+            )
+        if bias is not None:
+            bias_grad = None
+            if not bias.stop_gradient:
+                bias_grad = paddle.sum(fold_init_dims(out_grad), axis=0)
+            return input_grad, None, bias_grad
+        else:
+            return input_grad, None
+
+
+class SplitBWLinear(nn.Linear):
+    def forward(self, input):
+        return SplitBWMatmul.apply(input, self.weight, bias=self.bias)