add cache to tokenizer

Author: binary-husky

.gitignore

@@ -190,3 +190,4 @@ cocktail_vs_separate*
 cocktail_results_*
 PRINCIPLE.md
 TODO.md
+plot_mean_reward.py

ajet/backbone/main_verl.py

@@ -132,9 +132,16 @@ def run(self, config):
 
         # Instantiate the tokenizer and processor.
         from verl.utils import hf_processor, hf_tokenizer
+        from ajet.tokenizer.service import start_tokenizer_service
 
         trust_remote_code = config.data.get("trust_remote_code", False)
-        tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)
+        local_tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)
+        # Cache hot tokenization calls (encode / decode / apply_chat_template)
+        # in a sidecar process; every other tokenizer attribute is served by
+        # the local instance directly.
+        tokenizer = start_tokenizer_service(
+            local_tokenizer, local_path, trust_remote_code=trust_remote_code
+        )
         # Used for multimodal LLM, could be None
         processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)
 

ajet/backbone/trainer_verl.py

@@ -511,15 +511,16 @@ def fit(self):  # noqa: C901
                                 ]
                             )
                         )
-                        logger.info("start fit rollout")
+                        logger.info("start batch rollout")
                         self.parallel_env.current_global_steps = self.global_steps
+                        # rollout stage begin ✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨
                         context_tracker_arr: List[SingleAgentContextTracker] = self.parallel_env.rollout(
                             tasks, mode="sample", epoch=f"train.{epoch}"
                         )
 
                         # from ajet import bp; bp("BATCH")
 
-                        logger.info("end fit rollout")
+                        logger.info("end batch rollout")
                         gen_batch_output = self.parallel_env.to_dataproto(context_tracker_arr)
                         logger.info("end dataproto convertion")
 
@@ -710,7 +711,7 @@ def fit(self):  # noqa: C901
 
                     # implement critic warmup
                     if self.config.trainer.critic_warmup <= self.global_steps:
-                        # update actor
+                        # update actor ✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨✨
                         with marked_timer("update_actor", timing_raw, color="red"):
                             actor_output = self._update_actor(batch)
                         actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])

ajet/context_tracker/multiagent_tracking.py

@@ -130,6 +130,8 @@ def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_to
         if disable_toolcalls:
             consider_roles.remove("tool")
 
+        previous_message_encounter_user_role = False
+
         for i, msg in enumerate(messages):
 
             if (disable_toolcalls) and (not isinstance(msg["content"], str)):
@@ -166,6 +168,11 @@ def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_to
             else:
                 author = "env"
 
+            if msg["role"] == "user":
+                previous_message_encounter_user_role = True
+
+            any_later_msg_has_user_role = any((m["role"] == "user") for m in messages[i+1:])
+
             # extract content block from openai-competible messages and convert to ExtendedMessage
             timeline += [
                 ExtendedMessage(
@@ -179,8 +186,11 @@ def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_to
                     token_generator="auto",
                     name = (msg["name"] if "name" in msg else ""),
                     first_message=(i == 0),
+                    before_last_query=any_later_msg_has_user_role
                 )
             ]
+            if ("<think>" in msg["content"]) and (not previous_message_encounter_user_role):
+                logger.warning(f"Warning! Message content contains <think> tag, but no prior message has `user` role! This is not a common scenario. Please check your agent loop carefully.")
 
         return timeline
 
@@ -393,9 +403,9 @@ def patch_prompt_tokens(
     ) -> List[ExtendedMessage]:
         """
         fix retokenization drift
-        prompt_text: [this llm call] the prompt in text format used in generation
-        prompt_token_ids: [this llm call] the prompt token ids used in generation (prompt_text->prompt_token_ids using tokenizer)
-        previous_ext_context: [from previous context] the context history
+        prompt_text = llm_output["prompt_text"]:            [this llm call] the prompt in text format used in generation
+        prompt_token_ids = llm_output["prompt_token_ids"]:  [this llm call] the prompt token ids used in generation (prompt_text->prompt_token_ids using tokenizer)
+        previous_ext_context:                               [from previous context] the context history
         """
 
         # remove tailing, usually `<|im_start|> assistant`
@@ -467,21 +477,22 @@ def ensure_retokenization_perfect_match(self, previous_ext_context, split_prompt
                 # good, everything is perfect
                 continue
             else:
+                from ajet import bp; bp("SWARM")
                 # otherwise, we throw a warning (do not worry, this causes almost no influence in the training)
                 print_dict(
                     {
-                        "expected_prompt_text": prompt_text_split[j],
-                        "current_prompt_text": current_prompt_text[j],
-                        "expected_token_ids": split_prompt_token_ids[j],
-                        "current_token_ids": previous_ext_context[j].token_arr,
+                        "expected_prompt_text": prompt_text_split[j],       # from llm_output["prompt_text"]
+                        "current_prompt_text": current_prompt_text[j],      # history prompt text converted from token_arr to text using tokenizer
+                        "expected_token_ids": vllm_token_array,         # from llm_output["prompt_token_ids"]
+                        "current_token_ids": tracker_token_array,       # from previous_ext_context[j].token_arr
                     },
                     mod="exception",
                     header="Prompt token ids mismatch.",
                 )
-                # fix drift
-                previous_ext_context[j].token_arr = self.tokenizer(
-                    prompt_text_split[j], return_tensors="pt", padding=False
-                )["input_ids"]
+                # # fix drift
+                # previous_ext_context[j].token_arr = self.tokenizer(
+                #     prompt_text_split[j], return_tensors="pt", padding=False
+                # )["input_ids"].tolist()
 
     def process_reward(self, reward_structure: Reward):
         self.reward_structure = reward_structure

ajet/context_tracker/single_agent_tracking.py

@@ -91,7 +91,7 @@ def get_token_inc_from_llm_response(
         self.generated_token_cnt += len(vllm_output_raw_token)
         if not self.generation_prompt_token:
             self.generation_prompt_token = self.get_generation_prompt_token()
-        final_token_arr, token_logprob_arr, loss_mask, lack_normal_eos = replace_token_ids(
+        final_token_arr, token_logprob_arr, loss_mask, lack_normal_eos = replace_token_ids( # pad tokens and logprobs with begin_ids / other_ids / NA
             token_container=completion_token_arr,
             precise_token=vllm_output_raw_token,
             precise_logprob=vllm_output_raw_logprob,

ajet/copilot/create-keep-think-model-chat-template/SKILL.md

@@ -0,0 +1,16 @@
+
+
+Your task is to investigate the chat template of given model, go to its tokenizer config and check whether the following behavior exists:
+
+>
+> Remove history <think> block from the input when apply chat template when converting messages.
+>
+
+This behavior will make RL training slower, if this behavior exists, please change the chat template to forbid such behavior.
+
+You must not do this in-place, instead, please create another model.
+E.g., "/mnt/data_cpfs/xielipeng.xlp/models/Qwen3-8B" -> "/mnt/data_cpfs/xielipeng.xlp/models/Qwen3-8B-Keep-History"
+For all files within the original model path, please create symbolic links instead of copying files.
+With only one exception, the tokenizer config file, which should be copied and modified to change the chat template.
+
+

ajet/schema/extended_msg.py

@@ -20,7 +20,7 @@
     "memory",
     "llm(do_not_train)",
 ]
-DUMMY_MSG = [{"role": "user", "content": "dummy text"}]
+
 
 
 def find_sublist_indices(large_list, small_list, reverse=False):
@@ -77,6 +77,7 @@ def __init__(
         token_logprob_arr=[],
         name="",    # preserved field, not used currently
         first_message=False,
+        before_last_query=True,   # whether this message is before the last user query in the conversation, used for auto tokenization logic
     ):
         self.author = author
         self.role = role
@@ -107,6 +108,7 @@ def __init__(
         self.eos_token_id = tokenizer.eos_token_id
 
         if token_generator == "auto":
+            self.before_last_query = before_last_query
             self.token_arr = self.auto_tokenize(
                 tokenizer=tokenizer,
                 tools=tools,
@@ -137,6 +139,12 @@ def auto_tokenize(self, tokenizer, tools):
         return self.token_arr
 
     def auto_tokenize_non_first_message(self, tokenizer, tools):
+        if self.before_last_query:
+            # for example, this will remove the <thinking> block for qwen3's chat template
+            dummy_msg = [{"role": "assistant", "content": "dummy text"}]
+        else:
+            dummy_msg = [{"role": "user", "content": "dummy text"}]
+
         try:
             # completion_token_arr will contain generation_prompt header
             auto_tokenize_target:dict = {
@@ -149,7 +157,7 @@ def auto_tokenize_non_first_message(self, tokenizer, tools):
                 auto_tokenize_target.update({"tool_call_id": self.tool_call_id})
             text_frag_to = ajet_apply_chat_template(
                 tokenizer=tokenizer,
-                conversation=DUMMY_MSG + [auto_tokenize_target],
+                conversation=dummy_msg + [auto_tokenize_target],
                 tokenize=False,
                 tools=tools,
             )
@@ -160,7 +168,7 @@ def auto_tokenize_non_first_message(self, tokenizer, tools):
         self.token_arr, _ = self.get_inc_simple(
             text_frag_from=ajet_apply_chat_template(
                 tokenizer=tokenizer,
-                conversation=DUMMY_MSG,
+                conversation=dummy_msg,
                 tokenize=False,
                 tools=tools,
             ),
@@ -314,21 +322,23 @@ def merge_tool_group(group, tokenizer):
                 token_logprob_arr=msg0.token_logprob_arr,
                 first_message=msg0.first_message,
             )
+            # a dummy msg, not necessary, can be []
+            dummy_msg = [{"role": "user", "content": "dummy text"}]
             # re-compute token_arr
             auto_tokenize_targets = [
                 {"role": msg.role, "content": msg.text_content_for_compare} for msg in group
             ]
             merged.token_arr, _ = merged.get_inc_simple(
                 text_frag_from=ajet_apply_chat_template(
                     tokenizer=tokenizer,
-                    conversation=DUMMY_MSG,
+                    conversation=dummy_msg,
                     tokenize=False,
                     tools=merged.tools,
                     add_generation_prompt=False,
                 ),
                 text_frag_to=ajet_apply_chat_template(
                     tokenizer,
-                    conversation=DUMMY_MSG + auto_tokenize_targets,
+                    conversation=dummy_msg + auto_tokenize_targets,
                     tokenize=False,
                     tools=merged.tools,
                     add_generation_prompt=False,

ajet/tokenizer/__init__.py

@@ -0,0 +1,7 @@
+# Intentionally empty. Import directly from ``ajet.tokenizer.service``:
+#
+#     from ajet.tokenizer.service import RemoteTokenizer, start_tokenizer_service
+#
+# Re-exporting here would cause ``python -m ajet.tokenizer.service`` to load
+# the package eagerly before running the service as __main__, which trips a
+# RuntimeWarning from runpy.