diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 226b011320..50bb0f5d3b 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -32,6 +32,7 @@ from typing import Any, Callable, Optional from core.training.diffusion_train_common import ( DEFAULT_LORA_FILENAME, + DEFAULT_LORA_TARGETS, DiffusionLoraConfig, EventCb, StopCb, @@ -63,9 +64,11 @@ def _select_lora_targets( ) -> tuple[str, ...]: """Pick the LoRA target modules for a DiT run. - An empty ``cfg_targets`` means "unset": use the family's ``spec.lora_targets``. Any - explicit tuple is a deliberate override and still wins.""" - if not tuple(cfg_targets): + ``normalized()`` always fills ``lora_target_modules`` with the generic + ``DEFAULT_LORA_TARGETS`` when a caller does not set it, so that value means "unset" + here: prefer the family's ``spec.lora_targets`` (which add the DiT-specific + projections). Any OTHER explicit tuple is a deliberate override and still wins.""" + if tuple(cfg_targets) == DEFAULT_LORA_TARGETS: return tuple(spec_targets) return tuple(cfg_targets) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 3424df7dc3..db2eca1011 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -211,15 +211,13 @@ def run_diffusion_lora_training( for m in (unet, *text_encoders): m.to(device, dtype = weight_dtype) - # Empty (unset) config means use the family default. - unet_targets = list(cfg.lora_target_modules) or list(DEFAULT_LORA_TARGETS) unet.add_adapter( LoraConfig( r = cfg.lora_rank, lora_alpha = cfg.lora_alpha, lora_dropout = cfg.lora_dropout, init_lora_weights = "gaussian", - target_modules = unet_targets, + target_modules = list(cfg.lora_target_modules), ) ) if cfg.gradient_checkpointing: diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index e4a83b8c50..cdf78995b3 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -30,8 +30,9 @@ from core.inference.diffusion_families import ( trainable_family_names, ) -# Default LoRA target modules: the SDXL U-Net attention projections. DiT trainers supply -# their own wider set instead. +# Default LoRA target modules: the attention projections common to the SDXL U-Net and the +# DiT transformers (the diffusers/kohya convention). A family whose trainer wants a wider +# set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback. DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0") # diffusers' SchedulerType names (diffusers.optimization.get_scheduler). @@ -235,8 +236,7 @@ class DiffusionLoraConfig: lora_rank: int = 16 lora_alpha: Optional[int] = None # defaults to lora_rank lora_dropout: float = 0.0 - # Empty = "unset": each trainer supplies its own family default. - lora_target_modules: tuple[str, ...] = () + lora_target_modules: tuple[str, ...] = DEFAULT_LORA_TARGETS seed: int = 42 mixed_precision: str = "bf16" # "bf16" | "fp16" | "no" snr_gamma: Optional[float] = 5.0 # min-SNR loss weighting; None disables @@ -300,8 +300,7 @@ class DiffusionLoraConfig: if learning_rate <= 0: raise ValueError("learning_rate must be > 0") alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank - # Leave an unset (empty) target list empty so the trainer fills the family default. - targets = tuple(self.lora_target_modules) + targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS # A blank Hub token (the Studio default when none is configured) must load # anonymously, not as an explicit empty credential. token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 224caad35d..3883b9847f 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -698,11 +698,11 @@ class DiffusionTrainingStartRequest(BaseModel): lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank") lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0) # Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that - # sets them is not silently trained with defaults. Empty = "unset": the trainer fills in - # its family default. + # sets them is not silently trained with defaults. Default the target list to the SDXL + # attention projections (the trainer's DEFAULT_LORA_TARGETS) so it is never None. lora_target_modules: List[str] = Field( - default_factory = list, - description = "Modules to attach LoRA to; empty = the trainer's family default", + default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"], + description = "U-Net modules to attach LoRA to", ) max_grad_norm: float = Field(1.0, gt = 0, description = "Gradient clipping max-norm") seed: int = Field(42) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 79a3777113..63cc01922c 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -41,25 +41,23 @@ def test_specs_cover_the_three_dit_families(): def test_select_lora_targets_uses_family_default_for_generic_config(): - # normalized() leaves lora_target_modules empty when a caller doesn't set it, so an empty - # tuple must resolve to the family's targets (which add the DiT-specific joint-attention - # projections), not the generic SDXL list. - assert _select_lora_targets((), _FLUX_TARGETS) == _FLUX_TARGETS - assert _select_lora_targets((), _QWEN_TARGETS) == _QWEN_TARGETS - assert _select_lora_targets((), _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS - # The generic SDXL default is NOT treated as unset here: an explicit list wins. - assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == DEFAULT_LORA_TARGETS + # normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a + # caller doesn't set it, so that value must resolve to the family's targets (which add + # the DiT-specific projections), not stay stuck on the generic SDXL list. + assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == _FLUX_TARGETS + assert _select_lora_targets(DEFAULT_LORA_TARGETS, _QWEN_TARGETS) == _QWEN_TARGETS + assert _select_lora_targets(DEFAULT_LORA_TARGETS, _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS def test_select_lora_targets_explicit_override_wins(): - # Any explicit tuple is a deliberate override and must win over the family spec. + # Any OTHER explicit tuple is a deliberate override and must win over the family spec. override = ("to_q", "to_k") assert _select_lora_targets(override, _FLUX_TARGETS) == override - # The default request path (config leaving targets unset) reaches the spec. + # The default request path (config carrying the generic default) reaches the spec. cfg = DiffusionLoraConfig( base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o" ).normalized() - assert cfg.lora_target_modules == () + assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS assert ( _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) == _FLUX_TARGETS diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index dc84b70f17..d9078bf0e7 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -94,8 +94,7 @@ def test_discover_missing_dir_raises(tmp_path): def test_config_normalized_defaults(): cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o").normalized() assert cfg.lora_alpha == cfg.lora_rank # alpha defaults to rank - # Targets stay empty (unset) after normalize; each trainer fills its own family default. - assert cfg.lora_target_modules == () + assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS @pytest.mark.parametrize(