GradCache

梯度缓存

梯度缓存是一种简单的技术，用于无限地扩展对比度学习批次，远远超出了GPU/TPU内存约束。这意味着可以在单个GPU上进行训练，例如8 V100 GPU，例如8 V100 GPU。此外，梯度缓存允许用户用更具成本效益的高绒布低RAM系统替换大型RAM GPU/TPU。

该仓库具有在我们的纸张缩放量表中描述的梯度缓存的通用实现，在内存有限的设置下进行了深层对比度学习批次大小。支持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在使用下。

lose_fn-损耗函数，其表示数字的表示张量等于models中的模型数量和任意数字的关键字参数。它应该根据输入张量计算损失，并且在任何情况下都不应修改自动射击图中的输入张量的关系，后来依赖于创建梯度缓存。

split_input_fn-一个可选功能，该功能将通用模型输入基于定义的Chunk_sizes将其拆分为块。如果未提供，此类将尽力拆分受支持类型的输入。请参阅split_inputs功能。

GET_REP_FN-可选功能，该功能采用通用模型输出并返回表示张量。如果未提供，则假定通用输出为表示张量。

FP16-如果是真的，请运行混合精度训练，这也需要设置缩放器。

鳞片- 用于自动混合精度训练的毕业生对象。

缓存梯度步骤

要运行缓存的梯度Computatoin步骤，请致电cache_step函数，

 cache_step(  
  *model_inputs,  
  no_sync_except_last: bool = False,  
  **loss_kwargs  
)

运行一个梯度缓存步骤。函数返回后，为自self.models中的每个模型计算更新，重量为梯度，就好像model_inputs是在足够大的硬件上作为巨大的单批次运行的。使用__call__调用gradcache对象也将调用此功能。

model_inputs-每个编码器模型的输入列表。应该按类似于self.models顺序。

no_sync_except_last-如果在分布式设置下为true，对于每个模型，则仅触发跨进程的跨进程的梯度减少，用于最后一个子批次的前回溯通行证。处理a）大型模型时可能会派上用场，/或b）非琐碎数量的子批次。

lose_kwargs-到损失函数loss_fn的其他关键字参数。这旨在启用灵活的损耗计算（由于Pytorch中的动态图），例如减少，加权等。可能是，使用loss_kwargs ，您可以从高速缓存跟踪的那些编码模型中合并输出。

返回- 损耗，当前的步骤损失量表张量（从图形上分离）。

本机支持的输入类型

• X：张量 - 将作为model(x)传递
• X：列表[张量] - 将作为model(*x)传递
• X：dict [str，张量]（或用户dict [str，张量]） - 将作为model(**x)
• x：元组[列表[张量]，dict [str，张量] - 将作为model(*x[0], **x[1])传递。

其他通用输入不完全支持，我们使用以下启发式方法执行模型调用，

• X：列表[任何] - 将作为model(*x)传递
• X：dict [str，任何] - 将作为model(**x)
• X：元组[列表[Any]，dict [str，Any] - 将作为model(*x[0], **x[1])传递。

要与它们一起运行，应在缓存初始化期间指定split_input_fn ，以将这些输入分解为较小的批次。在某些极少数情况下，您可能还需要覆盖get_input_tensors ，当它的启发式量无法捕获足够的张量，以覆盖所有cuda设备，这些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'作为损失_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
)

在引擎盖下，将注册不同的钩子以进行正确的梯度计算。

有多个GPU分布式培训？

我们预计表示表示形式的交叉流程通信将由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])

将类型列表的位置和关键字参数（张量）连接到0位维度的单个张量中的装饰器。这可能会派上用场，以处理来自多个缓存前进的代表性张量结果。

功能- 损失功能

返回- 缓存结果的装饰损失功能。

 grad_cache.functional.gather_input_tensor(func: Callable[..., Tensor], axis=0)

张量的所有聚集位置和关键字参数的装饰器，并在轴上串联。旨在用于创造分布式对比度学习损失。

功能- 损失功能

返回- 分布式培训的装饰损失功能。

用法

如果您的数据加载程序发射小批量，则功能装饰器特别有用，您可以从中构造大批批次。假设您也想进行自动混合精度，我们首先定义模型调用功能和丢失功能，

 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 ，该加载程序会散发出一小批元组(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）（ALL-）跨设备收集表示形式，以及跨设备的梯度。这两个步骤都将发生在缓存的装饰功能之外。

通过包装编码器（例如bert ），在DistributedDataParallel中可以易于实现。

 bert = DistributedDataParallel(
	bert, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)

前者需要在损失函数中进行额外的分布式操作，应根据原始损失定义进行。例如，

 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-定义装饰器为从普通模型调用功能和损失功能中梯度缓存创建高阶功能。