GradCache
1.0.0
그라디언트 캐시는 GPU/TPU 메모리 제약을 넘어서는 대조적 인 학습 배치를 무제한으로 확장하는 간단한 기술입니다. 이것은 단일 GPU에서 무거운 하드웨어 (예 : 8 V100 GPU)를 복용하는 데 사용되는 훈련을 의미합니다. 또한 그라디언트 캐시를 통해 사용자는 대형 RAM GPU/TPU를 훨씬 더 비용 효율적인 높은 플롭 저 RAM 시스템으로 교체 할 수 있습니다.
이 repo는 메모리 제한 설정에서 대조적 인 학습 배치 크기를 스케일링하는 논문에 설명 된 구배 캐시의 일반적인 구현을 보유하고 있습니다. 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 의 모델 수와 임의의 키워드 인수 수와 동일한 숫자의 표현 텐서를 취하는 손실 함수. 입력 텐서를 기반으로 손실을 계산해야하며 어떠한 경우에도 AutoGrad 그래프에서 입력 텐서의 관계를 수정해야하며 나중에 그라디언트 캐시를 생성하기 위해 의존합니다.
split_input_fn- 정의 된 chunk_sizes를 기반으로 일반 모델 입력을 청크로 분할하는 선택적 함수. 제공되지 않으면이 클래스는 지원되는 유형의 입력을 분할하기 위해 최선을 다합니다. split_inputs 함수를 참조하십시오.
get_rep_fn- 일반 모델 출력 및 반환 표현 텐서를 취하는 선택적 함수. 제공되지 않으면 일반 출력은 표현 텐서라고 가정합니다.
FP16- 사실 인 경우 혼합 정밀 훈련을 실행하여 스케일러도 설정해야합니다.
스케일러 - 자동 혼합 정밀 훈련을위한 GradScaler 객체.
캐시 된 그라디언트 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- 실제 인 경우 분산 설정에서 각 모델에 대해 마지막 하위 배치의 전방 백 패스에 대한 프로세스에서 구배 감소 만 트리거됩니다. 이것은 a) 대형 모델 및/또는 b) 비 사소한 수의 하위 배치를 다룰 때 유용 할 수 있습니다.
LOSS_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 지정하여 이러한 입력을 더 작은 배치로 분해해야합니다. 일부 드문 경우, Heuristic이 입력에 일부 텐서를 보유하는 모든 CUDA 장치를 덮는 충분한 텐서를 잡을 수없는 경우 get_input_tensors 무시해야 할 수도 있습니다.
레이블과 텍스트의 임베딩 공간을 배우고 싶다고 가정 해 봅시다. 다음 네 쌍을 고려하십시오. (실제로, 당신은 더 많은 텍스트 항목을 가지고 있습니다.)
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 인수를 사용하여 일반적인 huggingface 모델 출력을 취하고 실제 표현 텐서를 반환하는 함수를 지정합니다.
모델 입력 생성,
xx = tokenizer(tt, return_tensors='pt', padding=True)
yy = tokenizer(tt2, return_tensors='pt', padding=True)
캐시 단계를 실행하고
gc(xx, yy, reduction='mean')
여기서 우리는 reduction='mean' Loss_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)
크기 (m * n)의 튜플 (xx, yy) 의 작은 배치를 방출하는 데이터 로더 loader 있다고 가정 해 봅시다.
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) (모든 레디스) 기기 전체의 기울기. 두 단계는 캐시 된 장식 기능 밖에서 발생합니다.
후자는 Eg a bert , DistributedDataParallel DataparAllel을 포장하여 쉽게 달성 할 수 있습니다.
bert = DistributedDataParallel(
bert, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)
전자는 손실 함수에서 추가 분산 OP를 필요로하며, 원래 손실 정의에 따라 수행되어야합니다. 예를 들어,
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- 데코레이터를 정의하여 일반 모델 호출 기능 및 손실 기능에서 구배 캐싱에 대한 고차 기능을 만듭니다.
