fix: RoFormer on DirectML — CPU state-dict load + complex-op CPU hops (#292) - #295
Merged
Merged
Conversation
torch-directml's torch.load deserialization hook expects integer device
ids and raises TypeError ("'>=' not supported between instances of
'torch.device' and 'int'") when map_location targets a privateuseone
device. This made the new Roformer implementation fail on every DirectML
load and silently fall back to the legacy path (whose CPU map_location
got further, then died on ComplexFloat at inference — separate fix).
Load the state dict on CPU and move the model to the device afterwards,
gated on DML so cuda/mps/cpu keep byte-identical behavior.
Part 3 of the RoFormer-on-Windows plan (see #294 for the plan doc).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
torch-directml has no complex tensor support — torch.stft with return_complex, view_as_complex, complex multiplies and torch.istft all fail with 'Invalid or unsupported data type ComplexFloat'. Hop exactly those ops to CPU when the input lives on a privateuseone device, keeping the transformer stack (the heavy compute) on the DML device. This is the same pattern the MDX arch already uses in uvr_lib_v5/stft.py (which is why MDX works on DirectML today) and mirrors the existing MPS istft hop. All hops are gated on _is_dml_device(); cuda/mps/cpu behavior is unchanged — the MPS branches are untouched. The helper is module-level so tests can force the DML branches on CPU tensors: new equivalence tests assert the hopped path produces identical output to the normal path for both BSRoformer and MelBandRoformer. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
beveradb
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The two code fixes from the RoFormer-on-Windows plan (#292; plan doc in #294), both strictly gated on DirectML (
privateuseone) devices so CUDA/MPS/CPU behavior is byte-identical.Fix 1 — loader
TypeError(the "'>=' not supported between torch.device and int" error)RoformerLoader._load_with_new_implementationcalledtorch.load(map_location="privateuseone:0"); torch-directml's deserialization hook expects integer device ids and throws. Result: the new implementation failed on every DirectML load and silently fell back to legacy. Fix: DML-gatedmap_location="cpu"+ existingmodel.to(device)— the same pattern the non-Roformer MDXC path already uses.Fix 2 —
Invalid or unsupported data type ComplexFloattorch-directml has no complex tensor support.
bs_roformer/mel_band_roformernow hop exactly the complex ops (stft → view_as_real, view_as_complex + mask multiply/scatter, istft) to CPU on DML, keeping the transformer stack — the heavy compute — on the GPU. Same pattern as the MDX arch'suvr_lib_v5/stft.py(which is why MDX already works on DirectML) and the existing MPS istft hop. All MPS branches untouched.Testing
_is_dml_devicetruth tables; loader tests assertmap_locationis CPU for DML and unchanged for cpu/cuda/mps withmodel.to(device)still receiving the original device_is_dml_deviceis module-level so tests force the DML branches on CPU tensors (where every hop is a no-op) and assert bit-for-bit identical output vs the normal path for both models — proving the hop plumbing doesn't alter resultstests/unit+tests/contract: 301 passed, 4 skippedwindows-directmlCI job once the Windows runner image exists (feat(runner-manager): Windows GPU ephemeral runner family (gpu-windows) karaoke-gen#881); that job will assert new-implementation load (no legacy fallback) and CPU-vs-DML output qualityKnown remaining risk on real DML hardware:
F.scaled_dot_product_attention(flash path inattend.py) may be unsupported — if the first real-DML run shows that, the einsum fallback gets wired up in a follow-up.@coderabbitai ignore
🤖 Generated with Claude Code