fix(examples): compute LR scheduler steps from num_train_epochs consi… · adi-IL/diffusers@d81879a · GitHub
Skip to content

Commit d81879a

Browse files
adi-ILadi-IL
committed
fix(examples): compute LR scheduler steps from num_train_epochs consistently
Same fix across training examples where num_train_epochs under-counted scheduler steps, especially in distributed runs. Supersedes huggingface#14527 huggingface#14528 huggingface#14540 huggingface#14546 huggingface#14597.
1 parent d30c748 commit d81879a

8 files changed

Lines changed: 159 additions & 40 deletions

File tree

‎examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py‎

Lines changed: 19 additions & 8 deletions

‎examples/t2i_adapter/test_t2i_adapter.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,22 @@ def test_t2i_adapter_sdxl(self):
4949
run_command(self._launch_args + test_args)
5050

5151
assert os.path.isfile(os.path.join(tmpdir, "diffusion_pytorch_model.safetensors"))
52+
53+
def test_t2i_adapter_sdxl_num_train_epochs(self):
54+
with tempfile.TemporaryDirectory() as tmpdir:
55+
test_args = f"""
56+
examples/t2i_adapter/train_t2i_adapter_sdxl.py
57+
--pretrained_model_name_or_path=hf-internal-testing/tiny-stable-diffusion-xl-pipe
58+
--adapter_model_name_or_path=hf-internal-testing/tiny-adapter
59+
--dataset_name=hf-internal-testing/fill10
60+
--output_dir={tmpdir}
61+
--resolution=64
62+
--train_batch_size=1
63+
--gradient_accumulation_steps=1
64+
--num_train_epochs=1
65+
--checkpointing_steps=2
66+
""".split()
67+
68+
run_command(self._launch_args + test_args)
69+
70+
assert os.path.isfile(os.path.join(tmpdir, "diffusion_pytorch_model.safetensors"))

‎examples/t2i_adapter/train_t2i_adapter_sdxl.py‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@
6161
import wandb
6262

6363
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
64-
check_min_version("0.41.0.dev0")
64+
check_min_version("0.40.0.dev0")
6565

6666
logger = get_logger(__name__)
6767

@@ -1080,17 +1080,22 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
10801080
)
10811081

10821082
# Scheduler and math around the number of training steps.
1083-
overrode_max_train_steps = False
1084-
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
1083+
# Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
1084+
num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
10851085
if args.max_train_steps is None:
1086-
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
1087-
overrode_max_train_steps = True
1086+
len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
1087+
num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
1088+
num_training_steps_for_scheduler = (
1089+
args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
1090+
)
1091+
else:
1092+
num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes
10881093

10891094
lr_scheduler = get_scheduler(
10901095
args.lr_scheduler,
10911096
optimizer=optimizer,
1092-
num_warmup_steps=args.lr_warmup_steps,
1093-
num_training_steps=args.max_train_steps,
1097+
num_warmup_steps=num_warmup_steps_for_scheduler,
1098+
num_training_steps=num_training_steps_for_scheduler,
10941099
num_cycles=args.lr_num_cycles,
10951100
power=args.lr_power,
10961101
)
@@ -1102,8 +1107,14 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
11021107

11031108
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
11041109
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
1105-
if overrode_max_train_steps:
1110+
if args.max_train_steps is None:
11061111
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
1112+
if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:
1113+
logger.warning(
1114+
f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
1115+
f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
1116+
f"This inconsistency may result in the learning rate scheduler not functioning properly."
1117+
)
11071118
# Afterwards we recalculate our number of training epochs
11081119
args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
11091120

‎examples/text_to_image/test_text_to_image.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,3 +345,26 @@ def test_text_to_image_sdxl(self):
345345
# save_pretrained smoke test
346346
assert os.path.isfile(os.path.join(tmpdir, "unet", "diffusion_pytorch_model.safetensors"))
347347
assert os.path.isfile(os.path.join(tmpdir, "scheduler", "scheduler_config.json"))
348+
349+
def test_text_to_image_sdxl_num_train_epochs(self):
350+
with tempfile.TemporaryDirectory() as tmpdir:
351+
test_args = f"""
352+
examples/text_to_image/train_text_to_image_sdxl.py
353+
--pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-xl-pipe
354+
--dataset_name hf-internal-testing/dummy_image_text_data
355+
--resolution 64
356+
--center_crop
357+
--random_flip
358+
--train_batch_size 1
359+
--gradient_accumulation_steps 1
360+
--num_train_epochs 1
361+
--learning_rate 5.0e-04
362+
--scale_lr
363+
--lr_scheduler constant
364+
--lr_warmup_steps 0
365+
--output_dir {tmpdir}
366+
""".split()
367+
368+
run_command(self._launch_args + test_args)
369+
assert os.path.isfile(os.path.join(tmpdir, "unet", "diffusion_pytorch_model.safetensors"))
370+
assert os.path.isfile(os.path.join(tmpdir, "scheduler", "scheduler_config.json"))

‎examples/text_to_image/test_text_to_image_lora.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,28 @@ def test_text_to_image_lora_sdxl(self):
220220
is_lora = all("lora" in k for k in lora_state_dict.keys())
221221
assert is_lora
222222

223+
def test_text_to_image_lora_sdxl_num_train_epochs(self):
224+
with tempfile.TemporaryDirectory() as tmpdir:
225+
test_args = f"""
226+
examples/text_to_image/train_text_to_image_lora_sdxl.py
227+
--pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-xl-pipe
228+
--dataset_name hf-internal-testing/dummy_image_text_data
229+
--resolution 64
230+
--train_batch_size 1
231+
--gradient_accumulation_steps 1
232+
--num_train_epochs 1
233+
--learning_rate 5.0e-04
234+
--scale_lr
235+
--lr_scheduler constant
236+
--lr_warmup_steps 0
237+
--output_dir {tmpdir}
238+
""".split()
239+
240+
run_command(self._launch_args + test_args)
241+
assert os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
242+
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
243+
assert all("lora" in k for k in lora_state_dict.keys())
244+
223245
def test_text_to_image_lora_sdxl_with_text_encoder(self):
224246
with tempfile.TemporaryDirectory() as tmpdir:
225247
test_args = f"""

‎examples/text_to_image/train_text_to_image_lora_sdxl.py‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@
6868
import wandb
6969

7070
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
71-
check_min_version("0.41.0.dev0")
71+
check_min_version("0.40.0.dev0")
7272

7373
logger = get_logger(__name__)
7474
if is_torch_npu_available():
@@ -1011,17 +1011,22 @@ def collate_fn(examples):
10111011
)
10121012

10131013
# Scheduler and math around the number of training steps.
1014-
overrode_max_train_steps = False
1015-
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
1014+
# Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
1015+
num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
10161016
if args.max_train_steps is None:
1017-
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
1018-
overrode_max_train_steps = True
1017+
len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
1018+
num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
1019+
num_training_steps_for_scheduler = (
1020+
args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
1021+
)
1022+
else:
1023+
num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes
10191024

10201025
lr_scheduler = get_scheduler(
10211026
args.lr_scheduler,
10221027
optimizer=optimizer,
1023-
num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
1024-
num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
1028+
num_warmup_steps=num_warmup_steps_for_scheduler,
1029+
num_training_steps=num_training_steps_for_scheduler,
10251030
)
10261031

10271032
# Prepare everything with our `accelerator`.
@@ -1036,8 +1041,14 @@ def collate_fn(examples):
10361041

10371042
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
10381043
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
1039-
if overrode_max_train_steps:
1044+
if args.max_train_steps is None:
10401045
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
1046+
if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:
1047+
logger.warning(
1048+
f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
1049+
f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
1050+
f"This inconsistency may result in the learning rate scheduler not functioning properly."
1051+
)
10411052
# Afterwards we recalculate our number of training epochs
10421053
args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
10431054

‎examples/text_to_image/train_text_to_image_sdxl.py‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@
5555

5656

5757
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
58-
check_min_version("0.41.0.dev0")
58+
check_min_version("0.40.0.dev0")
5959

6060
logger = get_logger(__name__)
6161
if is_torch_npu_available():
@@ -979,17 +979,22 @@ def collate_fn(examples):
979979
)
980980

981981
# Scheduler and math around the number of training steps.
982-
overrode_max_train_steps = False
983-
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
982+
# Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
983+
num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
984984
if args.max_train_steps is None:
985-
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
986-
overrode_max_train_steps = True
985+
len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
986+
num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
987+
num_training_steps_for_scheduler = (
988+
args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
989+
)
990+
else:
991+
num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes
987992

988993
lr_scheduler = get_scheduler(
989994
args.lr_scheduler,
990995
optimizer=optimizer,
991-
num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
992-
num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
996+
num_warmup_steps=num_warmup_steps_for_scheduler,
997+
num_training_steps=num_training_steps_for_scheduler,
993998
)
994999

9951000
# Prepare everything with our `accelerator`.
@@ -1002,8 +1007,14 @@ def collate_fn(examples):
10021007

10031008
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
10041009
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
1005-
if overrode_max_train_steps:
1010+
if args.max_train_steps is None:
10061011
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
1012+
if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:
1013+
logger.warning(
1014+
f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
1015+
f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
1016+
f"This inconsistency may result in the learning rate scheduler not functioning properly."
1017+
)
10071018
# Afterwards we recalculate our number of training epochs
10081019
args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
10091020

‎examples/textual_inversion/textual_inversion.py‎

Lines changed: 19 additions & 8 deletions

0 commit comments

Comments
 (0)