0
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?

@kotetsu_yama(Kotetsu Yamamoto)

NVIDIA NeMo の日本語 ASR を AMD ROCm で動かす — 移植に必要だったのは 4 行だった

0
Posted at

Summary
NVIDIA の音声認識モデル nemotron-3.5-asr-streaming-0.6b を、AMD Ryzen AI MAX+ 395 (gfx1151) の ROCm 上で動かした。
NeMo は ROCm を公式サポートしていないが、推論に必要だったコード変更は 1 ファイル 4 行だった。
日本語 27 件の認識結果は CPU 版と完全に一致、速度は CPU 比 3.2 倍。
事前に身構えていた CUDA Graphs / Flash Attention / Numba-CUDA はいずれも対応不要で、実際の障壁は別のところにあった。

はじめに

NVIDIA NeMo は当然ながら CUDA 前提のツールキットで、ROCm は公式サポート外です。
一方で対象の nvidia/nemotron-3.5-asr-streaming-0.6b は 40 言語対応・cache-aware ストリーミング対応の 0.6B ASR モデルで、日本語も transcription-ready 扱い。ローカルで動かせるなら魅力的です。

「どうせ CUDA べったりで移植は大仕事だろう」と思って始めたのですが、実際に必要だった変更は 4 行でした。
そして事前にリスクとして挙げていた項目は、ほぼ全部が空振りでした。

事前の懸念 実際
CUDA Graphs を無効化しないと動かない 半分だけ正解。conditional node を使うモードだけ落ちる。部分グラフは HIP で動く
Flash Attention を ROCm 版に差し替える必要がある 空振り。このモデルは flash-attn を使っていない
RNNT loss の Numba-CUDA がデコードで呼ばれるかも 空振り。ロード時に構築されるだけで推論では呼ばれない
（想定外） uv の依存解決が一番の関門だった

時間を溶かしたのは ROCm 対応そのものではなく、エラーメッセージが原因を指していない箇所と、ROCm と無関係な NeMo 側の落とし穴でした。この記事はその両方を書きます。

差分・検証スクリプト・記録一式は fork に置いてあります: https://github.com/kotetsuy/Speech/tree/rocm-inference/rocm-inference

検証環境

項目 内容
機種 GMKtec NucBox EVO X2
CPU/GPU AMD Ryzen AI MAX+ 395 w/ Radeon 8060S (gfx1151)
メモリ 96GB unified (VRAM 48GiB + システム 45GiB)
OS Ubuntu 26.04 LTS
ROCm 7.14.0 (/opt/rocm, gfx1151 ネイティブビルド)
PyTorch 2.9.1+rocm7.13.0 (Python 3.12)
NeMo Speech 3.1.0.dev (ef41369156)
検証日 2026-08-13

HSA_OVERRIDE_GFX_VERSION は設定しません。
gfx1151 ネイティブビルドの ROCm では、アーキテクチャを上書きすると逆に動かなくなります。
ネット上の gfx1151 の記事には 11.5.1 を設定する手順が多いですが、それは gfx1151 向けビルドが無かった頃の話です。

結果

CPU ROCm (gfx1151)
バッチ（オフライン）推論 ✅ ✅
cache-aware ストリーミング 160ms チャンク ✅ ✅
cache-aware ストリーミング 1120ms チャンク ✅ ✅
出力の一致 — 27/27 件が CPU と完全一致
純粋な推論時間の中央値 0.170〜0.284s 0.052〜0.086s
高速化 1.0x 約 3.2x

第 1 部 — 移植に必要だった 4 行

症状

ROCm 環境でモデルをロードしようとすると、推論以前にモデル生成の時点で落ちます。

[NeMo E] Model instantiation failed!
    Target class: nemo.collections.asr.models.rnnt_bpe_models_prompt.EncDecRNNTBPEModelWithPrompt
    Error(s): Failed to dlopen libcuda.so.1
...
TypeError: Can't instantiate abstract class ASRModel without an implementation for
abstract methods 'setup_training_data', 'setup_validation_data'

このメッセージ、最後の TypeError が完全にミスリードです。
「抽象クラスを実装しろ」と言われると自分のコードの書き方を疑いますが、これは具象クラスの生成に失敗した後の二次被害にすぎません。本当の原因はその上の Failed to dlopen libcuda.so.1 です。

原因

nemo/core/utils/cuda_python_utils.py の CUDA Graphs サポート判定にあります。

def check_cuda_python_cuda_graphs_conditional_nodes_supported():
    # for CPU-only environment we need to raise an exception, otherwise cuda-python library will fail
    if not torch.cuda.is_available():
        raise EnvironmentError("CUDA is not available")

    try:
        from cuda.bindings import driver as cuda
    ...
    error, driver_version = cuda.cuDriverGetVersion()   # ← ここで死ぬ

ポイントは、ROCm ビルドの PyTorch では torch.cuda.is_available() が True を返すことです。
HIP は CUDA の API 面をそのまま被せてくるので、torch.cuda.* はひととおり動きます。だから CPU 用のガードを素通りしてしまう。

その先で cuda-python が libcuda.so.1 を dlopen しに行きますが、AMD マシンに NVIDIA のドライバはありません。ここで RuntimeError が飛びます。

そして呼び出し側がこうなっています。

# nemo/collections/asr/parts/submodules/transducer_decoding/label_looping_base.py
try:
    check_cuda_python_cuda_graphs_conditional_nodes_supported()
    self.cuda_graphs_mode = self.CudaGraphsMode.FULL_GRAPH
except (ImportError, ModuleNotFoundError, EnvironmentError) as e:
    logging.warning("No conditional node support for Cuda. ...")
    self.cuda_graphs_mode = self.CudaGraphsMode.NO_WHILE_LOOPS

捕捉しているのは (ImportError, ModuleNotFoundError, EnvironmentError) の 3 つだけ。
RuntimeError は素通りしてデコーダのコンストラクタを突き抜け、モデル生成ごと巻き添えにします。

CPU 環境ではこの経路は正しく動きます（torch.cuda.is_available() が False なので EnvironmentError になり、フォールバックする）。ROCm だけが、CPU でも CUDA でもない第三の状態に落ちるわけです。

修正

やることは単純で、「ROCm は未対応環境である」と、呼び出し側が期待している例外型で伝えるだけです。

 if not torch.cuda.is_available():
     raise EnvironmentError("CUDA is not available")

+    # On a ROCm/HIP build `torch.cuda.is_available()` is True but there is no CUDA driver:
+    # cuda-python would fail to dlopen libcuda.so.1 and raise a bare RuntimeError, which
+    # callers such as `maybe_enable_cuda_graphs` do not catch. Report it as an unsupported
+    # environment instead, so they fall back to the non-graph decoding path.
+    if torch.version.hip is not None:
+        raise EnvironmentError("CUDA graphs with conditional nodes are not supported on ROCm/HIP")
+
     try:
         from cuda.bindings import driver as cuda

torch.version.hip は ROCm ビルドでのみ非 None になるので、CUDA 環境の挙動には一切影響しません。

これだけで、バッチ推論もストリーミング推論も通るようになりました。移植で書いたコードは、実質この 4 行だけです。

おまけ：テストも直った

この修正の副作用として、NeMo のテストスイートも ROCm 上でまともになりました。

tests/collections/asr/decoding/test_cuda_graph_rnnt_greedy_decoding.py
  → 4 passed, 14 skipped

conditional node 系のテストは skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported() を使っていて、これも同じ 3 例外しか見ていません。修正前はテストが skip されずにエラーになっていました。

第 2 部 — 移植不要だったもの

事前に「ここが大変だろう」と挙げていた項目が、調べたらどれも該当しませんでした。移植の見積もりを外した記録として残しておきます。

CUDA Graphs は「全部無効化」ではない

仕様検討時は「ROCm では CUDA Graphs を切って eager にフォールバックさせる」と書いていました。実際は違いました。

NeMo の RNNT デコーダには 3 モードあります。

モード 内容 ROCm
full_graph CUDA graph の conditional node を使う最速実装 ❌ 使えない
no_while_loops PyTorch の while ループ + 部分的な CUDA graph ✅ 動く
no_graphs グラフなし（デバッグ用） ✅

