diff --git a/README.md b/README.md index 4928824..cc2eb24 100755 --- a/README.md +++ b/README.md @@ -101,11 +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`, 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 @@ -113,11 +144,7 @@ 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 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