TinyComputers.io (Posts about pytorch)

Processing 51,000 Photos with AI on AMD Strix Halo

A.C. Jokela — Sat, 14 Mar 2026 17:00:00 GMT

I have roughly 20 years of photos sitting on a home fileserver. They span 2001 to 2020, shot on everything from a Minolta DiMAGE F100 to a Nikon D5100 to various iPhones over the years. A mix of 21,554 JPEGs and 29,860 Nikon RAW files (51,414 images total) organized in a Lightroom backup directory by year, month, and date. Most were shot handheld, many in a hurry. The kind of archive that accumulates when you take photos for two decades without ever going back to curate them.

The Lightroom catalog that once made sense of all this was long gone, lost to a drive migration somewhere around 2018. What remained was a directory tree of raw files with no organization beyond the date folders. No star ratings, no keywords, no collections. Just files. Thousands of them, some sideways, some crooked, all unlabeled.

I wanted to fix that. Not manually (I don't have a month to spend in Lightroom) but programmatically. The goals were straightforward: correct orientation issues, straighten crooked horizons, generate AI descriptions of every photo's content, and catalog the whole archive in a queryable database. The kind of batch processing job that would have been impractical five years ago but is now entirely doable with the right hardware and a weekend of scripting.

The Hardware

Two machines on the local network, each with a distinct role:

Machine	Role	Key Specs
Fileserver	NAS / photo storage	28TB RAID (`/md0`), 125GB RAM, NFS exports
GPU workstation	ML inference	AMD Ryzen AI Max+ 395, Radeon 8060S, 121GB RAM

The fileserver is a straightforward storage box. The interesting machine is the GPU workstation running an AMD Strix Halo APU, specifically the AI Max+ 395 with its integrated Radeon 8060S. I've written about this chip before, and it continues to impress for inference workloads. The RDNA 3.5 integrated GPU shares system memory, giving it access to 65.2 GB of VRAM without the typical constraints of a discrete card. For a model like BLIP that needs maybe 2 GB, that's absurdly generous, but it means you never have to think about VRAM budgets, which is a luxury when you're iterating on a processing pipeline.

The fileserver already had NFS configured, exporting /md0 to the local subnet. One mount command on the GPU workstation and both machines could see the same filesystem:

sudo mount -t nfs fileserver.localnet:/md0 /md0

No file copying, no rsync scripts, no staging directories. The photos live on the NAS and get processed in-place over the network. Gigabit Ethernet introduces some I/O overhead (each 25 MB NEF file takes 200–300ms to read across the wire), but for an overnight batch job, the simplicity of a single shared filesystem is worth the throughput trade-off. If this were a recurring workflow, I'd invest in 10GbE, but for a one-time archive processing run, gigabit got it done.

The Software Stack

Everything runs in a Python virtual environment on the GPU workstation:

PyTorch 2.9.1+rocm6.3: ML framework with AMD ROCm backend
BLIP (Salesforce/blip-image-captioning-large): vision-language model for image captioning
OpenCV 4.13: horizon detection via Canny edge detection and Hough transforms
rawpy 0.26.1: Nikon NEF/NRW decoding (wraps LibRaw)
piexif: EXIF metadata extraction for JPEGs
exiftool: EXIF extraction for RAW files (called as a subprocess)
SQLite: metadata and results database

The gfx1151 Situation

If you've followed my previous posts on Strix Halo, you know the drill. The Radeon 8060S reports as gfx1151 in ROCm, which is newer than what PyTorch's ROCm wheels officially target. The fix is the same environment variable override:

export HSA_OVERRIDE_GFX_VERSION=11.0.0

This maps the GPU to a generic gfx11 target. In practice, it works without issues, with no compute errors and no performance penalties. ROCm 6.16 on this machine also reports amdgcn-amd-amdhsa--gfx11-generic as a supported ISA, which is likely why the override works cleanly. I've been running production workloads with this flag for months now without incident.

The Processing Pipeline

Each photo passes through five stages: EXIF extraction, orientation correction, horizon detection and straightening, AI captioning, and finally saving the corrected image and cataloging everything in SQLite.

EXIF Metadata Extraction

For JPEGs, piexif reads the embedded EXIF data directly; it's a pure Python library that parses the binary EXIF structure without needing any external dependencies. For NEF/NRW files, piexif can't handle Nikon's proprietary container format, so I shell out to exiftool with JSON output (exiftool -json -n <file>). The -n flag is important; it returns numeric values instead of human-readable strings, which makes downstream processing much cleaner.

The extracted fields cover the full gamut: camera make and model, lens, dates, exposure settings (shutter speed, aperture, ISO, focal length), flash, white balance, metering mode, GPS coordinates, and the original orientation tag.

EXIF data is notoriously inconsistent across two decades of cameras. I'll come back to this; it became a debugging story of its own.

Orientation Correction

The EXIF orientation tag (values 1 through 8) encodes how the camera was held when the photo was taken. A value of 1 means the image is right-side up. A value of 6 means the camera was rotated 90 degrees clockwise. Value 3 means 180 degrees. Some values encode horizontal or vertical flips. The full matrix looks like this:

ops = {
    2: [Image.FLIP_LEFT_RIGHT],
    3: [Image.ROTATE_180],
    4: [Image.FLIP_TOP_BOTTOM],
    5: [Image.FLIP_LEFT_RIGHT, Image.ROTATE_270],
    6: [Image.ROTATE_270],
    7: [Image.FLIP_LEFT_RIGHT, Image.ROTATE_90],
    8: [Image.ROTATE_90],
}

Out of the 51,411 successfully processed photos, 8,797 (17.1%) needed orientation correction. The majority came from the Nikon D5100 and iPhone 4, both of which set the orientation tag but don't bake the rotation into the pixel data itself. Without this correction, nearly one in five photos would display sideways or upside-down in any viewer that doesn't respect EXIF orientation.

Here's what that looks like in practice. The raw pixel data from this iPhone photo is stored sideways; the camera recorded an EXIF orientation tag of 6, meaning "rotate 90 degrees clockwise to display correctly." Any viewer that ignores that tag renders the image on its side:

Before: raw pixel data (EXIF orientation 6, displayed sideways)

After: orientation corrected based on EXIF tag

Horizon Detection and Straightening

This stage uses classical computer vision, no neural network needed. The approach:

Downscale the image to 1200px on the long side for speed
Convert to grayscale, apply Gaussian blur
Run Canny edge detection
Crop to the vertical middle 50%, since the horizon is rarely at the extreme top or bottom of a frame
Apply the Hough Line Transform to find line segments, requiring a minimum length of one-quarter the image width
Filter to near-horizontal lines (within 20 degrees of level)
Compute a weighted average of the detected angles, weighted by line length

The key is the threshold window. If the detected angle is less than 0.5 degrees, it's not worth correcting, since you'd introduce interpolation artifacts for no visible benefit. If it's greater than 15 degrees, it's probably not a tilted horizon at all; it's either intentional composition or the algorithm latching onto a staircase railing. The correction itself uses cv2.warpAffine with Lanczos interpolation and a reflective border mode, followed by an inward crop to eliminate any border artifacts:

crop_factor = 1.0 / (cos(angle) + sin(angle) * (min(h, w) / max(h, w)))

The initial implementation used Canny edge detection and Hough line transforms, classical CV techniques from the 1980s. Fast, deterministic, 100ms per image. But it had a fatal flaw: it couldn't distinguish between a tilted horizon and a roofline receding toward a vanishing point. Architecture, roads, staircases, any strong line in the middle band of the image would register as a "tilted horizon," and the algorithm would dutifully rotate the image to "correct" it. In practice, this meant a significant number of photos were being made worse, not better.

The fix was to replace Hough line detection with semantic segmentation. SegFormer (nvidia/segformer-b2-finetuned-ade-512-512), trained on the ADE20K dataset, segments each image into 150 classes, including sky. The approach is simple: find the sky pixels, trace the bottom edge of the sky region, fit a line to that boundary, and measure its angle. If there's no sky (less than 5% of the image), or the sky boundary is too fragmented (fewer than 20 points), skip the correction entirely.

This eliminates false positives on indoor shots, close-ups, architecture, and anything without a visible sky. SegFormer runs on CPU at about 0.4 seconds per image; the model is only 25M parameters, so it doesn't need the GPU. The GPU stays dedicated to BLIP captioning.

Two examples from the corrected archive. This bridge over a river had a 2.68-degree clockwise tilt, and the bridge deck and far shore are visibly leveled:

Before: 2.68° clockwise tilt

After: horizon straightened via sky boundary detection

This rocky Lake Superior shore had a 3.85-degree clockwise tilt, and the far horizon is leveled:

Before: 3.85° clockwise tilt

After: horizon straightened

AI Captioning with BLIP

The Salesforce/blip-image-captioning-large model generates natural language descriptions of each photo. It runs in float16 on the Radeon 8060S. Each image is resized to a maximum of 1024px before inference. Beam search with 5 beams and a 75-token limit generates the caption:

output_ids = model.generate(
    **inputs,
    max_new_tokens=75,
    num_beams=5,
    early_stopping=True,
)

Caption inference takes about 0.5–0.7 seconds per image, consistent regardless of whether the input was a JPEG or a decoded NEF. The model handles a wide variety of subjects surprisingly well. Some examples from the archive:

"a brown and white dog standing next to a blue chair"
"two silos sitting in the middle of a field"
"a bird sitting on a branch of a tree"
"a wooden sign that says hoban road in front of some trees"
"a blurry photo of a car driving down a snowy road"
"a dog being groomed by a woman in a salon"

The captions tend toward a "there is a..." pattern, and they occasionally get details wrong (BLIP once described a photo of my living room as "a hotel lobby," which is generous). But for searchability and cataloging purposes, they're remarkably useful. Being able to query WHERE caption LIKE '%dog%' across 51,000 photos and get meaningful results is something that would have required manual tagging before models like BLIP existed. For an archive this size, "good enough" captions on every photo are vastly more useful than perfect captions on none of them.

Save and Catalog

Corrected images are saved as high-quality JPEGs (quality 92) to /md0/photos_processed/images/, mirroring the original directory structure. NEF and NRW files are converted to JPEG in the process; the corrected archive is a uniform format. All metadata flows into a SQLite database with WAL journaling, tracking 40+ fields per photo: every piece of EXIF data, processing flags (was orientation corrected? was the horizon straightened? by how many degrees?), the AI caption, file hashes, dimensions, and processing timestamps.

The database makes the archive queryable in ways that were never possible before:

-- What cameras did I use, and when?
SELECT camera_model, MIN(date_taken), MAX(date_taken), COUNT(*)
FROM photos GROUP BY camera_model ORDER BY COUNT(*) DESC;

-- Photos with GPS data
SELECT filename, caption, gps_latitude, gps_longitude
FROM photos WHERE gps_latitude IS NOT NULL;

-- How crooked were my photos, by camera?
SELECT camera_model, ROUND(AVG(ABS(horizon_angle_degrees)), 2) as avg_tilt, COUNT(*)
FROM photos WHERE horizon_corrected = 1
GROUP BY camera_model ORDER BY avg_tilt DESC;

The EXIF Tuple Bug

The first processing pass completed 51,414 photos, but with 2,146 errors. All of them were TypeError: type tuple doesn't define __round__ method. For a pipeline that had been running cleanly on thousands of Nikon D5100 and D60 photos, this was unexpected.

The root cause turned out to be a two-part problem with how certain budget cameras from the 2008–2012 era write EXIF rational numbers.

Part 1: Malformed Tuples

The EXIF standard stores rational numbers as (numerator, denominator) pairs. Most cameras follow this. But some, particularly a batch of older point-and-shoots, wrote the ExposureBiasValue field as a 4-element tuple like (36, 0, 18, 0) instead of the expected 2-element (36, 0).

My _rational_to_float helper only handled 2-tuples:

def _rational_to_float(val):
    if isinstance(val, tuple) and len(val) == 2:
        if val[1] == 0:
            return None
        return val[0] / val[1]
    return val  # passes through 4-tuples as raw tuples

When a 4-tuple fell through, the downstream round() call choked on it. The fix was simple: return None for any tuple that isn't a standard rational pair.

Part 2: None Propagation

Even after fixing Part 1, many of these same cameras had written (36, 0), a rational with a zero denominator. The function correctly returned None for division by zero, but the calling code then did round(None, 2), triggering the same TypeError with a slightly different message.

The fix was a _safe_round wrapper:

def _safe_round(val, digits=1):
    if val is None:
        return None
    try:
        return round(val, digits)
    except TypeError:
        return None

After both fixes, the second pass recovered all 2,143 photos. The remaining 3 errors were genuine file corruption: a truncated JPEG, a NEF that LibRaw couldn't parse, and a NEF with filesystem-level I/O errors. Probably bad sectors on the source drive. Those can't be fixed in code.

This is one of those bugs that only surfaces at scale. Run the pipeline on a hundred Nikon photos and everything works perfectly. Run it on 51,000 photos spanning 15 different camera models over 20 years, and every edge case in the EXIF spec comes out to play. The lesson, which I should have internalized long ago: never trust external data formats at scale without defensive parsing on every field. The EXIF spec is a suggestion, not a contract, and camera manufacturers have been interpreting it creatively since the early 2000s.

Resumability

A 15-hour batch job will inevitably need to be restarted: bugs, system updates, a random hound disconnects the magsafe power cord from my MacBook Pro. The script tracks progress in SQLite and skips completed files on restart:

def is_already_processed(conn, source_path):
    row = conn.execute(
        "SELECT id FROM photos WHERE source_path = ? AND error IS NULL",
        (source_path,),
    ).fetchone()
    return row is not None

Photos that failed with errors are intentionally not skipped, so fixing a bug and re-running automatically retries them. This made the EXIF debugging cycle painless: fix the parser, clear the failed rows from the database, relaunch, and only the 2,143 affected photos get reprocessed.

Performance

The pipeline sustained 1.0–1.8 photos per second, depending on file format:

Stage	Time per Photo
JPEG load	~10ms
NEF decode (rawpy)	~400ms
MD5 hash	~5ms (JPEG), ~100ms (NEF)
Horizon detection	~100ms
BLIP inference	~500–700ms
JPEG save	~50ms

BLIP inference dominates the runtime. NEF decoding is the second bottleneck; each RAW file is 20–30 MB and requires full demosaicing through LibRaw. The NFS overhead for reading large NEFs over gigabit Ethernet is noticeable but not the primary constraint.

Total wall time: 15.5 hours across two passes for 51,414 photos. The BLIP model uses roughly 2 GB of the 65.2 GB available VRAM on the Strix Halo. Memory was never a concern.

Final Results

Metric	Count	Percentage
Total photos	51,414	100%
Successfully processed	51,411	99.99%
Orientation corrected	8,797	17.1%
Horizon straightened	15,251	29.7%
AI captioned	51,411	99.99%
Unrecoverable errors	3	0.006%

The top cameras in the archive tell the story of 20 years of gear:

Camera	Photos
Nikon D5100	24,073
Nikon D60	8,734
iPhone 4	2,664
Nikon D3100	1,698
Panasonic DMC-FX07	975
Minolta DiMAGE F100	870
iPad	803
iPhone 5s	698
Samsung SCH-I500	645

The output lives on the NAS:

Corrected images: /md0/photos_processed/images/, 51,411 JPEGs preserving the original year/month/date folder structure, all NEFs converted, all orientation and horizon corrections applied.
SQLite database: /md0/photos_processed/photos.db, 40+ fields per photo with full EXIF metadata, processing results, and AI-generated captions.
Processing log: /md0/photos_processed/processing.log, timestamped record of the entire run.

Takeaways

AMD's Strix Halo continues to earn its keep for ML inference. The HSA_OVERRIDE_GFX_VERSION=11.0.0 workaround remains necessary, but once set, PyTorch and ROCm run without complaints. The 65 GB shared VRAM pool means you can load models without thinking about memory budgets, a workflow advantage that's easy to underestimate until you've experienced it.

Classical computer vision still has its place. The horizon detection pipeline uses Canny edge detection and Hough transforms, techniques from the 1980s. No training data, no GPU needed, deterministic results, and the whole thing runs in 100ms per image. For geometric corrections on photographic images, you don't need a neural network. You need line detection.

EXIF is a minefield. Twenty years of cameras from different manufacturers means every edge case in the spec gets exercised. Tuple lengths vary, denominators are zero, fields are missing or repurposed. If you're parsing EXIF at scale, assume nothing about the data's shape and validate everything.

Resumability is non-negotiable for long-running jobs. Tracking progress in the database and skipping completed work made it trivial to iterate on bugs. Without this, every fix would mean reprocessing 51,000 photos from scratch.

NFS over gigabit is fine for batch processing. Not optimal, but for an overnight job, the network overhead from NAS-attached storage is acceptable. The real bottleneck was ML inference at 0.6 seconds per photo. If I were doing this regularly, 10GbE would be worth the upgrade, but for a one-time archive processing run, gigabit got the job done.

The whole project, from first SSH to final database entry, took about a day of wall time, most of which was unattended processing. The scripting itself was maybe three hours of work. Twenty years of photos, cataloged and corrected overnight. Not bad for a Strix Halo and some Python. The full source is available on GitHub.

What I didn't expect was how useful the database would be after the fact. Being able to ask "show me every photo I took with the D5100 at ISO 3200 or higher" or "find photos with GPS data from 2015" turns a pile of files into something that actually tells a story. The AI captions add another dimension; I can now search my own photo archive by content, not just metadata. It's the kind of capability that makes you wonder why photo management software hasn't done this for years. The models have been available. The hardware has been affordable. Someone just needed to wire it together.

Upgrading ROCm 7.0 to 7.2 on AMD Strix Halo (gfx1151)

A.C. Jokela — Wed, 18 Feb 2026 16:00:00 GMT

Introduction

If you're running AMD's Strix Halo hardware -- specifically the Ryzen AI MAX+ 395 with its integrated Radeon 8060S GPU -- you already know the software ecosystem is a moving target. The gfx1151 architecture sits in an awkward spot: powerful hardware that isn't officially listed on AMD's ROCm support matrix, yet functional enough to run real workloads with the right driver stack. When ROCm 7.2 landed in early 2026, upgrading from 7.0.2 was a priority. The newer stack brings an updated HSA runtime, a refreshed amdgpu kernel module, and broader compatibility improvements that matter on bleeding-edge silicon.

This post documents the complete upgrade procedure from ROCm 7.0.2 to 7.2 on a production Ubuntu 24.04 system. It's not a theoretical exercise -- this was performed on a live server running QEMU virtual machines and network services, with the expectation that everything would come back online after a single reboot.

AMD's official documentation states that in-place ROCm upgrades are not supported. The recommended path is a full uninstall followed by a clean reinstall. That's exactly what we did, and the entire process took about 20 minutes of wall-clock time (excluding the reboot).

System Overview

The target system is a Bosgame mini PC running the Ryzen AI MAX+ 395 APU. If you've read the earlier review of this hardware, you'll be familiar with the specs. For context on this upgrade, here's what matters:

Hardware

CPU: AMD Ryzen AI MAX+ 395, 16 cores / 32 threads, Zen 5
GPU: Integrated Radeon 8060S, 40 Compute Units, RDNA 3.5 (gfx1151)
Memory: 32 GB DDR5, unified architecture with 96 GB allocatable to GPU
Peak GPU Clock: 2,900 MHz

Software (Pre-Upgrade)

OS: Ubuntu 24.04.3 LTS (Noble Numbat)
Kernel: 6.14.0-37-generic (HWE, pinned)
ROCm: 7.0.2
amdgpu-dkms: 6.14.14 (from repo.radeon.com/amdgpu/30.10.2)
ROCk Module: 6.14.14

Running Services

The system was actively serving several roles during the upgrade:

Five QEMU virtual machines (three x86, two aarch64)
A PXE boot server (dnsmasq) for the local network
Docker daemon with various containers

None of these services are tied to the GPU driver stack, so the plan was to perform the upgrade and reboot without shutting them down first. The VMs and network services would come back automatically after the reboot.

Why Upgrade

ROCm 7.0.2 worked on this hardware. Models loaded, inference ran, rocminfo detected the GPU. So why bother upgrading?

Three reasons:

Driver maturity for gfx1151: The amdgpu kernel module jumped from 6.14.14 to 6.16.13 between the two releases. That's not a minor revision -- it represents months of kernel driver development. On hardware that isn't officially supported, newer drivers tend to bring meaningful stability improvements as AMD's internal teams encounter and fix issues on adjacent architectures.
HSA Runtime improvements: ROCm 7.2 ships HSA Runtime Extension version 1.15, up from 1.11 in ROCm 7.0.2. The HSA (Heterogeneous System Architecture) runtime is the lowest layer of the ROCm software stack -- it handles device discovery, memory management, and kernel dispatch. Improvements here affect everything built on top of it.
Ecosystem alignment: PyTorch wheels, Ollama builds, and other ROCm-dependent tools increasingly target 7.2 as the baseline. Running 7.0.2 was becoming an exercise in version pinning and compatibility workarounds.

The Kernel Hold: Why It Matters

Before diving into the procedure, a note on kernel management. This system runs the Ubuntu HWE (Hardware Enablement) kernel, which provides newer kernel versions on LTS releases. At the time of this upgrade, the HWE kernel was 6.14.0-37-generic. The upstream kernel had already moved to 6.17, but we didn't want the ROCm upgrade to pull in a kernel that AMD's DKMS module might not build against.

The solution is apt-mark hold:

sudo apt-mark hold linux-generic-hwe-24.04 linux-headers-generic-hwe-24.04 linux-image-generic-hwe-24.04

This prevents apt from upgrading the kernel meta-packages, effectively pinning the system to 6.14.0-37-generic. The hold was already in place before the upgrade and remained untouched throughout. After the upgrade, we confirmed it was still active:

apt-mark showhold

linux-generic-hwe-24.04
linux-headers-generic-hwe-24.04
linux-image-generic-hwe-24.04

If you're running Strix Halo or any other hardware where kernel compatibility with amdgpu-dkms is uncertain, kernel holds are essential. A kernel upgrade that breaks the DKMS build means no GPU driver after reboot.

Upgrade Procedure

Step 1: Uninstall the Current ROCm Stack

AMD provides the amdgpu-uninstall script for exactly this purpose. It removes all ROCm userspace packages and the amdgpu-dkms kernel module in a single operation:

sudo amdgpu-uninstall -y

This command removed approximately 120 packages, including the full HIP runtime, rocBLAS, MIOpen, MIGraphX, ROCm SMI, the LLVM-based compiler toolchain, and the Mesa graphics drivers that ship with ROCm. The DKMS module was purged, which means the amdgpu kernel module was removed from the 6.14.0-37-generic kernel's module tree.

After the ROCm stack was removed, we purged the amdgpu-install meta-package itself:

sudo apt purge -y amdgpu-install

This also cleaned up the APT repository entries that amdgpu-install had configured in /etc/apt/sources.list.d/. The old repos -- repo.radeon.com/amdgpu/30.10.2, repo.radeon.com/rocm/apt/7.0.2, and repo.radeon.com/graphics/7.0.2 -- were all removed automatically.

Step 2: Clean Up Leftover Files

The package removal was thorough but not perfect. A few leftover directories remained in /opt/:

ls /opt/ | grep rocm

rocm-7.0.0
rocm-7.0.2
rocm-7.9.0

The rocm-7.0.0 directory was from a previous installation attempt. The rocm-7.9.0 was from an earlier experiment with a release candidate build. The rocm-7.0.2 directory contained a single orphaned shared library (libamdhip64.so.6) that dpkg couldn't remove because the directory wasn't empty. All three were cleaned up manually:

sudo rm -rf /opt/rocm-7.0.0 /opt/rocm-7.0.2 /opt/rocm-7.9.0

It's worth checking for stale ROCm directories after any uninstall. They consume negligible disk space but can confuse build systems and scripts that scan /opt/rocm* for active installations.

Step 3: Install the ROCm 7.2 Installer

AMD distributes ROCm through a meta-package called amdgpu-install. Each ROCm release has its own version of this package, which configures the appropriate APT repositories. The 7.2 installer was downloaded directly from AMD's repository:

cd /tmp
wget https://repo.radeon.com/amdgpu-install/7.2/ubuntu/noble/amdgpu-install_7.2.70200-1_all.deb
sudo apt install -y ./amdgpu-install_7.2.70200-1_all.deb
sudo apt update

After installation and apt update, three new repositories were active:

https://repo.radeon.com/amdgpu/30.30/ubuntu noble -- the kernel driver and Mesa components
https://repo.radeon.com/rocm/apt/7.2 noble -- the ROCm userspace stack
https://repo.radeon.com/graphics/7.2/ubuntu noble -- graphics libraries

The version numbering can be confusing. The amdgpu-install package version is 30.30.0.0.30300000-2278356.24.04, which maps to the amdgpu driver release 30.30. The ROCm version is 7.2.0. These are different version tracks that AMD maintains in parallel.

Step 4: Install ROCm 7.2

With the repositories configured, the actual installation was a single command:

sudo amdgpu-install -y --usecase=graphics,rocm

The --usecase=graphics,rocm flag tells the installer to include both the Mesa graphics drivers and the full ROCm compute stack. This is the right choice for a system that needs both display output and GPU compute capabilities.

The installation took approximately 10 minutes and included:

amdgpu-dkms 6.16.13: The kernel module, compiled via DKMS against the running kernel
Full ROCm 7.2 stack: HIP runtime, hipcc compiler, rocBLAS, rocFFT, MIOpen, MIGraphX, RCCL, ROCm SMI, ROCProfiler, and dozens of other libraries
Mesa graphics: Updated EGL, OpenGL, and Vulkan drivers from the amdgpu Mesa fork
ROCm LLVM toolchain: The LLVM-based compiler infrastructure that HIP uses for kernel compilation

The DKMS build is the critical step. During installation, DKMS compiled the amdgpu module against the kernel headers for 6.14.0-37-generic. The output confirmed a successful build:

depmod...
update-initramfs: Generating /boot/initrd.img-6.14.0-37-generic

The initramfs was regenerated to include the new module, ensuring it would be loaded at boot.

Step 5: Verify DKMS

Before rebooting, we confirmed the DKMS status:

dkms status

amdgpu/6.16.13-2278356.24.04, 6.14.0-37-generic, x86_64: installed
virtualbox/7.0.16, 6.14.0-36-generic, x86_64: installed
virtualbox/7.0.16, 6.14.0-37-generic, x86_64: installed
virtualbox/7.0.16, 6.8.0-100-generic, x86_64: installed

The new amdgpu module (6.16.13) was built and installed for 6.14.0-37-generic. Note that it only built for the currently running kernel, unlike VirtualBox which had modules built for older kernels as well. This is expected -- DKMS builds against available kernel headers, and the old kernel headers for 6.14.0-36 and 6.8.0-100 were still present from earlier installations.

Step 6: Reboot

sudo reboot

The server came back online in approximately 50 seconds.

Post-Reboot Verification

rocminfo

The first check after reboot was rocminfo, which queries the HSA runtime for available agents:

rocminfo

ROCk module version 6.16.13 is loaded
=====================
HSA System Attributes
=====================
Runtime Version:         1.18
Runtime Ext Version:     1.15
...
==========
HSA Agents
==========
Agent 1: AMD RYZEN AI MAX+ 395 w/ Radeon 8060S (CPU)
Agent 2: gfx1151 (GPU)
  Marketing Name:          AMD Radeon Graphics
  Compute Unit:            40
  Max Clock Freq. (MHz):   2900
  Memory Properties:       APU
  ISA 1: amdgcn-amd-amdhsa--gfx1151
  ISA 2: amdgcn-amd-amdhsa--gfx11-generic

Key observations:

ROCk module 6.16.13: The new kernel module loaded successfully.
Runtime Ext Version 1.15: Upgraded from 1.11 in ROCm 7.0.2.
gfx1151 detected: The GPU was recognized with its correct ISA identifier.
gfx11-generic ISA: ROCm 7.2 also exposes a generic gfx11 ISA, which allows software compiled for the broader RDNA 3 family to run on this device without gfx1151-specific builds.
APU memory: The memory properties correctly identify this as an APU with unified memory.

ROCm SMI

rocm-smi

Device  Node  Temp    Power     SCLK  MCLK     Fan  Perf  VRAM%  GPU%
0       1     33.0C   9.087W    N/A   1000Mhz  0%   auto  0%     0%

The GPU was visible and reporting telemetry. The 0% VRAM reading is expected on an APU -- rocm-smi reports dedicated VRAM usage, but on a unified memory architecture, GPU memory allocations come from system RAM and aren't reflected in this counter.

ROCm Version

cat /opt/rocm/.info/version

7.2.0

DKMS

dkms status

Confirmed amdgpu/6.16.13 remained installed for 6.14.0-37-generic after reboot.

PyTorch Validation

With the driver stack verified, the next step was confirming that PyTorch could see and use the GPU. ROCm 7.2 ships with prebuilt PyTorch wheels on AMD's repository.

Installing PyTorch for ROCm 7.2

We set up a Python virtual environment and installed the ROCm-specific wheels:

python3 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip

The PyTorch wheel for ROCm 7.2 requires a matching ROCm-specific build of Triton. Both are available from AMD's manylinux repository. The order matters -- Triton must be installed first, since the PyTorch wheel declares it as a dependency with a specific version that doesn't exist on PyPI:

pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp312-cp312-linux_x86_64.whl
pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp312-cp312-linux_x86_64.whl
pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp312-cp312-linux_x86_64.whl

These are the ROCm 7.2 builds for Python 3.12. AMD also provides wheels for Python 3.10, 3.11, and 3.13.

Smoke Test

import torch
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("Device:", torch.cuda.get_device_name(0))
print("VRAM:", round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1), "GB")

PyTorch: 2.9.1+rocm7.2.0.git7e1940d4
CUDA available: True
Device: AMD Radeon Graphics
VRAM: 103.1 GB

PyTorch detected the GPU through ROCm's HIP-to-CUDA translation layer. The 103.1 GB figure represents the total addressable memory on this unified-memory APU, which includes both the 96 GB GPU allocation and additional system memory accessible through the HSA runtime.

Note the use of torch.cuda despite this being an AMD GPU. ROCm's HIP runtime presents itself through PyTorch's CUDA interface, so all CUDA API calls in PyTorch (device selection, memory management, kernel launches) work transparently with AMD hardware.

Before and After Summary

Component	ROCm 7.0.2	ROCm 7.2.0
ROCm Version	7.0.2	7.2.0
amdgpu-dkms	6.14.14	6.16.13
ROCk Module	6.14.14	6.16.13
HSA Runtime Ext	1.11	1.15
amdgpu Repo	30.10.2	30.30
PyTorch	N/A	2.9.1+rocm7.2.0
Triton	N/A	3.5.1+rocm7.2.0
Kernel	6.14.0-37-generic	6.14.0-37-generic
Kernel Holds	In place	In place

Notes on gfx1151 Support

It's worth being explicit about the support situation. As of February 2026, gfx1151 (Strix Halo) is not listed on AMD's official ROCm support matrix. The supported RDNA 3 targets are gfx1100 (Navi 31, RX 7900 XTX) and gfx1101 (Navi 32). Strix Halo's gfx1151 is an RDNA 3.5 derivative that shares much of the ISA with gfx1100 but has architectural differences in the memory subsystem and compute unit layout.

In practice, ROCm 7.2 works on gfx1151. The kernel driver loads, rocminfo detects the GPU, and PyTorch can allocate tensors and dispatch compute kernels. The gfx11-generic ISA target in ROCm 7.2 is particularly helpful -- it provides a compatibility path for software that hasn't been explicitly compiled for gfx1151.

However, "works" and "fully supported" are different things. There are known quirks:

rocm-smi VRAM reporting: Always shows 0% on the APU since it only tracks discrete VRAM
No official PyTorch gfx1151 builds: The ROCm PyTorch wheels target gfx1100. They run on gfx1151 through ISA compatibility, but performance may not be optimal
Large model loading latency: Moving large models to the GPU device can be slow on the unified memory architecture, as the HSA runtime handles page migration differently than discrete GPU DMA transfers

If you're considering this hardware for production AI workloads, treat ROCm support as "functional but experimental." It works well enough for development, testing, and moderate inference workloads. For production training or latency-sensitive deployment, stick with hardware on AMD's official support list.

Rollback Plan

If the upgrade fails -- the DKMS module doesn't build, the GPU isn't detected after reboot, or something else goes wrong -- the rollback path is straightforward:

Uninstall ROCm 7.2:

sudo amdgpu-uninstall -y
sudo apt purge -y amdgpu-install

Reinstall ROCm 7.0.2:

wget https://repo.radeon.com/amdgpu-install/30.10.2/ubuntu/noble/amdgpu-install_30.10.2.0.30100200-2226257.24.04_all.deb
sudo apt install -y ./amdgpu-install_30.10.2.0.30100200-2226257.24.04_all.deb
sudo apt update
sudo amdgpu-install -y --usecase=graphics,rocm
sudo reboot

The entire rollback takes about 15 minutes. Keep the old amdgpu-install deb URL handy -- it's not linked from AMD's current download pages once a newer version is published.

Conclusion

Upgrading ROCm on hardware that isn't officially supported always carries some risk, but this upgrade from 7.0.2 to 7.2 on gfx1151 was uneventful. The procedure follows AMD's documented uninstall-reinstall approach with no deviations. The kernel hold strategy kept the kernel stable, the DKMS module built cleanly against 6.14.0-37-generic, and all post-reboot checks passed.

The improvements in ROCm 7.2 -- particularly the HSA runtime bump to 1.15 and the introduction of the gfx11-generic ISA target -- represent meaningful progress for Strix Halo users. The ecosystem is slowly catching up to the hardware. It's not there yet, but each release closes the gap.

For anyone running a Ryzen AI MAX+ 395 or similar Strix Halo hardware on Ubuntu 24.04, this upgrade is worth doing. The procedure is well-defined, the rollback path is clear, and the newer driver stack brings tangible benefits. Just remember to hold your kernel first.

Recommended Resources

Hardware

Bosgame M5 AI Mini PC (Ryzen AI MAX+ 395) - The system used in this post
GMKtec EVO X2 (Ryzen AI MAX+ 395) - Another Strix Halo mini PC option on Amazon

Books

Deep Learning with PyTorch by Stevens, Antiga, Huang, Viehmann - Comprehensive guide to building, training, and tuning neural networks with PyTorch
Programming PyTorch for Deep Learning by Ian Pointer - Practical guide to creating and deploying deep learning applications
Understanding Deep Learning by Simon Prince - Modern treatment of deep learning fundamentals

Image Editing on 10-Year-Old GPUs: NVIDIA P40 vs AMD Strix Halo

A.C. Jokela — Tue, 17 Feb 2026 18:00:00 GMT

Introduction

There's a certain satisfaction in making old hardware do new tricks. When NVIDIA released the Tesla P40 in 2016, deep learning was still finding its footing. ImageNet classification was the benchmark everyone cared about, GANs were generating blurry faces, and the idea of a 57-billion-parameter image editing model would have seemed like science fiction.

Around the middle of 2017, when the P40 would have been seeing peak adoption in datacenters, I found myself in an advanced pattern recognition course, my final credits needed for a masters in computer science (the name hadn't been updated to reflect more contemporary terminology like "machine learning," let alone "deep learning"). The textbook was Bishop's Pattern Recognition and Machine Learning, a book that managed to make Bayesian inference feel both rigorous and approachable. We spent the last two weeks of the course looking at deep learning using TensorFlow, but we didn't even have GPU infrastructure available. Everything ran on CPU. It would have been great to have experienced the P40 in its prime, when 24 GB of VRAM and 3,840 CUDA cores made it one of the most capable inference GPUs money could buy. Instead, I'm getting acquainted with it a decade later, asking it to do things its designers never imagined.

Fast forward to 2026, and here I am, running a 57-billion-parameter model on four of these decade-old GPUs, and comparing the results against AMD's latest Strix Halo APU, a chip that didn't exist until 2025.

The model in question is FireRed-Image-Edit-1.0 from FireRedTeam, a 57.7GB diffusion model built on the QwenImageEditPlusPipeline architecture. It takes an input image and a text prompt, then produces an edited version. The kind of thing that would have required a massive cloud GPU a couple of years ago.

This post documents the full journey: the precision pitfalls of running modern diffusion models on Pascal-era GPUs, the quantization trade-offs that make or break image quality, and the head-to-head performance comparison that produced some genuinely surprising results. All of the inference scripts and output images are available on GitHub.

The Hardware

NVIDIA Tesla P40 (2016)

The P40 was NVIDIA's inference-focused datacenter GPU from the Pascal generation. The key specs for our purposes:

Architecture: Pascal (sm_6.1)
CUDA Cores: 3,840
Memory: 24 GB GDDR5X
Memory Bandwidth: 346 GB/s
FP32 Performance: 12 TFLOPS
FP16 Performance: Limited, no native FP16 tensor cores
BF16 Support: None
Price today: ~$100-200 per card on the secondary market

I have four of these cards in a server, giving me 96 GB of total VRAM, but spread across four separate memory spaces, which introduces its own challenges.

AMD Ryzen AI MAX+ 395 / Strix Halo (2025)

AMD's Strix Halo is a different beast entirely. It's an APU (CPU and GPU on the same die, sharing the same memory pool):

GPU Architecture: RDNA 3.5 (gfx1151)
Compute Units: 40 CUs
Memory: 128 GB unified LPDDR5X (32 GB for CPU, 96 GB for VRAM)
Memory Bandwidth: ~256 GB/s (shared)
BF16 Support: Yes
FP16 Support: Yes (Fast F16 Operation)
ROCm: 7.9.0
Price: ~$2,000+ for the complete system

The unified memory architecture means all 96 GB is accessible to the GPU without any PCIe transfer overhead, and the entire model can live in a single memory space.

The Model: FireRed-Image-Edit-1.0

FireRed-Image-Edit is a diffusion-based image editing model with three major components:

Component	Size	Description
Transformer	40.9 GB	QwenImageTransformer2DModel, 60 layers
Text Encoder	16.6 GB	Qwen2.5-VL 7B vision-language model
VAE	~0.3 GB	AutoencoderKL for encoding/decoding images

Total: 57.7 GB of model weights. The scheduler is FlowMatchEulerDiscreteScheduler, and the pipeline uses true classifier-free guidance (CFG), which roughly doubles the memory needed during inference since it runs both conditional and unconditional passes.

The test task: take this input image and apply the prompt "Add a red hat on the cat"; the model draws a cat wearing a red hat onto the book cover, rendered in the style of the O'Reilly animal illustrations.

The P40 Challenge: When FP16 Breaks Everything

The Precision Problem

The first, and biggest, challenge with the P40s is numerical precision. Modern diffusion models are designed for BF16 (bfloat16), which has the same exponent range as FP32 (8 exponent bits, range ±3.4×10³⁸) but with reduced mantissa precision. The P40, being a Pascal-era GPU, supports neither BF16 nor proper FP16 tensor operations.

FP16 has only 5 exponent bits, giving it a range of ±65,504. This might seem sufficient, but the diffusion scheduler's internal sigma values and the VAE's convolution operations routinely produce intermediate values that overflow this range. The FlowMatchEulerDiscreteScheduler, in particular, works with sigma schedules that can produce large intermediate values during the noise prediction and scaling steps. When these overflow FP16's limited range, they become NaN or Inf, and these corrupt values propagate through every subsequent operation (matrix multiplications, attention computations, residual connections) until the entire tensor is garbage.

The result: NaN propagation that silently corrupts the entire pipeline, producing an all-black output image.

This was the most time-consuming discovery in the entire project. The model would load, the progress bar would advance through all 40 denoising steps without any indication of trouble, and then the output would be perfectly black: mean=0.0, min=0, max=0. No error messages. No warnings. No NaN detection exceptions. Just silent numerical corruption that only becomes visible when you look at the final image.

The debugging process was particularly frustrating because the corruption happens gradually. Partial NaN contamination in early steps doesn't crash anything; the attention mechanisms and residual connections continue to produce tensor outputs of the expected shapes. The model appears to be working normally right up until the final image is decoded from all-zero latents.

The FP32 Solution (and a Speed Surprise)

The fix was to run the entire pipeline in FP32: scheduler, VAE, and all non-quantized transformer layers. The quantized weights themselves stay compressed (INT8 or NF4), but every arithmetic operation uses full 32-bit precision.

It wasn't enough to just set the quantization compute dtype to FP32; that only fixes the dequantized matmul operations inside the quantized layers. The scheduler's sigma arithmetic, the VAE's convolution operations, and the non-quantized components (layer norms, biases, attention scaling) all needed FP32 as well. Similarly, loading the pipeline with torch_dtype=torch.float32 but leaving the transformer's non-quantized layers in FP16 caused a dtype mismatch in the attention mechanism; PyTorch's scaled dot-product attention requires query, key, and value tensors to share the same dtype. Every component in the computational chain needed to be FP32.

The one exception is the text encoder, which runs once before the denoising loop begins. It stays in FP16 on its own GPU, and its output embeddings are upcast to FP32 when transferred to the main device. This is safe because the text encoder doesn't participate in the iterative process where precision errors compound.

Here's where things got interesting: FP32 was actually faster than FP16 on the P40. The first attempts with FP16 ran at approximately 9 minutes per denoising step. After switching to FP32, the same operations completed in about 2.4 minutes per step with NF4, and 1.5 minutes per step with INT8. The P40's FP32 throughput is its native strength; it was designed for FP32 datacenter inference, after all. FP16 on Pascal is handled through slower pathways that add overhead rather than saving it.

Multi-GPU Device Orchestration

With 57.7 GB of model weights and only 24 GB per GPU, some form of model sharding or quantization is mandatory. After extensive testing, the optimal configuration for the P40s turned out to be:

GPU 0: INT8-quantized transformer (~22 GB)
GPU 1: Text encoder in FP16 (~16.6 GB)
GPU 2: VAE in FP32 (~6.6 GB including decode workspace)
GPU 3: Unused

This layout requires significant monkey-patching of the diffusers pipeline. The _execution_device property must be overridden to ensure latents are created on the correct GPU. The encode_prompt method needs patching to route inputs to the text encoder's GPU and move the resulting embeddings back. And for the INT8 configuration, the VAE's encode and decode methods need wrappers to handle cross-device tensor transfers.

The text encoder stays in FP16 because it fits on a single GPU and its outputs are immediately upcast to FP32 when moved to the main device. This is safe because the text encoder runs once at the beginning; it doesn't participate in the iterative denoising loop where precision matters most.

Quantization Quality: INT8 vs NF4

With the FP32 pipeline in place, I tested both INT8 (8-bit) and NF4 (4-bit) quantization for the transformer:

NF4 (4-bit quantization): The NF4 approach uses bitsandbytes' normalized float 4-bit quantization with double quantization enabled. The transformer compresses from 40.9 GB to roughly 10 GB, easily fitting on a single P40 alongside the VAE. However, the output quality was significantly degraded, with heavy noise and grain throughout the image, even at the full 40 denoising steps. Each denoising step introduces small numerical errors from the 4-bit weight approximations, and these errors compound across 40 iterations.

INT8 (8-bit quantization): INT8 produced dramatically better results. The output was clean and sharp, visually comparable to what you'd expect from full-precision inference on a modern GPU. The 8-bit precision preserves enough information in the weights that the per-step errors don't accumulate into visible artifacts.

The trade-off is memory: the INT8 transformer occupies ~22 GB, nearly filling an entire P40. This is why the VAE had to move to a third GPU; there wasn't enough headroom on GPU 0 for the VAE's convolution workspace during the decode phase. An early attempt that kept the VAE on GPU 0 ran all 40 denoising steps successfully, only to crash with an out-of-memory error at the very last operation.

The Strix Halo Experience: Simplicity Wins

BF16 Full Precision

Running the same model on the Strix Halo was refreshingly simple. With 96 GB of unified VRAM and native BF16 support, the entire pipeline loads in a few lines:

pipe = QwenImageEditPlusPipeline.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda:0")

No quantization. No multi-GPU patching. No device transfer hooks. No FP32 workarounds. The model loads in BF16 and runs natively.

During inference, the pipeline consumed approximately 75 GB of VRAM (the true CFG doubles the workspace requirements), well within the 96 GB budget.

The first run did take about 35 minutes of JIT kernel compilation before producing any inference steps; ROCm compiles HIP kernels for the gfx1151 architecture on first use. During this phase, the GPU sits at 100% utilization with no visible progress, which can be alarming if you're not expecting it. The GPU temperature climbed from 31°C idle to 69°C, and power draw went from 9W to 119W as the compiler worked through the hundreds of unique kernel configurations needed by a 60-layer transformer. These compiled kernels are cached, so subsequent runs skip this overhead entirely.

Quantization on Strix Halo: Does It Help?

Given the surprising performance parity between the two systems at full precision, I tested whether quantization could speed up the Strix Halo by reducing memory traffic. The theory was that if the workload is memory-bandwidth-limited, smaller model weights should mean faster inference.

The results were definitive:

Configuration	Per-Step Time	40-Step Estimate	VRAM Used
BF16 (full precision)	82.6s	55 min	~75 GB
NF4 (4-bit)	83.5s	56 min	~30 GB
INT8 (8-bit)	94.9s	63 min	~44 GB

NF4 quantization produced virtually identical speed to full BF16. The model shrank from 75 GB to 30 GB of VRAM usage, but inference time didn't improve at all. INT8 was actually slower; the bitsandbytes INT8 matmul path adds dequantization overhead that more than offsets any memory bandwidth savings.

This tells us something important about the Strix Halo's performance profile for this workload: it's compute-bound, not memory-bound. The RDNA 3.5 GPU's 40 compute units are the bottleneck, not the LPDDR5X memory bandwidth. Reducing the model size doesn't help because the GPU is already busy with arithmetic, not waiting on memory.

This contrasts with LLM inference workloads (text generation), where the Strix Halo's large memory pool is a genuine advantage. LLM token generation is almost entirely memory-bound, making quantization directly translate to speed improvements. Each token generation pass reads the entire model's weights but performs relatively little computation per weight. Diffusion models are the opposite: each denoising step runs a full forward pass through 60 transformer layers with dense matrix multiplications, attention computations, and residual connections. The arithmetic intensity is much higher, putting the pressure squarely on the GPU's compute units rather than its memory subsystem.

Head-to-Head: The Numbers

Here's the complete performance comparison across all tested configurations:

System	Configuration	Per-Step	40 Steps	Image Quality
Strix Halo	BF16 full precision	82.6s	55 min	Clean, sharp
Strix Halo	NF4 (4-bit)	83.5s	56 min	Clean (10-step test)
Strix Halo	INT8 (8-bit)	94.9s	63 min	Clean (10-step test)
4× P40	INT8 + FP32 pipeline	87.5s	58 min	Clean, sharp
4× P40	NF4 + FP32 pipeline	145.9s	97 min	Heavy noise/grain

The headline result: a single AMD Strix Halo APU from 2025 is about 6% faster per step than four NVIDIA P40s from 2016 running INT8-quantized inference. That's not exactly the generational leap you might expect from a decade of GPU evolution.

To be fair, the comparison isn't entirely apples-to-apples. The P40 is running an 8-bit quantized model (less computation per step but with dequantization overhead), while the Strix Halo runs the full BF16 model. The P40's dedicated GDDR5X provides 346 GB/s of bandwidth to a single GPU, while the Strix Halo's LPDDR5X shares its ~256 GB/s between the CPU and GPU. And the P40 setup requires three GPUs working in concert, while the Strix Halo uses a single unified memory space.

Lessons Learned

Old GPUs Are Surprisingly Capable

Four P40s at ~$500 total produce inference quality and speed that's competitive with a $2,000+ modern APU system. The P40's 346 GB/s memory bandwidth per card and strong FP32 throughput remain relevant even for models that were designed for hardware two generations newer. The main challenge is software engineering: working around the precision limitations and multi-GPU complexity takes significant effort.

Precision Matters More Than Speed

The single most impactful discovery in this project was that FP16 silently corrupts diffusion model outputs on Pascal GPUs. There are no error messages, no NaN warnings during inference, just a black image at the end. The fix (using FP32 everywhere) actually improved performance, which was counterintuitive. The lesson: when dealing with older hardware, always validate your numerical precision assumptions before optimizing for speed.

Quantization Is Not Free

On the P40s, INT8 quantization was essential (the model simply wouldn't fit otherwise) and produced excellent results. NF4 was too aggressive; the 4-bit precision degraded output quality visibly.

On the Strix Halo, quantization was unnecessary and even counterproductive. INT8 added overhead without any speed benefit, and NF4 didn't save time despite dramatically reducing memory usage. The takeaway: quantization's value depends entirely on your bottleneck. If you're compute-bound, smaller weights don't help.

Unified Memory Is Underrated

The Strix Halo's greatest advantage wasn't raw performance; it was simplicity. Loading a 57.7 GB model into a single 96 GB memory space eliminates an entire category of engineering problems: no device placement, no cross-GPU tensor transfers, no monkey-patching encode/decode methods, no VAE OOM surprises at the decode step. The inference script for the Strix Halo is about 50 lines. The P40 version is over 150, most of it careful device orchestration code.

For anyone who values development velocity and code maintainability over squeezing the last dollar of cost-efficiency out of used datacenter hardware, unified memory APUs have a compelling argument even when they don't win on raw throughput.

What About Newer NVIDIA GPUs?

It's worth putting these numbers in context. An NVIDIA RTX 4090 with 24 GB of VRAM and native BF16/FP16 tensor core support would likely run this model (with INT8 quantization) at roughly 10-15 seconds per step, 5-8x faster than either system tested here. An A100 with 80 GB could run it unquantized in BF16 at similar or better speeds. The P40 and Strix Halo are both firmly in the "budget/accessible" tier of AI hardware.

The more interesting comparison is cost-per-step. Four P40s from eBay cost about $500 total (plus a server that can host them). The Strix Halo system runs about $2,000+. Both produce essentially the same result at the same speed. The P40 route demands more engineering knowledge; the Strix Halo route demands more money.

Conclusion

Both systems successfully ran a 57.7 GB diffusion model that would have been considered impossibly large for consumer hardware just a few years ago. The P40s did it through clever quantization and multi-GPU orchestration. The Strix Halo did it by brute force: 96 GB of memory and native BF16 support.

The performance story is more nuanced than "newer is always better." For diffusion model inference, the NVIDIA P40 (a card you can buy for $100 on eBay) remains remarkably competitive when properly configured. It requires more engineering effort, and you need to know the precision pitfalls, but the results speak for themselves.

The Strix Halo's strength lies not in raw speed but in its unified memory architecture and modern instruction set support. It eliminates the multi-GPU complexity entirely, runs native BF16 without precision hacks, and provides a development experience that's orders of magnitude simpler. For iterating on models, testing new architectures, or just avoiding the headaches of cross-device tensor management, that simplicity has real value.

If you're considering hardware for running large diffusion models locally, the choice comes down to how you value your time versus your budget. Four P40s and a weekend of debugging will get you to roughly the same place as a Strix Halo system that just works out of the box. Both paths lead to a cat in a red hat.

Recommended Resources

Hardware

NVIDIA Tesla P40 24GB - The GPU used in this post. Available on eBay for a fraction of the original price. You'll need a server with PCIe x16 slots and adequate cooling.
GMKtec EVO-X2 (AMD Ryzen AI MAX+ 395) - A compact Strix Halo mini PC with 128GB unified LPDDR5X 8000MHz, WiFi 7, and USB4. A representative platform for running large models on Strix Halo.

Books

Pattern Recognition and Machine Learning by Christopher M. Bishop - The classic that introduced many to Bayesian methods and kernel machines. Still one of the best foundations for understanding the statistical principles behind modern ML.
Hands-On Generative AI with Transformers and Diffusion Models by Omar Sanseviero, Pedro Cuenca, Apolinário Passos, and Jonathan Whitaker - A practical guide to building and fine-tuning diffusion models using the Hugging Face ecosystem, including the diffusers library used in this post.
Understanding Deep Learning by Simon J.D. Prince - A thorough modern treatment of deep learning fundamentals through diffusion models, with excellent visualizations and mathematical rigor.

Partial LLM Loading: Running Models Too Big for VRAM

A.C. Jokela — Thu, 05 Feb 2026 16:00:00 GMT

What happens when you want to run a 70B parameter model but only have 24GB of VRAM? Traditionally, you either quantize aggressively, rent cloud GPUs, or accept that the model is simply out of reach. But there's a third option that's becoming increasingly viable: partial loading, where you keep some layers on the CPU or disk and stream them to the GPU on demand.

I spent a couple days testing partial loading strategies on an AMD Strix Halo APU with 128GB of unified memory, configured with 96GB allocated to VRAM, trying to answer a simple question: can you actually run models that don't fit in VRAM, and if so, how much performance do you sacrifice?

The answer turns out to be: yes, you can, and the performance penalty is more nuanced than I expected.

The Memory Problem

Large language models are memory hogs. A 7B parameter model in bfloat16 needs roughly 14GB just for the weights. A 70B model needs 140GB. An 80B MoE model might need 160GB or more. Most consumer GPUs max out at 24GB, with only a handful of prosumer cards reaching 48GB.

The traditional solutions each have trade-offs:

Quantization reduces memory requirements by storing weights in lower precision formats. INT8 cuts memory in half. INT4 cuts it to a quarter. But quantization also reduces quality, sometimes significantly for complex reasoning tasks.

Model sharding across multiple GPUs works if you have multiple GPUs. Most people don't. Early on (e.g. two years ago), this is how I experimented with models, a handful of Pascal chipset NVIDIA GPUs in a former crypto mining server.

Cloud inference works but adds latency, costs money per token, and means your data leaves your machine.

Partial loading offers a fourth path: keep the model weights somewhere other than VRAM (CPU RAM, disk, NVMe) and load them into the GPU only when needed. You take a latency hit on every layer that needs to be fetched, but you can run models that would otherwise be impossible.

Understanding Transformer Layer Architecture

To understand why partial loading works, you need to understand how transformers process information. A typical LLM consists of:

Embedding layer: Converts input tokens to vectors. Relatively small.
Decoder layers: The bulk of the model. A 70B parameter model might have 80+ decoder layers, each containing attention heads and a feed-forward network (FFN).
Final layer norm and output projection: Converts the final hidden states back to token probabilities. Relatively small.

The key insight is that inference is sequential through the layers. When processing a token, you go through layer 0, then layer 1, then layer 2, and so on. You never need layer 5 while you're processing layer 3. This means you can theoretically keep only one layer's weights in VRAM at a time, loading the next layer while processing the current one.

In practice, keeping all layers streaming from disk adds too much latency. The sweet spot is typically keeping some layers resident in VRAM (usually the first few and last few, which see the most traffic) while streaming the middle layers on demand.

The Hardware Setup

My test machine is an AMD Strix Halo APU:

Hardware Configuration:
- AMD Radeon 8060S (gfx1151)
- 128GB unified memory (96GB VRAM / 32GB system)
- ROCm 7.0 with HSA_OVERRIDE_GFX_VERSION=11.0.0
- PyTorch 2.9.1+rocm6.3
- NVMe SSD: Samsung 990 Pro 2TB (PCIe 4.0, 7450 MB/s read)

The unified memory architecture is interesting for this experiment. On a discrete GPU, moving weights from CPU RAM to VRAM requires crossing the PCIe bus, which tops out at around 64GB/s for PCIe 4.0 x16. On the Strix Halo APU, both "GPU memory" and "CPU memory" share the same physical RAM; it's just a question of which pages are mapped for GPU access.

This should give partial loading an advantage on APUs, since there's no physical data movement, just page table updates. The actual numbers bear this out, as we'll see.

Three Approaches to Partial Loading

I tested three different strategies for loading models that exceed VRAM:

1. llama.cpp with Partial GPU Offloading

The simplest approach uses llama.cpp's -ngl (number of GPU layers) flag. This lets you specify exactly how many transformer layers go on the GPU, with the rest staying on CPU.

./main -m models/llama-70b-chat.gguf -ngl 35 \
    -p "The capital of France is" -n 50

With a 70B model that has 80 layers, setting -ngl 35 puts roughly 44% of the model on the GPU and 56% on CPU. The GPU handles the compute-intensive matrix multiplications, while the CPU layers run on AMD's Zen cores.

Advantages:

Simple to configure
Automatic handling of which layers go where
Works with GGUF quantized models
CPU layers use optimized AVX-512 implementations

Disadvantages:

Static partitioning: layers stay where they're assigned
CPU inference is much slower than GPU
Limited to llama.cpp's supported architectures

2. HuggingFace Accelerate Disk Offloading

HuggingFace's Accelerate library provides device_map="auto" with disk offloading:

from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import infer_auto_device_map

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-70B-Instruct",
    device_map="auto",
    offload_folder="./offload",
    torch_dtype=torch.bfloat16
)

When VRAM is insufficient, Accelerate automatically spills layers to disk (or CPU RAM if you use offload_buffers=True). During inference, layers are loaded as needed.

Advantages:

Works with any HuggingFace model
Automatic layer placement decisions
Can use disk for infinite capacity
Integrates with the broader HuggingFace ecosystem

Disadvantages:

Disk I/O is slow (even NVMe)
Layer loading happens synchronously
Each token generation can require full model traversal
Memory peaks during layer swaps

3. oLLM Layer Streaming

The most sophisticated approach I tested was oLLM, a library designed specifically for layer-by-layer streaming from SSD to GPU. Unlike HuggingFace's approach, oLLM implements asynchronous layer prefetching: while one layer is processing on the GPU, the next layer is being loaded.

from ollm import Inference

o = Inference("llama3-1B-chat", device="cuda:0")
o.ini_model(models_dir="./models/")

# Offload half the layers to CPU
o.offload_layers_to_cpu(layers_num=8)

# Generate
response = o.generate("The capital of France is", max_tokens=50)

The library instruments each layer load, giving you visibility into the streaming behavior:

layer_load: [0.002, 0.004, 0.004, 0.004, 0.004] t:0.242

This tells you each layer took about 4ms to load, and the total token generation time was 242ms.

Advantages:

Asynchronous prefetching reduces latency
Per-layer timing instrumentation
Designed specifically for memory-constrained scenarios
Can leverage GPU Direct Storage (GDS) for faster NVMe-to-GPU transfers

Disadvantages:

Limited model architecture support
Requires transformers 4.x (incompatible with 5.0)
Less mature than llama.cpp or HuggingFace

Benchmarking Partial Loading

I ran a series of tests with Llama 3.2 1B (16 layers, 2.8GB model size) to measure the impact of partial loading:

Test Configuration

model_id = "llama3-1B-chat"
prompt = "The capital of France is"
max_tokens = 50
configurations = [
    {"gpu_layers": 16, "cpu_layers": 0},   # Full GPU
    {"gpu_layers": 12, "cpu_layers": 4},   # 75% GPU
    {"gpu_layers": 8, "cpu_layers": 8},    # 50% GPU
    {"gpu_layers": 4, "cpu_layers": 12},   # 25% GPU
]

Results: oLLM Layer Streaming

With the oLLM library and 8 of 16 layers offloaded to CPU:

Configuration: 8 GPU layers, 8 CPU layers
Model loading: 2.3 seconds
First token latency: 242ms
Per-layer load time: ~4ms average
Output quality: Correct ("The capital of France is Paris.")

The layer load times are interesting. At 4ms per layer, you might expect significant overhead when 8 layers need to be fetched from CPU RAM. But because oLLM prefetches the next layer while the current one is executing, the effective latency impact is much smaller.

On a discrete GPU with PCIe transfers, these numbers would be different. Loading a 200MB layer across PCIe 4.0 x16 takes about 3ms at full bandwidth. But PCIe rarely achieves full bandwidth due to protocol overhead, so real-world numbers are typically 4-6ms per layer, similar to what I measured on the APU.

The Quality Question

A critical question with partial loading: does offloading layers affect output quality?

The answer is no, with an important caveat. Partial loading doesn't change the weights; it just changes where they're stored. The same matrices participate in the same computations. The outputs are bit-identical to full GPU inference.

The caveat is that some partial loading implementations use reduced precision for CPU layers (FP32 instead of bfloat16, or even FP16) to speed up CPU computation. This can introduce small numerical differences. In my testing with oLLM, both GPU and CPU layers used the same bfloat16 precision, so outputs matched exactly.

Practical Performance Analysis

Let's break down what partial loading actually costs in terms of latency.

Layer Loading Overhead

For a model with N layers, where K layers are on CPU: - Each token generation requires K layer loads - If each load takes T_load milliseconds - The total added latency per token is approximately K * T_load

With oLLM's prefetching, the effective latency is lower because loads overlap with computation. In my tests: - K = 8 layers on CPU - T_load = 4ms per layer - Naive overhead = 32ms per token - Actual overhead (with prefetching) = ~10-15ms per token

Memory Bandwidth Bottleneck

The real constraint isn't CPU speed; it's memory bandwidth. A single transformer layer in a 70B model might be 800MB-1.2GB. Loading this from:

NVMe SSD: 7.4GB/s = 108-162ms per layer
DDR5 RAM: 80GB/s = 10-15ms per layer
PCIe 4.0 x16: 32GB/s = 25-37ms per layer (in practice)

This is why oLLM's authors recommend fast NVMe SSDs (Samsung 990 Pro, WD SN850X) and ideally GPU Direct Storage, which bypasses the CPU entirely for disk-to-GPU transfers.

Token Generation Speed Comparison

For the Llama 3.2 1B model (16 layers total), I ran benchmarks across multiple prompts and averaging the results:

Configuration	Avg Tokens/sec	Avg Inference Time	Load Time
Full GPU (0 offloaded)	1.92 tok/s	13.90s	0.46s
4 layers on CPU	2.23 tok/s	11.09s	0.55s
8 layers on CPU	2.26 tok/s	10.87s	0.65s
12 layers on CPU	3.36 tok/s	7.30s	0.75s

Wait, that's backwards from what you'd expect. More layers on CPU resulted in faster inference?

This counterintuitive result makes sense when you consider the Strix Halo's unified memory architecture. Unlike a discrete GPU where CPU-to-GPU transfers cross the PCIe bus, the APU's "CPU memory" and "GPU memory" are the same physical RAM. Moving layers between them is essentially just a page table operation, not a data copy.

The performance improvement with more offloading likely comes from reduced memory bandwidth contention. When all layers are "on GPU," they're competing for the same memory channels. With layer streaming, only the active layer's weights occupy high-bandwidth GPU memory paths, while inactive layers sit in lower-priority memory regions.

This finding suggests that on unified memory systems (AMD APUs, Apple Silicon), partial loading might actually be preferable to full GPU loading for memory-bandwidth-bound workloads. The conventional wisdom (that GPU is always faster) doesn't hold when there's no physical separation between GPU and CPU memory.

Transformer Version Compatibility Issues

One challenge I encountered was library compatibility. oLLM was designed for transformers 4.x, and when I initially ran it with transformers 5.0, I hit several errors:

TypeError: 'Qwen3NextExperts' object is not iterable

This error occurred because transformers 5.0 changed how Mixture of Experts (MoE) layers expose their expert modules. The oLLM library's layer streaming code assumed it could iterate over self.mlp.experts, but the new implementation uses a different structure.

There were also weight shape mismatches:

model.layers.0.input_layernorm.weight: found shape torch.Size([2048])
in the checkpoint and torch.Size([0]) in the model instantiated

This happened because oLLM creates placeholder layers with zero-size tensors to save memory, then loads the actual weights on demand. The new transformers version changed how these placeholder shapes were inferred.

The solution was straightforward: pin transformers to version 4.57.6:

pip install 'transformers<5.0'

This is a common pattern with cutting-edge ML libraries. The ecosystem moves fast, and specialized tools often lag behind major version updates.

Storage and System Requirements

Before diving into partial loading, it's worth understanding the storage requirements. Unlike full GPU loading where you only need enough VRAM, partial loading requires sufficient storage capacity and bandwidth.

Disk Space Calculations

Model files on disk are typically stored in safetensors or GGUF format. A rough calculation: - 7B model (bfloat16): ~14GB - 13B model (bfloat16): ~26GB - 70B model (bfloat16): ~140GB - 70B model (GGUF Q4_K_M): ~40GB

For oLLM's layer streaming, you also need the model to be split into per-layer shards, which the library handles automatically during the first load. This adds temporary storage overhead during the conversion process.

RAM Requirements

CPU offloading means the offloaded layers live in system RAM. If you're offloading 40 of 80 layers from a 70B model, you need roughly 70GB of system RAM available, in addition to whatever the operating system and other applications need.

On my Strix Halo system with 128GB unified memory (96GB allocated to VRAM, 32GB to system), this gets interesting. The "CPU" portion of memory and the "GPU" portion share the same physical DIMMs. Allocating layers to "CPU" really just means they're in a different memory region that the GPU can still access, but through a different (slower) path.

SSD Endurance Considerations

If you're streaming weights from disk rather than CPU RAM, consider your SSD's endurance. A 70B model with 80 layers means moving roughly 1.75GB per token generated (all layers traversed once). Generate 1000 tokens and you've read 1.75TB from the SSD.

For occasional use, this is fine. For continuous operation (like a chatbot running 24/7), you might wear out a consumer SSD within months. Enterprise SSDs with higher TBW (Terabytes Written) ratings are worth considering for heavy use cases, or preferring CPU RAM offloading over disk offloading.

Memory Mapping and Page Tables

Under the hood, partial loading relies on the operating system's memory management. When a layer is "loaded" to the GPU, this typically involves:

Reading the layer from storage (if not already in RAM)
Pinning the memory pages so they can't be swapped
Mapping the pages into GPU-accessible memory space
Synchronizing to ensure the GPU sees the updated data

On Linux, this uses mmap() and mlock() syscalls. The vm.max_map_count sysctl may need to be increased for very large models:

# Check current value
cat /proc/sys/vm/max_map_count

# Increase if needed
sudo sysctl -w vm.max_map_count=1048576

I hit this limit when testing 70B+ models and saw cryptic "cannot allocate memory" errors until increasing the map count.

When Partial Loading Makes Sense

Based on my testing, here's when partial loading is a good fit:

Good Use Cases:

Batch processing where latency isn't critical (overnight analysis, embedding generation)
Interactive use with smaller models where the overhead is manageable
Running larger models occasionally without investing in more VRAM
Testing different model sizes before committing to hardware
APU systems where CPU-GPU transfer costs are minimal

Poor Use Cases:

Real-time applications (chatbots, live transcription)
High-throughput production systems
When quantization gives acceptable quality with lower overhead
Systems with slow storage (spinning disks, older SSDs)

The break-even point depends heavily on your specific hardware. On my APU system with unified memory, offloading 50% of layers costs about 33% of throughput. On a discrete GPU with PCIe 3.0, the same configuration might cost 60-70%.

Future Directions

Several developments could make partial loading more practical:

GPU Direct Storage (GDS): NVIDIA's GDS and AMD's equivalent allow direct SSD-to-GPU transfers, bypassing the CPU and PCIe. Early implementations show 3-4x improvements in layer load times.

Better Prefetching Algorithms: Current implementations use simple next-layer prefetching. More sophisticated approaches could predict multiple layers ahead, or prioritize layers that are accessed most frequently (relevant for some architectures with skip connections).

Hardware Evolution: Unified memory architectures like Apple Silicon and AMD APUs eliminate the CPU-GPU transfer bottleneck entirely. As these architectures gain more memory capacity, partial loading becomes increasingly attractive.

Compression: Applying neural compression to stored weights (not quantization, but actual neural codecs) could reduce the bandwidth requirements by 2-4x without quality loss.

Building a Benchmark Framework

For those who want to measure partial loading on their own hardware, here's the framework I developed:

from dataclasses import dataclass
from typing import List, Optional
import time
import torch

@dataclass
class BenchmarkResult:
    model_name: str
    total_layers: int
    gpu_layers: int
    cpu_layers: int
    loading_mode: str  # 'partial' or 'full'
    load_time_seconds: float
    tokens_generated: int
    total_inference_time: float
    tokens_per_second: float
    per_layer_load_times: Optional[List[float]] = None
    peak_vram_gb: Optional[float] = None

def benchmark_ollm_inference(
    model_id: str,
    offload_layers: int,
    prompt: str,
    max_tokens: int = 50
) -> BenchmarkResult:
    from ollm import Inference

    # Measure loading time
    load_start = time.perf_counter()
    o = Inference(model_id, device="cuda:0", logging=False)
    o.ini_model(models_dir="./models/", force_download=False)

    total_layers = len(o.model.model.layers)

    if offload_layers > 0:
        o.offload_layers_to_cpu(layers_num=offload_layers)

    load_time = time.perf_counter() - load_start

    # Measure inference time
    torch.cuda.synchronize()
    infer_start = time.perf_counter()

    output = o.generate(prompt, max_tokens=max_tokens)

    torch.cuda.synchronize()
    infer_time = time.perf_counter() - infer_start

    # Count tokens
    tokens = len(o.tokenizer.encode(output))

    return BenchmarkResult(
        model_name=model_id,
        total_layers=total_layers,
        gpu_layers=total_layers - offload_layers,
        cpu_layers=offload_layers,
        loading_mode='partial' if offload_layers > 0 else 'full',
        load_time_seconds=load_time,
        tokens_generated=tokens,
        total_inference_time=infer_time,
        tokens_per_second=tokens / infer_time if infer_time > 0 else 0,
        peak_vram_gb=torch.cuda.max_memory_allocated() / 1e9
    )

This framework measures the key metrics: loading time, inference time, tokens per second, and VRAM usage. Run it with different offload_layers values to map out the performance curve for your specific hardware.

Conclusion

Partial LLM loading isn't a silver bullet, but it's a valuable technique for expanding what's possible on memory-constrained hardware. On my 128GB APU system, I found something unexpected: partial loading with 12 of 16 layers on CPU actually outperformed full GPU loading by 75% (3.36 tok/s vs 1.92 tok/s).

The key takeaways:

Unified memory changes everything. On APUs and Apple Silicon, the conventional wisdom that "GPU is always faster" doesn't hold. Reduced memory bandwidth contention can make partial loading preferable even when you have enough VRAM.
Prefetching is essential. Naive layer loading is too slow. Libraries like oLLM that prefetch the next layer during current layer computation can reduce overhead by 50% or more.
Memory bandwidth matters more than CPU speed. The bottleneck is getting bytes from storage/RAM to the GPU, not processing them once they're there.
Library maturity varies. Expect compatibility issues with newer transformers versions. Pin your dependencies; oLLM requires transformers 4.x.
Quality is preserved. Partial loading changes where weights live, not what they are. Outputs match full GPU inference exactly (assuming matching precision).
Benchmark your specific hardware. My results on a Strix Halo APU won't match discrete GPU performance. The only way to know what works best is to measure it.

For batch processing and experimentation, partial loading lets you access models that would otherwise require more expensive hardware. For unified memory systems specifically, partial loading might be the optimal configuration, not just a fallback.

The era of "if it doesn't fit in VRAM, you can't run it" is ending. With the right techniques, nearly any model becomes accessible, and on the right hardware, you might even get a performance bonus for your trouble.

Running Qwen TTS on AMD Strix Halo: A Complete Guide to Local Text-to-Speech

A.C. Jokela — Sat, 24 Jan 2026 18:00:00 GMT

The rise of high-quality text-to-speech models has opened new possibilities for content creators, accessibility advocates, and developers alike. Qwen3-TTS, developed by Alibaba's Qwen team, represents a significant leap forward in neural TTS technology, offering natural-sounding speech synthesis with multiple speaker voices. In this guide, we'll walk through setting up Qwen3-TTS on AMD's Strix Halo platform (specifically the AI Max+ 395 with its integrated Radeon 8060S graphics) and demonstrate how we use it to generate audio narrations for blog posts right here on TinyComputers.

Why Qwen3-TTS?

The text-to-speech landscape has evolved dramatically over the past few years. While cloud-based services like Amazon Polly, Google Cloud TTS, and ElevenLabs offer impressive quality, they come with ongoing costs, privacy considerations, and internet dependency. Local TTS solutions have historically lagged behind in quality, often producing robotic or unnatural speech.

Qwen3-TTS changes this equation. The model produces remarkably natural speech with proper intonation, pacing, and emphasis. It supports multiple pre-trained speaker voices (including options like Eric, Aiden, Dylan, Serena, and others) each with distinct characteristics suitable for different content types. For technical content like our blog posts, the Eric voice provides clear, professional narration that listeners find easy to follow.

The model we're using, Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice, weighs in at 1.7 billion parameters. While not small, this is manageable on modern hardware and runs efficiently on GPU. The 12Hz designation refers to the audio frame rate used during generation, balancing quality with computational requirements.

The Hardware: AMD AI Max+ 395

AMD's Strix Halo architecture represents their latest push into the high-performance APU market, combining powerful CPU cores with substantial integrated graphics. Our test system features:

CPU: AMD Ryzen AI Max+ 395 with 16 Zen 5 cores (32 threads)
GPU: Integrated Radeon 8060S (RDNA 3.5 architecture)
Memory: 128GB unified DDR5, configured with 96GB VRAM and 32GB system RAM
Compute Units: 40 CUs dedicated to graphics/compute workloads

Our test system is the Bosgame M5 AI Mini Desktop, one of the first mini PCs to ship with AMD's Strix Halo silicon. The GMKtec EVO-X2 is an extremely similar system if you're looking to replicate this setup. The unified memory architecture is particularly relevant for machine learning workloads. Unlike discrete GPUs with their own VRAM, the Radeon 8060S shares system memory with the CPU. This means no PCIe bottleneck for data transfers, and with 96GB allocated as VRAM, even large models fit comfortably.

For our TTS workload, the 8060S provides adequate performance. The 1.7B parameter model fits comfortably in memory, and inference runs entirely on GPU once loaded. We see 100% GPU utilization during speech synthesis, indicating the hardware is being fully leveraged.

Setting Up the Environment

The first challenge with AMD GPUs is getting PyTorch working correctly with ROCm, AMD's open-source GPU compute stack. The Strix Halo uses a newer GPU architecture (gfx1151) that requires ROCm 6.x and some environment variable overrides.

Step 1: Create a Python Virtual Environment

We'll use a dedicated virtual environment to isolate our TTS dependencies:

mkdir -p ~/qwen-tts
cd ~/qwen-tts
python3 -m venv venv
source venv/bin/activate

Step 2: Install PyTorch with ROCm Support

The standard PyTorch installation won't work; we need the ROCm-enabled build. As of this writing, ROCm 6.4 is the latest stable release:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.4

This downloads PyTorch builds compiled specifically for AMD GPUs. The installation is larger than the standard CUDA builds due to the different compute libraries involved.

Step 3: Install Qwen-TTS

With PyTorch in place, install the Qwen TTS package:

pip install qwen-tts soundfile numpy

The soundfile library handles WAV file I/O, while numpy is needed for audio array manipulation.

Step 4: Install xformers for ROCm (Optional but Recommended)

The xformers library provides optimized attention implementations that can improve performance:

pip install xformers --index-url https://download.pytorch.org/whl/rocm6.4

While Qwen-TTS will work without xformers, having it available enables more efficient memory-attention patterns during inference.

Step 5: Configure Environment Variables

The Strix Halo's gfx1151 architecture isn't explicitly recognized by all ROCm components yet. We need to tell the system to treat it as a compatible architecture:

export HSA_OVERRIDE_GFX_VERSION=11.0.0
export GPU_MAX_ALLOC_PERCENT=100
export GPU_MAX_HEAP_SIZE=100
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

Let's break down what these do:

HSA_OVERRIDE_GFX_VERSION=11.0.0: Tells the HSA runtime to report the GPU as gfx1100, which has broader library support
GPU_MAX_ALLOC_PERCENT=100: Allows the GPU to use up to 100% of available memory for allocations
GPU_MAX_HEAP_SIZE=100: Similar memory allocation setting for heap operations
TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1: Enables experimental efficient attention implementations for AMD GPUs

Add these to your .bashrc or create an activation script for convenience.

Step 6: Verify GPU Detection

Before proceeding, confirm PyTorch can see your GPU:

import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device count: {torch.cuda.device_count()}")
print(f"Device name: {torch.cuda.get_device_name(0)}")

You should see output like:

CUDA available: True
Device count: 1
Device name: AMD Radeon 8060S

Note that PyTorch uses "CUDA" terminology even for AMD GPUs when using ROCm; this is for API compatibility.

Basic TTS Usage

With the environment configured, let's test basic speech synthesis:

from qwen_tts import Qwen3TTSModel
import soundfile as sf
import torch

# Load model on GPU with bfloat16 precision
model = Qwen3TTSModel.from_pretrained(
    'Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice',
    attn_implementation='sdpa',
    device_map='cuda:0',
    dtype=torch.bfloat16
)

# Check available speakers
print(f"Available speakers: {model.get_supported_speakers()}")

# Generate speech
text = "Hello, and welcome to TinyComputers. Today we're exploring text-to-speech on AMD hardware."
audios, sample_rate = model.generate_custom_voice(
    text,
    speaker='eric',
    language='english'
)

# Save to file
sf.write('output.wav', audios[0], sample_rate)
print(f"Saved audio at {sample_rate}Hz")

A few important notes:

We use attn_implementation='sdpa' for scaled dot-product attention, which works on ROCm
The device_map='cuda:0' explicitly places the model on the GPU
Using dtype=torch.bfloat16 reduces memory usage while maintaining quality
The language parameter must be the full word 'english', not the abbreviation 'en'

Building a Blog-to-Speech Pipeline

For our use case (generating audio versions of blog posts) we need more than basic TTS. Blog posts contain markdown formatting, code blocks, images, and other elements that shouldn't be read aloud. We built a complete pipeline that handles these challenges.

The Blog Cleaner

Our cleaning process strips out non-spoken content while preserving the narrative flow:

import re

def clean_markdown(content):
    # Remove YAML frontmatter
    if content.startswith('---'):
        end = content.find('---', 3)
        if end != -1:
            content = content[end + 3:]

    # Strip HTML tags (audio, video, images)
    content = re.sub(r'<audio[^>]*>[\s\S]*?</audio>', '', content, flags=re.IGNORECASE)
    content = re.sub(r'<video[^>]*>[\s\S]*?</video>', '', content, flags=re.IGNORECASE)
    content = re.sub(r'<img[^>]*/?>', '', content, flags=re.IGNORECASE)
    content = re.sub(r'<[^>]+>', '', content)

    # Remove markdown images and convert links to just text
    content = re.sub(r'!\[[^\]]*\]\([^)]+\)', '', content)
    content = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', content)

    # Remove code blocks
    content = re.sub(r'```[\s\S]*?```', '', content)
    content = re.sub(r'`[^`]+`', '', content)

    # Convert headers to sentences
    content = re.sub(r'^(#{1,6})\s+(.+)$', r'\2.', content, flags=re.MULTILINE)

    # Remove emphasis markers
    content = re.sub(r'\*\*([^*]+)\*\*', r'\1', content)
    content = re.sub(r'\*([^*]+)\*', r'\1', content)

    return content

Unit Conversion for Speech

Technical content often includes abbreviations that sound awkward when read literally. We convert common units to their spoken forms:

def convert_units_for_speech(text):
    text = re.sub(r'(\d+)\s*GB\b', r'\1 gigabytes', text)
    text = re.sub(r'(\d+)\s*MB\b', r'\1 megabytes', text)
    text = re.sub(r'(\d+)\s*GHz\b', r'\1 gigahertz', text)
    text = re.sub(r'(\d+)\s*MHz\b', r'\1 megahertz', text)
    text = re.sub(r'(\d+)\s*KB\b', r'\1 kilobytes', text)
    return text

Chunking Long Content

TTS models work best with moderate-length inputs. Very long passages can cause quality degradation or memory issues. We split content into chunks at sentence boundaries:

def chunk_text(text, max_chars=500):
    chunks = []
    sentences = re.split(r'(?<=[.!?])\s+', text)
    current = ""

    for sentence in sentences:
        if len(current) + len(sentence) < max_chars:
            current += sentence + " "
        else:
            if current:
                chunks.append(current.strip())
            current = sentence + " "

    if current:
        chunks.append(current.strip())

    return chunks

The Complete Script

Putting it all together, here's our blog_to_speech.py script:

#!/usr/bin/env python3
import argparse
import re
import torch
from pathlib import Path
from qwen_tts import Qwen3TTSModel
import soundfile as sf
import numpy as np

def clean_blog_post(source):
    content = Path(source).read_text()
    # Apply cleaning functions...
    return cleaned_text

def synthesize_speech(text, output_file, speaker="eric"):
    model = Qwen3TTSModel.from_pretrained(
        'Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice',
        attn_implementation='sdpa',
        device_map='cuda:0',
        dtype=torch.bfloat16
    )

    chunks = chunk_text(text)
    all_audio = []

    for i, chunk in enumerate(chunks):
        print(f"Processing chunk {i+1}/{len(chunks)}")
        audios, sample_rate = model.generate_custom_voice(
            chunk,
            speaker=speaker,
            language='english'
        )
        all_audio.append(audios[0])

    combined = np.concatenate(all_audio)
    sf.write(output_file, combined, sample_rate)
    duration = len(combined) / sample_rate
    print(f"Saved {duration:.1f}s audio to: {output_file}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('source', help='Blog post markdown file')
    parser.add_argument('-o', '--output', default='output.wav')
    parser.add_argument('--speaker', default='eric')
    args = parser.parse_args()

    text = clean_blog_post(args.source)
    synthesize_speech(text, args.output, args.speaker)

Choosing the Right Speaker Voice

Qwen3-TTS ships with nine pre-trained speaker voices, each with distinct characteristics:

Speaker	Characteristics	Best For
Eric	Clear, professional male voice with measured pacing	Technical content, tutorials, documentation
Aiden	Younger male voice, slightly more casual	Blog posts, conversational content
Dylan	Deeper male voice with authoritative tone	Formal presentations, announcements
Ryan	Energetic male voice	Marketing content, product demos
Serena	Clear female voice, professional	Corporate content, tutorials
Vivian	Warm female voice	Storytelling, narrative content
Ono Anna	Female voice with distinct character	Creative content
Sohee	Female voice, versatile	General purpose
Uncle Fu	Character voice	Specialized applications

For our technical blog content, we primarily use Eric. His clear enunciation and measured pacing work well for complex technical explanations. The voice handles acronyms, numbers, and technical terminology naturally, making it ideal for content about hardware, programming, and system administration.

You can easily switch voices by changing the speaker parameter:

audios, sample_rate = model.generate_custom_voice(
    text,
    speaker='serena',  # Try different voices
    language='english'
)

Consider matching voice characteristics to content type. A hardware review might work better with Eric's authoritative tone, while a personal essay might benefit from Aiden's more conversational style.

Comparing TTS Options

Before settling on Qwen3-TTS, we evaluated several alternatives. Here's how they compare for our use case:

Cloud Services

Amazon Polly and Google Cloud TTS offer excellent quality with minimal setup. However, costs accumulate quickly for long-form content. At roughly \$4-16 per million characters (depending on voice quality), a 3000-word blog post costs \$0.10-0.40 per generation. For a site with dozens of posts requiring periodic regeneration, this adds up.

ElevenLabs produces arguably the most natural voices available, with impressive emotional range. But their pricing model (based on character quotas) makes it expensive for regular content generation. The quality is exceptional, but overkill for straightforward narration.

Local Alternatives

Coqui TTS (now deprecated) was a popular open-source option but development has stalled. Bark from Suno produces impressive results but runs slowly and lacks fine-grained control. XTTS offers voice cloning but requires more setup and compute resources.

Piper deserves special mention as a lightweight option. It runs quickly even on CPU and produces acceptable quality for many applications. However, the voices sound noticeably synthetic compared to Qwen3-TTS, fine for notifications or short snippets, but fatiguing for 30-minute narrations.

Qwen3-TTS hits a sweet spot: quality approaching cloud services, reasonable compute requirements, and fully local operation. The 1.7B parameter model is large enough for natural prosody but small enough to run on consumer hardware.

Batch Processing for Multiple Posts

When generating audio for multiple blog posts, efficiency matters. Loading the model takes 15-30 seconds, so we keep it loaded while processing multiple files:

#!/usr/bin/env python3
"""Batch TTS processing for multiple blog posts"""

import os
from pathlib import Path
from qwen_tts import Qwen3TTSModel
import torch
import soundfile as sf
import numpy as np

# Load model once
print("Loading model...")
model = Qwen3TTSModel.from_pretrained(
    'Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice',
    attn_implementation='sdpa',
    device_map='cuda:0',
    dtype=torch.bfloat16
)

posts = [
    'post1.md',
    'post2.md',
    'post3.md',
]

for post in posts:
    print(f"\n{'='*50}")
    print(f"Processing: {post}")
    print(f"{'='*50}")

    text = clean_blog_post(post)
    output = f"/tmp/{Path(post).stem}_tts.wav"

    # Process chunks
    chunks = chunk_text(text)
    all_audio = []

    for i, chunk in enumerate(chunks):
        print(f"  Chunk {i+1}/{len(chunks)}")
        audios, sr = model.generate_custom_voice(
            chunk, speaker='eric', language='english'
        )
        all_audio.append(audios[0])

    combined = np.concatenate(all_audio)
    sf.write(output, combined, sr)
    print(f"Saved: {output}")

This approach processes our five-post backlog overnight, with results ready for review in the morning.

Performance Characteristics

On the AI Max+ 395, speech synthesis runs at roughly real-time to 0.5x real-time speed, meaning a 30-minute audio file takes 30-60 minutes to generate. This is slower than high-end discrete GPUs but perfectly acceptable for batch processing.

For reference, here's how different content lengths performed in our testing:

Content	Characters	Chunks	Audio Duration	Generation Time
Short post	5,000	12	~5 min	~15 min
Medium post	15,000	35	~15 min	~45 min
Long post	25,000	55	~27 min	~90 min
Very long	40,000	85	~45 min	~150 min

The relationship between content length and generation time is roughly linear after the initial model warmup.

Some observations from our testing:

First chunk latency: The first chunk takes longer due to GPU kernel compilation and caching
Memory usage: Peak usage around 8-10GB during inference
GPU utilization: Consistent 100% during active synthesis
Quality: Indistinguishable from cloud TTS services for most content

The MIOpen library sometimes logs workspace warnings during execution. These don't affect output quality and can be safely ignored:

MIOpen(HIP): Warning [IsEnoughWorkspace] Solver <GemmFwdRest>, workspace required: 103133184

Integrating Audio into Blog Posts

Once we have the WAV file, we convert to MP3 for web delivery and embed an HTML5 audio player:

ffmpeg -i blog_post.wav -codec:a libmp3lame -qscale:a 2 blog_post.mp3

For reviewing TTS output quality, we recommend using studio monitor headphones that reveal any artifacts or unnatural tones in the generated speech.

The player HTML is straightforward:

<div style="background: #f8f9fa; border: 1px solid #e9ecef;
            border-radius: 8px; padding: 16px 20px; margin: 20px 0;">
  <div class="audio-widget-header">
    <span class="audio-widget-icon">🎧</span>
    <span style="color: #495057; font-weight: 600;">Listen to this article</span>
  </div>
  <audio controls preload="metadata" style="width: 100%;">
    <source src="/audio/blog_post.mp3" type="audio/mpeg">
  </audio>
  <div class="audio-widget-footer">
    27 min · AI-generated narration
  </div>
</div>

Why We're Doing This

Adding audio narration to blog posts serves multiple purposes:

Accessibility: Readers with visual impairments or reading difficulties can consume content aurally
Convenience: Listeners can enjoy posts during commutes, workouts, or other activities
Engagement: Audio content creates a more personal connection with the audience
Reach: Some audiences prefer audio format, expanding our potential readership

Running TTS locally rather than using cloud services gives us:

Cost control: No per-character or per-minute fees
Privacy: Content never leaves our infrastructure
Consistency: Same voice and quality across all posts
Flexibility: Full control over processing pipeline

Troubleshooting Common Issues

"CUDA not available" despite GPU present

Ensure you've installed the ROCm version of PyTorch, not the standard build:

pip uninstall torch torchvision torchaudio
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.4

Model runs on CPU instead of GPU

Check that device_map='cuda:0' is specified when loading the model. Also verify the environment variables are set before starting Python.

"Unsupported language 'en'"

Use the full language name: language='english' not language='en'.

Out of memory errors

Try reducing chunk size or using a smaller batch. The model should fit in 16GB, but very long chunks can spike memory usage.

Slow first chunk

This is normal; ROCm compiles GPU kernels on first use. Subsequent chunks process faster.

Future Improvements

Our current pipeline works well but has room for enhancement. Some improvements we're considering:

Voice cloning: Qwen3-TTS supports custom voice training. With sufficient audio samples, we could create a unique voice for TinyComputers rather than using the stock speakers. This would provide brand consistency and differentiation.

Automatic post detection: Currently we manually select posts for TTS generation. A CI/CD integration could automatically generate audio for new posts when they're published, keeping the audio library current without manual intervention.

Chapter markers: For longer posts, embedding chapter markers in the audio file would allow listeners to skip to specific sections. This requires parsing the markdown headers and mapping them to audio timestamps.

Multiple format export: Beyond MP3, offering Opus or AAC formats could reduce file sizes while maintaining quality, benefiting listeners on metered connections.

Speed adjustment: Some listeners prefer 1.25x or 1.5x playback speed. Pre-generating speed-adjusted versions could provide better quality than real-time speed adjustment in the browser.

Conclusion

Running Qwen3-TTS on AMD's Strix Halo platform demonstrates that high-quality local TTS is now accessible beyond NVIDIA hardware. While setup requires some ROCm-specific configuration, the results are impressive: natural-sounding narration suitable for professional content.

The democratization of AI capabilities continues apace. What once required expensive cloud subscriptions or high-end NVIDIA GPUs now runs on integrated graphics. The AI Max+ 395's Radeon 8060S, primarily designed for gaming and general compute tasks, handles a 1.7-billion parameter language model without breaking a sweat.

We're actively using this pipeline to generate audio versions of posts across TinyComputers, making our technical content more accessible and convenient for our readers. As of this writing, we've processed our retrocomputing series, hardware reviews, and technical tutorials, dozens of hours of content generated entirely on local hardware.

The combination of AMD's capable integrated graphics and Qwen's excellent TTS model proves that you don't need expensive discrete GPUs or cloud subscriptions to achieve broadcast-quality speech synthesis. For content creators, educators, and accessibility advocates, this opens new possibilities for enriching written content with audio without ongoing service costs.

If you're running AMD hardware and want to add audio narration to your own content, this guide should get you started. The initial setup investment pays dividends in ongoing cost savings and the satisfaction of running capable AI models entirely on your own infrastructure. And if you encounter issues along the way, the troubleshooting section above addresses the most common pitfalls we discovered during our own setup process.

The audio player at the top of many TinyComputers posts now represents a small but meaningful step toward making technical content more accessible. Every post you can listen to while commuting, exercising, or doing dishes is content that might otherwise go unread. That's the real value of local TTS: not just cost savings, but expanded reach for the ideas we share.

Rockchip RK3588 NPU Deep Dive: Real-World AI Performance Across Multiple Platforms

A.C. Jokela — Fri, 07 Nov 2025 16:02:55 GMT

Introduction

The Rockchip RK3588 has emerged as one of the most compelling ARM System-on-Chips (SoCs) for edge AI applications in 2024-2025, featuring a dedicated 6 TOPS Neural Processing Unit (NPU) integrated alongside powerful Cortex-A76/A55 CPU cores. This SoC powers a growing ecosystem of single-board computers and system-on-modules from manufacturers worldwide, including Orange Pi, Radxa, FriendlyElec, Banana Pi, and numerous industrial board makers.

But how does the RK3588's NPU perform in real-world scenarios? In this comprehensive deep dive, I'll share detailed benchmarks of the RK3588 NPU testing both Large Language Models (LLMs) and computer vision workloads, with primary testing on the Orange Pi 5 Max and comparative analysis against the closely-related RK3576 found in the Banana Pi CM5-Pro.

The RK3588 Ecosystem: Devices and Availability

The Rockchip RK3588 powers a diverse range of single-board computers (SBCs) and system-on-modules (SoMs) from multiple manufacturers in 2024-2025:

Consumer SBCs:

Orange Pi 5 Max - Full-featured SBC with up to 16GB RAM, M.2 NVMe, WiFi 6
Radxa ROCK 5B/5B+ - Available with up to 32GB RAM, PCIe 3.0, 8K video output
FriendlyElec NanoPC-T6 - Compact form factor with AV1 hardware acceleration
Firefly ROC-RK3588S-PC - Budget-friendly option starting at $219

Industrial and Embedded Modules:

Geniatech DB3588V2 - Industrial-grade development kit with wide temperature range (-40°C to 85°C)
Forlinx OK3588-C - SoM + carrier board design for custom integration
Vantron VT-SBC-3588 - AIoT-focused platform for edge applications
Boardcon Idea3588 - Compute module with up to 16GB RAM and 256GB eMMC
Theobroma Systems TIGER/JAGUAR - High-reliability modules for robotics and industrial automation

Recent Developments:

RK3588S2 (2024-2025) - Updated variant with modernized memory controllers and platform I/O while maintaining the same 6 TOPS NPU performance

The RK3576, found in devices like the Banana Pi CM5-Pro, shares the same 6 TOPS NPU architecture as the RK3588 but features different CPU cores (Cortex-A72/A53 vs. A76/A55), making it an interesting comparison point for NPU-focused workloads.

Hardware Overview

RK3588 SoC Specifications

Built on an 8nm process, the Rockchip RK3588 integrates:

CPU:

4x ARM Cortex-A76 @ 2.4 GHz (high-performance cores)
4x ARM Cortex-A55 @ 1.8 GHz (efficiency cores)

NPU:

6 TOPS total performance
3-core architecture (2 TOPS per core)
Shared memory architecture
Optimized for INT8 operations
Supports INT4/INT8/INT16/BF16/TF32 quantization formats
Device path: /sys/kernel/iommu_groups/0/devices/fdab0000.npu

GPU:

ARM Mali-G610 MP4 (quad-core)
8K@30fps H.265/VP9 decoding
4K@60fps H.264/H.265 encoding

Architecture: ARM64 (aarch64)

Test Platform: Orange Pi 5 Max

For these benchmarks, we used the Orange Pi 5 Max with:

16GB LPDDR5 RAM
1TB M.2 NVMe SSD
WiFi 6 (802.11ax)
Debian-based Linux distribution

Software Stack:

RKNPU Driver: v0.9.8
RKLLM Runtime: v1.2.2 (for LLM inference)
RKNN Runtime: v1.6.0 (for general AI models)
RKNN-Toolkit-Lite2: v2.3.2

Test Setup

I conducted two separate benchmark suites:

Large Language Model (LLM) Testing using RKLLM
Computer Vision Model Testing using RKNN-Toolkit2

Both tests used a two-system approach:

Conversion System: AMD RYZEN AI MAX+ 395 (32 cores, x86_64) running Ubuntu 24.04.3 LTS
Inference System: Orange Pi 5 Max (ARM64) with RK3588 NPU

This reflects the real-world workflow where model conversion happens on powerful workstations, and inference runs on edge devices.

Part 1: Large Language Model Performance

Model: TinyLlama 1.1B Chat

Source: Hugging Face (TinyLlama-1.1B-Chat-v1.0)

Parameters: 1.1 billion

Original Size: ~2.1 GB (505 MB model.safetensors)

Conversion Performance (x86_64)

Converting the Hugging Face model to RKNN format on the AMD RYZEN AI MAX+ 395:

Phase	Time	Details
Load	0.36s	Loading Hugging Face model
Build	22.72s	W8A8 quantization + NPU optimization
Export	56.38s	Export to .rkllm format
Total	79.46s	~1.3 minutes

Output Model:

File: tinyllama_W8A8_rk3588.rkllm
Size: 1142.9 MB (1.14 GB)
Compression: 54% of original size
Quantization: W8A8 (8-bit weights, 8-bit activations)

Note: The RK3588 only supports W8A8 quantization for LLM inference, not W4A16.

NPU Inference Results

Hardware Detection:

I rkllm: rkllm-runtime version: 1.2.2, rknpu driver version: 0.9.8, platform: RK3588
I rkllm: rkllm-toolkit version: 1.2.2, max_context_limit: 2048, npu_core_num: 3
I rkllm: Enabled cpus: [4, 5, 6, 7]
I rkllm: Enabled cpus num: 4

Key Observations:

✅ NPU successfully detected and initialized
✅ All 3 NPU cores utilized
✅ 4 CPU cores (Cortex-A76) enabled for coordination
✅ Model loaded and text generation working
✅ Coherent English text output

Expected Performance (from Rockchip official benchmarks):

TinyLlama 1.1B W8A8 on RK3588: ~10-15 tokens/second
First token latency: ~200-500ms

Is This Fast Enough for Real-Time Conversation?

To put the 10-15 tokens/second performance in perspective, let's compare it to human reading speeds:

Human Reading Rates:

Silent reading: 200-300 words/minute (3.3-5 words/second)
Reading aloud: 150-160 words/minute (2.5-2.7 words/second)
Speed reading: 400-700 words/minute (6.7-11.7 words/second)

Token-to-Word Conversion:

LLM tokens ≈ 0.75 words on average (1.33 tokens per word)
10-15 tokens/sec = ~7.5-11.25 words/second

Performance Analysis:

✅ 2-4x faster than reading aloud (2.5-2.7 words/sec)
✅ 2-3x faster than comfortable silent reading (3.3-5 words/sec)
✅ Comparable to speed reading (6.7-11.7 words/sec)

Verdict: The RK3588 NPU running TinyLlama 1.1B generates text significantly faster than most humans can comfortably read, making it well-suited for real-time conversational AI, chatbots, and interactive applications at the edge.

This is particularly impressive for a $180 device consuming only 5-6W of power. Users won't be waiting for the AI to "catch up" - instead, the limiting factor is human reading speed, not the NPU's generation capability.

Output Quality Verification

To verify the model produces meaningful, coherent responses, I tested it with several prompts:

Test 1: Factual Question

Prompt: "What is the capital of France?"
Response: "The capital of France is Paris."

✅ Result: Correct and concise answer.

Test 2: Simple Math

Prompt: "What is 2 plus 2?"
Response: "2 + 2 = 4"

✅ Result: Correct mathematical calculation.

Test 3: List Generation

Prompt: "List 3 colors: red,"
Response: "Here are three different color options for your text:
1. Red
2. Orange
3. Yellow"

✅ Result: Logical completion with proper formatting.

Observations:

Responses are coherent and grammatically correct
Factual accuracy is maintained after W8A8 quantization
The model understands context and provides relevant answers
Text generation is fluent and natural
No obvious degradation from quantization

Note: The interactive demo tends to continue generating after the initial response, sometimes repeating patterns. This appears to be a demo interface issue rather than a model quality problem - the initial responses to each prompt are consistently accurate and useful.

LLM Findings

Strengths:

Fast model conversion (~1.3 minutes for 1.1B model)
Successful NPU detection and initialization
Good compression ratio (54% size reduction)
Verified high-quality output: Factually correct, grammatically sound responses
Text generation faster than human reading speed (7.5-11.25 words/sec)
All 3 NPU cores actively utilized
No noticeable quality degradation from W8A8 quantization

Limitations:

RK3588 only supports W8A8 quantization (no W4A16 for better compression)
1.14 GB model size may be limiting for memory-constrained deployments
Max context length: 2048 tokens

RK3588 vs RK3576: NPU Performance Comparison

The RK3576, found in the Banana Pi CM5-Pro, shares the same 6 TOPS NPU architecture as the RK3588 but differs in CPU configuration (Cortex-A72/A53 vs. A76/A55). This provides an interesting comparison for understanding NPU-specific performance versus overall platform capabilities.

LLM Performance (Official Rockchip Benchmarks):

Model	RK3588 (W8A8)	RK3576 (W4A16)	Notes
Qwen2 0.5B	~42.58 tokens/sec	34.24 tokens/sec	RK3588 ~1.24x faster
MiniCPM4 0.5B	N/A	35.8 tokens/sec	-
TinyLlama 1.1B	~10-15 tokens/sec	21.32 tokens/sec	RK3576 faster (different quant)
InternLM2 1.8B	N/A	13.65 tokens/sec	-

Key Observations:

RK3588 supports W8A8 quantization only for LLMs
RK3576 supports W4A16 quantization (4-bit weights, 16-bit activations)
W4A16 models are smaller (645MB vs 1.14GB for TinyLlama) but may run slower on some models
The NPU architecture is fundamentally the same (6 TOPS, 3 cores), but software stack differences affect performance
For 0.5B models, RK3588 shows ~20% better performance
Larger models benefit from W4A16's memory efficiency on RK3576

Computer Vision Performance:

Both RK3588 and RK3576 share the same NPU architecture for computer vision workloads:

MobileNet V1 on RK3576 (Banana Pi CM5-Pro): ~161.8ms per image (~6.2 FPS)
ResNet18 on RK3588 (Orange Pi 5 Max): 4.09ms per image (244 FPS)

The dramatic performance difference here is primarily due to model complexity (ResNet18 is better optimized for NPU execution than older MobileNet V1) rather than NPU hardware differences.

Practical Implications:

For NPU-focused workloads, both the RK3588 and RK3576 deliver similar AI acceleration capabilities. The choice between platforms should be based on:

CPU performance needs: RK3588's A76 cores are significantly faster
Quantization requirements: RK3576 offers W4A16 for LLMs, RK3588 only W8A8
Model size constraints: W4A16 (RK3576) produces smaller models
Cost considerations: RK3576 platforms (like CM5-Pro at $103) vs RK3588 platforms ($150-180)

Part 2: Computer Vision Model Performance

Model: ResNet18 (PyTorch Converted)

Source: PyTorch pretrained ResNet18

Parameters: 11.7 million

Original Size: 44.6 MB (ONNX format)

Can PyTorch Run on RK3588 NPU?

Short Answer: Yes, but through conversion.

Workflow: PyTorch → ONNX → RKNN → NPU Runtime

PyTorch/TensorFlow models cannot execute directly on the NPU. They must be converted through an AOT (Ahead-of-Time) compilation process. However, this conversion is fast and straightforward.

Conversion Performance (x86_64)

Converting PyTorch ResNet18 to RKNN format:

Phase	Time	Size	Details
PyTorch → ONNX	0.25s	44.6 MB	Fixed batch size, opset 11
ONNX → RKNN	1.11s	-	INT8 quantization, operator fusion
Export	0.00s	11.4 MB	Final .rknn file
Total	1.37s	11.4 MB	25.7% of ONNX size

Model Optimizations:

INT8 quantization (weights and activations)
Automatic operator fusion
Layout optimization for NPU
Target: 3 NPU cores on RK3588

Memory Usage:

Internal memory: 1.1 MB
Weight memory: 11.5 MB
Total model size: 11.4 MB

NPU Inference Performance

Running ResNet18 inference on Orange Pi 5 Max (10 iterations after 2 warmup runs):

Results:

Average Inference Time: 4.09 ms
Min Inference Time: 4.02 ms
Max Inference Time: 4.43 ms
Standard Deviation: ±0.11 ms
Throughput: 244.36 FPS

Initialization Overhead:

NPU initialization: 0.350s (one-time)
Model load: 0.008s (one-time)

Input/Output:

Input: 224×224×3 images (INT8)
Output: 1000 classes (Float32)

Performance Comparison

Platform	Inference Time	Throughput	Notes
RK3588 NPU	4.09 ms	244 FPS	3 NPU cores, INT8
ARM A76 CPU (est.)	~50 ms	~20 FPS	Single core
Desktop RTX 3080	~2-3 ms	~400 FPS	Reference
NPU Speedup	12x faster than CPU	-	Same hardware

Computer Vision Findings

Strengths:

Extremely fast conversion (<2 seconds)
Excellent inference performance (4.09ms, 244 FPS)
Very consistent latency (±0.11ms)
Efficient quantization (74% size reduction)
12x speedup vs CPU cores on same SoC
Simple Python API for inference

Trade-offs:

INT8 quantization may reduce accuracy slightly
AOT conversion required (no dynamic model execution)
Fixed input shapes required

Technical Deep Dive

NPU Architecture

The RK3588 NPU is based on a 3-core design with 6 TOPS total performance:

Each core contributes 2 TOPS
Shared memory architecture
Optimized for INT8 operations
Direct DRAM access for large models

Memory Layout

For ResNet18, the NPU memory allocation:

Feature Tensor Memory:
- Input (224×224×3):     147 KB
- Layer activations:     776 KB (peak)
- Output (1000 classes): 4 KB

Constant Memory (Weights):
- Conv layers:    11.5 MB
- FC layers:      2.0 MB
- Total:          11.5 MB

Operator Support

The RKNN runtime successfully handled all ResNet18 operators:

Convolution layers: ✅ Fused with ReLU activation
Batch normalization: ✅ Folded into convolution
MaxPooling: ✅ Native support
Global average pooling: ✅ Converted to convolution
Fully connected: ✅ Converted to 1×1 convolution

All 26 operators executed on NPU (no CPU fallback needed).

Power Efficiency

While I didn't measure power consumption directly, the RK3588 NPU is designed for edge deployment:

Estimated Power Draw:

Idle: ~2-3W (entire SoC)
NPU active: +2-3W
Total under AI load: ~5-6W

Performance per Watt:

ResNet18 @ 244 FPS / ~5W = ~49 FPS per Watt
Compare to desktop GPU: RTX 3080 @ 400 FPS / ~320W = ~1.25 FPS per Watt

The RK3588 NPU delivers approximately 39x better performance per watt than a high-end desktop GPU for INT8 inference workloads.

Real-World Applications

Based on these benchmarks, the RK3588 NPU is well-suited for:

✅ Excellent Performance:

Real-time object detection: 244 FPS for ResNet18-class models
Image classification: Sub-5ms latency
Face recognition: Multiple faces per frame at 30+ FPS
Pose estimation: Real-time tracking
Edge AI cameras: Low power, high throughput

✅ Good Performance:

Small LLMs: 1B-class models at 10-15 tokens/second
Chatbots: Acceptable latency for edge applications
Text classification: Fast inference for short sequences

⚠️ Limited Performance:

Large LLMs: 7B+ models may not fit in memory or run slowly
High-resolution video: 4K processing may require frame decimation
Transformer models: Attention mechanism less optimized than CNNs

Developer Experience

Pros:

Clear documentation and examples
Python API is straightforward
Automatic NPU detection
Fast conversion times
Good error messages

Cons:

Requires separate x86_64 system for conversion
Some dependency conflicts (PyTorch versions)
Limited dynamic shape support
Debugging NPU issues can be challenging

Getting Started

Here's a minimal example for running inference:

from rknnlite.api import RKNNLite
import numpy as np

# Initialize
rknn = RKNNLite()

# Load model
rknn.load_rknn('model.rknn')
rknn.init_runtime()

# Run inference
input_data = np.random.randint(0, 256, (1, 3, 224, 224), dtype=np.uint8)
outputs = rknn.inference(inputs=[input_data])

# Cleanup
rknn.release()

That's it! The NPU is automatically detected and utilized.

Cost Analysis

Orange Pi 5 Max: ~$150-180 (16GB RAM variant)

Performance per Dollar:

244 FPS / $180 = 1.36 FPS per dollar (ResNet18)
10-15 tokens/s / $180 = 0.055-0.083 tokens/s per dollar (TinyLlama 1.1B)

Compare to:

Raspberry Pi 5 (8GB): $80, ~5 FPS CPU → 0.063 FPS per dollar
NVIDIA Jetson Orin Nano: $499, ~400 FPS → 0.80 FPS per dollar
Desktop RTX 3080: $699+, ~400 FPS → 0.57 FPS per dollar

The RK3588 NPU offers excellent value for edge AI applications, especially for INT8 workloads.

Comparison to Other Edge AI Platforms

Platform	NPU/GPU	TOPS	Price	ResNet18 FPS	Notes
Orange Pi 5 Max (RK3588)	3-core NPU	6	$180	244	Best value
Raspberry Pi 5	CPU only	-	$80	~5	No accelerator
Google Coral Dev Board	Edge TPU	4	$150	~400	INT8 only
NVIDIA Jetson Orin Nano	GPU (1024 CUDA)	40	$499	~400	More flexible
Intel NUC with Neural Compute Stick 2	VPU	4	$300+	~150	Requires USB

The RK3588 stands out for offering strong NPU performance at a very competitive price point.

Limitations and Gotchas

1. Conversion System Required

You cannot convert models directly on the Orange Pi. You need an x86_64 Linux system with RKNN-Toolkit2 for model conversion.

2. Quantization Constraints

LLMs: Only W8A8 supported (no W4A16)
Computer vision: INT8 quantization required for best performance
Floating-point models will run slower

3. Memory Limitations

Large models (>2GB) may not fit
Context length limited to 2048 tokens for LLMs
Batch sizes are constrained by NPU memory

4. Framework Support

PyTorch/TensorFlow: Supported via conversion
Direct framework execution: Not supported
Some operators may fall back to CPU

5. Software Maturity

RKNN-Toolkit2 is actively developed but not as mature as CUDA
Some edge cases and exotic operators may not be supported
Version compatibility between toolkit and runtime must match

Best Practices

Based on my testing, here are recommendations for optimal RK3588 NPU usage:

1. Model Selection

Choose models designed for mobile/edge: MobileNet, EfficientNet, SqueezeNet
Start small: Test with smaller models before scaling up
Consider quantization-aware training: Better accuracy with INT8

2. Optimization

Use fixed input shapes: Dynamic shapes have overhead
Batch carefully: Batch size 1 often optimal for latency
Leverage operator fusion: Design models with fusible ops (Conv+BN+ReLU)

3. Deployment

Pre-load models: Model loading takes ~350ms
Use separate threads: Don't block main application during inference
Monitor memory: Large models can cause OOM errors

4. Development Workflow

1. Train on workstation (GPU)
2. Export to ONNX with fixed shapes
3. Convert to RKNN on x86_64 system
4. Test on Orange Pi 5 Max
5. Iterate based on accuracy/performance

Conclusion

The RK3588 NPU on the Orange Pi 5 Max delivers impressive performance for edge AI applications. With 244 FPS for ResNet18 (4.09ms latency) and 10-15 tokens/second for 1.1B LLMs, it's well-positioned for real-time computer vision and small language model inference.

Key Takeaways:

✅ Excellent computer vision performance: 244 FPS for ResNet18, <5ms latency

✅ Good LLM support: 1B-class models run at usable speeds

✅ Outstanding value: $180 for 6 TOPS of NPU performance

✅ Easy to use: Simple Python API, automatic NPU detection

✅ Power efficient: ~5-6W under AI load, 39x better than desktop GPU

✅ PyTorch compatible: Via conversion workflow

⚠️ Conversion required: Cannot run PyTorch/TensorFlow directly

⚠️ Quantization needed: INT8 for best performance

⚠️ Memory constrained: Large models (>2GB) challenging

The RK3588 NPU is an excellent choice for edge AI applications where power efficiency and cost matter. It's not going to replace high-end GPUs for training or large-scale inference, but for deploying computer vision models and small LLMs at the edge, it's one of the best options available today.

Recommended for:

Edge AI cameras and surveillance
Robotics and autonomous systems
IoT devices with AI requirements
Embedded AI applications
Prototyping and development

Not recommended for:

Large language model training
7B+ LLM inference
High-precision (FP32) inference
Dynamic model execution
Cloud-scale deployments

Getting YOLOv8 Training Working on AMD Ryzen™ AI Max+ 395

A.C. Jokela — Wed, 22 Oct 2025 14:54:43 GMT

Introduction

Machine learning on AMD GPUs has always been... interesting. With NVIDIA's CUDA dominating the landscape, AMD's ROCm platform remains the underdog: powerful, but often requiring patience and persistence to get working properly. This is the story of how I got YOLOv8 object detection training working on an AMD Radeon 8060S integrated GPU (gfx1151) in the AMD RYZEN AI MAX+ 395 after encountering batch normalization failures, version mismatches, and a critical bug in MIOpen.

The goal was simple: train a bullet hole detection model for a ballistics application using YOLOv8. The journey? Anything but simple.

The Hardware

System Specifications:

CPU: AMD RYZEN AI MAX+ 395
GPU: AMD Radeon 8060S (integrated, RDNA 3.5 architecture, gfx1151)
VRAM: 96GB shared system memory
ROCm Version: 7.0.2
ROCk module: 6.14.14
PyTorch: 2.8.0+rocm7.0.0.git64359f59
MIOpen: Initially 3.0.5.1 (version code 3005001), later custom build
OS: Linux (conda environment: pt2.8-rocm7)

The AMD Radeon 8060S is an integrated GPU in the AMD RYZEN AI MAX+ 395 based on AMD's RDNA 3.5 architecture (gfx1151). What makes this system particularly interesting for machine learning is the massive 96GB of shared system memory available to the GPU, far more VRAM than typical consumer discrete GPUs. While machine learning support on RDNA 3.5 is still maturing compared to older RDNA 2 architectures, the memory capacity makes it compelling for AI workloads.

But, for about $1,699, you can get up to 96GB of VRAM in a whisper-quiet form factor. This setup beats the pants off of my old GPU rig.

Why YOLOv8 and Ultralytics?

Before diving into the technical challenges, it's worth explaining why we chose YOLOv8 from Ultralytics for this project.

YOLOv8 (You Only Look Once, version 8) is the latest iteration of one of the most popular object detection architectures. Developed and maintained by Ultralytics, it offers several advantages:

Why Ultralytics YOLOv8?

State-of-the-art Accuracy: YOLOv8 achieves excellent detection accuracy while maintaining real-time inference speeds, critical for practical applications.
Ease of Use: Ultralytics provides a clean, well-documented Python API that makes training custom models remarkably straightforward:

from ultralytics import YOLO
model = YOLO("yolov8n.pt")
results = model.train(data="dataset.yaml", epochs=100)

Active Development: Ultralytics is actively maintained with frequent updates, bug fixes, and community support. This proved invaluable during debugging.
Model Variants: YOLOv8 comes in multiple sizes (nano, small, medium, large, extra-large), allowing us to balance accuracy vs. speed for our specific use case.
Built-in Data Augmentation: The framework includes extensive data augmentation capabilities out of the box, essential for training robust detection models with limited training data.
PyTorch Native: Being built on PyTorch meant it should theoretically work with ROCm (AMD's CUDA equivalent)... in theory.

For our bullet hole detection application, YOLOv8's ability to accurately detect small objects (bullet holes in paper targets) while training efficiently made it the obvious choice. Little did I know that "training efficiently" would require a week-long debugging odyssey.

The Initial Setup (ROCm 7.0.0)

I started with ROCm 7.0.0, following AMD's official installation guide. Everything installed cleanly:

$ python -c "import torch; print(torch.cuda.is_available())"
True

$ python -c "import torch; print(torch.cuda.get_device_name(0))"
AMD Radeon Graphics

Perfect! PyTorch recognized the GPU. Time to train some models, right?

The First Failure: Batch Normalization

I loaded a simple YOLOv8 nano model and kicked off training:

from ultralytics import YOLO

model = YOLO("yolov8n.pt")
results = model.train(
    data="data/bullet_hole_dataset_combined/data.yaml",
    epochs=100,
    imgsz=416,
    batch=16,
    device="cuda:0"
)

Within seconds, the training crashed:

RuntimeError: miopenStatusUnknownError

The error was cryptic, but digging deeper revealed the real issue: MIOpen was failing to compile batch normalization kernels with inline assembly errors:

<inline asm>:14:20: error: not a valid operand.
v_add_f32 v4 v4 v4 row_bcast:15 row_mask:0xa
                   ^

Batch normalization. The most common operation in modern deep learning, and it was failing spectacularly on gfx1151. The inline assembly instructions (row_bcast and row_mask) appeared incompatible with the RDNA 3.5 architecture.

What is Batch Normalization?

Batch normalization (BatchNorm) is a technique that normalizes layer inputs across a mini-batch, helping neural networks train faster and more stably. It's used in virtually every modern CNN architecture, including YOLO.

The error message pointed to MIOpen, AMD's equivalent of NVIDIA's cuDNN, a library of optimized deep learning primitives.

Attempt 1: Upgrade to ROCm 7.0.2

My first instinct was to upgrade ROCm. Version 7.0.0 was relatively new, and perhaps 7.0.2 had fixed the batch normalization issues.

# Upgraded PyTorch to ROCm 7.0.2
pip install --upgrade torch --index-url https://download.pytorch.org/whl/rocm7.0

Result? Same error. Batch normalization still failed.

RuntimeError: miopenStatusUnknownError

With the same inline assembly compilation errors about invalid row_bcast and row_mask operands. At this point, I realized this wasn't a simple version mismatch; there was something fundamentally broken with MIOpen's batch normalization implementation for the gfx1151 architecture.

The Revelation: It's MIOpen, Not ROCm

After hours of testing different PyTorch versions, driver configurations, and kernel parameters, I turned to the ROCm community for help.

I posted my issue on Reddit's r/ROCm subreddit, describing the inline assembly compilation failures and miopenStatusUnknownError on gfx1151. Within a few hours, a knowledgeable Redditor responded with a crucial piece of information:

"There's a known issue with MIOpen 3.0.x and gfx1151 batch normalization. The inline assembly instructions use operands that aren't compatible with RDNA 3. A fix was recently merged into the develop branch. Try using a nightly build of MIOpen or build from source."

This was the breakthrough I needed. The issue wasn't with ROCm itself or PyTorch; it was specifically MIOpen version 3.0.5.1 that shipped with ROCm 7.0.x. The maintainers had already fixed the gfx1151 batch normalization bug in a recent pull request, but it hadn't made it into a stable release yet.

The Reddit user suggested two options:

Use a nightly Docker container with the latest MIOpen build
Build MIOpen 3.5.1 from source using the develop branch

Testing the Theory: Docker Nightly Builds

Before committing to building from source, I wanted to verify that a newer MIOpen would actually fix the problem. AMD provides nightly Docker images with bleeding-edge ROCm builds:

docker pull rocm/pytorch-nightly:latest

docker run --rm \
    --device=/dev/kfd \
    --device=/dev/dri \
    --group-add video \
    -v ~/ballistics_training:/workspace \
    -w /workspace \
    rocm/pytorch-nightly:latest \
    bash -c 'pip install ultralytics && python3 test_yolo.py'

The nightly container included MIOpen 3.5.1 from the develop branch.

# test_yolo.py
from ultralytics import YOLO
import torch

print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device: {torch.cuda.get_device_name(0)}")

model = YOLO("yolov8n.pt")
results = model.train(
    data="data_docker.yaml",
    epochs=1,
    imgsz=416,
    batch=2,
    device="cuda:0"
)

Result:

✅ SUCCESS! Nightly build FIXES gfx1151 batch normalization!

It worked! The miopenStatusUnknownError was gone, no more inline assembly compilation failures. Training completed successfully with MIOpen 3.5.1 from the develop branch. The newer version had updated the batch normalization kernels to use instructions compatible with RDNA 3.5's gfx1151 architecture.

This confirmed the Reddit user's tip: the fix was indeed in the newer MIOpen code that hadn't been released in a stable version yet.

The Solution: Building MIOpen from Source

Docker was great for testing, but I needed a permanent solution for my native conda environment. That meant building MIOpen 3.5.1 from source.

Step 1: Clone the Repository

cd ~/ballistics_training
git clone https://github.com/ROCm/MIOpen.git rocm-libraries/projects/miopen
cd rocm-libraries/projects/miopen
git checkout develop  # Latest development branch with gfx1151 fixes

Step 2: Build MIOpen

mkdir build && cd build

cmake \
    -DCMAKE_PREFIX_PATH="/opt/rocm" \
    -DCMAKE_INSTALL_PREFIX="$HOME/ballistics_training/rocm-libraries/projects/miopen/build" \
    -DMIOPEN_BACKEND=HIP \
    -DCMAKE_BUILD_TYPE=Release \
    ..

make -j$(nproc)

[ 98%] Building CXX object src/CMakeFiles/MIOpen.dir/softmax_api.cpp.o
[ 99%] Linking CXX shared library libMIOpen.so
[100%] Built target MIOpen

Success! MIOpen 3.5.1 was built from source.

Step 3: Install Custom MIOpen to Conda Environment

Now came the tricky part: replacing the system MIOpen (version 3.0.5.1) with my custom-built version 3.5.1.

CONDA_LIB=~/anaconda3/envs/pt2.8-rocm7/lib

# Backup the original MIOpen
cp $CONDA_LIB/libMIOpen.so.1.0 $CONDA_LIB/libMIOpen.so.1.0.backup_system

# Install custom MIOpen
cp ~/ballistics_training/rocm-libraries/projects/miopen/build/lib/libMIOpen.so.1.0 $CONDA_LIB/

# Update symlinks
cd $CONDA_LIB
ln -sf libMIOpen.so.1.0 libMIOpen.so.1
ln -sf libMIOpen.so.1 libMIOpen.so

Step 4: Verify the Installation

conda activate pt2.8-rocm7
python -c "import torch; print(f'MIOpen version: {torch.backends.cudnn.version()}')"

Output:

MIOpen version: 3005001

Wait, 3005001? That's version 3.5.1! (MIOpen uses an integer versioning scheme: major1000000 + minor1000 + patch)

The custom MIOpen was successfully loaded.

The Final Test: YOLOv8 Training

Time for the moment of truth. Could I finally train YOLOv8 on my AMD GPU?

from ultralytics import YOLO
import torch

print("=" * 60)
print("Testing YOLOv8 Training with Custom MIOpen 3.5.1")
print("=" * 60)
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"MIOpen version: {torch.backends.cudnn.version()}")
print()

model = YOLO("yolov8n.pt")
print("Starting training...")

results = model.train(
    data="data/bullet_hole_dataset_combined/data.yaml",
    epochs=100,
    imgsz=416,
    batch=16,
    device="cuda:0",
    name="bullet_hole_detector"
)

Output:

============================================================
Testing YOLOv8 Training with Custom MIOpen 3.5.1
============================================================
PyTorch: 2.8.0+rocm7.0.0.git64359f59
CUDA available: True
MIOpen version: 3005001

Starting training...

Ultralytics 8.3.217 🚀 Python-3.12.11 torch-2.8.0+rocm7.0.0 CUDA:0 (AMD Radeon Graphics, 98304MiB)

Model summary: 129 layers, 3,011,043 parameters, 3,011,027 gradients, 8.2 GFLOPs

Transferred 319/355 items from pretrained weights
AMP: running Automatic Mixed Precision (AMP) checks...
AMP: checks passed ✅

Starting training for 1 epochs...

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
        1/1     0.172G      3.022      3.775      1.215         29        416
        1/1     0.174G      2.961      4.034      1.147         46        416
        1/1     0.203G      3.133       4.08      1.251         36        416
        1/1     0.205G       3.14      4.266       1.25         60        416
        1/1     0.205G      3.028      4.194      1.237         18        416
        1/1     0.205G      2.995      4.114      1.235         28        416
        1/1     0.205G      3.029      4.118      1.226         41        416
        1/1     0.205G      2.961      4.031      1.209         26        416
        1/1     0.205G      2.888      3.998      1.193         22        416
        1/1     0.205G      2.861      3.823      1.185         49        416
        1/1     0.205G      2.812      3.657      1.169         46        416
        1/1     0.205G      2.821      3.459      1.149         78        416
        1/1     0.205G      2.776      3.253      1.134         26        416
        1/1     0.217G      2.784      3.207      1.131        122        416
        1/1     0.217G      2.772      3.074      1.121         40        416
        1/1     0.217G      2.774       2.98      1.114         13        416
        1/1     0.217G      2.763      2.914      1.118         37        416
        1/1     0.217G       2.75      2.876      1.113         81        416
        1/1     0.217G      2.731      2.799      1.104         31        416
        1/1     0.217G      2.736      2.732      1.101         30        416: 100% 14.8it/s

                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95)
                   all         60        733      0.653      0.473       0.53      0.191

1 epochs completed in 0.002 hours.

==============================================================
✅ SUCCESS! Training completed without errors!
==============================================================

Speed: 0.0ms preprocess, 1.9ms inference, 0.0ms loss, 0.5ms postprocess per image
Results saved to runs/detect/bullet_hole_detector/

It worked! Batch normalization executed flawlessly. The training progressed smoothly from epoch to epoch, with GPU utilization staying high, memory management remaining stable, and losses converging as expected. The model achieved 53.0% mAP50 and trained without a single error.

After a week of debugging, version wrangling, and source code compilation, I finally had GPU-accelerated YOLOv8 training working on my AMD RDNA 3.5 GPU. The custom MIOpen 3.5.1 build resolved the inline assembly compatibility issues, and training now runs as smoothly on gfx1151 as it would on any other supported GPU.

Performance Notes

With the custom MIOpen build, training performance was excellent:

Training Speed: 70.5 images/second (batch size 16, 416×416 images)
Training Time: 32.6 seconds for 10 epochs (2,300 total images)
Throughput: 9.7-9.9 iterations/second
GPU Utilization: ~95% during training with no throttling
Memory Usage: ~1.2 GB VRAM for YOLOv8n with batch size 16

The GPU utilization stayed consistently high with no performance degradation across epochs. Each epoch averaged approximately 3.3 seconds with solid consistency. For comparison, CPU-only training on the same dataset would be roughly 15-20x slower. The GPU acceleration was well worth the effort.

Lessons Learned

This debugging journey taught me several valuable lessons:

1. The ROCm Community is Invaluable

The Reddit r/ROCm community proved to be the key to solving this issue. When official documentation fails, community knowledge fills the gap. Don't hesitate to ask for help; chances are someone has encountered your exact issue before.

2. MIOpen ≠ ROCm

I initially assumed upgrading ROCm would fix the problem. In reality, MIOpen (the deep learning library) had a separate bug that was independent of the ROCm platform version. Understanding the component architecture of ROCm saved hours of debugging time.

3. RDNA 3.5 (gfx1151) Support is Still Maturing

AMD's latest integrated GPU architecture is powerful, but ML support lags behind older architectures like RDNA 2 (gfx1030) and Vega. If you're doing serious ML work on AMD, consider that newer hardware may require more troubleshooting.

4. Nightly Builds Can Be Production-Ready

There's often hesitation to use nightly/development builds in production. However, in this case, the develop branch of MIOpen was actually more stable than the official release for my specific GPU. Sometimes bleeding-edge code is exactly what you need.

5. Docker is Great for Testing

The ROCm nightly Docker containers were instrumental in proving my hypothesis. Being able to test a newer MIOpen version without committing to a full rebuild saved significant time.

6. Source Builds Give You Control

Building from source is time-consuming and requires understanding the build system, but it gives you complete control over your environment. When binary distributions fail, source builds are your safety net.

Tips for AMD GPU Machine Learning

If you're attempting to do machine learning on AMD GPUs, here are some recommendations:

Environment Setup

Use conda/virtualenv: Isolate your Python environment to avoid system package conflicts
Pin your versions: Lock PyTorch, ROCm, and MIOpen versions once you have a working setup
Keep backups: Always backup working library files before swapping them out

Debugging Strategy

Verify GPU detection first: Ensure torch.cuda.is_available() returns True
Test simple operations: Try basic tensor operations before complex models
Check MIOpen version: torch.backends.cudnn.version() can reveal version mismatches
Monitor logs: ROCm logs (MIOPEN_ENABLE_LOGGING=1) provide valuable debugging info
Try Docker first: Test potential fixes in Docker before modifying your system

Hardware Considerations

RDNA 2 (gfx1030) is more mature than RDNA 3.5 (gfx1151) for ML workloads
Server GPUs (MI series) have better ROCm support than consumer cards
Integrated GPUs with large shared memory (like the Radeon 8060S with 96GB) offer unique advantages for ML
Check compatibility: Always verify your specific GPU (gfx code) is supported before purchasing

Conclusion

Getting YOLOv8 training working on an AMD RDNA 3.5 GPU wasn't easy, but it was achievable. The combination of:

Community support from r/ROCm pointing me to the right solution
Docker testing to verify the fix
Building MIOpen 3.5.1 from source
Carefully replacing system libraries

...resulted in a fully functional GPU-accelerated machine learning training environment.

AMD's ROCm platform still has rough edges compared to NVIDIA's CUDA ecosystem, but it's improving rapidly. With some patience, persistence, and willingness to dig into source code, AMD GPUs can absolutely be viable for machine learning workloads.

The bullet hole detection model trained successfully, achieved excellent accuracy, and now runs in production. Sometimes the journey is as valuable as the destination; I learned more about ROCm internals, library dependencies, and GPU computing in this week than I would have in months of smooth sailing.

If you're facing similar issues with AMD GPUs and ROCm, I hope this guide helps. And remember: when in doubt, check r/ROCm. The community might just have the answer you're looking for.

System Details (for reference):

CPU: AMD RYZEN AI MAX+ 395
GPU: AMD Radeon 8060S (integrated, gfx1151)
VRAM: 96GB shared system memory
ROCm: 7.0.2
ROCk module: 6.14.14
PyTorch: 2.8.0+rocm7.0.0.git64359f59
MIOpen: 3.5.1 (custom build from develop branch)
Conda Environment: pt2.8-rocm7
YOLOv8: Ultralytics 8.3.217

Key Files:

MIOpen source: https://github.com/ROCm/MIOpen
Ultralytics YOLOv8: https://github.com/ultralytics/ultralytics
ROCm installation: https://rocm.docs.amd.com/

Special thanks to the r/ROCm community for pointing me toward the MIOpen develop branch fix!

Getting PyTorch Working with AMD Radeon Pro W7900 (MAX+ 395): A Comprehensive Guide

A.C. Jokela — Sat, 11 Oct 2025 23:08:14 GMT

Your browser does not support the audio element.

Getting PyTorch Working with AMD Radeon Pro W7900 (MAX+ 395): A Comprehensive Guide

Introduction

The AMD Radeon Pro W7900 represents a significant leap forward in professional GPU computing. With 96GB of unified memory and 20 compute units, this workstation-class GPU brings serious computational power to tasks like machine learning, scientific computing, and data analysis. However, getting deep learning frameworks like PyTorch to work with AMD GPUs has historically been more challenging than with NVIDIA's CUDA ecosystem.

Here's a complete walkthrough of setting up PyTorch with ROCm support on the AMD MAX+ 395, including installation, verification, and real-world testing. By the end, you'll have a fully functional PyTorch environment capable of leveraging your AMD GPU's computational power.

Understanding ROCm and PyTorch

What is ROCm?

ROCm (Radeon Open Compute) is AMD's open-source software platform for GPU computing. It serves as AMD's answer to NVIDIA's CUDA, providing:

Low-level GPU programming interfaces
Optimized libraries for linear algebra, FFT, and other operations
Deep learning framework support
Compatibility with CUDA-based code through HIP (Heterogeneous-compute Interface for Portability)

PyTorch and ROCm Integration

PyTorch has officially supported ROCm since version 1.8, and support has matured significantly over subsequent releases. The ROCm version of PyTorch uses the same API as the CUDA version, making it straightforward to port existing PyTorch code to AMD GPUs. In fact, most PyTorch code written for CUDA will work without modification on ROCm, as the framework abstracts away the underlying GPU platform.

System Specifications

Testing was performed on a system with the following specifications:

GPU: AMD Radeon Pro W7900 (MAX+ 395)
GPU Memory: 96 GB
Compute Units: 20
CUDA Capability: 11.5 (ROCm compatibility level)
Operating System: Linux
Python: 3.12.11
PyTorch Version: 2.8.0+rocm7.0.0
ROCm Version: 7.0.0

Installation and Setup

This section provides detailed, step-by-step instructions for bootstrapping a complete ROCm 7.0 + PyTorch 2.8 environment on Ubuntu 24.04.3 LTS. These instructions are based on successful installations on the AMD Ryzen AI Max+395 platform.

Prerequisites

Ubuntu 24.04.3 LTS (Server or Desktop)
Administrator/sudo access
Internet connection for downloading packages

Step 1: Update Linux Kernel

ROCm 7.0 works best with Linux kernel 6.14 or later. Update your kernel:

sudo apt-get install linux-generic-hwe-24.04

Verify the installation:

cat /proc/version

You should see output similar to:

Linux version 6.14.0-33-generic (buildd@lcy02-amd64-026)...

Reboot to load the new kernel:

sudo reboot

Step 2: Install AMDGPU Driver

First, set up the AMD repository:

# Create keyring directory if it doesn't exist
sudo mkdir --parents --mode=0755 /etc/apt/keyrings

# Download and install AMD GPG key
wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
  gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null

# Add AMDGPU repository
sudo tee /etc/apt/sources.list.d/amdgpu.list << EOF
deb [arch=amd64,i386 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/amdgpu/latest/ubuntu noble main
EOF

Install the AMDGPU DKMS driver:

sudo apt update
sudo apt install amdgpu-dkms
sudo reboot

Verify the driver installation:

sudo dkms status

You should see output like:

amdgpu/6.14.14-2212064.24.04, 6.14.0-33-generic, x86_64: installed

Step 3: Install ROCm 7.0

Install prerequisites:

sudo apt install python3-setuptools python3-wheel
sudo apt update

Download and install the AMD GPU installer:

wget https://repo.radeon.com/amdgpu-install/7.0/ubuntu/noble/amdgpu-install_7.0.70000-1_all.deb
sudo apt install ./amdgpu-install_7.0.70000-1_all.deb

Install ROCm with the compute use case (choose Y when prompted to overwrite amdgpu.list):

amdgpu-install -y --usecase=rocm
sudo reboot

Add your user to the required groups:

sudo usermod -a -G render,video $LOGNAME
sudo reboot

Verify ROCm installation:

rocminfo

You should see your GPU listed as an agent with detailed properties.

Step 4: Configure ROCm Libraries

Configure the system to find ROCm shared libraries:

# Add ROCm library paths
sudo tee --append /etc/ld.so.conf.d/rocm.conf <<EOF
/opt/rocm/lib
/opt/rocm/lib64
EOF

sudo ldconfig

# Set library path environment variable (add to ~/.bashrc for persistence)
export LD_LIBRARY_PATH=/opt/rocm-7.0.0/lib:$LD_LIBRARY_PATH

Install and verify OpenCL runtime:

sudo apt install rocm-opencl-runtime
clinfo

The clinfo command should display information about your AMD GPU.

Step 5: Install PyTorch with ROCm Support

Create a conda environment and install PyTorch:

# Create conda environment
conda create -n pt2.8-rocm7 python=3.12
conda activate pt2.8-rocm7

# Install PyTorch 2.8.0 with ROCm 7.0 from AMD's repository
pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0/pytorch_triton_rocm-3.2.0%2Brocm7.0.0.4d510c3a44-cp312-cp312-linux_x86_64.whl
pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0/torch-2.8.0%2Brocm7.0.0-cp312-cp312-linux_x86_64.whl
pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0/torchvision-0.23.0%2Brocm7.0.0-cp312-cp312-linux_x86_64.whl
pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0/torchaudio-2.8.0%2Brocm7.0.0-cp312-cp312-linux_x86_64.whl

# Install GCC 12.1 (required for some operations)
conda install -c conda-forge gcc=12.1.0

Important Notes: - The URLs above are for Python 3.12 (cp312). Adjust for your Python version if different. - These wheels are built specifically for ROCm 7.0 and may not work with other ROCm versions. - The LD_LIBRARY_PATH must be set correctly, or PyTorch won't find ROCm libraries.

Verifying Installation

After installation, perform a quick verification:

import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device count: {torch.cuda.device_count()}")
if torch.cuda.is_available():
    print(f"Device name: {torch.cuda.get_device_name(0)}")

Note that despite using ROCm, PyTorch still refers to the GPU API as "CUDA" for compatibility reasons. This is intentional and allows CUDA-based code to run on AMD GPUs without modification.

Comprehensive GPU Testing

To thoroughly validate that PyTorch is working correctly with the MAX+ 395, we developed a comprehensive test suite that exercises various aspects of GPU computing.

Test Suite Overview

Our test suite includes five major components:

Installation Verification: Confirms PyTorch version and GPU detection
ROCm Availability Check: Validates GPU properties and capabilities
Tensor Operations: Tests basic tensor creation and mathematical operations
Neural Network Operations: Validates deep learning functionality
Memory Management: Tests GPU memory allocation and deallocation

Test Script

Here's the complete test script we developed:

#!/usr/bin/env python3
"""
ROCm PyTorch GPU Test POC
Tests if ROCm PyTorch can successfully detect and use AMD GPUs
"""

import torch
import sys

def print_section(title):
    """Print a formatted section header"""
    print(f"\n{'='*60}")
    print(f" {title}")
    print(f"{'='*60}")

def test_pytorch_installation():
    """Test basic PyTorch installation"""
    print_section("PyTorch Installation Info")
    print(f"PyTorch Version: {torch.__version__}")
    print(f"Python Version: {sys.version}")

def test_rocm_availability():
    """Test ROCm/CUDA availability"""
    print_section("ROCm/CUDA Availability")

    cuda_available = torch.cuda.is_available()
    print(f"CUDA Available: {cuda_available}")

    if cuda_available:
        print(f"CUDA Device Count: {torch.cuda.device_count()}")
        print(f"Current Device: {torch.cuda.current_device()}")
        print(f"Device Name: {torch.cuda.get_device_name(0)}")

        props = torch.cuda.get_device_properties(0)
        print(f"\nDevice Properties:")
        print(f"  - Total Memory: {props.total_memory / 1024**3:.2f} GB")
        print(f"  - Multi Processor Count: {props.multi_processor_count}")
        print(f"  - CUDA Capability: {props.major}.{props.minor}")
    else:
        print("No CUDA/ROCm devices detected!")
        return False

    return True

def test_tensor_operations():
    """Test basic tensor operations on GPU"""
    print_section("Tensor Operations Test")

    try:
        cpu_tensor = torch.randn(1000, 1000)
        print(f"CPU Tensor created: {cpu_tensor.shape}")
        print(f"CPU Tensor device: {cpu_tensor.device}")

        gpu_tensor = cpu_tensor.cuda()
        print(f"\nGPU Tensor created: {gpu_tensor.shape}")
        print(f"GPU Tensor device: {gpu_tensor.device}")

        print("\nPerforming matrix multiplication on GPU...")
        result = torch.matmul(gpu_tensor, gpu_tensor)
        print(f"Result shape: {result.shape}")
        print(f"Result device: {result.device}")

        cpu_result = result.cpu()
        print(f"Moved result back to CPU: {cpu_result.device}")

        print("\n✓ Tensor operations successful!")
        return True

    except Exception as e:
        print(f"\n✗ Tensor operations failed: {e}")
        return False

def test_simple_neural_network():
    """Test a simple neural network operation on GPU"""
    print_section("Neural Network Test")

    try:
        model = torch.nn.Sequential(
            torch.nn.Linear(100, 50),
            torch.nn.ReLU(),
            torch.nn.Linear(50, 10)
        )

        print("Model created on CPU")
        print(f"Model device: {next(model.parameters()).device}")

        model = model.cuda()
        print(f"Model moved to GPU: {next(model.parameters()).device}")

        input_data = torch.randn(32, 100).cuda()
        print(f"\nInput data shape: {input_data.shape}")
        print(f"Input data device: {input_data.device}")

        print("Performing forward pass...")
        output = model(input_data)
        print(f"Output shape: {output.shape}")
        print(f"Output device: {output.device}")

        print("\n✓ Neural network test successful!")
        return True

    except Exception as e:
        print(f"\n✗ Neural network test failed: {e}")
        return False

def test_memory_management():
    """Test GPU memory management"""
    print_section("GPU Memory Management Test")

    try:
        if torch.cuda.is_available():
            print(f"Allocated Memory: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB")
            print(f"Cached Memory: {torch.cuda.memory_reserved(0) / 1024**2:.2f} MB")

            tensors = []
            for i in range(5):
                tensors.append(torch.randn(1000, 1000).cuda())

            print(f"\nAfter allocating 5 tensors:")
            print(f"Allocated Memory: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB")
            print(f"Cached Memory: {torch.cuda.memory_reserved(0) / 1024**2:.2f} MB")

            del tensors
            torch.cuda.empty_cache()

            print(f"\nAfter clearing cache:")
            print(f"Allocated Memory: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB")
            print(f"Cached Memory: {torch.cuda.memory_reserved(0) / 1024**2:.2f} MB")

            print("\n✓ Memory management test successful!")
            return True
        else:
            print("No GPU available for memory test")
            return False

    except Exception as e:
        print(f"\n✗ Memory management test failed: {e}")
        return False

def main():
    """Run all tests"""
    print("\n" + "="*60)
    print(" ROCm PyTorch GPU Test POC")
    print("="*60)

    test_pytorch_installation()

    if not test_rocm_availability():
        print("\n" + "="*60)
        print(" FAILED: No ROCm/CUDA devices available")
        print("="*60)
        sys.exit(1)

    results = []
    results.append(("Tensor Operations", test_tensor_operations()))
    results.append(("Neural Network", test_simple_neural_network()))
    results.append(("Memory Management", test_memory_management()))

    print_section("Test Summary")
    all_passed = True
    for test_name, passed in results:
        status = "✓ PASSED" if passed else "✗ FAILED"
        print(f"{test_name}: {status}")
        if not passed:
            all_passed = False

    print("\n" + "="*60)
    if all_passed:
        print(" SUCCESS: All tests passed! ROCm GPU is working.")
    else:
        print(" PARTIAL SUCCESS: Some tests failed.")
    print("="*60 + "\n")

    return 0 if all_passed else 1

if __name__ == "__main__":
    sys.exit(main())

Test Results and Analysis

Running our comprehensive test suite on the MAX+ 395 yielded excellent results across all categories.

GPU Detection and Properties

The first test confirmed that PyTorch successfully detected the AMD GPU:

CUDA Available: True
CUDA Device Count: 1
Current Device: 0
Device Name: AMD Radeon Graphics

Device Properties:
  - Total Memory: 96.00 GB
  - Multi Processor Count: 20
  - CUDA Capability: 11.5

The 96GB of memory is particularly impressive, far exceeding what's available on most consumer or even professional NVIDIA GPUs. This massive memory capacity opens up possibilities for:

Training larger models without splitting across multiple GPUs
Processing high-resolution images or long sequences
Handling larger batch sizes for improved training efficiency
Running multiple models simultaneously

Tensor Operations Performance

Basic tensor operations executed flawlessly:

CPU Tensor created: torch.Size([1000, 1000])
CPU Tensor device: cpu

GPU Tensor created: torch.Size([1000, 1000])
GPU Tensor device: cuda:0

Performing matrix multiplication on GPU...
Result shape: torch.Size([1000, 1000])
Result device: cuda:0
Moved result back to CPU: cpu

✓ Tensor operations successful!

The seamless movement of tensors between CPU and GPU memory, along with successful matrix multiplication, confirms that the fundamental PyTorch operations work correctly on ROCm.

Neural Network Operations

Our neural network test validated that PyTorch's high-level APIs work correctly:

Model created on CPU
Model device: cpu
Model moved to GPU: cuda:0

Input data shape: torch.Size([32, 100])
Input data device: cuda:0
Performing forward pass...
Output shape: torch.Size([32, 10])
Output device: cuda:0

✓ Neural network test successful!

This test confirms that: - Models can be moved to GPU with the .cuda() method - Forward passes execute correctly on GPU - All layers (Linear, ReLU) are properly accelerated

Memory Management

The memory management test showed efficient allocation and deallocation:

Allocated Memory: 32.00 MB
Cached Memory: 54.00 MB

After allocating 5 tensors:
Allocated Memory: 52.00 MB
Cached Memory: 54.00 MB

After clearing cache:
Allocated Memory: 32.00 MB
Cached Memory: 32.00 MB

PyTorch's memory management on ROCm works identically to CUDA, with proper caching behavior and the ability to manually clear cached memory when needed.

Performance Considerations

Memory Bandwidth

The MAX+ 395's 96GB of memory is a significant advantage, but memory bandwidth is equally important for deep learning workloads. The W7900's memory subsystem provides substantial bandwidth for data transfers between GPU memory and compute units.

Compute Performance

With 20 compute units, the MAX+ 395 provides substantial parallel processing capability. While direct comparisons to NVIDIA GPUs depend on the specific workload, ROCm's optimization for AMD architectures ensures efficient utilization of available compute resources.

Software Maturity

ROCm has matured significantly over recent years. Most PyTorch operations that work on CUDA now work seamlessly on ROCm. However, some edge cases and newer features may still have better support on CUDA, so testing your specific workload is recommended.

Practical Tips and Best Practices

Code Portability

To write code that works on both CUDA and ROCm:

# Use device-agnostic code
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
inputs = inputs.to(device)

Monitoring GPU Utilization

Use rocm-smi to monitor GPU utilization:

watch -n 1 rocm-smi

This provides real-time information about GPU usage, memory consumption, temperature, and power draw.

Optimizing Memory Usage

With 96GB available, you might be tempted to use very large batch sizes. However, optimal batch size depends on many factors:

# Experiment with batch sizes
for batch_size in [32, 64, 128, 256]:
    # Train and measure throughput
    # Find the sweet spot between memory usage and performance

Debugging

Enable PyTorch's anomaly detection during development:

torch.autograd.set_detect_anomaly(True)

Troubleshooting Common Issues

GPU Not Detected

If torch.cuda.is_available() returns False:

Verify ROCm installation: rocm-smi
Check PyTorch was installed with ROCm support: print(torch.__version__) should show +rocm
Ensure ROCm drivers match PyTorch's ROCm version

Out of Memory Errors

Even with 96GB, you can run out of memory:

# Clear cache periodically
torch.cuda.empty_cache()

# Use gradient checkpointing for large models
from torch.utils.checkpoint import checkpoint

Performance Issues

If training is slower than expected:

Profile your code: torch.profiler.profile()
Check for CPU-GPU transfer bottlenecks
Verify data loading isn't the bottleneck
Consider using mixed precision training with torch.cuda.amp

Conclusion

The AMD Radeon Pro W7900 (MAX+ 395) with ROCm provides a robust, capable platform for PyTorch-based machine learning workloads. Our comprehensive testing demonstrated that:

PyTorch 2.8.0 with ROCm 7.0.0 works seamlessly with the MAX+ 395
All tested operations (tensors, neural networks, memory management) function correctly
The massive 96GB memory capacity enables unique use cases
Code written for CUDA generally works without modification

For organizations invested in AMD hardware or looking for alternatives to NVIDIA's ecosystem, the MAX+ 395 with ROCm represents a viable option for deep learning workloads. The open-source nature of ROCm and PyTorch's strong support for the platform ensure that AMD GPUs are first-class citizens in the deep learning community.

As ROCm continues to evolve and PyTorch support deepens, AMD's GPU offerings will only become more compelling for machine learning practitioners. The MAX+ 395, with its exceptional memory capacity and solid compute performance, stands ready to tackle demanding deep learning tasks.

Acknowledgments

The detailed ROCm 7.0 installation procedure is based on Wei Lu's excellent article "Ultralytics YOLO/SAM with ROCm 7.0 on AMD Ryzen AI Max+395 'Strix Halo'" published on Medium in October 2025. Wei Lu's pioneering work in documenting the complete bootstrapping process for ROCm 7.0 on the Max+395 platform made this possible.

Resources

Based on real-world testing performed on October 10, 2025, using PyTorch 2.8.0 with ROCm 7.0.0 on an AMD Radeon Pro W7900 GPU with 96GB memory. Installation instructions adapted from Wei Lu's documentation of the AMD Ryzen AI Max+395 platform.

Transfer Learning for Predictive Custom Drag Modeling: Automated Generation of Drag Coefficient Curves Using Multi-Modal AI

A.C. Jokela — Fri, 10 Oct 2025 18:10:00 GMT

TL;DR

We built a neural network that predicts full drag coefficient curves (41 Mach points from 0.5 to 4.5) for rifle bullets using only basic specifications like weight, caliber, and ballistic coefficient. The system achieves 3.15% mean absolute error and has been serving predictions in production since September 2025. This post walks through the technical implementation details, architecture decisions, and lessons learned building a real-world ML system for ballistic physics.

Read the full whitepaper: Transfer Learning for Predictive Custom Drag Modeling (17 pages)

The Problem: Drag Curves Are Scarce, But Critical

If you've ever built a ballistic calculator, you know the challenge: accurate drag modeling is everything. Standard drag models (G1, G7, G8) work okay for "average" bullets, but modern precision shooting demands better. Custom Drag Models (CDMs), full drag coefficient curves measured with doppler radar, are the gold standard. They capture the unique aerodynamic signature of each bullet design.

The catch? Getting a CDM requires: - Access to a doppler radar range (≈$500K+ equipment) - Firing 50-100 rounds at various velocities - Expert analysis to process the raw data - Cost: $5,000-$15,000 per bullet

For manufacturers like Hornady and Lapua, this is routine. For smaller manufacturers or custom bullet makers? Not happening. We had 641 bullets with real radar-measured CDMs and thousands of bullets with only basic specs. Could we use machine learning to bridge the gap?

The Vision: Transfer Learning from Radar Data

The core insight: bullets with similar physical characteristics have similar drag curves. A 168gr .308 boattail match bullet from Manufacturer A will drag similarly to one from Manufacturer B. We could train a neural network on our 641 radar-measured bullets and use transfer learning to predict CDMs for bullets we've never measured.

But we faced an immediate data problem: 641 samples isn't much for deep learning. Enter synthetic data augmentation.

Part 1: Automating Data Extraction with Claude Vision

Applied Ballistics publishes ballistic data for 704+ bullets as JPEG images. Manual data entry would take 1,408 hours (704 bullets × 2 hours each). We needed automation.

The Vision Processing Pipeline

We built an extraction pipeline using Claude 3.5 Sonnet's vision capabilities:

import anthropic
import base64
from pathlib import Path

def extract_bullet_data(image_path: str) -> dict:
    """Extract bullet specifications from AB datasheet JPEG."""
    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

    # Load and encode image
    with open(image_path, "rb") as f:
        image_data = base64.standard_b64encode(f.read()).decode("utf-8")

    # Vision extraction prompt
    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/jpeg",
                        "data": image_data,
                    },
                },
                {
                    "type": "text",
                    "text": """Extract the following from this Applied Ballistics bullet datasheet:
                    - Caliber (inches, decimal format)
                    - Bullet weight (grains)
                    - G1 Ballistic Coefficient
                    - G7 Ballistic Coefficient
                    - Bullet length (inches, if visible)
                    - Ogive radius (calibers, if visible)

                    Return as JSON with keys: caliber, weight_gr, bc_g1, bc_g7, length_in, ogive_radius_cal"""
                }
            ],
        }]
    )

    # Parse response
    data = json.loads(message.content[0].text)

    # Physics validation
    validate_bullet_physics(data)

    return data

def validate_bullet_physics(data: dict):
    """Sanity checks for extracted data."""
    caliber = data['caliber']
    weight = data['weight_gr']

    # Caliber bounds
    assert 0.172 <= caliber <= 0.50, f"Invalid caliber: {caliber}"

    # Weight-to-caliber ratio (sectional density proxy)
    ratio = weight / (caliber  3)
    assert 0.5 <= ratio <= 2.0, f"Implausible weight for caliber: {weight}gr @ {caliber}in"

    # BC sanity
    assert 0.1 <= data['bc_g1'] <= 1.2, f"Invalid G1 BC: {data['bc_g1']}"
    assert 0.1 <= data['bc_g7'] <= 0.9, f"Invalid G7 BC: {data['bc_g7']}"

Figure 2: Claude Vision extraction pipeline - from JPEG datasheets to structured bullet specifications

Results: - 704/704 successful extractions (100% success rate) - 2.3 seconds per bullet (average) - 27 minutes total vs. 1,408 hours manual - 99.97% time savings

We validated against a manually-verified subset of 50 bullets: - 100% match on caliber - 98% match on weight (±0.5 grain tolerance) - 96% match on BC values (±0.002 tolerance)

The vision model occasionally struggled with hand-drawn or low-quality scans, but the physics validation caught these errors before they corrupted our dataset.

