From 27d1921d9cb48becfebd8157fdfe0c7954c8ae5f Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Fri, 3 Jul 2026 22:03:36 +0800
Subject: [PATCH] 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
---
 astrai/trainer/schedule.py | 12 +++++++++---
 astrai/trainer/strategy.py |  2 +-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/astrai/trainer/schedule.py b/astrai/trainer/schedule.py
index 496c930..c9db9bd 100644
--- a/astrai/trainer/schedule.py
+++ b/astrai/trainer/schedule.py
@@ -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
diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py
index aca7e08..65bc310 100644
--- a/astrai/trainer/strategy.py
+++ b/astrai/trainer/strategy.py
@@ -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"]