diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index 8575301..3a50d44 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -47,14 +47,18 @@ class TrainConfig(BaseConfig): # checkpoint setting start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."}) - start_batch: int = field( - default=0, metadata={"help": "Start batch iteration for training."} + start_samples: int = field( + default=0, + metadata={ + "help": "Start samples count (per rank). Superseded by checkpoint consumed_samples." + }, ) ckpt_dir: str = field( default="./checkpoint", metadata={"help": "Checkpoint directory."} ) ckpt_interval: int = field( - default=5000, metadata={"help": "Number of iterations between checkpoints."} + default=5000, + metadata={"help": "Number of optimizer steps between checkpoints."}, ) # lora setting diff --git a/astrai/serialization.py b/astrai/serialization.py index 9537fe9..9be039c 100644 --- a/astrai/serialization.py +++ b/astrai/serialization.py @@ -136,7 +136,7 @@ def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict: class Checkpoint: state_dict: Dict[str, Any] = field(default_factory=dict) epoch: int = 0 - iteration: int = 0 + consumed_samples: int = 0 extra: Dict[str, Any] = field(default_factory=dict) meta: Dict[str, Any] = field(default_factory=dict) config: Dict[str, Any] = field(default_factory=dict) @@ -150,7 +150,7 @@ class Checkpoint: meta = { "epoch": self.epoch, - "iteration": self.iteration, + "consumed_samples": self.consumed_samples, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), **self.meta, } @@ -176,7 +176,7 @@ class Checkpoint: return cls( state_dict=state_dict, epoch=meta.get("epoch", 0), - iteration=meta.get("iteration", 0), + consumed_samples=meta.get("consumed_samples", 0), extra=extra, config=config, ) diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index ab22421..c743a9b 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -139,34 +139,35 @@ class CheckpointCallback(TrainCallback): self.interval = interval self.weight_only = weight_only self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra - self.last_ckpt_iter = 0 + self.last_ckpt_step = 0 def _save_checkpoint(self, context: TrainContext): state_dict = context.executor.unwrap_model(context.model) - self.last_ckpt_iter = context.iteration + self.last_ckpt_step = context.optimizer_step if get_rank() == 0: save_path = os.path.join( - self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}" + self.save_dir, + f"epoch_{context.epoch}_step_{context.optimizer_step}", ) extra = self.save_extra_fn(context) meta = context.config.to_dict() context.checkpoint = Checkpoint( state_dict=state_dict, epoch=context.epoch, - iteration=context.iteration, + consumed_samples=context.consumed_samples, + config=context.model_config, extra=extra, meta=meta, - config=context.model_config, ) context.checkpoint.save(save_path) def on_batch_end(self, context: TrainContext): - if context.iteration - self.last_ckpt_iter >= self.interval: + if context.optimizer_step - self.last_ckpt_step >= self.interval: self._save_checkpoint(context) def on_train_end(self, context: TrainContext): - if context.iteration != self.last_ckpt_iter: + if context.optimizer_step != self.last_ckpt_step: self._save_checkpoint(context) def on_error(self, context: TrainContext): @@ -232,7 +233,7 @@ class MetricLoggerCallback(TrainCallback): log_interval: int = 1, metrics: List[str] = None, ): - self.last_log_iter = 0 + self.last_log_flush_step = 0 self._last_val_loss = None self._last_log_step = 0 self.save_interval = save_interval @@ -264,31 +265,30 @@ class MetricLoggerCallback(TrainCallback): "type": event_type, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), "epoch": context.epoch, - "step": context.iteration // context.config.grad_accum_steps, - "iter": context.iteration, + "step": context.optimizer_step, + "consumed_samples": context.consumed_samples, **extra, } self.log_cache.append(entry) @only_on_rank(0) - def _flush(self, epoch, iter): - log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl" + def _flush(self, epoch, consumed): + log_file = self.log_dir / f"epoch_{epoch}_consumed_{consumed}_metric.jsonl" log_file.parent.mkdir(parents=True, exist_ok=True) with open(log_file, "w") as f: for log in self.log_cache: f.write(json.dumps(log) + "\n") def on_batch_end(self, context): - if context.iteration - self.last_log_iter >= self.save_interval: - self._flush(context.epoch, context.iteration) - self.last_log_iter = context.iteration + if context.optimizer_step - self.last_log_flush_step >= self.save_interval: + self._flush(context.epoch, context.optimizer_step) + self.last_log_flush_step = context.optimizer_step def on_optimizer_step(self, context): - step = context.iteration // context.config.grad_accum_steps - if step - self._last_log_step >= self.log_interval: + if context.optimizer_step - self._last_log_step >= self.log_interval: step_metrics = [m for m in self.metrics if m != "val_loss"] self._append("step", context, **self._metrics(context, step_metrics)) - self._last_log_step = step + self._last_log_step = context.optimizer_step if context.val_loss is not None and context.val_loss != self._last_val_loss: self._append("validation", context, val_loss=context.val_loss) self._last_val_loss = context.val_loss @@ -297,11 +297,11 @@ class MetricLoggerCallback(TrainCallback): self._append("epoch", context) def on_train_end(self, context): - if context.iteration != self.last_log_iter: - self._flush(context.epoch, context.iteration) + if context.optimizer_step != self.last_log_flush_step: + self._flush(context.epoch, context.optimizer_step) def on_error(self, context): - self._flush(context.epoch, context.iteration) + self._flush(context.epoch, context.optimizer_step) @CallbackFactory.register("validation") @@ -328,9 +328,9 @@ class ValidationCallback(TrainCallback): context.val_loss = avg_loss context.model.train() - step_count = context.iteration // context.config.grad_accum_steps logger.info( - f"Epoch {context.epoch + 1}, Step {step_count}, Val Loss: {avg_loss:.4f}" + f"Epoch {context.epoch + 1}, Step {context.optimizer_step}, " + f"Val Loss: {avg_loss:.4f}" ) def on_optimizer_step(self, context: TrainContext): @@ -339,6 +339,5 @@ class ValidationCallback(TrainCallback): cfg = context.config if cfg.val_step <= 0: return - step_count = context.iteration // cfg.grad_accum_steps - if step_count % cfg.val_step == 0: + if context.optimizer_step % cfg.val_step == 0: self._run_validation(context) diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 783cd68..d2df191 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -29,7 +29,7 @@ class TrainContext: executor: BaseExecutor = field(default=None) epoch: int = field(default=0) - iteration: int = field(default=0) + consumed_samples: int = field(default=0) loss: float = field(default=0.0) grad_norm: Optional[float] = field(default=None) val_dataloader: Optional[DataLoader] = field(default=None) @@ -39,6 +39,14 @@ class TrainContext: rank: int = field(default=0) kwargs: Dict[str, Any] = field(default_factory=dict) + @property + def optimizer_step(self) -> int: + return self.consumed_samples // ( + self.config.batch_per_device + * self.world_size + * self.config.grad_accum_steps + ) + class TrainContextBuilder: def __init__( @@ -90,7 +98,10 @@ class TrainContextBuilder: if checkpoint.config: context.model_config = checkpoint.config context.epoch = checkpoint.epoch or cfg.start_epoch - context.iteration = checkpoint.iteration or cfg.start_batch + if checkpoint.consumed_samples > 0: + context.consumed_samples = checkpoint.consumed_samples + else: + context.consumed_samples = cfg.start_samples * context.world_size context.checkpoint = checkpoint if cfg.lora is not None: @@ -116,7 +127,7 @@ class TrainContextBuilder: cfg.dataset, [n_train, n_val], generator=generator ) - sampler_offset = context.iteration * cfg.batch_per_device + sampler_offset = context.consumed_samples // context.world_size sampler = ResumableDistributedSampler( data_source=train_dataset, start_epoch=context.epoch, diff --git a/astrai/trainer/trainer.py b/astrai/trainer/trainer.py index c8c73a1..1c76fc5 100644 --- a/astrai/trainer/trainer.py +++ b/astrai/trainer/trainer.py @@ -74,7 +74,9 @@ class Trainer: context.loss = loss.item() stand_loss = loss / executor.grad_accum_steps executor.backward(stand_loss) - context.iteration += 1 + context.consumed_samples += ( + context.config.batch_per_device * context.world_size + ) self._call_callbacks("on_batch_end", context) if executor.sync_gradients: diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 10ad134..c9acf8a 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -175,7 +175,10 @@ def parse_args() -> argparse.Namespace: "--start_epoch", type=int, default=0, help="Start epoch for training." ) parser.add_argument( - "--start_batch", type=int, default=0, help="Start batch for training." + "--start_samples", + type=int, + default=0, + help="Start samples (per rank) for training.", ) parser.add_argument( @@ -317,7 +320,7 @@ def train( n_epoch: int, batch_per_device: int, start_epoch: int, - start_batch: int, + start_samples: int, grad_accum_steps: int, warmup_ratio: float, ckpt_interval: int, @@ -444,7 +447,7 @@ def train( n_epoch=n_epoch, batch_per_device=batch_per_device, start_epoch=start_epoch, - start_batch=start_batch, + start_samples=start_samples, ckpt_interval=ckpt_interval, grad_accum_steps=grad_accum_steps, max_grad_norm=max_grad_norm, diff --git a/tests/conftest.py b/tests/conftest.py index 47b484a..087b87c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -75,7 +75,7 @@ class MultiTurnDataset(Dataset): class EarlyStoppingDataset(Dataset): - """Dataset that triggers early stopping after a specified number of iterations.""" + """Dataset that triggers early stopping after consuming a specified number of samples.""" def __init__(self, length=10, stop_after=5): self.length = length diff --git a/tests/data/test_checkpoint.py b/tests/data/test_checkpoint.py index bfac737..9d8c858 100644 --- a/tests/data/test_checkpoint.py +++ b/tests/data/test_checkpoint.py @@ -25,7 +25,9 @@ def test_single_process(): scheduler.step() - checkpoint = Checkpoint(state_dict=model.state_dict(), epoch=3, iteration=30) + checkpoint = Checkpoint( + state_dict=model.state_dict(), epoch=3, consumed_samples=120 + ) with tempfile.TemporaryDirectory() as tmpdir: checkpoint.save(tmpdir) @@ -33,7 +35,7 @@ def test_single_process(): loaded_checkpoint = Checkpoint.load(tmpdir) assert loaded_checkpoint.epoch == 3 - assert loaded_checkpoint.iteration == 30 + assert loaded_checkpoint.consumed_samples == 120 def test_checkpoint_with_extra(): @@ -46,7 +48,10 @@ def test_checkpoint_with_extra(): "scheduler": {"last_epoch": 5}, } checkpoint = Checkpoint( - state_dict=model.state_dict(), epoch=1, iteration=10, extra=extra + state_dict=model.state_dict(), + epoch=1, + consumed_samples=40, + extra=extra, ) with tempfile.TemporaryDirectory() as tmpdir: @@ -77,7 +82,7 @@ def simple_training(): checkpoint = Checkpoint( state_dict=model.state_dict(), epoch=2, - iteration=10, + consumed_samples=40, ) rank = get_rank() diff --git a/tests/trainer/conftest.py b/tests/trainer/conftest.py index e7cca22..bc02316 100644 --- a/tests/trainer/conftest.py +++ b/tests/trainer/conftest.py @@ -52,7 +52,7 @@ def create_train_config( batch_per_device: Batch size per device (default: 2) grad_accum_steps: Gradient accumulation steps (default: 1) max_grad_norm: Maximum gradient norm for clipping (default: 1.0) - ckpt_interval: Checkpoint save interval in iterations (default: 5) + ckpt_interval: Checkpoint save interval in optimizer steps (default: 5) random_seed: Random seed for reproducibility (default: 42) **kwargs: Additional arguments passed to TrainConfig diff --git a/tests/trainer/test_early_stopping.py b/tests/trainer/test_early_stopping.py index 729d069..81bd490 100644 --- a/tests/trainer/test_early_stopping.py +++ b/tests/trainer/test_early_stopping.py @@ -44,14 +44,14 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset): pass # Resume from latest checkpoint - load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_iter_2") + load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_step_1") trainer = Trainer(train_config) trainer.train(resume_dir=load_dir) - # Verify checkpoint was saved at expected iteration - load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_iter_10") + # Verify checkpoint was saved at expected step + load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_step_5") import json with open(os.path.join(load_dir, "meta.json")) as f: meta = json.load(f) - assert meta["iteration"] == 10 + assert meta["consumed_samples"] == 20