ADD SD3 batch_parallel

ADD SD3 batch_parallel

Author: nemonameless

paddlemix/triton_ops/triton_ops.py

@@ -1711,9 +1711,18 @@ def split_concat(x, y):
         out1 = paddle.empty(shape=[batch, seq_qkv + seq_eqkv, ouput_hidden], dtype=x.dtype)
         out2 = paddle.empty(shape=[batch, seq_qkv + seq_eqkv, ouput_hidden], dtype=x.dtype)
         grid = ("3", "batch", "seq_qkv + seq_eqkv")
-
+        # -1 means this value does not matter for triton compilation
         split_concat_kernel[(op_name, grid)](
-            out0, out1, out2, x, y, batch, seq_qkv, seq_eqkv, ouput_hidden, BLOCK_SIZE=BLOCK_SIZE
+            out0,
+            out1, 
+            out2, 
+            x, 
+            y,
+            -1, # batch,
+            seq_qkv, 
+            seq_eqkv, 
+            ouput_hidden, 
+            BLOCK_SIZE=BLOCK_SIZE
         )
 
     if in_dynamic_or_pir_mode():

ppdiffusers/deploy/sd3/README.md

@@ -29,3 +29,29 @@ python  text_to_image_generation-stable_diffusion_3.py  --dtype float16 --height
 | Paddle Inference|    PyTorch   | Paddle 动态图 |
 | --------------- | ------------ | ------------ |
 |       1.2 s     |     1.78 s   |    4.202 s   |
+
+
+## Paddle Stable Diffusion 3 模型多卡推理： 
+### batch parallel 实现原理  
+- 在SD3中，对于输入是一个prompt时，使用CFG需要同时进行unconditional guide和text guide的生成，此时 MM-DiT-blocks 的输入batch_size=2；  
+所以我们考虑在多卡并行的方案中，将batch为2的输入拆分到两张卡上进行计算，这样单卡的计算量就减少为原来的一半，降低了单卡所承载的浮点计算量。  
+计算完成后，我们再把两张卡的计算结果 聚合在一起，结果与单卡计算完全一致。  
+### 开启多卡推理方法 
+- Paddle Inference 提供了SD3模型的多卡推理功能，用户可以通过设置 `--inference_optimize_bp 1` 来开启这一功能，  
+使用 `python -m paddle.distributed.launch --gpus 0,1` 指定使用哪些卡进行推理。
+高性能多卡推理指令：
+```shell
+# 执行多卡推理指令
+python -m paddle.distributed.launch --gpus 0,1 text_to_image_generation-stable_diffusion_3.py \
+--dtype float16 \
+--height 512 --width 512 \
+--num-inference-steps 50 \
+--inference_optimize 1 \
+--inference_optimize_bp 1 \
+--benchmark 1
+```
+## 在 NVIDIA A800-SXM4-80GB 上测试的性能如下：
+
+| Paddle batch parallel | Paddle Single Card |  PyTorch  | Paddle 动态图 |
+| --------------------- | ------------------ | --------- | ------------ |
+|          0.86 s       |        1.2 s       |   1.78 s  |    4.202 s   |

ppdiffusers/examples/inference/text_to_image_generation-stable_diffusion_3.py

@@ -12,16 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-
-os.environ["FLAGS_use_cuda_managed_memory"] = "true"
 import argparse
-import datetime
-
 import paddle
-
-from ppdiffusers import StableDiffusion3Pipeline
-
-
 def parse_args():
     parser = argparse.ArgumentParser(
         description=" Use PaddleMIX to accelerate the Stable Diffusion3 image generation model."
@@ -30,13 +22,19 @@ def parse_args():
         "--benchmark",
         type=(lambda x: str(x).lower() in ["true", "1", "yes"]),
         default=False,
-        help="if benchmark is set to True, measure inference performance",
+        help="if set to True, measure inference performance",
     )
     parser.add_argument(
         "--inference_optimize",
         type=(lambda x: str(x).lower() in ["true", "1", "yes"]),
         default=False,
-        help="If inference_optimize is set to True, all optimizations except Triton are enabled.",
+        help="If set to True, all optimizations except Triton are enabled.",
+    )
+    parser.add_argument(
+        "--inference_optimize_bp",
+        type=(lambda x: str(x).lower() in ["true", "1", "yes"]),
+        default=False,
+        help="If set to True, batch parallel is enabled in DIT and dual-GPU acceleration is used.",
     )
     parser.add_argument("--height", type=int, default=512, help="Height of the generated image.")
     parser.add_argument("--width", type=int, default=512, help="Width of the generated image.")
