Option 2 for dynamic batching OOM safety

src/dataloaders/build_dataloader.py

@@ -26,6 +26,7 @@ def build_dataloader(
         return torch_data_loader
 
     def build_torch_dataloader(self, torch_dataset):
+        self.dataset = torch_dataset
         sampler = self.get_torch_dataloader_sampler(torch_dataset)
         if "train" in self.args.mode:
             torch_data_loader = DataLoader(
@@ -64,6 +65,8 @@ def collate_fn(self, batch):
         batch = [item for item in batch if item is not None]
         if len(batch) == 0:
             return None
+        max_len = max(item["elm_input_ids"].shape[0] for item in batch)
+        batch = [self.dataset.pad_to_batch(item, max_len) for item in batch]
         self._assert_same_structure_and_shapes(batch)
         return torch.utils.data.dataloader.default_collate(batch)
 

src/dataloaders/data_representation/base.py

@@ -158,6 +158,19 @@ def pad_input(self, tokens: list) -> list:
         padding_len = self.args.llm_input_len - len(tokens)
         return [self.llm_tokenizer.pad_token_id] * padding_len + tokens  # left side padding
 
+    def pad_to_batch(self, item: dict, target_len: int) -> dict:
+        pad_len = target_len - item["elm_input_ids"].shape[0]
+        if pad_len == 0:
+            return item
+        pad_values = {"elm_input_ids": self.llm_tokenizer.pad_token_id, "elm_labels": -100, "elm_attention_mask": 0}
+        for key, value in pad_values.items():
+            if key in item:
+                item[key] = torch.nn.functional.pad(item[key], (pad_len, 0), value=value)  # left side padding
+        if "signal_id_indices" in item:
+            idx = item["signal_id_indices"]
+            item["signal_id_indices"] = torch.where(idx >= 0, idx + pad_len, idx)  # left-pad shifts every real position by pad_len
+        return item
+
     def make_prompt(
         self,
         text: str,

src/dataloaders/data_representation/rgb.py

@@ -72,8 +72,8 @@ def prepare_training_set(
         signal_id_indices = self.find_signal_token_indices(truncated_padded_input)
         attention_mask = self.create_attention_mask(truncated_padded_input)
         labels = self.create_labels(truncated_padded_input)
-        assert len(truncated_padded_input) == len(attention_mask) == len(labels) == self.args.llm_input_len, (
-            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)} != {self.args.llm_input_len}"
+        assert len(truncated_padded_input) == len(attention_mask) == len(labels), (
+            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)}"
         )
         elm = {
             "elm_input_ids": torch.tensor(truncated_padded_input, dtype=torch.int64),
@@ -113,12 +113,6 @@ def augment_image(self, image: np.array):
 
     def trunc_pad_input(self, prompt: str):
         prompt_tokens = self.llm_tokenizer.encode(prompt, add_special_tokens=False)
-        if "train" in self.args.mode:
-            prompt_len = len(prompt_tokens)
-            if prompt_len == self.args.llm_input_len:
-                return prompt_tokens
-            elif prompt_len < self.args.llm_input_len:
-                return self.pad_input(prompt_tokens)
+        if "train" in self.args.mode and len(prompt_tokens) > self.args.llm_input_len:
             return self.truncate_input_preserving_signal_tokens(prompt_tokens)
-        else:
-            return prompt_tokens
+        return prompt_tokens

src/dataloaders/data_representation/signal.py

@@ -47,8 +47,8 @@ def prepare_training_set(
         labels = self.create_labels(truncated_padded_input)
         # print("signal_id_indices", len(signal_id_indices), "\n")
         assert len(signal_id_indices) == self.args.num_encoder_tokens
-        assert len(truncated_padded_input) == len(attention_mask) == len(labels) == self.args.llm_input_len, (
-            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)} != {self.args.llm_input_len}"
+        assert len(truncated_padded_input) == len(attention_mask) == len(labels), (
+            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)}"
         )
         elm = {
             "elm_input_ids": torch.tensor(truncated_padded_input, dtype=torch.int64),
@@ -76,16 +76,9 @@ def prepare_eval_inference_set(
 
     def trunc_pad_input(self, prompt: str):
         prompt_tokens = self.llm_tokenizer.encode(prompt, add_special_tokens=False)
-        if "train" in self.args.mode:
-            prompt_len = len(prompt_tokens)
-            # print("prompt len", prompt_len, "\n")
-            if prompt_len == self.args.llm_input_len:
-                return prompt_tokens
-            elif prompt_len < self.args.llm_input_len:
-                return self.pad_input(prompt_tokens)
+        if "train" in self.args.mode and len(prompt_tokens) > self.args.llm_input_len:
             return self.truncate_input_preserving_signal_tokens(prompt_tokens)
-        else:
-            return prompt_tokens
+        return prompt_tokens
 
     def transform_ecg_signal(self, ecg_signal):
         if self.args.elm == "base_elf":

src/dataloaders/data_representation/stacked_signal.py

@@ -55,8 +55,8 @@ def prepare_training_set(
         signal_id_indices = self.find_signal_token_indices(truncated_padded_input)
         attention_mask = self.create_attention_mask(truncated_padded_input)
         labels = self.create_labels(truncated_padded_input)
-        assert len(truncated_padded_input) == len(attention_mask) == len(labels) == self.args.llm_input_len, (
-            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)} != {self.args.llm_input_len}"
+        assert len(truncated_padded_input) == len(attention_mask) == len(labels), (
+            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)}"
         )
         elm = {
             "elm_input_ids": torch.tensor(truncated_padded_input, dtype=torch.int64),
@@ -90,12 +90,6 @@ def signal_to_stacked_signal(self, signal):
 
     def trunc_pad_input(self, prompt: str):
         prompt_tokens = self.llm_tokenizer.encode(prompt, add_special_tokens=False)
-        if "train" in self.args.mode:
-            prompt_len = len(prompt_tokens)
-            if prompt_len == self.args.llm_input_len:
-                return prompt_tokens
-            elif prompt_len < self.args.llm_input_len:
-                return self.pad_input(prompt_tokens)
+        if "train" in self.args.mode and len(prompt_tokens) > self.args.llm_input_len:
             return self.truncate_input_preserving_signal_tokens(prompt_tokens)
-        else:
-            return prompt_tokens
+        return prompt_tokens

src/dataloaders/data_representation/symbolic.py

@@ -55,8 +55,8 @@ def prepare_training_set(
             self.check_labels(labels)
             self.check_attention_mask(truncated_padded_input, attention_mask)
 
-        assert len(truncated_padded_input) == len(attention_mask) == len(labels) == self.args.llm_input_len, (
-            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)} != {self.args.llm_input_len}"
+        assert len(truncated_padded_input) == len(attention_mask) == len(labels), (
+            f"Length mismatch: {len(truncated_padded_input)} != {len(attention_mask)} != {len(labels)}"
         )
         # print("truncated_padded_ecg_tokens", truncated_padded_ecg_tokens)
         # print("signal_id_indices", signal_id_indices)
@@ -86,12 +86,8 @@ def trunc_pad_input(self, ecg_tokens: np.ndarray, prompt: str):
             min_ecg_token_len = int(self.args.min_ecg_tokens_len)
             before_len, after_len, ecg_token_len = len(before), len(after), len(ecg_tokens)
 
-            if before_len + after_len + ecg_token_len == self.args.llm_input_len:
-                # return before + ecg_tokens + after, self.convert_ecg_tokens(ecg_tokens)
+            if before_len + after_len + ecg_token_len <= self.args.llm_input_len:
                 return before + ecg_tokens + after
-            elif before_len + after_len + ecg_token_len < self.args.llm_input_len:
-                # return self.pad_input(before + ecg_tokens + after), self.convert_ecg_tokens(ecg_tokens)
-                return self.pad_input(before + ecg_tokens + after)
 
             if before_len + min_ecg_token_len > self.args.llm_input_len:
                 raise ValueError("before + min_ecg exceeds llm_input_len; lower min_ecg_tokens_len.")

src/runners/trainer.py

@@ -2,7 +2,9 @@
 from tqdm import tqdm
 import wandb
 
-from utils.gpu_manager import is_main, train_dev_break, batch_to_device
+from utils.gpu_manager import is_main, train_dev_break, batch_to_device, pad_batch_to_len
+
+WARMUP_FULL_BATCHES = 20  # first N batches padded to llm_input_len so a worst-case OOM surfaces early
 
 def run_train(
     nn,
@@ -32,8 +34,17 @@ def run_train(
     optimizer.zero_grad()
     accum_loss_for_log = 0.0
 
+    # Fixed-pad the first N batches to llm_input_len so a worst-case-shape OOM surfaces
+    # early (dynamic padding would otherwise hide it until a long batch appears mid-run).
+    warmup = WARMUP_FULL_BATCHES if epoch == 0 else 0
+    pad_id = dataloader.dataset.llm_tokenizer.pad_token_id if warmup else None
+    if warmup and is_main():
+        print(f"[warmup] first {warmup} batches padded to llm_input_len={args.llm_input_len} to surface OOM early")
+
     for step, batch in enumerate(progress):
         batch = {k: batch_to_device(v, device) for k, v in batch.items()}
+        if step < warmup:
+            batch = pad_batch_to_len(batch, args.llm_input_len, pad_id)
 
         out = nn(**batch)
         raw_loss = out.loss

src/utils/gpu_manager.py

@@ -24,6 +24,22 @@ def batch_to_device(v, device):
         return {k: batch_to_device(x, device) for k, x in v.items()}
     return v
 
+
+def pad_batch_to_len(batch: dict, target_len: int, pad_id: int) -> dict:
+    """Left-pad a collated batch up to target_len (used to force worst-case-shape warmup steps)."""
+    pad = target_len - batch["elm_input_ids"].shape[1]
+    if pad <= 0:
+        return batch
+    F = torch.nn.functional
+    batch = dict(batch)
+    batch["elm_input_ids"] = F.pad(batch["elm_input_ids"], (pad, 0), value=pad_id)
+    batch["elm_attention_mask"] = F.pad(batch["elm_attention_mask"], (pad, 0), value=0)
+    if "elm_labels" in batch:
+        batch["elm_labels"] = F.pad(batch["elm_labels"], (pad, 0), value=-100)
+    if "signal_id_indices" in batch:
+        batch["signal_id_indices"] = torch.where(batch["signal_id_indices"] >= 0, batch["signal_id_indices"] + pad, batch["signal_id_indices"])
+    return batch
+
 def get_world_size() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_world_size()
@@ -82,7 +98,7 @@ def setup_gpu(self, model: torch.nn.Module, find_unused_parameters) -> torch.nn.
         if is_main():
             print(f"find_unused_parameters: {find_unused_parameters}")
         if self.args.torch_compile:
-            model = torch.compile(model)
+            model = torch.compile(model, dynamic=True)
         return model
 
     def get_device(self) -> torch.device: