大規模言語モデル（LLM）を自前の環境で動かす際、最も頭を悩ませるのがメモリ不足の問題ではないでしょうか。特に gpt-oss のようなオープンソース実装では、モデルサイズが数 GB から数十 GB に及ぶため、VRAM（GPU メモリ）や RAM が枯渇して OOM（Out Of Memory） エラーが発生することが珍しくありません。

本記事では、gpt-oss でモデルを実行する際に直面する OOM／VRAM 枯渇の原因を深掘りし、モデル分割、ページング（オフロード）、バッチ制御 という 3 つの解決策を具体的なコード例とともに解説します。これらのテクニックを活用すれば、限られたハードウェアリソースでも大規模モデルを安定稼働させることが可能になるでしょう。

背景

LLM とメモリの関係

大規模言語モデルは、数億から数千億のパラメータを持つニューラルネットワークです。これらのパラメータは学習済みの重み（weights）として保存され、推論時には GPU の VRAM や CPU の RAM に展開されます。

モデルのメモリ使用量を理解するために、まず基本的な計算式を見てみましょう。

typescriptCopy code/**
 * モデルのメモリ使用量計算
 * パラメータ数とデータ型から必要なメモリを推定
 */
interface ModelMemoryEstimate {
  parameters: number; // パラメータ数
  precision: string; // データ型（float32, float16, int8 など）
  memoryGB: number; // 必要なメモリ（GB）
}

typescriptCopy code/**
 * データ型ごとのバイト数
 */
const BYTES_PER_PARAMETER: Record<string, number> = {
  float32: 4, // 32ビット浮動小数点
  float16: 2, // 16ビット浮動小数点
  int8: 1, // 8ビット整数
  int4: 0.5, // 4ビット整数
};

typescriptCopy code/**
 * モデルサイズを計算する関数
 */
function calculateModelMemory(
  parameters: number,
  precision: string = 'float32'
): ModelMemoryEstimate {
  const bytesPerParam = BYTES_PER_PARAMETER[precision];
  const totalBytes = parameters * bytesPerParam;
  const memoryGB = totalBytes / 1024 ** 3; // GB に変換

  return {
    parameters,
    precision,
    memoryGB,
  };
}

たとえば、7B（70 億）パラメータのモデルを float32 で読み込むと約 28GB、float16 でも約 14GB のメモリが必要になります。

以下の図は、モデルサイズとメモリ使用量の関係を示しています。

mermaidCopy codeflowchart TB
  model["モデルファイル<br/>(数GB〜数十GB)"] --> load["メモリへ読み込み"]
  load --> vram["VRAM<br/>(GPU メモリ)"]
  load --> ram["RAM<br/>(システムメモリ)"]
  vram --> inference["推論処理"]
  ram --> swap["スワップ領域<br/>(ディスク)"]

  style vram fill:#ff9999
  style ram fill:#ffcc99
  style swap fill:#ffff99

gpt-oss の特徴とメモリ要件

gpt-oss は、GPT 系モデルを PyTorch や Transformers ライブラリを使って実装したオープンソースプロジェクトです。柔軟性が高く学習にも利用できる反面、メモリ管理は開発者側で適切に行う必要があります。

一般的な構成では以下のようなメモリが消費されます。

#	項目	メモリ消費	説明
1	モデル重み	最大	パラメータ数 × データ型のバイト数
2	アクティベーション	大	推論中の中間層出力（バッチサイズに比例）
3	KV キャッシュ	中	Attention 機構で使用するキャッシュ
4	オプティマイザ	大（学習時のみ）	Adam などの状態保存
5	システムオーバーヘッド	小	PyTorch やドライバの管理領域

このように、モデル本体だけでなく推論中の中間データも大きなメモリを消費するため、スペックギリギリの環境では OOM が発生しやすくなるのです。

課題

OOM エラーの発生パターン

gpt-oss で OOM エラーが発生する典型的なパターンを見ていきましょう。

パターン 1：モデル読み込み時の OOM

モデルファイルを読み込む段階で VRAM が不足し、以下のようなエラーが発生します。

