diff --git a/areal/core/remote_inf_engine.py b/areal/core/remote_inf_engine.py index b835785357..f713a83a77 100644 --- a/areal/core/remote_inf_engine.py +++ b/areal/core/remote_inf_engine.py @@ -739,9 +739,10 @@ async def agenerate(self, req: ModelRequest) -> ModelResponse: # Deal with rollout interruption stop_reason = None + ori_max_new_tokens = gconfig.max_new_tokens while ( stop_reason not in ["stop", "tool_calls", "length"] - and len(accumulated_output_tokens) < gconfig.max_new_tokens + and len(accumulated_output_tokens) < ori_max_new_tokens ): # Request is interrupted, wait for some time to avoid interfering # with update weights requests