Patch EOD out of inference results

Author: tdene

megatron/core/inference/data_parallel_inference_coordinator.py

@@ -504,9 +504,15 @@ def detokenize(self, finished_request):
             finished_request["prompt"] = self.tokenizer.detokenize(
                 finished_request["prompt_tokens"][1]
             )
-        finished_request["generated_text"] = self.tokenizer.detokenize(
-            finished_request["generated_tokens"]
-        )
+        generated_tokens = finished_request["generated_tokens"]
+        termination_id = (finished_request.get("sampling_params", {}) or {}).get("termination_id")
+        while (
+            generated_tokens
+            and termination_id is not None
+            and generated_tokens[-1] == termination_id
+        ):
+            generated_tokens = generated_tokens[:-1]
+        finished_request["generated_text"] = self.tokenizer.detokenize(generated_tokens)
 
     @classmethod
     def entrypoint(

megatron/core/inference/engines/dynamic_engine.py

@@ -1638,9 +1638,15 @@ async def async_bookkeep(
                         request.prompt = self.controller.tokenizer.detokenize(
                             request.prompt_tokens.tolist()
                         )
-                    request.generated_text = self.controller.tokenizer.detokenize(
-                        request.generated_tokens
-                    )
+                    generated_tokens = request.generated_tokens
+                    termination_id = request.sampling_params.termination_id
+                    while (
+                        generated_tokens
+                        and termination_id is not None
+                        and generated_tokens[-1] == termination_id
+                    ):
+                        generated_tokens = generated_tokens[:-1]
+                    request.generated_text = self.controller.tokenizer.detokenize(generated_tokens)
             range_pop()
 
         # Handle necessary ZMQ DP coordinator communication.