ML4GW / ML4GW/aframe

Add support for AMD GPUs (ROCm backend) or non-NVIDIA hardware

Open
#443 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
18
Forks
28
Avg merge
2d 2h
Merged PRs (30d)
7

Description

This came up while attempting to run training jobs on MI100 node on Delta, where CUDA-based containers do not work. Currently, we assume NVIDIA environment such as use of --nv in containers, CUDA-specific dependencies here. This limits the use of the AMD GPUs (MI100) using --rocm.

Proposed changes

  • Add support for ROCm-enabled PyTorch builds
  • Update container workflow to support --rocm (Apptainer)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with projects/train/pyproject.toml at the CUDA-specific dependencies, then trace the container workflow that currently uses Apptainer's --nv option. Check how PyTorch builds and container options are selected for training jobs. Done means ROCm-enabled PyTorch dependencies and --rocm container execution support are documented and usable on AMD MI100 hardware.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
devops, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.