VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-08 22:35:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix most of ZeRO stage uses optimizer partitioning

Browse Source 
- we have to prepare optimizer and ds_model at the same time.
 - pull/1139#issuecomment-1986790007

Signed-off-by: BootsofLagrangian <hard2251@yonsei.ac.kr>
This commit is contained in:
BootsofLagrangian
2024-03-20 20:52:59 +09:00
parent fbb98f144e
commit d9456020d7
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
sdxl_train.py
View File

@@ -412,7 +412,10 @@ def train(args):
text_encoder1=text_encoder1 if train_text_encoder1 else None, text_encoder1=text_encoder1 if train_text_encoder1 else None,
text_encoder2=text_encoder2 if train_text_encoder2 else None, text_encoder2=text_encoder2 if train_text_encoder2 else None,
) )
ds_model = accelerator.prepare(ds_model) # most of ZeRO stage uses optimizer partitioning, so we have to prepare optimizer and ds_model at the same time. # pull/1139#issuecomment-1986790007
ds_model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
ds_model, optimizer, train_dataloader, lr_scheduler
)
training_models = [ds_model] training_models = [ds_model]
else: else:
@@ -423,8 +426,7 @@ def train(args):
text_encoder1 = accelerator.prepare(text_encoder1) text_encoder1 = accelerator.prepare(text_encoder1)
if train_text_encoder2: if train_text_encoder2:
text_encoder2 = accelerator.prepare(text_encoder2) text_encoder2 = accelerator.prepare(text_encoder2)
optimizer, train_dataloader, lr_scheduler = accelerator.prepare(optimizer, train_dataloader, lr_scheduler)
optimizer, train_dataloader, lr_scheduler = accelerator.prepare(optimizer, train_dataloader, lr_scheduler)
# TextEncoderの出力をキャッシュするときにはCPUへ移動する # TextEncoderの出力をキャッシュするときにはCPUへ移動する
if args.cache_text_encoder_outputs: if args.cache_text_encoder_outputs: