From 8be76d87e9c9226e52c181773644ec3f2d8b9d4f Mon Sep 17 00:00:00 2001 From: GHOpenonic Date: Fri, 10 Jul 2026 13:43:35 -0400 Subject: [PATCH] implement mid-train-chkpting --- src/samudra/config.py | 8 ++ src/samudra/datasets.py | 4 + src/samudra/train.py | 168 ++++++++++++++++++++++++-- src/samudra/utils/logging.py | 18 ++- src/samudra/utils/train.py | 20 +++ tests/test_mid_train_checkpointing.py | 121 +++++++++++++++++++ tests/test_trainer.py | 26 ++++ 7 files changed, 350 insertions(+), 15 deletions(-) create mode 100644 tests/test_mid_train_checkpointing.py diff --git a/src/samudra/config.py b/src/samudra/config.py index 159128d1..bc3e7f5b 100644 --- a/src/samudra/config.py +++ b/src/samudra/config.py @@ -1107,6 +1107,14 @@ class TrainConfig(TopLevelConfig): disk_mode: bool = True pin_mem: bool = True save_freq: int = 5 + mid_train_checkpoint_interval_minutes: float = Field( + default=0.0, + ge=0.0, + description=( + "How often to overwrite the mid-train checkpoint during an active " + "training epoch, in wall-clock minutes. Set to 0 to disable." + ), + ) validation_image_log_freq: int = Field( default=10, ge=1, diff --git a/src/samudra/datasets.py b/src/samudra/datasets.py index 25562b6c..8dd1260d 100644 --- a/src/samudra/datasets.py +++ b/src/samudra/datasets.py @@ -731,6 +731,10 @@ def __getitem__(self, index: int) -> TrainData: def dataset(self): return self._dataloader.dataset + @property + def datasets(self) -> list[TorchTrainDataset]: + return list(self._datasets.values()) + @property def sampler(self): return self._dataloader.sampler diff --git a/src/samudra/train.py b/src/samudra/train.py index 97533d4b..a2b4cc9d 100644 --- a/src/samudra/train.py +++ b/src/samudra/train.py @@ -25,6 +25,7 @@ DataLoader, DistributedSampler, RandomSampler, + Sampler, ) from samudra import config @@ -101,6 +102,22 @@ def should_log_validation_images(epoch: int, frequency: int) -> bool: return (epoch - 1) % frequency == 0 +class _OffsetBatchSampler(Sampler[list[int]]): + """Wrap a batch sampler and skip batches already completed in an epoch.""" + + def __init__(self, batch_sampler: Sampler[list[int]], start_batch: int): + self._batch_sampler = batch_sampler + self._start_batch = max(start_batch, 0) + + def __iter__(self): + for batch_index, batch in enumerate(self._batch_sampler): + if batch_index >= self._start_batch: + yield batch + + def __len__(self) -> int: + return max(len(self._batch_sampler) - self._start_batch, 0) # type: ignore[arg-type] + + class Trainer: """Orchestrates the full model training loop. @@ -256,9 +273,9 @@ def __init__(self, cfg: TrainConfig) -> None: # Check for preemption if cfg.preemptible: assert not cfg.finetune, "Finetune is not supported with preemptible" - preempted = os.path.isfile(self.ckpt_paths.latest_checkpoint_path) - if preempted: - cfg.resume_ckpt_path = str(self.ckpt_paths.latest_checkpoint_path) + resumable_checkpoint = self.ckpt_paths.latest_resumable_checkpoint_path() + if resumable_checkpoint is not None: + cfg.resume_ckpt_path = str(resumable_checkpoint) # Set up wandb run self.wandb_id, self.wandb_name = self.wandb_logger.setup_run( @@ -284,6 +301,7 @@ def __init__(self, cfg: TrainConfig) -> None: ) self.num_batches_seen = 0 + self.start_batch_in_epoch = 0 loaded_checkpoint = False if cfg.resume_ckpt_path is not None: if cfg.finetune: @@ -341,6 +359,18 @@ def __init__(self, cfg: TrainConfig) -> None: self.normalize_before_mask: bool = cfg.data.normalize_before_mask self.normalize_fill_value: float = cfg.data.masked_fill_value self.delayed_loss_estimate: bool = cfg.delayed_loss_estimate + self.mid_train_checkpoint_interval_seconds = ( + cfg.mid_train_checkpoint_interval_minutes * 60.0 + ) + self._next_mid_train_checkpoint_time: float | None = None + if self.mid_train_checkpoint_interval_seconds > 0: + self._next_mid_train_checkpoint_time = ( + time.perf_counter() + self.mid_train_checkpoint_interval_seconds + ) + logger.info( + "Mid-train checkpointing enabled: every %.1f minutes", + cfg.mid_train_checkpoint_interval_minutes, + ) self.profiler = cfg.profiler.build(self.output_dir, self.device) self.validation_images_enabled = self._sync_flag_from_main( @@ -441,7 +471,17 @@ def run(self) -> None: self.val_sampler.set_epoch(epoch) start_epoch_train_time = time.perf_counter() - train_stats = self.train_one_epoch(epoch) + start_batch_in_epoch = ( + self.start_batch_in_epoch if epoch == self.start_epoch else 0 + ) + if start_batch_in_epoch > 0: + logger.info( + "Resuming epoch %s from batch %s", + epoch, + start_batch_in_epoch, + ) + train_stats = self.train_one_epoch(epoch, start_batch_in_epoch) + self.start_batch_in_epoch = 0 end_epoch_train_time = time.perf_counter() val_stats = self.validate_one_epoch(epoch) end_epoch_val_time = time.perf_counter() @@ -490,7 +530,7 @@ def run(self) -> None: logger.info(f"Training time {total_time_str}") self.finish() - def train_one_epoch(self, epoch): + def train_one_epoch(self, epoch, start_batch_in_epoch: int = 0): self.model.train(True) train_aggregator = Aggregator.get_train_aggregator(self.tensor_map) metric_logger = MetricLogger(delimiter=" ") @@ -498,6 +538,11 @@ def train_one_epoch(self, epoch): header = f"Training Epoch: [{epoch}]" total_batches = len(self.train_loader) + train_loader = ( + self._train_loader_with_batch_offset(start_batch_in_epoch) + if start_batch_in_epoch > 0 + else self.train_loader + ) # Ensure gradients are zeroed at the start of the epoch so we don't # accidentally accumulate leftovers from checkpoint/loading. @@ -512,13 +557,20 @@ def train_one_epoch(self, epoch): ) for data_iter_step, data in enumerate( - metric_logger.log_every(self.train_loader, 1, header) + metric_logger.log_every( + train_loader, + 1, + header, + start_index=start_batch_in_epoch, + total_steps=total_batches, + ) ): + global_data_iter_step = start_batch_in_epoch + data_iter_step if self.debug and (data_iter_step + 1) % 5 == 0: break in_final_cycle = ( - data_iter_step + 1 > final_cycle_start + global_data_iter_step + 1 > final_cycle_start ) and remaining_batches > 0 # Determine the actual number of microbatches in this accumulation cycle @@ -540,8 +592,10 @@ def train_one_epoch(self, epoch): self.num_batches_seen += 1 - is_last = data_iter_step + 1 == total_batches - should_step = (data_iter_step + 1) % self.gradient_accumulation_steps == 0 + is_last = global_data_iter_step + 1 == total_batches + should_step = ( + global_data_iter_step + 1 + ) % self.gradient_accumulation_steps == 0 # Step optimizer after accumulating enough batches or at the end if should_step or is_last: torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) @@ -629,6 +683,11 @@ def train_one_epoch(self, epoch): self._maybe_update_loss(TO, data) self.profiler.after_batch(self.num_batches_seen) + if should_step or is_last: + self._maybe_save_mid_train_checkpoint( + epoch=epoch, + batch_in_epoch=global_data_iter_step, + ) if self.scheduler is not None: self.scheduler.step() @@ -636,6 +695,29 @@ def train_one_epoch(self, epoch): logger.info(f"Aggregating train logs") return train_aggregator.get_logs() + def _train_loader_with_batch_offset( + self, + start_batch_in_epoch: int, + ) -> TrainDataLoader: + raw_loader = self.train_loader._dataloader + if raw_loader.batch_sampler is None: + raise RuntimeError("Cannot resume mid-epoch without a batch sampler.") + + batch_sampler = _OffsetBatchSampler( + raw_loader.batch_sampler, + start_batch=start_batch_in_epoch, + ) + dataloader = DataLoader( + raw_loader.dataset, + batch_sampler=batch_sampler, + num_workers=self.num_workers, + persistent_workers=self.persistent_workers and self.num_workers > 0, + pin_memory=self.pin_mem, + collate_fn=raw_loader.collate_fn, + multiprocessing_context=self.mp_context, + ) + return TrainDataLoader(dataloader, self.train_loader.datasets, self.device) + def _maybe_update_loss(self, output: TrainBatchOutput, data: TrainData): if (update := getattr(self.loss_fn, "update", None)) is None: return @@ -999,11 +1081,65 @@ def save_all_checkpoints(self, epoch: int, v_loss: float, inf_loss: float): for_inference=True, ) + def _maybe_save_mid_train_checkpoint( + self, + epoch: int, + batch_in_epoch: int, + ) -> None: + if self._next_mid_train_checkpoint_time is None: + return + + now = time.perf_counter() + if now < self._next_mid_train_checkpoint_time: + return + + if is_main_process(): + self.save_mid_train_checkpoint( + epoch=epoch, + batch_in_epoch=batch_in_epoch, + reason="periodic_interval", + ) + + while self._next_mid_train_checkpoint_time <= now: + self._next_mid_train_checkpoint_time += ( + self.mid_train_checkpoint_interval_seconds + ) + + def save_mid_train_checkpoint( + self, + epoch: int, + batch_in_epoch: int, + reason: str = "periodic_interval", + ) -> Path | None: + checkpoint_path = self.ckpt_paths.latest_mid_train_checkpoint_path + try: + self.save_checkpoint( + epoch=epoch, + checkpoint_path=checkpoint_path, + batch_in_epoch=max(batch_in_epoch, -1), + epoch_complete=False, + save_reason=f"mid_train_{reason}", + ) + except Exception: + logger.exception("Failed to save mid-train checkpoint to %s", checkpoint_path) + return None + + logger.info( + "Saved mid-train checkpoint to %s after epoch %s batch %s", + checkpoint_path, + epoch, + batch_in_epoch, + ) + return checkpoint_path + def save_checkpoint( self, epoch: int, checkpoint_path: Path, for_inference: bool = False, + batch_in_epoch: int | None = None, + epoch_complete: bool = True, + save_reason: str | None = None, ): if for_inference: with self._ema_context(): @@ -1025,7 +1161,12 @@ def save_checkpoint( "num_batches_seen": self.num_batches_seen, "wandb_id": self.wandb_id, "wandb_name": self.wandb_name, + "epoch_complete": epoch_complete, } + if batch_in_epoch is not None: + checkpoint["batch_in_epoch"] = batch_in_epoch + if save_reason is not None: + checkpoint["save_reason"] = save_reason loss_state: dict[str, Any] | None = None if state_dict_fn := getattr(self.loss_fn, "state_dict", None): loss_state = state_dict_fn() @@ -1076,12 +1217,19 @@ def remove_module_prefix(state_dict, prefix="module."): ) load_state_dict_fn(checkpoint["loss_fn_state"]) - self.start_epoch = checkpoint["epoch"] + 1 + epoch_complete = checkpoint.get("epoch_complete", True) + if epoch_complete: + self.start_epoch = checkpoint["epoch"] + 1 + self.start_batch_in_epoch = 0 + else: + self.start_epoch = checkpoint["epoch"] + self.start_batch_in_epoch = checkpoint.get("batch_in_epoch", -1) + 1 self.wandb_id = checkpoint.get("wandb_id") self.wandb_name = checkpoint.get("wandb_name") self.num_batches_seen = checkpoint.get("num_batches_seen", 0) logger.info(f"Start Epoch: {self.start_epoch}") + logger.info(f"Start Batch In Epoch: {self.start_batch_in_epoch}") logger.info(f"Wandb id: {self.wandb_id}") logger.info(f"Wandb name: {self.wandb_name}") logger.info(f"Optimizer LR: {self.optimizer.param_groups[-1]['lr']}") diff --git a/src/samudra/utils/logging.py b/src/samudra/utils/logging.py index 27b214ae..72f66456 100644 --- a/src/samudra/utils/logging.py +++ b/src/samudra/utils/logging.py @@ -164,10 +164,14 @@ def log_every( data_loader: "TrainDataLoader", print_freq, header=None, + start_index: int = 0, + total_steps: int | None = None, ): i = 0 if not header: header = "" + if total_steps is None: + total_steps = len(data_loader) start_time = time.perf_counter() end = time.perf_counter() iter_time = SmoothedValue(fmt="{value:.3f}({avg:.3f})", window_size=print_freq) @@ -180,7 +184,7 @@ def log_every( self.meters["iter_time"] = iter_time self.meters["data_wait_time"] = data_wait_time self.meters["data_load_time"] = data_load_time - space_fmt = ":" + str(len(str(len(data_loader)))) + "d" + space_fmt = ":" + str(len(str(total_steps))) + "d" log_msg_list: list[str] = [ header, "[{0" + space_fmt + "}/{1}]", @@ -199,8 +203,10 @@ def log_every( data_load_time.update(obj.load_stats.load_time_seconds) yield obj iter_time.update(time.perf_counter() - end) - if i % print_freq == 0 or i == len(data_loader) - 1: - eta_seconds = iter_time.global_avg * (len(data_loader) - i) + display_index = start_index + i + if i % print_freq == 0 or display_index == total_steps - 1: + remaining_steps = max(total_steps - display_index - 1, 0) + eta_seconds = iter_time.global_avg * remaining_steps eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) named_metrics = dict( eta=eta_string, @@ -210,7 +216,7 @@ def log_every( if torch.cuda.is_available(): named_metrics["gpu_memory"] = torch.cuda.max_memory_allocated() / MB - logger.info(log_msg.format(i, len(data_loader), **named_metrics)) + logger.info(log_msg.format(display_index, total_steps, **named_metrics)) if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() @@ -218,9 +224,11 @@ def log_every( end = time.perf_counter() total_time = time.perf_counter() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) + iterations = len(data_loader) + seconds_per_iteration = total_time / iterations if iterations > 0 else 0.0 logger.info( f"{header} Total time: {total_time_str} " - f"({total_time / len(data_loader):.4f} s / it)" + f"({seconds_per_iteration:.4f} s / it)" ) diff --git a/src/samudra/utils/train.py b/src/samudra/utils/train.py index 37c5aff4..b7fef2e2 100644 --- a/src/samudra/utils/train.py +++ b/src/samudra/utils/train.py @@ -57,6 +57,26 @@ def __init__(self, checkpoint_dir: Path): def latest_checkpoint_path(self) -> Path: return self.checkpoint_dir / "ckpt.pt" + @property + def latest_mid_train_checkpoint_path(self) -> Path: + return self.checkpoint_dir / "ckpt_mid_train.pt" + + def latest_resumable_checkpoint_path(self) -> Path | None: + paths = [ + self.latest_checkpoint_path, + self.latest_mid_train_checkpoint_path, + ] + existing_paths = [path for path in paths if path.exists()] + if not existing_paths: + return None + return max( + existing_paths, + key=lambda path: ( + path.stat().st_mtime, + path == self.latest_mid_train_checkpoint_path, + ), + ) + def latest_checkpoint_path_with_epoch(self, epoch: int) -> Path: return self.checkpoint_dir / f"ckpt_{epoch}.pt" diff --git a/tests/test_mid_train_checkpointing.py b/tests/test_mid_train_checkpointing.py new file mode 100644 index 00000000..8bfaf06a --- /dev/null +++ b/tests/test_mid_train_checkpointing.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import logging +from dataclasses import dataclass + +from torch.utils.data import BatchSampler, SequentialSampler + +from samudra.train import Trainer, _OffsetBatchSampler +from samudra.utils.logging import MetricLogger +from samudra.utils.train import CheckpointPaths + + +@dataclass +class _FakeTrainData: + load_stats = type("LoadStats", (), {"load_time_seconds": 0.0})() + + +class _FakeTrainLoader: + def __init__(self, size: int): + self._items = [_FakeTrainData() for _ in range(size)] + + def __iter__(self): + return iter(self._items) + + def __len__(self) -> int: + return len(self._items) + + +def test_offset_batch_sampler_skips_prefix_batches(): + sampler = BatchSampler(SequentialSampler(range(10)), batch_size=2, drop_last=False) + offset_sampler = _OffsetBatchSampler(sampler, start_batch=2) + + assert list(offset_sampler) == [[4, 5], [6, 7], [8, 9]] + assert len(offset_sampler) == 3 + + +def test_offset_batch_sampler_is_empty_when_offset_past_end(): + sampler = BatchSampler(SequentialSampler(range(4)), batch_size=2, drop_last=False) + offset_sampler = _OffsetBatchSampler(sampler, start_batch=5) + + assert list(offset_sampler) == [] + assert len(offset_sampler) == 0 + + +def test_metric_logger_log_every_supports_global_start_index(caplog): + caplog.set_level(logging.INFO) + metric_logger = MetricLogger(delimiter=" ") + + for _data in metric_logger.log_every( + _FakeTrainLoader(size=2), + print_freq=1, + header="Training Epoch: [3]", + start_index=5, + total_steps=7, + ): + metric_logger.update(loss=1.0) + + assert any("[5/7]" in record.message for record in caplog.records) + assert any("[6/7]" in record.message for record in caplog.records) + + +def test_mid_train_checkpoint_path_and_resumable_selection(tmp_path): + paths = CheckpointPaths(tmp_path) + assert paths.latest_mid_train_checkpoint_path == tmp_path / "ckpt_mid_train.pt" + assert paths.latest_resumable_checkpoint_path() is None + + paths.latest_checkpoint_path.write_text("epoch") + assert paths.latest_resumable_checkpoint_path() == paths.latest_checkpoint_path + + paths.latest_mid_train_checkpoint_path.write_text("mid") + assert ( + paths.latest_resumable_checkpoint_path() + == paths.latest_mid_train_checkpoint_path + ) + + +def test_mid_train_checkpoint_interval_saves_and_advances_timer(monkeypatch): + trainer = Trainer.__new__(Trainer) + trainer._next_mid_train_checkpoint_time = 10.0 + trainer.mid_train_checkpoint_interval_seconds = 5.0 + calls = [] + + def save_mid_train_checkpoint(**kwargs): + calls.append(kwargs) + + trainer.save_mid_train_checkpoint = save_mid_train_checkpoint + monkeypatch.setattr("samudra.train.time.perf_counter", lambda: 21.0) + monkeypatch.setattr("samudra.train.is_main_process", lambda: True) + + trainer._maybe_save_mid_train_checkpoint(epoch=3, batch_in_epoch=7) + + assert calls == [ + { + "epoch": 3, + "batch_in_epoch": 7, + "reason": "periodic_interval", + } + ] + assert trainer._next_mid_train_checkpoint_time == 25.0 + + +def test_save_mid_train_checkpoint_uses_mid_train_path_and_metadata(tmp_path): + trainer = Trainer.__new__(Trainer) + trainer.ckpt_paths = CheckpointPaths(tmp_path) + captured = {} + + def save_checkpoint(**kwargs): + captured.update(kwargs) + + trainer.save_checkpoint = save_checkpoint + + path = trainer.save_mid_train_checkpoint(epoch=2, batch_in_epoch=4) + + assert path == tmp_path / "ckpt_mid_train.pt" + assert captured["checkpoint_path"] == path + assert captured["epoch"] == 2 + assert captured["batch_in_epoch"] == 4 + assert captured["epoch_complete"] is False + assert captured["save_reason"] == "mid_train_periodic_interval" diff --git a/tests/test_trainer.py b/tests/test_trainer.py index fec405f9..099e98d8 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -142,6 +142,32 @@ def test_checkpoint_inference(trainer_pair: TrainPair, caplog): assert torch.allclose(out, out2) +@pytest.mark.parametrize( + "data_source,config_name", + [("mock-om4", "test/train_default_2step.yaml")], + indirect=True, +) +@pytest.mark.parametrize("backend", ["cpu"], indirect=True) +def test_checkpoint_incomplete_mid_train_resume_batch(trainer_pair: TrainPair, tmp_path): + _, trainer = trainer_pair + + trainer.best_val_loss = 10 + trainer.best_inf_loss = 10 + checkpoint_path = tmp_path / "mid_train.pt" + + trainer.save_checkpoint( + epoch=3, + checkpoint_path=checkpoint_path, + batch_in_epoch=5, + epoch_complete=False, + save_reason="mid_train_periodic_interval", + ) + trainer.load_checkpoint(checkpoint_path) + + assert trainer.start_epoch == 3 + assert trainer.start_batch_in_epoch == 6 + + @pytest.mark.parametrize( "data_source,config_name,extra_config_args", [