Part 2: Generating Synthetic CDM Curves

Now we had 704 bullets with BC values but no full CDM curves. We needed to synthesize them.

The BC-to-CDM Transformation Algorithm

The relationship between ballistic coefficient and drag coefficient is straightforward:

BC = m / (C_d × d²)

Rearranging:
C_d(M) = m / (BC(M) × d²)

But BC values are typically single scalars, not curves. We developed a 5-step hybrid algorithm combining standard drag model references with BC-derived corrections:

Step 1: Base Reference Curve

Start with the G7 standard drag curve as a baseline (better for modern boattail bullets than G1):

def get_g7_reference_curve(mach_points: np.ndarray) -> np.ndarray:
    """G7 standard drag curve from McCoy (1999)."""
    # Precomputed G7 curve at 41 Mach points
    return interpolate_standard_curve("G7", mach_points)

Step 2: BC-Based Scaling

Scale the reference curve using extracted BC values:

def scale_by_bc(cd_base: np.ndarray, bc_actual: float, bc_reference: float = 0.221) -> np.ndarray:
    """Scale drag curve to match actual BC.

    BC_G7_ref = 0.221 (G7 standard projectile)
    """
    scaling_factor = bc_reference / bc_actual
    return cd_base * scaling_factor

Step 3: Multi-Regime Interpolation

When both G1 and G7 BCs are available, blend them based on Mach regime:

def blend_drag_models(mach: np.ndarray, cd_g1: np.ndarray, cd_g7: np.ndarray) -> np.ndarray:
    """Blend G1 and G7 curves based on flight regime.

    - Supersonic (M > 1.2): Use G1 (better for shock wave region)
    - Transonic (0.8 < M < 1.2): Cubic spline interpolation
    - Subsonic (M < 0.8): Use G7 (better for low-speed)
    """
    cd_blended = np.zeros_like(mach)

    for i, M in enumerate(mach):
        if M > 1.2:
            # Supersonic: G1 better captures shock effects
            cd_blended[i] = cd_g1[i]
        elif M < 0.8:
            # Subsonic: G7 better for boattail bullets
            cd_blended[i] = cd_g7[i]
        else:
            # Transonic: smooth interpolation
            t = (M - 0.8) / 0.4  # Normalize to [0, 1]
            cd_blended[i] = cubic_interpolate(cd_g7[i], cd_g1[i], t)

    return cd_blended

Step 4: Transonic Peak Generation

Model the transonic drag spike using a Gaussian kernel:

def add_transonic_peak(cd_base: np.ndarray, mach: np.ndarray,
                       bc_g1: float, bc_g7: float) -> np.ndarray:
    """Add realistic transonic drag spike.

    Peak amplitude calibrated from BC ratio (G1 worse than G7 in transonic).
    """
    # Estimate peak amplitude from BC discrepancy
    bc_ratio = bc_g1 / bc_g7
    peak_amplitude = 0.15 * (bc_ratio - 1.0)  # Empirically tuned

    # Gaussian centered at critical Mach
    M_crit = 1.0
    sigma = 0.15

    transonic_spike = peak_amplitude * np.exp(-((mach - M_crit)  2) / (2 * sigma  2))

    return cd_base + transonic_spike

Step 5: Monotonicity Enforcement

Apply Savitzky-Golay smoothing to prevent unphysical oscillations:

from scipy.signal import savgol_filter

def enforce_smoothness(cd_curve: np.ndarray, window_length: int = 7, polyorder: int = 3) -> np.ndarray:
    """Smooth drag curve while preserving transonic peak.

    Savitzky-Golay filter preserves peak shape better than moving average.
    """
    # Must have odd window length
    if window_length % 2 == 0:
        window_length += 1

    return savgol_filter(cd_curve, window_length, polyorder, mode='nearest')

Validation Against Ground Truth

We validated synthetic curves against 127 bullets where both BC values and full CDM curves were available:

Metric	Value	Notes
Mean Absolute Error	3.2%	Across all Mach points
Transonic Error	4.8%	Mach 0.8-1.2 (most challenging)
Supersonic Error	2.1%	Mach 1.5-3.0 (best performance)
Shape Correlation	r = 0.984	Pearson correlation

The synthetic curves satisfied all physics constraints: - Monotonic decrease in supersonic regime - Realistic transonic peaks (1.3-2.0× baseline) - Smooth transitions between regimes

Figure 3: Validation of synthetic CDM curves against ground truth radar measurements

Total training data: 1,345 bullets (704 synthetic + 641 real), 2.1x data augmentation.

Part 3: Architecture Exploration

With data ready, we explored four neural architectures:

1. Multi-Layer Perceptron (Baseline)

Simple feedforward network:

import torch
import torch.nn as nn

class CDMPredictor(nn.Module):
    """MLP for CDM prediction: 13 features → 41 Cd values."""

    def __init__(self, dropout: float = 0.2):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(13, 256),
            nn.ReLU(),
            nn.Dropout(dropout),

            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Dropout(dropout),

            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Dropout(dropout),

            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(dropout),

            nn.Linear(256, 41)  # Output: 41 Mach points
        )

    def forward(self, x):
        return self.network(x)

Input Features (13 total):

features = [
    'caliber',           # inches
    'weight_gr',         # grains
    'bc_g1',            # G1 ballistic coefficient
    'bc_g7',            # G7 ballistic coefficient
    'length_in',        # bullet length (imputed if missing)
    'ogive_radius_cal', # ogive radius in calibers
    'meplat_diam_in',   # meplat diameter
    'boat_tail_angle',  # boattail angle (degrees)
    'bearing_length',   # bearing surface length
    'sectional_density', # weight / caliber²
    'form_factor_g1',   # i / BC_G1
    'form_factor_g7',   # i / BC_G7
    'length_to_diameter' # L/D ratio
]

Figure 4: MLP architecture - 13 input features through 4 hidden layers to 41 output Mach points

2. Physics-Informed Neural Network (PINN)

Added physics loss term enforcing drag model constraints:

class PINN_CDMPredictor(nn.Module):
    """Physics-Informed NN with drag equation constraints."""

    def __init__(self):
        super().__init__()
        # Same architecture as MLP
        self.network = build_mlp_network()

    def physics_loss(self, cd_pred: torch.Tensor, features: torch.Tensor, mach: torch.Tensor) -> torch.Tensor:
        """Enforce physics constraints on predictions.

        Constraints:
        1. Drag increases with Mach in subsonic
        2. Transonic peak exists near M=1
        3. Monotonic decrease in supersonic
        """
        # Constraint 1: Subsonic gradient
        subsonic_mask = mach < 0.8
        subsonic_cd = cd_pred[subsonic_mask]
        subsonic_grad = torch.diff(subsonic_cd)
        subsonic_violation = torch.relu(-subsonic_grad).sum()  # Penalize decreases

        # Constraint 2: Transonic peak
        transonic_mask = (mach >= 0.8) & (mach <= 1.2)
        transonic_cd = cd_pred[transonic_mask]
        peak_violation = torch.relu(1.1 - transonic_cd.max()).sum()  # Must exceed 1.1

        # Constraint 3: Supersonic monotonicity
        supersonic_mask = mach > 1.5
        supersonic_cd = cd_pred[supersonic_mask]
        supersonic_grad = torch.diff(supersonic_cd)
        supersonic_violation = torch.relu(supersonic_grad).sum()  # Penalize increases

        return subsonic_violation + peak_violation + supersonic_violation

def total_loss(cd_pred, cd_true, features, mach, lambda_physics=0.1):
    """Combined data + physics loss."""
    data_loss = nn.MSELoss()(cd_pred, cd_true)
    physics_loss = model.physics_loss(cd_pred, features, mach)

    return data_loss + lambda_physics * physics_loss

Result: Over-regularization. Physics loss was too strict, preventing the model from learning subtle variations. Performance degraded to 4.86% MAE.

3. Transformer Architecture

Treated the 41 Mach points as a sequence:

class TransformerCDM(nn.Module):
    """Transformer encoder for sequence-to-sequence CDM prediction."""

    def __init__(self, d_model=128, nhead=8, num_layers=4):
        super().__init__()

        self.feature_embedding = nn.Linear(13, d_model)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=512,
            dropout=0.1
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        self.output_head = nn.Linear(d_model, 41)

    def forward(self, x):
        # x: [batch, 13]
        embedded = self.feature_embedding(x)  # [batch, d_model]
        embedded = embedded.unsqueeze(1).expand(-1, 41, -1)  # [batch, 41, d_model]

        transformed = self.transformer(embedded)  # [batch, 41, d_model]

        cd_pred = self.output_head(transformed[:, :, :]).squeeze(-1)  # [batch, 41]

        return cd_pred

Result: Mismatch between architecture and problem. CDM prediction isn't a sequence modeling task; Mach points are independent given bullet features. Performance: 6.05% MAE.

4. Neural ODE

Attempted to model drag as a continuous ODE:

from torchdiffeq import odeint

class DragODE(nn.Module):
    """Neural ODE for continuous drag modeling."""

    def __init__(self, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1 + 13, hidden_dim),  # Mach + features
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1)  # dCd/dM
        )

    def forward(self, t, state):
        # t: current Mach number
        # state: [Cd, features...]
        return self.net(torch.cat([t, state], dim=-1))

def predict_cdm(features, mach_points):
    """Integrate ODE to get Cd curve."""
    initial_cd = torch.tensor([0.5])  # Initial guess
    state = torch.cat([initial_cd, features])

    solution = odeint(ode_func, state, mach_points)

    return solution[:, 0]  # Extract Cd values

Result: Failed to converge due to dimension mismatch errors and extreme sensitivity to initial conditions. Abandoned after 2 days of debugging.

Architecture Comparison Results

Architecture	MAE	Smoothness	Shape Correlation	Status
MLP Baseline	3.66%	90.05%	0.9380	✅ Best
Physics-Informed NN	4.86%	64.02%	0.8234	❌ Over-regularized
Transformer	6.05%	56.83%	0.7891	❌ Poor fit
Neural ODE	---	---	---	❌ Failed to converge

Figure 5: Performance comparison across four neural architectures - MLP baseline wins

Key Insight: Simple MLP with dropout outperformed complex physics-constrained models. The training data already contained sufficient physics signal; explicit constraints hurt generalization.

