GradCache
1.0.0
يعد التدرج ذاكرة التخزين المؤقت تقنية بسيطة لتقييم دفعة تعليمية متناقضة بشكل غير محدود إلى ما هو أبعد من قيد ذاكرة GPU/TPU. هذا يعني أن التدريب الذي يستخدم لأخذ أجهزة ثقيلة ، على سبيل المثال ، GPU 8 V100 ، يمكن القيام به على وحدة معالجة الرسومات الواحدة. بالإضافة إلى ذلك ، تتيح ذاكرة التخزين المؤقت التدرج للمستخدمين استبدال GPU/TPU الكبير الكبير مع أنظمة ذاكرة الوصول العشوائي المنخفضة ذات التكلفة العالية أكثر بكثير من حيث التكلفة.
يحمل هذا الريبو تنفيذًا عامًا لذاكرة التخزين المؤقت التدرج الموصوفة في تحجيم الورق الذي يقوم بتوسيع حجم الدفعة التعليمية العميقة للتناقض تحت الإعداد المحدود للذاكرة. يتم دعم كل من Pytorch و Jax Frameworks.
@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 المتاحة. القيمة الصغيرة جدا سوف تترك وحدة معالجة الرسومات تحت استخدام.
LOST_FN - وظيفة الخسارة التي تأخذ موترات تمثيل العدد تساوي عدد النماذج في models والأرقام التعسفية من وسيطات الكلمات الرئيسية. يجب أن تحسب الخسارة بناءً على موترات الإدخال ، وفي أي حال من الأحوال تعدل علاقات الإدخال في الرسم البياني الذاتي ، والتي تعتمد عليها لاحقًا لإنشاء ذاكرة التخزين المؤقت التدرج.
split_input_fn - وظيفة اختيارية تقسم إدخال النموذج العام إلى أجزاء استنادًا إلى chunk_sizes المحددة. إذا لم يتم توفيرها ، فستبذل هذا الفئة قصارى جهدها لتقسيم مدخلات الأنواع المدعومة. انظر وظيفة split_inputs .
GET_REP_FN - وظيفة اختيارية تأخذ توترات تمثيل النموذج العام وإرجاع. إذا لم يتم توفيره ، من المفترض أن يكون الإخراج العام هو موتر التمثيل.
FP16 - إذا كان ذلك صحيحًا ، قم بتشغيل التدريب الدقيق المختلط ، والذي يتطلب تعيين Scaler أيضًا.
Scaler - كائن GradScaler للتدريب التلقائي المختلط الدقيق.
لتشغيل خطوة computatoin التدرج المخبأة ، استدعاء وظيفة cache_step ،
cache_step(
*model_inputs,
no_sync_except_last: bool = False,
**loss_kwargs
)
تشغيل خطوة ذاكرة التخزين المؤقت التدرج واحد. model_inputs إرجاع الوظيفة ، يتم حساب التحديثات لكل نموذج في self.models . سيؤدي استدعاء كائن GradCache مع __call__ أيضًا إلى استدعاء هذه الوظيفة.
Model_inputs - قائمة المدخلات لكل نموذج تشفير. يجب أن يكون في ترتيب مماثل مثل self.models .
no_sync_except_last- إذا كان هذا صحيحًا ، ضمن الإعداد الموزع ، لكل نموذج ، يؤدي فقط إلى انخفاض التدرج عبر عمليات التمريرة الأمامية للدائم الفرعي الأخير. يمكن أن يكون هذا مفيدًا عند التعامل مع نموذج كبير ، و/أو ب) عدد غير تافهة للمزارع الفرعية.
LOST_KWARGS - وسيطات الكلمات الرئيسية الإضافية إلى وظيفة الخسارة loss_fn . يهدف هذا إلى تمكين حساب الخسارة المرن (بفضل الرسم البياني الديناميكي في 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 أثناء تهيئة ذاكرة التخزين المؤقت لكسر هذه المدخلات إلى دفعات أصغر. في بعض الحالات النادرة ، قد تحتاج أيضًا إلى تجاوز get_input_tensors عندما لا تتمكن الاستدلال من الاستيلاء على ما يكفي من الموترات التي تغطي جميع أجهزة CUDA التي تحمل بعض الموترات في المدخلات.
قل أننا نريد أن نتعلم مساحة تضمين من الملصقات والنص. النظر في أزواج الأربعة التالية. (في الممارسة العملية ، سيكون لديك المزيد من إدخالات النص أطول وأطول.)
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 لتحديد وظيفة تأخذ إخراج نموذج العناق العام وإرجاع موتر التمثيل الفعلي.
إنشاء إدخال النموذج ،
xx = tokenizer(tt, return_tensors='pt', padding=True)
yy = tokenizer(tt2, return_tensors='pt', padding=True)
قم بتشغيل خطوة ذاكرة التخزين المؤقت ،
gc(xx, yy, reduction='mean')
نحن هنا نستخدم reduction='mean' كخسارة_كوارج للسيطرة على سلوك الخسارة. مع 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) مع 1) كحجة بعد الاتصال للخلف على موتر الخسارة.
grad_cache.functional.cat_input_tensor(func: Callable[..., Tensor])
ديكور يسلط وسيطات الموضعية والكلمات الرئيسية من قائمة النوع [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)
لنفترض أن لديك loader dataloader ينبعث من دفعات صغيرة من tuple (xx, yy) من الحجم (m * n) وأنك تريد التدريب عن طريق تجميع 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) (All-Reduce) عبر الأجهزة. ستحدث كلتا الخطوتين خارج الترون المزخرف المخزنة مؤقتًا.
هذا الأخير يسهل تحقيقه عن طريق التفاف المشفرات ، على سبيل المثال bert ، في DistributedDataParallel .
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 - حدد الديكورات لإنشاء وظيفة ترتيب أعلى للتخزين المؤقت التدرج من وظائف استدعاء النموذج العادية ووظائف الخسارة.
