2
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@segfo

CUDA用のプログラムをROCmに対応させてみた話

2
Last updated at Posted at 2026-03-07

はじめに

というわけで、今回のテーマはCUDA用の推論プログラムをROCmに対応させてみる話です。
ROCm7.2がWindowsでも動くようになった

EVO-X2はUnified Memoryが128GBありますね。
ROCmが対応してるし、じゃあこれを使わないともったいなくない？

動機

LLMと自分の好きな音楽について語るというプロジェクトをなんとなく作り始めた。
で、ここらってMIR（Music Information Retrieval）や音響解析周りのお話ですよね。
なので専門的な知識は無いんですよ。ですが、近年LLM先生によってこの知識のなさを補完してくれるようになったので始動させました。
MIRに関していろんな論文を読んでたらこのツールが出てきたので、今回は以下のツールを対象にします。

で、別にPytorchのCPU版でも現実的な速度で推論はできるんですが、せっかくROCmあるし、LLMもあるし改造したら速くなるんじゃね？　と思って改造してみた話です。

というわけで早速始めていく。

環境

  • OS: Windows 11
  • GPU: AMD (ROCm)
  • PyTorch: ROCm ビルド（torch.version.hip がTrue）
  • demucs モデル: htdemucs（Hybrid Transformer Demucs）
  • パッケージ管理: uv

チェッカースクリプト

ROCmが正しく動いてるかをまずは判定するスクリプトを書いてもらった。

"""
ROCm 7.1 + PyTorch GPU 検出確認スクリプト
"""

import torch

print("=" * 70)
print("🔍 PyTorch GPU 検出確認 (ROCm 7.1)")
print("=" * 70)

print(f"\n✅ 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)}")
    props = torch.cuda.get_device_properties(0)
    print(f"✅ Total Memory: {props.total_memory / (1024**3):.1f} GB")
    print(f"✅ Compute Capability: {props.major}.{props.minor}")
    
    print("\n✅ GPU テスト実行中...")
    x = torch.randn(1000, 1000, device="cuda")
    y = torch.randn(1000, 1000, device="cuda")
    z = torch.matmul(x, y)
    print(f"✅ GPU 演算成功 - 出力形状: {z.shape}")
    print(f"   デバイス: {z.device}")
    
    print("\n✅ GPU は完全に動作しています！")
else:
    print("\n❌ GPU が検出されていません")
    print("   CPU モードで実行します")

print("\n" + "=" * 70)

実行結果

✅ PyTorch Version: 2.9.1+rocmsdk20260116
✅ CUDA Available: True
✅ Device Count: 1
✅ Device Name: AMD Radeon(TM) 8060S Graphics
✅ Total Memory: 107.9 GB
✅ Compute Capability: 11.5

✅ GPU テスト実行中...
✅ GPU 演算成功 - 出力形状: torch.Size([1000, 1000])
デバイス: cuda:0

動いとるね。
じゃあ続き。

伝家の宝刀「AI先生あとよしなによろしく！」

楽曲のセクション分析についてGPUを使うように改良してもらうことにしました。
で、楽曲のセクション分析について大まかな流れを書きますね。

  1. 楽曲をドラム・ベース・その他・ボーカルの4つにデミキシング（分離）します
  2. デミキシングされた楽曲それぞれの要素の特徴量から、イントロ・Aメロ/Bメロ（Verse）・サビ・間奏・アウトロを推定します

というものです。

よろしくした結果がこちら

image.png

image.png

正常に分析できている楽曲の場合

image.png

image.png

この分析ですが

よろしくした結果はですね、デミキシングした結果、無音のファイルが出来上がっています。
正常に分析できている楽曲は、正しく波形が分離されていて、問題なくセクション分類できていました。

うーん、これは何なんでしょうね？なんNaNでしょうか・・・ちょっとよくわからないです。
というわけでこっちもAI先生によろしくしてみることにしました。

ここからはAIさんが書いてくれたお話です。しばし俺自身は成りを潜めますのでお楽しみください。

問題の概要

ROCm GPU での実行時に、パイプラインの 2 か所 で独立した NaN 生成バグが存在した。

Phase 発生箇所 症状
Phase 1 HTDemucs（demix）の create_sin_embedding demix wav が全サンプル -1.0 になる
Phase 2 AllInOne モデル推論の nn.LayerNorm bpm=null, beats=[], segments が 0.01 秒刻みで爆発

どちらも ROCm GPU 上の大規模テンソル演算でNaNが生成されるという同一の根本原因を持つ。

Phase 1: HTDemucs demix の NaN

問題

GPU（ROCm）で demix を実行すると、出力 wav ファイルの全サンプルが -1.0 になり、音声として無効な状態になっていた。

# Phase 1 修正前の各ステム確認結果
bass    : peak=1.000000  rms=1.000000  unique_vals=1
drums   : peak=1.000000  rms=1.000000  unique_vals=1
other   : peak=1.000000  rms=1.000000  unique_vals=1
vocals  : peak=1.000000  rms=1.000000  unique_vals=1

unique_vals=1 はすべてのサンプルが同一値（-1.0）であることを示している。

NaN 伝播の仕組み:
モデル出力に NaN が混在 → demucs/audio.py の prevent_clip（rescale モード）内で NaN / max(1.01*NaN, 1) = NaN → soundfile.write で PCM_16 に変換する際 NaN が int16 の最小値にマップ → 正規化後 -1.0

調査

Step 1: CPU フォールバックの除去

src/allin1/demix.py にあった「ROCm では CPU で demix する」というフォールバックコードを削除し、GPU で動作するように変更した。

Step 2: Flash/Mem-Efficient SDP の無効化

nn.MultiheadAttention が ROCm で実験的な Flash SDP / Mem-Efficient SDP バックエンドを使用して NaN を生成する可能性があるため無効化した。

torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(False)

→ NaN は解消されなかった（nan=2751840 継続）。

Step 3: GroupNorm CPU パッチの適用

demucs/demucs.py の DConv レイヤー内の GroupNorm が ROCm で大規模 reduction 時に NaN を生成する仮説を立て、CPU フォールバックパッチを適用した。

→ NaN は解消されなかった。GroupNorm は demix NaN の直接原因ではなかった。

Step 4: MIOpen GEMM 無効化の試行

MIOPEN_DEBUG_CONV_GEMM=0 を設定してテスト。

→ NaN は解消されなかった。

Step 5: forward hook による段階別 NaN 追跡

debug_nan.py を作成し、各モジュールの出力を CPU に転送してから NaN をカウントするフックを設置した（ROCm では GPU 上の isnan() 自体が正しく動作しないため、CPU 転送が必須）。

  ok   tencoder[0]: nan=0  shape=(1, 48, 85995)
  ok   encoder[0]: nan=0  shape=(1, 48, 512, 336)
  ...（中略）
  ok   ct.norm_in: nan=0  shape=(1, 2688, 512)
  ok   ct.norm_in_t: nan=0  shape=(1, 1344, 512)
  ok   ct.layers[0]: nan=0  shape=(1, 2688, 512)        ← 周波数ブランチは正常
  NaN! ct.layers_t[0]: nan=688128  shape=(1, 1344, 512) ← ここで初めて NaN 発生
  NaN! ct.layers[1]: nan=1376256  ...
  （以降すべて NaN）

CrossTransformerEncoder の時間ブランチ（layers_t[0]）で最初に NaN が発生していることを特定。

Step 6: layers_t[0] の入力確認

pre-hook で layers_t[0] への入力テンソルを確認した結果：

  NaN! layers_t[0]_in[0]: nan=688126  shape=(1, 1344, 512)

layers_t[0] への入力自体に既に NaN があることが判明。NaN の発生源は CrossTransformerEncoder.forward 内の position embedding 加算処理だった。

Step 7: 根本原因の特定

CrossTransformerEncoder.forward（demucs/transformer.py）の該当箇所：

def forward(self, x, xt):
    # ... (周波数ブランチ x の処理)
    pos_emb = self._get_pos_embedding(T2, B, C, x.device)   # ← GPU で呼ばれる
    xt = self.norm_in_t(xt)
    xt = xt + self.weight_pos_embed * pos_emb               # ← NaN を含む pos_emb が加算される

