Send T5 tokens to the T5 encoder's own device when caching text embeddings

[rajakshay]

What

encode_prompts_flux in toolkit/train_tools.py computes a single device from the first text encoder (CLIP) and reuses it for the second encoder (T5). When caching text embeddings, it sends the T5 input ids to device (CLIP's device) instead of to T5's own device.

Why it breaks

On setups that place the large T5 encoder on a different device than CLIP (e.g. CPU offload / device_map on low-VRAM or unified-memory systems), the tokens land on CLIP's GPU while T5's weights are elsewhere, so the T5 forward raises:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu

This surfaces during "Caching text embeddings to disk" — the model, VAE, network, and CLIP all load fine first, then it dies when T5 is invoked.

Fix

Send the tokens to text_encoder[1].device (T5's own device).

Regression safety

When both encoders are co-located (the default single-GPU case), text_encoder[1].device == device, so this is a no-op — no behavior change for existing setups. It only corrects the split-device path.

Validation

• Before: a CLIP+T5 LoRA training run with the T5 encoder offloaded crashes at the caching step with the device mismatch above.
• After: caching completes and training proceeds; a full LoRA training run finished end-to-end and saved weights.

Refs #374 (the cache_text_embeddings device-mismatch report).

[rajakshay]

Reproduction (no models or data)

Tiny stand-in encoders — CLIP on GPU, T5 on CPU — which is exactly the state when the large T5 encoder is offloaded while CLIP stays on the GPU. Needs one CUDA device (to place the two encoders on different devices). Run from the repo root.

repro.py

import sys
import torch
import torch.nn as nn

if not torch.cuda.is_available():
    print("SKIP: needs a CUDA device to place CLIP and T5 on different devices")
    sys.exit(0)

from toolkit.train_tools import encode_prompts_flux


class Tok:
    model_max_length = 77
    def __call__(self, prompts, **kw):
        o = type("O", (), {})()
        o.input_ids = torch.zeros((len(prompts), 8), dtype=torch.long)
        return o


class CLIP(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(8, 8)
    @property
    def device(self): return next(self.parameters()).device
    @property
    def dtype(self): return next(self.parameters()).dtype
    def forward(self, ids, output_hidden_states=False):
        o = type("O", (), {})()
        o.pooler_output = self.lin(ids.float())
        return o


class T5(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(100, 16)
    @property
    def device(self): return next(self.parameters()).device
    @property
    def dtype(self): return next(self.parameters()).dtype
    def forward(self, ids, output_hidden_states=False):
        return (self.emb(ids),)


clip = CLIP().to("cuda")   # CLIP on GPU
t5 = T5().to("cpu")        # T5 offloaded to CPU (the real low-VRAM / device_map case)

try:
    embeds, pooled = encode_prompts_flux([Tok(), Tok()], [clip, t5], ["a photo of a cat"])
    print("PASSED: encode_prompts_flux returned embeds", tuple(embeds.shape),
          "pooled", tuple(pooled.shape))
except RuntimeError as e:
    print("REPRODUCED RuntimeError:", str(e).splitlines()[0])

Result:

# BEFORE (current main)
REPRODUCED RuntimeError: Expected all tensors to be on the same device, but got index is on cuda:0, different from other tensors on cpu (when checking argument in method wrapper_CUDA__index_select)

# AFTER (this PR)
PASSED: encode_prompts_flux returned embeds (1, 8, 16) pooled (1, 8)

[rajakshay]

Hi @jaretburkett
I'd appreciate it if you can point me in the right direction to find a reviewer for this pull request. Thanks!