From dc2e969b92c9bf5374ee873737261832ebc48b34 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Thu, 23 Oct 2025 12:37:13 +0800 Subject: [PATCH 1/2] more --- miles/backends/fsdp_utils/actor.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index eb52fa771fe..ff2d543fff6 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -215,7 +215,7 @@ def compute_log_prob( model_args["pixel_values"] = batch["pixel_values"] logits = self.model(**model_args).logits batch[f"{store_prefix}log_probs"] = gather_log_probs_packed( - logits, batch["tokens"], self.args.rollout_temperature + logits, batch["tokens"], temperature=self.args.rollout_temperature ) return rollout_data @@ -654,7 +654,10 @@ def gather_log_probs(logits: torch.Tensor, input_ids: torch.Tensor, rollout_temp def gather_log_probs_packed( - logits: torch.Tensor, input_ids: torch.Tensor, cu_seqlens: torch.Tensor | float | None = None + logits: torch.Tensor, + input_ids: torch.Tensor, + cu_seqlens: torch.Tensor | float | None = None, + temperature: torch.Tensor | None = None, ) -> torch.Tensor: """Gather next-token log probabilities for packed sequences. @@ -673,6 +676,9 @@ def gather_log_probs_packed( logits = logits.squeeze(0) input_ids = input_ids.squeeze(0) + if temperature is not None: + logits = logits.div(temperature) + # Shift for next-token prediction: logits[:-1] predicts input_ids[1:] log_probs = torch.log_softmax(logits[:-1], dim=-1) targets = input_ids[1:].to(device=log_probs.device) From 6f490cf116384f0891c2168cdf588d8a2b8df8ba Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Thu, 23 Oct 2025 21:58:46 +0800 Subject: [PATCH 2/2] add temperature for cur_log_probs --- miles/backends/fsdp_utils/actor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index ff2d543fff6..92e6151b461 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -371,7 +371,9 @@ def train(self, rollout_id: int, rollout_data_ref: Box) -> None: ).logits # Handle packed sequences - log_probs = gather_log_probs_packed(logits, packed_batch["tokens"], packed_batch["cu_seqlens"]) + log_probs = gather_log_probs_packed( + logits, packed_batch["tokens"], packed_batch["cu_seqlens"], temperature=self.args.rollout_temperature + ) packed_batch["cur_log_probs"] = log_probs unpacked_batches = unpack_sequences(packed_batch)