NVIDIA GPU

NVIDIA GPU は CUDA を通じて PyTorch を実行し、デバイス文字列は cuda です。Transformers におけるほとんどのアテンションカーネルやサービング最適化は CUDA をサポートしており、FlashAttention-3 や FlashAttention-4 のように CUDA でのみ実行されるものもあります。

インストール

GPU 用の CUDA ドライバーをインストールし、その後 PyTorch と Transformers をインストールしてください。NVIDIA Spark (ARM64) デバイスは NVIDIA のパッケージインデックスから PyTorch をインストールする必要があるため、インストールガイドを参照してください。

pip install "transformers[torch]"

PyTorchがGPUを認識できることを確認します。

import torch

print(torch.cuda.is_available(), torch.cuda.device_count())

True と GPU の数が出力されます。例えば True 4 のようにです。

GPUの選択と順序付け

CUDA_VISIBLE_DEVICES 環境変数を使用して、使用する GPU を選択し、その順序を設定します。エクスポートした変数を忘れてしまうと、後続の実行が意図せず誤った GPU に割り当てられてしまう可能性があるため、実行コマンドと同じコマンドライン上で設定してください。

例えば、4つのGPUから0番と2番を選択するには:

CUDA_VISIBLE_DEVICES=0,2 torchrun trainer-program.py ...

PyTorch は GPU 0 と 2 のみを認識し、それぞれ cuda:0 と cuda:1 にマッピングされます。順序を逆にする（GPU 2 を cuda:0 として、GPU 0 を cuda:1 として使用する）には:

CUDA_VISIBLE_DEVICES=2,0 torchrun trainer-program.py ...

GPUなしで実行するには:

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

CUDA はデフォルトで GPU を最速の順に番号付けしますが、nvidia-smi の表示と一致しない場合があります。CUDA_VISIBLE_DEVICES のインデックスが nvidia-smi と一致するように、CUDA_DEVICE_ORDER=PCI_BUS_ID を設定してください。

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0,2 torchrun trainer-program.py ...

アテンションバックエンドを選択

モデルを読み込む際にattn_implementationでアテンションバックエンドを設定し、推論と学習の両方に適用されます。モデルはデフォルトでSDPAを使用します。NVIDIA GPUは3つのFlashAttentionバージョンをサポートしています。FlashAttentionパッケージがインストールされていない場合、Transformersは代わりにHubから事前構築されたカーネルを読み込みますが、これにはkernelsエクストラ（pip install "transformers[kernels]"）が必要です。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype="bfloat16",
    device_map="cuda",
    attn_implementation="flash_attention_2",
)

推論を実行する

device_map="auto" を from_pretrained() に渡すことで、1つの GPU に収まらないモデルを利用可能な複数の GPU に分割して配置できます。単一の GPU に読み込む場合は device_map="cuda" を指定してください。device を指定しない場合、pipeline() は自動的に GPU を選択します。

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

inputs = tokenizer("NVIDIA GPUs are", return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_new_tokens=50, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Continuous batchingは、NVIDIA GPU 上で CUDA グラフと非同期バッチ処理をサポートし、それらを使用するかどうかを自動的に判断します。

Hubカーネルを使用する

Set use_kernels=True in from_pretrained() to replace supported layers with optimized kernels from the Hub. NVIDIA GPUs have the widest coverage, including activations, RMSNorm, rotary embeddings, MegaBlocks mixture-of-experts layers, and Mamba state-space ops. Install the kernels extra to use them.

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="cuda", use_kernels=True
)

Hub カーネルはトレーニングにも適用されます。model.train() または model.eval() を呼び出すと、そのモード用のカーネルに切り替わります。カーネルを自分で選択する場合は、KernelConfig を渡し、NVIDIA 固有のマッピングにはデバイスキーとして "cuda" を使用します。完全なマッピングについては、カーネルの読み込み を参照してください。

モデルの量子化

Quantizationにより、より大きなモデルをGPUメモリに収めることができます。TransformersのほとんどのQuantization手法はCUDAに対応しています。各手法を比較するには、Quantizationの概要のCUDA列を確認してください。

GPU メモリに収まらないモデルをファインチューニングするには、bitsandbytes を使用して 4-bit で読み込み、PEFT で LoRA アダプターを学習します。

Trainerで学習

TrainingArguments は NVIDIA GPU を検出し、デバイスを cuda に設定するため、Trainer はデバイス設定なしでそれら上で学習できます。Ampere 以降の GPU では bf16 混合精度を、それ以前の GPU では fp16 を使用してください。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
)
print(training_args.device)

Trainerが単一GPUで学習を行う場合、cuda:0が出力されます。

分散学習

マルチGPUでの学習には、NCCLの集約型バックエンドを使用します。torchrunまたはaccelerate launchを使用して、GPUごとに1つのプロセスを起動してください。

torchrun --nproc_per_node=4 train.py

DDP、FSDP2、DeepSpeed ZeRO、Tensor parallelismを参照して、並列化戦略を選択してください。

次のステップ

• PyTorchがデバイス、ストリーム、メモリをどのように管理しているかについては、PyTorchのCUDAセマンティクスのメモを参照してください。
• トレーニング中にGPUメモリがどこに消費されるかを理解するには、GPUメモリ使用量を参照してください。
• attn_implementation がどのようにアテンション関数を選択するかを学ぶには、Attention backends をお読みください。