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

refactor : replace iteration with consumed_samples 

- 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:33 +08:00
parent 44579ea6dc
commit aabb0d83e9
10 changed files with 72 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
astrai/config/train_config.py
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

6
astrai/serialization.py
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,
)

49
astrai/trainer/train_callback.py
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)

17
astrai/trainer/train_context.py
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,

4
astrai/trainer/trainer.py
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:

9
scripts/tools/train.py
View File

@ -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,

2
tests/conftest.py
View File

@ -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

13
tests/data/test_checkpoint.py
View File

@ -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()

2
tests/trainer/conftest.py
View File

@ -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

8
tests/trainer/test_early_stopping.py
View File

@ -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