Embedding documents on multi-GPU single-ode Docker using pretrained models of huggingface transformers and pytorch DistributedDataParallel
Author: ntaherkhani
Pull request: no

Hi, 
This is a question.
I am trying to embed some documents including a couple of sentences using huggingface transformers models. I have multi-gpu single-node and I want to do embedding parallel and distributed in all 8 gpus. I tried to use pytorch DistributedDataParallel, but I think all sentences are sending to all GPUs and for all sentences, it is returning one tensor. this is a sample code:

import torch
import torch.nn as nn
import torch.nn.functional as F
import time

import argparse
import os

from transformers import AlbertTokenizer, AlbertModel
import numpy
from tqdm import tqdm
from torch.utils.data import DataLoader,TensorDataset

def parse_args():
    parse = argparse.ArgumentParser()
    parse.add_argument(
            '--local_rank',
            dest = 'local_rank',
            type = int,
            default = 0,
            )
    parse.add_argument("--gpu", type=str, default='None',
                        help="choose gpu device.")
    return parse.parse_args()



def train():
    args = parse_args()

    if not args.gpu == 'None':
        device = torch.device("cuda")
        os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu
    else:
        device = torch.device("cpu")

    torch.cuda.set_device(args.local_rank)

    torch.distributed.init_process_group(
                backend='nccl',
                init_method='env://',
                )

    tokenizer = AlbertTokenizer.from_pretrained('albert-xxlarge-v2')


    sentences=['I love tea',    
               'He hates tea',
               'We love tea', 
               'python coder',
               'geeksforgeeks',
               'coder in geeksforgeeks']


    sentence_tokens = []
    for sent in (sentences):

        token_id = tokenizer.encode(sent, max_length=128, add_special_tokens=True, pad_to_max_length=True)

        sentence_tokens.append(token_id)
    original_sentences = torch.tensor(sentence_tokens)
    train_dataset = TensorDataset(original_sentences)

    #setup training sampler
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,num_replicas=len(sentences))

    #setup training data loader with the train sampler setup
    train_dataloader = DataLoader(train_dataset, batch_size=16,sampler=train_sampler, shuffle=False)


    model = AlbertModel.from_pretrained('albert-xxlarge-v2', return_dict=True)

    model = model.to(device)
    model = nn.parallel.DistributedDataParallel(model,
            device_ids = [args.local_rank, ],
            output_device = args.local_rank,\
            find_unused_parameters=True
            )
    for batch in (train_dataloader):

        batch_input_tensors = batch[0].to('cuda')
        outputs = model(batch_input_tensors)
        last_hidden_states = outputs.last_hidden_state

        average= torch.mean(last_hidden_states,dim=1)

if __name__ == "__main__":
    train()

all of sentences are sending to all 8 GPUs and output as last_hidden_states is only one tensor. I got the average of tensor elements because I thought at the end they should be same but they aren't. how can do it distributed and sentences distribute to GPUs and embed over there? and finally for each sentence or for my final case each Doc I have one tensor as feature vector? thanks


Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.