使えないのは full_graph だけです。no_while_loops が使う torch.cuda.CUDAGraph は HIP graphs にマップされて普通に動くので、グラフによる高速化の恩恵はちゃんと残ります。実行時に確認するとこうなります。

cuda_graphs_mode : no_while_loops
allow_cuda_graphs: True
Reason: CUDA graphs with conditional nodes are not supported on ROCm/HIP

CPU 比 3.2 倍という数字が出たのは、これが効いているからです。「ROCm だから CUDA Graphs は全部ダメ」と決め打ちして eager に落としていたら、だいぶ損をしていました。

Flash Attention は最初から使われていない

「CUDA 版 flash-attn を ROCm 版（composable_kernel ベース）に差し替える必要がある」と身構えていました。

$ grep -rn "flash_attn\|FlashAttention" nemo/collections/asr/ --include=*.py
（何も出ない）

このモデルの encoder は self_attention_model: rel_pos（相対位置エンコーディング）で、そもそも flash-attn を経由しません。ASR コレクション全体に import すら存在しませんでした。

Numba-CUDA の RNNT loss は推論では呼ばれない

ロード時のログにはしっかり出ます。

[NeMo I] Using RNNT Loss : warprnnt_numba

これを見て「デコードでも呼ばれるのでは」と疑っていたのですが、構築されるだけで推論では使われません。判定は簡単で、CPU 環境（Numba CUDA が動くはずがない）で推論が完走している時点で、推論パスには乗っていないと確定します。

学習には当然必要なので、学習の移植は別問題です。今回は推論のみがスコープでした。

その他の CUDA 専用 API

torch.cuda.CUDAGraph や cuda.bindings を使う箇所は他にもありますが、いずれも今回の経路に入りません。

ファイル 使われる条件 今回
streaming_encoder_cuda_graphs.py use_cuda_graphs=True（既定 off） 未使用
ctc_batched_beam_decoding.py CTC beam search 未使用（RNNT）
rnnt_malsd_batched_computer.py RNNT beam search 未使用（greedy_batch）

第 3 部 — 実際の関門は uv だった

コード移植より手こずったのが、Python 環境の構築です。2 つの罠がありました。

罠 1：マルチアーキ wheel は gfx1151 で実行時に落ちる

ROCm 版 PyTorch には複数の配布経路がありますが、gfx1151 では専用インデックスを使わないといけません。

# ✅ これ
--index-url https://repo.amd.com/rocm/whl/gfx1151/

# ❌ マルチアーキ wheel はインストールは通るが実行時に落ちる
#    hipErrorInvalidImage / kpack_load_code_object failed with error: 13

インストールは成功するのが厄介なところで、実際に GPU カーネルを起動する段になって初めて失敗します。

罠 2：index-strategy = "unsafe-best-match" が必須

これが一番わかりにくかった。uv sync がこう言って止まります。

× No solution found when resolving dependencies
╰─▶ Because there is no version of rocm[libraries]==7.13.0 and
    torch==2.9.1+rocm7.13.0 depends on rocm[libraries]==7.13.0, we can
    conclude that torch==2.9.1+rocm7.13.0 cannot be used.

torch==2.9.1+rocm7.13.0 は ROCm ランタイムを wheel として引き込む（rocm[libraries]）のですが、その sdist は AMD のインデックスにしか存在しません。
uv の既定は first-index 戦略で、パッケージが見つかった最初のインデックスで探索を打ち切ります。rocm を PyPI で探して見つからず、そこで諦めてしまうわけです。

[tool.uv]
index-strategy = "unsafe-best-match"

これで解決します。uv pip install に --index-strategy unsafe-best-match を付ける形でも同じです。

罠 3（NeMo 固有）：torch のインデックス指定が衝突する

ついでにもう一つ。NeMo Speech の pyproject.toml は torch のインデックスを自前で指定しています。

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu",   marker = "sys_platform != 'linux' and sys_platform != 'darwin'" },
  { index = "pytorch-cu129", extra = "cu12", marker = "sys_platform == 'linux'" },
  { index = "pytorch-cu132", extra = "cu13", marker = "sys_platform == 'linux'" },
  { index = "pypi",          marker = "sys_platform == 'darwin'" },
]

NeMo を editable な path 依存として自分のプロジェクトから参照すると、この指定が自分の指定とぶつかります。

× Requirements contain conflicting indexes for package `torch`:
  - https://pypi.org/simple
  - https://download.pytorch.org/whl/cpu

Linux でしか動かさないなら、解決対象の環境を絞るのが手っ取り早いです。

[tool.uv]
environments = ["sys_platform == 'linux'"]

これで darwin 向けの指定が解決から消えて衝突しなくなります。

最終的な pyproject.toml

[project]
name = "nemo-rocm-verify-rocm"
version = "0.1.0"
requires-python = "==3.12.*"
dependencies = [
    "nemo-toolkit[asr]",
    "torch==2.9.1+rocm7.13.0",
    "cython", "soundfile", "librosa", "huggingface-hub[cli]",
]

[tool.uv]
environments = ["sys_platform == 'linux'"]
index-strategy = "unsafe-best-match"

[tool.uv.sources]
nemo-toolkit = { path = "../..", editable = true }
torch = [{ index = "rocm-gfx1151", marker = "sys_platform == 'linux'" }]

[[tool.uv.index]]
name = "rocm-gfx1151"
url = "https://repo.amd.com/rocm/whl/gfx1151/"

動作確認:

$ .venv/bin/python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"
2.9.1+rocm7.13.0 True Radeon 8060S Graphics

torch.cuda.is_available() が True で Radeon が出てくるのは、何度見ても妙な気分です。

第 4 部 — clone から日本語 ASR が動くまで

ここまでの内容を踏まえた、ゼロからの手順です。この節の内容は、まっさらな clone から実際に通ることを確認しています。

前提

項目 備考
ROCm /opt/rocm にインストール済み。今回は 7.14.0
GPU gfx1151（他アーキは未検証）
Python 3.12（uv が勝手に用意するので事前準備は不要）
uv https://docs.astral.sh/uv/
ディスク clone に約 190MB + venv とモデルで数 GB

HSA_OVERRIDE_GFX_VERSION は設定しません（設定すると壊れます）。
Hugging Face のトークンも不要です。このモデルは gated ではありません。

システムパッケージを 2 つだけ入れておきます。

sudo apt-get update && sudo apt-get install -y libsndfile1 ffmpeg

STEP 1 — clone

git clone --depth 1 -b rocm-inference https://github.com/kotetsuy/Speech.git
cd Speech

ROCm 対応の差分（第 1 部の 4 行）は、このブランチに入っています。

STEP 2 — 環境構築

cd rocm-inference/env-rocm
uv venv --python 3.12
uv sync
cd ../..

uv sync は rocm-inference/env-rocm/pyproject.toml を読みます。第 3 部で書いた 3 つの罠（gfx1151 専用インデックス / unsafe-best-match / environments による衝突回避）は、この pyproject に書き込んであるので、そのまま通ります。

GPU が見えていることを確認します。

rocm-inference/env-rocm/.venv/bin/python -c \
  "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"

2.9.1+rocm7.13.0 True Radeon 8060S Graphics

STEP 3 — 音声を用意する

モデルの入力仕様は 16kHz / mono / WAV です。手持ちの音声を変換します。

mkdir -p audio
ffmpeg -i input.wav -af "aresample=resampler=soxr:out_sample_rate=16000,apad=pad_dur=0.5" \
  -ac 1 -sample_fmt s16 audio/test_ja.wav

apad=pad_dur=0.5 で末尾に 0.5 秒の無音を足しているのが地味に重要です（理由は第 5 部）。
複数ファイルをまとめて変換するなら同梱スクリプトが使えます。

rocm-inference/env-rocm/.venv/bin/python rocm-inference/scripts/prepare_audio.py \
  --in-dir raw --out-dir audio

STEP 4 — バッチ（オフライン）推論

rocm-inference/env-rocm/.venv/bin/python rocm-inference/scripts/transcribe_batch.py \
  --audio-dir audio --lang ja-JP --device cuda --out result.jsonl

