feat: distribute data parallel

README.md

@@ -101,23 +101,50 @@ The CLI entry point is `modeling/cli.py`, which provides two subcommands:
 - **`run`** -- load a YAML experiment config (recommended; all parameters come from the config file to ensure full reproducibility)
 - **`pretrain`** -- pass all parameters directly as CLI options
 
+Pre-training uses PyTorch DistributedDataParallel (DDP) via `torchrun`, which ships with PyTorch itself (no extra dependencies). The same training code runs on both multi-GPU and single-GPU nodes -- the HuggingFace `Trainer` auto-detects the distributed environment set up by `torchrun` and enables or disables DDP accordingly.
+
+#### Multi-GPU node
+
+Use `torchrun` to spawn one process per GPU. `--nproc_per_node=gpu` automatically uses all visible GPUs:
+
+```sh
+torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml
+```
+
+To select specific GPUs, set `CUDA_VISIBLE_DEVICES`:
+
+```sh
+CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml
+```
+
+There is also a convenience script that launches `torchrun` on all visible GPUs:
+
+```sh
+./run_pretrain.sh
+```
+
+#### Single-GPU node
+
+On a single-GPU machine, run with plain `python` -- no `torchrun` needed:
+
 ```sh
-# From a YAML config (recommended)
 python modeling/cli.py run experiments/base.yaml
+```
+
+`torchrun --nproc_per_node=1` also works if you prefer a uniform launch command across environments.
 
-# Or pass all parameters directly
+#### CLI examples
+
+```sh
+# Pass all parameters directly (without a YAML config)
 python modeling/cli.py pretrain --batch-size 64 --num-epochs 3 --model-name ./saved_models/ContraBERT_G
 
 # See all options
 python modeling/cli.py run --help
 python modeling/cli.py pretrain --help
 ```
 
-There is also a convenience script:
-
-```sh
-./run_pretrain.sh
-```
+The `batch_size` in config is the **total** batch size across all GPUs. It is automatically divided by the number of processes. For example, `batch_size: 128` on 2 GPUs gives 64 per GPU; with `gradient_accumulation_steps: 2` the effective batch size is 256.
 
 #### Experiment Configs
 

modeling/common.py

@@ -4,8 +4,6 @@
 import numpy as np
 import torch
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 MAX_NUM_PROC = 80
 
 

modeling/model.py

@@ -3,7 +3,6 @@
 from transformers import Trainer
 
 from ._types import ContraMode
-from .common import DEVICE
 
 
 def info_nce_loss(query, key, temperature=0.07):
@@ -227,13 +226,14 @@ def compute_loss(
         if self.contra_mode == ContraMode.GROUPED:
             return self._compute_grouped_loss(model, inputs, return_outputs)
 
-        # Move inputs to device
-        code_input_ids = inputs["code_input_ids"].to(DEVICE)
-        code_attention_mask = inputs["code_attention_mask"].to(DEVICE)
-        code_labels = inputs["code_labels"].to(DEVICE)
-        aug_input_ids = inputs["aug_input_ids"].to(DEVICE)
-        aug_attention_mask = inputs["aug_attention_mask"].to(DEVICE)
-        aug_labels = inputs["aug_labels"].to(DEVICE)
+        # Move inputs to the device the model lives on (supports DDP)
+        device = model.device
+        code_input_ids = inputs["code_input_ids"].to(device)
+        code_attention_mask = inputs["code_attention_mask"].to(device)
+        code_labels = inputs["code_labels"].to(device)
+        aug_input_ids = inputs["aug_input_ids"].to(device)
+        aug_attention_mask = inputs["aug_attention_mask"].to(device)
+        aug_labels = inputs["aug_labels"].to(device)
 
         # Forward pass for MLM
         # use bi-encoder training, encode code and augmentation separately using self.model
@@ -267,7 +267,7 @@ def compute_loss(
         # Compute contrastive loss between code and its augmentation
         if self.contra_mode == ContraMode.SUPCON:
             all_embeddings = torch.cat([code_embeddings, aug_embeddings], dim=0)
-            function_ids = inputs["function_id"].to(DEVICE)
+            function_ids = inputs["function_id"].to(device)
             all_function_ids = torch.cat([function_ids, function_ids], dim=0)
             contrastive_loss = supcon_loss(
                 all_embeddings, all_function_ids, self.temperature
@@ -294,13 +294,14 @@ def _compute_grouped_loss(self, model, inputs, return_outputs=False):
           - aug_attention_mask, aug_labels: same shape
           - group_sizes: [B]
         """
-        code_input_ids = inputs["code_input_ids"].to(DEVICE)
-        code_attention_mask = inputs["code_attention_mask"].to(DEVICE)
-        code_labels = inputs["code_labels"].to(DEVICE)
-        aug_input_ids = inputs["aug_input_ids"].to(DEVICE)
-        aug_attention_mask = inputs["aug_attention_mask"].to(DEVICE)
-        aug_labels = inputs["aug_labels"].to(DEVICE)
-        group_sizes = inputs["group_sizes"].to(DEVICE)
+        device = model.device
+        code_input_ids = inputs["code_input_ids"].to(device)
+        code_attention_mask = inputs["code_attention_mask"].to(device)
+        code_labels = inputs["code_labels"].to(device)
+        aug_input_ids = inputs["aug_input_ids"].to(device)
+        aug_attention_mask = inputs["aug_attention_mask"].to(device)
+        aug_labels = inputs["aug_labels"].to(device)
+        group_sizes = inputs["group_sizes"].to(device)
 
         # Forward anchor
         code_outputs = model(

modeling/pretrain.py

@@ -4,8 +4,9 @@
 import os
 from collections import defaultdict
 
+import torch
+import torch.distributed as dist
 from datasets import Dataset, Features, Value, load_dataset
-from torch.cuda import device_count
 from transformers import (
     DataCollatorForLanguageModeling,
     RobertaConfig,
@@ -15,11 +16,18 @@
 )
 
 from ._types import ContraMode
-from .common import DEVICE, default_num_proc, set_seed
+from .common import default_num_proc, set_seed
 from .dataloader import contra_data_collator, grouped_contra_data_collator
 from .model import ContrastiveTrainer
 
 
+def _get_world_size() -> int:
+    """Return DDP world size if distributed is initialized, else GPU count (min 1)."""
+    if dist.is_initialized():
+        return dist.get_world_size()
+    return max(torch.cuda.device_count(), 1)
+
+
 def compute_function_id(code: str) -> int:
     """Deterministic 63-bit hash of the code string (positive, fits int64)."""
     digest = hashlib.sha256(code.encode("utf-8")).digest()
@@ -234,7 +242,6 @@ def main(
         model_name if checkpoint is None else checkpoint,
         config=config,
     )  # load weights from stage 1
-    model.to(DEVICE)
 
     features = Features(
         {
@@ -290,7 +297,7 @@ def main(
     training_args = TrainingArguments(
         output_dir=f"./saved_models/{run_name}",
         overwrite_output_dir=True,
-        per_device_train_batch_size=batch_size // device_count(),
+        per_device_train_batch_size=batch_size // _get_world_size(),
         gradient_accumulation_steps=gradient_accumulation_steps,
         num_train_epochs=num_epochs,
         save_strategy="epoch",
@@ -304,7 +311,7 @@ def main(
         run_name=run_name,
         save_total_limit=3,
         load_best_model_at_end=True,
-        dataloader_num_workers=os.cpu_count(),
+        dataloader_num_workers=max(1, (os.cpu_count() or 1) // _get_world_size()),
     )
 
     trainer = ContrastiveTrainer(

run_pretrain.sh

@@ -3,6 +3,5 @@
 # export CUDA_VISIBLE_DEVICES=4,5,6,7
 export WANDB_PROJECT="InvPT"
 
-# Use a YAML config; override specific values with CLI options if needed:
-#   python modeling/cli.py run experiments/base.yaml --seed 42
-python modeling/cli.py run experiments/base.yaml
+# Use torchrun for DDP; uses all visible GPUs (control with CUDA_VISIBLE_DEVICES).
+torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml