flux2: add dual-GPU model-parallel path for FLUX.2 LoRA training

Three changes, all gated on environment variables so single-GPU users
see no behavior difference unless they opt in.

1) examples/flux2/model_training/flux2_dual_gpu_diffsynth.py (new):
   ~170 LOC helper. Splits Flux2DiT.single_transformer_blocks at the
   midpoint across cuda:0/cuda:1. Registers forward_pre_hook on every
   cuda:1 single block (not just the boundary) because Flux2DiT.forward
   passes loop-level constants (temb_mod_params, image_rotary_emb,
   joint_attention_kwargs) into every iteration -- a boundary-only hook
   would leave subsequent blocks receiving cuda:0 tensors. The norm_out
   hook bridges activations back to cuda:0 for the final output layers.

2) examples/flux2/model_training/train.py:
   When FLUX2_DUAL_GPU=true, force CPU model load (so the ~60 GB bf16
   transformer doesn't pre-allocate on cuda:0 before split), run
   torchao Float8WeightOnlyConfig quantize_ with a filter that
   excludes lora_A/lora_B Linear submodules from the quant pass
   (otherwise their requires_grad is stripped and backward fails),
   then call enable_flux2_dual_gpu(model.pipe.dit) after PEFT has
   injected LoRA so layer .to(device) carries LoRA params with their
   base layers.

3) diffsynth/diffusion/runner.py:
   launch_training_task: skip model.to(accelerator.device) when
   FLUX2_DUAL_GPU=true (would undo the split), use
   device_placement=[False, True, True, True] so accelerate doesn't
   re-place the model.
   launch_data_process_task: separate DIFFSYNTH_DATA_PROCESS_ON_CPU
   env-var-gated branch for users without a 48 GB card -- Mistral-24B
   is too big for 32 GB; combined with --initialize_model_on_cpu this
   lets the TE/VAE feature caching step run on CPU.

Bonus: diffsynth/models/flux2_text_encoder.py:
   transformers 5.8 trimmed Mistral3ForConditionalGeneration's
   positional args (dropped output_attentions, output_hidden_states,
   return_dict, cache_position from the signature; they're now
   TransformersKwargs-only). The existing code passed 15 positional
   args and crashed with "takes from 1 to 11 positional arguments but
   15 were given". Rewrote to pass everything by keyword and re-inject
   output_hidden_states/output_attentions via **kwargs so the existing
   get_mistral_3_small_prompt_embeds caller still receives a populated
   output.hidden_states. Independent fix; single-GPU users benefit too.

Validated end-to-end on 2x RTX 5090 with sumi v8 data + the unchanged
recipe from examples/flux2/model_training/lora/FLUX.2-dev.sh:

- sft:data_process (with --initialize_model_on_cpu + the env var):
  3 sumi images cached on CPU-offloaded Mistral-24B in 43s.
- sft:train (FLUX2_DUAL_GPU=true): 15/15 steps at 2.69 s/it sustained,
  40s wall-clock. Distribution lands 20.7 GB cuda:0 / 12.6 GB cuda:1
  after fp8 weight-only quant. LoRA checkpoint epoch-0.safetensors
  (270 MB at rank 32) saved.

The same patch shape has been validated on three other FLUX.2 LoRA
trainers (ai-toolkit, musubi-tuner, OneTrainer) -- documented at
https://github.com/genno-whittlery/flux2-dual-gpu-lora -- with the same
20.7/12.6 GB distribution shape across all four.

Author: genno-whittlery

diffsynth/diffusion/runner.py

@@ -27,8 +27,17 @@ def launch_training_task(
     optimizer = torch.optim.AdamW(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)
     scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
     dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=num_workers)
-    model.to(device=accelerator.device)
-    model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
+    # Dual-GPU model-parallel: skip the device move that would undo our
+    # manual split, and tell accelerate not to touch the model's device.
+    _flux2_dual_gpu = os.environ.get("FLUX2_DUAL_GPU", "false").lower() == "true"
+    if not _flux2_dual_gpu:
+        model.to(device=accelerator.device)
+        model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
+    else:
+        model, optimizer, dataloader, scheduler = accelerator.prepare(
+            model, optimizer, dataloader, scheduler,
+            device_placement=[False, True, True, True],
+        )
     initialize_deepspeed_gradient_checkpointing(accelerator)
     for epoch_id in range(num_epochs):
         for data in tqdm(dataloader):
@@ -59,8 +68,16 @@ def launch_data_process_task(
         num_workers = args.dataset_num_workers
 
     dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0], num_workers=num_workers)
-    model.to(device=accelerator.device)
-    model, dataloader = accelerator.prepare(model, dataloader)
+    # Keep TE/VAE on CPU when explicitly requested. FLUX.2's Mistral-24B TE
+    # is ~48 GB bf16; data_process OOMs on cards <48 GB without CPU offload.
+    _data_process_on_cpu = os.environ.get('DIFFSYNTH_DATA_PROCESS_ON_CPU', 'false').lower() == 'true'
+    if not _data_process_on_cpu:
+        model.to(device=accelerator.device)
+        model, dataloader = accelerator.prepare(model, dataloader)
+    else:
+        model, dataloader = accelerator.prepare(
+            model, dataloader, device_placement=[False, True],
+        )
 
     for data_id, data in enumerate(tqdm(dataloader)):
         with accelerator.accumulate(model):

diffsynth/models/flux2_text_encoder.py

@@ -54,5 +54,25 @@ def __init__(self):
         super().__init__(config)
 
     def forward(self, input_ids = None, pixel_values = None, attention_mask = None, position_ids = None, past_key_values = None, inputs_embeds = None, labels = None, use_cache = None, output_attentions = None, output_hidden_states = None, return_dict = None, cache_position = None, logits_to_keep = 0, image_sizes = None, **kwargs):
-        return super().forward(input_ids, pixel_values, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, logits_to_keep, image_sizes, **kwargs)
+        # transformers 5.8 trimmed Mistral3's positional args; pass everything
+        # as kwargs so the call survives across transformers versions. The dropped
+        # output_hidden_states / output_attentions are still honored via
+        # TransformersKwargs -> forward them through **kwargs.
+        if output_hidden_states is not None:
+            kwargs.setdefault("output_hidden_states", output_hidden_states)
+        if output_attentions is not None:
+            kwargs.setdefault("output_attentions", output_attentions)
+        return super().forward(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            logits_to_keep=logits_to_keep,
+            image_sizes=image_sizes,
+            **kwargs,
+        )
 

examples/flux2/model_training/flux2_dual_gpu_diffsynth.py

@@ -0,0 +1,173 @@
+"""Dual-GPU model-parallel for FLUX.2 in DiffSynth-Studio.
+
+Drop-in helper for DiffSynth-Studio
+(https://github.com/modelscope/DiffSynth-Studio) ``examples/flux2/model_training/``
+that splits the ``Flux2DiT`` transformer across two CUDA devices at the
+``single_transformer_blocks`` midpoint. Enables FLUX.2-dev LoRA training
+on pairs of 24+ GB consumer GPUs (2× RTX 3090, 2× RTX 4090, 2× RTX 5090)
+— on a single 24 GB card the FLUX.2 transformer can't fit alongside
+activations even with WDDM unified-memory paging.
+
+Companion to the validated ai-toolkit patch
+(https://github.com/genno-whittlery/flux2-dual-gpu-lora) and parallel
+ports for musubi-tuner, OneTrainer, and HuggingFace diffusers. The
+DiffSynth port is small because:
+
+- ``Flux2DiT`` is structurally identical to the diffusers
+  ``Flux2Transformer2DModel`` (same field names: ``x_embedder``,
+  ``context_embedder``, ``transformer_blocks``,
+  ``single_transformer_blocks``, ``norm_out``, ``proj_out``,
+  ``time_guidance_embed``, ``*_modulation*``, ``pos_embed``). The
+  identical helper structure transfers directly.
+- DiffSynth uses PEFT (``inject_adapter_in_model`` in
+  ``diffsynth.diffusion.training_module``) for LoRA injection, so
+  per-layer LoRA routing follows the base layer's device automatically.
+- The forward isn't overridden — pre-hooks bridge devices at the split
+  point and at ``norm_out`` (the boundary back to cuda:0).
+
+Env vars:
+    FLUX2_DUAL_GPU=true                    enable dual-GPU path
+    FLUX2_DUAL_GPU_SPLIT_AT=24             override split index
+                                           (default: num_single // 2)
+
+Usage in DiffSynth-Studio's training script
+(``examples/flux2/model_training/train.py``)::
+
+    from flux2_dual_gpu_diffsynth import enable_flux2_dual_gpu
+
+    # ...build pipeline / training module normally...
+    training_module = Flux2ImageTrainingModule(...)
+
+    # Activate the split (env-gated; no-op when FLUX2_DUAL_GPU is unset).
+    # Call AFTER LoRA injection (switch_pipe_to_training_mode) so PEFT
+    # has wrapped target modules. The split places the wrapped modules
+    # and PEFT LoRA params follow the base layer's device automatically.
+    enable_flux2_dual_gpu(training_module.pipe.dit)
+
+    # ...accelerator.prepare with device_placement=False for the
+    #    training_module, since the transformer is already distributed...
+
+Launch with ``--num_processes=1`` — this is *model* parallelism (both
+GPUs cooperate on a single training step), not data parallelism (which
+would spawn one process per GPU).
+"""
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+# ─── Env-gated public surface ───────────────────────────────────────────────
+
+def is_dual_gpu_enabled() -> bool:
+    """True iff ``FLUX2_DUAL_GPU=true`` in the environment."""
+    return os.getenv("FLUX2_DUAL_GPU", "false").lower() == "true"
+
+
+def get_split_at(num_single_blocks: int) -> int:
+    """Single-blocks split index. Override via ``FLUX2_DUAL_GPU_SPLIT_AT``."""
+    override = os.getenv("FLUX2_DUAL_GPU_SPLIT_AT")
+    if override is not None:
+        return int(override)
+    return num_single_blocks // 2
+
+
+def enable_flux2_dual_gpu(dit: nn.Module) -> nn.Module:
+    """Distribute the DiffSynth Flux2DiT across cuda:0 and cuda:1.
+
+    When ``FLUX2_DUAL_GPU`` is unset this is a no-op pass-through.
+
+    Call after the Flux2DiT (typically ``training_module.pipe.dit``) is
+    loaded and after any PEFT LoRA injection has run. PEFT places LoRA
+    params on the base layer's device automatically, so calling this
+    function after LoRA injection puts everything in the right place.
+
+    Returns the (in-place modified) dit.
+    """
+    if not is_dual_gpu_enabled():
+        return dit
+
+    if torch.cuda.device_count() < 2:
+        raise RuntimeError(
+            f"FLUX2_DUAL_GPU=true requires ≥2 CUDA devices, found "
+            f"{torch.cuda.device_count()}."
+        )
+
+    num_single = len(dit.single_transformer_blocks)
+    split_at = get_split_at(num_single)
+    if not 0 < split_at < num_single:
+        raise RuntimeError(
+            f"FLUX2_DUAL_GPU_SPLIT_AT={split_at} out of range "
+            f"(dit has {num_single} single blocks)."
+        )
+
+    cuda0 = torch.device("cuda:0")
+    cuda1 = torch.device("cuda:1")
+
+    # Place pre-blocks scaffolding + all double_blocks + first half of
+    # single_blocks on cuda:0. Output layers (norm_out + proj_out) stay
+    # on cuda:0 — the second pre-hook below brings the activation back
+    # from cuda:1 just in time for them.
+    dit.x_embedder.to(cuda0)
+    dit.context_embedder.to(cuda0)
+    dit.time_guidance_embed.to(cuda0)
+    dit.pos_embed.to(cuda0)
+    dit.double_stream_modulation_img.to(cuda0)
+    dit.double_stream_modulation_txt.to(cuda0)
+    dit.single_stream_modulation.to(cuda0)
+    for block in dit.transformer_blocks:
+        block.to(cuda0)
+    for block in dit.single_transformer_blocks[:split_at]:
+        block.to(cuda0)
+    for block in dit.single_transformer_blocks[split_at:]:
+        block.to(cuda1)
+    dit.norm_out.to(cuda0)
+    dit.proj_out.to(cuda0)
+
+    # Per-block hook on every cuda:1 single_block — the forward loop passes
+    # loop-level constants (temb_mod_params, image_rotary_emb,
+    # joint_attention_kwargs) to each block; a boundary-only hook bridges
+    # only the first block's inputs, so subsequent blocks receive the
+    # originals from cuda:0 and crash with a device-mismatch error.
+    for block in dit.single_transformer_blocks[split_at:]:
+        block.register_forward_pre_hook(
+            _make_device_bridge_hook(cuda1), with_kwargs=True
+        )
+    dit.norm_out.register_forward_pre_hook(
+        _make_device_bridge_hook(cuda0), with_kwargs=True
+    )
+
+    dit._flux2_dual_gpu_split_at = split_at
+    return dit
+
+
+# ─── Internals ──────────────────────────────────────────────────────────────
+
+def _move_to_device(obj: Any, device: torch.device) -> Any:
+    """Recursively move tensors in nested tuple/list/dict to ``device``.
+
+    No-op when a tensor is already on ``device`` (``Tensor.to`` is an
+    identity operation in that case).
+    """
+    if torch.is_tensor(obj):
+        return obj.to(device) if obj.device != device else obj
+    if isinstance(obj, tuple):
+        return tuple(_move_to_device(x, device) for x in obj)
+    if isinstance(obj, list):
+        return [_move_to_device(x, device) for x in obj]
+    if isinstance(obj, dict):
+        return {k: _move_to_device(v, device) for k, v in obj.items()}
+    return obj
+
+
+def _make_device_bridge_hook(target_device: torch.device):
+    """Forward pre-hook that moves all tensor inputs to ``target_device``."""
+    def hook(module, args, kwargs):
+        return (
+            _move_to_device(args, target_device),
+            _move_to_device(kwargs, target_device),
+        )
+    return hook

