Lora not learning anything

[AlexandreLrz]

To demonstrate this issue I ran a training with a minimal dataset composed of one picture repeated 12 times :

Here is the caption for the picture (generated by Qwen 7B) :

Man with brown hair and blue shirt smiling. Large wooden crate with question mark and "THIS WAY UP" arrows. Man in red shirt pointing at crate. Green grassy field and blue sky background Bright, cheerful lighting.

Here is my configuration file :

{    
    "data_backend_config": "config/multidatabackend.json",
    "checkpointing_steps": 250,
    "checkpoints_total_limit": 40,
    "disable_benchmark": false,
   
    "learning_rate": 6e-5,
    "lr_warmup_steps": 100,
    "lr_scheduler": "constant_with_warmup",
    "flow_schedule_shift": 2,
    "gradient_checkpointing": true,
    "max_grad_norm": 1.0,
    
    "use_ema": false,
    
    "hub_model_id": "yts",
    "model_family": "hidream",
    "model_flavour": "full",
    "mixed_precision": "bf16",
    "base_model_precision": "int8-quanto",
    "model_type": "lora",
    "lora_alpha": 16,
    "lora_rank": 16,
    "lora_type": "standard",

    "quantize_via": "cpu",
    "quantize_activations": false,
    
    "report_to": "wandb",
    "tracker_project_name": "hidream-training_29",
    "tracker_run_name": "training-run_1",

    "push_checkpoints_to_hub": false,
    "push_to_hub": false,

    "train_batch_size": 1,
    "resume_from_checkpoint": "latest",
    "seed": 42,
    "optimizer": "adamw_bf16",
    "max_train_steps": 30000,
    "num_train_epochs": 0,
    "ignore_final_epochs": true,
    "caption_dropout_probability": 0.15,
    "gradient_accumulation_steps": 1,

    "validation_guidance": 5.0,
    "validation_num_inference_steps": 30,
    "validation_resolution": "640x360",
    "validation_seed": 42,
    "validation_steps": 50,
    "validation_torch_compile": false,
    "validation_prompt_library": true,
    "user_prompt_library": "config/user_prompt_library.json",

    "output_dir": "output/",

    "minimum_image_size": 360,
    "maximum_image_size": 2000,
    "target_downsample_size": 2000,
    "resolution_type": "pixel_area",
    "resolution": 720,
    "override_dataset_config": true,

    "flux_lora_target": "all+ffs+embedder+controlnet"
}

multidatabackend :

[
  {
    "id": "test1",
    "type": "local",
    "instance_data_dir": "/workspace/dataset",
    "crop": false,
    "crop_style": "center",
    "crop_aspect": "preserve",
    "minimum_image_size": 360,
    "maximum_image_size": 2000,
    "target_downsample_size": 2000,
    "resolution": 720,
    "resolution_type": "pixel_area",
    "prepend_instance_prompt": false,
    "instance_prompt": "",
    "only_instance_prompt": false,
    "caption_strategy": "filename",
    "cache_dir_vae": "cache/vae/test",
    "vae_cache_clear_each_epoch": false,
    "probability": 1.0,
    "repeats": 1,
    "skip_file_discovery": "",
    "preserve_data_backend_cache": true,
    "text_embeds": "text-embed-cache"
  },
  {
    "id": "text-embed-cache",
    "dataset_type": "text_embeds",
    "default": true,
    "type": "local",
    "cache_dir": "cache/text/test",
    "write_batch_size": 128
  }
]

And prompt library :

{
  "mrbeast": "Man with brown hair and blue shirt smiling. Large wooden crate with question mark and \"THIS WAY UP\" arrows. Man in red shirt pointing at crate. Green grassy field and blue sky background. Bright, cheerful lighting."
}

I reused the exact same prompt of the only picture in training for my validation samples and here is the result at 400 steps :

As you can see it's not even close to the training picture, even if it's the exact same prompt.

I targeted all layers in the network with this line (to be sure to have the maximum influence) :

"flux_lora_target": "all+ffs+embedder+controlnet"

But even like that the model fails to be close to one single image.

From what I know about lora the model should be really close to that picture after few steps considering that there is only one to learn...

I'd be grateful if you could help me understand what's wrong here.

Thank you.

[bghira]

your learning rate isn't high enough for the low rank of 16 and the AdamW optimiser.

for rank-16 you may be able to use as high as 3e-4, but if your LR is not high enough, it will just be stuck in local minima.

you can see the model weights ARE changing - the image is definitely different.

alternatively, you can keep the LR the same and switch optimiser to optimi-lion (or another Lion variant) and it will definitely start to cook the model more rapidly.

once you increase batch size you will again have to increase the learning rate a little bit.

there's the Prodigy optimiser also supported, which bypasses most of the need for setting a learning rate. you'll want to set LR to 1.0 for that.

[bghira]

also i note the flow schedule shift is set to 2 which might not be training the high-noise timesteps enough to really change the image composition, you might want to try a value like 3 or 5.