12GB VRAMで動かすCosmos 3：メタデバイス問題の回避と高速推論の実践

はじめに

NVIDIAが発表した最新の高性能物理世界シミュレータ・画像動画生成モデル「Cosmos 3」は、その高い表現力と引き換えに、15B（150億パラメータ）スケールという巨大なモデルサイズを要求します。一般家庭や個人開発者のGPU環境（例：VRAM 12GBのRTX 3080 Tiなど）でこの巨大モデルを触ることは、従来では困難と考えられていました。

しかし、コミュニティによってトランスフォーマーモデルを4bit（NF4量子化）に圧縮した「SanDiegoDude/Cosmos3-Nano-nf4」が公開されたことで、12GB VRAM内での動作可能性が見えてきました。モデルカードによれば、ロード時のVRAM使用量は約11GBとされており、RTX 3080 Tiの限界ギリギリで動作させることができます。

だが、実際にHugging Face Diffusersライブラリを用いてモデルをロードしようとすると、メタデバイス（meta tensor）に起因するエラーに直面します。本稿では、このロード時バグのメカニズムを特定し、数行のモンキーパッチを用いて12GB VRAM環境でCosmos 3を正常ロードさせ、画像および動画生成を完結させるための具体的な手順と実装コードを解説します。

1. ロード時エラー「Cannot copy out of meta tensor」の解析

クラッシュのコールフロー

DiffusersでNF4量子化されたトランスフォーマーモデルをロードする際、config.jsonにquantization_configが含まれていると、システムは自動的にCPUメモリ消費を抑えるために low_cpu_mem_usage=True を強制します。この時、バックエンドの accelerate ライブラリの機能である init_empty_weights() コンテキストが実行され、モデル全体のパラメータが「実データを持たないダミーテンソル（デバイス: meta）」としてアロケートされます。その後、ディスク上のチェックポイントから実際のウェイトデータがロードされ、メタデバイス上のパラメータへと順次コピーされます。

しかし、今回のCosmos 3-Nano NF4モデルでは、ロード処理の最終段階（GPUへのメモリ転送やディスパッチ時）に以下のエラーを吐いてクラッシュします。

NotImplementedError: Cannot copy out of meta tensor; no data!
# または、別バージョンの accelerate の場合：
ValueError: action_modality_embed is on the meta device, we need a value to put in on cuda.

これは、モデルのクラス定義（Cosmos3OmniTransformer）内に存在するはずのいくつかのパラメータ（action_modality_embed や action_proj_in.fc.weight などのModality Adapter関連の重み）が、今回のテキスト画像生成用のチェックポイント内に存在しない（欠落している）ことが原因です。チェックポイントにデータが存在しないため、メタデバイス上にアロケートされたパラメータが「実データを持たないまま放置」され、最後のGPUデバイス移動処理 (.to("cuda")) のタイミングで値がないためにコピーできず、PyTorchが例外を投げます。

2. 解決策：`accelerate.init_empty_weights` に対するモンキーパッチ

このエラーは low_cpu_mem_usage=False に設定すれば回避できるように思えますが、Diffusersのバリデーションによって「量子化モデルのロード時は low_cpu_mem_usage=True が必須」と決められており、強制終了してしまいます。

そこで本検証では、accelerate.init_empty_weights コンテキストマネージャを動的にダミー化するモンキーパッチを適用しました。これにより、low_cpu_mem_usage=True の指定は通しつつも、内部アロケーション時にメタデバイスではなくCPUメモリ上に直接実体を確保させます。チェックポイントに存在しないパラメータは、PyTorchのデフォルト（ランダム初期化やゼロ初期化）で実データとして生成されるため、最後のGPU転送時にも「実データが存在する状態」となり、エラーを回避できます。

import contextlib
import accelerate

# accelerate.init_empty_weights を何もしないダミーコンテキストに差し替える
@contextlib.contextmanager
def dummy_init_empty_weights(*args, **kwargs):
    yield

accelerate.init_empty_weights = dummy_init_empty_weights

この数行の記述を、Diffusersのライブラリやパイプラインをインポートする前に実行するだけで、ロード時のメタデバイス問題を綺麗に解消できます。ロード後は、pipe.enable_model_cpu_offload() を適用することで、必要なモジュールだけを順次GPUへロードしながら推論を行い、VRAM使用量を10GB以下に抑制します。

3. テキストから画像生成（Text-to-Image）の実装と検証

モンキーパッチを組み込み、512x512解像度の画像を生成する検証スクリプトを作成しました。VRAM 12GBの範囲に収めるため、num_frames=1 を指定して画像生成モードとして走らせています。

