[research_projects] add shortened flux training script with quantization #11743
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| if __name__ == "__main__": | ||
| class Args: | ||
| pretrained_model_name_or_path = "black-forest-labs/FLUX.1-dev" | ||
| data_df_path = "embeddings_alphonse_mucha.parquet" |
There was a problem hiding this comment.
Add a note on where this is coming from.
| ) | ||
|
|
||
| # Compute loss | ||
| weighting = compute_loss_weighting_for_sd3("none", sigmas) |
There was a problem hiding this comment.
We have the weighting_scheme args. So, let's use it from there.
| noise = torch.randn_like(model_input) | ||
| bsz = model_input.shape[0] | ||
|
|
||
| u = compute_density_for_timestep_sampling("none", bsz, 0.0, 1.0, 1.29) |
There was a problem hiding this comment.
args.weighting_scheme
Let's also make constants for the magic numbers.
| model = accelerator.unwrap_model(model) | ||
| return model._orig_mod if is_compiled_module(model) else model | ||
|
|
||
| def save_model_hook(models, weights, output_dir): |
There was a problem hiding this comment.
If we don't have a load model hook, then I don't think it will make sense to have this either no? Are do we have a utility to resume from a checkpoint in this script?
There was a problem hiding this comment.
I kept save_model_hook to save intermediate checkpoints, but I probably didn't need to save optimizer states too. Though, yeah I think adding back load model hook to resume from checkpoints is a good idea
There was a problem hiding this comment.
Yeah either we don't support intermediate checkpoints at all or support it. I think okay without to prefer minimalism.
| cast_training_params([transformer], dtype=torch.float32) if args.mixed_precision == "fp16" else None | ||
|
|
||
| # Initialize tracking | ||
| accelerator.init_trackers("dreambooth-flux-dev-lora-alphonse-mucha", config=vars(args)) if accelerator.is_main_process else None |
There was a problem hiding this comment.
| accelerator.init_trackers("dreambooth-flux-dev-lora-alphonse-mucha", config=vars(args)) if accelerator.is_main_process else None | |
| if accelerator.is_main_proces | |
| accelerator.init_trackers("dreambooth-flux-dev-lora-alphonse-mucha", config=vars(args)) |
Can we do it while creating the output folder?
examples/research_projects/flux_lora_quantization/train_dreambooth_lora_flux_nano.py
Show resolved
Hide resolved
| init_lora_weights="gaussian", | ||
| target_modules=["to_k", "to_q", "to_v", "to_out.0"], | ||
| ) | ||
| transformer.add_adapter(transformer_lora_config) |
There was a problem hiding this comment.
Should cast the LoRA params to FP32. Do you have a full run with this script that works without FP32 upcasting?
There was a problem hiding this comment.
I think I was casting to FP32 below with
cast_training_params([transformer], dtype=torch.float32) if args.mixed_precision == "fp16" else None below (probably will move it over here and change it to match original training script better.
I do have a full run with this script with reasonable results without FP32 upcasting.
But, I noticed in the loss curves are slightly different between nano script (rare-voice-24 run) and original script (fanciful-totem-2) so I will need to find where the discrepancy is coming from.
There was a problem hiding this comment.
If it doesn't affect results, probably okay
|
@bot /style |
|
Style fixes have been applied. View the workflow run here. |
Co-authored-by: Sayak Paul <[email protected]>
adds shortened script to be referenced in this blogpost huggingface/blog#2888