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
This commit is contained in:
@@ -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]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user