[load] EncDecRNNTBPEModelWithPrompt on cuda in 25.5s
[lang] target_lang=ja-JP -> prompt id 10
[batch]   0.289s  test_ja.wav: こんにちは
[out] wrote 1 rows to result.jsonl

初回はモデルのダウンロードが走るので、ロードに時間がかかります。2 回目以降は HF のキャッシュから読まれます。

--device cuda で ROCm GPU を指します（HIP が CUDA の API を被せているので、指定は cuda のままです）。CPU で動かすなら --device cpu にしてください。

このスクリプトは、第 5 部で書く「言語プロンプトが固定されない問題」を回避するために、内部で lang と prompt_mode を入れた manifest を書き出しています。

STEP 5 — cache-aware ストリーミング推論

ストリーミングは NeMo 同梱のスクリプトを使います。入力は manifest で渡します（audio_file= にはバグがあります。第 5 部参照）。

manifest は 1 行 1 音源の JSON Lines です。

python3 -c "
import json, wave
p = 'audio/test_ja.wav'
with wave.open(p, 'rb') as w:
    d = w.getnframes() / w.getframerate()
open('manifest.json', 'w').write(json.dumps({'audio_filepath': p, 'duration': d, 'text': ''}) + '\n')
"

{"audio_filepath": "audio/test_ja.wav", "duration": 3.75, "text": ""}

実行します。

rocm-inference/env-rocm/.venv/bin/python \
  examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py \
  pretrained_name=nvidia/nemotron-3.5-asr-streaming-0.6b \
  dataset_manifest=manifest.json \
  output_path=stream_out \
  target_lang=ja-JP \
  strip_lang_tags=true \
  att_context_size="[56,13]" \
  cuda=0 \
  batch_size=1

[NeMo I] Final streaming transcriptions: ['こんにちは']

主なオプション:

オプション 意味
target_lang=ja-JP 言語プロンプト。auto にすると自動判定して <xx-XX> タグを付ける
strip_lang_tags=true 出力末尾の <ja-JP> タグを除去する
att_context_size="[56,13]" チャンク長 1120ms。160ms なら [56,1]
cuda=0 GPU 番号。CPU で動かすなら cuda=-1

バッチ推論と違い、こちらの manifest には lang / prompt_mode が要りません。
ストリーミングスクリプトは target_lang からモデル側のプロンプトを直接設定するためです（ログに Inference prompt set to 'ja-JP' (index 10) と出ます）。

結果は stream_out/ ディレクトリの中に JSON で出ます（ファイル名ではなくディレクトリなので注意）。

STEP 6 — ベンチマーク（任意）

rocm-inference/env-rocm/.venv/bin/python rocm-inference/scripts/benchmark.py \
  --audio-dir audio --device cuda --runs 20 --warmup 5 --out bench.json

ウォームアップ後の encoder forward + RNNT デコードのみを、torch.cuda.synchronize() 込みで計測します。

第 5 部 — ROCm と無関係にハマったところ

ここからは ROCm の話ではありませんが、このモデルを使うなら踏む地雷なので書いておきます。むしろこちらの方が時間を使いました。

model.transcribe(target_lang=...) では言語が固定されない

このモデルは language-ID prompt conditioning で言語を指定します。日本語のキーは ja-JP（ja 単体は prompt_dictionary に無く弾かれます）。

素直に書くとこうしたくなります。

model.transcribe(audio=["a.wav"], target_lang="ja-JP")   # ❌ 効かない

これは動きません。

ValueError: Unknown prompt key: 'None'. Available: ['en-US', 'en', 'en-GB', ...]

プロンプトは cut.supervisions[0].language、つまり manifest の lang フィールドから解決されます。音声パスだけ渡すと合成 manifest に lang が無いので None になる、という話です。

さらに厄介なのがこの先で、_setup_transcribe_dataloader は default_prompt_mode を設定しません。既定値は unified で、その中身がこうなっています。

# nemo/collections/asr/data/audio_to_text_lhotse_prompt_index.py
elif mode == 'unified':
    if random.random() < self.unified_auto_ratio:   # 既定 0.5
        return self.auto_index                       # ← 言語非依存の auto プロンプト
    return self._get_prompt_index(cut.supervisions[0].language)

言語が分かっていても、約半分のサンプルで auto プロンプトが選ばれます。
lang を埋めて ValueError を回避しただけだと、この挙動に気づかないまま「なんか結果が安定しないな」で終わります。学習時の意図（言語指定と auto を混ぜて学習する）がそのまま推論に漏れている形です。

NeMo 本体に手を入れずに回避するなら、manifest を自前で書いて prompt_mode を明示します。manifest の余剰キーは cut.custom に入り、per-cut の prompt_mode はそこから読まれます。

{
  "audio_filepath": "/path/to/a.wav",
  "duration": 3.25,
  "text": "",
  "lang": "ja-JP",         # -> cut.supervisions[0].language
  "prompt_mode": "langID"  # -> cut.custom、auto を引かせない
}

VOICEVOX 音源は末尾に無音を足さないと助詞が落ちる

評価用に、自分の肉声・VOICEVOX ずんだもん・VOICEVOX 波音リツの 3 種で音声を用意しました。最初のストリーミング実行で、VOICEVOX 音源だけ文末が欠けました。

原文 出力
日本で二番目に高い山は 日本で二番目に高い山
日本で一番長い川は 日本で一番長い川

末尾無音を測ると原因は明らかでした。

音源 末尾無音
肉声 0.812s
ずんだもん 0.105s
波音リツ 0.169s

cache-aware ストリーミングは発話終端の右 context が足りないと最後のトークンを出力できません。人間の録音は録音開始・停止の操作でどうしても前後に余白ができますが、TTS 出力にはそれが無いわけです。

ffmpeg -i in.wav -af "apad=pad_dur=0.5" -ac 1 -ar 16000 -sample_fmt s16 out.wav

0.5 秒足すだけで 3 件とも解消しました。TTS 音声で ASR を評価するときの一般的な注意点だと思います。

ストリーミング推論スクリプトの細かい罠

NeMo 同梱の examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py を使う場合:

  • output_path= はファイルではなくディレクトリとして扱われ、中に streaming_out_<model>_<manifest>.json が作られる
  • その JSON は pred_text / text / wer しか持たず、ファイル名を含まない。manifest の順序で対応付けるしかない
  • audio_file=（単一ファイル指定）は転写結果を表示した直後にクラッシュする
    （UnboundLocalError: cannot access local variable 'all_refs_text'。manifest 用の変数を単一ファイル経路でも参照している）。dataset_manifest= を使えば回避できる
  • 出力される WER% は日本語では無意味。単語区切りで測っているので、モデルカードにあるとおり日本語は CER で見るべき

チャンクサイズの指定方法

ストリーミングのレイテンシは att_context_size = [左context, 右context] で指定します。単位は 80ms フレームで、チャンク長 = (右context + 1) × 80ms。

チャンク長 att_context_size
80ms [56, 0]
160ms [56, 1]
320ms [56, 3]
560ms [56, 6]
1120ms [56, 13]

右 context は {0, 1, 3, 6, 13} から選びます。今回は 160ms と 1120ms の 2 点で検証しました。

認識精度と速度

精度：ROCm 版は CPU 版と完全一致

日本語 9 音源（3 話者 × 3 文）× 3 モード（バッチ / ストリーミング 160ms / 1120ms）= 27 件を両環境で実行しました。

27/27 件で出力が完全一致。 ROCm 移植による劣化はゼロです。

原文との比較では 25/27 が完全一致でした。

読み上げた 3 文は「こんにちは」「日本で二番目に高い山は」「日本で一番長い川は」です。

音源 バッチ 160ms 1120ms
肉声「こんにちは」 ✅ ✅ ✅
肉声「山」 ✅ ✅ ✅
肉声「川」 ✅ ✅ ✅
ずんだもん「こんにちは」 ✅ ✅ ✅
ずんだもん「山」 ✅ ✅ ✅
ずんだもん「川」 ✅ ✅ ✅
波音リツ「こんにちは」 ✅ ✅ ✅
波音リツ「山」 ✅ ✅ ✅
波音リツ「川」 ⚠️ 一本で一番長い川は ⚠️ みっとんで一番長い川は ✅