@@ -51,11 +49,38 @@ def parse_args():
 if args.inference_optimize:
     os.environ["INFERENCE_OPTIMIZE"] = "True"
     os.environ["INFERENCE_OPTIMIZE_TRITON"] = "True"
+if args.inference_optimize_bp:
+    os.environ["INFERENCE_OPTIMIZE_BP"] = "True"
 if args.dtype == "float32":
     inference_dtype = paddle.float32
 elif args.dtype == "float16":
     inference_dtype = paddle.float16
 
+
+if args.inference_optimize_bp:
+    from paddle.distributed import fleet
+    from paddle.distributed.fleet.utils import recompute
+    import numpy as np
+    import random
+    import paddle.distributed as dist
+    import paddle.distributed.fleet as fleet
+    strategy = fleet.DistributedStrategy()
+    model_parallel_size = 2
+    data_parallel_size = 1
+    strategy.hybrid_configs = {
+    "dp_degree": data_parallel_size,
+    "mp_degree": model_parallel_size,
+    "pp_degree": 1
+    }
+    fleet.init(is_collective=True, strategy=strategy)
+    hcg = fleet.get_hybrid_communicate_group()
+    mp_id = hcg.get_model_parallel_rank()
+    rank_id = dist.get_rank()
+
+import datetime
+from ppdiffusers import StableDiffusion3Pipeline
+
+
 pipe = StableDiffusion3Pipeline.from_pretrained(
     "stabilityai/stable-diffusion-3-medium-diffusers",
     paddle_dtype=inference_dtype,
@@ -67,6 +92,7 @@ def parse_args():
     enable_new_ir=True,
     cache_static_model=True,
     exp_enable_use_cutlass=True,
+    delete_pass_lists=["add_norm_fuse_pass"],
 )
 
 generator = paddle.Generator().manual_seed(42)
@@ -111,4 +137,8 @@ def parse_args():
     cuda_mem_after_used = paddle.device.cuda.max_memory_allocated() / (1024**3)
     print(f"Max used CUDA memory : {cuda_mem_after_used:.3f} GiB")
 
-image.save("text_to_image_generation-stable_diffusion_3-result.png")
+if args.inference_optimize_bp:
+    if rank_id == 0:
+        image.save("text_to_image_generation-stable_diffusion_3-result.png")
+else:
+    image.save("text_to_image_generation-stable_diffusion_3-result.png")

ppdiffusers/ppdiffusers/models/simplified_sd3.py

@@ -106,12 +106,13 @@ def forward(self, hidden_states, encoder_hidden_states, temb):
             qkv = self.qkv[i](norm_hidden_states)
             eqkv = self.eqkv[i](norm_encoder_hidden_states)
             q, k, v = paddlemix.triton_ops.split_concat(qkv, eqkv)
-            q = q.reshape([2, -1, 24, 64])
-            k = k.reshape([2, -1, 24, 64])
-            v = v.reshape([2, -1, 24, 64])
+            bs = hidden_states.shape[0]
+            q = q.reshape([bs, -1, 24, 64])
+            k = k.reshape([bs, -1, 24, 64])
+            v = v.reshape([bs, -1, 24, 64])
 
             norm_hidden_states1 = F.scaled_dot_product_attention_(q, k, v, dropout_p=0.0, is_causal=False)
-            norm_hidden_states1 = norm_hidden_states1.reshape([2, -1, self.dim])
+            norm_hidden_states1 = norm_hidden_states1.reshape([bs, -1, self.dim])
             attn_output, context_attn_output = paddle.split(norm_hidden_states1, num_or_sections=[seq1, seq2], axis=1)
 
             # attn_output, context_attn_output = paddlemix.triton_ops.triton_split(
@@ -155,7 +156,5 @@ def forward(self, hidden_states, encoder_hidden_states, temb):
                 last_context_ffn_output = context_ffn_output
                 last_context_hidden_states = encoder_hidden_states
                 last_context_gate_mlp = c_gate_mlp
-            else:
-                encoder_hidden_states = None
 
-        return encoder_hidden_states, hidden_states
+        return  hidden_states

ppdiffusers/ppdiffusers/models/transformer_sd3.py

@@ -329,10 +329,9 @@ def forward(
         encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
         if self.inference_optimize:
-            out = self.simplified_sd3(
+            hidden_states = self.simplified_sd3(
                 hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb
             )
-            hidden_states = out[1]
             encoder_hidden_states = None
         else:
             encoder_hidden_states, hidden_states = self.sd3_origin_transformer(

ppdiffusers/ppdiffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -12,11 +12,12 @@
 # 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 os
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import paddle
+import paddle.distributed as dist
 
 from ppdiffusers.transformers import (  # T5TokenizerFast,
     CLIPTextModelWithProjection,
@@ -195,6 +196,7 @@ def __init__(
             if hasattr(self, "transformer") and self.transformer is not None
             else 128
         )
+        self.inference_optimize_bp = os.getenv("INFERENCE_OPTIMIZE_BP") == "True"
 
     def _get_t5_prompt_embeds(
         self,
@@ -229,7 +231,6 @@ def _get_t5_prompt_embeds(
                 "The following part of your input was truncated because CLIP can only handle sequences up to"
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
-        # breakpoint()
         prompt_embeds = self.text_encoder_3(text_input_ids)[0]
 
         dtype = self.text_encoder_3.dtype
@@ -395,7 +396,6 @@ def encode_prompt(
 
             prompt_embeds = paddle.concat([clip_prompt_embeds, t5_prompt_embed], axis=-2)
             pooled_prompt_embeds = paddle.concat([pooled_prompt_embed, pooled_prompt_2_embed], axis=-1)
-
         if do_classifier_free_guidance and negative_prompt_embeds is None:
             negative_prompt = negative_prompt or ""
             negative_prompt_2 = negative_prompt_2 or negative_prompt
@@ -707,7 +707,6 @@ def __call__(
                 The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
                 will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
                 `._callback_tensor_inputs` attribute of your pipeline class.
-
         Examples:
 
         Returns:
@@ -801,22 +800,47 @@ def __call__(
                 latent_model_input = paddle.concat([latents] * 2) if self.do_classifier_free_guidance else latents
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latent_model_input.shape[0])
+                if self.inference_optimize_bp and self.do_classifier_free_guidance:
+                    latent_input ,latent_model_input_ = paddle.split(latent_model_input,2,axis=0)
+                    timestep_input ,timestep_ = paddle.split(timestep,2,axis=0)
+                    prompt_embeds_input ,prompt_embeds_ = paddle.split(prompt_embeds,2,axis=0)
+                    pooled_prompt_embeds_input ,pooled_prompt_embeds_ = paddle.split(pooled_prompt_embeds,2,axis=0)
+                    
+                    dist.scatter(latent_input,[latent_input,latent_model_input_])
+                    dist.scatter(timestep_input,[timestep_input,timestep_])
+                    dist.scatter(prompt_embeds_input,[prompt_embeds_input,prompt_embeds_])
+                    dist.scatter(pooled_prompt_embeds_input,[pooled_prompt_embeds_input,pooled_prompt_embeds_])
 
+                else:
+                    latent_input = latent_model_input
+                    timestep_input = timestep
+                    prompt_embeds_input = prompt_embeds
+                    pooled_prompt_embeds_input = pooled_prompt_embeds
+                
                 model_output = self.transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep,
-                    encoder_hidden_states=prompt_embeds,
-                    pooled_projections=pooled_prompt_embeds,
+                    hidden_states=latent_input,
+                    timestep=timestep_input,
+                    encoder_hidden_states=prompt_embeds_input,
+                    pooled_projections=pooled_prompt_embeds_input,
                     joint_attention_kwargs=self.joint_attention_kwargs,
                     return_dict=False,
                 )
-
                 if is_inference_mode(self.transformer):
                     # NOTE:(changwenbin,zhoukangkang)
                     # This is for paddle inference mode
-                    noise_pred = model_output
+                    output = model_output
+                else:
+                    output = model_output[0]
+                    
+                if self.inference_optimize_bp:
+                    tmp_shape = output.shape
+                    tmp_shape[0] *=2
+                    noise_pred = paddle.zeros(tmp_shape,dtype=output.dtype)
+                    dist.all_gather(noise_pred,output)
                 else:
-                    noise_pred = model_output[0]
+                    noise_pred = output
+
+
 
                 # perform guidance
                 if self.do_classifier_free_guidance: