PyTorchで`fatal : Unsupported .version 8.6; current version is '8.5'ptxas fatal `が発生したときの対処法

はじめに

PyTorchで機械学習の勉強をしようと思い、コードを書いていたら次のようなエラーが出るようになりました。
実際のエラーは非常に長いので、最後の部分だけキャプチャしています。

おそらく将来のアップデートで修正されると思いますが、2024/11/12現在においてはまだ発生するので、備忘録として残しておきます。
なお、今回の修正方法はPythonライブラリのコードを直接修正するので、あくまでも暫定的な措置と思ってください。

発生環境

自分の環境は次のとおりです。
機械学習の勉強を始めたばかりなので、比較的バージョンは新しい方だと思いますが、だからこそ発生しています。

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Thu_Sep_12_02:18:05_PDT_2024
Cuda compilation tools, release 12.6, V12.6.77
Build cuda_12.6.r12.6/compiler.34841621_0

＄python --version
Python 3.12.7

$ python
Python 3.12.7 (main, Oct 30 2024, 00:39:08) [GCC 13.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.__version__)
2.5.1+cu124

• CUDA 12.6
• Python 12.6
• PyTorch 2.5.1( CUDA12.4)

発生コード

細かい部分は省略しますが、次のようにモデルを作成しています(初心者用の本にあるコードをほぼそのまま使っています)。PyTorchのcompile機能を使い、学習を始めると発生します。なお、エラーは発生しますが、学習自体はちゃんと終わります(compieが正常に行われているのかどうかは不明)。

model.py

class IrisModel(nn.Module):
    def __init__(self) -> None:
        super(IrisModel, self).__init__()

        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)
        self.relu = torch.nn.ReLU(inplace=False)

    def forward(self, x):
        x = F.relu(self.linear1(x))
        x = self.linear2(x)
        return x

main.py

tmp_model = IrisModel()
model = torch.compile(tmp_model.to(self.device))

#下記コードだと発生しない
#model = tmp_model.to(self.device)

修正方法

根本的な原因はtritonというライブラリのようです。
ちなみに、自分はこのライブラリがなんのライブラリなのかもわからない程度の初心者です。
なので、自分ではどうしようもなかったので、ネットで検索していたら次のサイトに原因が書いてありました。
https://github.com/triton-lang/triton/issues/4576

見ての通り、tritonのgithubリポジトリのイシューとして上がっているので、近い将来反映されると思います。
なので、以下のようにtritonをアップデートすれば、それだけで直るかもしれません(3.1.0では発生しました)。

pip install -U triton

ただ、残念ながら直らなかったので、仕方ないのでtritonを直接修正します。
まず、tritonの場所を確認します。

$ pip show triton
Name: triton
Version: 3.1.0
Summary: A language and compiler for custom Deep Learning operations
Home-page: https://github.com/triton-lang/triton/
Author: Philippe Tillet
Author-email: [e-mail]
License: 
Location: /home/sano/.asdf/installs/python/3.12.7/lib/python3.12/site-packages
Requires: filelock
Required-by: torch

必要なのは、Location部分です。これに、triton/backends/nvidia/compiler.pyを加えたパスを開きます。下記はあくまでも自分の環境なので、適宜変更してください。

$ vim ~/.asdf/installs/python/3.12.7/lib/python3.12/site-packages/triton/backends/nvidia/compiler.py 

次の関数が修正しなければならない関数です。

compiler.py

@functools.lru_cache()
def ptx_get_version(cuda_version) -> int:
    '''
    Get the highest PTX version supported by the current CUDA driver.
    '''
    assert isinstance(cuda_version, str)
    major, minor = map(int, cuda_version.split('.'))
    if major == 12:
        return 80 + minor
    if major == 11:
        return 70 + minor
    if major == 10:
        return 63 + minor
    raise RuntimeError("Triton only support CUDA 10.0 or higher")

次のようにコードを追加します(イシューに書いてある修正コード通りです)

compiler.py

@functools.lru_cache()
def ptx_get_version(cuda_version) -> int:
    '''
    Get the highest PTX version supported by the current CUDA driver.
    '''
    assert isinstance(cuda_version, str)
    major, minor = map(int, cuda_version.split('.'))
    if major == 12:
+       if minor >= 5:
+           return 85
+       else:
+           return 80 + minor
-       return 80 + minor
    if major == 11:
        return 70 + minor
    if major == 10:
        return 63 + minor
    raise RuntimeError("Triton only support CUDA 10.0 or higher")

これでエラーはでなくなるはずです。
ライブラリを直接修正するのはあまりよくないと思うので、あくまでも暫定的な措置だと思ってください。

なぜ発生していたのか？

そもそもの原因として、CUDAのバージョンを12.6であることが発生原因のようです。
コードを見ると、cudaのバージョンのメジャーとマイナーを分割し、それに応じてなんらかのバージョンを返しています。
もともとのコードはCUDAのマイナーバージョンと一致していたようですが、どうもそれが不一致になったことで発生したようです。