ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor : replace iteration with consumed_samples

Browse Source 
- Replace context.iteration with consumed_samples (global sample count)
- Add optimizer_step property derived from consumed_samples
- Checkpoint meta.json stores consumed_samples, drops iteration
- CLI --start_batch renamed to --start_samples (per-rank samples)
- Checkpoint dir naming: epoch_X_step_Y instead of epoch_X_iter_Y
- Metric log entries use step and consumed_samples fields
- Backward compat removed (old iteration checkpoints unsupported)
This commit is contained in:
ViperEkura
2026-06-30 18:42:42 +08:00
parent 44579ea6dc
commit aabb0d83e9
10 changed files with 72 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/config/train_config.py
+7 -3
View File
@@ -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
astrai/serialization.py
+3 -3
View File
@@ -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,
)
astrai/trainer/train_callback.py
+24 -25
View File
@@ -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)
astrai/trainer/train_context.py
+14 -3
View File
@@ -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,
astrai/trainer/trainer.py
+3 -1
View File
@@ -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: