RTX 4060 Laptop 8GBでYuE2を動かす｜CUDA OOMをフェーズ別オフロードで回避した方法

音楽生成AI「YuE2」をローカル環境で試してみました。

使用したGPUは NVIDIA GeForce RTX 4060 Laptop GPU 8GB。

YuE2は標準では24GBクラスのVRAMを前提としているため、そのまま実行するとCUDA Out of Memoryが発生しました。

今回は、FP8量子化やCPUオフロードに加えて、ARとNARを処理フェーズごとにGPU／CPUへ退避するようYuE2のコードを変更することで、8GB VRAM環境でも最終的に

CoT：full
Semantic max_tokens：9000
ABC生成
Semantic生成
NAR音声生成
VAEデコード

まで完走できました。

この記事では、そのときに行った対応をまとめます。

使用環境

今回の検証環境は以下です。

OS          : Windows
GPU         : NVIDIA GeForce RTX 4060 Laptop GPU
VRAM        : 8GB
YuE         : YuE2
Python      : 仮想環境 .venv
GPU Compute : CUDA

YuE2のリポジトリは以下のような構成で配置しました。

C:\dev\YuE

仮想環境は、

C:\dev\YuE\.venv

です。

最初に発生したエラー

まず公式サンプルをそのまま実行しました。

python examples/generate.py --output outputs/first-song

すると、CUDA Out of Memoryが発生しました。

代表的なエラーは以下です。

torch.OutOfMemoryError: CUDA out of memory.

GPU 0 has a total capacity of 8.00 GiB
6.00 GiB allowed

Of the allocated memory 5.96 GiB is allocated by PyTorch

ポイントは、

8GB GPUなのに6GBまでしか使えない

という部分です。

YuE2ではGPUメモリをすべて使い切らないよう、内部で約2GiBを予約する実装になっています。

そのため、8GB GPUではPyTorchが利用できる量が実質約6GBになります。

まずCLI版へ変更

examples/generate.pyでは低VRAM向けオプションを指定できないため、YuE2のCLIを使用します。

最初に試したのが以下です。

python -m yue2.cli generate ^
  --request examples/song.json ^
  --output outputs/first-song-fp8 ^
  --budget 8 ^
  --quantization fp8 ^
  --offload-ar

RTX 4060はFP8を利用できるため、

--quantization fp8

を指定しています。

さらに、

--offload-ar

によって、使用していないAR側のモジュールをCPUへ退避させます。

しかし、これだけでは8GBには収まりませんでした。

CUDAメモリ断片化対策

念のためPyTorchのCUDAアロケータ設定も変更しました。

WindowsのCMDでは以下です。

set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

PowerShellの場合は書き方が異なりますが、今回はCMDを使用しています。

その後、

python -m yue2.cli generate ...

を実行します。

ただし今回の場合、最終的な原因はメモリ断片化ではなく、モデル自体のVRAM使用量でした。

Semanticの最大トークン数を削減

YuE2ではSemantic生成時にKV Cacheを確保します。

標準設定では、

semantic.max_tokens = 9000

です。

8GB GPUではこのKV Cacheもかなり大きいため、最初は低い値から試しました。

low-vram.jsonを作成します。

{
  "semantic": {
    "max_tokens": 1024
  }
}

実行コマンドは以下です。

python -m yue2.cli generate ^
  --request examples/song.json ^
  --output outputs/first-song ^
  --budget 8 ^
  --quantization fp8 ^
  --offload-ar ^
  --config low-vram.json

それでもCUDA OOMが発生しました。

VRAM上限を6GBから7GBへ変更してみる

YuE2内部では概ね以下のような処理でGPU利用量を制限しています。

budget = min(
    (self.memory_budget_gib - 2) * 2**30,
    total - 2 * 2**30
)

8GB GPUの場合、

8GB - 2GB = 6GB

となります。

検証として予約領域を1GiBまで減らし、

budget = min(
    (self.memory_budget_gib - 1) * 2**30,
    total - 1 * 2**30
)

に変更しました。

