torch DataParallel causes error

[moritzblum]

Hello all,

I'm using the BertAdapterModel for encoding in my model as follows:

class BertEmbeddingsLP(InductiveLinkPrediction):

    def __init__(self):
        super().__init__()

        self.encoder = BertAdapterModel.from_pretrained(
           "bert-base-cased",
            output_hidden_states=True
        )

        self.encoder.add_adapter("test_adapter", config="double_seq_bn")
        self.encoder.train_adapter("test_adapter")

        hidden_size = self.encoder.config.hidden_size
        self.enc_linear = nn.Linear(hidden_size, self.dim, bias=False)

When I make my model data parallel by model = torch.nn.DataParallel(model).to(device) and then call model.encoder(text_tok, text_mask), I get the following device error:

Original Traceback (most recent call last):
  ......
    embs = self.encoder(text_tok, text_mask)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/models/bert/adapter_model.py", line 69, in forward
    outputs, context = self.bert(
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/context.py", line 116, in wrapper_func
    results = f(self, *args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/model_mixin.py", line 1270, in forward
    return super().forward(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 1013, in forward
    encoder_outputs = self.encoder(
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 607, in forward
    layer_outputs = layer_module(
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 497, in forward
    self_attention_outputs = self.attention(
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 436, in forward
    attention_output = self.output(self_outputs[0], hidden_states)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/models/bert/modeling_bert.py", line 149, in forward
    hidden_states = self.bottleneck_layer_forward(hidden_states, input_tensor, self.LayerNorm)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/methods/bottleneck.py", line 348, in bottleneck_layer_forward
    state = self.compose(adapter_setup, state)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/methods/adapter_layer_base.py", line 472, in compose
    state = composition_func(adapter_setup, state, lvl=0)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/methods/adapter_layer_base.py", line 308, in compose_stack
    state = self.compose_single(adapter_stack_layer, state, lvl=lvl + 1)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/methods/bottleneck.py", line 230, in compose_single
    layer_output = adapter_layer(
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/adapters/methods/modeling.py", line 172, in forward
    down = self.adapter_down(x)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/container.py", line 217, in forward
    input = module(input)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/homes/mblum/miniconda3/envs/lp/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

Can anybody help me find the issue?

[calpt]

Resolved in #658.