Блог

Главная - Вопросы по программированию - Факельная распределенная трансляция и уменьшение между устройствами CPU / GPU

Факельная распределенная трансляция и уменьшение между устройствами CPU / GPU

#pytorch #distributed #torch #array-broadcasting #torchvision

#pytorch #распределенный #факел #массив-широковещательная передача #torchvision

Вопрос:

Использование пакета torch.distributed. Я пытаюсь переместить тензоры из CPU -> GPU0, GPU1 в два отдельных процесса и обновить основную версию (на CPU).

Предположим, у меня подключены два графических процессора. Один на Device0, другой на Device1.

  1. Например, сохраните очень большой массив на CPU (что-то, что не может поместиться на одном устройстве / gpu) X = [1,2,3,4,5,6] .
  2. Передайте часть массива на устройство GPU 0 и устройство GPU 1. 0,1 имеют разные фрагменты этого массива. GPU0 Inds = [0,1] GPU0 data = [1,2] GPU1 Inds = [2,3] GPU1 data = [2,3]
  3. Запустите процесс на GPU0 и GPU1 независимо. Для этой цели подойдет простая функция Add() .
  4. При необходимости обновите версию CPU данными GPU (для индексов, которые захватил GPU). Здесь я бы, вероятно, использовал reduce для получения обоих тензоров с устройств. Я бы, вероятно, сохранил его в формате ключ-значение, где ключ — это идентификатор устройства (0 для GPU 0, 1 для GPU 1), а индексы и данные хранятся в кортеже. Затем мне нужно обновить версию процессора и снова запустить весь процесс.

То, что я пытаюсь сделать, можно увидеть на этой диаграмме.введите описание изображения здесь

Я планирую использовать серверную часть NCCL, которая, по-видимому, поддерживает широковещательную передачу и сокращение.

Мой код должен выглядеть примерно так:

Функция Main() запускает два процесса и удерживает тензор процессора Foo() запускает два процесса и разрешает широковещательную передачу и обновление между ними (что я хочу сделать на диаграмме)

 def main():
  args.world_size=2
  args.backend='nccl'
  os.environ['MASTER_ADDR'] = '127.0.0.2'
  os.environ['MASTER_PORT'] = '29500'
  args.data = t.Tensor([1,2,3,4,5,6])
  mp.spawn(foo, nprocs=2, args=(args,), join=True)

def foo(rank):
  data = args.data # cpu data
  dist.init_process_group(backend=args.backend, init_method='env://', world_size=args.world_size, rank=rank)
  if rank==0:
    inds=[0,1]
  elif rank == 1:
    inds=[2,3]
  
  gpu_data = data[inds].cuda(rank) # will send to GPU0 or GPU1. # probably need to use the torch.dist.broadcast operation here but idk how.
  data[inds].data=gpu.data # this would be the update step. 
  
  gpu_data  =1
  dist.destroy_process_group()
  print('data: ', data) >> is not [2,3,4,5,5,6] which is what it should be

Ответ №1:

Судя по всему, вы хотите использовать три реплики: одну на CPU, две на GPU. Если это так, возможно, вам нужен другой сервер, чем nccl . Пожалуйста, взгляните на https://pytorch.org/docs/stable/distributed.html