fix: 断点续训恢复优化器/调度器状态及采样器剩余长度

- 使用Checkpoint.load()替代手动加载model.safetensors，恢复optimizer/scheduler状态
- TrainContextBuilder从checkpoint.extra恢复优化器和调度器state_dict
- ResumableDistributedSampler.__len__返回剩余样本数而非总数
- 训练前对state_dict置空避免mp.spawn pickle 7GB大对象

astrai/dataset/sampler.py

@@ -43,6 +43,7 @@ class ResumableDistributedSampler(Sampler[int]):
         offset = 0 if drop_last else self.num_replicas - 1
         self.num_samples_per_replica = (self.num_samples + offset) // self.num_replicas
         self.total_size = self.num_samples_per_replica * self.num_replicas
+        self.iter = self.iter % self.num_samples_per_replica
 
         self._indices = None
 
@@ -74,5 +75,10 @@ class ResumableDistributedSampler(Sampler[int]):
         self.epoch += 1
         self._indices = None
 
+    @property
+    def _remaining(self):
+        remaining = self.num_samples_per_replica - self.iter
+        return max(remaining, 0)
+
     def __len__(self):
-        return self.num_samples_per_replica
+        return self._remaining

astrai/trainer/train_context.py

@@ -71,7 +71,8 @@ class TrainContextBuilder:
         if self._checkpoint is not None:
             context.epoch = max(self._checkpoint.epoch, cfg.start_epoch)
             context.iteration = max(self._checkpoint.iteration, cfg.start_batch)
-            context.model.load_state_dict(self._checkpoint.state_dict)
+            if self._checkpoint.state_dict:
+                context.model.load_state_dict(self._checkpoint.state_dict)
             context.checkpoint = self._checkpoint
         else:
             context.checkpoint = Checkpoint(