`split_model_over_gpus` + LoRA training on multi-GPU: five independent device/signature bugs break FLUX LoRA runs (16 GB cards)

Summary

When training FLUX.1-dev LoRA with split_model_over_gpus: true on a machine whose 16 GB GPUs cannot hold the model as a whole, the run crashes. I tracked the failures one by one (each fix moved the crash point forward), and there are five independent bugs: four of them are the same root cause — the lazy GPU splitter only moves the transformer blocks' own module trees, while LoRA weights, the network-level per-batch multiplier, and EMA shadow parameters are all left on the default device (cuda:0) — and one is a stale block-wrapper signature that no longer matches the pinned diffusers build.

All five were verified on two independent machines with the same config and the same set of workaround patches:

• 4×16 GB Ada (1× RTX 4080 SUPER + 3× RTX 4070 Ti SUPER), i9-10900KF, 32 GB RAM, Ubuntu 26.04, driver 615.71.09 — bf16, canary 50/50 green at 3.42 s/it
• 4 of 16×V100 16 GB (sm_70), dual-socket, 512 GB RAM, Ubuntu 26.04 — fp16 (V100 has no bf16), canary 50/50 green at 71.96 s/it, loss 2.047, no NaN

First observed on 0.13.5 (2026-09-07), re-verified on 0.13.6 (2026-09-11), re-applied unchanged on current main (commit db8dbd6, 2026-09-15) — all anchors still match, so the bugs are present in the current release too.

Config used

model:
  name_or_path: "<path>/FLUX.1-dev"
  is_flux: true
  quantize: true
  quantize_te: true
  low_vram: false
  split_model_over_gpus: true
  split_model_other_module_param_count_scale: 0.3
train:
  dtype: bf16        # fp16 on sm_70
  lr: 0.0001
  rank: 16
  batch_size: 1
  gradient_checkpointing: true

Queue: gpu_ids: "0,1,2,3".

Bug 1 — LoRA modules never follow their host module's device

Symptom (first sampling forward):

RuntimeError: Expected all tensors to be on the same device,
got mat2 is on cuda:0, different from other tensors on cuda:1

Why: add_model_gpu_splitter_to_flux (toolkit/models/flux.py) only stubs the split: it assigns each block a _split_device, wraps block.forward for device relay, and replaces transformer.to with new_device_to. The physical move (block.to(block._split_device)) happens later, when the trainer calls transformer.to(...) — i.e. after the LoRA network has already been built. LoRA modules patch their host Linear's forward and live outside the block's module tree, so the block move carries the host weights but leaves lora_down/lora_up on cuda:0. Any block on another device crashes on its first forward. (0.13.5+ keeps the host as a weakref.ref on the LoRA module, which can also read back None after GC — worth guarding upstream as well.)

Workaround applied (in LoRASpecialNetwork.__init__, before optimizer creation — one-shot, avoids desyncing 8-bit AdamW state):

for _lora in self.text_encoder_loras + self.unet_loras:
    _host = _lora.orig_module_ref()
    if _host is None:
        continue
    _lora.to(next(iter(_host.parameters())).device)

This alone is necessary but not sufficient — the lazy move in new_device_to runs after __init__, so it breaks again on the first sample (bug 2).

Bug 2 — LoRA weights must be co-located on every block move

Symptom: the bug-1 one-shot sync passes init, but the first sample still crashes with the same cross-device matmul error.

Why: new_device_to moves blocks after the network exists; the LoRA modules are invisible to block.to().

Workaround applied (two halves):

lora_special.py — register every LoRA on its host module as a plain (non-Module-typed) container, which nn.Module.__setattr__ does not register as a child module:

for _lora in self.text_encoder_loras + self.unet_loras:
    _host = _lora.orig_module_ref()
    if _host is not None:
        _ref = getattr(_host, "_lora_ref", None)
        if _ref is None:
            _host._lora_ref = {_lora}
        else:
            _ref.add(_lora)

toolkit/models/flux.py — co-locate them inside new_device_to, after each block move:

def _co_locate_block_lora(block):
    try:
        for module in block.modules():
            refs = getattr(module, "_lora_ref", None)
            if not refs:
                continue
            for lora in list(refs):
                try:
                    if next(iter(lora.parameters())).device != block._split_device:
                        lora.to(block._split_device)
                except Exception:
                    pass
    except Exception:
        pass

# inside new_device_to:
for block in self.transformer_blocks:
    block.to(block._split_device, *args, **kwargs)
    _co_locate_block_lora(block)
for block in self.single_transformer_blocks:
    block.to(block._split_device, *args, **kwargs)
    _co_locate_block_lora(block)

Steady-state training has no .to() calls, so the hook is zero-cost. Cleaner upstream options: have the LoRA network expose a per-block hook the splitter invokes, or make the LoRA modules proper attributes that participate in nn.Module device moves without becoming named children (e.g. a _parameters-free _buffers-free weak container handled by .to).

Bug 3 — network-level torch_multiplier lives on the default device

Symptom:

RuntimeError: Expected all tensors to be on the same device, cuda:1 and cuda:0

in network_mixins.py at lora_output * multiplier (broadcast_and_multiply).

Why: torch_multiplier is a network-level per-batch tensor created on the default device (cuda:0), while lora_output is on the host block's cuda:N. It is rebuilt by Network._update_torch_multiplier on the default device too, so moving it once at init is not stable — the multiplication site is the only reliable anchor.

