diff --git a/src/dataloaders/build_dataloader.py b/src/dataloaders/build_dataloader.py index af229d3..2449b7c 100644 --- a/src/dataloaders/build_dataloader.py +++ b/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) diff --git a/src/dataloaders/data_representation/base.py b/src/dataloaders/data_representation/base.py index da1b19d..65b2709 100644 --- a/src/dataloaders/data_representation/base.py +++ b/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, diff --git a/src/dataloaders/data_representation/rgb.py b/src/dataloaders/data_representation/rgb.py index 621cbbf..1b31002 100644 --- a/src/dataloaders/data_representation/rgb.py +++ b/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 \ No newline at end of file + return prompt_tokens \ No newline at end of file diff --git a/src/dataloaders/data_representation/signal.py b/src/dataloaders/data_representation/signal.py index 6730544..7633ca2 100644 --- a/src/dataloaders/data_representation/signal.py +++ b/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": diff --git a/src/dataloaders/data_representation/stacked_signal.py b/src/dataloaders/data_representation/stacked_signal.py index 9f078c3..c9ccfbb 100644 --- a/src/dataloaders/data_representation/stacked_signal.py +++ b/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 diff --git a/src/dataloaders/data_representation/symbolic.py b/src/dataloaders/data_representation/symbolic.py index 1e9b0cc..4c88f2e 100644 --- a/src/dataloaders/data_representation/symbolic.py +++ b/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.") diff --git a/src/runners/trainer.py b/src/runners/trainer.py index da1de6f..99b9e5b 100644 --- a/src/runners/trainer.py +++ b/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 diff --git a/src/utils/gpu_manager.py b/src/utils/gpu_manager.py index 61ab8aa..691396f 100644 --- a/src/utils/gpu_manager.py +++ b/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: