Repository: deepset-ai/haystack

Issue:
Data pipeline parallelism doesn't work for multiple GPUs when trying to fine-tune DPR. 
**Bug description**
Data pipeline parallelism doesn't work for multiple GPUs when trying to fine-tune DPR. This is because the biadaptive model instance is not called as a module. The bug disappears if you call the biadaptive model as module.biadaptive. This caused us to reuse and override unnecessary code from haystack in our code base. 

**Error message**
It throws the error of "input and label ids not found to pass"

**Expected behavior**
Smooth fine-tuning without overhead of overriding functions

**To Reproduce**
Try to fine-tune the retriever in multi-gpu environment.

Our adapted code which solves the problem.
![image](https://user-images.githubusercontent.com/27398711/194812170-6f965ea8-67f9-462b-9981-b29be2c02660.png)




Discussion:
Hi @sinchanabhat thank you for reporting this bug and also for sharing how you fixed it! 👍 I would like to invite you to make a pull request with your fix if you have the time to do it. Just let us know. 🙂 
Hey @julian-risch I'd be happy to open a PR for this since this also affects the Pytorch AMP PR for multi-gpu training https://github.com/deepset-ai/haystack/pull/2827
That would be great! 👍 Please assign the issue to yourself once you start working on it and then remove the "Contributions wanted!" tag. Thank you!

Resolving patch:
diff --git a/haystack/modeling/evaluation/eval.py b/haystack/modeling/evaluation/eval.py
--- a/haystack/modeling/evaluation/eval.py
+++ b/haystack/modeling/evaluation/eval.py
@@ -3,12 +3,14 @@
 import logging
 import numbers
 import torch
+from torch.nn import DataParallel
 import numpy as np
 from tqdm import tqdm
 
 from haystack.modeling.evaluation.metrics import compute_metrics, compute_report_metrics
 from haystack.modeling.model.adaptive_model import AdaptiveModel
 from haystack.modeling.model.biadaptive_model import BiAdaptiveModel
+from haystack.modeling.model.optimization import WrappedDataParallel
 from haystack.utils.experiment_tracking import Tracker as tracker
 from haystack.modeling.visual import BUSH_SEP
 
@@ -70,9 +72,13 @@ def eval(
         for step, batch in enumerate(tqdm(self.data_loader, desc="Evaluating", mininterval=10)):
             batch = {key: batch[key].to(self.device) for key in batch}
 
-            with torch.no_grad():
+            if isinstance(model, (DataParallel, WrappedDataParallel)):
+                module = model.module
+            else:
+                module = model
 
-                if isinstance(model, AdaptiveModel):
+            with torch.no_grad():
+                if isinstance(module, AdaptiveModel):
                     logits = model.forward(
                         input_ids=batch.get("input_ids", None),
                         segment_ids=batch.get("segment_ids", None),
@@ -80,7 +86,7 @@ def eval(
                         output_hidden_states=batch.get("output_hidden_states", False),
                         output_attentions=batch.get("output_attentions", False),
                     )
-                elif isinstance(model, BiAdaptiveModel):
+                elif isinstance(module, BiAdaptiveModel):
                     logits = model.forward(
                         query_input_ids=batch.get("query_input_ids", None),
                         query_segment_ids=batch.get("query_segment_ids", None),
diff --git a/haystack/modeling/training/base.py b/haystack/modeling/training/base.py
--- a/haystack/modeling/training/base.py
+++ b/haystack/modeling/training/base.py
@@ -18,7 +18,7 @@
 from haystack.modeling.evaluation.eval import Evaluator
 from haystack.modeling.model.adaptive_model import AdaptiveModel
 from haystack.modeling.model.biadaptive_model import BiAdaptiveModel
-from haystack.modeling.model.optimization import get_scheduler
+from haystack.modeling.model.optimization import get_scheduler, WrappedDataParallel
 from haystack.modeling.utils import GracefulKiller
 from haystack.utils.experiment_tracking import Tracker as tracker
 from haystack.utils.early_stopping import EarlyStopping
@@ -292,12 +292,17 @@ def train(self):
 
     def compute_loss(self, batch: dict, step: int) -> torch.Tensor:
         # Forward & backward pass through model
-        if isinstance(self.model, AdaptiveModel):
+        if isinstance(self.model, (DataParallel, WrappedDataParallel)):
+            module = self.model.module
+        else:
+            module = self.model
+
+        if isinstance(module, AdaptiveModel):
             logits = self.model.forward(
                 input_ids=batch["input_ids"], segment_ids=None, padding_mask=batch["padding_mask"]
             )
 
-        elif isinstance(self.model, BiAdaptiveModel):
+        elif isinstance(module, BiAdaptiveModel):
             logits = self.model.forward(
                 query_input_ids=batch["query_input_ids"],
                 query_segment_ids=batch["query_segment_ids"],


Test patch:
diff --git a/test/modeling/test_dpr.py b/test/modeling/test_dpr.py
--- a/test/modeling/test_dpr.py
+++ b/test/modeling/test_dpr.py
@@ -1,6 +1,6 @@
+import os
 from typing import Tuple
 
-import os
 import logging
 from pathlib import Path
 
@@ -17,8 +17,11 @@
 from haystack.modeling.model.language_model import get_language_model, DPREncoder
 from haystack.modeling.model.prediction_head import TextSimilarityHead
 from haystack.modeling.model.tokenization import get_tokenizer
+from haystack.nodes.retriever.dense import DensePassageRetriever
 from haystack.modeling.utils import set_all_seeds, initialize_device_settings
 
+from ..conftest import SAMPLES_PATH
+
 
 def test_dpr_modules(caplog=None):
     if caplog:
@@ -970,6 +973,33 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa
     assert np.array_equal(all_embeddings["query"][0], all_embeddings3["query"][0])
 
 
+@pytest.mark.parametrize("document_store", ["memory"], indirect=True)
+def test_dpr_training(document_store, tmp_path):
+    retriever = DensePassageRetriever(
+        document_store=document_store,
+        query_embedding_model="facebook/dpr-question_encoder-single-nq-base",
+        passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base",
+        max_seq_len_query=8,
+        max_seq_len_passage=8,
+    )
+
+    save_dir = f"{tmp_path}/test_dpr_training"
+    retriever.train(
+        data_dir=str(SAMPLES_PATH / "dpr"),
+        train_filename="sample.json",
+        dev_filename="sample.json",
+        test_filename="sample.json",
+        n_epochs=1,
+        batch_size=1,
+        grad_acc_steps=1,
+        save_dir=save_dir,
+        evaluate_every=10,
+        embed_title=True,
+        num_positives=1,
+        num_hard_negatives=1,
+    )
+
+
 # TODO fix CI errors (test pass locally or on AWS, next steps: isolate PyTorch versions once FARM dependency is removed)
 # def test_dpr_training():
 #     batch_size = 1
@@ -982,8 +1012,6 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa
 #     use_fast = True
 #     similarity_function = "dot_product"
 #
-#
-#
 #     device, n_gpu = initialize_device_settings(use_cuda=False)
 #
 #     query_tokenizer = get_tokenizer(pretrained_model_name_or_path=question_lang_model,