Skip to content

Commit 4295f91

Browse files
BootsofLagrangianBootsofLagrangian
committed
fix all trainer about vae
1 parent 2824312 commit 4295f91

3 files changed

Lines changed: 37 additions & 36 deletions

File tree

‎fine_tune.py‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -221,10 +221,18 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
221221

222222
# 学習ステップ数を計算する
223223
if args.max_train_epochs is not None:
224-
args.max_train_steps = args.max_train_epochs * math.ceil(
225-
len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
226-
)
227-
accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
224+
if args.deepspeed:
225+
args.max_train_steps = args.max_train_epochs * math.ceil(
226+
len(train_dataloader) / args.gradient_accumulation_steps
227+
)
228+
accelerator.print(
229+
f"[DeepSpeed] override steps not dividing by {accelerator.num_processes}. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}"
230+
)
231+
else:
232+
args.max_train_steps = args.max_train_epochs * math.ceil(
233+
len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
234+
)
235+
accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
228236

229237
# データセット側にも学習ステップを送信
230238
train_dataset_group.set_max_train_steps(args.max_train_steps)
@@ -244,21 +252,16 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
244252
if args.deepspeed:
245253
# wrapping model
246254
class DeepSpeedModel(torch.nn.Module):
247-
def __init__(self, unet, text_encoder, vae) -> None:
255+
def __init__(self, unet, text_encoder) -> None:
248256
super().__init__()
249257
self.unet = unet
250258
self.text_encoders = self.text_encoder = torch.nn.ModuleList(text_encoder)
251-
self.vae = vae
252259
def get_models(self):
253-
return self.unet, self.text_encoders, self.vae
254-
255-
unet.to(accelerator.device, dtype=weight_dtype)
256-
[t_enc.to(accelerator.device, dtype=weight_dtype) for t_enc in text_encoders]
257-
ds_model = DeepSpeedModel(unet, text_encoders, vae)
260+
return self.unet, self.text_encoders
261+
ds_model = DeepSpeedModel(unet, text_encoders)
258262
ds_model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(ds_model, optimizer, train_dataloader, lr_scheduler)
259263
# Now, ds_model is an instance of DeepSpeedEngine.
260-
unet, text_encoders, vae = ds_model.get_models() # for compatiblility
261-
vae.to(vae_dtype)
264+
unet, text_encoders = ds_model.get_models() # for compatiblility
262265
text_encoder = text_encoders
263266

264267
else: # acceleratorがなんかよろしくやってくれるらしい

‎train_db.py‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -190,10 +190,18 @@ def train(args):
190190

191191
# 学習ステップ数を計算する
192192
if args.max_train_epochs is not None:
193-
args.max_train_steps = args.max_train_epochs * math.ceil(
194-
len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
195-
)
196-
accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
193+
if args.deepspeed:
194+
args.max_train_steps = args.max_train_epochs * math.ceil(
195+
len(train_dataloader) / args.gradient_accumulation_steps
196+
)
197+
accelerator.print(
198+
f"[DeepSpeed] override steps not dividing by {accelerator.num_processes}. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}"
199+
)
200+
else:
201+
args.max_train_steps = args.max_train_epochs * math.ceil(
202+
len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
203+
)
204+
accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
197205

198206
# データセット側にも学習ステップを送信
199207
train_dataset_group.set_max_train_steps(args.max_train_steps)
@@ -217,22 +225,17 @@ def train(args):
217225
if args.deepspeed:
218226
# wrapping model
219227
class DeepSpeedModel(torch.nn.Module):
220-
def __init__(self, unet, text_encoder, vae) -> None:
228+
def __init__(self, unet, text_encoder) -> None:
221229
super().__init__()
222230
self.unet = unet
223231
self.text_encoders = self.text_encoder = torch.nn.ModuleList(text_encoder)
224-
self.vae = vae
225232

226233
def get_models(self):
227-
return self.unet, self.text_encoders, self.vae
228-
229-
unet.to(accelerator.device, dtype=weight_dtype)
230-
[t_enc.to(accelerator.device, dtype=weight_dtype) for t_enc in text_encoders]
231-
ds_model = DeepSpeedModel(unet, text_encoders, vae)
234+
return self.unet, self.text_encoders
235+
ds_model = DeepSpeedModel(unet, text_encoders)
232236
ds_model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(ds_model, optimizer, train_dataloader, lr_scheduler)
233237
# Now, ds_model is an instance of DeepSpeedEngine.
234-
unet, text_encoders, vae = ds_model.get_models() # for compatiblility
235-
vae.to(vae_dtype) # to avoid explicitly half-vae
238+
unet, text_encoders = ds_model.get_models() # for compatiblility
236239
text_encoder = text_encoders
237240
else:
238241
if train_text_encoder:

‎train_network.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -364,7 +364,7 @@ def train(self, args):
364364
len(train_dataloader) / args.gradient_accumulation_steps
365365
)
366366
accelerator.print(
367-
f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}"
367+
f"[DeepSpeed] override steps not dividing by {accelerator.num_processes}. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}"
368368
)
369369
else:
370370
args.max_train_steps = args.max_train_epochs * math.ceil(
@@ -420,23 +420,18 @@ def train(self, args):
420420
if args.deepspeed:
421421
# wrapping model
422422
class DeepSpeedModel(torch.nn.Module):
423-
def __init__(self, unet, text_encoder, vae, network) -> None:
423+
def __init__(self, unet, text_encoder, network) -> None:
424424
super().__init__()
425425
self.unet = unet
426426
self.text_encoders = self.text_encoder = torch.nn.ModuleList(text_encoder)
427-
self.vae = vae
428427
self.network = network
429428

430429
def get_models(self):
431-
return self.unet, self.text_encoders, self.vae, self.network
432-
433-
unet.to(accelerator.device, dtype=unet_weight_dtype)
434-
[t_enc.to(accelerator.device, dtype=te_weight_dtype) for t_enc in text_encoders]
435-
ds_model = DeepSpeedModel(unet, text_encoders, vae, network)
430+
return self.unet, self.text_encoders, self.network
431+
ds_model = DeepSpeedModel(unet, text_encoders, network)
436432
ds_model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(ds_model, optimizer, train_dataloader, lr_scheduler)
437433
# Now, ds_model is an instance of DeepSpeedEngine.
438-
unet, text_encoders, vae, network = ds_model.get_models() # for compatiblility
439-
vae.to(vae_dtype) # to avoid explicitly half-vae
434+
unet, text_encoders, network = ds_model.get_models() # for compatiblility
440435
text_encoder = text_encoders
441436
else:
442437
if train_unet:

0 commit comments

Comments
 (0)