text: 
  
  TITLE
  Model Parallelism for Bert Models
  
  COMMENTS
  22
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  Hi,
  
  I'm trying to implement Model parallelism for BERT models by splitting and assigning layers across GPUs. I took DeBERTa as an example for this. 
  For DeBERTa, I'm able to split entire model into 'embedding', 'encoder', 'pooler', 'classifier' and 'dropout' layers as shown in below pic.
  
  ![Capture](https://user-images.githubusercontent.com/41769919/107767327-d3090500-6d5a-11eb-98b3-9dfd7b1c78fb.PNG)
  
  With this approach, I trained on IMDB classification task by assigning 'encoder' to second GPU and others to first 'GPU'. At the end of the training, second GPU consumed lot of memory when compared to first GPU and this resulted in 20-80 split of the entire model.
  
  So, I tried splitting encoder layers also as shown below but getting this error - **"TypeError: forward() takes 1 positional argument but 2 were given"**
  
  ```
  embed = dberta.deberta.embeddings.to('cuda:0')
  
  f6e = dberta.deberta.encoder.layer[:6].to('cuda:0')
  
  l6e = dberta.deberta.encoder.layer[6:].to('cuda:1')
  
  pooler = dberta.pooler.to('cuda:0')
  
  classifier = dberta.classifier.to('cuda:0')
  
  dropout = dberta.dropout.to('cuda:0')
  
  test = "this is to test deberta"
  
  inp_ids = tok_dberta(test, return_tensors='pt').input_ids
  att_mask = tok_dberta(test, return_tensors='pt').attention_mask
  
  emb_out = embed(inp_ids.to('cuda:0'))
  
  first_6_enc_lay_out = f6e(emb_out)
  ---------------------------------------------------------------------------
  TypeError                                 Traceback (most recent call last)
  <ipython-input-15-379d948e5ba5> in <module>
  ----> 1 first_6_enc_lay_out = f6e(emb_out)
  
  /usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
      725             result = self._slow_forward(*input, **kwargs)
      726         else:
  --> 727             result = self.forward(*input, **kwargs)
      728         for hook in itertools.chain(
      729                 _global_forward_hooks.values(),
  
  TypeError: forward() takes 1 positional argument but 2 were given
  
  ```
  
  Plz suggest how to proceed further..