_get_pos_embedding → create_sin_embedding(T, C, ..., device=x.device) の実装：

def create_sin_embedding(length, dim, shift=0, device="cpu", max_period=10000):
    adim = torch.arange(dim // 2, device=device).view(1, 1, -1)
    phase = pos / (max_period ** (adim / (half_dim - 1)))   # ← ROCm で NaN 発生
    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)

max_period ** (adim / (half_dim - 1)) という整数テンソルへのべき乗演算が ROCm GPU 上で NaN を生成することが根本原因。
周波数ブランチが正常だったのは create_2d_sin_embedding が最初から CPU 計算していたため。時間ブランチだけが device=x.device（GPU）で直接計算していた。

修正

src/allin1/demix.py に _patch_sin_embedding_for_rocm() を追加：

def _patch_sin_embedding_for_rocm():
    import demucs.transformer as _dt
    _orig = _dt.create_sin_embedding

    def _rocm_create_sin_embedding(length, dim, shift=0, device="cpu", max_period=10000):
        _dev = device if isinstance(device, torch.device) else torch.device(str(device))
        if _dev.type == 'cuda' and getattr(torch.version, 'hip', None) is not None:
            result_cpu = _orig(length, dim, shift=shift, device='cpu', max_period=max_period)
            return result_cpu.to(_dev)
        return _orig(length, dim, shift=shift, device=device, max_period=max_period)

    _dt.create_sin_embedding = _rocm_create_sin_embedding

結果

# Phase 1 修正後の各ステム確認結果
bass    : peak=0.789825  rms=0.173165  unique_vals=27
drums   : peak=0.990082  rms=0.165411  unique_vals=50
other   : peak=0.990082  rms=0.151473  unique_vals=50
vocals  : peak=0.990082  rms=0.167132  unique_vals=17
  • NaN count: 2,751,840 / 2,751,840 → 0 / 2,751,840 ✅
  • peak が 1.0 未満でクリッピングなし ✅
  • unique_vals が多数（多様なサンプル値） ✅

Phase 2: AllInOne モデル推論の NaN

問題

Phase 1 修正後も、特定の楽曲で分析結果が不正になる問題が残っていた。

// Phase 2 修正前の struct/01_01_インフェリア.json（破損例）
{
  "bpm": null,
  "beats": [],
  "downbeats": [],
  "segments": [
    { "start": 90.19, "end": 90.20, "label": "end" },
    { "start": 90.20, "end": 90.21, "label": "end" },
    ...  // 0.01秒刻みで2803件
  ]
}

加えて、beats=[] の状態で -s（sonify）オプションを使うとクラッシュが発生した。

症状と原因の対応:

症状 原因
bpm: null, beats: [] logits が NaN → DBN が何も検出できない
segments が 0.01 秒刻みで爆発 NaN → sigmoid → 全フレームが boundary 判定
クラッシュ sonify.py の _sonify_metronome で beats=[] に対し .min() を呼ぶと ValueError

調査

Step 1: spec キャッシュの NaN 確認

debug_nan.py に spec チェックブロックを追加して確認した結果、spec/01_01_インフェリア.npy 自体は正常だった。

[Phase 1] spec キャッシュ NaN チェック
  shape : (4, 21897, 81)
  NaN   : 0  Inf: 0
  min=0.0000  max=2.3392  mean=0.0020
  => OK: spec キャッシュは正常。問題はモデル推論側にあります

→ 問題はスペクトログラムではなく AllInOne モデルの推論にあると特定。

Step 2: AllInOne 推論の NaN 追跡（初回 — hook の欠陥発覚）

debug_nan2.py を AllInOne 用に改修し、spec/01_01_インフェリア.npy を実入力として推論を実行した。

当初、logit の NaN カウントを GPU 上で行っていたため（t != t）、ROCm の比較演算バグにより全フィールドが偽のガベージ値を返した。また forward hook が AllInOneOutput（dataclass）を tuple/list として扱えず、hook が素通りしていた。

ROCm 特有の注意点（今回の調査で判明）:

  • GPU 上の torch.isnan() や != 比較は ROCm で正しく動作しない
  • すべての NaN チェックは CPU に転送（.detach().cpu()）してから行う必要がある
  • dataclass 型の出力は isinstance(out, (tuple, list)) に引っかからないため hook の実装に注意が必要

Step 3: LayerNorm パッチ適用で原因特定

debug_nan2.py を修正し、以下を改善した上で再実行：

  • logit の NaN チェックを CPU 転送後に実施
  • hook が dataclass フィールドも走査するよう修正
  • nn.LayerNorm への ROCm パッチ（CPU 計算フォールバック）を試験的に追加
[PATCH] LayerNorm for ROCm (NEW)

[LOGITS] NaN check (CPU-transferred):
  logits_beat:     shape=[1, 21897]  NaN=0  min=-7.3520  max=-1.8679
  logits_downbeat: shape=[1, 21897]  NaN=0  min=-9.4485  max=-3.6900
  logits_section:  shape=[1, 21897]  NaN=0  min=-10.0945 max=-5.4153
  logits_function: shape=[1, 10, 21897]  NaN=0  min=-2.5694 max=4.8950

[HOOK] No NaN detected in any module output

==> SUCCESS: No NaN detected with LayerNorm patch applied!

nn.LayerNorm が ROCm 上で NaN を生成していることが根本原因と確定。

AllInOne モデルは LayerNorm を 11 層で多用しており（各 AllInOneBlock に layernorm_before / layernorm_after、最終段に norm）、GroupNorm と同じ根本原因（GPU 上の大規模 reduction 演算）で NaN が発生していた。

Step 4: 古い spec キャッシュの問題

LayerNorm パッチを本番コードに適用した後も最初の実行では結果が改善しなかった。
原因は spec/01_01_インフェリア.npy が Phase 1 修正前の壊れた demix（全 -1.0）から生成されたキャッシュだったため。
spec 自体に NaN はないが、DC 信号に近いほぼゼロのスペクトルを入力すると推論が正常化しなかった。

キャッシュを削除して新しい demix wav から spec を再生成したところ正常に動作した。

# demix wav（修正後・正常）の確認
bass    : peak=0.8831  unique=18  NaN=0  shape=(9656430, 2)
drums   : peak=0.9901  unique=19  NaN=0
other   : peak=0.9901  unique=38  NaN=0
vocals  : peak=0.9901  unique=13  NaN=0

修正

1. _patch_layer_norm_for_rocm() の追加（src/allin1/demix.py）

GroupNorm パッチと同じアプローチで LayerNorm をパッチ：

def _patch_layer_norm_for_rocm():
    _orig_forward = nn.LayerNorm.forward

    def _rocm_forward(self, x: torch.Tensor) -> torch.Tensor:
        if not (x.is_cuda and torch.version.hip is not None):
            return _orig_forward(self, x)
        cpu_out = F.layer_norm(
            x.cpu(),
            self.normalized_shape,
            self.weight.cpu() if self.weight is not None else None,
            self.bias.cpu() if self.bias is not None else None,
            self.eps,
        )
        return cpu_out.to(x.device)

    nn.LayerNorm.forward = _rocm_forward

2. パッチ適用タイミングの変更（src/allin1/demix.py）

以前はパッチを if todos: ブロック内（demix が必要な場合のみ）で適用していた。
demix キャッシュが存在する場合（todos が空）でも AllInOne 推論は実行されるため、パッチを if todos: の外に移動した。

# 修正前：demix が必要な場合のみパッチ適用（推論時にパッチ漏れ）
if todos:
    if str(device) == 'cuda' and ...:
        _patch_sin_embedding_for_rocm()
        ...
    demucs_main(...)

# 修正後：demix キャッシュ有無によらず常にパッチ適用
if str(device) == 'cuda' and getattr(torch.version, 'hip', None) is not None:
    torch.backends.cuda.enable_flash_sdp(False)
    torch.backends.cuda.enable_mem_efficient_sdp(False)
    _patch_sin_embedding_for_rocm()
    _patch_group_norm_for_rocm()
    _patch_layer_norm_for_rocm()   # 新規追加

