Thank you for reply! @mrshenli, I already tried with to("cuda:0") . and as far as I knew cuda() and to(device="cuda") is same.
Below is the issues.

In my model, I need to use some constant tensor, and I defined it in model’s forward as you can see top post.
So I give torch.device('cuda') to model, and in model, locate that tensor to the device. more specified,

device = torch.device('cuda')
...
self.model = Bert(6, 12, 513, 384*4, 64, 64, 2, 384, 
						self.base_task.max_vocab_indexes['input_ids'], device=device)
self.model = nn.DataParallel(self.model)
self.model = self.model.cuda()


I ran it with 2 gpus, and it shows up right.

    device_count = torch.cuda.device_count()
    print(f'gpu count: {torch.cuda.device_count()}')

gpu count: 2


dataset iterator is device=torch.device('cuda')

DataParallel wrapped model with cuda()

and model get returned, logits_clsf and logits_lm

loss was calculated.
and loss.backward() phase, I think stuck. print after loss.backward() wasn’t printed.
belows are logs and the model.

batch starts
logits_lm device: cuda:0, batch target device: cuda:0
loss_lm calculated
loss_clsf calculated

            print('epoch starts')
            for i, batch in tq_iter:
                self.model.zero_grad()
                print('batch starts')
                logits_lm, logits_clsf = self.model(batch.input_ids.transpose(0,1), batch.segment_ids.transpose(0,1), batch.masked_pos.transpose(0,1))
                print(f'logits_lm device: {logits_lm.device}, batch target device: {batch.masked_tokens.device}')
                loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
                print('loss_lm calculated')
                loss_lm = (loss_lm.float()).mean()
                loss_clsf = self.criterion(logits_clsf, batch.is_next) # for sentence classification
                print('loss_clsf calculated')
                loss = loss_lm + loss_clsf
                loss.backward()
                print('loss backwarded')

without any error message or logs, the log loss backwarded didn’t shows up. I think it’s stuck.
I don’t know what I have to do.
All the post in this thread, are my tries after this happened…
I assumed that,
in backward phase,
‘replica 0’, which is dedicated to cuda:0 is worked well(backward completed well), and waiting for ‘replica 1’, which is dedicated to cuda:1 to be completed with it’s backward, but in this case, some error occurred to him and failed to backward, so ‘replica 0’ keeps waiting and thus 'loss backwarded' log didn’t show up.
Further more assuming that,
even embedding tensor in each replica’s model, looks like splitted well and dedicated well like this,
class Embedding(nn.Module):
    def __init__ 
        ...
    def forward(self, x, seg):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long, device=self.device)
        pos = pos.unsqueeze(0).expand_as(x)  # (seq_len,) -> (batch_size, seq_len)
        pos.requires_grad = False
        print(f'pos tensor device: {pos.device}, shape: {pos.shape}')
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)


pos tensor device: cuda:1, shape: torch.Size([64, 513])
pos tensor device: cuda:0, shape: torch.Size([64, 513])

embedding tensor device: cuda:1, shape: torch.Size([64, 513, 384])
embedding tensor device: cuda:0, shape: torch.Size([64, 513, 384])

and about other cuda() defined tensor in the model, too
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])

and the last logit gathered well, too,
logtis_lm shape : torch.Size([128, 5, 6015])

There are some error could happened at backward phase, but I cannot even imagine about it…
Or maybe the sum about loss, loss = loss_lm + loss_clsf or loss_lm.float().mean() after get gathered outputs, affects to backward phase, I think but … I don’t know what should I do.
                loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
                print('loss_lm calculated')
                print(f'loss_lm tensor device: {loss_lm.device}, shape: {loss_lm.shape}')
                loss_lm = (loss_lm.float()).mean()
                print(f'loss_lm tensor device: {loss_lm.device}, shape: {loss_lm.shape}')
                loss_clsf = self.criterion(logits_clsf, batch.is_next) # for sentence classification
                print('loss_clsf calculated')
                print(f'loss_clsf tensor device: {loss_clsf.device}, shape: {loss_clsf.shape}')
                loss = loss_lm + loss_clsf
                print(f'sumloss tensor device: {loss.device}, shape: {loss.shape}')
                loss.backward()
                self.optimizer.step()
                print('stepped')