UGPUM™ and pipeline parallelism — making personal multi-GPU diffusion training work on Windows

Personal LoRA training, large diffusion models, and the consumer-GPU gap

A LoRA — Low-Rank Adaptation — lets you teach a large diffusion model a new style, a new subject or a new concept by training only a small adapter on top of its frozen weights. For text-to-image and text-to-video models, LoRAs are how individual creators bring their own data into a model they didn’t train. The legal frame is straightforward: you’re running an open-weights model you’re entitled to use, on a dataset you’re entitled to use, for your personal creative work.

The technical frame got harder in 2025–2026. The diffusion transformer (DiT) inside modern open-weights video and image models keeps growing. Wan 2.1 14B sits at 28 GB in bf16. Flux 2 Dev is even larger. HunyuanVideo is in the same league. A single consumer GPU — even a top-tier RTX 5080 at 16 GB — cannot hold the weights of one of these DiTs in VRAM, let alone the gradients, the optimizer state and the activations a training pass needs.

The standard escape hatches all have a catch:

• Distributed Data Parallel (DDP), as implemented in most diffusion training frameworks, replicates the model on each GPU and shards the data. Useful when the model fits on one card and you want to multiply throughput; useless when the model itself doesn’t fit on one card.
• Block-swap / CPU offload keeps the DiT weights on CPU RAM and pages individual transformer blocks in and out of VRAM for each forward pass. The second GPU sits idle; the PCIe bus becomes the bottleneck.
• Cross-GPU sharding (the one tool that does pool VRAM across cards for diffusion training) is Linux-only or requires WSL2 with a measurable performance penalty, and breaks when the cards have asymmetric VRAM (which is the common case on consumer rigs — many users have a fast newer card paired with an older one).

DiffusionForce exists for the trough: a personal LoRA trainer that pools two consumer GPUs on native Windows, with heterogeneous VRAM and compute handled by the splitter, with the open-source kohya-ss / musubi-tuner training engine doing the actual data, scheduler and optimizer work underneath.

Why pipeline parallelism, not block-swap

The choice that drives everything else is where the DiT weights live during training.

Block-swap puts them all on CPU RAM and moves a block at a time into VRAM when the forward pass needs it. The work is sequential. The PCIe bus is the bottleneck. A second GPU adds nothing.

Pipeline parallelism puts blocks permanently on their assigned GPU. Block 0 lives on cuda:0 forever; block 22 lives on cuda:1 forever. The forward pass walks the chain: cuda:0 computes blocks 0–N, hands an activation tensor (~4 MB) to cuda:1 across PCIe, cuda:1 computes blocks N+1 to the end. The backward pass goes the other way. Both GPUs do real compute on real weights; only activations cross the bus.

For Wan 2.1 14B on a 5080 + 4060 Ti rig (16 + 16 = 32 GB usable), the math works: 28 GB of weights split across two cards, with the remaining 4 GB shared between activations, gradients, optimizer state and LoRA adapters. Tight, but feasible.

The catch is that the entire diffusion training ecosystem assumes either single-GPU work or DDP, not a third option. To make pipeline parallelism work on top of existing Musubi code without forking it, we needed to intercept a number of assumptions baked into PyTorch and Hugging Face Accelerate (which Musubi uses under the hood). More on that below.

The splitter — VRAM and compute, not just VRAM

The naive split assigns blocks proportionally to VRAM: 16 GB cards get 50/50, an 8+16 GB rig gets 33/67. That’s wrong for heterogeneous rigs.

Compute speed matters at least as much as VRAM. On a 5080 + 4060 Ti, the 5080 has roughly twice the TFLOPS of the 4060. If you split 50/50 by block count, the 5080 finishes its half of the forward pass and then waits for the 4060 to finish, every single step. The whole rig moves at the speed of the slower card.

The DiffusionForce splitter reads torch.cuda.get_device_properties() at startup, looks at the per-card TFLOPS budget and the per-card memory bandwidth, and computes a ratio that balances finish times. On the 5080 + 4060 Ti pair, that lands around 65% of blocks on the 5080 and 35% on the 4060 Ti for Wan 14B — fewer blocks per card but those blocks compute faster. The slow card stops being the floor; both cards finish their leg of the pass at roughly the same time.

A small VRAM reserve (≈1 GB per device, scaling adaptively with total VRAM in upcoming releases) protects against allocation spikes during activations and gradient computation. The splitter refuses to launch and reports a clear actionable message — “need ~X GB more, free Y GB from n” — if no valid split fits.

The hook layer — making Musubi behave without forking it

This is the part of the project that took most of the calendar time, and the part that’s hardest to describe without going into ML-Twitter weeds. The short version: PyTorch supports cross-device autograd natively, but the ecosystem above it (Hugging Face Accelerate, model loading conventions, the LoRA injection scheme used by kohya-ss/sd-scripts and Musubi) makes a number of assumptions that are correct for single-GPU or DDP but incorrect for single-process pipeline parallelism.

DiffusionForce overrides those assumptions through a layer of hooks — small, surgical Python patches that sit between Musubi and the underlying frameworks. Each hook addresses one specific assumption that, left as-is, makes the partition fall apart. The first ten hooks were the path to the first successful Wan 14B training step; the more recent ones (notably hook #12, the UGPUM™ scheduler itself, and hook #22, the per-prompt extra-LoRA injector for sample generation) handle features that emerged once the foundation worked.

A flavour of the kind of thing each hook handles, without going through all of them:

• CUDA_VISIBLE_DEVICES=0,1 exposed to the Python child process — without it, the splitter sees one GPU and falls back to single-GPU mode that doesn’t fit.
• Initialising torch.distributed with world_size=1, backend="gloo" so Accelerate stops assuming DDP when it sees more than one device.
• Forcing loading_device='cpu' on the DiT load, then routing blocks to their target GPU ourselves — otherwise Musubi tries to put the whole 28 GB DiT on cuda:0 first and triggers Windows’ shared-memory fallback (50× slowdown).
• Populating model.hf_device_map with our per-block placement, so Accelerator’s prepare() doesn’t silently .to(cuda:0) everything and undo the partition.
• A pre-forward hook that re-syncs every block to its assigned device just before the first forward pass, to catch LoRA adapters that Musubi attaches after the partition is built (Musubi’s LoRA scheme replaces org_module.forward and discards the reference, so the LoRA tensors don’t move with the block when you do block.to(device)).
• A VAE CPU/fp32 conversion for sample generation, because the CPU bf16 kernels in PyTorch don’t up-promote inputs the way the CUDA kernels do — a mismatch that crashes inside conv2() if you don’t catch it.

The list isn’t glamorous. It’s the kind of engineering that exists because each upstream project made a reasonable assumption that happened to not hold in our exact intersection. Each hook is small, named, isolated to one file, and documented with the exact failure mode it prevents. Musubi itself isn’t forked: it ships in the installer as a pinned vendored snapshot, untouched, alongside the hook layer that intercepts its behaviour at known integration points.

UGPUM™ — three modes for three different constraints

Pipeline parallelism alone is enough for a rig where the DiT, the gradients and the optimizer state all fit across the two cards combined. Wan 14B on 16 + 16 GB is the canonical case for that.

Two situations break the assumption:

1. Even larger models where the weights alone don’t fit across two consumer cards (Flux 2 Dev with its full text encoder is the open case). Pipeline parallelism is no longer sufficient on its own.
2. HD video sample generation during training. Generating a 1024×1024×9-frame MP4 sample at the end of an epoch needs more activation memory on cuda:0 than the training loop normally has free, even though training itself is comfortably profitable.

UGPUM™ — Unified GPU Memory — is the multi-mode memory scheduler that handles all three cases under one roof. At launch, it estimates static and dynamic VRAM budgets per GPU and picks the minimal mode that fits the job:

• Mode A — pipeline pure. The fastest mode. Blocks permanently on their GPUs, weights stay in VRAM, no CPU staging. This is what runs when everything fits.
• Mode B — pipeline + permanent CPU swap. For training models so large the weights don’t fit across the two cards even with pipeline. Some blocks live on CPU permanently and are paged in for each forward. Slower than Mode A but unlocks training that would otherwise be impossible. Deferred to V1.1+; the V1 release reaches it as a refusal rather than an automatic fallback so the user knows what they’re asking for.
• Mode C — pipeline + temporary CPU swap during sample only. The interesting case. Training runs in Mode A; just before a sample generation step, the scheduler unwraps the boundary block, pauses the pipeline hooks, moves N blocks to CPU to free room for the HD sample pass, runs the sample, then restores the blocks back to their target devices and resumes the hooks. After the sample, training continues in Mode A as if nothing happened.

Mode C is the one that took the most engineering — five successive bug fixes before it survived an end-to-end run. The reasons range from picking the right primary device for the sample (cuda:0 is where the sample latent input arrives, so blocks need to be there) to noticing that a one-shot re-sync hook from earlier in the pipeline was firing in the middle of the swap and restoring the partition at exactly the wrong moment.

The result, the night Mode C first worked: a 1024×1024 × 9-frame video sample produced during a Wan 2.1 14B + LoRA training pass on a 5080 + 4060 Ti rig, on native Windows. To our knowledge that exact combination — that resolution, that frame count, that DiT size, that consumer GPU pair, on Windows without WSL2 — hadn’t worked before.

What we deliberately don’t support in V1

A multi-GPU diffusion trainer has an unbounded surface. Some scope-cuts that landed for the V1 release:

• No inference path. DiffusionForce only does training. Inference is well-served by other open-source toolchains for these models and adding it would have meant a parallel set of memory-budget calculations.
• Exactly two GPUs. V1 splits across exactly two devices. Three-or-more rigs would multiply the splitter combinatorics and introduce activation-routing scenarios we haven’t tested.
• FP8 / int8 quantization for training is on the V2 list. We expose an inference-time fp8 toggle (Musubi downcasts the bf16 weights to fp8 in VRAM at load time, freeing ~7 GB per GPU) which is enough to unblock HD sample generation in Mode A on most rigs.
• Wan 2.1 14B and Wan 2.2 (high & low) are the validated V1 models. Flux 2 Klein 9B, Flux 1, LTX and HunyuanVideo all have entries in the registry but aren’t covered by the V1 end-to-end validation. Each new model family requires 1–2 sessions of hook-edge-case debug to validate against the partition.
• Single-GPU users get vanilla Musubi. Pipeline parallelism is only invoked when two devices are visible; otherwise the launcher hands the run off to the unmodified Musubi path, which works as well as it does on its own.

Credits and license

DiffusionForce vendors a pinned snapshot of kohya-ss / musubi-tuner (Apache 2.0). All the heavy lifting — the data loaders, the schedulers, the optimizer wrappers, the safetensors save/load, the model implementations themselves — comes from there. The kohya-ss project and its contributors did the diffusion-training reverse engineering that everything in this ecosystem is built on. The original Musubi LICENSE travels intact in the installer’s Resources/musubi-tuner-vendored/ directory.

The cross-device autograd is straight PyTorch. The single-process distributed init that lets Hugging Face Accelerate behave is the same library; Accelerate itself is from the Hugging Face team and the work it does on top of PyTorch is what makes the bf16 / FP8 / mixed-precision plumbing tractable in the first place.

What DiffusionForce adds, by line count: the pipeline-engine subpackage (splitter, activation router, DiT wrapper, the hook layer including UGPUM™ — roughly 2 000 lines of Python), the C# / WPF Windows installer and Run-tab UI, the uv-based Python environment management, and the integration of all the above so a non-CLI user gets a one-click Setup + a one-click Run.

Take-aways

• For training models that don’t fit on a single consumer GPU, pipeline parallelism is the right tool — but on Windows you have to do it yourself, single-process, because the cross-machine MPI plumbing the cluster world uses isn’t available without WSL2.
• VRAM-only splitters are wrong for heterogeneous rigs. Compute-aware partitioning is the difference between a working multi-GPU setup and one where the fast card waits for the slow card every step.
• Hook-based integration beats forking. Patching at known integration points keeps the open-source engine pristine and lets upstream improvements flow in.
• Memory scheduling for a training run is multi-modal by nature — training memory needs and sample-generation memory needs are different shapes. Building the scheduler around modes (A / B / C) instead of one monolithic budget made every other decision easier.