How to correct TypeError: zip argument #1 must support iteration training in multiple GPU

pchhapolika: <h3>
<a name="describe-the-bug-1" class="anchor" href="#describe-the-bug-1"></a><img src="https://discuss.pytorch.org/images/emoji/apple/bug.png?v=12" title=":bug:" class="emoji" alt=":bug:" loading="lazy" width="20" height="20"> Describe the bug</h3>
<p>I am doing a creating custom pytorch layer and model training using <code>Trainer API</code> function on top of <code>Hugging face</code> model.</p>
<p>When I run on <code>single GPU</code>, it trains fine. But when I train it on <code>multiple GPU</code> it throws me error.</p>
<p><code>TypeError: zip argument #1 must support iteration training in multiple GPU</code></p>
<p><strong>Training Code</strong></p>
<pre><code>bert_model = BertForTokenClassification.from_pretrained( model_checkpoint,id2label=id2label,label2id=label2id)
bert_model.config.output_hidden_states=True


class BERT_CUSTOM(nn.Module):
    
    
    def __init__(self, bert_model,id2label,num_labels):
        
        
        
        super(BERT_CUSTOM, self).__init__()
        self.bert = bert_model
        self.config=self.bert.config
        self.dropout = nn.Dropout(0.25)
        self.classifier = nn.Linear(768, num_labels)
        self.crf = CRF(num_labels, batch_first = True)
        
    
    def forward(self, input_ids, attention_mask,  labels=None, token_type_ids=None):
        
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        sequence_output = torch.stack((outputs[1][-1], outputs[1][-2], outputs[1][-3], outputs[1][-4])).mean(dim=0)
        sequence_output = self.dropout(sequence_output)
        emission = self.classifier(sequence_output) # [32,256,21] logits
        
        if labels is not None:
            
            labels=labels.reshape(attention_mask.size()[0],attention_mask.size()[1])
            loss = -self.crf(log_soft(emission, 2), labels, mask=attention_mask.type(torch.uint8), reduction='mean')
            prediction = self.crf.decode(emission, mask=attention_mask.type(torch.uint8))
            return [loss, prediction]
                
        else:
            
            prediction = self.crf.decode(emission, mask=attention_mask.type(torch.uint8))
            prediction=[id2label[k] for k in prediction]
            return prediction
</code></pre>
<p><strong>Training API</strong></p>
<pre><code>model = BERT_CUSTOM(bert_model, id2label,num_labels=len(label2id))
model.to(device)

args = TrainingArguments(
    "model",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=2,
    weight_decay=0.01,
    per_device_train_batch_size=32,
    fp16=True
    
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_data,
    tokenizer=tokenizer)

trainer.train()
</code></pre>
<h3>
<a name="versions-2" class="anchor" href="#versions-2"></a>Versions</h3>
<p>‘1.7.1+cu110’</p>

pchhapolika: <p>Hi <a href="https://discuss.pytorch.org/u/ptrblck">ptrblck</a>, could you please help.?</p>

pchhapolika: <p>Here is the data creation code:</p>
<pre><code class="lang-auto">train_ex ={'texts':[x[0] for x in train_set],'tag_names':[x[1] for x in train_set]}
train_data = tokenize_and_align_labels(train_ex,label2id)
_=train_data.pop('offset_mapping')

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, examples):
        self.encodings = examples        
        self.labels = examples['labels']
    def __getitem__(self, idx):
        item = {k: torch.tensor(v[idx]) for k, v in self.encodings.items()}
        item["labels"] = torch.tensor([self.labels[idx]])
        return item    def __len__(self):
        return len(self.labels)
train_data=MyDataset(train_data)
</code></pre>
<p>And I use this <code>train_data</code> directly in model.</p>

ptrblck: <p>I’m not deeply familiar with the HuggingFace <code>Trainer</code> class, but could you post the full stacktrace here, please?</p>

pchhapolika: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "spanbert_model_check.py", line 263, in &lt;module&gt;
    trainer.train()
  File "/opt/conda/lib/python3.7/site-packages/transformers/trainer.py", line 1531, in train
    ignore_keys_for_eval=ignore_keys_for_eval,
  File "/opt/conda/lib/python3.7/site-packages/transformers/trainer.py", line 1775, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/lib/python3.7/site-packages/transformers/trainer.py", line 2523, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/lib/python3.7/site-packages/transformers/trainer.py", line 2555, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 162, in forward
    return self.gather(outputs, self.output_device)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 174, in gather
    return gather(outputs, output_device, dim=self.dim)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 68, in gather
    res = gather_map(outputs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 63, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 63, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py", line 63, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
  [Previous line repeated 1 more time]
TypeError: zip argument #1 must support iteration
</code></pre>

ptrblck: <p>Thanks for the stacktrace! I don’t see any obvious issue but it seems an internal <code>DataParallel</code> call fails. I would recommend updating PyTorch and <code>transformers</code> to the latest release to see if you are still running into the same issue.<br>
I also see you’ve cross-posted the issue in the <a href="https://discuss.huggingface.co/t/how-to-correct-typeerror-zip-argument-1-must-support-iteration-training-in-multiple-gpu/31489">HF discussion board</a> but didn’t receive any responses yet.<br>
A related issue can be found in the <a href="https://forums.fast.ai/t/how-to-use-multiple-gpus/26057/5">FastAI discussion board</a> where a user claimed:</p>
<blockquote>
<p>Turns out multi-gpu works perfectly with fastai. All one has to do is <code>model=torch.nn.DataParallel(model)</code> and then pass this model to the <code>Learner</code> object.<br>
The error was because I was passing a list instead of tensors which was causing the error reported.</p>
</blockquote>

ninetyfour2010: <aside class="quote no-group" data-username="pchhapolika" data-post="1" data-topic="172448">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/pchhapolika/48/57278_2.png" class="avatar"> pchhapolika:</div>
<blockquote>
<p><code>loss</code></p>
</blockquote>
</aside>
<p>let your  loss be tensor maybe solve your problem</p>