perf(diffusion): improve Flux training throughput

examples/diffusion/finetune/flux_t2i_flow.yaml

@@ -3,14 +3,19 @@ model:
   mode: "finetune"
   cache_dir: null
   attention_backend: "flash"
+  transformer_engine_fp8: true
+  transformer_engine_fp8_recipe: "delayed"
+  transformer_engine_fp8_amax_history_len: 1024
+  transformer_engine_fp8_amax_compute_algo: "max"
 
 optim:
   learning_rate: 2e-5
 
   optimizer:
     weight_decay: 0.01
     betas: [0.9, 0.999]
-    foreach: true
+    foreach: false
+    fused: true
 
 performance:
   check_loss: false
@@ -25,9 +30,10 @@ fsdp:
   activation_checkpointing: false
   cpu_offload: false
   defer_fsdp_grad_sync: true
+  enable_compile: false
   enable_fsdp2_prefetch: true
-  fsdp2_backward_prefetch_depth: 2
-  fsdp2_forward_prefetch_depth: 1
+  fsdp2_backward_prefetch_depth: 3
+  fsdp2_forward_prefetch_depth: 2
 
 flow_matching:
   adapter_type: "flux"
@@ -49,8 +55,8 @@ flow_matching:
 
 step_scheduler:
   num_epochs: 5000
-  local_batch_size: 1
-  global_batch_size: 8
+  local_batch_size: 4
+  global_batch_size: 32
   ckpt_every_steps: 500
   log_every: 1
   # max_steps: null  # Set to limit training to a specific number of steps

examples/diffusion/finetune/flux_t2i_flow_lora.yaml

@@ -17,6 +17,7 @@ model:
   mode: "finetune"
   cache_dir: null
   attention_backend: "flash"
+  transformer_engine_fp8: false
 
 # ── LoRA / PEFT configuration ─────────────────────────────────────────────────
 # Remove this section for full fine-tune.
@@ -45,22 +46,18 @@ optim:
   optimizer:
     weight_decay: 0.01
     betas: [0.9, 0.999]
+    foreach: false
+    fused: true
 
 lr_scheduler:
   lr_decay_style: cosine
   lr_warmup_steps: 100
   min_lr: 1e-6
 
-# FSDP2 with param_dtype=None when LoRA enabled:
-#   base weights (bf16): sharded, no grad comms
-#   LoRA weights (fp32): sharded, grad allreduced by FSDP2
-fsdp:
-  dp_size: 8
-  tp_size: 1
-  cp_size: 1
-  pp_size: 1
+# Best-performing Flux LoRA config from the local perf sweep uses replicated DDP.
+ddp:
+  backend: "nccl"
   activation_checkpointing: false
-  cpu_offload: false
 
 flow_matching:
   adapter_type: "flux"
@@ -82,8 +79,8 @@ flow_matching:
 
 step_scheduler:
   num_epochs: 5000
-  local_batch_size: 1
-  global_batch_size: 8
+  local_batch_size: 6
+  global_batch_size: 48
   ckpt_every_steps: 500
   log_every: 1
   # max_steps: null  # Set to limit training to a specific number of steps
@@ -111,4 +108,4 @@ checkpoint:
 
 ci:
   recipe_owner: pthombre
-  time: "00:30:00"
+  time: "00:30:00"

examples/diffusion/generate/generate.py

@@ -266,10 +266,30 @@ def load_lora_weights_into_pipeline(pipe, cfg):
         raise FileNotFoundError(f"LoRA weights directory not found: {lora_path}")
 
     with open(lora_path / "adapter_config.json") as f:
-        peft_config = PeftConfig.from_dict(json.load(f))
-    apply_lora_to_linear_modules(pipe.transformer, peft_config, skip_freeze=True)
+        peft_config_dict = json.load(f)
+    if "dim" not in peft_config_dict and "r" in peft_config_dict:
+        peft_config_dict["dim"] = peft_config_dict["r"]
+    if "alpha" not in peft_config_dict and "lora_alpha" in peft_config_dict:
+        peft_config_dict["alpha"] = peft_config_dict["lora_alpha"]
+    if "dropout" not in peft_config_dict and "lora_dropout" in peft_config_dict:
+        peft_config_dict["dropout"] = peft_config_dict["lora_dropout"]
+
+    peft_config = PeftConfig.from_dict(peft_config_dict)
+    num_lora_modules = apply_lora_to_linear_modules(pipe.transformer, peft_config, skip_freeze=True)
+    if num_lora_modules == 0:
+        raise RuntimeError(f"No transformer modules matched LoRA config from {lora_path / 'adapter_config.json'}")
+
     state_dict = load_file(lora_path / "adapter_model.safetensors", device="cuda")
-    pipe.transformer.load_state_dict(state_dict, strict=False)
+    state_dict = {key.removeprefix("base_model.model."): value for key, value in state_dict.items()}
+    load_result = pipe.transformer.load_state_dict(state_dict, strict=False)
+    missing_lora_keys = sorted(key for key in load_result.missing_keys if ".lora_" in key)
+    unexpected_lora_keys = sorted(key for key in load_result.unexpected_keys if ".lora_" in key)
+    if missing_lora_keys or unexpected_lora_keys:
+        raise RuntimeError(
+            "LoRA checkpoint did not load cleanly. "
+            f"missing_lora_keys={missing_lora_keys[:10]}, unexpected_lora_keys={unexpected_lora_keys[:10]}"
+        )
+    logger.info("Loaded LoRA adapter from %s (%d modules, %d tensors)", lora_path, num_lora_modules, len(state_dict))
 
 
 def _load_sharded_fsdp_checkpoint(transformer, sharded_dir, torch_dtype=torch.bfloat16):