sqlCopy codeRuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 8.00 GiB total capacity; 6.50 GiB already allocated; 1.20 GiB free; 6.80 GiB reserved in total by PyTorch)

エラーコード: RuntimeError: CUDA out of memory

発生条件:

• GPU の VRAM がモデルサイズより小さい
• 他のプロセスが VRAM を占有している
• float32 など高精度でモデルを読み込んでいる

パターン 2：推論実行時の OOM

モデルは読み込めたものの、実際に推論を実行すると VRAM が不足するケースです。

sqlCopy codetorch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 8.00 GiB total capacity; 7.20 GiB already allocated; 420.00 MiB free)

エラーコード: torch.cuda.OutOfMemoryError

発生条件:

• バッチサイズが大きすぎる
• シーケンス長が長い（KV キャッシュが肥大化）
• 複数リクエストの同時処理

パターン 3：メモリリークによる段階的枯渇

連続実行していると徐々にメモリが増え続け、最終的に OOM に至るパターンです。

pythonCopy codeMemoryError: Unable to allocate array with shape (1024, 1024, 4096) and data type float32

エラーコード: MemoryError

発生条件:

• キャッシュのクリアが不十分
• GPU メモリの断片化
• 参照が残り続けて GC が効かない

以下の図は、OOM が発生するまでのメモリ使用量の推移を示しています。

mermaidCopy codeflowchart LR
  start["起動"] --> load["モデル読み込み<br/>6GB 消費"]
  load --> req1["リクエスト 1<br/>+1GB"]
  req1 --> req2["リクエスト 2<br/>+1GB"]
  req2 --> oom["OOM エラー<br/>8GB 超過"]

  style oom fill:#ff6666

パフォーマンスへの影響

OOM を避けるために小さいモデルやバッチサイズを選択すると、今度はパフォーマンスが犠牲になります。

#	対策	メモリ削減効果	パフォーマンス影響
1	モデルサイズ削減	★★★	精度低下
2	バッチサイズ削減	★★	スループット低下
3	シーケンス長制限	★	長文対応不可
4	CPU オフロード	★★	レイテンシ増加

このジレンマを解決するためには、メモリ効率とパフォーマンスを両立させる工夫が必要です。

解決策

解決策 1：モデル分割（Model Parallelism）

モデル分割は、大きなモデルを複数の GPU やデバイスに分散配置する手法です。1 つの GPU に収まらないモデルでも、複数 GPU で分担すれば実行可能になります。

レイヤー単位の分割

最も基本的な分割方法は、Transformer の各レイヤーを異なる GPU に配置する方法です。

typescriptCopy code/**
 * モデル分割の設定インターフェース
 */
interface ModelParallelConfig {
  numGPUs: number; // 使用する GPU 数
  layersPerGPU: number[]; // 各 GPU に配置するレイヤー数
  deviceMap?: Record<string, number>; // レイヤー名 → GPU ID のマッピング
}

typescriptCopy code/**
 * Hugging Face Transformers での自動分割設定例
 */
const autoConfig: ModelParallelConfig = {
  numGPUs: 2,
  layersPerGPU: [16, 16], // 32 層を 2 つの GPU で均等分割
  deviceMap: 'auto', // 自動で最適配置
};

実際の Python コードでは以下のように実装します。

pythonCopy code"""
モデルを複数 GPU に分割して読み込む
Hugging Face の device_map を使用
"""
from transformers import AutoModelForCausalLM
import torch

# デバイスマップを指定してモデルを読み込み
model = AutoModelForCausalLM.from_pretrained(
    "gpt-oss-7b",
    device_map="auto",          # 自動でレイヤーを分散配置
    torch_dtype=torch.float16,  # メモリ節約のため float16 を使用
    low_cpu_mem_usage=True      # CPU メモリ使用量も削減
)

pythonCopy code"""
手動でデバイスマップを指定する場合
特定のレイヤーを特定の GPU に配置
"""
device_map = {
    "transformer.wte": 0,              # Embedding を GPU 0 へ
    "transformer.h.0": 0,              # レイヤー 0 を GPU 0 へ
    "transformer.h.1": 0,              # レイヤー 1 を GPU 0 へ
    # ... 中略 ...
    "transformer.h.16": 1,             # レイヤー 16 を GPU 1 へ
    "transformer.h.17": 1,             # レイヤー 17 を GPU 1 へ
    # ... 中略 ...
    "transformer.ln_f": 1,             # 最終レイヤーを GPU 1 へ
    "lm_head": 1                       # 出力層を GPU 1 へ
}

model = AutoModelForCausalLM.from_pretrained(
    "gpt-oss-7b",
    device_map=device_map,
    torch_dtype=torch.float16
)

以下の図は、モデル分割の概念を示しています。

mermaidCopy codeflowchart TB
  input["入力トークン"] --> emb["Embedding 層"]
  emb --> gpu0["GPU 0<br/>Layer 0-15"]
  gpu0 --> transfer["層間データ転送"]
  transfer --> gpu1["GPU 1<br/>Layer 16-31"]
  gpu1 --> output["出力層"]
  output --> result["生成結果"]

  style gpu0 fill:#99ccff
  style gpu1 fill:#99ffcc
  style transfer fill:#ffcc99

Tensor Parallelism（テンソル並列化）

より高度な手法として、1 つのレイヤー内の行列演算を複数 GPU で分割する Tensor Parallelism があります。

pythonCopy code"""
Megatron-LM スタイルのテンソル並列化
行列演算を列方向に分割
"""
import torch.distributed as dist

# 2 GPU でテンソル並列化を初期化
dist.init_process_group(backend="nccl", world_size=2)

pythonCopy code"""
線形層のテンソル分割実装例
weight 行列を列方向に分割
"""
class ColumnParallelLinear(torch.nn.Module):
    def __init__(self, input_size, output_size, world_size):
        super().__init__()
        # output_size を world_size で分割
        self.output_size_per_partition = output_size // world_size
        # 分割された重み行列
        self.weight = torch.nn.Parameter(
            torch.empty(self.output_size_per_partition, input_size)
        )

    def forward(self, x):
        # 各 GPU で部分的な行列積を計算
        output_parallel = torch.matmul(x, self.weight.t())
        return output_parallel

pythonCopy code"""
分割結果を集約する AllGather 操作
"""
def gather_outputs(output_parallel, world_size):
    # 全 GPU の出力を収集
    output_list = [torch.empty_like(output_parallel)
                   for _ in range(world_size)]
    dist.all_gather(output_list, output_parallel)
    # 連結して完全な出力を復元
    output = torch.cat(output_list, dim=-1)
    return output

解決策 2：ページング（CPU オフロード）

VRAM が不足する場合、使用頻度の低いレイヤーを CPU の RAM にオフロードし、必要なときだけ GPU に転送する手法が有効です。

基本的なオフロード設定

Accelerate ライブラリを使えば、簡単にオフロードを設定できます。

pythonCopy code"""
Accelerate を使った CPU オフロード
VRAM 不足時に自動で CPU へ退避
"""
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

# 空のモデルを初期化（メモリ消費なし）
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

pythonCopy code"""
チェックポイントを読み込みながらデバイスに配置
max_memory で各デバイスの上限を指定
"""
model = load_checkpoint_and_dispatch(
    model,
    checkpoint="gpt-oss-7b",
    device_map="auto",
    max_memory={
        0: "6GiB",      # GPU 0 は 6GB まで使用
        "cpu": "20GiB"  # CPU は 20GB まで使用
    },
    offload_folder="offload"  # ディスクキャッシュ用フォルダ
)

以下の図は、ページングの動作を示しています。

mermaidCopy codesequenceDiagram
  participant CPU as CPU RAM
  participant GPU as GPU VRAM
  participant Model as モデルレイヤー

  Note over GPU: VRAM 空き容量 2GB
  Model->>CPU: Layer 0-10 を CPU へ
  Model->>GPU: Layer 11-20 を GPU へ

  Note over Model: 推論開始
  CPU->>GPU: Layer 5 を転送
  GPU->>GPU: Layer 5 で計算
  GPU->>CPU: Layer 5 を退避

  CPU->>GPU: Layer 15 を転送
  GPU->>GPU: Layer 15 で計算

オフロードフックのカスタマイズ

より細かい制御が必要な場合、フックを使ってレイヤー単位でオフロードタイミングを制御できます。

pythonCopy code"""
カスタムオフロードフックの実装
forward 後に自動で CPU へ移動
"""
def offload_hook(module, input, output):
    # 計算後すぐに CPU へ移動
    module.to("cpu")
    # 出力は GPU に残す
    return output.to("cuda")

pythonCopy code"""
特定のレイヤーにフックを登録
"""
for i, layer in enumerate(model.transformer.h):
    if i < 10:  # 最初の 10 層のみオフロード
        layer.register_forward_hook(offload_hook)

pythonCopy code"""
推論前に必要なレイヤーを GPU へプリロード
"""
def preload_layers(model, layer_indices):
    for i in layer_indices:
        model.transformer.h[i].to("cuda")

# 次に使うレイヤーを事前に転送
preload_layers(model, [5, 6, 7])

解決策 3：バッチ制御とメモリ管理

推論時のバッチサイズや KV キャッシュを適切に管理することで、限られた VRAM を効率的に使用できます。

動的バッチサイズ調整

利用可能な VRAM に応じて、バッチサイズを動的に調整します。

typescriptCopy code/**
 * バッチ制御の設定
 */
interface BatchConfig {
  maxBatchSize: number; // 最大バッチサイズ
  dynamicBatching: boolean; // 動的調整の有効化
  memoryThreshold: number; // メモリ使用率の閾値（%）
}

pythonCopy code"""
VRAM 使用量を監視する関数
"""
import torch

def get_gpu_memory_usage():
    # 現在の VRAM 使用量を取得
    allocated = torch.cuda.memory_allocated() / (1024 ** 3)  # GB 単位
    reserved = torch.cuda.memory_reserved() / (1024 ** 3)
    total = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)

    usage_percent = (allocated / total) * 100

    return {
        'allocated': allocated,
        'reserved': reserved,
        'total': total,
        'usage_percent': usage_percent
    }

pythonCopy code"""
動的バッチサイズ計算
VRAM 使用率に応じて自動調整
"""
def calculate_batch_size(max_batch_size=32, threshold=80):
    memory_info = get_gpu_memory_usage()
    usage = memory_info['usage_percent']

    if usage > threshold:
        # 使用率が高い場合はバッチサイズを削減
        new_batch_size = max(1, max_batch_size // 2)
    else:
        # 余裕がある場合は最大サイズを使用
        new_batch_size = max_batch_size

    return new_batch_size

pythonCopy code"""
バッチ処理の実装例
メモリを監視しながら推論
"""
def process_with_dynamic_batching(inputs, model, max_batch_size=32):
    results = []

    for i in range(0, len(inputs), max_batch_size):
        # 現在のバッチサイズを計算
        batch_size = calculate_batch_size(max_batch_size)
        batch = inputs[i:i+batch_size]

        # 推論実行
        with torch.no_grad():
            outputs = model.generate(
                batch,
                max_length=100,
                do_sample=True
            )

        results.extend(outputs)

        # メモリをクリア
        torch.cuda.empty_cache()

    return results

KV キャッシュの最適化

Attention 機構で使用する KV キャッシュは、シーケンス長に比例して増大します。これを効率化する方法を見ていきましょう。

pythonCopy code"""
KV キャッシュサイズの計算
"""
def calculate_kv_cache_size(
    num_layers: int,
    hidden_size: int,
    num_heads: int,
    sequence_length: int,
    batch_size: int,
    dtype_bytes: int = 2  # float16
):
    # Key と Value それぞれのサイズ
    kv_size = 2 * num_layers * hidden_size * sequence_length * batch_size * dtype_bytes
    kv_size_gb = kv_size / (1024 ** 3)

    return kv_size_gb

pythonCopy code"""
使用例：7B モデルの KV キャッシュサイズを計算
"""
cache_size = calculate_kv_cache_size(
    num_layers=32,
    hidden_size=4096,
    num_heads=32,
    sequence_length=2048,
    batch_size=4
)
print(f"KV Cache size: {cache_size:.2f} GB")
# 出力：KV Cache size: 2.00 GB

pythonCopy code"""
KV キャッシュを制限する設定
過去のトークンを一部削除
"""
generation_config = {
    "max_length": 2048,
    "use_cache": True,           # キャッシュを有効化
    "cache_implementation": "dynamic",  # 動的キャッシュ
    "cache_config": {
        "max_cache_len": 1024,   # キャッシュの最大長
        "eviction_policy": "fifo"  # 古いものから削除
    }
}

outputs = model.generate(
    inputs,
    **generation_config
)

メモリリーク防止

推論後のメモリ解放を確実に行い、リークを防ぎます。

pythonCopy code"""
推論後のメモリクリーンアップ
"""
def inference_with_cleanup(model, inputs):
    try:
        # 推論実行
        with torch.no_grad():
            outputs = model.generate(inputs)

        return outputs

    finally:
        # 確実にメモリを解放
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()

pythonCopy code"""
定期的なガベージコレクション
長時間実行時にメモリ断片化を防ぐ
"""
import gc

def periodic_cleanup(interval=10):
    """interval 回の推論ごとにクリーンアップ"""
    counter = 0

    def cleanup_if_needed():
        nonlocal counter
        counter += 1

        if counter >= interval:
            gc.collect()
            torch.cuda.empty_cache()
            counter = 0

    return cleanup_if_needed

pythonCopy code"""
使用例
"""
cleanup = periodic_cleanup(interval=10)

for batch in data_loader:
    outputs = model.generate(batch)
    # ... 結果を処理 ...

    # 定期的にクリーンアップ
    cleanup()

以下の図は、バッチ制御とメモリ管理の流れを示しています。

mermaidCopy codeflowchart TD
  start["推論リクエスト"] --> check["VRAM 使用率チェック"]
  check -->|80%以上| reduce["バッチサイズ削減"]
  check -->|80%未満| normal["通常バッチサイズ"]

  reduce --> infer["推論実行"]
  normal --> infer

  infer --> cleanup["メモリクリーンアップ"]
  cleanup --> gc["GC 実行"]
  gc --> cache["CUDA キャッシュクリア"]
  cache --> next["次のバッチ"]

  style check fill:#ffffcc
  style cleanup fill:#ccffcc

具体例

実践例 1：8GB VRAM で 7B モデルを動かす

一般的なゲーミング GPU（RTX 3070 など）の 8GB VRAM で 7B パラメータモデルを動かす実装例です。

環境セットアップ

bashCopy code# 必要なライブラリをインストール
yarn add transformers torch accelerate bitsandbytes

pythonCopy code"""
ライブラリのインポート
"""
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig
)
import torch

量子化設定

8bit 量子化を使ってメモリ使用量を削減します。

pythonCopy code"""
8bit 量子化の設定
bitsandbytes を使用してメモリを 1/4 に削減
"""
quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,              # 8bit 量子化を有効化
    llm_int8_threshold=6.0,         # 量子化の閾値
    llm_int8_has_fp16_weight=False  # メモリ節約モード
)

pythonCopy code"""
モデルとトークナイザーの読み込み
"""
model_name = "gpt-oss-7b"

# トークナイザー
tokenizer = AutoTokenizer.from_pretrained(model_name)

# モデル（8bit 量子化 + CPU オフロード）
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto",               # 自動配置
    max_memory={
        0: "7GiB",                   # GPU 0 に 7GB 割り当て
        "cpu": "16GiB"               # CPU に 16GB 割り当て
    },
    offload_folder="offload_cache",  # オフロード用キャッシュ
    offload_state_dict=True          # 状態辞書もオフロード
)

推論実行