すると、

7.00 GiB allowed

まで利用できるようになりました。

しかし今度は、

GPU 0 has a total capacity of 8.00 GiB
0 bytes is free

7.00 GiB allowed
6.93 GiB allocated

となりました。

つまり、YuE2側の6GB制限ではなく、物理VRAM 8GBそのものを使い切っている状態です。

単純に使用可能VRAMを増やす方法では解決できません。

根本原因はARとNARの同時GPU常駐

YuE2のモデルを見ると、Transformer Layerには、

AR Attention
AR MLP

NAR Attention
NAR MLP

がそれぞれ存在しています。

問題は、AR生成中にもNAR側のモジュールがGPU上に存在していることです。

逆に、NAR生成時にもAR側が必要以上にGPUへ残るタイミングがあります。

そこで、

AR処理中はNARをCPUへ退避

し、

NAR処理中はARをCPUへ退避

するように変更しました。

これが今回最も効果の大きかった対応です。

NARモジュールをCPUへ退避する関数を追加

src\yue2\nar.pyへ以下の関数を追加しました。

def _move_nar_solve_modules(model, device):
    """Move NAR-only compute modules between CPU and GPU."""

    modules = [
        model.llm2vae,
        model.time_embedder,
    ]

    for layer in model.model.layers:
        modules.extend([
            layer.nar_input_layernorm,
            layer.nar_self_attn,
            layer.nar_pre_mlp_layernorm,
            layer.nar_mlp,
        ])

    for module in modules:
        module.to(device=device)

これでNARだけに使用するモジュールをまとめてCPU／GPU間で移動できます。

AR生成中はNARをCPUへ移動

次にpipeline.py側を変更しました。

モデルロード後、AR生成を行う段階ではNAR側のモジュールをCPUへ移動します。

イメージとしては以下です。

self._model.to(self.device)

if self.offload_ar and self.device.type == "cuda" and not for_nar:
    from .nar import _move_nar_solve_modules

    _move_nar_solve_modules(
        self._model,
        "cpu"
    )

    torch.cuda.empty_cache()

これによって、

AR生成時

AR  → GPU
NAR → CPU

という状態にします。

NAR生成時はARをCPUへ移動

NAR生成時には逆に、

AR  → CPU
NAR → GPU

へ切り替えます。

synthesize()の処理を以下のような構成にしました。

engine = CachedNAR(
    model,
    chunk,
    attention,
    query_chunk_size,
)

with _offload_ar(model, offload_ar):

    if offload_ar:
        _move_nar_solve_modules(
            model,
            engine.device
        )

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    try:
        progress = None

        if on_progress is not None:
            def progress(completed, total):
                on_progress(
                    chunk_index * total + completed,
                    total * len(chunks),
                )

        output.append(
            engine.solve(
                steps,
                cancelled,
                on_progress=progress,
            )
        )

    finally:
        engine.close()

        if offload_ar:
            _move_nar_solve_modules(
                model,
                "cpu"
            )

            if torch.cuda.is_available():
                torch.cuda.empty_cache()

del engine

重要なのは、

同時にARとNARをGPUへ置かない

という考え方です。

Pythonの構文チェック

コードを変更したあとは、いきなりYuE2を実行せずに構文チェックを行いました。

python -m py_compile C:\dev\YuE\src\yue2\nar.py

続いて、

python -m py_compile C:\dev\YuE\src\yue2\pipeline.py

何も表示されなければPythonの構文としては問題ありません。

途中、インデントを間違えて、

IndentationError:
expected an indented block after 'with' statement

も発生したため、コード変更後のpy_compileはおすすめです。

まずSemantic 1024で成功

修正後、まずはかなり低い設定でテストしました。

{
  "semantic": {
    "max_tokens": 1024
  }
}

さらにVRAMを節約するため、

cot = off
cfg-scale = 1

を指定しました。

コマンドは以下です。

python -m yue2.cli generate ^
  --request examples/song.json ^
  --output outputs/first-song-phase-offload ^
  --budget 8 ^
  --quantization fp8 ^
  --offload-ar ^
  --cot off ^
  --cfg-scale 1 ^
  --config low-vram.json

結果、

[YuE2] Finished (generation limit reached): 41.0s audio in 84.4s

となり、ついに生成に成功しました。

結果は、

{
  "status": "complete",
  "truncated": {
    "abc": false,
    "semantic": true
  }
}

でした。

semantic=trueはエラーではなく、設定した1024トークンの上限まで生成したことを意味します。

Semantic上限を徐々に増やす

1024で成功したので、少しずつ最大トークン数を増やしました。

結果は以下です。

1024  → 成功
1536  → 成功
2048  → 成功
2560  → 成功
3072  → 成功
3584  → 成功
4096  → 成功
9000  → 成功

最終的には、

semantic.max_tokens = 9000

というYuE2標準値でも完走しました。

つまり、KV Cacheを削らなければ動かないわけではなく、

AR/NARの同時GPU常駐を避けたことが決定的だった

と考えられます。

最後にcot=fullもテスト

次に、簡易的な

cot=off

ではなく、

cot=full

へ戻しました。

つまり、

歌詞・スタイル
↓
ABC生成
↓
Semantic生成
↓
NAR生成
↓
VAEデコード
↓
音声

というYuE2本来のフルパイプラインです。

最終的な実行コマンドは以下です。

set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

python -m yue2.cli generate ^
  --request examples/song.json ^
  --output outputs/song ^
  --budget 8 ^
  --quantization fp8 ^
  --offload-ar ^
  --cot full

この構成でも正常に完走しました。

最終的に成功した構成

今回RTX 4060 Laptop GPU 8GBで成功した構成をまとめると、

GPU
RTX 4060 Laptop GPU 8GB

YuE2
YuE2-3B

Quantization
FP8

CoT
full

Semantic max_tokens
9000

AR offload
ON

追加対応
AR/NARフェーズ別CPU/GPUオフロード

です。

処理イメージは以下です。

【AR / ABC / Semantic生成】

AR
↓
GPU

NAR
↓
CPU


【NAR音声生成】

AR
↓
CPU

NAR
↓
GPU


【VAEデコード】

不要になったモデル
↓
CPU

VAE
↓
GPU

このように、その瞬間に必要なモジュールだけをGPUへ置くことで8GB VRAMでも動かせました。

修正内容は必ず保存しておく

今回の変更はYuE2公式コードに対するローカル修正です。

git pullや再インストールによって消える可能性があるため、差分を保存しました。

git diff > yue2-rtx4060-8gb.patch

必要になった場合は、

git apply yue2-rtx4060-8gb.patch

で再適用できます。

専用ブランチを作る方法でも良いと思います。

git switch -c rtx4060-8gb

git add src\yue2\nar.py src\yue2\pipeline.py

git commit -m "Add phase offloading for 8GB GPUs"

注意点

今回の方法は、RTX 4060 Laptop GPU 8GB環境で実際に動作確認したものですが、YuE2公式の8GB対応方法ではありません。

GPU、CUDA、PyTorch、ドライバ、モデルのバージョンによって結果が変わる可能性があります。

特に、

--quantization fp8

についてはGPU側のFP8対応が必要です。

またCPUへモデルを頻繁に退避するため、

VRAM使用量
↓
減少

RAM使用量
↓
増加

CPU⇔GPU転送
↓
増加

というトレードオフがあります。

速度よりも、

「8GB GPUでまず動かす」

ことを優先した方法です。

まとめ

最初は、

6.00 GiB allowed
5.97 GiB allocated
CUDA Out of Memory

となり、RTX 4060 Laptop GPU 8GBではYuE2の実行は難しい状態でした。

しかし、

FP8
+
AR offload
+
AR/NARフェーズ別オフロード

を組み合わせることで、

cot=full
semantic.max_tokens=9000

まで正常に動作しました。

今回特に効果が大きかったのは、

VRAM上限を無理に増やすのではなく、使っていないモデル部分をCPUへ逃がすこと