fix: supcon oom

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"

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)
@@ -212,6 +295,14 @@ def grouped_contrastive_loss(
 
 
 class ContrastiveTrainer(Trainer):
+    """HF Trainer subclass for contrastive pre-training.
+
+    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]``.
+    """
+
     def __init__(
         self, alpha=1.0, temperature=0.07, contra_mode="info_nce", *args, **kwargs
     ):
@@ -226,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)
@@ -241,21 +329,19 @@ 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)
 
-        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,
-            labels=all_labels,
-            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, :]
-
-        mlm_loss = outputs.loss
+        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:
@@ -274,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.
@@ -297,26 +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)
-        all_labels = torch.cat([code_labels, aug_labels], dim=0)
 
-        outputs = model(
+        mlm_loss, last_hidden = model(
             input_ids=all_input_ids,
             attention_mask=all_attention_mask,
-            labels=all_labels,
-            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 (padding augs have labels=-100, contribute 0)
-        mlm_loss = outputs.loss
+        code_embeddings = last_hidden[:B, 0, :]
+        aug_embeddings = last_hidden[B:, 0, :]
 
         # Contrastive loss
         contrastive_loss = grouped_contrastive_loss(
@@ -325,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)

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)