Advanced training has an issue when saving checkpoints

[landmann]

Describe the bug

I'm running both

https://github.com/huggingface/diffusers/blob/main/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py

and

https://github.com/huggingface/diffusers/blob/main/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

and they both have issues around saving the checkpoint -different issues, but issues nonetheless.

Reproduction

Just train it

Logs

For sdxl:


Traceback (most recent call last):
  File "/home/ubuntu/notebooks/../scripts/train_lora_sdxl.py", line 2234, in <module>
    main(args)
  File "/home/ubuntu/notebooks/../scripts/train_lora_sdxl.py", line 2007, in main
    accelerator.save_state(save_path)
  File "/home/ubuntu/myenv/lib/python3.10/site-packages/accelerate/accelerator.py", line 2767, in save_state
    hook(self._models, weights, output_dir)
  File "/home/ubuntu/notebooks/../scripts/train_lora_sdxl.py", line 1420, in save_model_hook
    embedding_handler.save_embeddings(f"{output_dir}/{args.output_dir}_emb.safetensors")
  File "/home/ubuntu/notebooks/../scripts/train_lora_sdxl.py", line 787, in save_embeddings
    save_file(tensors, file_path)
  File "/home/ubuntu/myenv/lib/python3.10/site-packages/safetensors/torch.py", line 281, in save_file
    serialize_file(_flatten(tensors), filename, metadata=metadata)
safetensors_rust.SafetensorError: Error while serializing: IoError(Os { code: 2, kind: NotFound, message: "No such file or directory" })

For 1.5:

Traceback (most recent call last):
  File "/home/ubuntu/myenv/bin/accelerate", line 8, in <module>
    sys.exit(main())
  File "/home/ubuntu/myenv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py", line 47, in main
    args.func(args)
  File "/home/ubuntu/myenv/lib/python3.10/site-packages/accelerate/commands/launch.py", line 1023, in launch_command
    simple_launcher(args)
  File "/home/ubuntu/myenv/lib/python3.10/site-packages/accelerate/commands/launch.py", line 643, in simple_launcher
    raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/home/ubuntu/myenv/bin/python3', '../scripts/train_lora_sd_15.py', '--pretrained_model_name_or_path=/home/ubuntu/models/sd15/fml', '--dataset_name=/home/ubuntu/nate_pics_768/', '--output_dir=/home/ubuntu/nate_models/lora_fml', "--instance_prompt='a TOK man'", '--gradient_accumulation_steps=1', '--caption_column=prompt', '--train_batch_size=4', '--repeats=1', '--mixed_precision=bf16', '--resolution=768', '--gradient_checkpointing', '--learning_rate=1.0', '--text_encoder_lr=1.0', '--adam_beta2=0.99', '--optimizer=prodigy', '--train_text_encoder_ti', '--train_text_encoder_ti_frac=0.5', '--token_abstraction=TOK', '--snr_gamma=5.0', '--lr_scheduler=constant', '--lr_warmup_steps=0', '--rank=32', '--max_train_steps=2160', '--checkpointing_steps=10', '--seed=0', '--with_prior_preservation', '--prior_generation_precision=bf16', '--sample_batch_size=1', "--class_prompt='a man'", '--class_data_dir=/home/ubuntu/notebooks/man_4321_imgs_768x768px', '--report_to=wandb', '--validation_prompt=a TOK man, professional headshot, hyperdetailed photography, soft light, head and shoulders portrait, cover', '--num_validation_images=3', '--validation_epochs=200']' returned non-zero exit status 1.

### System Info

linux distro - nothing special.

### Who can help?

@sayakpaul

[sayakpaul]

What training commands are you using?

Cc: @linoytsaban

[landmann]

Very simple

process = [
    "accelerate", 
    "launch",
    f"{dreambooth_filename}",
    f"--pretrained_model_name_or_path=SG161222/RealVisXL_V4.0",
    f"--dataset_name={img_folder_path}",
    # f"--instance_data_dir={img_folder_path}",
    f"--output_dir={output_dir}",
    f"--instance_prompt='{class_prompt}'",
    f"--gradient_accumulation_steps=1",
    f"--caption_column=prompt",
    "--train_batch_size=4",
    "--repeats=1",
    "--mixed_precision=bf16",
    f"--resolution={IMG_SIZE}",
    "--gradient_checkpointing",
    "--learning_rate=1.0",
    "--text_encoder_lr=1.0",
    "--adam_beta2=0.99",
    f"--optimizer=prodigy",
    "--train_text_encoder_ti",
    "--train_text_encoder_ti_frac=0.5",
    "--token_abstraction=TOK",
    "--snr_gamma=5.0", # Default 0. SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. 
    f"--lr_scheduler=constant",
    f"--lr_warmup_steps=0",
    f"--rank={rank}",
    f"--max_train_steps={TRAINING_STEPS}",
    "--checkpointing_steps=10",
    "--seed=0"
]

process.extend(
    [
        f"--with_prior_preservation",
        f"--prior_generation_precision=bf16",
        f"--sample_batch_size=1",
        f"--class_prompt='a man'",
        # f"--num_class_images=100", # Not necessary - we have too many class_data_dir images so they will be taken straight from there.
        f"--class_data_dir={class_data_dir}",
    ]
)

process.extend(
    [
        f"--report_to=wandb",
        f'--validation_prompt={validation_prompt}',
        f"--num_validation_images=3",
        f"--validation_epochs=200",
    ]
)

try:
    # Run the subprocess and capture stdout and stderr
    print("Running Subprocess...")
    subprocess.run(process)
    print("Finished!")
except subprocess.CalledProcessError as e:
    # Print more detailed error information
    print(f"Command failed with exit status {e.returncode}")
    print(f"Command: {' '.join(e.cmd)}")
    print(f"Stdout: {e.stdout.decode('utf-8')}")
    print(f"Stderr: {e.stderr.decode('utf-8')}")
    traceback.print_exc()
except Exception as e:
    # Catch-all for other exceptions
    print("An unexpected error occurred.")
    traceback.print_exc()

[landmann]

Ah found the issue in the XL file - I'll make a pr.

Seems like save_embeddings only worked if the model was in the current directory. If the model was in another directory it added the path multiple times. Pushing now.

[landmann]

Oof I love you guys but I found another bug in the sdxl training script potentially when it does the text encoder pivot.

Any ideas, @linoytsaban ?

PIVOT HALFWAY 2
/home/ubuntu/nate_models/lora_RealVis4/checkpoint-5/lora_RealVis4_emb.safetensors
The file path /home/ubuntu/nate_models/lora_RealVis4/checkpoint-5/lora_RealVis4_emb.safetensors does not exist.
/home/ubuntu/nate_models/lora_RealVis4/checkpoint-10/lora_RealVis4_emb.safetensors
The file path /home/ubuntu/nate_models/lora_RealVis4/checkpoint-10/lora_RealVis4_emb.safetensors does not exist.
/home/ubuntu/nate_models/lora_RealVis4/lora_RealVis4_emb.safetensors
The file path /home/ubuntu/nate_models/lora_RealVis4/lora_RealVis4_emb.safetensors does not exist.
Loaded text_encoder_2 as CLIPTextModelWithProjection from `text_encoder_2` subfolder of SG161222/RealVisXL_V4.0.

Loading pipeline components...:  14%|█▍        | 1/7 [00:01<00:06,  1.11s/it]{'attention_type', 'dropout', 'reverse_transformer_layers_per_block'} was not found in config. Values will be initialized to default values.
Loaded unet as UNet2DConditionModel from `unet` subfolder of SG161222/RealVisXL_V4.0.

Loading pipeline components...:  29%|██▊       | 2/7 [00:20<00:57, 11.59s/it]Loaded tokenizer_2 as CLIPTokenizer from `tokenizer_2` subfolder of SG161222/RealVisXL_V4.0.
{'sigma_min', 'timestep_type', 'rescale_betas_zero_snr', 'sigma_max'} was not found in config. Values will be initialized to default values.
Loaded scheduler as EulerDiscreteScheduler from `scheduler` subfolder of SG161222/RealVisXL_V4.0.
Loaded text_encoder as CLIPTextModel from `text_encoder` subfolder of SG161222/RealVisXL_V4.0.

Loading pipeline components...:  71%|███████▏  | 5/7 [00:23<00:08,  4.02s/it]Loaded tokenizer as CLIPTokenizer from `tokenizer` subfolder of SG161222/RealVisXL_V4.0.
Loading pipeline components...: 100%|██████████| 7/7 [00:23<00:00,  3.31s/it]
{'variance_type', 'final_sigmas_type', 'algorithm_type', 'use_lu_lambdas', 'lambda_min_clipped', 'lower_order_final', 'solver_order', 'dynamic_thresholding_ratio', 'euler_at_final', 'solver_type', 'thresholding'} was not found in config. Values will be initialized to default values.
Loading unet.

Traceback (most recent call last):
  File "/home/ubuntu/notebooks/../scripts/train_lora_sdxl.py", line 2242, in <module>
    main(args)
  File "/home/ubuntu/notebooks/../scripts/train_lora_sdxl.py", line 2176, in main
    state_dict["clip_l"],
KeyError: 'clip_l'

The sd 15 file also has a bug in the pivot. Maybe I'm using the wrong version of a library?

Steps:  50%|█████     | 5/10 [00:21<00:19,  3.83s/it, loss=0.0958, lr=1]03/06/2024 15:34:14 - INFO - accelerate.accelerator - Saving current state to /home/ubuntu/nate_models/lora_fml/checkpoint-5
Traceback (most recent call last):
  File "/home/ubuntu/notebooks/../scripts/train_lora_sd_15.py", line 1998, in <module>
    main(args)
  File "/home/ubuntu/notebooks/../scripts/train_lora_sd_15.py", line 1791, in main
    accelerator.save_state(save_path)
  File "/home/ubuntu/myenv/lib/python3.10/site-packages/accelerate/accelerator.py", line 2767, in save_state
    hook(self._models, weights, output_dir)
  File "/home/ubuntu/notebooks/../scripts/train_lora_sd_15.py", line 1290, in save_model_hook
    raise ValueError(f"unexpected save model: {model.__class__}")
ValueError: unexpected save model: <class 'transformers.models.clip.modeling_clip.CLIPTextModel'>

[linoytsaban]

Hi! can you share your environment please?

[landmann]

Hi @linoytsaban! Thank you so much for your help 🤗
I think I made it work. Installed new environment for it and after a few tries and wiggles it fixed the issue. And I just realized all experiments do have a requirements.txt! Thank you so much guys for your patience 😁

[chideraidabor]

Ah found the issue in the XL file - I'll make a pr.

Seems like save_embeddings only worked if the model was in the current directory. If the model was in another directory it added the path multiple times. Pushing now.

hi what was the solution to this?