From 59cebfcf0db821cab70e7a252e6df1028d75eb0b Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Fri, 6 Feb 2026 14:41:31 -0800 Subject: [PATCH 1/3] feat: migrate from DataParallel to DDP via torchrun (#45) Replace hardcoded DEVICE=cuda:0 with model.device so each DDP process uses its assigned GPU. Switch launch script to torchrun. Fixes GPU 0 OOM on multi-GPU setups. --- modeling/common.py | 2 -- modeling/model.py | 33 +++++++++++++++++---------------- modeling/pretrain.py | 17 ++++++++++++----- run_pretrain.sh | 5 ++--- 4 files changed, 31 insertions(+), 26 deletions(-) diff --git a/modeling/common.py b/modeling/common.py index a3859d7..a42665f 100644 --- a/modeling/common.py +++ b/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 diff --git a/modeling/model.py b/modeling/model.py index 8a91414..97ab6e3 100644 --- a/modeling/model.py +++ b/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( diff --git a/modeling/pretrain.py b/modeling/pretrain.py index 0b78c76..a3c5ead 100644 --- a/modeling/pretrain.py +++ b/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( diff --git a/run_pretrain.sh b/run_pretrain.sh index 80e76e5..d67ff65 100755 --- a/run_pretrain.sh +++ b/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 From 35f8834fc082fb0214af2f96ea50cdb3c6752255 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Fri, 6 Feb 2026 14:43:38 -0800 Subject: [PATCH 2/3] doc: torchrun --- README.md | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 4928824..f49d22f 100755 --- a/README.md +++ b/README.md @@ -101,11 +101,19 @@ 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`. By default it uses all visible GPUs; control GPU selection with `CUDA_VISIBLE_DEVICES`. + ```sh -# From a YAML config (recommended) +# Multi-GPU (uses all visible GPUs) +torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml + +# Select specific GPUs +CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml + +# Single-GPU (also works with plain python) python modeling/cli.py run experiments/base.yaml -# Or pass all parameters directly +# Pass all parameters directly python modeling/cli.py pretrain --batch-size 64 --num-epochs 3 --model-name ./saved_models/ContraBERT_G # See all options @@ -113,12 +121,14 @@ python modeling/cli.py run --help python modeling/cli.py pretrain --help ``` -There is also a convenience script: +There is also a convenience script that launches `torchrun` on all visible GPUs: ```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 YAML experiment configs live in `experiments/`. To create a new experiment, copy an existing file and modify the parameters: From 0ceb19e84c30bc62011e5eb8b4cf7ceaaf86bfe0 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Fri, 6 Feb 2026 15:33:35 -0800 Subject: [PATCH 3/3] add instruction for multi-gpu --- README.md | 39 ++++++++++++++++++++++++++++----------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index f49d22f..cc2eb24 100755 --- a/README.md +++ b/README.md @@ -101,19 +101,42 @@ 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`. By default it uses all visible GPUs; control GPU selection with `CUDA_VISIBLE_DEVICES`. +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 -# Multi-GPU (uses all visible GPUs) torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml +``` + +To select specific GPUs, set `CUDA_VISIBLE_DEVICES`: -# Select specific GPUs +```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 -# Single-GPU (also works with plain python) +On a single-GPU machine, run with plain `python` -- no `torchrun` needed: + +```sh python modeling/cli.py run experiments/base.yaml +``` -# Pass all parameters directly +`torchrun --nproc_per_node=1` also works if you prefer a uniform launch command across environments. + +#### 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 @@ -121,12 +144,6 @@ python modeling/cli.py run --help python modeling/cli.py pretrain --help ``` -There is also a convenience script that launches `torchrun` on all visible GPUs: - -```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