GradCache
1.0.0
Gradientキャッシュは、GPU/TPUメモリの制約をはるかに超えて、コントラストの学習バッチを無制限にスケーリングするための簡単な手法です。これは、1つのGPUで重いハードウェア（8 V100 GPUなど）を摂取するために使用されるトレーニングを意味します。さらに、グラデーションキャッシュにより、ユーザーは大きなRAM GPU/TPUを、よりコスト効率の高い高フロップ低RAMシステムに置き換えることができます。
このレポは、メモリ限定セットアップの下で深いコントラスト学習バッチサイズをスケーリングするペーパーで説明されているグラデーションキャッシュの一般的な実装を保持します。 PytorchとJaxの両方のフレームワークがサポートされています。
@inproceedings{gao2021scaling,
title={Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup},
author={Luyu Gao, Yunyi Zhang, Jiawei Han, Jamie Callan},
booktitle ={Proceedings of the 6th Workshop on Representation Learning for NLP},
year={2021},
}
新規：JaxとTPUをサポートしました！
勾配キャッシュは、密な通過回収（DPR）にも統合されています。 GC-DPRツールキットをチェックアウトします。
最初に、PytorchまたはJaxのいずれかの希望のディープラーニングバックエンドをインストールします。 GradCacheをインストールするには、このレポをクローンしてPIPを実行します。
git clone https://github.com/luyug/GradCache
cd GradCache
pip install .
開発のために、
pip install --editable .
グラジエントキャッシング機能は、 GradCacheクラスで実装されています。古いプロジェクトにパッチを適用する代わりに新しいプロジェクトを開発している場合は、努力を削減するための機能的アプローチもチェックアウトしてください。
Jax/Flaxユーザーの場合は、こちらの単純な列車機能をご覧ください。
クラスの__init__メソッドはキャッシュを定義し、モデル動作を簡単に調整するためにいくつかの機能パラメーター*_fnを備えています。または、Gradcacheをサブクラスすることもできます。
grad_cache.GradCache(
models: List[nn.Module],
chunk_sizes: Union[int, List[int]],
loss_fn: Callable[..., Tensor],
split_input_fn: Callable[[Any, int], Any] = None,
get_rep_fn: Callable[..., Tensor] = None,
fp16: bool = False,
scaler: GradScaler = None,
)
モデル- グラデーションキャッシュを使用して更新されるエンコーダーモデルのリスト。
Chunk_sizes-チャンクサイズを示す整数。または、各モデルのチャンクサイズの整数のリスト。これにより、各モデルのサブバッチサイズを制御して、前方バックワードパスを実行し、使用可能なGPUメモリに基づいて設定する必要があります。値が小さすぎると、GPUは利用されます。
loss_fn-数の表現テンソルをmodelsの数のモデル数と任意の数のキーワード引数に等しくする損失関数。入力テンソルに基づいて損失を計算する必要があり、自動グラフの入力テンソルの関係を変更しないでください。
split_input_fn-定義されたchunk_sizesに基づいて汎用モデル入力をチャンクに分割するオプションの関数。提供されていない場合、このクラスは、サポートされているタイプの入力を分割するために最善を尽くします。 split_inputs関数を参照してください。
get_rep_fn-ジェネリックモデルの出力と戻りの表現テンソルを取得するオプションの関数。提供されていない場合、一般的な出力は表現テンソルであると想定されます。
FP16- Trueの場合、スケーラーも設定する必要がある混合精度トレーニングを実行します。
スカラー- 自動混合精密トレーニングのグラッドスカラーオブジェクト。
キャッシュグラデーションComputatoinステップを実行するには、 cache_step関数を呼び出します。
cache_step(
*model_inputs,
no_sync_except_last: bool = False,
**loss_kwargs
)
単一の勾配キャッシュステップを実行します。機能が戻ってくると、 model_inputs十分に大きなハードウェアで巨大な単一バッチとして実行されるかのように、重みに入力された勾配を持つself.modelsの各モデルの更新が計算されます。 __call__を使用してGradcacheオブジェクトを呼び出すと、この関数も呼び出します。
model_inputs各エンコーダーモデルへの入力のリスト。 self.modelsと同様の順序である必要があります。
no_sync_except_last -trueの場合、分散セットアップの下で、各モデルの場合、最後のサブバッチのフォワードバックワードパスのプロセス間での勾配削減のみをトリガーします。これは、a）大規模なモデル、および/またはb）非些細なサブバッチを扱うときに役立つ可能性があります。
loss_kwargs-損失関数loss_fnに対する追加のキーワード引数。これは、柔軟な損失計算（Pytorchの動的グラフのおかげで）を有効にすることを目的としています。削減、重み付けなど、Pytorchの動的グラフのおかげです）。潜在的に、 loss_kwargsを使用して、キャッシュで追跡されないエンコーダーモデルからの出力を組み込むことができます。
戻り- 損失、現在のステップ損失スケーラーテンソル（グラフから分離）。
model(x)として渡されますmodel(*x)として渡されますmodel(**x)として渡されますmodel(*x[0], **x[1])他の一般的な入力は完全にはサポートされていません。次のヒューリスティックを使用してモデルコールを実行します。
model(*x)として渡されますmodel(**x)として渡されますmodel(*x[0], **x[1])それらと一緒に実行するには、これらの入力を小さなバッチに分割するために、キャッシュ初期化中にsplit_input_fn指定する必要があります。まれな場合には、Heuristicが入力にいくつかのテンソルを保持するすべてのCUDAデバイスをカバーする十分なテンソルをつかむことができない場合、 get_input_tensorsオーバーライドする必要がある場合があります。
ラベルとテキストの埋め込みスペースを学びたいとします。次の4つのペアを検討してください。 （実際には、もっと多くのテキストエントリがあります。）
labels = ['fruit', 'meat', 'school', 'company']
texts = [
'this is an apple',
'steak should be cooked medium rare',
'cmu is pittsburgh',
'apple sells laptop'
]
エンコーダーモデルを初期化し、
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
encoder1 = AutoModel.from_pretrained("bert-base-uncased").cuda()
encoder2 = AutoModel.from_pretrained("bert-base-uncased").cuda()
Gradcacheオブジェクトを初期化し、
from grad_cache import GradCache
from grad_cache.loss import SimpleContrastiveLoss
loss_fn = SimpleContrastiveLoss()
gc = GradCache(
models=[encoder1, encoder2],
chunk_sizes=2,
loss_fn=loss_fn,
get_rep_fn=lambda v: v.pooler_output
)
ここでは、 GET_REP_FN引数を使用して、汎用ハグFaceモデル出力を取得し、実際の表現テンソルを返す関数を指定します。
モデル入力を作成し、
xx = tokenizer(tt, return_tensors='pt', padding=True)
yy = tokenizer(tt2, return_tensors='pt', padding=True)
キャッシュステップを実行し、
gc(xx, yy, reduction='mean')
ここでは、損失の動作を制御するために、 reduction='mean'をlos_kwargsとして使用します。定義されたoptimizerを使用すると、完全なグラデーションアップデートを次のように実行できます。
optimizer.zero_grad()
gc(xx, yy, reduction='mean')
optimizer.step()
これは、（魔法の）動的グラフによって自然に処理されます。同じエンコーダモデルの浅いコピーをGradCache INITメソッドに渡します。
tied_encoder = AutoModel.from_pretrained("bert-base-uncased").cuda()
gc = GradCache(
models=[tied_encoder , tied_encoder],
chunk_sizes=2,
loss_fn=loss_fn,
get_rep_fn=lambda v: v.pooler_output
)
フードの下では、正しいグラデーション計算を行うために異なるフックが登録されます。
表現のクロスプロセス通信は、 loss_fnによって処理されると予想されます。
from grad_cache.loss import DistributedContrastiveLoss
loss_fn_dist = DistributedContrastiveLoss()
勾配削減のためにエンコーダモデルを適切にラップし、
encoder1_ddp = DistributedDataParallel(
encoder1, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)
encoder2_ddp = DistributedDataParallel(
encoder2, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)
分散損失とDDPモデルを使用して、キャッシュを初期化できます。
gc = GradCache(
models=[encoder1_ddp, encoder2_ddp],
chunk_sizes=2,
loss_fn=loss_fn_dist,
get_rep_fn=lambda v: v.pooler_output
)
キャッシュステップを実行し、
gc(xx, yy, no_sync_except_last=True, reduction='mean')
no_sync_except_last=True不必要な勾配削減を避けます。
新しいプロジェクトを開発している場合は、キャッシュ用の高次関数を作成するために提供したデコレータをチェックアウトすることもお勧めします。
grad_cache.functional.cached(func: Callable[..., Tensor])
モデルコール関数をキャッシュされた互換性のあるバージョンに持ち込むデコレーター。
FUNC-モデルを呼び出し、表現テンソルを返す関数。
返品- 1）キャッシュ構造の表現葉のテンソルを返す関数、2）2番目の前方とキャッシュされた後方の閉鎖関数。コール2）1）損失テンソルを後方に呼び出した後の引数として。
grad_cache.functional.cat_input_tensor(func: Callable[..., Tensor])
タイプリスト[テンソル]の位置とキーワードの引数を0番目の次元の単一のテンソルに連結するデコレーター。これは、複数のキャッシュフォワードからの表現テンソルの結果に便利に対処することができます。
FUNC-損失関数
返品- キャッシュ結果の装飾損失関数。
grad_cache.functional.gather_input_tensor(func: Callable[..., Tensor], axis=0)
タイプテンソルの全集まりの位置とキーワード引数が軸上でそれらを連結するデコレーター。分散型の対照学習損失を作成するために使用することを目的としています。
FUNC-損失関数
返品- 分散トレーニングのための装飾損失関数。
機能的なデコレーターは、データローダーが小さなバッチを発している場合に特に便利です。そこから大きなバッチを構築できます。また、自動混合精度を実行したいとします。まず、モデルコール関数と損失関数を定義します。
from grad_cache.functional import cached, cat_input_tensor
import torch
import torch.nn.functional as F
from torch.cuda.amp import autocast
@cached
@autocast()
def call_model(model, input):
return model(**input).pooler_output
@cat_input_tensor
@autocast()
def contrastive_loss(x, y):
target = torch.arange(0, y.size(0), int(y.size(0) / x.size(0)), device=x.device)
scores = torch.matmul(x, y.transpose(0, 1))
return F.cross_entropy(scores, target=target)
サイズのタプル(xx, yy)の小さなバッチ（m * n）の小さなバッチを発するデータローダーloaderがあり、16の小さなバッチを集約して（16m * 16n）のバッチを取得してトレーニングしたいとします。
cache_x = []
cache_y = []
closures_x = []
closures_y = []
for step, sub_batch in enumerate(loader):
xx, yy = sub_batch
rx, cx = call_model(bert, xx)
ry, cy = call_model(bert, yy)
cache_x.append(rx)
cache_y.append(ry)
closuresx.append(cx)
closuresy.append(cy)
if (step + 1) % 16 == 0:
loss = contrastive_loss(cache_x, cache_y)
scaler.scale(loss).backward()
for f, r in zip(closuresx, cache_x):
f(r)
for f, r in zip(closuresy, cache_y):
f(r)
cache_x = []
cache_y = []
closures_x = []
closures_y = []
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
分散マルチプロセストレーニングを実行するには、次のものが必要です。1）（すべて）デバイス間の表現を収集し、デバイス間で2）（全レディュース）グラデーションが必要です。両方のステップは、キャッシュされた装飾された機能の外で発生します。
後者は、 DistributedDataParallelでエンコーダーをbertすることで簡単に実現できます。
bert = DistributedDataParallel(
bert, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)
前者は、損失関数に追加の分散OPSを必要とします。これは、元の損失定義に従って行う必要があります。例えば、
from torch import distributed as dist
from grad_cache.functional import cat_input_tensor, gather_input_tensor
@cat_input_tensor
@gather_input_tensor
@autocast()
def contrastive_loss(x, y):
target = torch.arange(0, y.size(0), int(y.size(0) / x.size(0)), device=x.device)
scores = torch.matmul(x, y.transpose(0, 1))
# scale the loss as DistributedDataParallel will do mean reduce
return F.cross_entropy(scores, target=target) * dist.get_world_size()
grad_cache/grad_cache.py- GradCacheクラスを定義します。コードはコメントを含む300行未満です。開発のために、私たちはあなたがそれを読むことをお勧めします。
grad_cache/functional.py-デコレータを定義して、通常のモデルコール関数と損失関数から勾配キャッシュのための高次関数を作成します。
