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

apply offloading method runable for all trainer

Browse Source
This commit is contained in:
BootsofLagrangian
2024-02-05 22:42:06 +09:00
parent 3970bf4080
commit 7d2a9268b9
3 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
fine_tune.py
View File

@@ -251,6 +251,11 @@ def train(args):
if args.deepspeed:
# wrapping model
import deepspeed
if args.offload_optimizer_device is not None:
accelerator.print('[DeepSpeed] start to manually build cpu_adam.')
deepspeed.ops.op_builder.CPUAdamBuilder().load()
accelerator.print('[DeepSpeed] building cpu_adam done.')
class DeepSpeedModel(torch.nn.Module):
def __init__(self, unet, text_encoder) -> None:
super().__init__()

5
train_db.py
View File

@@ -224,6 +224,11 @@ def train(args):
# acceleratorがなんかよろしくやってくれるらしい
if args.deepspeed:
# wrapping model
import deepspeed
if args.offload_optimizer_device is not None:
accelerator.print('[DeepSpeed] start to manually build cpu_adam.')
deepspeed.ops.op_builder.CPUAdamBuilder().load()
accelerator.print('[DeepSpeed] building cpu_adam done.')
class DeepSpeedModel(torch.nn.Module):
def __init__(self, unet, text_encoder) -> None:
super().__init__()

5
train_network.py
View File

@@ -419,6 +419,11 @@ class NetworkTrainer:
# acceleratorがなんかよろしくやってくれるらしい / accelerator will do something good
if args.deepspeed:
# wrapping model
import deepspeed
if args.offload_optimizer_device is not None:
accelerator.print('[DeepSpeed] start to manually build cpu_adam.')
deepspeed.ops.op_builder.CPUAdamBuilder().load()
accelerator.print('[DeepSpeed] building cpu_adam done.')
class DeepSpeedModel(torch.nn.Module):
def __init__(self, unet, text_encoder, network) -> None:
super().__init__()