SentenceTransformersで始めるGradient Checkpointing：仕組みと実践のポイント

この記事は Uzabase Advent Calendar 2025 の 25 日目の記事です。

はじめに

ププロダクトチームでは、Speeda AI Agentを開発していますが、端々でベクトル検索を用いています。また、ベクトル検索するためのモデル開発や、そのホスティングも自前でおこなっています。

ベクトル検索用のモデルは社内ファインチューニングしているため、社内のベンチマークデータでは、Geminiのベクトルを精度で大きく上回っており、ユーザー体験に貢献していると考えています。また、基盤も安定しており、ニュースのような頻繁かつ断続的に追加されるデータも安定してさばけるような基盤を構築しています。

本シリーズではここに至るまでの以下の技術的な工夫を紹介して行きたいと思います。

• HPAでスケールできるインフラの構築
• vLLM化による速度と安定性の向上
• Sentence Transformersで省計算資源で学習を行う工夫

今回は上記のイントロにある「Sentence Transformersで省計算資源で学習を行う工夫」の1つとして、SentenceTransformersを使ったGradient Checkpointingでの学習方法を紹介します。

Gradient Checkpointingとは

深層学習では順伝播で得た中間アクティベーション（層ごとの出力）を逆伝播に使うためメモリに保持します。大規模モデル・長いシーケンス・大きいバッチでは、この保持がGPUメモリを圧迫してOOMになりがちです。
この問題を解決する手段の一つが、Gradient Checkpointingです。Gradient Checkpointingはアクティベーションを「保存せず破棄」します。その後、逆伝播時に必要箇所だけ「もう一度順伝播を再計算」して勾配を求めます。そのため、アクティベーションを保存するメモリを減らことが可能です。代わりに再計算の時間（計算量）が増えるアプローチになります。

仕組みのイメージ

ここで、深層学習のモデルのアーキテクチャーを以下だと考えます。

cybertronai/gradient-checkpointingより図を引用

• 通常の学習:
  • 順伝播で各層出力を全て保存 → 逆伝播で保存済み出力から勾配計算
    • メモリ消費は大きいが高速。
  • 以下の図の通り、順伝播が進むにつれて、紫色の部分がメモリに保存される。

cybertronai/gradient-checkpointingより図を引用

• Gradient Checkpointingあり:
  • ネットワークをブロックに分割し、ブロック境界（チェックポイント）のみ保存。内部の中間出力は捨てる → 逆伝播で必要範囲のみ再計算
    • メモリ消費は小さいが、その分遅くなる。
  • 以下の図のようにメモリのチェックポイントのみを保存することにより、あるノードを再計算する際にはそのノードとその直前のチェックポイントの間のノードのみを再計算すればよいことになる。

cybertronai/gradient-checkpointingより図を引用

この仕組みにより、中間値を2回再計算する計算コストは増加しますが、中間値を事前に格納するメモリコストを削減できます。結果として、より大きなモデルやバッチサイズでのトレーニングが可能になります。

メリット・デメリット

• メリット:
  • GPUメモリ使用量の削減。より大きなモデル・長い入力・大きいバッチサイズが扱いやすくなる。
  • メモリがボトルネックで学習が行えないケースの解消（OOMの回避）。
• デメリット:
  • 再計算コストで1ステップの処理時間が増え、スループットは低下しやすい。

適用が有効な場面

• モデル規模が大きくアクティベーションがメモリを圧迫する。
• 入力シーケンスが長い（NLPで長文、音声で長フレームなど）。
• バッチサイズを上げて学習を安定させたいが、現状メモリが足りない。

SentenceTransformersでGradient Checkpointingを使って学習する方法

まずは普通に学習してどの程度のメモリサイズが必要になるかを確認してみましょう。このコードは参考文献に載せたSentenceTransformersの公式のサンプルコードを利用しています。

from datasets import load_dataset

from sentence_transformers import SentenceTransformer, losses
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
from sentence_transformers.similarity_functions import SimilarityFunction
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import BatchSamplers, SentenceTransformerTrainingArguments

train_batch_size = 128
max_seq_length = 75
num_epochs = 1
model_name = "distilbert/distilroberta-base"
model = SentenceTransformer(model_name)

train_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="train").select(range(10000))
eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="dev").select(range(1000))

train_loss = losses.MultipleNegativesRankingLoss(model)
stsb_eval_dataset = load_dataset("sentence-transformers/stsb", split="validation")

args = SentenceTransformerTrainingArguments(
    output_dir=".",
    num_train_epochs=1,
    per_device_train_batch_size=train_batch_size,
    per_device_eval_batch_size=train_batch_size,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if you get an error that your GPU can't run on FP16
    bf16=False,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,
    eval_strategy="steps",
    eval_steps=10,
    save_strategy="steps",
    save_steps=10,
    save_total_limit=2,
    logging_steps=100,
)

dev_evaluator = EmbeddingSimilarityEvaluator(
    sentences1=stsb_eval_dataset["sentence1"][:],
    sentences2=stsb_eval_dataset["sentence2"][:],
    scores=stsb_eval_dataset["score"][:],
    main_similarity=SimilarityFunction.COSINE,
    name="sts-dev",
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=train_loss,
    evaluator=dev_evaluator,
)
trainer.train()

test_dataset = load_dataset("sentence-transformers/stsb", split="test")
test_evaluator = EmbeddingSimilarityEvaluator(
    sentences1=test_dataset["sentence1"][:],
    sentences2=test_dataset["sentence2"][:],
    scores=test_dataset["score"][:],
    main_similarity=SimilarityFunction.COSINE,
    name="sts-test",
)
test_evaluator(model)

次にSentenceTransformersでGradient Checkpointingを利用して学習してみましょう。
「SentenceTransformersでGradient Checkpointingを利用する場合、SentenceTransformerTrainingArgumentsにはgradient_checkpointing: bool = Falseという項目があります。これをgradient_checkpointing=Trueと設定することで利用できます。

先ほどのサンプルコードをSentenceTransformerTrainingArgumentsの部分を以下のように編集して実際に動作するかを確認してみます。

...
args = SentenceTransformerTrainingArguments(
    output_dir=".",
    num_train_epochs=1,
    per_device_train_batch_size=train_batch_size,
    per_device_eval_batch_size=train_batch_size,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if you get an error that your GPU can't run on FP16
    bf16=False,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,
    eval_strategy="steps",
    eval_steps=10,
    save_strategy="steps",
    save_steps=10,
    save_total_limit=2,
    logging_steps=100,
    gradient_checkpointing=True
)
...

ここまでの結果からGradient Checkpointingを利用した場合と利用しなかった場合の比較表を作ると以下になります。

項目	通常	Gradient Checkpointing	備考
GPUメモリ使用量	6296 MiB	3368 MiB	約-47%
train_runtime	65.60 s	55.95 s	-9.65 s
train_samples_per_second	152.43	178.72	+26.29
train_steps_per_second	1.204	1.412	+0.208
train_loss	1.7810	1.7809	ほぼ同等
eval_loss	1.5949	1.5937	ほぼ同等
STS Pearson	0.7645	0.7647	ほぼ同等
STS Spearman	0.7925	0.7926	ほぼ同等

なぜかGradient Checkpointingが走ることでかなり早くなっているのですが、評価値はそのままにGPUメモリの消費量は大幅に削減できたことは確認できたかと思います。
別で学習時のバッチサイズを12と10分の1くらいに小さくして実験してみたのですが、train_runtimeがGradient Checkpointing無しの場合は222秒だったのが、Gradient Checkpointingありにした時に250秒だったので、はっきりとはわかりませんが、バッチサイズが大きい場合は再計算によるコストよりもアクティベーションの保持や読み出しの方がコストがかかっていたため、学習が早くなったのかもしれません。

Gradient Checkpointing + LoRAを使いたいとき

以下の部分を追加してから実行すればそのままGradient Checkpointing + LoRAが使えます。

from peft import LoraConfig, TaskType
...
model = SentenceTransformer(model_name)

peft_config = LoraConfig(
    task_type=TaskType.FEATURE_EXTRACTION,
    inference_mode=False,
    r=64,
    lora_alpha=128,
    lora_dropout=0.1,
)
model.add_adapter(peft_config)

...

まとめ

• 要点: アクティベーションを保存せず必要時に再計算することでVRAMを節約する手法。計算量増とのトレードオフ。
• 効果: 実運用では30–50%程度のVRAM削減例が多い。帯域ボトルネックが強い条件では、再計算増よりI/O削減が勝ち速度が上がることもある。
• 使いどころ: OOM回避、大きなバッチや長いシーケンス、LoRA併用での大規模モデル微調整など。
• SentenceTransformersの有効化: TrainingArgumentsでgradient_checkpointing=True。Transformers側の一部モデルではuse_cache=Falseが必要な点に注意。
• 本記事の観測: 精度は同等のままVRAM約-47%。条件によってはtrain_runtimeが短縮された。

参考文献

• https://residentmario.github.io/pytorch-training-performance-guide/gradient-checkpoints.html
• https://github.com/huggingface/sentence-transformers/blob/main/examples/sentence_transformer/training/nli/training_nli_v2.py
• https://github.com/cybertronai/gradient-checkpointing