Workaround applied (toolkit/network_mixins.py):

multiplier = self.network_ref().torch_multiplier
multiplier = multiplier.to(lora_output.device)   # same-device .to() is a no-op

Upstream option: keep torch_multiplier device-agnostic (CPU scalar broadcast) or rebuild/keep it on the same device as its consumer.

Bug 4 — split_gpu_single_block_forward wrapper signature is stale vs pinned diffusers

Symptom:

TypeError: FluxSingleTransformerBlock.forward() got multiple values for argument 'encoder_hidden_states'

(first sample; training path "worked" only by positional coincidence).

Why: the pinned diffusers build (requirements pin, commit c943837) has

FluxSingleTransformerBlock.forward(hidden_states, encoder_hidden_states, temb,
                                   image_rotary_emb, joint_attention_kwargs)

but the splitter's wrapper still declares the older 4-arg order (hidden_states, temb, image_rotary_emb, joint_attention_kwargs, **kwargs):

• keyword-style call (eval/sampling): encoder_hidden_states from kwargs lands in **kwargs, while the real temb is passed positionally into the temb param and relayed… the wrapper then passes the old params positionally into the real 5-arg signature → encoder_hidden_states gets bound twice → crash.
• positional-style call (training, grad checkpointing): happens to line up by value order, but is semantically fragile.
• Additionally, the wrapper only relays hidden_states/temb/rope to _split_device. The real forward does torch.cat([enc, hidden]), so the un-relayed encoder_hidden_states would cross devices even after the signature is fixed.
• Finally, the pinned single block returns a (hidden, enc) tuple; the wrapper's _split_output_device relay only handles a single tensor.

Workaround applied (whole-function replacement in toolkit/models/flux.py):

def split_gpu_single_block_forward(
    self,
    hidden_states: torch.FloatTensor = None,
    encoder_hidden_states: torch.FloatTensor = None,
    temb: torch.FloatTensor = None,
    image_rotary_emb=None,
    joint_attention_kwargs=None,
    **kwargs
):
    if hidden_states is not None and hidden_states.device != self._split_device:
        hidden_states = hidden_states.to(device=self._split_device)
    if encoder_hidden_states is not None and encoder_hidden_states.device != self._split_device:
        encoder_hidden_states = encoder_hidden_states.to(device=self._split_device)
    if temb is not None and temb.device != self._split_device:
        temb = temb.to(device=self._split_device)
    if image_rotary_emb is not None and image_rotary_emb[0].device != self._split_device:
        image_rotary_emb = tuple([t.to(self._split_device) for t in image_rotary_emb])

    hidden_state_out = self._pre_gpu_split_forward(
        hidden_states=hidden_states,
        encoder_hidden_states=encoder_hidden_states,
        temb=temb,
        image_rotary_emb=image_rotary_emb,
        joint_attention_kwargs=joint_attention_kwargs,
        **kwargs,
    )
    if hasattr(self, "_split_output_device"):
        if isinstance(hidden_state_out, (tuple, list)):   # pinned diffusers single block returns (hidden, enc)
            return tuple(t.to(self._split_output_device) for t in hidden_state_out)
        return hidden_state_out.to(self._split_output_device)
    return hidden_state_out

Upstream suggestion: name wrapper params exactly like the real forward, relay every tensor input, call through by keyword, and handle tuple outputs. Since the wrapper and the pinned diffusers signature are a coupled pair, a test that asserts the two signatures stay in sync would prevent regressions on every diffusers bump.

Bug 5 — EMA shadow parameters stay on the default device

Symptom (first training step, EMA.update):

RuntimeError: Expected all tensors to be on the same device,
but found at least two devices, cuda:0 and cuda:1
  at  tmp = (s_param_float - param_float)

Why: shadow_params are cloned at EMA init, which happens before the lazy split move — so the shadows are all on cuda:0 while the live parameters then move to their per-block devices.

Workaround applied (toolkit/ema.py, top of the update loop):

for s_param, param in zip(self.shadow_params, parameters):
    if s_param.device != param.device:
        s_param.data = s_param.data.to(param.device)
    s_param_float = s_param.float()

(Split assignment is static, so this is a no-op after the first update. copy_to is already cross-device safe via copy_.)

Upstream option: clone shadow parameters from the live parameters' devices (e.g. param.detach().clone() without forcing a device), which is correct regardless of when the splitter runs.

Verification after all five fixes

• 50-step canary: 50/50 completed, no errors; final checkpoint + optimizer + step-0/step-50 samples written; all GPUs return to 0 MB after the run.
• Machine 1 (Ada, bf16): 3.42 s/it. Machine 2 (V100, fp16): 71.96 s/it (V100 is expected to be slow here).
• Re-ran the identical five patches on 0.13.6 (2026-09-11) and on main@db8dbd6 (2026-09-12): all anchors byte-identical, canary green both times, so the fix set is stable across versions and the bugs are still un-fixed upstream.
• low_vram: true with split_model_over_gpus is not compatible and must remain false.

Suggested upstream design (short form)

1. Make the splitter the single owner of device moves: when a block moves, move everything attached to it (LoRA modules registered on it) too — the registration/co-location pair above is the minimal shape.
2. Keep network-level per-batch tensors device-agnostic or follow their consumer's device.
3. Clone EMA shadows from the live parameter device.
4. Keep the block wrapper signature in lockstep with the pinned diffusers blocks (add a signature-pair regression test).