Flux performance optimization (mindspore-lab#966)

* flux performance optimization

* add comment

* change default value

* change the default value

* fix ci

* fix typo

Author: alien-0119

examples/diffusers/controlnet/train_controlnet_flux.py

@@ -49,6 +49,7 @@
     set_seed,
 )
 from mindone.transformers import CLIPTextModel, T5EncoderModel
+from mindone.utils.config import str2bool
 
 logger = logging.getLogger(__name__)
 
@@ -419,6 +420,12 @@ def parse_args(input_args=None):
             " https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html"
         ),
     )
+    parser.add_argument(
+        "--dataset_iterator_no_copy",
+        default=True,
+        type=str2bool,
+        help="dataset iterator optimization strategy. Whether dataset iterator creates a Tensor without copy.",
+    )
     parser.add_argument(
         "--dataset_name",
         type=str,
@@ -1165,7 +1172,13 @@ def __len__(self):
         # Only show the progress bar once on each machine.
         disable=not is_master(args),
     )
-    train_dataloader_iter = train_dataloader.create_tuple_iterator(num_epochs=args.num_train_epochs - first_epoch)
+    # do_copy=False enables the dataset iterator to not do copy when creating a tensor which takes less time.
+    # Currently the default value of do_copy is True,
+    # it is expected that the default value of do_copy will be changed to False in MindSpore 2.7.0.
+    train_dataloader_iter = train_dataloader.create_tuple_iterator(
+        num_epochs=args.num_train_epochs - first_epoch,
+        do_copy=not args.dataset_iterator_no_copy,
+    )
 
     for epoch in range(first_epoch, args.num_train_epochs):
         flux_controlnet.set_train(True)

examples/diffusers/dreambooth/train_dreambooth_lora_flux.py

@@ -201,6 +201,20 @@ def parse_args(input_args=None):
         required=False,
         help="A folder containing the training data of class images.",
     )
+    parser.add_argument(
+        "--jit_level",
+        type=str,
+        default="O1",
+        choices=["O0", "O1", "O2"],
+        help=(
+            "Used to control the compilation optimization level, supports [O0, O1, O2]. The framework automatically "
+            "selects the execution method. O0: All optimizations except those necessary for functionality are "
+            "disabled, using an operator-by-operator execution method. O1: Enables common optimizations and automatic "
+            "operator fusion optimizations, using an operator-by-operator execution method. This is an experimental "
+            "optimization level, which is continuously being improved. O2: Enables extreme performance optimization, "
+            "using a sinking execution method."
+        ),
+    )
     parser.add_argument(
         "--instance_prompt",
         type=str,
@@ -908,7 +922,11 @@ def encode_prompt(
 
 def main(args):
     args = parse_args()
-    ms.set_context(mode=ms.GRAPH_MODE, jit_syntax_level=ms.STRICT)
+    ms.set_context(
+        mode=ms.GRAPH_MODE,
+        jit_syntax_level=ms.STRICT,
+        jit_config={"jit_level": args.jit_level},
+    )
     init_distributed_device(args)
 
     logging_dir = Path(args.output_dir, args.logging_dir)

mindone/diffusers/models/embeddings.py

@@ -1237,7 +1237,7 @@ def construct(self, ids: ms.Tensor) -> ms.Tensor:
         cos_out = []
         sin_out = []
         pos = ids.float()
-        freqs_dtype = ms.float64
+        freqs_dtype = ms.float32
         for i in range(n_axes):
             cos, sin = get_1d_rotary_pos_embed(
                 self.axes_dim[i],