From 2b46ab1401db3977529aa2f5cf2a39a922fc3d3b Mon Sep 17 00:00:00 2001 From: YeAnbang Date: Thu, 18 Sep 2025 18:28:36 +0800 Subject: [PATCH] simplify _run_agentic_pipeline; fix old_log_probs --- .../coati/distributed/agent/agentic_producer.py | 12 ++++-------- .../ColossalChat/coati/distributed/agent/base.py | 16 +++++++--------- .../ColossalChat/coati/distributed/loss.py | 4 ++-- applications/ColossalChat/rl_example.py | 10 ++++------ 4 files changed, 17 insertions(+), 25 deletions(-) diff --git a/applications/ColossalChat/coati/distributed/agent/agentic_producer.py b/applications/ColossalChat/coati/distributed/agent/agentic_producer.py index 48b07ffad..4f7dc3c9f 100644 --- a/applications/ColossalChat/coati/distributed/agent/agentic_producer.py +++ b/applications/ColossalChat/coati/distributed/agent/agentic_producer.py @@ -224,9 +224,7 @@ class AgenticProducer(BaseAgenticProducer): if llm_call_count > self.llm_call_budget: print(f"LLM call budget exceeded: {llm_call_count} > {self.llm_call_budget}. Stopping.") del self.async_llm_engine_map[request_id] - while messages[-1]["role"] == "tool": - messages.pop() - return messages, logprobs + return messages, response_input_ids, logprobs inputs = self._build_prompt(messages, return_dict=True, return_tensors="pt") if num_prompt_tokens == 0: num_prompt_tokens = inputs["input_ids"].size(-1) @@ -235,9 +233,7 @@ class AgenticProducer(BaseAgenticProducer): f"Max tokens exceeded: Current have generated {inputs['input_ids'].size(-1) - num_prompt_tokens} tokens > {self.generate_config.get('max_tokens', 512)}. Stopping." ) del self.async_llm_engine_map[request_id] - while messages[-1]["role"] == "tool": - messages.pop() - return messages, logprobs + return messages, response_input_ids, logprobs async_producer = self._select_async_producer(request_id=request_id) agentic_generate_config = copy.deepcopy(self.generate_config) agentic_generate_config["max_tokens"] = self.agentic_config.get("max_tokens", 2048) @@ -262,7 +258,7 @@ class AgenticProducer(BaseAgenticProducer): if tool_call_count > self.tool_call_budget: print(f"Tool call budget exceeded: {tool_call_count} > {self.tool_call_budget}. Stopping.") del self.async_llm_engine_map[request_id] - return messages, logprobs + return messages, response_input_ids, logprobs tool_call_count += len(assistant_message["tool_calls"]) handlers = [] for tool_call in assistant_message["tool_calls"]: @@ -277,4 +273,4 @@ class AgenticProducer(BaseAgenticProducer): else: # no further tool call, return the messages del self.async_llm_engine_map[request_id] - return messages, logprobs + return messages, response_input_ids, logprobs diff --git a/applications/ColossalChat/coati/distributed/agent/base.py b/applications/ColossalChat/coati/distributed/agent/base.py index 6ea5b17e1..e5ff9ffc5 100644 --- a/applications/ColossalChat/coati/distributed/agent/base.py +++ b/applications/ColossalChat/coati/distributed/agent/base.py @@ -123,24 +123,22 @@ class BaseAgenticProducer(BaseProducer): ) for i in range(self.num_generations): - _messages, logprobs = results[i] - response_input_ids = self._build_prompt( - _messages, return_dict=True, return_tensors="pt", add_generation_prompt=False - )["input_ids"] + # due to the multiround feature, action_mask and attention_mask need to be recomputed + _messages, response_input_ids, logprobs = results[i] # truncate if too long - response_input_ids = response_input_ids[:, : self.grpo_config["max_length"] - to_pad_left] + response_input_ids = response_input_ids[0, :, : self.grpo_config["max_length"] - to_pad_left] # add left right padding - to_pad_right = self.grpo_config["max_length"] - response_input_ids.shape[1] - to_pad_left - response_length = response_input_ids.shape[1] - prompt_length + to_pad_right = self.grpo_config["max_length"] - response_input_ids.size(-1) - to_pad_left input_ids = torch.nn.functional.pad( response_input_ids, (to_pad_left, to_pad_right), "constant", value=self.tokenizer.pad_token_id ) # [1, max_length] attention_mask = input_ids.ne(self.tokenizer.pad_token_id).int() # [1, max_length] action_mask = input_ids[:, max_prompt_length:].ne(self.tokenizer.pad_token_id).int() + response_length = action_mask.sum().item() rollouts["attention_mask"].append(attention_mask) rollouts["action_mask"].append(action_mask) truncated_logprobs = logprobs[ - :, :, prompt_length : prompt_length + self.generate_config["max_tokens"] + 0, :, prompt_length : prompt_length + self.generate_config["max_tokens"] ] # truncate to max_new_tokens logprobs_padded = torch.nn.functional.pad( truncated_logprobs, @@ -148,7 +146,7 @@ class BaseAgenticProducer(BaseProducer): "constant", value=0.0, ) # [1, max_new_tokens] - rollouts["action_log_probs"].append(logprobs_padded[0]) + rollouts["action_log_probs"].append(logprobs_padded) rollouts["response_idx"].append( torch.tensor( [ diff --git a/applications/ColossalChat/coati/distributed/loss.py b/applications/ColossalChat/coati/distributed/loss.py index ab38f987f..7fcfdba31 100644 --- a/applications/ColossalChat/coati/distributed/loss.py +++ b/applications/ColossalChat/coati/distributed/loss.py @@ -37,9 +37,9 @@ class PolicyLoss(nn.Module): total_effective_tokens_in_batch: torch.Tensor = None, ) -> torch.Tensor: if action_mask is None: - ratio = (log_probs - log_probs.detach()).exp() + ratio = (log_probs - old_log_probs.detach()).exp() else: - ratio = ((log_probs - log_probs.detach()) * action_mask).exp() + ratio = ((log_probs - old_log_probs.detach()) * action_mask).exp() surr1 = ratio * advantages surr2 = ratio.clamp(1 - self.clip_eps_low, 1 + self.clip_eps_high) * advantages diff --git a/applications/ColossalChat/rl_example.py b/applications/ColossalChat/rl_example.py index 5af6fb792..5770bfde6 100644 --- a/applications/ColossalChat/rl_example.py +++ b/applications/ColossalChat/rl_example.py @@ -429,18 +429,16 @@ if __name__ == "__main__": "max_tokens": 2048, } grpo_config["forced_patterns"] = [ - r"\n.+\n" + r"\n.+\n" # please modify based on your tool response format ] # force at least one correct tool call else: raise ValueError(f"Unsupported agentic model type: {args.agentic_type}") else: agentic_config = None - tokenizer_config = { - "path": args.model, - "trust_remote_code": True, - "chat_template": args.chat_template, - } + tokenizer_config = {"path": args.model, "trust_remote_code": True} + if args.chat_template is not None: + tokenizer_config["chat_template"] = args.chat_template launch_distributed( num_producers=args.num_inferencer,