Huge loss spikes and no convergence while training SD3 medium

[r-aristov]

Hello. Having trouble training "stable-diffusion-3-medium-diffusers" a new style via train_dreambooth_lora_sd3.py script.
I have 65 style reference pictures (and 1000 regularization pictures, generated from class prompt), tried different learning rates, batch sizes, tried training with and without --with_prior_preservation.

accelerate launch train_dreambooth_sd3.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-3-medium-diffusers"  \
  --output_dir="kyr-trained-sd3-lora" \
  --instance_data_dir="kyr-dataset" \
  --class_data_dir="kyr-regset" \
  --instance_prompt="kyrpa pixel art" \
  --class_prompt="pixel art" \
  --mixed_precision="fp16" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=8000 \
  --checkpointing_steps=250 \
  --num_class_images=1000 \
  --with_prior_preservation \
  --report_to="wandb"

This config gives huge spikes in loss:

I mean really huge, with "exclude outliers" option disabled it looks like this:

When I disable --with_prior_preservation and try batch sizes of 2 and 4, gradient_accumulation_steps of 2, 4, 8 and learning rates 1e-6, 3e-6, 5e-6, I do not get spikes in loss, but no convergence either.

Training was performed on A100.

Am I missing something?

I also tried train_dreambooth_sd3.py script, but never managed to run it - there are seems to be bugs in it.
First one is on the line 1582:

prompt_embeds, pooled_prompt_embeds = encode_prompt(
                        text_encoders=[text_encoder_one, text_encoder_two, text_encoder_three],
                        tokenizers=None,
                        prompt=None,
                        text_input_ids_list=[tokens_one, tokens_two, tokens_three],
                    )

There is no such parameter 'text_input_ids_list' in encode_prompt - script fails if --train_text_encoder is set.
I fixed it (probably), but there is another one, causing this error:

File "/root/train_dreambooth_sd3.py", line 1782, in <module>
    main(args)
  File "/root/train_dreambooth_sd3.py", line 1570, in main
    model_pred = transformer(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/accelerate/utils/operations.py", line 822, in forward
    return model_forward(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/accelerate/utils/operations.py", line 810, in call
    return convert_to_fp32(self.model_forward(*args, **kwargs))
  File "/opt/conda/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 14, in decorate_autocast
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/diffusers-0.30.0.dev0-py3.10.egg/diffusers/models/transformers/transformer_sd3.py", line 309, in forward
    hidden_states = torch.utils.checkpoint.checkpoint(
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 251, in checkpoint
    return _checkpoint_without_reentrant(
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 432, in _checkpoint_without_reentrant
    output = function(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/diffusers-0.30.0.dev0-py3.10.egg/diffusers/models/transformers/transformer_sd3.py", line 304, in custom_forward
    return module(*inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/diffusers-0.30.0.dev0-py3.10.egg/diffusers/models/attention.py", line 162, in forward
    norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/diffusers-0.30.0.dev0-py3.10.egg/diffusers/models/normalization.py", line 83, in forward
    x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/normalization.py", line 190, in forward
    return F.layer_norm(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/functional.py", line 2515, in layer_norm
    return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)
TypeError: layer_norm(): argument 'input' (position 1) must be Tensor, not tuple

[sayakpaul]

Can try again after pulling in the latest changes?

What weighting scheme are you using? Have you tried training with "logit_normal"?

There is no such parameter 'text_input_ids_list' in encode_prompt - script fails if --train_text_encoder is set.

This is being fixed in #8630

[r-aristov]

Thanks! I've moved to SDXL, it has much more mature pipeline and I managed to train it without any problems. So, I'll stick to SDXL for now.