examples/flux2/model_training/train.py

@@ -2,6 +2,7 @@
 from diffsynth.core import UnifiedDataset
 from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
 from diffsynth.diffusion import *
+from flux2_dual_gpu_diffsynth import enable_flux2_dual_gpu, is_dual_gpu_enabled
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
@@ -133,8 +134,41 @@ def flux2_parser():
         template_model_id_or_path=args.template_model_id_or_path,
         enable_lora_hot_loading=args.enable_lora_hot_loading,
         task=args.task,
-        device="cpu" if args.initialize_model_on_cpu else accelerator.device,
+        # Dual-GPU: force CPU load so the 60 GB bf16 transformer doesn't
+        # pre-allocate on cuda:0 before we get a chance to split it.
+        device="cpu" if (args.initialize_model_on_cpu or is_dual_gpu_enabled()) else accelerator.device,
     )
+    # Dual-GPU split: gated on FLUX2_DUAL_GPU env var. Distributes
+    # pipe.dit across cuda:0 and cuda:1 at the single_transformer_blocks
+    # midpoint. Called AFTER LoRA injection (which happened inside
+    # Flux2ImageTrainingModule.__init__) so PEFT LoRA params follow
+    # the base layer's device automatically when block.to() runs.
+    if is_dual_gpu_enabled():
+        for _i in range(torch.cuda.device_count()):
+            _free, _total = torch.cuda.mem_get_info(_i)
+            print(f"[dual-gpu] pre-distribute cuda:{_i} free={_free/1024**3:.1f}G / {_total/1024**3:.1f}G")
+        # fp8 weight-only quant on CPU BEFORE distribute. Without this the
+        # bf16 transformer (~60 GB) overflows a 32 GB card on either side.
+        # Filter excludes LoRA's lora_A/lora_B Linear submodules -- quantizing
+        # them strips requires_grad and breaks backward.
+        import gc as _gc
+        from torchao.quantization import quantize_, Float8WeightOnlyConfig
+        def _quant_filter(module, name):
+            if not isinstance(module, torch.nn.Linear):
+                return False
+            return "lora_A" not in name and "lora_B" not in name
+        quantize_(model.pipe.dit, Float8WeightOnlyConfig(), filter_fn=_quant_filter)
+        _gc.collect()
+        print("[dual-gpu] fp8 weight-only quant done (LoRA params unmodified)")
+        enable_flux2_dual_gpu(model.pipe.dit)
+        # Pipeline tracks a logical device for input transfer (forward()
+        # calls transfer_data_to_device(inputs, pipe.device, ...)). Pin it
+        # to cuda:0 since x_embedder lives there.
+        model.pipe.device = torch.device("cuda:0")
+        for _i in range(torch.cuda.device_count()):
+            _free, _total = torch.cuda.mem_get_info(_i)
+            print(f"[dual-gpu] post-distribute cuda:{_i} free={_free/1024**3:.1f}G / {_total/1024**3:.1f}G")
+
     model_logger = ModelLogger(
         args.output_path,
         remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,