Run Polyaxon workloads on AMD GPUs

The AMD device plugin makes GPUs schedulable as `amd.com/gpu`; ROCm and the workload image determine whether an allocated device is actually usable. Polyaxon keeps the hardware request, placement, image, and run evidence together.

Outcome

An AMD-specific readiness component that verifies device allocation and PyTorch ROCm access before higher-level workloads are enabled.

Start the guide Scaling documentation

Infrastructure responsibilities

Kubernetes discovers

Drivers and a vendor device integration expose healthy accelerator capacity to the scheduler.

Polyaxon places

A component or preset requests the custom resource and targets the compatible node pool and image.

The workload verifies

A smoke test proves device access before training or serving is promoted to that hardware pool.

Prerequisites

ROCm-capable AMD GPUs with supported host drivers and firmware.
The AMD GPU device plugin exposing `amd.com/gpu` on the Kubernetes nodes.
A tested ROCm framework image and an AMD-specific Polyaxon queue or preset.

Validate the infrastructure

Request an AMD GPU

Target the AMD pool and request the resource name published by the official Kubernetes device plugin.

version: 1.1kind: componentname: amd-gpu-readinessrun:  kind: job  environment:    nodeSelector:      workload.polyaxon.com/accelerator: amd  container:    image: rocm/pytorch:rocm7.14.1_ubuntu24.04_py3.12_pytorch_release_2.12.0    command: ["python", "-c"]    args: ["import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name())"]    resources:      limits:        amd.com/gpu: "1"

ROCm exposes accelerators through PyTorch's `torch.cuda` API for compatibility; that name does not mean the workload is using CUDA.

2
Submit the operation
Run the component through the target Polyaxon project, queue, preset, and approval path.
polyaxon run -f amd-readiness.yaml
3
Verify the ROCm execution path
Confirm the pod lands on the AMD pool, receives one device, loads the expected framework build, and emits logs and artifacts normally.
polyaxon ops dashboard

Production checklist

Pin a compatible driver, device integration, container runtime, framework image, and firmware combination.

Publish accelerator selection through centrally managed Polyaxon presets instead of user-specific labels.

Track allocatable devices, hardware health, utilization, memory, thermals, and workload failures.

Validate one device, one node, and multi-node communication separately before production adoption.

Run RCCL tests before enabling multi-node AMD workloads and keep CUDA and ROCm images in separate presets.

Troubleshooting

The resource does not appear

Inspect the node driver, firmware, device integration, kubelet registration, node labels, and allocatable resources.

The container cannot access the device

Verify the requested custom resource, runtime class, permissions, device mounts, and image compatibility.

Performance is unexpectedly low

Check topology, NUMA placement, clocks, memory pressure, host-device transfer, communication libraries, and framework build.

Sources

Official platform, library, model, and Polyaxon references used by this guide.

Continue

Compare AI platforms Explore GPU orchestration