Получите те же потери, но другую оценку на нескольких графических процессорах при использовании DDP (DistributedDataPara ⇐ Python

[Anonymous]

При добавлении torch.nn.parallel.DistributedDataParallel в код обучения с одним графическим процессором я столкнулся с проблемой: я получаю одинаковые потери, но разный градиент на разных графических процессорах. Сравнивая с предыдущим кодом обучения с одним графическим процессором, я уверен, что потеря правильная, но после loss.backward() я наблюдаю за градацией весов и смещения слоев и обнаружил, что они разные на разных графических процессорах до all_gather градация слоев между all_gather и потерями вычислений одинакова на разных графических процессорах.

Это код контрастного обучения, поэтому я собираю тензоры all_gather на всех графических процессорах для вычисления общего окончательная потеря.

Вот часть кода модели:
import torch.nn as nn
import torch
from config.base_config import Config
from modules.transformer import Transformer
from modules.stochastic_module import StochasticText
from modules.basic_utils import AllGather
allgather = AllGather.apply
from modules.tokenization_clip import SimpleTokenizer

class CLIPStochastic(nn.Module):
def __init__(self, config: Config):
super(CLIPStochastic, self).__init__()
self.config = config

from transformers import CLIPModel
if config.clip_arch == 'ViT-B/32':
self.clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
elif config.clip_arch == 'ViT-B/16':
self.clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
else:
raise ValueError

self.task_config = config
config.pooling_type = 'transformer'
self.pool_frames = Transformer(config)
self.stochastic = StochasticText(config)

def forward(self, data, return_all_frames=False, is_train=True):
batch_size = data['video'].shape[0]
text_data = data['text'] # text_data["input_ids"].shape = torch.Size([16, 17])
video_data = data['video'] # [16, 12, 3, 224, 224]
video_data = video_data.reshape(-1, 3, self.config.input_res, self.config.input_res) # [192, 3, 224, 224]

if is_train:

text_features = self.clip.get_text_features(**text_data)
video_features = self.clip.get_image_features(video_data)

video_features = video_features.reshape(batch_size, self.config.num_frames, -1) # [bs, #F, 512]

text_features = allgather(text_features,self.task_config)
video_features = allgather(video_features,self.task_config)
torch.distributed.barrier()

video_features_pooled = self.pool_frames(text_features, video_features)

# @WJM: perform stochastic text
text_features_stochstic, text_mean, log_var = self.stochastic(text_features, video_features)

if return_all_frames:
return text_features, video_features, video_features_pooled, text_features_stochstic, text_mean, log_var

return text_features, video_features_pooled, text_features_stochstic, text_mean, log_var

else:

text_features = self.clip.get_text_features(**text_data)
video_features = self.clip.get_image_features(video_data)

video_features = video_features.reshape(batch_size, self.config.num_frames, -1)
video_features_pooled = self.pool_frames(text_features, video_features)

# @WJM: re-parameterization for text (independent of the text-cond pooling)
text_features_stochstic, _, _ = self.stochastic(text_features, video_features)

if return_all_frames:
return text_features, video_features, video_features_pooled, text_features_stochstic

return text_features, video_features_pooled, text_features_stochstic


и функция allgather выглядит так:
class AllGather(torch.autograd.Function):
"""An autograd function that performs allgather on a tensor."""

@staticmethod
def forward(ctx, tensor, args):
output = [torch.empty_like(tensor) for _ in range(args.world_size)]
torch.distributed.all_gather(output, tensor)
ctx.rank = local_rank
ctx.batch_size = tensor.shape[0]
return torch.cat(output, dim=0)

@staticmethod
def backward(ctx, grad_output):
local_grad = grad_output[ctx.batch_size * ctx.rank : ctx.batch_size * (ctx.rank + 1)]
return local_grad, None

Я пытался добавить all_reduce в класс AllGather задом наперед, как показано в следующих кодах, но, похоже, это не работает, возможно, потому, что DDP поставляется со встроенной функцией синхронного градиента?
def backward(ctx, *grads):
all_gradients = torch.stack(grads)
torch.distributed.all_reduce(all_gradients)
return all_gradients[torch.distributed.get_rank()]


Подробнее здесь: https://stackoverflow.com/questions/790 ... uteddatapa