Part 4: Production System Design

The POC model (3.66% MAE) validated the approach. Now we needed production hardening.

Training Pipeline Improvements

import pytorch_lightning as pl
from torch.utils.data import DataLoader, TensorDataset

class ProductionCDMModel(pl.LightningModule):
    """Production-ready CDM predictor with monitoring."""

    def __init__(self, learning_rate=1e-3, weight_decay=1e-4):
        super().__init__()
        self.save_hyperparameters()

        self.model = CDMPredictor(dropout=0.2)
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay

        # Metrics tracking
        self.train_mae = []
        self.val_mae = []

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        features, cd_true = batch
        cd_pred = self(features)

        # Weighted MSE loss (emphasize transonic region)
        weights = self._get_mach_weights()
        loss = (weights * (cd_pred - cd_true)  2).mean()

        # Metrics
        mae = torch.abs(cd_pred - cd_true).mean()
        self.log('train_loss', loss)
        self.log('train_mae', mae)

        return loss

    def validation_step(self, batch, batch_idx):
        features, cd_true = batch
        cd_pred = self(features)

        loss = nn.MSELoss()(cd_pred, cd_true)
        mae = torch.abs(cd_pred - cd_true).mean()

        self.log('val_loss', loss)
        self.log('val_mae', mae)

        # Physics validation
        smoothness = self._calculate_smoothness(cd_pred)
        transonic_quality = self._check_transonic_peak(cd_pred)

        self.log('smoothness', smoothness)
        self.log('transonic_quality', transonic_quality)

        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=self.learning_rate,
            weight_decay=self.weight_decay
        )

        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            mode='min',
            factor=0.5,
            patience=5,
            verbose=True
        )

        return {
            'optimizer': optimizer,
            'lr_scheduler': scheduler,
            'monitor': 'val_loss'
        }

    def _get_mach_weights(self):
        """Weight transonic region more heavily."""
        weights = torch.ones(41)
        transonic_indices = (self.mach_points >= 0.8) & (self.mach_points <= 1.2)
        weights[transonic_indices] = 2.0  # 2x weight in transonic
        return weights / weights.sum()

    def _calculate_smoothness(self, cd_pred):
        """Measure curve smoothness (low = better)."""
        second_derivative = torch.diff(cd_pred, n=2, dim=-1)
        return 1.0 / (1.0 + second_derivative.abs().mean())

    def _check_transonic_peak(self, cd_pred):
        """Verify transonic peak exists and is realistic."""
        transonic_mask = (self.mach_points >= 0.8) & (self.mach_points <= 1.2)
        peak_cd = cd_pred[:, transonic_mask].max(dim=1)[0]
        baseline_cd = cd_pred[:, 0]  # Subsonic baseline

        return (peak_cd / baseline_cd).mean()  # Should be > 1.0

Training Configuration

# Data preparation
X_train, X_val, X_test = prepare_features()  # 1,039 → 831 / 104 / 104
y_train, y_val, y_test = prepare_targets()

train_dataset = TensorDataset(X_train, y_train)
val_dataset = TensorDataset(X_val, y_val)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

# Model training
model = ProductionCDMModel(learning_rate=1e-3, weight_decay=1e-4)

trainer = pl.Trainer(
    max_epochs=100,
    callbacks=[
        pl.callbacks.EarlyStopping(monitor='val_loss', patience=10, mode='min'),
        pl.callbacks.ModelCheckpoint(monitor='val_mae', mode='min', save_top_k=3),
        pl.callbacks.LearningRateMonitor(logging_interval='epoch')
    ],
    accelerator='gpu',
    devices=1,
    log_every_n_steps=10
)

trainer.fit(model, train_loader, val_loader)

Figure 6: Training and validation loss convergence over 60 epochs

Training Results: - Converged at epoch 60 (early stopping) - Final validation loss: 0.0023 - Production model MAE: 3.15% (13.9% improvement over POC) - Smoothness: 88.81% (close to ground truth 89.6%) - Shape correlation: 0.9545

Figure 7: Example predicted CDM curves compared to ground truth measurements

API Integration

# ballistics/ml/cdm_transfer_learning.py

import torch
import pickle
from pathlib import Path

class CDMTransferLearning:
    """Production CDM prediction service."""

    def __init__(self, model_path: str = "models/cdm_transfer_learning/production_mlp.pkl"):
        self.model = self._load_model(model_path)
        self.model.eval()

        # Feature statistics for normalization
        with open(model_path.replace('.pkl', '_stats.pkl'), 'rb') as f:
            self.feature_stats = pickle.load(f)

    def predict(self, bullet_data: dict) -> dict:
        """Predict CDM curve from bullet specifications.

        Args:
            bullet_data: Dict with keys: caliber, weight_gr, bc_g1, bc_g7, etc.

        Returns:
            Dict with mach_numbers, drag_coefficients, validation_metrics
        """
        # Feature engineering
        features = self._extract_features(bullet_data)
        features_normalized = self._normalize_features(features)

        # Prediction
        with torch.no_grad():
            cd_pred = self.model(torch.tensor(features_normalized, dtype=torch.float32))

        # Denormalize
        cd_values = cd_pred.numpy()

        # Validation
        validation = self._validate_prediction(cd_values)

        return {
            'mach_numbers': self.mach_points.tolist(),
            'drag_coefficients': cd_values.tolist(),
            'source': 'ml_transfer_learning',
            'method': 'mlp_prediction',
            'validation': validation
        }

    def _validate_prediction(self, cd_values: np.ndarray) -> dict:
        """Physics validation of predicted curve."""
        return {
            'smoothness': self._calculate_smoothness(cd_values),
            'transonic_quality': self._check_transonic_peak(cd_values),
            'negative_cd_count': (cd_values < 0).sum(),
            'physical_plausibility': self._check_plausibility(cd_values)
        }

REST API Endpoint

# routes/bullets_unified.py

@bp.route('/search', methods=['GET'])
def search_bullets():
    """Search unified bullet database with optional CDM prediction."""
    query = request.args.get('q', '')
    use_cdm_prediction = request.args.get('use_cdm_prediction', 'true').lower() == 'true'

    # Search database
    results = search_database(query)

    cdm_predictions_made = 0

    if use_cdm_prediction:
        cdm_predictor = CDMTransferLearning()

        for bullet in results:
            if bullet.get('cdm_data') is None:
                # Predict CDM if not available
                try:
                    cdm_data = cdm_predictor.predict({
                        'caliber': bullet['caliber'],
                        'weight_gr': bullet['weight_gr'],
                        'bc_g1': bullet.get('bc_g1'),
                        'bc_g7': bullet.get('bc_g7'),
                        'length_in': bullet.get('length_in'),
                        'ogive_radius_cal': bullet.get('ogive_radius_cal')
                    })

                    bullet['cdm_data'] = cdm_data
                    bullet['cdm_predicted'] = True
                    cdm_predictions_made += 1

                except Exception as e:
                    logger.warning(f"CDM prediction failed for bullet {bullet['id']}: {e}")

    return jsonify({
        'results': results,
        'cdm_prediction_enabled': use_cdm_prediction,
        'cdm_predictions_made': cdm_predictions_made
    })

Example Response:

{
  "results": [
    {
      "id": 1234,
      "manufacturer": "Sierra",
      "model": "MatchKing",
      "caliber": 0.308,
      "weight_gr": 168,
      "bc_g1": 0.462,
      "bc_g7": 0.237,
      "cdm_data": {
        "mach_numbers": [0.5, 0.55, 0.6, ..., 4.5],
        "drag_coefficients": [0.287, 0.289, 0.295, ..., 0.312],
        "source": "ml_transfer_learning",
        "method": "mlp_prediction",
        "validation": {
          "smoothness": 91.2,
          "transonic_quality": 1.45,
          "negative_cd_count": 0,
          "physical_plausibility": true
        }
      },
      "cdm_predicted": true
    }
  ],
  "cdm_prediction_enabled": true,
  "cdm_predictions_made": 18
}

Part 5: Deployment and Monitoring

Model Serving Architecture

┌─────────────────┐
│   Client App    │
└────────┬────────┘
         │
         ▼
┌─────────────────────────┐
│  Google Cloud Function  │
│  (Python 3.12)          │
│  - Flask routing        │
│  - Request validation   │
│  - Response formatting  │
└────────┬────────────────┘
         │
         ▼
┌─────────────────────────┐
│  CDMTransferLearning    │
│  - PyTorch model (2.1MB)│
│  - CPU inference (<10ms)│
│  - Feature engineering  │
└────────┬────────────────┘
         │
         ▼
┌─────────────────────────┐
│  Physics Validation     │
│  - Smoothness check     │
│  - Peak detection       │
│  - Plausibility gates   │
└─────────────────────────┘

Performance Characteristics

Model Size: - PyTorch state dict: 2.1 MB - TorchScript (optional): 2.3 MB - ONNX (optional): 1.8 MB

Inference Speed (CPU): - Single prediction: 6-8 ms - Batch of 10: 12-15 ms (1.2-1.5 ms per bullet) - Batch of 100: 80-100 ms (0.8-1.0 ms per bullet)

Cold Start: - Model load time: 150-200 ms - First prediction: 220-280 ms (including load) - Subsequent predictions: 6-8 ms

Memory Footprint: - Model in memory: ~15 MB - Peak during inference: ~30 MB

Figure 8: Production inference performance metrics across different batch sizes

Monitoring and Observability

import newrelic.agent

class MonitoredCDMPredictor:
    """CDM predictor with New Relic monitoring."""

    def __init__(self):
        self.predictor = CDMTransferLearning()
        self.prediction_count = 0
        self.error_count = 0

    @newrelic.agent.function_trace()
    def predict(self, bullet_data: dict) -> dict:
        """Predict with telemetry."""
        self.prediction_count += 1

        try:
            # Track prediction time
            with newrelic.agent.FunctionTrace(name='cdm_prediction'):
                result = self.predictor.predict(bullet_data)

            # Custom metrics
            newrelic.agent.record_custom_metric('CDM/Predictions/Total', self.prediction_count)
            newrelic.agent.record_custom_metric('CDM/Validation/Smoothness',
                                               result['validation']['smoothness'])
            newrelic.agent.record_custom_metric('CDM/Validation/TransonicQuality',
                                               result['validation']['transonic_quality'])

            # Track feature availability
            features_available = sum(1 for k, v in bullet_data.items() if v is not None)
            newrelic.agent.record_custom_metric('CDM/Features/Available', features_available)

            return result

        except Exception as e:
            self.error_count += 1
            newrelic.agent.record_custom_metric('CDM/Errors/Total', self.error_count)
            newrelic.agent.notice_error()
            raise

Key Metrics Tracked: - Prediction latency (p50, p95, p99) - Validation scores (smoothness, transonic quality) - Feature availability (how many inputs provided) - Error rate and types - Cache hit rate (if caching enabled)

Lessons Learned

1. Simple Architectures Often Win

We spent a week exploring Transformers and Neural ODEs, only to find the vanilla MLP performed best. Why?

Data alignment: Our problem is function approximation, not sequence modeling
Inductive bias mismatch: Transformers expect temporal dependencies; drag curves don't have them
Regularization sufficiency: Dropout + weight decay provided enough regularization without physics constraints

Lesson: Start simple. Add complexity only when data clearly demands it.

2. Physics Validation > Physics Loss

Hard-coded physics loss functions became a liability: - Over-constrained the model - Required manual tuning of loss weights - Didn't generalize to all bullet types

Better approach: Validate predictions post-hoc and flag anomalies. Let the model learn physics from data.

3. Synthetic Data Quality Matters More Than Quantity

We generated 704 synthetic CDMs, but spent equal time validating them. Key insight: One bad synthetic sample can poison dozens of real samples during training.

Validation process: 1. Compare synthetic vs. real CDMs (where both exist) 2. Physics plausibility checks 3. Cross-validation with different BC values 4. Manual inspection of outliers

4. Feature Engineering > Model Complexity

The most impactful changes weren't architectural: - Adding sectional_density as a feature: -0.8% MAE - Computing form_factor_g1 and form_factor_g7: -0.6% MAE - Imputing missing features (length, ogive) using physics-based defaults: -0.5% MAE

Figure 9: Feature importance analysis showing impact of each input feature on prediction accuracy

Combined improvement: -1.9% MAE with zero code changes to the model.

5. Production Deployment ≠ POC

Our POC model worked great in notebooks. Production required: - Input validation and sanitization - Graceful degradation when features missing - Physics validation gates - Monitoring and alerting - Model versioning and rollback capability - A/B testing infrastructure

Time split: 30% research, 70% production engineering.

What's Next?

Phase 2: Uncertainty Quantification

Current model outputs point estimates. We're implementing Bayesian Neural Networks to provide confidence intervals:

class BayesianCDMPredictor(nn.Module):
    """Bayesian NN with dropout as approximate inference."""

    def predict_with_uncertainty(self, features, n_samples=100):
        """Monte Carlo dropout for uncertainty estimation."""
        self.train()  # Enable dropout during inference

        predictions = []
        for _ in range(n_samples):
            with torch.no_grad():
                pred = self(features)
                predictions.append(pred)

        predictions = torch.stack(predictions)

        mean = predictions.mean(dim=0)
        std = predictions.std(dim=0)

        return {
            'cd_mean': mean,
            'cd_std': std,
            'cd_lower': mean - 1.96 * std,  # 95% CI
            'cd_upper': mean + 1.96 * std
        }

Use case: Flag predictions with high uncertainty for manual review or experimental validation.

Conclusion

Building a production ML system for ballistic drag prediction required more than just training a model: - Data engineering (Claude Vision automation saved countless hours) - Synthetic data generation (2.1× data augmentation) - Architecture exploration (simple MLP won) - Real-world validation (94% physics check pass rate)

The result: 1,247 bullets now have accurate drag models that didn't exist before. Not bad for a side project.

Read the full technical whitepaper for mathematical derivations, validation details, and complete bibliography: cdm_transfer_learning.pdf

Resources

References: 1. McCoy, R. L. (1999). Modern Exterior Ballistics. Schiffer Publishing. 2. Litz, B. (2016). Applied Ballistics for Long Range Shooting (3rd ed.).

AMD AI Max+ 395 System Review: A Comprehensive Analysis

A.C. Jokela — Sun, 21 Sep 2025 20:25:28 GMT

Executive Summary

The AMD AI Max+ 395 system represents AMD's latest entry into the high-performance computing and AI acceleration market, featuring the company's cutting-edge Strix Halo architecture. This comprehensive review examines the system's performance characteristics, software compatibility, and overall viability for AI workloads and general computing tasks. While the hardware shows impressive potential with its 16-core CPU and integrated Radeon 8060S graphics, significant software ecosystem challenges, particularly with PyTorch/ROCm compatibility for the gfx1151 architecture, present substantial barriers to immediate adoption for AI development workflows.

Note: An Orange Pi 5 Max was photobombing this photograph

System Specifications and Architecture Overview

CPU Specifications

Processor: AMD RYZEN AI MAX+ 395 w/ Radeon 8060S
Architecture: x86_64 with Zen 5 cores
Cores/Threads: 16 cores / 32 threads
Base Clock: 599 MHz (minimum)
Boost Clock: 5,185 MHz (maximum)
Cache Configuration:
L1d Cache: 768 KiB (16 instances, 48 KiB per core)
L1i Cache: 512 KiB (16 instances, 32 KiB per core)
L2 Cache: 16 MiB (16 instances, 1 MiB per core)
L3 Cache: 64 MiB (2 instances, 32 MiB per CCX)
Instruction Set Extensions: Full AVX-512, AVX-VNNI, BF16 support

Memory Subsystem

Total System Memory: 32 GB DDR5
Memory Configuration: Unified memory architecture with shared GPU/CPU access
Memory Bandwidth: Achieved ~13.5 GB/s in multi-threaded tests

Graphics Processing Unit

GPU Architecture: Strix Halo (RDNA 3.5 based)
GPU Designation: gfx1151
Compute Units: 40 CUs (80 reported in ROCm, likely accounting for dual SIMD per CU)
Peak GPU Clock: 2,900 MHz
VRAM: 96 GB shared system memory (103 GB total addressable) - Note: This allocation was intentionally configured to maximize GPU memory for large language model inference
Memory Bandwidth: Shared with system memory
OpenCL Compute Units: 20 (as reported by clinfo)

Platform Details

Operating System: Ubuntu 24.04.3 LTS (Noble)
Kernel Version: 6.8.0-83-generic
Architecture: x86_64
Virtualization: AMD-V enabled

Performance Benchmarks

Figure 1: Comprehensive performance analysis and compatibility overview of the AMD AI Max+ 395 system

CPU Performance Analysis

Single-Threaded Performance

The sysbench CPU benchmark with prime number calculation revealed strong single-threaded performance:

Events per second: 6,368.92
Average latency: 0.16 ms
95th percentile latency: 0.16 ms

This performance places the AMD AI Max+ 395 in the upper tier of modern processors for single-threaded workloads, demonstrating the effectiveness of the Zen 5 architecture's IPC improvements and high boost clocks.

Multi-Threaded Performance

Multi-threaded testing across all 32 threads showed excellent scaling:

Events per second: 103,690.35
Scaling efficiency: 16.3x improvement over single-threaded (theoretical maximum 32x)
Thread fairness: Excellent distribution with minimal standard deviation

The scaling efficiency of approximately 51% indicates good multi-threading performance, though there's room for optimization in workloads that can fully utilize all available threads.

Memory Performance

Memory Bandwidth Testing

Memory performance testing using sysbench revealed:

Single-threaded bandwidth: 9.3 GB/s
Multi-threaded bandwidth: 13.5 GB/s (16 threads)
Latency characteristics: Sub-millisecond access times

The memory bandwidth results suggest the system is well-balanced for most workloads, though AI applications requiring extremely high memory bandwidth may find this a limiting factor compared to discrete GPU solutions with dedicated VRAM.

GPU Performance and Capabilities

Hardware Specifications

The integrated Radeon 8060S GPU presents impressive specifications on paper:

Architecture: RDNA 3.5 (Strix Halo)
Compute Units: 40 CUs with 2 SIMDs each
Memory Access: Full 96 GB of shared system memory
Clock Speed: Up to 2.9 GHz

OpenCL Capabilities

OpenCL enumeration reveals solid compute capabilities:

Device Type: GPU with full OpenCL 2.1 support
Max Compute Units: 20 (OpenCL reporting)
Max Work Group Size: 256
Image Support: Full 2D/3D image processing capabilities
Memory Allocation: Up to 87 GB maximum allocation

Network Performance Testing

Network infrastructure testing using iperf3 demonstrated excellent localhost performance:

Loopback Bandwidth: 122 Gbits/sec sustained
Latency: Minimal retransmissions (0 retries)
Consistency: Stable performance across 10-second test duration

This indicates robust internal networking capabilities suitable for distributed computing scenarios and high-bandwidth data transfer requirements.

PyTorch/ROCm Compatibility Analysis

Current State of ROCm Support

We installed ROCm 7.0 and related components: - ROCm Version: 7.0.0 - HIP Version: 7.0.51831 - PyTorch Version: 2.5.1+rocm6.2

gfx1151 Compatibility Issues

The most significant finding of this review centers on the gfx1151 architecture compatibility with current AI software stacks. Testing revealed critical limitations:

PyTorch Compatibility Problems

rocBLAS error: Cannot read TensileLibrary.dat: Illegal seek for GPU arch : gfx1151
List of available TensileLibrary Files:
- TensileLibrary_lazy_gfx1030.dat
- TensileLibrary_lazy_gfx906.dat
- TensileLibrary_lazy_gfx908.dat
- TensileLibrary_lazy_gfx942.dat
- TensileLibrary_lazy_gfx900.dat
- TensileLibrary_lazy_gfx90a.dat
- TensileLibrary_lazy_gfx1100.dat

This error indicates that PyTorch's ROCm backend lacks pre-compiled optimized kernels for the gfx1151 architecture. The absence of gfx1151 in the TensileLibrary files means:

No Optimized BLAS Operations: Matrix multiplication, convolutions, and other fundamental AI operations cannot leverage GPU acceleration
Training Workflows Broken: Most deep learning training pipelines will fail or fall back to CPU execution
Inference Limitations: Even basic neural network inference is compromised

Root Cause Analysis

The gfx1151 architecture represents a newer GPU design that hasn't been fully integrated into the ROCm software stack. While the hardware is detected and basic OpenCL operations function, the optimized compute libraries essential for AI workloads are missing.

Workaround Attempts

Testing various workarounds yielded limited success:

HSA_OVERRIDE_GFX_VERSION=11.0.0: Failed to resolve compatibility issues
CPU Fallback: PyTorch operates normally on CPU, but defeats the purpose of GPU acceleration
Basic GPU Operations: Simple tensor allocation succeeds, but compute operations fail

Software Ecosystem Gaps

Beyond PyTorch, the gfx1151 compatibility issues extend to:

TensorFlow: Likely similar rocBLAS dependency issues
JAX: ROCm backend compatibility uncertain
Scientific Computing: NumPy/SciPy GPU acceleration unavailable
Machine Learning Frameworks: Most frameworks dependent on rocBLAS will encounter issues

AMD GPU Software Support Ecosystem Analysis

Current State Assessment

AMD's GPU software ecosystem has made significant strides but remains fragmented compared to NVIDIA's CUDA platform:

Strengths

Open Source Foundation: ROCm's open-source nature enables community contributions
Standard API Support: OpenCL 2.1 and HIP provide industry-standard interfaces
Linux Integration: Strong kernel-level support through AMDGPU drivers
Professional Tools: rocm-smi and related utilities provide comprehensive monitoring

Weaknesses

Fragmented Architecture Support: New architectures like gfx1151 lag behind in software support
Limited Documentation: Less comprehensive than CUDA documentation
Smaller Developer Community: Fewer third-party tools and optimizations
Compatibility Matrix Complexity: Different software versions support different GPU architectures

Long-term Viability Concerns

The gfx1151 compatibility issues highlight broader ecosystem challenges:

Release Coordination Problems

Hardware releases outpace software ecosystem updates
Critical libraries (rocBLAS, Tensile) require architecture-specific optimization
Coordination between AMD hardware and software teams appears insufficient

Market Adoption Barriers

Developers hesitant to adopt platform with uncertain software support
Enterprise customers require guaranteed compatibility
Academic researchers need stable, well-documented platforms

Recommendations for AMD

Accelerated Software Development: Prioritize gfx1151 support in rocBLAS and related libraries
Pre-release Testing: Ensure software ecosystem readiness before hardware launches
Better Documentation: Comprehensive compatibility matrices and migration guides
Community Engagement: More responsive developer relations and support channels

Network Infrastructure and Connectivity

The system demonstrates excellent network performance characteristics suitable for modern computing workloads:

Internal Performance

Memory-to-Network Efficiency: 122 Gbps loopback performance indicates minimal bottlenecks
System Integration: Unified memory architecture benefits network-intensive applications
Scalability: Architecture suitable for distributed computing scenarios

External Connectivity Assessment

While specific external network testing wasn't performed, the system's infrastructure suggests:

Support for high-speed Ethernet (2.5GbE+)
Low-latency interconnects suitable for cluster computing
Adequate bandwidth for data center deployment scenarios

Power Efficiency and Thermal Characteristics

Limited thermal data was available during testing:

Idle Temperature: 29°C (GPU sensor)
Idle Power: 8.059W (GPU subsystem)
Thermal Management: Appears well-controlled under light loads

The unified architecture's power efficiency represents a significant advantage over discrete GPU solutions, particularly for mobile and edge computing applications.

Competitive Analysis

Comparison with Intel Arc

Intel's Arc GPUs face similar software ecosystem challenges, though Intel has made more aggressive investments in AI software stack development. The Arc series benefits from Intel's deeper software engineering resources but still lags behind NVIDIA in AI framework support.

Comparison with NVIDIA

NVIDIA maintains a substantial advantage in:

Software Maturity: CUDA ecosystem is mature and well-supported
AI Framework Integration: Native support across all major frameworks
Developer Tools: Comprehensive profiling and debugging tools
Documentation: Extensive, well-maintained documentation

AMD's advantages include:

Open Source Approach: More flexible licensing and community development
Unified Memory: Simplified programming model for certain applications
Cost: Potentially more cost-effective solutions

Market Positioning

The AMD AI Max+ 395 occupies a unique position as a high-performance integrated solution, but software limitations significantly impact its competitiveness in AI-focused markets.

Use Case Suitability Analysis

Recommended Use Cases

General Computing: Excellent performance for traditional computational workloads
Development Platforms: Strong for general software development (non-AI)
Edge Computing: Unified architecture benefits power-constrained deployments
Future AI Workloads: When software ecosystem matures

Not Recommended For

Current AI Development: gfx1151 compatibility issues are blocking
Production AI Inference: Unreliable software support
Machine Learning Research: Limited framework compatibility
Time-Critical Projects: Uncertain timeline for software fixes

Large Language Model Performance and Stability

Ollama LLM Inference Testing

Testing with Ollama reveals a mixed picture for LLM inference on the AMD AI Max+ 395 system. The platform successfully runs various models through CPU-based inference, though GPU acceleration faces significant challenges.

Performance Metrics

Testing with various model sizes revealed the following performance characteristics:

GPT-OSS 20B Model Performance:

Prompt evaluation rate: 61.29 tokens/second
Text generation rate: 8.99 tokens/second
Total inference time: ~13 seconds for 117 tokens
Memory utilization: ~54 GB VRAM usage

Llama 4 (67B) Model:

Successfully loads and runs
Generation coherent and accurate

The system demonstrates adequate performance for smaller models (20B parameters and below) when running through Ollama, though performance significantly lags behind NVIDIA GPUs with proper CUDA acceleration. The large unified memory configuration (96 GB VRAM, deliberately maximized for this testing) allows loading of substantial models that would typically require multiple GPUs or extensive system RAM on other platforms. This conscious decision to allocate maximum memory to the GPU was specifically made to evaluate the system's potential for large language model workloads.

Critical Stability Issues with Large Models

Driver Crashes with Advanced AI Workloads

Testing revealed severe stability issues when attempting to run larger models or when using AI-accelerated development tools:

Affected Scenarios:

Large Model Loading: GPT-OSS 120B model causes immediate amdgpu driver crashes
AI Development Tools: Continue.dev with certain LLMs triggers GPU reset
OpenAI Codex Integration: Consistent driver failures with models exceeding 70B parameters

GPU Reset Events

System logs reveal frequent GPU reset events during AI workload attempts:

[ 1030.960155] amdgpu 0000:c5:00.0: amdgpu: GPU reset begin!
[ 1033.972213] amdgpu 0000:c5:00.0: amdgpu: MODE2 reset
[ 1034.002615] amdgpu 0000:c5:00.0: amdgpu: GPU reset succeeded, trying to resume
[ 1034.003141] [drm] VRAM is lost due to GPU reset!
[ 1034.037824] amdgpu 0000:c5:00.0: amdgpu: GPU reset(1) succeeded!

These crashes result in:

Complete loss of VRAM contents
Application termination
Potential system instability requiring reboot
Interrupted workflows and data loss

Root Cause Analysis

The driver instability appears to stem from the same underlying issue as the PyTorch/ROCm incompatibility: immature driver support for the gfx1151 architecture. The drivers struggle with:

Memory Management: Large model allocations exceed driver's tested parameters
Compute Dispatch: Complex kernel launches trigger unhandled edge cases
Power State Transitions: Rapid load changes cause driver state machine failures
Synchronization Issues: Multi-threaded inference workloads expose race conditions

Implications for AI Development

The combination of LLM testing results and driver stability issues reinforces that the AMD AI Max+ 395 system, despite impressive hardware specifications, remains unsuitable for production AI workloads. The platform shows promise for future AI applications once driver maturity improves, but current limitations include:

Unreliable Large Model Support: Models over 70B parameters risk system crashes
Limited Tool Compatibility: Popular AI development tools cause instability
Workflow Interruptions: Frequent crashes disrupt development productivity
Data Loss Risk: VRAM resets can lose unsaved work or model states

Future Outlook and Development Roadmap

Short-term Expectations (3-6 months)

ROCm updates likely to address gfx1151 compatibility
PyTorch/TensorFlow support should improve
Community-driven workarounds may emerge

Medium-term Prospects (6-18 months)

Full AI framework support expected
Optimization improvements for Strix Halo architecture
Better documentation and developer resources

Long-term Considerations (18+ months)

AMD's commitment to open-source ecosystem should pay dividends
Potential for superior price/performance ratios
Growing developer community around ROCm platform

Conclusions and Recommendations

The AMD AI Max+ 395 system represents impressive hardware engineering with its unified memory architecture, strong CPU performance, and substantial GPU compute capabilities. However, critical software ecosystem gaps, particularly the gfx1151 compatibility issues with PyTorch and ROCm, severely limit its immediate utility for AI and machine learning workloads.

Key Findings Summary

Hardware Strengths:

Excellent CPU performance with 16 Zen 5 cores
Innovative unified memory architecture with 96 GB addressable
Strong integrated GPU with 40 compute units
Efficient power management and thermal characteristics

Software Limitations:

Critical gfx1151 architecture support gaps in ROCm ecosystem
PyTorch integration completely broken for GPU acceleration
Limited AI framework compatibility across the board
Insufficient documentation for troubleshooting

Market Position:

Competitive hardware specifications
Unique integrated architecture advantages
Significant software ecosystem disadvantages versus NVIDIA
Uncertain timeline for compatibility improvements

Purchasing Recommendations

Buy If: - Primary use case is general computing or traditional HPC workloads - Willing to wait 6-12 months for AI software ecosystem maturity - Value open-source software development approach - Need power-efficient integrated solution

Avoid If:

Immediate AI/ML development requirements
Production AI inference deployments planned
Time-critical project timelines
Require guaranteed software support

Final Verdict

The AMD AI Max+ 395 system shows tremendous promise as a unified computing platform, but premature software ecosystem development makes it unsuitable for current AI workloads. Organizations should monitor ROCm development progress closely, as this hardware could become highly competitive once software support matures. For general computing applications, the system offers excellent performance and value, representing AMD's continued progress in processor design and integration.

The AMD AI Max+ 395 represents a glimpse into the future of integrated computing platforms, but early adopters should be prepared for software ecosystem growing pains. As AMD continues investing in ROCm development and the open-source community contributes solutions, this platform has the potential to become a compelling alternative to NVIDIA's ecosystem dominance.

Accelerating Large-Scale Ballistic Simulations with torchdiffeq and PyTorch

A.C. Jokela — Sat, 10 May 2025 20:44:52 GMT

Introduction

Simulating the motion of projectiles is a classic problem in physics and engineering, with applications ranging from ballistics and aerospace to sports analytics and educational demonstrations. However, in modern computational workflows, it's rarely enough to simulate a single trajectory. Whether for Monte Carlo analysis to estimate uncertainties, parameter sweeps to optimize launch conditions, or robustness checks under variable drag and mass, practitioners often need to compute thousands or even tens of thousands of trajectories, each with distinct initial conditions and parameters.

Solving ordinary differential equations (ODEs) governing these trajectories becomes a computational bottleneck in such “large batch” scenarios. Traditional scientific Python tools like scipy.integrate.solve_ivp are excellent for solving ODEs in serial, one scenario at a time, making them ideal for interactive exploration or detailed studies of individual systems. However, when the number of parameter sets grows, the time required to loop over each one can quickly become prohibitive, especially when running on standard CPUs.

Recent advances in scientific machine learning and GPU computing have opened new possibilities for accelerating these kinds of simulations. The torchdiffeq library extends PyTorch’s ecosystem with differentiable ODE solvers, supporting batch-mode integration and seamless hardware acceleration via CUDA GPUs. By leveraging vectorized operations and batched computation, torchdiffeq makes it possible to simulate thousands of parameterized systems orders of magnitude faster than traditional approaches.

This article empirically compares scipy.solve_ivp and torchdiffeq on a realistic, parameterized ballistic projectile problem. We'll see how modern, batch-oriented tools unlock dramatic speedups, making large-scale simulation, optimization, and uncertainty quantification far more practical and scalable.

The Ballistics Problem: ODEs and Parameters

At the heart of projectile motion lies a classic set of equations: the Newtonian laws of motion under the influence of gravity. In real-world scenarios (be it sports, military science, or atmospheric research), it's crucial to account not just for gravity but also for aerodynamic drag, which resists motion and varies with both the speed and shape of the object. For fast-moving projectiles like baseballs, artillery shells, or drones, drag is well-approximated as quadratic in velocity.

The trajectory of a projectile under both gravity and quadratic drag is described by the following system of ODEs:

$ \frac{d\mathbf{r}}{dt} = \mathbf{v} $

$ \frac{d\mathbf{v}}{dt} = -g \hat{z} - \frac{k}{m} |\mathbf{v}| \mathbf{v} $

Here, $\mathbf{r}$ is the position vector, $\mathbf{v}$ is the velocity vector, $g$ is the gravitational acceleration (9.81 m/s², directed downward), $m$ is the projectile's mass, and $k$ is the drag coefficient, a parameter incorporating air density, projectile shape, and cross-sectional area. The term $-\frac{k}{m} |\mathbf{v}| \mathbf{v}$ captures the quadratic (speed-squared) air resistance opposing motion.

This model supports a range of relevant parameters:

Initial speed ($v_0$): How fast the projectile is launched.
Launch angle ($\theta$): The elevation above the horizontal.
Azimuth ($\phi$): The compass direction of the launch in the x-y plane.
Drag coefficient ($k$): Varies by projectile type and environment (e.g., bullets, baseballs, or debris).
Mass ($m$): Generally constant for a given projectile, but can vary in sensitivity analyses.

By randomly sampling these parameters, we can simulate broad families of real-world projectile trajectories, quantifying variations due to weather, launch conditions, or design tolerances. This approach is vital in engineering (for safety margins and optimization), defense (for targeting uncertainty), and physics education (visualizing parameter effects). With these governing principles defined, we’re equipped to systematically simulate and analyze thousands of projectile scenarios.

Vectorized Batch Simulation: Why It Matters

In classical physics instruction or simple engineering analyses, simulating a single projectile (perhaps varying its launch angle or speed by hand) was once sufficient to gain insight into trajectory behavior. But the demands of modern computational science and industry go far beyond this. Today, engineers, data scientists, and researchers routinely confront tasks like uncertainty quantification, statistical analysis, design optimization, or machine learning, all of which require running the same model across thousands or even millions of parameter combinations. For projectile motion, that might mean sampling hundreds of drag coefficients, launch angles, and initial velocities to estimate failure probabilities, optimize for maximum range under real-world disturbances, or quantify the uncertainty in a targeting system.

Attempting to tackle these large-scale parameter sweeps with traditional serial Python code quickly exposes severe performance limitations. Standard Python scripts iterate through scenarios using simple loops, solving the ODE for one set of inputs, then moving to the next. While such code is easy to write and understand, it suffers from significant overhead: each call to an ODE solver like scipy.solve_ivp carries the cost of repeatedly allocating memory, reinterpreting Python functions, and performing calculations on a single set of parameters without leveraging efficiencies of scale.

Moreover, CPUs themselves have limited capacity for parallel execution. Although some scientific computing libraries exploit multicore CPUs for modest speedups, true high-throughput workloads outstrip what a desktop processor can provide. This is where vectorization and hardware acceleration revolutionize scientific computing. By formulating simulations so that many parameter sets are processed in tandem, vectorized code can amortize memory access and computation over entire batches.

This paradigm is taken even further with the introduction of modern hardware accelerators, particularly Graphics Processing Units (GPUs). GPUs are designed for massive parallel processing, capable of performing thousands of operations simultaneously. Frameworks like PyTorch make it straightforward to move simulation data to the GPU and exploit this parallelism using batch operations and tensor arithmetic. Libraries such as torchdiffeq, built on PyTorch, allow entire ensembles of ODE initial conditions and parameters to be integrated at once, often achieving one or even two orders of magnitude speedup over standard serial approaches.

By harnessing vectorized and accelerated computation, we shift from thinking about trajectories one at a time to simulating entire probability distributions of outcomes, enabling robust analysis and real-time feedback that serial methods simply cannot deliver.

Setting Up the Experiment

To rigorously compare batch ODE solvers in a realistic context, we construct an experiment that simulates a large family of projectiles, each with unique initial conditions and drag parameters. Here, we demonstrate how to generate the complete dataset for such an experiment, scaling easily to $N=10,000$ scenarios or more.

First, we select which parameters to randomize:

Initial speed ($v_0$): uniformly sampled between 100 and 140 m/s.
Launch angle ($\theta$): uniformly distributed between 20° and 70° (converted to radians).
Azimuth ($\phi$): uniformly distributed from 0 to $2\pi$, representing all compass directions.
Drag coefficient ($k$): uniformly sampled between 0.03 and 0.07; these bounds reflect different projectile shapes or environmental conditions.
Mass ($m$): held constant at 1.0 kg for simplicity.

The initial position for each projectile is set at $(x, y, z) = (0, 0, 1)$, representing launches from a height of 1 meter above ground.

Here is the core code to generate these parameters and construct the state vectors:

N = 10000  # Number of projectiles
np.random.seed(42)
r0 = np.zeros((N, 3))
r0[:, 2] = 1  # start at z=1m

speeds = np.random.uniform(100, 140, size=N)
angles = np.random.uniform(np.radians(20), np.radians(70), size=N)
azimuths = np.random.uniform(0, 2*np.pi, size=N)
k = np.random.uniform(0.03, 0.07, size=N)
m = 1.0
g = 9.81

# Compute velocity components from speed, angle, and azimuth
v0 = np.zeros((N, 3))
v0[:, 0] = speeds * np.cos(angles) * np.cos(azimuths)
v0[:, 1] = speeds * np.cos(angles) * np.sin(azimuths)
v0[:, 2] = speeds * np.sin(angles)

# Combine into state vector: [x, y, z, vx, vy, vz]
y0 = np.hstack([r0, v0])

With this setup, each row of y0 fully defines the position and velocity of one simulated projectile, and associated arrays (k, m, etc.) capture the unique drag and physical parameters. This approach ensures our batch simulations cover a broad, realistic spread of possible projectile behaviors.

Serial Approach: scipy.solve_ivp

The scipy.integrate.solve_ivp function is a standard tool in scientific Python for numerically solving initial value problems for ordinary differential equations (ODEs). Designed for flexibility and usability, it allows users to specify the right-hand side function, initial conditions, time span, and integration tolerances. It's ideal for scenarios where you need to inspect or visualize a single trajectory in detail, perform stepwise integration, or analyze systems with events (such as ground impact in our ballistics context).

However, solve_ivp is fundamentally serial in nature: each call integrates one ODE system, with one set of inputs and parameters. To simulate a batch of projectiles with varying initial conditions and drag parameters, a typical approach is to loop over all $N$ cases, calling solve_ivp anew each time. This approach is straightforward, but comes with key drawbacks: overhead from repeated Python function calls, redundant setup within each call, and no built-in way to leverage vectorization or parallel computation on CPUs or GPUs.

Here’s how the serial batch simulation is performed for our random projectiles:

from scipy.integrate import solve_ivp

def ballistic_ivp_factory(ki):
    def fn(t, y):
        vel = y[3:]
        speed = np.linalg.norm(vel)
        acc = np.zeros_like(vel)
        acc[2] = -g
        acc -= (ki/m) * speed * vel
        return np.concatenate([vel, acc])
    return fn

def hit_ground_event(t, y):
    return y[2]
hit_ground_event.terminal = True
hit_ground_event.direction = -1

t_eval = np.linspace(0, 15, 400)

trajectories = []
for i in range(N):
    sol = solve_ivp(
        ballistic_ivp_factory(k[i]), (0, 15), y0[i],
        t_eval=t_eval, rtol=1e-5, atol=1e-7, events=hit_ground_event)
    trajectories.append(sol.y)

To extract and plot the $i$-th projectile’s trajectory (for example, $x$ vs. $z$):

x = trajectories[i][0]
z = trajectories[i][2]
plt.plot(x, z)

While this method is robust and works for small $N$, it scales poorly for large batches. Each ODE integration runs one after the other, keeping all computation on the CPU, and does not exploit the potential speedup from modern hardware or batch processing. For workflows involving thousands of projectiles, these limitations quickly become significant.

Batched & Accelerated: torchdiffeq and PyTorch

Recent advances in machine learning frameworks have revolutionized scientific computing, and PyTorch is at the forefront. While best known for deep learning, PyTorch offers powerful tools for general numerical tasks, including automatic differentiation, GPU acceleration, and, critically for large-scale simulations, native support for batched and vectorized computation. Building on this, the torchdiffeq library brings state-of-the-art ODE solvers to the PyTorch ecosystem. This unlocks not only scalable and differentiable simulations, but also unprecedented throughput for large parameter sweeps thanks to efficient batching.

Unlike scipy.solve_ivp, which solves one ODE system per call, torchdiffeq.odeint can handle entire batches simultaneously. If you stack $N$ initial conditions into a tensor of shape $(N, D)$ (with $D$ being the state dimension, e.g., position and velocity components), and you write your ODE’s right-hand-side function to process these $N$ states in parallel, odeint will integrate all of them in one go. This batched approach is highly efficient, especially when offloading the computation to a CUDA-enabled GPU, which can process thousands of simple ODE systems at once.

A custom ODE function in PyTorch for batched ballistics looks like this:

import torch
from torchdiffeq import odeint

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class BallisticsODEBatch(torch.nn.Module):
    def __init__(self, k, m, g):
        super().__init__()
        self.k = torch.tensor(k, device=device).view(-1,1)
        self.m = m
        self.g = g
    def forward(self, t, y):
        vel = y[:, 3:]
        speed = torch.norm(vel, dim=1, keepdim=True)
        acc = torch.zeros_like(vel)
        acc[:, 2] -= self.g
        acc -= (self.k / self.m) * speed * vel
        return torch.cat([vel, acc], dim=1)

After preparing the initial states (y0_torch, shape $(N, 6)$), you launch the batch integration with:

odefunc = BallisticsODEBatch(k, m, g).to(device)
y0_torch = torch.tensor(y0, dtype=torch.float32, device=device)
t_torch = torch.linspace(0, 15, 400).to(device)

sol_batch = odeint(odefunc, y0_torch, t_torch, rtol=1e-5, atol=1e-7)  # (T, N, 6)

By processing every $N$ parameter set in a single tensor operation, batching reduces memory and Python overhead substantially compared to looping with solve_ivp. When running on a GPU, these speedups are often dramatic (sometimes orders of magnitude) due to massive parallelism and reduced per-call Python latency. For researchers and engineers running uncertainty analyses or global optimizations, batched ODE integration with torchdiffeq makes large-scale simulation not only practical, but fast.

Cropping and Plotting Trajectories

When visualizing or comparing projectile trajectories, it's important to stop each curve exactly when the projectile reaches ground level ($z = 0$). Without this cropping, some trajectories would artificially continue below ground due to numerical integration, making visualizations misleading and length-biased. To ensure all plots fairly represent real-world impact, we truncate each trajectory at its ground crossing, interpolating between the last above-ground and first below-ground points to find the precise impact location.

The following function performs this interpolation:

def crop_trajectory(x, z, t):
    idx = np.where(z <= 0)[0]
    if len(idx) == 0:
        return x, z
    i = idx[0]
    if i == 0:
        return x[:1], z[:1]
    frac = -z[i-1] / (z[i] - z[i-1])
    x_crop = x[i-1] + frac * (x[i] - x[i-1])
    return np.concatenate([x[:i], [x_crop]]), np.concatenate([z[:i], [0.0]])

Using this, we can generate “spaghetti plots” for both solvers, showcasing dozens or hundreds of realistic, ground-terminated trajectories for direct comparison.
Example:

for i in range(100):
    x_t, z_t = crop_trajectory(sol_batch_np[:,i,0], sol_batch_np[:,i,2], t_np)
    plt.plot(x_t, z_t, color='tab:blue', alpha=0.2)

Performance Benchmarking: Timing the Solvers

To quantitatively compare the efficiency of scipy.solve_ivp against the batched, accelerator-aware torchdiffeq, we systematically measured simulation runtimes across a range of batch sizes ($N$): 100, 1,000, 5,000, and 10,000. We timed both solvers under identical conditions, measuring total wall-clock time and deriving the average simulation throughput (trajectories per second).

All experiments were run on a workstation equipped with an Intel i7 CPU and NVIDIA Pascal GPUs), with PyTorch configured for CUDA acceleration. The same ODE system and tolerance settings ($\text{rtol}=1\text{e-5}$, $\text{atol}=1\text{e-7}$) were used for both solvers.

The script below shows the core timing procedure:

import numpy as np
import torch
from torchdiffeq import odeint
from scipy.integrate import solve_ivp
import time
import matplotlib.pyplot as plt

# For reproducibility
np.random.seed(42)

# Physics constants
g = 9.81
m = 1.0

def generate_initial_conditions(N):
    r0 = np.zeros((N, 3))
    r0[:, 2] = 1  # z=1m
    speeds = np.random.uniform(100, 140, size=N)
    angles = np.random.uniform(np.radians(20), np.radians(70), size=N)
    azimuths = np.random.uniform(0, 2*np.pi, size=N)
    v0 = np.zeros((N, 3))
    v0[:, 0] = speeds * np.cos(angles) * np.cos(azimuths)
    v0[:, 1] = speeds * np.cos(angles) * np.sin(azimuths)
    v0[:, 2] = speeds * np.sin(angles)
    k = np.random.uniform(0.03, 0.07, size=N)
    y0 = np.hstack([r0, v0])
    return y0, k

def ballistic_ivp_factory(ki):
    def fn(t, y):
        vel = y[3:]
        speed = np.linalg.norm(vel)
        acc = np.zeros_like(vel)
        acc[2] = -g
        acc -= (ki/m) * speed * vel
        return np.concatenate([vel, acc])
    return fn

def hit_ground_event(t, y):
    return y[2]
hit_ground_event.terminal = True
hit_ground_event.direction = -1

class BallisticsODEBatch(torch.nn.Module):
    def __init__(self, k, m, g, device):
        super().__init__()
        self.k = torch.tensor(k, device=device).view(-1, 1)
        self.m = m
        self.g = g
    def forward(self, t, y):
        vel = y[:,3:]
        speed = torch.norm(vel, dim=1, keepdim=True)
        acc = torch.zeros_like(vel)
        acc[:,2] -= self.g
        acc -= (self.k/self.m) * speed * vel
        return torch.cat([vel, acc], dim=1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"PyTorch device: {device}")

N_list = [100, 1000, 5000, 10000]
t_points = 400
t_eval = np.linspace(0, 15, t_points)
t_torch = torch.linspace(0, 15, t_points)

timings = {'solve_ivp':[], 'torchdiffeq':[]}

for N in N_list:
    print(f"\n=== Benchmarking N = {N} ===")
    y0, k = generate_initial_conditions(N)

    # --- torchdiffeq batched solution
    odefunc = BallisticsODEBatch(k, m, g, device=device).to(device)
    y0_torch = torch.tensor(y0, dtype=torch.float32, device=device)
    t_torch_dev = t_torch.to(device)
    torch.cuda.synchronize() if device.type == "cuda" else None
    start = time.perf_counter()
    sol = odeint(odefunc, y0_torch, t_torch_dev, rtol=1e-5, atol=1e-7)  # shape (T,N,6)
    torch.cuda.synchronize() if device.type == "cuda" else None
    time_torch = time.perf_counter() - start
    print(f"torchdiffeq (batch): {time_torch:.2f}s")
    timings['torchdiffeq'].append(time_torch)

    # --- solve_ivp serial solution
    start = time.perf_counter()
    for i in range(N):
        solve_ivp(
            ballistic_ivp_factory(k[i]),
            (0, 15),
            y0[i],
            t_eval=t_eval,
            rtol=1e-5, atol=1e-7,
            events=hit_ground_event
        )
    time_ivp = time.perf_counter() - start
    print(f"solve_ivp (serial):  {time_ivp:.2f}s")
    timings['solve_ivp'].append(time_ivp)

# ---- Plot results
plt.figure(figsize=(8,5))
plt.plot(N_list, timings['solve_ivp'], label='solve_ivp (serial, CPU)', marker='o')
plt.plot(N_list, timings['torchdiffeq'], label=f'torchdiffeq (batch, {device.type})', marker='s')
plt.yscale('log')
plt.xscale('log')
plt.xlabel('Batch Size N')
plt.ylabel('Total Simulation Time (seconds, log scale)')
plt.title('ODE Solver Performance: solve_ivp vs torchdiffeq')
plt.grid(True, which='both', ls='--')
plt.legend()
plt.tight_layout()
plt.show()

Benchmark Results

PyTorch device: cuda

=== Benchmarking N = 100 ===
torchdiffeq (batch): 0.35s
solve_ivp (serial):  0.60s

=== Benchmarking N = 1000 ===
torchdiffeq (batch): 0.29s
solve_ivp (serial):  5.84s

=== Benchmarking N = 5000 ===
torchdiffeq (batch): 0.31s
solve_ivp (serial):  29.84s

=== Benchmarking N = 10000 ===
torchdiffeq (batch): 0.31s
solve_ivp (serial):  59.74s

As shown in the table and the bar chart below, torchdiffeq achieves orders of magnitude speedup, especially when run on GPU. While solve_ivp's wall time scales linearly with batch size, torchdiffeq’s increase is much more gradual due to highly efficient batch parallelism on both CPU and GPU.

Visualization

These results decisively demonstrate the advantage of batched, hardware-accelerated ODE integration for large-scale uncertainty quantification and parametric studies. For modern simulation workloads, torchdiffeq turns otherwise intractable analyses into routine computations.

Practical Insights & Limitations

The dramatic performance advantage of torchdiffeq for large-batch ODE integration is a game-changer for certain classes of scientific and engineering simulations. However, like any advanced computational tool, its real-world utility depends on the problem context, user preferences, and technical constraints.

When torchdiffeq Shines

Large Batch Sizes: The most compelling case for torchdiffeq is when you need to simulate many similar ODE systems in parallel. If your workflow naturally involves analyzing thousands of parameter sets (such as in Monte Carlo uncertainty quantification, global sensitivity analysis, optimization sweeps, or high-volume forward simulations), torchdiffeq can turn days of computation into minutes, especially when exploiting a modern GPU.
Homogeneous ODE Forms: torchdiffeq excels when the differential equations are structurally identical across all batch members (e.g., all projectiles differ only in launch parameters, mass, or drag, not in governing equations). This allows vectorized tensor operations and maximizes parallel hardware utilization.
GPU Acceleration: If you have access to CUDA hardware, the batch approach provided by PyTorch integrates seamlessly. For highly parallelizable problems, the speedup can be more than an order of magnitude compared to CPU execution alone.

Where scipy’s solve_ivp Is Preferable

Single or Few Simulations: If your workload only involves single or a handful of trajectories (or you need results interactively), scipy.solve_ivp is still highly convenient. It’s light on dependencies, simple to use, and well-integrated with the broader SciPy ecosystem.
Out-of-the-box Event Handling: solve_ivp integrates event location cleanly, making it straightforward to stop integration at complex conditions (like ground impact, threshold crossings, or domain boundaries) with minimal setup.
No PyTorch/Deep Learning Stack Needed: For users not otherwise relying on PyTorch, keeping everything in NumPy/SciPy can mean a lighter, more transparent setup and easier integration into classic scientific workflows.

Accuracy and Tolerances

Both torchdiffeq and solve_ivp allow setting relative and absolute tolerances for error control. In most practical applications, both provide comparable accuracy if configured similarly, though always test with your specific ODEs and parameters, as subtle differences can arise in stiff or highly nonlinear regimes.

Limitations of torchdiffeq

Complex Events and Custom Solvers: While torchdiffeq supports batching and GPU execution, its event handling isn’t as automatic or flexible as in solve_ivp. If you need advanced stopping criteria, adaptive step event targeting, or integration using custom/obscure methods, PyTorch-based solvers may require more custom code or workarounds.
Smaller Scientific Ecosystem: While PyTorch is hugely popular in machine learning, the larger SciPy ecosystem offers more “out-of-the-box” scientific routines and examples. Some users may need to roll their own utilities in PyTorch.
Learning Curve/Code Complexity: Writing vectorized, batched ODE functions (especially for newcomers to PyTorch or GPU programming) can pose an initial hurdle. For seasoned scientists accustomed to “for-loop” logic, adapting to a tensor-based, batch-first paradigm may require unlearning older habits.

Maintainability

For codebases built on PyTorch or targeted at high-throughput, the benefits are worth the upfront learning cost. For one-off or small-scale science projects, the classic SciPy stack may remain more maintainable and accessible for most users. Ultimately, the choice depends on the problem scale, user expertise, and requirements for future extensibility and hardware performance.

Conclusions

This benchmark study highlights the substantial performance gains attainable by leveraging torchdiffeq and PyTorch for batched ODE integration in Python. While scipy.solve_ivp remains robust and user-friendly for single or low-volume simulations, it quickly becomes a bottleneck when working with thousands of parameter variations common in uncertainty quantification, optimization, or high-throughput design. By contrast, torchdiffeq (especially when combined with GPU acceleration) enables orders-of-magnitude faster simulations thanks to its inherent support for vectorized batching and parallel computation.

Such speedups are transformative for both research and industry. Rapid batch simulations make Monte Carlo analyses, parametric studies, and iterative design far more feasible, allowing deeper exploration and faster time-to-insight across fields from engineering to quantitative science. For machine learning scientists, batched ODE integration can even be incorporated into differentiable pipelines for neural ODEs or model-based reinforcement learning.

If you face large-scale ODE workloads, we strongly encourage experimenting with the supplied example code and adapting torchdiffeq to your own applications. Additional documentation, tutorials, and PyTorch resources are available at the torchdiffeq repository and PyTorch documentation. Embracing modern computational tools can unlock dramatic gains in productivity, capability, and discovery.

Appendix: Code Listing

TorchDiffEq contains an HTML rendering of the complete code listing for this article, including all imports, functions, and plotting routines. For the actual Jupyter notebook, see torchdiffeq.ipynb. You can run it directly in a Jupyter notebook or adapt it to your own projects.