Fixed issue kohya-ss#1002; add default ddp arguments

ngitnenlim · Dec 13, 2023 · dc056ab · dc056ab
1 parent c366f0a
commit dc056ab
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/sdxl_train.py b/sdxl_train.py
@@ -463,7 +463,9 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
 
     # For --sample_at_first
     sdxl_train_util.sample_images(
-        accelerator, args, 0, global_step, accelerator.device, vae, [tokenizer1, tokenizer2], [text_encoder1, text_encoder2], unet
+        accelerator, args, 0, global_step, accelerator.device, vae, [tokenizer1, tokenizer2],
+        [accelerator.unwrap_model(text_encoder1), accelerator.unwrap_model(text_encoder2)],
+        accelerator.unwrap_model(unet),
     )
 
     loss_recorder = train_util.LossRecorder()
@@ -611,8 +613,8 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
                     accelerator.device,
                     vae,
                     [tokenizer1, tokenizer2],
-                    [text_encoder1, text_encoder2],
-                    unet,
+                    [accelerator.unwrap_model(text_encoder1), accelerator.unwrap_model(text_encoder2)],
+                    accelerator.unwrap_model(unet),
                 )
 
                 # 指定ステップごとにモデルを保存
@@ -693,8 +695,8 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
             accelerator.device,
             vae,
             [tokenizer1, tokenizer2],
-            [text_encoder1, text_encoder2],
-            unet,
+            [accelerator.unwrap_model(text_encoder1), accelerator.unwrap_model(text_encoder2)],
+            accelerator.unwrap_model(unet),
         )
 
     is_main_process = accelerator.is_main_process

diff --git a/xl_model_finetuner_training.py b/xl_model_finetuner_training.py
@@ -305,7 +305,8 @@ def get_folder_name_and_num_repeats(folder):
         "gradient_checkpointing"        : gradient_checkpointing,
         "gradient_accumulation_steps"   : gradient_accumulation_steps,
         "mixed_precision"               : mixed_precision,
-        "ddp_timeout"                   : 999999,
+        "ddp_gradient_as_bucket_view"   : True,
+        "ddp_static_graph"              : True,
     },
     "logging_arguments": {
         "log_with"          : "wandb" if wandb_api_key else "tensorboard",