if todos:
    demucs_main(...)

3. sonify クラッシュの修正（src/allin1/sonify.py）

_sonify_metronome で beats または downbeats が空配列のとき ValueError が発生していた。早期リターンを追加：

def _sonify_metronome(result, length, sr=44100):
    downbeats = np.asarray(result.downbeats)
    beats = np.asarray(result.beats)

    # beats または downbeats が空の場合（NaN 等で検出失敗）は無音を返す
    if len(beats) == 0 or len(downbeats) == 0:
        return np.zeros((2, length), dtype='float32')

    dists = np.abs(downbeats[:, np.newaxis] - beats).min(axis=0)
    ...

結果

# Phase 2 修正後の struct/01_01_インフェリア.json
bpm=133  beats=477  downbeats=120  segments=15
first beats: 0.42, 0.86, 1.30, 1.77, 2.21
first segment: start=0  end=11.16  label=intro
  • bpm が数値として検出される ✅
  • beats / downbeats が正常検出される ✅
  • segments が意味のある区間に分割される（intro/verse/chorus 等）✅
  • sonify / visualize が正常に完了する ✅

ROCm 共通課題まとめ

ROCm GPU では以下の演算が NaN や誤った結果を生成することが確認された。いずれも CPU で計算して GPU に転送することで回避できる。

演算 / API 発生箇所 対処
整数テンソルへのべき乗 ** demucs/transformer.py create_sin_embedding _patch_sin_embedding_for_rocm()
nn.GroupNorm 大規模 reduction demucs/demucs.py DConv _patch_group_norm_for_rocm()
nn.LayerNorm 大規模 reduction src/allin1/models/ AllInOne 全層 _patch_layer_norm_for_rocm()
Flash SDP / Mem-Efficient SDP nn.MultiheadAttention enable_flash_sdp(False) / enable_mem_efficient_sdp(False)
GPU 上の isnan() / != 比較 デバッグ・NaN チェック全般 .detach().cpu() してからチェック

注意: GPU 上の torch.isnan() や t != t は ROCm で正しく動作しない。NaN の有無を確認する際は必ず CPU に転送してから判定すること。

既存パッチ（demucs パッケージ内）

以下は demucs パッケージ自体にすでに組み込まれていた ROCm 対応パッチ（本プロジェクトの修正対象外）：

ファイル パッチ内容
demucs/spec.py STFT/iSTFT を CPU で計算して GPU に転送
demucs/htdemucs.py mean/std の大型テンソル reduction を CPU で計算

変更ファイル一覧

ファイル フェーズ 変更内容
src/allin1/demix.py Phase 1 _patch_sin_embedding_for_rocm() 追加
src/allin1/demix.py Phase 1 _patch_group_norm_for_rocm() 追加（予防的）
src/allin1/demix.py Phase 2 _patch_layer_norm_for_rocm() 追加
src/allin1/demix.py Phase 2 パッチ適用を if todos: の外に移動（demixキャッシュ済み時の漏れ修正）
src/allin1/sonify.py Phase 2 _sonify_metronome に beats=[] ガード追加（クラッシュ修正）

まとめ

AIさんの解説はなかなか理解しづらいのでざっとまとめるとですね、以下のようになります。

今回は、GPUの計算でNaNが発生する問題が根本原因でした。
なので、GPUの計算でNaNが発生する箇所だけCPUに持ってきて計算してGPUに送り返す
ということをひたすらやっています。

あと、ROCmでは整数テンソルへのべき乗がNaNになったり、torch.isnan()をサポートしていない問題もあります。

なので、CPU側に一回転送して、この判定をしないといけないということでした。
一見、GPUとCPU両方とも使うという「聞こえは良い」んですが、実際のところロスの多い実装になっています。（CPU単体よりは全然高速だけど）
ROCmというかハードウェアになるのかな、ネイティブ実装してくれ～～～～～
とはいえ、UnifiedMemoryのおかげでデータの転送ロスは多分免れてるはず。

このレベルのデバッグが物の数時間で終わる世界マジで便利すぎる。

というわけで勝手にForkしてROCm版つくる話でした。

おわり。

2
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
2
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?