distribute train

Author: Wonters

README.md

@@ -22,3 +22,12 @@ Credentials are stored in ~/.config/ovhai/context.json
 ```bash
 uv pip install boto3 awscli ovhai
 ```
+
+## Run on multi GPU
+DEBUG
+```bash
+export TORCH_DISTRIBUTED_DEBUG=DETAIL
+```
+```bash
+python -m torch.distributed.run --nproc_per_node=2 train.py
+```

src/mixins.py

@@ -7,8 +7,9 @@
 import logging
 import optuna
 import random
-from torch.utils.data import Subset, DataLoader
+from torch.utils.data import Subset, DataLoader, DistributedSampler
 import torch.distributed as dist
+import pandas as pd
 
 logger = logging.getLogger(__name__)
 
@@ -22,11 +23,15 @@ class TorchModelTrainMixin:
     lr: float = 2e-5
     device: torch.device
 
-    def get_sampled_dataloader(self, frac=0.1):
+    def sample_dataset(self, frac=0.1):
         dataset_size = len(self.dataset)
         sample_size = int(frac * dataset_size)
         indices = random.sample(range(dataset_size), sample_size)
         sampled_dataset = Subset(self.dataset, indices)
+        return sampled_dataset
+
+    def get_sampled_dataloader(self, frac=0.1):
+        sampled_dataset = self.sample_dataset(frac=frac)
         return DataLoader(sampled_dataset, batch_size=self.batch_size, shuffle=True)
 
     def optuna_train(self, run_name:str = "", n_trials:int=30, frac=0.1):
@@ -53,6 +58,13 @@ def objective(self, trial):
             self.reinit_scheduler_optimizer(**kwargs)
             acc = self.train()
         return acc
+    
+    def get_ddp_dataloader(self, frac=1.0):
+        sampled_dataset = self.sample_dataset(frac=frac)
+        sampler = DistributedSampler(sampled_dataset)
+        dataloader = DataLoader(sampled_dataset, batch_size=self.batch_size, sampler=sampler)
+        return dataloader, sampler
+    
 
     def _train_batch(self, x, y):
         inputs = self.tokenizer(x, return_tensors="pt", truncation=True, padding=True)
@@ -103,7 +115,6 @@ def train(self):
                     current_acc = self._train_batch(tweets, labels.float())
                     acc.append(current_acc)
                 except RuntimeError as e:
-                    raise e
                     logger.error(e)
                     del tweets, labels, self.optimizer
                     gc.collect()

src/ml.py

@@ -14,7 +14,7 @@
 import torch.nn as nn
 import torch.distributed as dist
 from tqdm import tqdm
-from torch.utils.data import DataLoader, DistributedSampler
+from torch.utils.data import DataLoader
 import torch.nn.functional as F
 from sklearn.linear_model import LogisticRegression
 import lightgbm as lgm
@@ -206,6 +206,13 @@ class TorchBaseModel(TorchModelTrainMixin, BaseModelABC):
     """
     Base class to train and predict on a dataset and register data on MLFLow
     """
+
+    def __init__(self, dataset: pd.DataFrame):
+        super().__init__(dataset)
+        if dist.is_initialized():
+            self.dataloader, self.sampler = self.get_ddp_dataloader()
+            logger.info(f"Rank {dist.get_rank()} using DDP")
+
     def preprocessing(self, data):
         return self.tokenizer(list(data), return_tensors="pt", truncation=True, padding=True)
 
@@ -441,10 +448,6 @@ class LSTMModel(TorchBaseModel):
     device = DEVICE
     # torch.nn.CrossEntropyLoss()
 
-    def __init__(self, dataset: pd.DataFrame):
-        super().__init__(dataset)
-        self.dataset = DistributedSampler(self.dataset)
-
     @property
     def get_metrics(self) -> dict:
         for k, v in self.model.state_dict().items():