外したのは波音リツの 1 音源だけで、しかも冒頭の「日本で」限定。同じ話者の別の文では「日本で」を正しく認識しているので、その音源の読み上げ固有の問題のようです。1120ms チャンクでは正解しているのも、右 context が長いほど有利という cache-aware の挙動と整合します。

速度：CPU 比 3.2 倍

計測は「モデルロード済み・ウォームアップ 5 回後・20 回の中央値」。torch.cuda.synchronize() を計測区間に含めています。

transcribe() は呼び出しごとに Lhotse の dataloader を作り直すので、そこは推論時間ではないと判断し、encoder forward + RNNT デコードのみを測りました。

音源 長さ CPU ROCm 高速化
肉声「こんにちは」 3.75s 0.2397s 0.0765s 3.13x
肉声「山」 4.00s 0.2840s 0.0861s 3.30x
肉声「川」 4.00s 0.2826s 0.0837s 3.38x
リツ「こんにちは」 1.35s 0.1721s 0.0519s 3.32x
リツ「山」 2.65s 0.2439s 0.0799s 3.05x
リツ「川」 2.41s 0.2210s 0.0624s 3.54x
ずんだもん「こんにちは」 1.56s 0.1703s 0.0523s 3.26x
ずんだもん「山」 2.95s 0.2412s 0.0799s 3.02x
ずんだもん「川」 2.81s 0.2310s 0.0781s 2.96x

実行ごとのばらつきは min/max で 1% 以内と安定していました。

計測方法で数字が倍以上変わる話
最初は transcribe() 全体を測っていて CPU 0.46 秒という値を得ていました。dataloader 構築を除くと 0.24 秒。約 2 倍の差です。
ベンチマークを載せるときは、何を計測区間に含めたかを書かないと意味が無い、という当たり前のことを再確認しました。

なお cache-aware モデルは compute_dtype != float32 を明示的に NotImplementedError で弾くので、上記は全て float32 での値です。fp16/bf16 による追加の高速化余地はありません（これは ROCm に限らず CUDA でも同じです）。

制限事項

正直に書いておきます。

  • 推論のみ。 学習は移植していないし試してもいません。Numba-CUDA の RNNT loss、Apex、Transformer Engine、Megatron-core はすべて対象外です
  • float32 固定。 上記のとおり
  • conditional node を使う full_graph デコードは使えない。 no_while_loops にフォールバックします
  • 短尺音声のみ。 検証したのは 1.35〜4.00 秒の発話だけです。長尺のストリーミングにおける cache の挙動・安定性は未検証。同じマシンの whisperX では 60 秒超の音声で Memory access fault by GPU node-1 が出る事例があるので、長尺は別途確認が要ります
  • gfx1151 のみ、このモデルのみ。 他のアーキテクチャや他のモデルは未検証です

まとめ

  • NeMo の ASR 推論は、1 ファイル 4 行の変更で ROCm (gfx1151) 上で動いた
  • 原因は「ROCm では torch.cuda.is_available() が True を返す」ことに尽きる。CPU 用のガードを素通りして、CUDA 前提のコードに突入する
  • そこで飛ぶ RuntimeError が、呼び出し側の except (ImportError, ModuleNotFoundError, EnvironmentError) に引っかからず、エラーメッセージが原因を指さない形で落ちる
  • 事前に挙げた懸念（CUDA Graphs / Flash Attention / Numba-CUDA）はほぼ全部空振りだった。実際の関門は uv の依存解決と、ROCm と無関係な NeMo 側の仕様
  • 認識精度は CPU と完全一致、速度は CPU 比 3.2 倍

同じ構成で torch.cuda.is_available() が True を返すことに起因する不具合は、NeMo に限らず色々なところにありそうです。「CPU でも CUDA でもない第三の状態」を想定していないコードを見つけたら、だいたいこのパターンだと思います。

差分と検証スクリプト一式はこちらに置いてあります。
https://github.com/kotetsuy/Speech/tree/rocm-inference/rocm-inference

0
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
0
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?