pythonCopy code"""
メモリ効率的な推論関数
"""
def generate_text(prompt, max_length=100):
    # トークン化
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=512
    ).to("cuda")

    # メモリ使用量を表示
    print(f"Memory before: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

    # 推論実行
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )

    # メモリ使用量を表示
    print(f"Memory after: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

    # デコード
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # メモリクリーンアップ
    del inputs, outputs
    torch.cuda.empty_cache()

    return generated_text

pythonCopy code"""
実行例
"""
prompt = "大規模言語モデルとは"
result = generate_text(prompt, max_length=200)
print(result)

# 出力例：
# Memory before: 6.82 GB
# Memory after: 7.15 GB
# 大規模言語モデルとは、膨大なテキストデータを学習した...

この構成により、7B モデルが約 7GB の VRAM で動作し、1GB の余裕を確保できています。

実践例 2：複数 GPU での並列推論

2 つの GPU（各 8GB）を使って 13B モデルを実行する例です。

マルチ GPU セットアップ

pythonCopy code"""
複数 GPU の確認
"""
import torch

num_gpus = torch.cuda.device_count()
print(f"Available GPUs: {num_gpus}")

for i in range(num_gpus):
    props = torch.cuda.get_device_properties(i)
    print(f"GPU {i}: {props.name}, {props.total_memory / 1e9:.2f} GB")

モデル分割設定

pythonCopy code"""
13B モデルを 2 GPU で分割
各 GPU に約半分のレイヤーを配置
"""
model_name = "gpt-oss-13b"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="balanced",           # バランス配置
    torch_dtype=torch.float16,       # float16 で容量削減
    max_memory={
        0: "7GiB",                   # GPU 0
        1: "7GiB"                    # GPU 1
    },
    low_cpu_mem_usage=True
)

pythonCopy code"""
配置状況の確認
"""
def print_device_map(model):
    for name, module in model.named_modules():
        if hasattr(module, 'weight'):
            device = module.weight.device
            print(f"{name}: {device}")

print_device_map(model)
# 出力例：
# transformer.wte: cuda:0
# transformer.h.0: cuda:0
# ...
# transformer.h.20: cuda:1
# transformer.h.39: cuda:1

パイプライン並列推論

pythonCopy code"""
パイプライン並列で複数リクエストを処理
"""
from torch.nn.parallel import DataParallel

def pipeline_inference(prompts, model, tokenizer):
    results = []

    for prompt in prompts:
        # トークン化（自動で適切な GPU へ配置）
        inputs = tokenizer(prompt, return_tensors="pt")

        # 推論（レイヤー間で自動転送）
        with torch.no_grad():
            outputs = model.generate(
                input_ids=inputs.input_ids,
                max_length=150,
                num_return_sequences=1
            )

        # デコード
        text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        results.append(text)

        # GPU 間の同期
        torch.cuda.synchronize()

    return results

pythonCopy code"""
実行例
"""
prompts = [
    "人工知能の未来について",
    "量子コンピュータの原理は",
    "機械学習の応用例を教えて"
]

results = pipeline_inference(prompts, model, tokenizer)

for i, result in enumerate(results):
    print(f"\n=== Result {i+1} ===")
    print(result)

以下の図は、マルチ GPU での処理フローを示しています。

mermaidCopy codeflowchart LR
  input["入力"] --> tok["トークン化"]
  tok --> gpu0["GPU 0<br/>Layer 0-19<br/>4GB 使用"]
  gpu0 -->|中間データ| transfer["GPU 間転送<br/>PCIe"]
  transfer --> gpu1["GPU 1<br/>Layer 20-39<br/>4GB 使用"]
  gpu1 --> decode["デコード"]
  decode --> output["出力"]

  style gpu0 fill:#99ccff
  style gpu1 fill:#99ffcc
  style transfer fill:#ffcc99

実践例 3：バッチ処理の最適化

大量のテキストを効率的に処理するバッチ最適化の例です。

適応的バッチサイズ

