Allow mixed-batch sampling in dynamic inference

examples/inference/gpt/gpt_dynamic_inference.py

@@ -117,8 +117,11 @@ def get_model() -> MegatronModule:
     return model
 
 
-def get_inference_context(requests: List[Request], sampling_params: SamplingParams, 
-                          calculate_max_sequence_length_from_requests: bool =True):
+def get_inference_context(
+    requests: List[Request],
+    sampling_params: SamplingParams,
+    calculate_max_sequence_length_from_requests: bool = True
+):
     """The inference context manages the KV cache and other inference state."""
 
     args = get_args()
@@ -199,13 +202,12 @@ def get_inference_controller(
 
 
 def run_inference(
-    requests: List[Request], sampling_params: SamplingParams, engine: DynamicInferenceEngine
+    requests: List[Request], engine: DynamicInferenceEngine
 ) -> List[Dict[str, float]]:
     """Add requests to engine and generate tokens.
 
     Args:
         requests (List[Request]): Requests that are to be added and processed.
-        sampling_params (SamplingParams): Sampling params for the logits.
         engine (DynamicInferenceEngine): Inference engine that manages generating tokens.
 
     Return:
@@ -244,7 +246,7 @@ def _add_request():
         engine.add_request(
             num_requests_added,
             _request.prompt_text,
-            sampling_params.num_tokens_to_generate,
+            _request.sampling_params,
         )
         _request.time_start = get_curr_time()
         _request.state = "started"
@@ -301,7 +303,7 @@ def _add_request():
                 request.output_text = finished_request.generated_text
                 request.state = "finished"
                 request.request_id = finished_request.request_id
-                if sampling_params.return_log_probs:
+                if finished_request.sampling_params.return_log_probs:
                     request.log_probs = (
                         finished_request.prompt_log_probs + finished_request.generated_log_probs
                     )
@@ -349,11 +351,12 @@ def main():
         top_p=args.top_p,
         return_log_probs=args.return_log_probs,
         num_tokens_to_generate=args.num_tokens_to_generate,
+        termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,
     )
 
     # Requests, context, conroller.
     model = get_model()
-    requests = build_requests(args, tokenizer)
+    requests = build_requests(args, tokenizer, sampling_params)
     context = get_inference_context(requests, sampling_params)
     controller = get_inference_controller(model, context)
 
@@ -371,7 +374,6 @@ def main():
     engine = DynamicInferenceEngine(
         controller,
         context,
-        termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,
         enable_cuda_graph=args.cuda_graph_impl == "local",
         random_seed=args.seed,
         track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
@@ -387,7 +389,7 @@ def main():
     throughputs = []
     for _ in range(args.inference_repeat_n):
         t = get_curr_time()
-        result = run_inference(requests, sampling_params, engine)
+        result = run_inference(requests, engine)
         step_times = result["step_times"]
         add_times = result["add_times"]
         output_times = result["output_times"]
@@ -458,7 +460,7 @@ def escape_str(s):
                         "cuda_graph_request_count_map" : result["cuda_graph_request_count_map"],
                         "step_count" : engine.step_count,
                     }
-                    if sampling_params.return_log_probs:
+                    if req.sampling_params.return_log_probs:
                         response_logprobs = req.log_probs
                         result_dict["logprobs"] = response_logprobs
                     json_results[req.request_id] = result_dict

examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py

@@ -18,13 +18,13 @@
 from megatron.training.arguments import parse_args
 from megatron.core import parallel_state
 
-async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling_params: SamplingParams, port: int):
+async def main(engine: DynamicInferenceEngine, requests: List[Request], port: int):
     # once you call engine.start_listening_to_data_parallel_coordinator,
     # the engine will start accepting requests from the data parallel coordinator.
     # and processing them in an asyncio coroutine. 
-    await engine.start_listening_to_data_parallel_coordinator(sampling_params, 
-                                                inference_coordinator_port=port,
-                                                launch_inference_coordinator=True)   
+    await engine.start_listening_to_data_parallel_coordinator( 
+        inference_coordinator_port=port, launch_inference_coordinator=True
+    )
     # if you want to use your own inference coordinator - 
     # 1. set launch_inference_coordinator to False
     # 2. setup a router socket at tcp://MASTER_ADDR:PORT
@@ -50,8 +50,7 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
                 # These add-request calls will queue up the request on a zmq socket and return
                 # instantaneously. They will return an asyncio future which can be awaited for
                 # request completion.
-                futures.append(client.add_request(request.prompt_text, 
-                                                        sampling_params))
+                futures.append(client.add_request(request.prompt_text, request.sampling_params))
                 num_requests_added += 1
                 #tbar.update(1)
             if num_requests_added == num_requests_total:
@@ -74,7 +73,7 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
                     "generated_tokens": req.generated_tokens,
                     "latency": req.latency, #InferenceClient populates this field in the returned future.
                 }
-                if sampling_params.return_log_probs:
+                if req.sampling_params.return_log_probs:
                     result_dict["logprobs"] = req.prompt_log_probs + req.generated_log_probs
                 json_results[req.request_id] = result_dict
             with open(args.output_path, "w") as fp:
@@ -119,9 +118,8 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
 
         # Requests, context, conroller.
         model = get_model()
-        requests = build_requests(args, tokenizer) if dist.get_rank() == 0 else None
+        requests = build_requests(args, tokenizer, sampling_params) if dist.get_rank() == 0 else None
 
-
         context = get_inference_context(None, 
                                         None,
                                         calculate_max_sequence_length_from_requests=False)
@@ -147,6 +145,5 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
 
         asyncio.run(main(engine, 
                         requests,
-                        sampling_params,
                         args.inference_coordinator_port))
 

examples/inference/gpt/utils.py

@@ -13,6 +13,8 @@
 from megatron.core.inference.contexts import DynamicInferenceContext
 from megatron.core.transformer.module import MegatronModule
 
+from megatron.core.inference.sampling_params import SamplingParams
+
 
 
 def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
@@ -130,6 +132,16 @@ def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
     return parser
 
 
+def get_default_sampling_params(termination_id: int = None):
+    return SamplingParams(
+    temperature=1.0,
+    top_k=1,
+    top_p=0.0,
+    return_log_probs=False,
+    num_tokens_to_generate=30,
+    termination_id = termination_id,
+)
+
 def get_curr_time() -> float:
     """Get synchronized time across ranks."""
     curr_time = torch.cuda.LongTensor([time.time_ns()])
@@ -153,7 +165,7 @@ class Request:
         tokenizer (Any): Tokenizer for tokenizing the prompt.
     """
 
-    def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any):
+    def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any, sampling_params: SamplingParams = None):
         self.prompt_text = prompt_text
         self.prompt_tokens = tokenizer.tokenize(prompt_text)
         self.output_text = None
@@ -163,6 +175,7 @@ def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any):
         self.time_start = None
         self.time_end = None
         self.state = "not-started"
+        self.sampling_params: SamplingParams = sampling_params if sampling_params is not None else get_default_sampling_params(tokenizer.eod)
 
     def __str__(self) -> str:
         return "state '%s'; toffset %.1e; prompt len %d; output len %d; '%s'" % (
@@ -231,7 +244,9 @@ def get_cli_requests(args: Namespace, tokenizer: Any) -> list[Request]:
     return requests
 
 
-def get_synthetic_requests(args: Namespace, tokenizer: Any) -> list[Request]:
+def get_synthetic_requests(
+    args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
     """Get example requests."""
 
     # Get time offsets.
@@ -251,7 +266,9 @@ def get_synthetic_requests(args: Namespace, tokenizer: Any) -> list[Request]:
     return requests
 
 
-def get_requests_from_file(args: Namespace, tokenizer: Any) -> list[Request]:
+def get_requests_from_file(
+    args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
     """Get requests from a file."""
     if not args.prompt_file:
         raise ValueError("Prompt file is required to read requests from a file.")
@@ -275,23 +292,25 @@ def get_requests_from_file(args: Namespace, tokenizer: Any) -> list[Request]:
 
     # Init requests.
     requests = [
-        Request(p, t, tokenizer)
+        Request(p, t, tokenizer, sampling_params)
         for p, t in tqdm(zip(prompts, time_offsets), "init requests", total=len(prompts))
     ]
 
     return requests
 
 
-def build_requests(args: Namespace, tokenizer: Any) -> list[Request]:
+def build_requests(
+    args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
     # Check if we have any prompts (from command line or JSONL)
     if args.prompts:
         if args.prompt_file:
             raise ValueError("Cannot use both --prompts and --prompt-file")
         return get_cli_requests(args, tokenizer)
     elif args.prompt_file:
-        return get_requests_from_file(args, tokenizer)
+        return get_requests_from_file(args, tokenizer, sampling_params)
     else:
-        return get_synthetic_requests(args, tokenizer)
+        return get_synthetic_requests(args, tokenizer, sampling_params)
 
 
 def get_model_size_str(model):