# test_cosmos3.py
import torch
import contextlib
import time
import os

# モンキーパッチの適用
import accelerate
@contextlib.contextmanager
def dummy_init_empty_weights(*args, **kwargs):
    yield
accelerate.init_empty_weights = dummy_init_empty_weights

from diffusers import Cosmos3OmniPipeline

def main():
    model_id = "SanDiegoDude/Cosmos3-Nano-nf4"
    prompt = "A professional cinematic shot of a cute robotic dog sitting in a lush green park, looking curious, sunny day, shallow depth of field, 8k resolution."
    
    print("Loading pipeline components...")
    pipe = Cosmos3OmniPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        enable_safety_checker=False,
        low_cpu_mem_usage=True,
        device_map=None,
    )
    
    print("Enabling model CPU offload to stay within 12GB VRAM...")
    pipe.enable_model_cpu_offload()
    
    print("Generating image...")
    result = pipe(
        prompt,
        height=512,
        width=512,
        num_frames=1,
        num_inference_steps=20,
        guidance_scale=4.0,
    )
    
    # 結果の保存 (単一フレームのImageか確認して保存)
    if hasattr(result, "video") and result.video is not None:
        frames = result.video[0]
        output_path = "output_robot_dog.png"
        if isinstance(frames, list):
            frames[0].save(output_path)
        else:
            frames.save(output_path)
        print(f"Saved to: {output_path}")

if __name__ == "__main__":
    main()

検証結果の画像

実行時、モデルのロードに約115秒かかりましたが、画像推論処理（20ステップ）は順調に進行しました。生成された画像（512x512）は以下の通りです。質感やディテールも良好であり、モデルの破綻なく生成されています。

Generated Robot Dog
実際に12GB VRAM環境下で生成されたロボット犬の画像

4. テキストから動画生成（Text-to-Video）の実装と検証

次に、本質的な動画生成モデルとしての性能を検証するため、フレーム数を9に引き上げ、動画ファイル（mp4）への書き出し検証を行いました。Cosmos 3の動画アテンション構造は8n+1フレーム（9, 17, 25...）の単位で処理されます。

動画エクスポート時の注意点として、Diffusersの result.video はバッチサイズ1の場合、要素がフレーム数分の PIL.Image になっている平坦なリスト（List[PIL.Image.Image]）として返されます。また、エクスポート処理には imageio-ffmpeg が必要となるため、あらかじめ仮想環境にインストールしておく必要があります。

# test_cosmos3_video.py の一部
# (モデルロードおよびモンキーパッチ部分は共通)
        
        # 9フレームの動画生成を指示
        result = pipe(
            prompt,
            height=512,
            width=512,
            num_frames=9,
            num_inference_steps=20,
            guidance_scale=4.0,
        )
        
        # 保存処理
        if hasattr(result, "video") and result.video is not None:
            # result.video 自体が 9枚の Image が入ったリスト
            frames = result.video
            output_path = "output_robot_dog.mp4"
            
            from diffusers.utils import export_to_video
            export_to_video(frames, output_path, fps=8)
            print(f"Saved generated video to: {output_path}")

推論パフォーマンスの最適化

動画生成実行時のコンソールログを確認すると、最初のロード（ディスクからの読み込み）こそ1.5分程度を要するものの、デノイジングステップが始まると、CUDAグラフのコンパイルやVRAMの配置が最適化されたため、1ステップあたり約1.1秒（計約26秒）で動画全体の推論が完了しました。3D VAEデコード処理を含めてもロード完了から約36秒でmp4動画の出力に成功しました。

メモリ消費の最適化（gc.collect() と torch.cuda.empty_cache()）を随所に挿入したことで、他ツールがバックグラウンドで起動している状態であっても、VRAMがオーバーフローすることなく安定動作することを確認しました。

まとめ

今回の検証により、NVIDIA Cosmos 3（NF4量子化版）を12GB VRAMという限定されたGPU環境でも実用的な速度で動作させることが可能であることが確認されました。

メタデバイス回避: チェックポイント欠落パラメータに起因する NotImplementedError を accelerate.init_empty_weights のモンキーパッチにより回避。
VRAM制御: enable_model_cpu_offload() の活用によりピークVRAMを10GB前後に抑え込み。
高速推論: キャッシュとコンパイル最適化により、T2V（9フレーム、512x512）の生成ステップ処理を平均1.1秒/itという実用レベルの速度で達成。

DiffusersとAccelerateの組み合わせにおけるエッジケースのバグであるため、将来的に公式アップデートで修正される可能性はありますが、現状の量子化モデルを即座に動かしたい場合にはこのモンキーパッチ手法が有効なアプローチとなります。