GradCache - 突破GPU/TPU内存限制，实现对比学习无限扩展

GradCache 项目介绍

GradCache 简介

GradCache 是一种简单的方法，可以将对比学习批次的扩展提升到超过 GPU/TPU 内存限制。这意味着过去需要重型硬件（例如 8 个 V100 GPU）进行的训练，现在可以在单个 GPU 上完成。此外，GradCache 还允许用户用更加经济高效的高 FLOP 低 RAM 系统替代高 RAM 的 GPU/TPU。

该项目提供了一个通用的 GradCache 实现，支持 Pytorch 和 JAX 框架。这一技术已经在论文《在内存限制条件下扩展深度对比学习的批次大小》中进行过描述，并被集成到 Dense Passage Retrieval (DPR) 系统中。

安装指南

要安装 GradCache，首先您需要安装所需的深度学习后端（Pytorch 或 JAX）。然后克隆项目库并运行 pip 进行安装：

git clone https://github.com/luyug/GradCache
cd GradCache
pip install .

对于开发用途，可使用以下安装方式：

pip install --editable .

如何使用

GradCache 的功能通过 GradCache 类实现。如果您正在开发新项目，而不是修补旧项目，还可以查看我们提供的具备简化工作量的功能方法。

初始化

要使用 GradCache 首先需要初始化 GradCache 类，其 __init__ 方法需要对缓存进行定义，并且包含多个功能参数以方便调整模型行为。您也可以通过继承方式来使用。

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,  
)

models：要通过 GradCache 更新的编码器模型列表。
chunk_sizes：块大小的整数或每个模型的块大小整数列表。该值基于可用的 GPU 内存设定，不宜过小以避免 GPU 未充分利用。
loss_fn：计算模型表示的损失函数。
split_input_fn：可选函数，用于根据定义的 chunk_sizes 将输入切分为小块。
get_rep_fn：可选函数，用于模型输出获取表示张量。
fp16 和 scaler：是否使用混合精度训练。

缓存梯度步骤

可调用 cache_step 来运行缓存梯度计算步骤：

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

在大多数情况下，通过该方法可以让模型在一个虚拟较大批次下运行，就如同在足够大的硬件上运行。执行此函数之后，模型的参数将更新。

使用 Huggingface Transformers 的示例

以下是一个简单示例：使用 BERT 模型创建双编码器来学习标签和文本的嵌入空间。

from transformers import AutoTokenizer, AutoModel
from grad_cache import GradCache
from grad_cache.loss import SimpleContrastiveLoss

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
encoder1 = AutoModel.from_pretrained("bert-base-uncased").cuda()
encoder2 = AutoModel.from_pretrained("bert-base-uncased").cuda()

loss_fn = SimpleContrastiveLoss()
gc = GradCache(
  models=[encoder1, encoder2],
  chunk_sizes=2,
  loss_fn=loss_fn,
  get_rep_fn=lambda v: v.pooler_output
)

创建模型输入并运行缓存步骤：

xx = tokenizer(["this is an apple"], return_tensors='pt', padding=True)
yy = tokenizer(["apple sells laptop"], return_tensors='pt', padding=True)

gc(xx, yy, reduction='mean')

分布式训练和多 GPU 支持

GradCache 可以与分布式数据并行 (Distributed Data Parallel) 模型结合使用，实现跨设备的梯度计算与通信。

from torch.nn.parallel import DistributedDataParallel

encoder1_ddp = DistributedDataParallel(encoder1, device_ids=[local_rank], output_device=local_rank)
encoder2_ddp = DistributedDataParallel(encoder2, device_ids=[local_rank], output_device=local_rank)

loss_fn_dist = DistributedContrastiveLoss()
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')

功能方法

项目还提供了便捷的功能装饰器，例如 cached 和 cat_input_tensor，简化缓存的模型和损失函数调用过程。这些装饰器特别适用于处理小批量数据构建大批量进行训练。

总之，GradCache 为那些内存受限的深度对比学习任务提供了一种高效且易于实施的方法，极大地降低了硬件门槛，具有广泛的应用前景。