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

fix: resume optimizer LR, step display, and consumed_samples alignment

This commit is contained in:
ViperEkura 2026-07-15 08:52:15 +08:00
parent 13998da15a
commit bb175fda91
3 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
astrai/trainer/train_callback.py
View File

@ -210,7 +210,7 @@ class ProgressBarCallback(TrainCallback):
@only_on_rank(0) @only_on_rank(0)
def on_optimizer_step(self, context: TrainContext): def on_optimizer_step(self, context: TrainContext):
postfix = { postfix = {
"step": context.optimizer_step, "step": f"{context.optimizer_step:d}",
"loss": f"{context.loss:.4f}", "loss": f"{context.loss:.4f}",
"lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}", "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
} }

7
astrai/trainer/train_context.py
View File

@ -99,7 +99,12 @@ class TrainContextBuilder:
context.model_config = checkpoint.config context.model_config = checkpoint.config
context.epoch = checkpoint.epoch or cfg.start_epoch context.epoch = checkpoint.epoch or cfg.start_epoch
if checkpoint.consumed_samples > 0: if checkpoint.consumed_samples > 0:
context.consumed_samples = checkpoint.consumed_samples per_step = (
cfg.batch_per_device * context.world_size * cfg.grad_accum_steps
)
context.consumed_samples = (
checkpoint.consumed_samples // per_step
) * per_step
else: else:
context.consumed_samples = cfg.start_samples * context.world_size context.consumed_samples = cfg.start_samples * context.world_size
context.checkpoint = checkpoint context.checkpoint = checkpoint

1
scripts/tools/train.py
View File

@ -90,6 +90,7 @@ class MuonMix(optim.Optimizer):
def load_state_dict(self, state_dict: Dict[str, Any]): def load_state_dict(self, state_dict: Dict[str, Any]):
self.muon.load_state_dict(state_dict["muon"]) self.muon.load_state_dict(state_dict["muon"])
self.adamw.load_state_dict(state_dict["adamw"]) self.adamw.load_state_dict(state_dict["adamw"])
self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
def parse_args() -> argparse.Namespace: def parse_args() -> argparse.Namespace: