Move to using the Inference OpenAI API server (bis)

examples/rl/README.md

@@ -172,7 +172,6 @@ torchrun \
     --save $CHECKPOINT_DIR \
     --load $CHECKPOINT_DIR \
     --tensorboard-dir $TB_DIR \
-    --langrl-inference-server-type inplace_megatron \
     --seed $SEED \
     --sequence-parallel \
     --finetune \

examples/rl/environment_configs/gsm8k_nanov3.yaml

@@ -2,8 +2,6 @@
   agent_args:
     answer_format: "boxed"
     format_reward: 0.5
-    assistant_suffix: "Assistant: "
-    chat_mode: true
     negative_reward: 0.0
     partial_end_reward: 0.75
   weight: 1.0

examples/rl/environments/countdown/countdown_agent.py

@@ -12,15 +12,8 @@
 class CountdownAgent(RewardOnlyAgent, HFDatasetAgent):
 
     def make_prefix(self, target, nums) -> str:
-        if self.chat_mode:
-            prefix = f"""Using the numbers {nums}, create an equation that equals {target}. You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. 
+        prefix = f"""Using the numbers {nums}, create an equation that equals {target}. You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. 
         Return the final answer in <answer> </answer> tags, for example <answer> (1 + 2) / 3 </answer>. Do not include an = sign."""
-        else:
-            prefix = f"""A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer.
-        User: Using the numbers {nums}, create an equation that equals {target}. You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. Show your work in <think> </think> tags. 
-        And return the final answer in <answer> </answer> tags, for example <answer> (1 + 2) / 3 </answer>. Do not include an = sign.
-        Assistant: Let me solve this step by step.
-        <think>"""
         return prefix
 
     def get_dataset(self, validation: bool = False):

examples/rl/environments/math/gsm8k_agent.py

@@ -25,16 +25,12 @@
 class GSM8KAgent(MathAgent):
     def __init__(self,
         answer_format: str = "boxed",
-        chat_mode: bool = False,
-        assistant_suffix: str = "Assistant: Let me solve this step by step.\n<think>",
         format_reward: float = 0.0,
         negative_reward: float = 0.0,
         partial_end_reward: float = 0.0,
         **kwargs):
         super().__init__(
             answer_format=answer_format,
-            chat_mode=chat_mode,
-            assistant_suffix=assistant_suffix,
             format_reward=format_reward,
             negative_reward=negative_reward,
             partial_end_reward=partial_end_reward,

examples/rl/environments/math/math_agent.py

@@ -3,7 +3,6 @@
 import re
 import traceback
 
-from megatron.rl.agent.pass_at_evaluation_agent import PassAtEvaluationAgent
 from megatron.rl.agent.reward_only_agent import RewardOnlyAgent
 
 try:
@@ -25,8 +24,6 @@ class MathAgent(RewardOnlyAgent):
     def __init__(self,
         format_reward: float = 0.0,
         answer_format: str = "tagged",
-        assistant_suffix: str = "Assistant: Let me solve this step by step.\n<think>",
-        chat_mode: bool = False,
         negative_reward: float = 0.0,
         partial_end_reward: float = 0.0,
         **kwargs):
@@ -36,9 +33,6 @@ def __init__(self,
                 even if the answer is incorrect or is missing the end-of-text token.
             answer_format (str): Which answer format is expected: "tagged" for <answer> tags,
                 or "boxed" for \boxed{} LaTeX formatting.
-            assistant_suffix (str): The suffix string included in the assistant's response, typically to
-                guide the assistant's output format and "persona". For example, "Let me solve this step by step."
-            chat_mode (bool): If True, agent operates in a chat (conversational) context.
             negative_reward (float): Reward assigned for a clearly incorrect or unparseable answer.
             partial_end_reward (float): Reward when the answer is correct but an expected end token is not matched exactly.
             **kwargs: Additional arguments for the base RewardOnlyAgent.
@@ -49,8 +43,6 @@ def __init__(self,
 
         self.format_reward = format_reward
         self.answer_format = answer_format
-        self.assistant_suffix = assistant_suffix
-        self.chat_mode = chat_mode
         self.negative_reward = negative_reward
         self.partial_end_reward = partial_end_reward
 
@@ -134,12 +126,6 @@ def make_prefix(self, problem_key: str = "problem", **kwargs) -> str:
         else:
             raise ValueError(f"Invalid answer format: {self.answer_format}")
 
-        if self.chat_mode:
-            prefix = f"""{kwargs[problem_key]}\n{answer_format}"""
-        else:
-            prefix = f"""A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer.
-    The question will be a word math problem. Show your work in <think> </think> tags. 
-    {answer_format}
-    User: {kwargs[problem_key]}
-    {self.assistant_suffix}"""
+        prefix = f"""{kwargs[problem_key]}\n{answer_format}"""
+
         return prefix

examples/rl/model_configs/common.sh

@@ -22,7 +22,7 @@ COMMON_OPTIONS="\
     --attention-backend flash \
     --timing-log-level 1 \
     --log-timers-to-tensorboard \
-    --save-retain-interval 120 \
+    --save-retain-interval 160 \
     --inference-dynamic-batching-num-cuda-graphs 1 \
     --inference-dynamic-batching-unified-memory-level 1 \
     --adam-beta1 0.9 \

examples/rl/model_configs/llama3p1_8b_instruct.sh

@@ -101,9 +101,6 @@ MODEL_OPTIONS="\
   --max-position-embeddings 131072  \
   --tokenizer-type HuggingFaceTokenizer \
   --tokenizer-model unsloth/Meta-Llama-3.1-8B-Instruct \
-  --legacy-tokenizer \
-  --langrl-inference-server-type "inplace_megatron_chat" \
-  --langrl-inference-server-conversation-template "unsloth/Meta-Llama-3.1-8B-Instruct" \
   --lr 3e-7 \
   --make-vocab-size-divisible-by 128 \
   --clip-grad 1.0 \

examples/rl/model_configs/qwen3_4b.sh

@@ -63,8 +63,6 @@ MODEL_OPTIONS="\
   --attention-softmax-in-fp32 \
   --tokenizer-type HuggingFaceTokenizer \
   --tokenizer-model Qwen/Qwen3-4B \
-  --langrl-inference-server-type "inplace_megatron_chat" \
-  --langrl-inference-server-conversation-template "Qwen/Qwen3-4B" \
   --vocab-size 151936 \
   --make-vocab-size-divisible-by 128 \
   --optimizer adam \

examples/rl/model_configs/qwen3_8b.sh

@@ -64,8 +64,6 @@ MODEL_OPTIONS="\
   --attention-softmax-in-fp32 \
   --tokenizer-type HuggingFaceTokenizer \
   --tokenizer-model Qwen/Qwen3-8B \
-  --langrl-inference-server-type "inplace_megatron_chat" \
-  --langrl-inference-server-conversation-template "Qwen/Qwen3-8B" \
   --vocab-size 151936 \
   --make-vocab-size-divisible-by 128 \
   --optimizer adam \

examples/rl/model_configs/qwen_2p5_3b.sh

@@ -22,7 +22,7 @@ if [ "$(basename "$ENV_CONFIG")" = "dapo.yaml" ]; then
   GRPO_KL_BETA=${GRPO_KL_BETA:-"0.0"}
   ENTROPY_WEIGHT=${ENTROPY_WEIGHT:-"0.0"}
   TRAINING_BATCH_SIZE=${TRAINING_BATCH_SIZE:-1024}
-  MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-2}
+  MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
   MAX_SEQ_LENGTH=${MAX_SEQ_LENGTH:-8192}
   EXIT_INTERVAL=${EXIT_INTERVAL:-16}
   CHKPT_SAVE_INTERVAL=${CHKPT_SAVE_INTERVAL:-16}
@@ -38,7 +38,7 @@ else
   GRPO_KL_BETA=${GRPO_KL_BETA:-"0.0"}
   ENTROPY_WEIGHT=${ENTROPY_WEIGHT:-"0.0"}
   TRAINING_BATCH_SIZE=${TRAINING_BATCH_SIZE:-512}
-  MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-2}
+  MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
   MAX_SEQ_LENGTH=${MAX_SEQ_LENGTH:-8192}
   EXIT_INTERVAL=${EXIT_INTERVAL:-16}
   CHKPT_SAVE_INTERVAL=${CHKPT_SAVE_INTERVAL:-16}

examples/rl/model_configs/qwen_2p5_distill_7b.sh

@@ -70,8 +70,6 @@ MODEL_OPTIONS="\
   --max-position-embeddings 131072  \
   --tokenizer-type HuggingFaceTokenizer \
   --tokenizer-model "unsloth/DeepSeek-R1-Distill-Qwen-7B" \
-  --langrl-inference-server-type "inplace_megatron_chat" \
-  --langrl-inference-server-conversation-template "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" \
   --lr 0.000001 \
   --lr-warmup-samples 0 \
   --make-vocab-size-divisible-by 128 \

megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py

@@ -3,8 +3,10 @@
 import asyncio
 import logging
 import time
+import traceback
 
 from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.tokenizers.text.parsers import PARSER_MAPPING
 
 logger = logging.getLogger(__name__)
 
@@ -19,6 +21,7 @@ async def chat_completions():
         """Handles async POST requests for chat completions."""
         client = current_app.config['client']
         tokenizer = current_app.config['tokenizer']
+        parsers = current_app.config['parsers']
 
         req = request.get_json()
 
@@ -31,14 +34,17 @@ async def chat_completions():
 
         try:
             prompt_tokens = tokenizer.apply_chat_template(
-                messages, tokenize=True, add_generation_prompt=True
+                messages, tokenize=True, add_generation_prompt=True, tools=req.get("tools", None)
+            )
+        except (AttributeError, AssertionError):
+            logger.warning(
+                "Tokenizer does not support 'apply_chat_template'. Using tokenize instead."
+            )
+            prompt_tokens = tokenizer.tokenize(
+                "\n".join([message["content"] for message in messages])
             )
-        except AttributeError:
-            return (
-                "Tokenizer does not support 'apply_chat_template'. "
-                "Chat completions requires a tokenizer with a configured chat template."
-            ), 500
         except Exception as e:
+            logger.error(f"{traceback.format_exc()}")
             return f"Error processing 'messages': {e}", 500
 
         # --- 2. Parse Sampling Params ---
@@ -55,14 +61,31 @@ async def chat_completions():
             # Check for 'logprobs' (bool) and 'top_logprobs' (int)
             return_log_probs = bool(req.get("logprobs", False))
             top_n_logprobs = int(req.get("top_logprobs", 0)) if return_log_probs else 0
+            skip_prompt_log_probs = bool(req.get("skip_prompt_log_probs", False))
+            add_BOS = bool(req.get("add_BOS", False))
+
+            # The engine only handles add_BOS for string prompts, not pre-tokenized
+            # input. Since we pre-tokenize via apply_chat_template, we must handle
+            # BOS ourselves, matching the logic in tokenize_prompt().
+            if hasattr(tokenizer, 'bos') and tokenizer.bos is not None:
+                while prompt_tokens and prompt_tokens[0] == tokenizer.bos:
+                    prompt_tokens.pop(0)
+                if add_BOS:
+                    prompt_tokens = [tokenizer.bos] + prompt_tokens
 
             sampling_params = SamplingParams(
                 temperature=temperature,
                 top_k=top_k,
                 top_p=top_p,
                 return_log_probs=return_log_probs,
                 top_n_logprobs=top_n_logprobs,
-                num_tokens_to_generate=int(req.get("max_tokens", 16)),
+                num_tokens_to_generate=(
+                    int(max_tokens)
+                    if ((max_tokens := req.get("max_tokens", None)) is not None)
+                    else None
+                ),
+                skip_prompt_log_probs=skip_prompt_log_probs,
+                add_BOS=add_BOS,
             )
         except ValueError as e:
             return f"Invalid sampling parameter: {e}", 400
@@ -71,20 +94,13 @@ async def chat_completions():
         # For chat, we run the *same* prompt 'n' times.
         tasks = []
         for _ in range(n):
-            per_req_params = SamplingParams(
-                temperature=sampling_params.temperature,
-                top_k=sampling_params.top_k,
-                top_p=sampling_params.top_p,
-                return_log_probs=sampling_params.return_log_probs,
-                top_n_logprobs=sampling_params.top_n_logprobs,
-                num_tokens_to_generate=sampling_params.num_tokens_to_generate,
-            )
-            tasks.append(client.add_request(prompt_tokens, per_req_params))
+            tasks.append(client.add_request(prompt_tokens, sampling_params))
 
         start_time = time.perf_counter()
         try:
             batch_results = await asyncio.gather(*tasks)
         except Exception as e:
+            logger.error(f"Error during inference: {e}")
             return f"Error during inference: {e}", 500
 
         logger.info(
@@ -95,21 +111,29 @@ async def chat_completions():
         # --- 4. Format OpenAI Response ---
         choices = []
         total_completion_tokens = 0
-        prompt_token_count = len(prompt_tokens)  # Calculated once
+        prompt_tokens_counts = []
 
         request_idx = 0
         for record in batch_results:
             assert len(record.requests) == 1, "Each record should contain one request result."
-            result = record.merge()
-            text_output = result.generated_text
+            result = record.merge().serialize()
+            # Unwrap ("tensor", [...]) tuples from serialize() into plain lists.
+            result = {
+                k: v[1] if isinstance(v, (list, tuple)) and len(v) == 2 and v[0] == "tensor" else v
+                for k, v in result.items()
+            }
+            prompt_tokens = result["prompt_tokens"]  # The engine can modify prompt_tokens.
+            text_output = result["generated_text"]
+            prompt_tokens_count = len(prompt_tokens) if prompt_tokens is not None else 0
+            prompt_tokens_counts.append(prompt_tokens_count)
 
             logprobs_content = None
             if sampling_params.return_log_probs:
-                token_logprobs = getattr(result, 'log_probs', [])
-                tokens = [tokenizer.detokenize([tok]) for tok in result.generated_tokens]
+                token_logprobs = result.get('log_probs', [])
+                tokens = [tokenizer.detokenize([tok]) for tok in result["generated_tokens"]]
 
                 # Get top_n_logprobs if available
-                generated_top_n_logprobs = getattr(result, 'generated_top_n_logprobs', None)
+                generated_top_n_logprobs = result.get('generated_top_n_logprobs')
 
                 logprobs_content = []
                 for i, (tok, lp) in enumerate(zip(tokens, token_logprobs)):
@@ -134,25 +158,50 @@ async def chat_completions():
                     }
                     logprobs_content.append(entry)
 
+            metadata = {}
+            message_text = text_output
+            if parsers:
+                for parser in parsers:
+                    if parser not in PARSER_MAPPING:
+                        raise ValueError(f"Parser {parser} not found in PARSER_MAPPING")
+                    message_text, new_info = PARSER_MAPPING[parser].parse(
+                        message_text, tools=req.get("tools", None)
+                    )
+                    assert not (
+                        metadata.keys() & new_info.keys()
+                    ), "Multiple parsers found the same information."
+                    metadata.update(new_info)
+            message = {"role": "assistant", "content": message_text}
+            if "tool_calls" in metadata:
+                message["tool_calls"] = metadata["tool_calls"]
+            if "reasoning" in metadata:
+                message["reasoning"] = metadata["reasoning"]
+
             choice_data = {
-                "index": 0,
-                "message": {"role": "assistant", "content": text_output},
+                "index": request_idx,
+                "message": message,
+                "prompt_token_ids": result["prompt_tokens"],
+                "generation_token_ids": result["generated_tokens"],
+                "generation_log_probs": result["generated_log_probs"],
+                "raw_text": result["prompt"] + result["generated_text"],
                 # 'logprobs' in chat API is an object containing 'content'
                 "logprobs": {"content": logprobs_content} if logprobs_content else None,
-                "finish_reason": "length",  # Original code hardcoded this.
+                "finish_reason": (
+                    "tool_calls" if metadata.get("tool_calls", []) else "stop"
+                ),  # Original code hardcoded this.
             }
             logging.info(result)
-            if result.routing_indices is not None:
-                choice_data["moe_topk_indices"] = result.routing_indices.tolist()
-                prompt_length = len(result.prompt_tokens) if result.prompt_tokens is not None else 0
-                if prompt_length:
-                    choices[-1]["prompt_moe_topk_indices"] = result.routing_indices[
-                        :prompt_length
-                    ].tolist()
+            if result["routing_indices"] is not None:
+                choice_data["moe_topk_indices"] = result["routing_indices"]
+                if prompt_tokens_count:
+                    choices[-1]["prompt_moe_topk_indices"] = result["routing_indices"][
+                        :prompt_tokens_count
+                    ]
             choices.append(choice_data)
-            total_completion_tokens += len(result.generated_tokens)
-            request_idx += 0
+            total_completion_tokens += len(result["generated_tokens"])
+            request_idx += 1
 
+        prompt_token_count = max(prompt_tokens_counts)
         response = {
             "choices": choices,
             "usage": {