fix: ddp

Author: sigridjineth

src/tevatron/reranker/driver/train.py

@@ -1,20 +1,29 @@
 import logging
 import os
 import sys
+import torch
 from transformers import AutoTokenizer
 from transformers import (
     HfArgumentParser,
     set_seed,
 )
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
 from tevatron.reranker.arguments import ModelArguments, DataArguments, TevatronTrainingArguments
 from tevatron.reranker.modeling import RerankerModel
 from tevatron.reranker.dataset import RerankerTrainDataset
 from tevatron.reranker.collator import RerankerTrainCollator
-from tevatron.reranker.trainer import RerankerTrainer  # Make sure this is your updated RerankerTrainer
+from tevatron.reranker.trainer import RerankerTrainer
 
 logger = logging.getLogger(__name__)
 
 
+def setup_ddp():
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+
 def main():
     parser = HfArgumentParser((ModelArguments, DataArguments, TevatronTrainingArguments))
 
@@ -23,15 +32,8 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    if (
-            os.path.exists(training_args.output_dir)
-            and os.listdir(training_args.output_dir)
-            and training_args.do_train
-            and not training_args.overwrite_output_dir
-    ):
-        raise ValueError(
-            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
-        )
+    if training_args.local_rank != -1:
+        setup_ddp()
 
     # Setup logging
     logging.basicConfig(
@@ -67,10 +69,12 @@ def main():
         cache_dir=model_args.cache_dir,
     )
 
+    if training_args.local_rank != -1:
+        model = DDP(model, device_ids=[training_args.local_rank], output_device=training_args.local_rank)
+
     train_dataset = RerankerTrainDataset(data_args)
     train_collator = RerankerTrainCollator(data_args, tokenizer)
 
-    # Add GradCache-specific arguments to training_args
     training_args.gc_chunk_size = getattr(training_args, 'gc_chunk_size', 2)
 
     trainer = RerankerTrainer(
@@ -81,7 +85,7 @@ def main():
     )
     train_dataset.trainer = trainer
 
-    trainer.train()  # TODO: resume training
+    trainer.train()
     trainer.save_model()
     if trainer.is_world_process_zero():
         tokenizer.save_pretrained(training_args.output_dir)

src/tevatron/reranker/trainer.py

@@ -7,46 +7,43 @@
 from transformers.trainer_utils import PredictionOutput
 
 from grad_cache import GradCache
-
 from grad_cache.functional import cached, cat_input_tensor
 from torch.cuda.amp import autocast
 
 logger = logging.getLogger(__name__)
 
-
 @cached
 @autocast()
 def get_model_rep(model, inputs):
     outputs = model(**inputs)
     return outputs.scores
 
-
 @cat_input_tensor
 @autocast()
 def contrastive_loss(scores):
     batch_size = scores.size(0) // 2
     labels = torch.arange(batch_size, device=scores.device)
     return nn.CrossEntropyLoss()(scores, labels)
 
-
 def split_inputs(model_input, chunk_size):
     logger.debug(f"Splitting inputs with chunk size: {chunk_size}")
     keys = list(model_input.keys())
     chunked_tensors = [model_input[k].split(chunk_size, dim=0) for k in keys]
     return [dict(zip(keys, tt)) for tt in zip(*chunked_tensors)]
 
-
 class RerankerTrainer(Trainer):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         logger.info("Initializing RerankerTrainer with GradCache")
         self.args: TrainingArguments
 
-        # Add these lines to include the necessary parameters
-        self.gc_chunk_size = getattr(self.args, 'gc_chunk_size', 4)  # default to 4 if not provided
+        self.gc_chunk_size = getattr(self.args, 'gc_chunk_size', 4)
+
+        # If the model is wrapped in DDP, we need to use the .module attribute
+        model_for_gc = self.model.module if hasattr(self.model, 'module') else self.model
 
         self.gc = GradCache(
-            models=[self.model],
+            models=[model_for_gc],
             chunk_sizes=self.gc_chunk_size,
             loss_fn=contrastive_loss,
             split_input_fn=split_inputs,
@@ -68,17 +65,17 @@ def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor,
         logger.debug("Entering training step")
         model.train()
         inputs = self._prepare_inputs(inputs)
-        _distributed = self.args.local_rank > -1
+        _distributed = self.args.local_rank != -1
         loss = self.gc(inputs, no_sync_except_last=_distributed)
         logger.debug(f"Training step loss: {loss.item()}")
         return loss
 
     def prediction_step(
-            self,
-            model: nn.Module,
-            inputs: Dict[str, Union[torch.Tensor, Any]],
-            prediction_loss_only: bool,
-            ignore_keys: bool = None,
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: bool = None,
     ) -> PredictionOutput:
         logger.debug("Entering prediction step")
         inputs = self._prepare_inputs(inputs)
@@ -87,4 +84,4 @@ def prediction_step(
         scores = outputs.scores
         loss = contrastive_loss(scores)
         logger.debug(f"Prediction step loss: {loss.item() if loss is not None else 'N/A'}")
-        return PredictionOutput(predictions=scores, label_ids=None, metrics=None)
+        return PredictionOutput(predictions=scores, label_ids=None, metrics=None)