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

fix: scheduler division-by-zero, loss_mask bool 

- schedule.py: guard warmup_steps/lr_decay_steps against zero
- strategy.py: use ~loss_mask instead of loss_mask==0 on bool tensor
This commit is contained in:
ViperEkura 2026-07-03 22:03:36 +08:00
parent 70c0e5de90
commit 27d1921d9c
2 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
astrai/trainer/schedule.py
View File

@ -65,11 +65,15 @@ class CosineScheduler(BaseScheduler):
def get_lr(self) -> List[float]:
# warmup
if self.last_epoch < self.warmup_steps:
warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
warmup_factor = max(
self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
)
return [base_lr * warmup_factor for base_lr in self.base_lrs]
# cosine decay
decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
decay_progress = (self.last_epoch - self.warmup_steps) / max(
self.lr_decay_steps, 1
)
decay_progress = min(decay_progress, 1.0)
cosine_decay = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
decay_factor = max(self.min_rate, cosine_decay)
@ -118,7 +122,9 @@ class SGDRScheduler(BaseScheduler):
def get_lr(self):
# warmup
if self.last_epoch < self.warmup_steps:
warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
warmup_factor = max(
self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
)
return [base_lr * warmup_factor for base_lr in self.base_lrs]
# SGDR

2
astrai/trainer/strategy.py
View File

@ -196,7 +196,7 @@ class SFTStrategy(BaseStrategy):
ignore_index = -100
input_mask = make_doc_boundary_mask(position_ids)
target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
target_ids = target_ids.masked_fill(~loss_mask, ignore_index)
logits = self.model(
input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
)["logits"]