The Problem

Fine-tuning LLMs on AMD GPUs typically requires Linux and ROCm. Windows users with AMD cards are locked out of mainstream training stacks like Axolotl, which hard-depend on torch.cuda. This shim makes Windows + AMD viable by intercepting PyTorch's CUDA checks and routing compute to DirectML instead.

What This Does

dml_shim.py is a Python interception layer that spoofs CUDA availability to PyTorch. It patches torch._C._cuda_init() to a no-op, forces cuda.is_available() to return True, and redirects device calls to privateuseone:0, which maps to DirectML. Training software never sees the difference.

The patches/apply_all_patches.sh script automates the sed edits to torch/cuda/__init__.py that neutralize the CUDA initialization and device checks. The README documents manual verification steps.

How It Is Wired

The internal call graph is minimal: 1 internal module, 0 import edges, no circular dependencies. dml_shim has Ca=0 (nothing imports it) and Ce=0 (it imports nothing internally) — it's a standalone patch.

Entry point: Python startup via sitecustomize.py, which imports dml_shim. The shim hooks torch.cuda functions at import time. The critical path is:

  1. sitecustomize.py imports dml_shim
  2. dml_shim patches torch._C._cuda_init() → pass
  3. dml_shim patches _cuda_setDevice() → pass
  4. dml_shim forces cuda.is_available() → True
  5. Device calls route to privateuseone:0 (DirectML)

What it touches outside itself: PyTorch's CUDA module internals, DirectML runtime, and sitecustomize.py. The blast radius is contained — only dml_shim.py (Python) and apply_all_patches.sh (Shell) carry logic.

File-by-file map:

  • dml_shim.py — the interception layer; owns all the CUDA spoofing logic
  • patches/apply_all_patches.sh — applies the sed patches to torch/cuda/__init__.py
  • README.md — setup and verification instructions

How To Use It

Setup (from README, verbatim):

pip install torch==2.4.1 torchvision==0.19.1 torch-directml
# Copy dml_shim.py to site-packages, add import to sitecustomize.py
sed -i 's/torch._C._cuda_init()/pass/' path/to/torch/cuda/__init__.py
sed -i 's/torch._C._cuda_setDevice(device)/pass/' path/to/torch/cuda/__init__.py
sed -i 's/if device < 0 or device >= device_count():/if False:/' path/to/torch/cuda/__init__.py

Running it:

python3 -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

Expected output: True and your AMD card name.

Configuration: No config files or env vars. The shim is activated by the sitecustomize.py import.

Real-World Use

On a Windows 10 machine with an AMD RX 7700 XT, install WSL2 + Ubuntu 22.04, apply the patches, then run Axolotl as usual. The shim makes torch.cuda.is_available() return True, so Axolotl proceeds with training on the DirectML backend. The README reports successful Mistral 7B LoRA fine-tuning.

Code Health & Issues

Static analysis (measured, not opinion):

  • Medium — No test files — untested code paths — repository-wide
  • Medium — No CI/CD pipeline — no automated build/test gate — .github/ or CI config
  • Medium — No LICENSE file — unclear usage/redistribution rights — root

The README claims MIT license, but no LICENSE file is committed. The repo is 3 files total; no dependency lockfile, Dockerfile, or secrets detected.

The Bottom Line

This is a pragmatic hack that works for a narrow use case: Windows + AMD + Axolotl. It's fragile — it patches PyTorch internals at runtime and relies on specific torch versions. Use it for a quick experiment, not production. The missing LICENSE file is a real blocker for redistribution.