chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import torch.nn as nn
+
+from deepspeed import module_inject
+from .diffusers_attention import DeepSpeedDiffusersAttention
+from .bias_add import nhwc_bias_add
+from .diffusers_2d_transformer import Diffusers2DTransformerConfig
+from deepspeed.utils.types import ActivationFuncType
+from .op_binding.gated_activation import GatedActivationOp
+from .op_binding.layer_norm import LayerNormOp
+
+
+class DeepSpeedDiffusersTransformerBlock(nn.Module):
+
+    def __init__(self, equivalent_module: nn.Module, config: Diffusers2DTransformerConfig):
+        super(DeepSpeedDiffusersTransformerBlock, self).__init__()
+        self.quantizer = module_inject.GroupQuantizer(q_int8=config.int8_quantization)
+        # Ensure ops are built by the time we start running
+        self.config = config
+
+        self.ff1_w = self.quantizer.quantize(
+            nn.Parameter(equivalent_module.ff.net[0].proj.weight.data, requires_grad=False))
+        self.ff1_b = nn.Parameter(equivalent_module.ff.net[0].proj.bias.data, requires_grad=False)
+        self.ff2_w = self.quantizer.quantize(nn.Parameter(equivalent_module.ff.net[2].weight.data,
+                                                          requires_grad=False))
+        self.ff2_b = nn.Parameter(equivalent_module.ff.net[2].bias.data, requires_grad=False)
+
+        self.norm1_g = nn.Parameter(equivalent_module.norm1.weight.data, requires_grad=False)
+        self.norm1_b = nn.Parameter(equivalent_module.norm1.bias.data, requires_grad=False)
+        self.norm1_eps = equivalent_module.norm1.eps
+
+        self.norm2_g = nn.Parameter(equivalent_module.norm2.weight.data, requires_grad=False)
+        self.norm2_b = nn.Parameter(equivalent_module.norm2.bias.data, requires_grad=False)
+        self.norm2_eps = equivalent_module.norm2.eps
+
+        self.norm3_g = nn.Parameter(equivalent_module.norm3.weight.data, requires_grad=False)
+        self.norm3_b = nn.Parameter(equivalent_module.norm3.bias.data, requires_grad=False)
+        self.norm3_eps = equivalent_module.norm3.eps
+
+        self.attn_1 = equivalent_module.attn1
+        self.attn_2 = equivalent_module.attn2
+
+        # Pull the bias in if we can
+        if isinstance(self.attn_1, DeepSpeedDiffusersAttention):
+            self.attn_1.do_out_bias = False
+            self.attn_1_bias = self.attn_1.attn_ob
+        else:
+            self.attn_1_bias = nn.Parameter(torch.zeros_like(self.norm2_g), requires_grad=False)
+
+        # Pull the bias in if we can
+        if isinstance(self.attn_2, DeepSpeedDiffusersAttention):
+            self.attn_2.do_out_bias = False
+            self.attn_2_bias = self.attn_2.attn_ob
+        else:
+            self.attn_2_bias = nn.Parameter(torch.zeros_like(self.norm3_g), requires_grad=False)
+
+        self.gated_activation = GatedActivationOp()
+        self.layer_norm = LayerNormOp()
+
+    def forward(self, hidden_states, context=None, timestep=None, **kwargs):
+        # In v0.12.0 of diffuser, several new kwargs were added. Capturing
+        # those with kwargs to maintain backward compatibility
+
+        # In v0.11.0 of diffusers, the kwarg was changed from 'context' to 'encoder_hidden_states'
+        # This is so we can support older and newer versions of diffusers
+        if "encoder_hidden_states" in kwargs and kwargs["encoder_hidden_states"] is not None:
+            context = kwargs["encoder_hidden_states"]
+
+        out_norm_1 = self.layer_norm(hidden_states, self.norm1_g, self.norm1_b, self.norm1_eps)
+        out_attn_1 = self.attn_1(out_norm_1)
+
+        out_norm_2, out_attn_1 = self.layer_norm.layer_norm_residual_store_pre_ln_res(
+            out_attn_1, self.attn_1_bias, hidden_states, self.norm2_g, self.norm2_b, self.norm2_eps)
+        out_attn_2 = self.attn_2(out_norm_2, context=context)
+        out_norm_3, out_attn_2 = self.layer_norm.layer_norm_residual_store_pre_ln_res(
+            out_attn_2, self.attn_2_bias, out_attn_1, self.norm3_g, self.norm3_b, self.norm3_eps)
+
+        out_ff1 = nn.functional.linear(out_norm_3, self.ff1_w)
+        out_geglu = self.gated_activation(out_ff1, self.ff1_b, ActivationFuncType.GATED_GELU)
+
+        out_ff2 = nn.functional.linear(out_geglu, self.ff2_w)
+        return nhwc_bias_add(out_ff2, self.ff2_b, other=out_attn_2)