diff --git a/experiments/grouped/codebert.yaml b/experiments/grouped/codebert.yaml index addef07..8bcf1e9 100644 --- a/experiments/grouped/codebert.yaml +++ b/experiments/grouped/codebert.yaml @@ -5,9 +5,9 @@ dataset_path: "data/aug_csn.jsonl" model_name: "microsoft/codebert-base" -batch_size: 16 +batch_size: 64 num_epochs: 3 -gradient_accumulation_steps: 16 +gradient_accumulation_steps: 4 learning_rate: 2.0e-5 seed: 0 diff --git a/experiments/grouped/contrabert_c.yaml b/experiments/grouped/contrabert_c.yaml index 1f8e87c..d5f8f4d 100644 --- a/experiments/grouped/contrabert_c.yaml +++ b/experiments/grouped/contrabert_c.yaml @@ -6,9 +6,9 @@ dataset_path: "data/aug_csn.jsonl" model_name: "./saved_models/ContraBERT_C" tokenizer_name: "microsoft/codebert-base" -batch_size: 16 +batch_size: 64 num_epochs: 3 -gradient_accumulation_steps: 16 +gradient_accumulation_steps: 4 learning_rate: 2.0e-5 seed: 0 diff --git a/experiments/grouped/contrabert_g.yaml b/experiments/grouped/contrabert_g.yaml index 0b8db32..392420a 100644 --- a/experiments/grouped/contrabert_g.yaml +++ b/experiments/grouped/contrabert_g.yaml @@ -6,9 +6,9 @@ dataset_path: "data/aug_csn.jsonl" model_name: "./saved_models/ContraBERT_G" tokenizer_name: "microsoft/graphcodebert-base" -batch_size: 16 +batch_size: 64 num_epochs: 3 -gradient_accumulation_steps: 16 +gradient_accumulation_steps: 4 learning_rate: 2.0e-5 seed: 0 diff --git a/experiments/grouped/graphcodebert.yaml b/experiments/grouped/graphcodebert.yaml index 0e2d204..7e235df 100644 --- a/experiments/grouped/graphcodebert.yaml +++ b/experiments/grouped/graphcodebert.yaml @@ -5,9 +5,9 @@ dataset_path: "data/aug_csn.jsonl" model_name: "microsoft/graphcodebert-base" -batch_size: 16 +batch_size: 64 num_epochs: 3 -gradient_accumulation_steps: 16 +gradient_accumulation_steps: 4 learning_rate: 2.0e-5 seed: 0 diff --git a/modeling/common.py b/modeling/common.py index a42665f..ed9cbe3 100644 --- a/modeling/common.py +++ b/modeling/common.py @@ -9,7 +9,7 @@ def default_num_proc() -> int: """Return the default number of parallel workers, capped at available CPUs.""" - return min(os.cpu_count() or 1, MAX_NUM_PROC) + return os.cpu_count() or 1 def set_seed(seed): diff --git a/modeling/dataloader.py b/modeling/dataloader.py index 70d2ca7..5ddf034 100644 --- a/modeling/dataloader.py +++ b/modeling/dataloader.py @@ -66,7 +66,7 @@ def contra_data_collator(mlm_collator, features): return batch -def grouped_contra_data_collator(mlm_collator, features, max_num_augs): +def grouped_contra_data_collator(mlm_collator, max_num_augs, features): """Collate grouped samples where each item has 1 anchor + variable-count augmentations. Each feature dict contains: diff --git a/modeling/model.py b/modeling/model.py index 97ab6e3..e4aaa5f 100644 --- a/modeling/model.py +++ b/modeling/model.py @@ -226,7 +226,9 @@ def compute_loss( if self.contra_mode == ContraMode.GROUPED: return self._compute_grouped_loss(model, inputs, return_outputs) - # Move inputs to the device the model lives on (supports DDP) + # Concatenate code and aug inputs into a single batch so that DDP + # sees exactly one forward pass per backward (two separate forwards + # through DDP cause in-place version errors on internal buffers). device = model.device code_input_ids = inputs["code_input_ids"].to(device) code_attention_mask = inputs["code_attention_mask"].to(device) @@ -235,34 +237,25 @@ def compute_loss( 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 - code_outputs = model( - input_ids=code_input_ids, - attention_mask=code_attention_mask, - labels=code_labels, - output_hidden_states=True, - return_dict=True, - ) - code_hidden_states = code_outputs.hidden_states[-1] - code_embeddings = code_hidden_states[:, 0, :] + B = code_input_ids.size(0) - aug_outputs = model( - input_ids=aug_input_ids, - attention_mask=aug_attention_mask, - labels=aug_labels, + all_input_ids = torch.cat([code_input_ids, aug_input_ids], dim=0) + all_attention_mask = torch.cat([code_attention_mask, aug_attention_mask], dim=0) + all_labels = torch.cat([code_labels, aug_labels], dim=0) + + outputs = model( + input_ids=all_input_ids, + attention_mask=all_attention_mask, + labels=all_labels, output_hidden_states=True, return_dict=True, ) - aug_hidden_states = aug_outputs.hidden_states[-1] - aug_embeddings = aug_hidden_states[:, 0, :] - # Average MLM losses so the combined MLM term is on the same scale - # as the single contrastive term (~3-8 each), letting alpha express - # a genuine preference rather than compensating for a 2x scale artifact. - code_mlm_loss = code_outputs.loss - aug_mlm_loss = aug_outputs.loss - mlm_loss = (code_mlm_loss + aug_mlm_loss) / 2 + hidden_states = outputs.hidden_states[-1] + code_embeddings = hidden_states[:B, 0, :] + aug_embeddings = hidden_states[B:, 0, :] + + mlm_loss = outputs.loss # Compute contrastive loss between code and its augmentation if self.contra_mode == ContraMode.SUPCON: @@ -279,10 +272,9 @@ def compute_loss( self.temperature, ) - # Total loss with weighting (adjust alpha as needed) total_loss = mlm_loss + self.alpha * contrastive_loss - return (total_loss, code_outputs) if return_outputs else total_loss + return (total_loss, outputs) if return_outputs else total_loss def _compute_grouped_loss(self, model, inputs, return_outputs=False): """Compute loss for grouped multi-key contrast mode. @@ -303,28 +295,28 @@ def _compute_grouped_loss(self, model, inputs, return_outputs=False): aug_labels = inputs["aug_labels"].to(device) group_sizes = inputs["group_sizes"].to(device) - # Forward anchor - code_outputs = model( - input_ids=code_input_ids, - attention_mask=code_attention_mask, - labels=code_labels, - output_hidden_states=True, - return_dict=True, - ) - code_embeddings = code_outputs.hidden_states[-1][:, 0, :] # [B, D] + B = code_input_ids.size(0) - # Forward all augmentations (flattened [B*max_K, seq_len]) - aug_outputs = model( - input_ids=aug_input_ids, - attention_mask=aug_attention_mask, - labels=aug_labels, + # Single forward pass: concatenate anchors and augmentations to avoid + # DDP in-place buffer errors from two separate forward calls. + all_input_ids = torch.cat([code_input_ids, aug_input_ids], dim=0) + all_attention_mask = torch.cat([code_attention_mask, aug_attention_mask], dim=0) + all_labels = torch.cat([code_labels, aug_labels], dim=0) + + outputs = model( + input_ids=all_input_ids, + attention_mask=all_attention_mask, + labels=all_labels, output_hidden_states=True, return_dict=True, ) - aug_embeddings = aug_outputs.hidden_states[-1][:, 0, :] # [B*max_K, D] + + hidden_states = outputs.hidden_states[-1] + code_embeddings = hidden_states[:B, 0, :] # [B, D] + aug_embeddings = hidden_states[B:, 0, :] # [B*max_K, D] # MLM loss (padding augs have labels=-100, contribute 0) - mlm_loss = (code_outputs.loss + aug_outputs.loss) / 2 + mlm_loss = outputs.loss # Contrastive loss contrastive_loss = grouped_contrastive_loss( @@ -333,7 +325,7 @@ def _compute_grouped_loss(self, model, inputs, return_outputs=False): total_loss = mlm_loss + self.alpha * contrastive_loss - return (total_loss, code_outputs) if return_outputs else total_loss + return (total_loss, outputs) if return_outputs else total_loss def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): """ diff --git a/modeling/pretrain.py b/modeling/pretrain.py index a3c5ead..7679ad8 100644 --- a/modeling/pretrain.py +++ b/modeling/pretrain.py @@ -3,6 +3,7 @@ import hashlib import os from collections import defaultdict +from functools import partial import torch import torch.distributed as dist @@ -270,25 +271,26 @@ def main( if contra_mode == "grouped": # Regroup flat rows by function_id into {code, [aug_1, ..., aug_K]} grouped_dataset = regroup_dataset(dataset, max_num_augs=max_num_augs) - tokenized_datasets = grouped_dataset.shuffle(seed=seed).map( - lambda example: tokenize_grouped( - tokenizer, example, max_seq_length, max_num_augs + tokenized_datasets = grouped_dataset.map( + partial( + tokenize_grouped, + tokenizer, + max_seq_length=max_seq_length, + max_num_augs=max_num_augs, ), batched=True, num_proc=num_proc, - ) + ).shuffle(seed=seed) - collator_fn = lambda features: grouped_contra_data_collator( - mlm_collator, features, max_num_augs - ) + collator_fn = partial(grouped_contra_data_collator, mlm_collator, max_num_augs) else: - tokenized_datasets = dataset.shuffle(seed=seed).map( - lambda example: tokenize(tokenizer, example, max_seq_length=max_seq_length), + tokenized_datasets = dataset.map( + partial(tokenize, tokenizer, max_seq_length=max_seq_length), batched=True, num_proc=num_proc, - ) + ).shuffle(seed=seed) - collator_fn = lambda features: contra_data_collator(mlm_collator, features) + collator_fn = partial(contra_data_collator, mlm_collator) split_dataset = tokenized_datasets.train_test_split(test_size=0.1) train_dataset = split_dataset["train"] diff --git a/tests/test_grouped.py b/tests/test_grouped.py index b77e0da..37c1fb1 100644 --- a/tests/test_grouped.py +++ b/tests/test_grouped.py @@ -162,7 +162,7 @@ def test_output_shapes(self): mlm_collator = self._get_mlm_collator() seq_len = 16 features = _make_grouped_feature(seq_len, aug_counts=[2, 3]) - batch = grouped_contra_data_collator(mlm_collator, features, max_num_augs=6) + batch = grouped_contra_data_collator(mlm_collator, 6, features) B = 2 max_K = 3 # max(2, 3) @@ -177,7 +177,7 @@ def test_padding_has_no_mlm_labels(self): seq_len = 16 # Group 0: 1 aug, Group 1: 3 augs → max_K=3, group 0 has 2 padding slots features = _make_grouped_feature(seq_len, aug_counts=[1, 3]) - batch = grouped_contra_data_collator(mlm_collator, features, max_num_augs=6) + batch = grouped_contra_data_collator(mlm_collator, 6, features) # Group 0's padding slots are indices 1 and 2 in the flattened aug batch # (group 0 occupies slots 0..2, real=1, padding=slots 1,2) @@ -190,14 +190,14 @@ def test_group_sizes_correct(self): """group_sizes should reflect actual aug counts.""" mlm_collator = self._get_mlm_collator() features = _make_grouped_feature(16, aug_counts=[1, 2, 4]) - batch = grouped_contra_data_collator(mlm_collator, features, max_num_augs=6) + batch = grouped_contra_data_collator(mlm_collator, 6, features) assert batch["group_sizes"].tolist() == [1, 2, 4] def test_max_num_augs_truncation(self): """Features with more augs than max_num_augs get truncated.""" mlm_collator = self._get_mlm_collator() features = _make_grouped_feature(16, aug_counts=[5, 3]) - batch = grouped_contra_data_collator(mlm_collator, features, max_num_augs=2) + batch = grouped_contra_data_collator(mlm_collator, 2, features) B = 2 max_K = 2 @@ -208,7 +208,7 @@ def test_function_id_passed_through(self): """function_id should be present in the batch.""" mlm_collator = self._get_mlm_collator() features = _make_grouped_feature(16, aug_counts=[2, 1]) - batch = grouped_contra_data_collator(mlm_collator, features, max_num_augs=6) + batch = grouped_contra_data_collator(mlm_collator, 6, features) assert "function_id" in batch assert batch["function_id"].shape == (2,)