From b9b9b0a305bf44db2c9eae766d27827b35ab4552 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sat, 7 Feb 2026 22:03:51 +0000 Subject: [PATCH 1/2] fix: avoid OOM by computing MLM loss per-half instead of on concatenated batch The DDP single-forward fix doubled the logit tensor [2B, seq, vocab] causing OOM on single GPU. Split MLM head computation into two halves so only [B, seq, vocab] logits are materialized at a time. --- experiments/supcon/codebert.yaml | 1 - modeling/model.py | 40 ++++++++++++++++++++++++++------ 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/experiments/supcon/codebert.yaml b/experiments/supcon/codebert.yaml index 477d1f4..6831157 100644 --- a/experiments/supcon/codebert.yaml +++ b/experiments/supcon/codebert.yaml @@ -16,6 +16,5 @@ run_name: "InvCodeBERT-supcon" alpha: 1.0 temperature: 0.1 max_seq_length: 512 -sample_rate: 0.2 contra_mode: "supcon" diff --git a/modeling/model.py b/modeling/model.py index e4aaa5f..6e24978 100644 --- a/modeling/model.py +++ b/modeling/model.py @@ -211,6 +211,20 @@ def grouped_contrastive_loss( return loss +def _mlm_loss_from_hidden(model, hidden_states, labels): + """Compute MLM loss from hidden states without materializing logits for the full concatenated batch. + + Runs the LM head on a subset of hidden states so the [N, seq_len, vocab_size] + logit tensor is only as large as the subset, not the full 2B (or B+B*K) batch. + """ + logits = model.lm_head(hidden_states) + return F.cross_entropy( + logits.view(-1, logits.size(-1)), + labels.view(-1), + ignore_index=-100, + ) + + class ContrastiveTrainer(Trainer): def __init__( self, alpha=1.0, temperature=0.07, contra_mode="info_nce", *args, **kwargs @@ -241,12 +255,12 @@ def compute_loss( 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) + # Don't pass labels — avoids materializing [2B, seq_len, vocab_size] + # logits in one tensor. We compute MLM loss per-half below. outputs = model( input_ids=all_input_ids, attention_mask=all_attention_mask, - labels=all_labels, output_hidden_states=True, return_dict=True, ) @@ -255,7 +269,14 @@ def compute_loss( code_embeddings = hidden_states[:B, 0, :] aug_embeddings = hidden_states[B:, 0, :] - mlm_loss = outputs.loss + # Compute MLM loss on each half separately so the [B, seq_len, vocab] + # logit tensor is only half as large and freed between the two calls. + base_model = model.module if hasattr(model, "module") else model + code_mlm_loss = _mlm_loss_from_hidden( + base_model, hidden_states[:B], code_labels + ) + aug_mlm_loss = _mlm_loss_from_hidden(base_model, hidden_states[B:], aug_labels) + mlm_loss = (code_mlm_loss + aug_mlm_loss) / 2 # Compute contrastive loss between code and its augmentation if self.contra_mode == ContraMode.SUPCON: @@ -301,12 +322,12 @@ def _compute_grouped_loss(self, model, inputs, return_outputs=False): # 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) + # Don't pass labels — compute MLM loss per-half to avoid the huge + # [B + B*max_K, seq_len, vocab_size] logit tensor. outputs = model( input_ids=all_input_ids, attention_mask=all_attention_mask, - labels=all_labels, output_hidden_states=True, return_dict=True, ) @@ -315,8 +336,13 @@ def _compute_grouped_loss(self, model, inputs, return_outputs=False): 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 = outputs.loss + # MLM loss on each half (padding augs have labels=-100, contribute 0) + base_model = model.module if hasattr(model, "module") else model + code_mlm_loss = _mlm_loss_from_hidden( + base_model, hidden_states[:B], code_labels + ) + aug_mlm_loss = _mlm_loss_from_hidden(base_model, hidden_states[B:], aug_labels) + mlm_loss = (code_mlm_loss + aug_mlm_loss) / 2 # Contrastive loss contrastive_loss = grouped_contrastive_loss( From 9bcbbc1fb83c20cb033b0aa55dca307a2bd65b57 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sat, 7 Feb 2026 22:17:07 +0000 Subject: [PATCH 2/2] fix: OOM on single GPU by splitting LM head computation per-chunk RobertaForMaskedLM.forward() always materializes [N, seq, vocab] logits. With the concatenated 2B batch from the DDP fix, this doubled peak memory (~13 GB logits on batch_size=64, seq=512) causing OOM on A100-80GB. Introduce SplitHeadWrapper that runs the encoder on the full batch but applies the LM head per-chunk, halving peak logit memory. DDP wraps the wrapper so gradient sync works correctly. --- modeling/model.py | 190 +++++++++++++++++++++++++++---------------- modeling/pretrain.py | 18 +++- 2 files changed, 134 insertions(+), 74 deletions(-) diff --git a/modeling/model.py b/modeling/model.py index 6e24978..1cb240e 100644 --- a/modeling/model.py +++ b/modeling/model.py @@ -1,10 +1,93 @@ import torch +import torch.nn as nn import torch.nn.functional as F -from transformers import Trainer +from transformers import RobertaForMaskedLM, Trainer from ._types import ContraMode +class SplitHeadWrapper(nn.Module): + """Wraps ``RobertaForMaskedLM`` so that the LM head is applied per-chunk. + + ``RobertaForMaskedLM.forward()`` always materializes a + ``[N, seq_len, vocab_size]`` logit tensor. When the input is a + concatenated code+aug batch (``N = 2B`` or ``N = B + B*K``), this + doubles/multiplies peak GPU memory and causes OOM. + + This wrapper runs the **encoder** on the full concatenated input (one + forward pass — required for DDP), then applies the **LM head** on each + chunk separately so the logit tensor is never larger than + ``[B, seq_len, vocab_size]``. + + The wrapper is what DDP wraps, so all parameters (encoder + lm_head) + participate in the single ``forward()`` and gradient sync works normally. + """ + + def __init__(self, roberta_mlm: RobertaForMaskedLM): + super().__init__() + self.roberta_mlm = roberta_mlm + + @property + def config(self): + return self.roberta_mlm.config + + @property + def device(self): + return self.roberta_mlm.device + + def forward( + self, + input_ids, + attention_mask, + labels_a, + labels_b, + split_at, + output_hidden_states=True, + ): + """Run encoder on full batch, compute MLM loss on each half. + + Args: + input_ids: ``[N, seq_len]`` concatenated code + aug tokens. + attention_mask: ``[N, seq_len]``. + labels_a: ``[split_at, seq_len]`` MLM labels for the first chunk. + labels_b: ``[N - split_at, seq_len]`` MLM labels for the second chunk. + split_at: integer index where to split (typically ``B``). + output_hidden_states: whether to return encoder hidden states. + + Returns: + ``(mlm_loss, last_hidden_state)`` where ``mlm_loss`` is the + average of the per-chunk MLM losses. + """ + encoder_outputs = self.roberta_mlm.roberta( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + last_hidden = encoder_outputs.last_hidden_state # [N, seq, D] + + # Run lm_head per-chunk to avoid [N, seq, vocab] peak memory. + logits_a = self.roberta_mlm.lm_head(last_hidden[:split_at]) + mlm_loss_a = F.cross_entropy( + logits_a.view(-1, logits_a.size(-1)), + labels_a.view(-1), + ignore_index=-100, + ) + del logits_a + + logits_b = self.roberta_mlm.lm_head(last_hidden[split_at:]) + mlm_loss_b = F.cross_entropy( + logits_b.view(-1, logits_b.size(-1)), + labels_b.view(-1), + ignore_index=-100, + ) + del logits_b + + mlm_loss = (mlm_loss_a + mlm_loss_b) / 2 + return mlm_loss, last_hidden + + def info_nce_loss(query, key, temperature=0.07): device = query.device query = F.normalize(query, dim=1) @@ -211,21 +294,15 @@ def grouped_contrastive_loss( return loss -def _mlm_loss_from_hidden(model, hidden_states, labels): - """Compute MLM loss from hidden states without materializing logits for the full concatenated batch. +class ContrastiveTrainer(Trainer): + """HF Trainer subclass for contrastive pre-training. - Runs the LM head on a subset of hidden states so the [N, seq_len, vocab_size] - logit tensor is only as large as the subset, not the full 2B (or B+B*K) batch. + Expects ``model`` to be a :class:`SplitHeadWrapper` (or DDP-wrapped + ``SplitHeadWrapper``). The wrapper's ``forward()`` runs the encoder on + the full concatenated batch but applies the LM head per-chunk, so the + ``[N, seq, vocab]`` logit tensor is never larger than ``[B, seq, vocab]``. """ - logits = model.lm_head(hidden_states) - return F.cross_entropy( - logits.view(-1, logits.size(-1)), - labels.view(-1), - ignore_index=-100, - ) - -class ContrastiveTrainer(Trainer): def __init__( self, alpha=1.0, temperature=0.07, contra_mode="info_nce", *args, **kwargs ): @@ -240,9 +317,6 @@ def compute_loss( if self.contra_mode == ContraMode.GROUPED: return self._compute_grouped_loss(model, inputs, return_outputs) - # 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) @@ -256,27 +330,18 @@ def compute_loss( 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) - # Don't pass labels — avoids materializing [2B, seq_len, vocab_size] - # logits in one tensor. We compute MLM loss per-half below. - outputs = model( + # Single forward through SplitHeadWrapper: encoder on full batch, + # LM head per-chunk to avoid [2B, seq, vocab] peak memory. + mlm_loss, last_hidden = model( input_ids=all_input_ids, attention_mask=all_attention_mask, - output_hidden_states=True, - return_dict=True, + labels_a=code_labels, + labels_b=aug_labels, + split_at=B, ) - hidden_states = outputs.hidden_states[-1] - code_embeddings = hidden_states[:B, 0, :] - aug_embeddings = hidden_states[B:, 0, :] - - # Compute MLM loss on each half separately so the [B, seq_len, vocab] - # logit tensor is only half as large and freed between the two calls. - base_model = model.module if hasattr(model, "module") else model - code_mlm_loss = _mlm_loss_from_hidden( - base_model, hidden_states[:B], code_labels - ) - aug_mlm_loss = _mlm_loss_from_hidden(base_model, hidden_states[B:], aug_labels) - mlm_loss = (code_mlm_loss + aug_mlm_loss) / 2 + code_embeddings = last_hidden[:B, 0, :] + aug_embeddings = last_hidden[B:, 0, :] # Compute contrastive loss between code and its augmentation if self.contra_mode == ContraMode.SUPCON: @@ -295,7 +360,7 @@ def compute_loss( total_loss = mlm_loss + self.alpha * contrastive_loss - return (total_loss, outputs) if return_outputs else total_loss + return (total_loss, None) if return_outputs else total_loss def _compute_grouped_loss(self, model, inputs, return_outputs=False): """Compute loss for grouped multi-key contrast mode. @@ -318,31 +383,19 @@ def _compute_grouped_loss(self, model, inputs, return_outputs=False): B = code_input_ids.size(0) - # 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) - # Don't pass labels — compute MLM loss per-half to avoid the huge - # [B + B*max_K, seq_len, vocab_size] logit tensor. - outputs = model( + mlm_loss, last_hidden = model( input_ids=all_input_ids, attention_mask=all_attention_mask, - output_hidden_states=True, - return_dict=True, + labels_a=code_labels, + labels_b=aug_labels, + split_at=B, ) - 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 on each half (padding augs have labels=-100, contribute 0) - base_model = model.module if hasattr(model, "module") else model - code_mlm_loss = _mlm_loss_from_hidden( - base_model, hidden_states[:B], code_labels - ) - aug_mlm_loss = _mlm_loss_from_hidden(base_model, hidden_states[B:], aug_labels) - mlm_loss = (code_mlm_loss + aug_mlm_loss) / 2 + code_embeddings = last_hidden[:B, 0, :] + aug_embeddings = last_hidden[B:, 0, :] # Contrastive loss contrastive_loss = grouped_contrastive_loss( @@ -351,33 +404,28 @@ def _compute_grouped_loss(self, model, inputs, return_outputs=False): total_loss = mlm_loss + self.alpha * contrastive_loss - return (total_loss, outputs) if return_outputs else total_loss + return (total_loss, None) if return_outputs else total_loss def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): - """ - Override the default prediction_step to handle custom inputs during evaluation. - """ - # Move inputs to device + """Override prediction_step for SplitHeadWrapper evaluation.""" device = self.args.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) - # Prepare inputs for the model - # Since evaluation usually focuses on the MLM task, we can use code inputs - inputs_for_model = { - "input_ids": code_input_ids, - "attention_mask": code_attention_mask, - "labels": code_labels, - } - with torch.no_grad(): - outputs = model(**inputs_for_model) + # For eval we only need code (no aug), so both halves are the + # same chunk. Pass an empty second chunk. + base = model.module if hasattr(model, "module") else model + roberta_mlm = base.roberta_mlm + outputs = roberta_mlm( + input_ids=code_input_ids, + attention_mask=code_attention_mask, + labels=code_labels, + return_dict=True, + ) + loss = outputs.loss if prediction_loss_only: - loss = outputs.loss return (loss, None, None) - else: - loss = outputs.loss - logits = outputs.logits - return (loss, logits, code_labels) + return (loss, outputs.logits, code_labels) diff --git a/modeling/pretrain.py b/modeling/pretrain.py index 7679ad8..7165671 100644 --- a/modeling/pretrain.py +++ b/modeling/pretrain.py @@ -19,7 +19,7 @@ from ._types import ContraMode from .common import default_num_proc, set_seed from .dataloader import contra_data_collator, grouped_contra_data_collator -from .model import ContrastiveTrainer +from .model import ContrastiveTrainer, SplitHeadWrapper def _get_world_size() -> int: @@ -239,11 +239,15 @@ def main( config = RobertaConfig.from_pretrained(tokenizer_name) # model = RobertaForMaskedLM.from_pretrained(model_name) - model = RobertaForMaskedLM.from_pretrained( + roberta_mlm = RobertaForMaskedLM.from_pretrained( model_name if checkpoint is None else checkpoint, config=config, ) # load weights from stage 1 + # Wrap so the LM head is applied per-chunk (avoids [2B, seq, vocab] OOM). + # DDP will wrap SplitHeadWrapper, keeping all params in one forward(). + model = SplitHeadWrapper(roberta_mlm) + features = Features( { "repo": Value("string"), @@ -328,4 +332,12 @@ def main( ) trainer.train(resume_from_checkpoint=resume) - trainer.save_model(f"saved_models/{run_name}/final") + + # Save the inner RobertaForMaskedLM so downstream tasks can load it + # directly with RobertaForMaskedLM.from_pretrained(). + save_path = f"saved_models/{run_name}/final" + unwrapped = trainer.model + if hasattr(unwrapped, "module"): # DDP + unwrapped = unwrapped.module + unwrapped.roberta_mlm.save_pretrained(save_path) + tokenizer.save_pretrained(save_path)