Skip to content

Add split-device text encoder routing - #38

Draft
Dillflix wants to merge 12 commits into
ideogram-oss:mainfrom
Dillflix:feature/split-text-diffusion-devices
Draft

Add split-device text encoder routing#38
Dillflix wants to merge 12 commits into
ideogram-oss:mainfrom
Dillflix:feature/split-text-diffusion-devices

Conversation

@Dillflix

Copy link
Copy Markdown

Summary

  • add backward-compatible text_device routing while retaining device as the diffusion-device API and alias
  • place Qwen3-VL on the text device and both diffusion transformers, VAE, latent normalization, sampler state, and decoder state on the diffusion device
  • run Qwen only over the common left-padded text block, transfer compact float32 text features once, and reconstruct zero image-token features on the diffusion device before denoising
  • add guarded device/memory/timing diagnostics plus direct GPU transfer with a CPU-staged fallback
  • add focused regression tests, an offline structured-caption smoke script, and experimental ROCm documentation

Motivation

The existing single-device pipeline loads Qwen3-VL alongside both diffusion transformers and the VAE. On memory-constrained diffusion GPUs, Qwen's one-time forward pass can exhaust the remaining VRAM before sampling begins. Keeping Qwen on a secondary GPU frees diffusion-device headroom while avoiding cross-device traffic inside the denoising loop.

Existing callers that omit text_device retain the original single-device behavior.

Validation

  • 8 CPU unit tests pass, covering default-device compatibility, component and input ownership, text-sequence trimming, left-padding semantics, full conditioning reconstruction, and transfer fallback behavior
  • Python compilation passes for the modified pipeline, tests, and smoke script
  • targeted Ruff formatting and lint checks pass
  • both worktree and range diff checks pass

Hardware validation status

The development host available for this patch is Windows with CPU-only PyTorch (HIP=None, zero visible GPUs), so model-weight and ROCm hardware validation has not been run here. The included smoke script performs device preflight, reports peer access and peak memory, accepts local structured captions, and supports the official V4_DEFAULT_20 preset for the target 1024x1024 acceptance run.

ROCm multi-GPU support is documented as experimental. The main remaining validation risks are bitsandbytes NF4 support across both gfx1100 and gfx1151 and any shape-dependent numerical differences from trimming trailing image placeholders.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant