fix: 修复 CLI 参数缺失/重复、device_ids 越界、generate 参数名不一致、scheduler 时序、非流式截断等 bug

- train.py: 补上 --batch_size、--grpo_clip_eps，删除 3 处重复 --group_size
- generate.py: --model_dir 改为 --param_path 对齐 README
- automodel.py: from_pretrained 新增 strict 参数（默认 True）
- parallel/setup.py: 修复 device_ids 索引越界
- train_callback.py: scheduler.step() 移至 on_step_end
- test_train_strategy.py: 测试中补 optimizer.step()
- engine.py: 非流式改为循环等待所有任务完成，补 remove_task 清理
- scheduler.py: Task 添加 _pages_freed 标志，杜绝双重释放
- trainer.py: accumulation_steps=0 时 clamp 为 1
- tokenizer.py: save_pretrained 添加 _tokenizer is None 检查
- benchmark.py: 修复 ModelConfig 过时 import 路径
- inference/__init__.py: 修复 stale docstring

astrai/inference/__init__.py

@@ -3,7 +3,7 @@
 Layers:
   - engine.py:    Facade (InferenceEngine), Value Object (GenerationParams, GenerationRequest)
   - scheduler.py: Continuous-batching loop, Task state machine, TaskStatus enum
-  - cache.py:     Object Pool (SlotAllocator), PrefixCacheManager
+  - cache.py:     PagedCache (page-table-indirected KV cache with alloc/free)
   - sampling.py:  Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
   - server.py:    FastAPI HTTP server (OpenAI-compatible endpoints)
 """

astrai/inference/engine.py

@@ -408,13 +408,14 @@ class InferenceEngine:
             Single string for one prompt, list of strings for batch.
         """
         result = _Result(count=len(prompts))
+        task_ids = []
 
         for i, p in enumerate(prompts):
 
             def make_cb(idx):
                 return lambda tok: result.append(tok, idx)
 
-            self.scheduler.add_task(
+            task_id = self.scheduler.add_task(
                 prompt=p,
                 max_tokens=max_tokens,
                 temperature=temperature,
@@ -422,8 +423,14 @@ class InferenceEngine:
                 top_k=top_k,
                 stream_callback=make_cb(i),
             )
+            task_ids.append(task_id)
+
+        while result._completed < result._total:
+            result.wait(timeout=1.0)
+
+        for task_id in task_ids:
+            self.scheduler.remove_task(task_id)
 
-        result.wait()
         res = result.get_results()
         return res if is_batch else res[0]
 

astrai/inference/scheduler.py

@@ -56,6 +56,7 @@ class Task:
         self.arrival_time = time.time()
         self.finish_time: Optional[float] = None
         self.stream_callback = stream_callback
+        self._pages_freed: bool = False
 
     @property
     def next_pos(self) -> int:
@@ -167,9 +168,11 @@ class InferenceScheduler:
             self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
 
         for task in removed_active:
-            self._free_pages(task.page_table)
-            task.page_table.clear()
-            task.n_pages = 0
+            if not task._pages_freed:
+                self._free_pages(task.page_table)
+                task.page_table.clear()
+                task.n_pages = 0
+                task._pages_freed = True
 
     def _free_pages(self, indices: List[int]) -> None:
         for idx in indices:
@@ -185,9 +188,11 @@ class InferenceScheduler:
                 self._total_tokens += task.output_tokens
 
         for task in finished:
-            self._free_pages(task.page_table)
-            task.page_table.clear()
-            task.n_pages = 0
+            if not task._pages_freed:
+                self._free_pages(task.page_table)
+                task.page_table.clear()
+                task.n_pages = 0
+                task._pages_freed = True
 
         self.active_tasks = [
             t for t in self.active_tasks if t.status != TaskStatus.FINISHED

astrai/model/automodel.py

@@ -84,6 +84,7 @@ class AutoModel(nn.Module):
         cls,
         path: Union[str, Path],
         disable_random_init: bool = True,
+        strict: bool = True,
     ) -> nn.Module:
 
         model_path = Path(path)
@@ -106,7 +107,7 @@ class AutoModel(nn.Module):
         weights_path = model_path / "model.safetensors"
         if weights_path.exists():
             state_dict = st.load_file(str(weights_path))
-            model.load_state_dict(state_dict, strict=False)
+            model.load_state_dict(state_dict, strict=strict)
 
         return model
 

astrai/parallel/setup.py

@@ -48,8 +48,9 @@ def setup_parallel(
     if device_ids is None:
         device_ids = [i for i in range(world_size)]
 
-    rank = device_ids[rank % len(device_ids)]
-    device_id = torch.device(device_type, device_ids[rank])
+    effective_rank = rank % len(device_ids)
+    device_id = torch.device(device_type, device_ids[effective_rank])
+    rank = device_ids[effective_rank]
 
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port

astrai/tokenize/tokenizer.py

@@ -64,6 +64,11 @@ class AutoTokenizer:
             save_path: Path to save the tokenizer
         """
 
+        if self._tokenizer is None:
+            raise RuntimeError(
+                "Tokenizer not initialized. Load or create a tokenizer first."
+            )
+
         save_path = Path(save_path)
         save_path.mkdir(parents=True, exist_ok=True)
 

astrai/trainer/train_callback.py

@@ -104,7 +104,7 @@ class SchedulerCallback(TrainCallback):
             if "initial_lr" not in group:
                 group["initial_lr"] = group["lr"]
 
-    def on_batch_end(self, context: TrainContext):
+    def on_step_end(self, context: TrainContext):
         if context.scheduler:
             context.scheduler.step()
 

astrai/trainer/trainer.py

@@ -68,8 +68,9 @@ class Trainer:
                 context.epoch = epoch
                 self._call_callbacks("on_epoch_begin", context)
 
+                accumulation_steps = max(self.train_config.accumulation_steps, 1)
                 for batch in context.dataloader:
-                    if context.iteration % self.train_config.accumulation_steps == 0:
+                    if context.iteration % accumulation_steps == 0:
                         # 2. step
                         self._call_callbacks("on_step_begin", context)
                         context.optimizer.step()
@@ -83,7 +84,7 @@ class Trainer:
                     context.iteration += 1
 
                     # to make the loss normalized by accumulation steps
-                    stand_loss = loss / self.train_config.accumulation_steps
+                    stand_loss = loss / accumulation_steps
                     stand_loss.backward()
 
                     self._call_callbacks("on_batch_end", context)