[BlenderBot2] add model.module to make BB2 DDP compatible

projects/blenderbot2/agents/blenderbot2.py

@@ -800,7 +800,10 @@ def _set_batch_skip_search(self, valid_exs: List[Message], batch: Batch) -> Batc
 
     def eval_step(self, batch):
         output = super().eval_step(batch)
-        if output is None or not hasattr(self.model, 'retriever'):
+        model = self.model
+        if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
+            model = self.model.module
+        if output is None or not hasattr(model, 'retriever'):
             return output
         if hasattr(self.model_api.retriever, 'top_docs'):
             output.top_docs = self.model_api.retriever.top_docs
@@ -855,7 +858,10 @@ def compute_loss(
         Override Rag.compute_loss to add some additional metrics.
         """
         loss, output = super().compute_loss(batch, return_output=True)
-        assert isinstance(self.model, BlenderBot2RagModel)
+        model = self.model
+        if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
+            model = self.model.module
+        assert isinstance(model, BlenderBot2RagModel)
         if (
             KnowledgeAccessMethod(self.opt['knowledge_access_method'])
             is KnowledgeAccessMethod.CLASSIFY

projects/msc/agents/memory_agent.py

@@ -205,6 +205,10 @@ def __init__(self, opt, dictionary, retriever_shared=None):
         if opt.get('fid_ddp_compatible', True):
             for param in self.long_term_memory.query_encoder.parameters():
                 param.requires_grad = False
+            for param in self.long_term_memory.memory_encoder.parameters():
+                param.requires_grad = False
+            for param in self.retriever.parameters():
+                param.requires_grad = False
 
 
 class MemoryLongFidAgent(LongFidAgent, MemoryRagAgent):