pythonCopy code"""
適応的バッチ処理クラス
VRAM 使用率に応じてバッチサイズを自動調整
"""
class AdaptiveBatchProcessor:
    def __init__(self, model, tokenizer, max_batch_size=32):
        self.model = model
        self.tokenizer = tokenizer
        self.max_batch_size = max_batch_size
        self.current_batch_size = max_batch_size

    def get_optimal_batch_size(self):
        """VRAM 使用率からバッチサイズを計算"""
        if not torch.cuda.is_available():
            return 1

        allocated = torch.cuda.memory_allocated() / 1e9
        total = torch.cuda.get_device_properties(0).total_memory / 1e9
        usage_ratio = allocated / total

        if usage_ratio > 0.85:
            # 使用率 85% 超：半減
            self.current_batch_size = max(1, self.current_batch_size // 2)
        elif usage_ratio < 0.60:
            # 使用率 60% 未満：増加
            self.current_batch_size = min(
                self.max_batch_size,
                self.current_batch_size * 2
            )

        return self.current_batch_size

pythonCopy code"""
バッチ処理メソッド
"""
def process_batch(self, texts):
    """テキストのリストをバッチ処理"""
    results = []
    i = 0

    while i < len(texts):
        # 現在の最適バッチサイズを取得
        batch_size = self.get_optimal_batch_size()
        batch = texts[i:i+batch_size]

        # トークン化
        inputs = self.tokenizer(
            batch,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512
        ).to("cuda")

        try:
            # 推論実行
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_length=100,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )

            # デコード
            batch_results = [
                self.tokenizer.decode(out, skip_special_tokens=True)
                for out in outputs
            ]
            results.extend(batch_results)

            i += batch_size

        except RuntimeError as e:
            if "out of memory" in str(e):
                # OOM 発生時はバッチサイズを削減
                print(f"OOM at batch size {batch_size}, reducing...")
                self.current_batch_size = max(1, batch_size // 2)
                torch.cuda.empty_cache()
                continue
            else:
                raise

        finally:
            # メモリクリーンアップ
            del inputs, outputs
            torch.cuda.empty_cache()

    return results

使用例

pythonCopy code"""
実行例
"""
# プロセッサの初期化
processor = AdaptiveBatchProcessor(
    model=model,
    tokenizer=tokenizer,
    max_batch_size=16
)

# 大量のテキストを処理
texts = [
    f"これはテスト文章 {i} です。"
    for i in range(100)
]

results = processor.process_batch(texts)

print(f"Processed {len(results)} texts")
print(f"Final batch size: {processor.current_batch_size}")

この実装により、VRAM の状況に応じて自動的にバッチサイズが調整され、OOM を回避しながら最大限のスループットを実現できます。

まとめ

本記事では、gpt-oss で発生する OOM／VRAM 枯渇の問題に対する 3 つの解決策を詳しく解説しました。

解決策のまとめ:

#	解決策	適用場面	メモリ削減効果	実装難易度
1	モデル分割	複数 GPU 環境	★★★	中
2	ページング	単一 GPU、VRAM 不足	★★	低
3	バッチ制御	推論時の安定化	★	低

重要なポイント:

まず、モデル分割は複数の GPU を活用して大規模モデルを分散配置する手法です。レイヤー単位の分割やテンソル並列化により、単一 GPU では扱えないモデルサイズにも対応できます。

次に、ページング（CPU オフロード）は、使用頻度の低いレイヤーを CPU の RAM に退避させることで、限られた VRAM を効率的に使用する方法です。Accelerate ライブラリを活用すれば、わずか数行のコードで実装できます。

最後に、バッチ制御とメモリ管理では、動的なバッチサイズ調整や KV キャッシュの最適化により、推論時の安定性を向上させます。定期的なメモリクリーンアップも忘れずに実施しましょう。

実践的な組み合わせ:

実際のプロジェクトでは、これらの手法を組み合わせることで最大の効果が得られます。たとえば、8GB VRAM の環境で 13B モデルを動かす場合、8bit 量子化 + CPU オフロード + 動的バッチサイズの 3 つを併用することで、安定した推論が可能になるでしょう。

パフォーマンスとのバランス:

メモリ削減とパフォーマンスはトレードオフの関係にあります。CPU オフロードを多用するとレイテンシが増加し、量子化は精度にわずかな影響を与える可能性があります。運用環境の要件に応じて、適切なバランスを見つけることが重要です。

これらのテクニックを活用すれば、限られたハードウェアリソースでも大規模言語モデルを実用的に運用できます。ぜひ、皆さんのプロジェクトでも試してみてください。

関連リンク

• Hugging Face Transformers - Model Parallelism
• Accelerate - Big Model Inference
• PyTorch Distributed Training
• bitsandbytes - 8bit Quantization
• NVIDIA CUDA Best Practices Guide
• Megatron-LM: Training Multi-Billion Parameter Language Models
• DeepSpeed: Extreme-scale model training