I’m currently training a word2vec model with a large vocabulary (50 million). When I use a embedding layer with 64 dimension, the volume of the embedding weight (64x50milliton) will exceed the memory of a single GPU, so I wanna put the weight on different GPUs.
I’ve already tried DataParallel but it still doesn’t work.
        self.in_embed = nn.Embedding(self.num_classes, self.embed_size, sparse=True)
        self.in_embed.weight = Parameter(t.FloatTensor(self.num_classes, self.embed_size).uniform_(-1, 1))
        self.in_embed = DataParallel(self.in_embed,device_ids=[0,1,2])
        self.in_embed.cuda()