Feature Request: RTX 5090 Support with ms-swift docker image with CUDA 12.8

[Frank-Y-Zheng]

RTX 5090 support is needed for users with NVIDIA's latest Blackwell architecture GPU. The RTX 5090 uses CUDA capability sm_120, which requires CUDA 12.8 or higher for GitHubSuper User proper support. Current software installations show compatibility errors like "NVIDIA GeForce RTX 5090 with CUDA capability sm_120 is not compatible with the current PyTorch installation" TensorFlow 25.01 + CUDA 12.8 + RTX 5090 on WSL2: "CUDA failed to initialize" (Error 500) Issue - TensorFlow - Google AI Developers Forum +2.

[dconsorte]

See if this helps. I used nightly builds: https://github.com/dconsorte/pytorch-tensorflow-gpu

[Frank-Y-Zheng]

Thank you very much to create a dual-framework comptable image.

Yet I have used the offical torch image to build the ms-swift i.e. pytorch/pytorch:2.7.1-cuda12.8-cudnn9-devel.

Docker Setup for MS-Swift with RTX 5090 Support

Overview

This guide provides instructions for setting up a dual-framework compatible Docker image for MS-Swift using the official PyTorch base image.

Base Image

I have successfully used the official PyTorch Docker image to build MS-Swift:

docker pull pytorch/pytorch:2.7.1-cuda12.8-cudnn9-devel

For runtime environments, you can also use:

docker pull pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime

Setup Instructions

For users who want to utilize MS-Swift with RTX 5090 support:

1. Start with the PyTorch base image mentioned above
2. Clone and compile vLLM from the official Git repository
3. Install MS-Swift following the standard installation process

This configuration has been tested and confirmed to work with RTX 5090 GPUs.

Requirements

• NVIDIA RTX 50xx GPU
• CUDA 12.8
• cuDNN 9
• Docker with NVIDIA container runtime

Notes

This setup ensures compatibility between the PyTorch framework, vLLM, and MS-Swift while validate on RTX 5090 hardware.