s123hree/green-code-optimizer-a100
0
1"""22026.4.932026.4.844.57.250.23.06__UNSLOTH_VERSIONING__7"""8 9# Unsloth auto generated code10# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.11#12# This program is free software: you can redistribute it and/or modify13# it under the terms of the GNU Lesser General Public License as published by14# the Free Software Foundation, either version 3 of the License, or15# (at your option) any later version.16#17# This program is distributed in the hope that it will be useful,18# but WITHOUT ANY WARRANTY; without even the implied warranty of19# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the20# GNU General Public License for more details.21#22# You should have received a copy of the GNU Lesser General Public License23# along with this program. If not, see <https://www.gnu.org/licenses/>.24 25from torch import Tensor26import torch27import torch.nn as nn28from torch.nn import functional as F29from unsloth_zoo.temporary_patches.common import torch_compile30from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable31from trl.trainer.grpo_trainer import (Any, AutoConfig, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, DataLoader, Dataset, FSDP, GRPOConfig, GRPOTrainer, GenerationConfig, IterableDataset, Optional, Path, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RepeatSampler, RewardFunc, Sampler, SyncRefModelCallback, Trainer, TrainerCallback, Union, VLLMClient, _ForwardRedirection, apply_chat_template, broadcast_object_list, copy, datasets, defaultdict, deque, disable_dropout_in_model, gather, gather_object, generate_model_card, get_comet_experiment_url, identity, inspect, is_conversational, is_datasets_available, is_flash_attn_2_available, is_liger_kernel_available, is_peft_model, is_rich_available, is_vllm_available, is_wandb_available, logger, logging, maybe_apply_chat_template, nanmax, nanmin, nanstd, nn, nullcontext, os, pad, partial, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, print_prompt_completions_sample, profiling_context, profiling_decorator, re, seed_worker, selective_log_softmax, set_seed, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, textwrap, torch, transformers, truncate_with_protected_tokens, unsplit_pixel_values_by_grid, unwrap_model_for_generation, AutoConfig, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, Dataset, GRPOConfig, GRPOTrainer, GenerationConfig, IterableDataset, Optional, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RewardFunc, SyncRefModelCallback, Trainer, TrainerCallback, Union, VLLMClient, datasets, defaultdict, deque, disable_dropout_in_model, identity, inspect, is_liger_kernel_available, is_peft_model, is_vllm_available, logger, nn, os, pad, prepare_deepspeed, prepare_fsdp, re, set_seed, torch, transformers, Any, FSDP, Union, apply_chat_template, broadcast_object_list, copy, gather, gather_object, is_conversational, is_flash_attn_2_available, logging, maybe_apply_chat_template, nanmax, nanmin, nanstd, nullcontext, os, pad, prepare_multimodal_messages, profiling_context, re, torch, transformers, truncate_with_protected_tokens, unwrap_model_for_generation, os, pad, re, selective_log_softmax, torch, transformers, re, Any, Union, profiling_decorator, re, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, torch, unsplit_pixel_values_by_grid, Optional, PreTrainedModel, Trainer, logger, os, re, torch, FSDP, nn, os, re, FSDP, nn, re, torch, GRPOTrainer, Trainer, gather, nanmax, nanmin, os, pad, re, torch)32 33 34import os35import math36import logging37from typing import *38from dataclasses import dataclass, field39from packaging.version import Version40import torch41import numpy as np42from contextlib import nullcontext43from torch.nn import functional as F44import inspect45from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling46from transformers.training_args import ParallelMode47from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize48 49# Wrap trainer with padding to right and enable training mode50import functools51from types import MethodType52try:53 from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers54except:55 def reset_unsloth_gradient_checkpointing_buffers(): pass56def prepare_for_training_mode(f):57 @functools.wraps(f)58 def wrapper(self, *args, **kwargs):59 # Finish the previous W&B run if this is a subsequent train() call.60 # We do this at the START of train() (not the end) so that61 # evaluate() / log() still work after train() completes.62 # HF's WandbCallback.setup() will call wandb.init() for the new run.63 # See: https://github.com/unslothai/unsloth/issues/395464 if getattr(self, '_unsloth_training_completed', False):65 try:66 import wandb67 if wandb.run is not None:68 wandb.finish()69 # Reset HF's WandbCallback so it calls wandb.init() for the new run70 for cb in self.callback_handler.callbacks:71 if type(cb).__name__ == 'WandbCallback':72 cb._initialized = False73 break74 except:75 pass76 # Enable training mode77 _was_training = None78 # Get gradient checkpointing setting from training arguments79 use_gc = getattr(self.args, 'gradient_checkpointing', True)80 if hasattr(self, 'model') and hasattr(self.model, "training"):81 _was_training = self.model.training82 if hasattr(self, 'model') and hasattr(self.model, "for_training"):83 self.model.for_training(use_gradient_checkpointing=use_gc)84 output = f(self, *args, **kwargs)85 # Restore previous mode when possible86 if hasattr(self, 'model') and hasattr(self.model, "for_inference"):87 if _was_training is False:88 self.model.for_inference()89 elif _was_training is True and hasattr(self.model, "for_training"):90 self.model.for_training(use_gradient_checkpointing=use_gc)91 # Reset gradient checkpointing buffers to free memory while staying ready for next run92 try:93 reset_unsloth_gradient_checkpointing_buffers()94 except:95 pass96 # Mark that training completed so the next train() call can97 # finish this W&B run before starting a new one98 self._unsloth_training_completed = True99 return output100 return wrapper101pass102 103torch_compile_options = {104 "epilogue_fusion" : True,105 "max_autotune" : False,106 "shape_padding" : True,107 "trace.enabled" : False,108 "triton.enable_persistent_tma_matmul": torch.cuda.get_device_capability()[0] >= 9,109 "cuda.compile_opt_level" : "-O2",110 "cuda.enable_cuda_lto" : True,111 }112 113@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,)114def chunked_hidden_states_selective_log_softmax(115 hidden_states: torch.Tensor,116 lm_head: torch.Tensor,117 index: torch.Tensor,118 chunks: int = 4,119 logit_scale_multiply: float = 0.0,120 logit_scale_divide: float = 0.0,121 logit_softcapping: float = 0.0,122 temperature: float = 1.0,123) -> torch.Tensor:124 # All Unsloth Zoo code licensed under AGPL3125 flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])126 flat_index = index.reshape(-1)127 128 chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0)129 chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0)130 131 all_per_token_logps = []132 133 for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index):134 chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t()135 136 if logit_scale_multiply != 0.0:137 chunk_logits = chunk_logits * logit_scale_multiply138 if logit_scale_divide != 0.0:139 chunk_logits = chunk_logits / logit_scale_divide140 if logit_softcapping != 0.0:141 chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping)142 143 chunk_logits = chunk_logits.to(torch.float32)144 145 if temperature != 1.0:146 chunk_logits = chunk_logits / temperature147 148 selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1)149 logsumexp_values = torch.logsumexp(chunk_logits, dim=-1)150 per_token_logps = selected_logits - logsumexp_values151 all_per_token_logps.append(per_token_logps)152 153 all_per_token_logps = torch.concat(all_per_token_logps)154 155 all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1]))156 return all_per_token_logps157 158@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,)159def chunked_selective_log_softmax(logits, index, temperature: float = 1.0):160 # Split into 4 chunks only161 chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = 4, dim = 0)162 chunked_index = torch.chunk(index.reshape(-1), chunks = 4, dim = 0)163 all_per_token_logps = []164 # Below loop does the same as selective_log_softmax(chunk_logits, chunk_index)165 for chunk_logits, chunk_index in zip(chunked_logits, chunked_index):166 chunk_logits = chunk_logits.to(torch.float32)167 if temperature != 1.0:168 chunk_logits = chunk_logits / temperature169 selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1)170 logsumexp_values = torch.logsumexp(chunk_logits, dim = -1)171 per_token_logps = selected_logits - logsumexp_values172 all_per_token_logps.append(per_token_logps)173 pass174 all_per_token_logps = torch.concat(all_per_token_logps)175 all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1]))176 return all_per_token_logps177 178def calculate_pad_tokens_in_prompt(179 input_ids: torch.Tensor,180 logits_to_keep: int,181 pad_token_id: int182) -> torch.Tensor:183 """184 Given prompt tensor, it returns all the left padded tokens in that sequence. so [pad, pad, pad, cat] = 3 tokens185 """186 if logits_to_keep >= input_ids.shape[1]:187 raise ValueError("logits_to_keep must be smaller than the sequence length.")188 189 prompt_section = input_ids[:, :-logits_to_keep]190 191 padding_mask = (prompt_section == pad_token_id)192 193 pad_token_counts = padding_mask.sum(dim=1)194 195 return pad_token_counts196 197def create_completion_attention_mask(198 completion_input_ids: torch.Tensor,199 left_pad_tokens_per_prompt: torch.Tensor,200 max_left_pad: int,201 pad_token_id: int202) -> torch.Tensor:203 """204 Given that we have a sequence, [p,p,p,c,c,c,pad,pad,pad]205 206 Where p are extra prompt tokens we got from slicing the torch tensor, c is completion tokens207 and pad are pad tokens, this function would make a completion mask that would 0 out the pad208 and p tokens. so in this example [0,0,0,1,1,1,0,0,0]209 """210 batch_size, completion_len = completion_input_ids.shape211 device = completion_input_ids.device212 213 num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt214 215 indices = torch.arange(completion_len, device=device).unsqueeze(0)216 shift_mask = indices >= num_tokens_to_mask.unsqueeze(1)217 218 non_padding_mask = (completion_input_ids != pad_token_id)219 220 final_mask = shift_mask & non_padding_mask221 222 return final_mask223 224def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor:225 """226 Moves all padding tokens in each sequence of a batch to the right.227 """228 mask = (tensor != pad_id)229 # Must do stable=True since binary mark is unordered230 sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True)231 packed_tensor = torch.gather(tensor, 1, sorted_indices)232 return packed_tensor233 234def align_logprobs_with_mask(235 logprob_tensor: torch.Tensor,236 attention_mask: torch.Tensor,237 pad_value: float = 0.0238) -> torch.Tensor:239 """240 Aligns a log probability tensor with a given attention mask.241 """242 243 device = logprob_tensor.device244 batch_size, logprob_seq_len = logprob_tensor.shape245 mask_seq_len = attention_mask.shape[1]246 247 padded_logprobs = torch.full(248 attention_mask.shape,249 fill_value=pad_value,250 dtype=logprob_tensor.dtype,251 device=device252 )253 254 left_pad_counts = torch.argmax(attention_mask, dim=1)255 256 cols = torch.arange(logprob_seq_len, device=device)257 dest_indices = left_pad_counts.unsqueeze(1) + cols258 259 # Create destination row indices260 # Shape: [batch_size, logprob_seq_len]261 row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices)262 263 # --- 4. Filter out-of-bounds indices and perform assignment ---264 # Create a mask to identify only the indices that are within the bounds265 # of the target tensor's sequence length.266 valid_mask = dest_indices < mask_seq_len267 268 # Use this mask to select only the valid row indices, column indices,269 # and the corresponding values from the logprob tensor.270 # This flattens the selected elements into 1D tensors.271 valid_rows = row_indices[valid_mask]272 valid_cols = dest_indices[valid_mask]273 valid_vals = logprob_tensor[valid_mask]274 275 # Place the valid values into their correct positions in the padded tensor276 # using a single, efficient advanced indexing operation.277 padded_logprobs[valid_rows, valid_cols] = valid_vals278 279 return padded_logprobs280 281def autotune_batch_and_chunks(282 total_input_rows,283 seq_len,284 hidden_size,285 vocab_size,286 dtype_bytes=16,287 multiplier=None288):289 if multiplier is None:290 final_m = max(4, seq_len // 4096)291 else:292 final_m = multiplier293 294 if torch.cuda.is_available():295 free_bytes, _ = torch.cuda.mem_get_info()296 limit_gb = (free_bytes / (1024**3))*.80297 elif hasattr(torch, "xpu") and torch.xpu.is_available():298 # For XPU: estimate free memory from total - reserved299 total_mem = torch.xpu.get_device_properties(0).total_memory300 reserved_mem = torch.xpu.memory_reserved()301 free_bytes = total_mem - reserved_mem302 limit_gb = (free_bytes / (1024**3)) * 0.80303 else:304 # Fallback: assume 8GB available305 limit_gb = 8.0306 307 bytes_to_gb = 1024**3308 309 b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32)310 311 hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb312 313 base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb314 logits_gb = base_logits / final_m315 316 total_mem_gb = hidden_gb + logits_gb317 318 valid_mask = total_mem_gb <= limit_gb319 valid_indices = torch.nonzero(valid_mask, as_tuple=False)320 321 if valid_indices.shape[0] == 0:322 #This means your GPU will OOM323 return 4, final_m324 325 best_idx = valid_indices[0].item()326 final_b = int(b_vals[best_idx].item())327 328 return final_b, final_m329 330def sanitize_logprob(logprob):331 """Local port of trl.scripts.vllm_serve.sanitize_logprob.332 Filters NaN logprobs from vLLM outputs."""333 value = logprob.logprob334 if math.isnan(value):335 logging.getLogger(__name__).warning(336 f"Generated NaN logprob, token logprob '{logprob}' will be ignored"337 )338 return None339 return value340def _unsloth_get_final_logit_softcapping(config):341 """Return final_logit_softcapping for a model config, falling back to the342 nested text sub-config for composite models. Handles both:343 - Gemma-4-style configs where the attribute lives on ``config.text_config``344 - T5Gemma-style composite configs where the text sub-config is only345 reachable via ``config.get_text_config()``346 Returns 0 if unset, matching the previous behaviour.347 """348 softcap = getattr(config, "final_logit_softcapping", None)349 if softcap is None:350 text_cfg = getattr(config, "text_config", None)351 if text_cfg is None:352 get_text_config = getattr(config, "get_text_config", None)353 if callable(get_text_config):354 try:355 text_cfg = get_text_config()356 except (TypeError, ValueError):357 text_cfg = None358 if text_cfg is not None and text_cfg is not config:359 softcap = getattr(text_cfg, "final_logit_softcapping", None)360 return 0 if softcap is None else softcap361 362def grpo_compute_loss(363 ref,364 new,365 old,366 sampling_per_token_logps,367 input_ids,368 mask,369 beta,370 advantages,371 **kwargs372):373 # All Unsloth Zoo code licensed under AGPL3374 # Set defaults for optional arguments375 loss_type = kwargs.get("loss_type", "grpo")376 epsilon_low = kwargs.get("epsilon_low", 0.2)377 epsilon_high = kwargs.get("epsilon_high", 0.2)378 max_completion_length = kwargs.get("max_completion_length", 8192)379 delta = kwargs.get("delta", None)380 importance_sampling_level = kwargs.get("importance_sampling_level", "token")381 num_items_in_batch = kwargs.get("num_items_in_batch", None)382 current_gradient_accumulation_steps = kwargs.get("current_gradient_accumulation_steps", 1)383 num_processes = kwargs.get("num_processes", 1)384 use_vllm = kwargs.get("use_vllm", False)385 vllm_importance_sampling_cap = kwargs.get("vllm_importance_sampling_cap", 2.0)386 get_sapo_token_loss = kwargs.get("get_sapo_token_loss", None)387 sapo_temperature_pos = kwargs.get("sapo_temperature_pos", 1.0)388 sapo_temperature_neg = kwargs.get("sapo_temperature_neg", 1.05)389 get_off_policy_mask = kwargs.get("get_off_policy_mask", None)390 off_policy_mask_threshold = kwargs.get("off_policy_mask_threshold", None)391 input_ids = input_ids.unsqueeze(-1)392 393 if advantages.dim() == 1:394 advantages = advantages.unsqueeze(1)395 396 if off_policy_mask_threshold is not None:397 off_policy_mask = get_off_policy_mask(398 advantages=advantages,399 per_token_logps=new,400 old_per_token_logps=old,401 mask=mask,402 off_policy_threshold=off_policy_mask_threshold,403 )404 405 with torch.no_grad():406 if use_vllm and sampling_per_token_logps is not None:407 #must filter out extra prompt tokens in begining after making input_ids left padded408 importance_sampling_ratio = torch.exp((old * mask) - sampling_per_token_logps)409 importance_sampling_ratio = torch.clamp(410 importance_sampling_ratio, max=vllm_importance_sampling_cap411 )412 pass413 414 # Must detach - otherwise gradients are not propagated correctly!415 # exp(x - x) == 1416 # loss_i = torch.exp(new - new.detach()) * advantages.unsqueeze(1)417 if old is not None:418 log_ratio = new - old419 else:420 log_ratio = new - new.detach()421 422 if importance_sampling_level == "token":423 log_importance_weights = log_ratio424 elif importance_sampling_level == "sequence":425 log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)426 log_importance_weights = log_importance_weights.unsqueeze(-1)427 else:428 raise ValueError(429 f"Unknown importance sampling level: {importance_sampling_level}. Possible values are 'token' "430 "and 'sequence'."431 )432 433 coef_1 = torch.exp(log_importance_weights)434 435 # Reverse KL436 # Note that this is a low variance low bias estimator for the KL divergence as used in GRPO paper437 if beta != 0.0:438 kl_i = torch.exp(ref - new) - (ref - new) - 1.0439 440 else:441 # set kl_i to a tensor of zeros with the correct shape442 if importance_sampling_level == "sequence":443 kl_i = new.new_zeros(new.size(0), 1)444 else:445 kl_i = torch.zeros_like(new)446 # Full correct reverse KL divergence?? Missing term maybe?447 # kl_i = torch.exp(new) * kl_i448 449 # Below is forward KL (normal KL)450 # kl_i = torch.exp(old) * (old - new)451 if loss_type == "cispo":452 clamped_ratios = torch.clamp(coef_1, max=epsilon_high).detach()453 loss_i = -clamped_ratios * advantages * new454 #breakpoint()455 elif loss_type in ["grpo", "bnpo", "dr_grpo", "dapo"]:456 coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high)457 458 if delta is not None:459 loss_1 = torch.clamp(coef_1, max=delta) * advantages460 else:461 loss_1 = coef_1 * advantages462 pass463 loss_2 = coef_2 * advantages464 loss_i = -torch.min(loss_1, loss_2)465 elif loss_type == "sapo":466 if get_sapo_token_loss is None:467 raise Exception(f"sapo is only available in TRL 0.26.0+")468 loss_i = torch.empty_like(coef_1)469 positive_advantages_mask = advantages.repeat([1, coef_1.shape[1]]) > 0470 #since we have n_chunks some tensors may error if they dont have elements in them471 if coef_1[positive_advantages_mask].numel() != 0:472 loss_i[positive_advantages_mask] = get_sapo_token_loss(473 coef_1[positive_advantages_mask], sapo_temperature_pos474 )475 if coef_1[~positive_advantages_mask].numel() != 0:476 loss_i[~positive_advantages_mask] = get_sapo_token_loss(477 coef_1[~positive_advantages_mask], sapo_temperature_neg478 )479 loss_i = -loss_i * advantages480 else:481 raise ValueError(f"Unknown loss type: {loss_type}")482 483 if off_policy_mask_threshold is not None:484 loss_i = loss_i * off_policy_mask485 486 if use_vllm and sampling_per_token_logps is not None:487 loss_i = loss_i * importance_sampling_ratio488 #delta for metric489 with torch.no_grad():490 delta = torch.abs(old - sampling_per_token_logps)491 delta = delta * mask492 flat_is_ratio = importance_sampling_ratio * mask493 else:494 delta = torch.tensor([]).detach()495 flat_is_ratio = torch.tensor([]).detach()496 if beta != 0.0:497 loss_i = loss_i + beta * kl_i498 499 mask = mask.to(torch.float32)500 n_mask_per_reward = mask.sum(1)501 502 # https://github.com/huggingface/trl/blob/e8b8499f1f8d76838155b515e414ee98f757d6d5/trl/trainer/grpo_trainer.py#L1624503 if loss_type in ["grpo", "sapo"]:504 loss = ((loss_i * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean()505 loss = loss / current_gradient_accumulation_steps506 elif loss_type == "bnpo":507 loss = (loss_i * mask).sum() / mask.sum().clamp(min=1.0)508 loss = loss / current_gradient_accumulation_steps509 elif loss_type == "dr_grpo":510 loss = (loss_i * mask).sum() / (loss_i.size(0) * max_completion_length)511 loss = loss / current_gradient_accumulation_steps512 elif loss_type in ["cispo", "dapo"]:513 normalizer = num_items_in_batch/ num_processes514 loss = (loss_i * mask).sum() / normalizer515 else:516 raise ValueError(f"Unknown loss type: {loss_type}")517 518 # loss = (loss_i * mask).sum() / mask.sum()519 520 # Get metrics as well which are folded521 def masked_batch_mean(x):522 with torch.inference_mode():523 completion_length = n_mask_per_reward.mean()524 if x.shape[1] == 1: # when importance_sampling_level == "sequence"525 return completion_length, x.mean()526 else:527 mean_kl_per_reward = (x * mask).sum(1) / n_mask_per_reward528 mean_kl = mean_kl_per_reward.mean()529 return completion_length, mean_kl530 completion_length, mean_kl = masked_batch_mean(kl_i)531 return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, mask532 533class UnslothEfficientGRPO(torch.autograd.Function):534 # All Unsloth Zoo code licensed under AGPL3535 @staticmethod536 def forward(ctx, _new_logps, _old_logps, _ref_logps, _sampling_per_token_logps, lm_head, _input_ids, _mask, _advantages, beta, scaler = None, n_chunks = 1, extra_kwargs=None):537 if extra_kwargs is None:538 extra_kwargs = {}539 def compute_loss(new_logps, old_logps, ref_logps, sampling_per_token_logps, input_ids, mask, advantages, scaling):540 loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, _mask = grpo_compute_loss(541 ref_logps,542 new_logps,543 old_logps,544 sampling_per_token_logps,545 input_ids,546 mask,547 beta,548 advantages,549 **extra_kwargs,550 )551 552 # Scale loss if needed for mixed precision training553 scaled_loss = loss * scaling554 # Must add .loss.detach otherwise autograd uses 2x VRAM555 return scaled_loss, (loss.detach(), completion_length, mean_kl, delta, flat_is_ratio, coef_1)556 pass557 558 device =_new_logps.device559 grad_inputs = torch.empty_like(_new_logps)560 accumulated_loss = torch.zeros(1, device = device)[0]561 accumulated_completion_length = torch.zeros(1, device = device)[0]562 accumulated_mean_kl = torch.zeros(1, device = device)[0]563 accumulated_delta = []564 accumulated_flat_is_ratio = []565 accumulated_coef_1 = []566 567 def accumulate_chunk(568 new_logps_j,569 old_logps_j,570 ref_logps_j,571 sampling_per_token_logps_j,572 input_ids_j,573 mask_j,574 advantages_j,575 scaling,576 grad_inputs_j,577 ):578 (chunk_grad_input,), (chunk_loss, (unscaled_loss, chunk_completion_length, chunk_mean_kl, chunk_delta, chunk_flat_is_ratio, chunk_coef_1)) = torch.func.grad_and_value(579 compute_loss,580 argnums = (0,),581 has_aux = True,582 )(new_logps_j, old_logps_j, ref_logps_j, sampling_per_token_logps_j, input_ids_j, mask_j, advantages_j, scaling)583 accumulated_loss .add_(unscaled_loss)584 accumulated_completion_length.add_(chunk_completion_length)585 accumulated_mean_kl .add_(chunk_mean_kl)586 accumulated_delta .append(chunk_delta)587 accumulated_flat_is_ratio .append(chunk_flat_is_ratio)588 accumulated_coef_1 .append(chunk_coef_1)589 grad_inputs_j[:] = chunk_grad_input590 pass591 592 accumulate_chunk = torch.compile(593 accumulate_chunk,594 fullgraph = True,595 # [TODO] Dynamic marking causes torch.compile errors if sequence length is long596 dynamic = True,597 options = torch_compile_options,598 )599 600 grad_inputs_chunks = torch.chunk(grad_inputs, chunks = n_chunks, dim = 0)601 new_logps = torch.chunk(_new_logps, chunks = n_chunks, dim = 0)602 if _old_logps is not None:603 old_logps = torch.chunk(_old_logps, chunks = n_chunks, dim = 0)604 else:605 old_logps = [None] * n_chunks606 if _ref_logps is not None:607 ref_logps = torch.chunk(_ref_logps, chunks = n_chunks, dim = 0)608 else:609 ref_logps = [None] * n_chunks610 if _sampling_per_token_logps is not None:611 sampling_per_token_logps = torch.chunk(_sampling_per_token_logps, chunks = n_chunks, dim = 0)612 else:613 sampling_per_token_logps = [None] * n_chunks614 input_ids = torch.chunk(_input_ids, chunks = n_chunks, dim = 0)615 mask = torch.chunk(_mask, chunks = n_chunks, dim = 0)616 advantages = torch.chunk(_advantages, chunks = n_chunks, dim = 0)617 618 # Get mixed precision scaling if seen619 scaling = scaler.get_scale() if scaler is not None else 1.0620 621 # Force torch.compile to use dynamic shapes for seqlen dim622 # mark_dynamic = lambda x: torch._dynamo.mark_dynamic(x, 1)623 624 for (grad_inputs_j, new_logps_j, old_logps_j, ref_logps_j, sampling_per_token_logps_j, input_ids_j, mask_j, advantages_j, ) in \625 zip(grad_inputs_chunks, new_logps, old_logps, ref_logps, sampling_per_token_logps, input_ids, mask, advantages):626 627 # [TODO] Dynamic marking causes torch.compile errors if sequence length is long628 629 # mark_dynamic(new_hidden_states_j)630 # mark_dynamic(ref_hidden_states_j)631 # if old_hidden_states_j is not None:632 # mark_dynamic(old_hidden_states_j)633 # mark_dynamic(input_ids_j)634 # mark_dynamic(mask_j)635 accumulate_chunk(636 new_logps_j,637 old_logps_j,638 ref_logps_j,639 sampling_per_token_logps_j,640 input_ids_j,641 mask_j,642 advantages_j,643 scaling,644 grad_inputs_j,645 )646 pass647 648 grad_inputs .div_(n_chunks)649 accumulated_loss .div_(n_chunks)650 accumulated_completion_length.div_(n_chunks)651 accumulated_mean_kl .div_(n_chunks)652 653 if _sampling_per_token_logps is not None:654 accumulated_delta = torch.cat(accumulated_delta, dim=0)655 accumulated_flat_is_ratio = torch.cat(accumulated_flat_is_ratio, dim=0)656 else:657 accumulated_delta = None658 accumulated_flat_is_ratio = None659 accumulated_coef_1 = torch.cat(accumulated_coef_1, dim=0)660 ctx.save_for_backward(grad_inputs)661 return (662 accumulated_loss,663 accumulated_completion_length,664 accumulated_mean_kl,665 accumulated_delta,666 accumulated_flat_is_ratio,667 accumulated_coef_1668 )669 pass670 671 @staticmethod672 def backward(ctx, grad_output, dcompletion_length, dmean_kl, ddelta, ddflat_is_ratio, dcoef_1):673 (grad_input,) = ctx.saved_tensors674 return (grad_input, None, None, None, None, None, None, None, None, None, None, None)675 pass676 677def grpo_accumulated_loss(678 trainer,679 input_ids,680 attention_mask,681 logits_to_keep,682 completion_mask,683 advantages,684 old_logps,685 ref_logps,686 n_chunks = -1,687 **kwargs,688):689 # All Unsloth Zoo code licensed under AGPL3690 bsz, qlen = input_ids.shape691 692 pixel_values = kwargs.get('pixel_values',None)693 image_grid_thw = kwargs.get('image_grid_thw',None)694 pixel_attention_mask = kwargs.get('pixel_attention_mask',None)695 image_sizes = kwargs.get('image_sizes',None)696 # Transformers 5.x requires token_type_ids/mm_token_type_ids for some vision models697 token_type_ids = kwargs.get('token_type_ids',None)698 mm_token_type_ids = kwargs.get('mm_token_type_ids',None)699 sampling_per_token_logps = kwargs.get("sampling_per_token_logps", None) if getattr(trainer, "vllm_importance_sampling_correction", False) else None700 temperature = kwargs.get("temperature", 1.0)701 logit_scale_multiply = kwargs.get("logit_scale_multiply", 0.0)702 logit_scale_divide = kwargs.get("logit_scale_divide", 0.0)703 logit_softcapping = kwargs.get("logit_softcapping", 0.0)704 prev_max_left_pad = kwargs.get("max_left_pad", 0) #Always get max_left_pad for when training LLMs, enabled by deafult.705 706 #Delete this from kwargs so less issues707 _ = kwargs.pop("sampling_per_token_logps", None)708 kwargs["vllm_importance_sampling_cap"] = trainer.vllm_importance_sampling_cap if sampling_per_token_logps is not None else None709 kwargs["get_sapo_token_loss"] = trainer.get_sapo_token_loss if hasattr(trainer, "get_sapo_token_loss") else None710 kwargs["sapo_temperature_pos"] = trainer.args.sapo_temperature_pos if hasattr(trainer.args, "sapo_temperature_pos") else None711 kwargs["sapo_temperature_neg"] = trainer.args.sapo_temperature_neg if hasattr(trainer.args, "sapo_temperature_neg") else None712 kwargs["get_off_policy_mask"] = trainer.get_off_policy_mask if hasattr(trainer, "get_off_policy_mask") else None713 kwargs["off_policy_mask_threshold"] = trainer.args.off_policy_mask_threshold if hasattr(trainer.args, "off_policy_mask_threshold") else None714 kwargs["use_vllm"] = trainer.use_vllm715 # Find closest multiple716 factors = [i for i in range(1, bsz + 1) if bsz % i == 0]717 if n_chunks == -1: n_chunks = bsz718 n_chunks = factors[min(np.searchsorted(factors, n_chunks), len(factors)-1)]719 720 if not hasattr(trainer, '_autocast_dtype'):721 trainer._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16722 if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': trainer._autocast_dtype = None723 pass724 os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"725 726 lm_head = trainer.model.get_output_embeddings().weight727 dtype_bytes = 16 if trainer._autocast_dtype in [torch.float16, torch.bfloat16] else 32728 729 total_rows = input_ids.shape[0]730 seq_len = input_ids.shape[1]731 hidden_dim = lm_head.shape[1]732 vocab_dim = lm_head.shape[0]733 734 if trainer.args.unsloth_grpo_mini_batch is None:735 if not hasattr(trainer, "_has_autotuned"):736 trainer._has_autotuned = True737 B, multiplier = autotune_batch_and_chunks(738 total_rows, seq_len, hidden_dim, vocab_dim, dtype_bytes, trainer.args.unsloth_logit_chunk_multiplier739 )740 trainer.args.unsloth_grpo_mini_batch = max(1, total_rows//B)741 trainer.args.unsloth_logit_chunk_multiplier = multiplier742 B = trainer.args.unsloth_grpo_mini_batch743 multiplier = trainer.args.unsloth_logit_chunk_multiplier744 elif trainer._step % trainer.current_gradient_accumulation_steps == 0:745 B = trainer.args.unsloth_grpo_mini_batch746 multiplier = trainer.args.unsloth_logit_chunk_multiplier747 del trainer._has_autotuned748 del trainer.args.unsloth_grpo_mini_batch749 del trainer.args.unsloth_logit_chunk_multiplier750 else:751 B = trainer.unsloth_grpo_mini_batch752 multiplier = trainer.args.unsloth_logit_chunk_multiplier753 else:754 if trainer.args.unsloth_grpo_mini_batch > total_rows:755 B = total_rows756 else:757 B = trainer.args.unsloth_grpo_mini_batch758 759 if trainer.args.unsloth_logit_chunk_multiplier is None:760 multiplier = max(4, seq_len // 4096)761 else:762 multiplier = trainer.args.unsloth_logit_chunk_multiplier763 764 if pixel_values is None:765 left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(input_ids, logits_to_keep, trainer.processing_class.pad_token_id)766 767 # Determine max_left_pad from precomputed logprobs shape for consistency768 if old_logps is not None:769 max_left_pad = old_logps.shape[1] - logits_to_keep770 elif ref_logps is not None:771 max_left_pad = ref_logps.shape[1] - logits_to_keep772 else:773 max_left_pad = torch.max(left_pad_tokens_per_prompt).item()774 775 input_ids = left_pack_padding(input_ids, trainer.processing_class.pad_token_id)776 777 completion_input_ids = input_ids[:, -(logits_to_keep +max_left_pad):]778 779 completion_mask = create_completion_attention_mask(completion_input_ids, left_pad_tokens_per_prompt, max_left_pad, trainer.processing_class.pad_token_id).to(attention_mask.dtype)780 781 if trainer.use_vllm and sampling_per_token_logps is not None and getattr(trainer, "vllm_importance_sampling_correction", False):782 sampling_per_token_logps = align_logprobs_with_mask(sampling_per_token_logps, completion_mask)783 else:784 sampling_per_token_logps = None785 attention_mask = input_ids != trainer.processing_class.pad_token_id786 attention_mask = attention_mask.to(attention_mask.dtype)787 else:788 completion_input_ids = input_ids[:, -logits_to_keep:]789 790 unwrapped_model = trainer.accelerator.unwrap_model(trainer.model, keep_fp32_wrapper = False)791 792 for module in unwrapped_model.modules():793 if hasattr(module, "_hf_hook") and hasattr(module._hf_hook, "io_same_decice"):794 module._hf_hook.io_same_decice = False795 pass796 797 all_logprobs_list = []798 799 attention_mask_chunks = torch.chunk(attention_mask, chunks=B, dim=0)800 completion_ids_chunks = torch.chunk(completion_input_ids, chunks=B, dim=0)801 802 def chunk_optional(tensor, chunks):803 if tensor is None:804 return [None] * chunks805 return torch.chunk(tensor, chunks=chunks, dim=0)806 807 import math808 total_samples = input_ids.shape[0]809 batch_size = math.ceil(total_samples / B)810 811 input_ids_chunks = []812 attention_mask_chunks = []813 pixel_values_chunks = []814 image_grid_thw_chunks = []815 pixel_attention_mask_chunks = []816 817 current_pixel_idx = 0818 #TRL 0.23.0 batching logic819 for start in range(0, total_samples, batch_size):820 end = start + batch_size821 822 input_ids_chunks.append(input_ids[start:end])823 attention_mask_chunks.append(attention_mask[start:end])824 825 if image_grid_thw is not None and pixel_values is not None:826 827 grid_slice = image_grid_thw[start:end]828 image_grid_thw_chunks.append(grid_slice)829 batch_pixel_count = grid_slice.prod(dim=-1).sum().item()830 831 start_pixel_idx = current_pixel_idx832 end_pixel_idx = current_pixel_idx + batch_pixel_count833 834 pixel_values_chunks.append(pixel_values[start_pixel_idx:end_pixel_idx])835 836 if pixel_attention_mask is not None:837 pixel_attention_mask_chunks.append(838 pixel_attention_mask[start_pixel_idx:end_pixel_idx]839 )840 else:841 pixel_attention_mask_chunks.append(None)842 843 current_pixel_idx = end_pixel_idx844 845 else:846 pixel_values_chunks.append(None)847 image_grid_thw_chunks.append(None)848 pixel_attention_mask_chunks.append(None)849 850 if image_sizes is not None and not isinstance(image_sizes, torch.Tensor):851 image_sizes_chunks = [[size] for size in image_sizes]852 else:853 image_sizes_chunks = chunk_optional(image_sizes, B)854 855 # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models856 token_type_ids_chunks = chunk_optional(token_type_ids, B)857 mm_token_type_ids_chunks = chunk_optional(mm_token_type_ids, B)858 859 zipped_inputs = zip(860 input_ids_chunks,861 attention_mask_chunks,862 pixel_values_chunks,863 image_grid_thw_chunks,864 pixel_attention_mask_chunks,865 image_sizes_chunks,866 token_type_ids_chunks,867 mm_token_type_ids_chunks,868 completion_ids_chunks869 )870 871 if trainer._autocast_dtype is None:872 autocaster = nullcontext()873 else:874 autocaster = torch.amp.autocast(device_type = trainer.model.device.type, dtype = trainer._autocast_dtype)875 876 def to_device(tensor, device, non_blocking=True):877 if tensor is None: return None878 return tensor.to(device, non_blocking=non_blocking)879 880 class Unsloth_Offloaded_Log_Softmax(torch.autograd.Function):881 """882 Manual Gradient Checkpointing/CPU Offloading for Log Softmax.883 """884 @staticmethod885 def forward(ctx, hidden_states, lm_head, index, chunks,886 logit_scale_multiply, logit_scale_divide,887 logit_softcapping, temperature):888 #Only the activations are needed so if we keep entire computational graph, keeps unnecessary memory on CPU so we detach it889 ctx.saved_hidden_states = hidden_states.detach().contiguous().to("cpu", non_blocking=True)890 ctx.device = hidden_states.device891 ctx.dtype = hidden_states.dtype892 893 ctx.lm_head = lm_head894 ctx.lm_head_requires_grad = lm_head.requires_grad895 ctx.index = index896 ctx.args = (chunks, logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature)897 898 with torch.no_grad():899 output = chunked_hidden_states_selective_log_softmax(900 hidden_states, lm_head, index, *ctx.args901 )902 903 return output904 905 @staticmethod906 def backward(ctx, grad_output):907 hidden_states = to_device(ctx.saved_hidden_states, ctx.device)908 hidden_states = hidden_states.to(ctx.dtype)909 hidden_states.requires_grad_(True)910 911 lm_head = ctx.lm_head912 # #Possibly redundant lines913 # if ctx.lm_head_requires_grad:914 # hidden_states.requires_grad_(True)915 # else:916 # lm_head = lm_head.detach()917 918 index = ctx.index919 920 with torch.enable_grad():921 output = chunked_hidden_states_selective_log_softmax(922 hidden_states, lm_head, index, *ctx.args923 )924 925 torch.autograd.backward(output, grad_output)926 927 return (928 hidden_states.grad,929 lm_head.grad if ctx.lm_head_requires_grad else None,930 None,931 None,932 None,933 None,934 None,935 None,936 )937 938 def efficient_log_softmax(hidden_states, lm_head, index, chunks=32,939 logit_scale_multiply=0.0, logit_scale_divide=0.0,940 logit_softcapping=0.0, temperature=1, batch_size=8):941 if (index.shape[1] <= 1024 and batch_size <= 8) or batch_size==1:942 #We save a gigabyte or speed with the normal path under these specific conditions943 return chunked_hidden_states_selective_log_softmax(944 hidden_states,945 lm_head,946 index,947 chunks,948 logit_scale_multiply,949 logit_scale_divide,950 logit_softcapping,951 temperature952 )953 else:954 return Unsloth_Offloaded_Log_Softmax.apply(955 hidden_states, lm_head, index, chunks,956 logit_scale_multiply, logit_scale_divide,957 logit_softcapping, temperature958 )959 for (960 input_ids_chunk,961 attention_mask_chunk,962 pixel_values_chunk,963 image_grid_thw_chunk,964 pixel_attention_mask_chunk,965 image_sizes_chunk,966 token_type_ids_chunk,967 mm_token_type_ids_chunk,968 completion_ids969 ) in zipped_inputs:970 _extra_vision_kwargs = {}971 if token_type_ids_chunk is not None:972 _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk973 if mm_token_type_ids_chunk is not None:974 _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk975 with autocaster:976 if pixel_values is None:977 new_hidden_states_chunk = unwrapped_model(978 input_ids = input_ids_chunk,979 attention_mask = attention_mask_chunk,980 pixel_values = pixel_values_chunk,981 image_grid_thw = image_grid_thw_chunk,982 pixel_attention_mask = pixel_attention_mask_chunk,983 image_sizes = image_sizes_chunk,984 **_extra_vision_kwargs,985 ).logits986 987 new_hidden_states_chunk = new_hidden_states_chunk[:, -(logits_to_keep + max_left_pad + 1): , :]988 new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :]989 logprobs_chunk = efficient_log_softmax(990 new_hidden_states_chunk,991 lm_head,992 completion_ids,993 chunks=input_ids_chunk.shape[0]*multiplier,994 logit_scale_multiply=logit_scale_multiply,995 logit_scale_divide=logit_scale_divide,996 logit_softcapping=logit_softcapping,997 temperature=temperature,998 batch_size = B999 )1000 else:1001 new_hidden_states_chunk = unwrapped_model(1002 input_ids = input_ids_chunk,1003 attention_mask = attention_mask_chunk,1004 pixel_values = pixel_values_chunk,1005 image_grid_thw = image_grid_thw_chunk,1006 pixel_attention_mask = pixel_attention_mask_chunk,1007 image_sizes = image_sizes_chunk,1008 logits_to_keep = logits_to_keep + 1,1009 **_extra_vision_kwargs,1010 ).logits1011 1012 new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :]1013 # Guard: check if model returned hidden states or logits1014 if new_hidden_states_chunk.shape[-1] == lm_head.shape[1]:1015 logprobs_chunk = efficient_log_softmax(1016 new_hidden_states_chunk,1017 lm_head,1018 completion_ids,1019 chunks=input_ids_chunk.shape[0]*multiplier,1020 logit_scale_multiply=logit_scale_multiply,1021 logit_scale_divide=logit_scale_divide,1022 logit_softcapping=logit_softcapping,1023 temperature=temperature,1024 batch_size = B1025 )1026 else:1027 # Model returned logits directly - scaling/softcapping already applied by model forward1028 logprobs_chunk = chunked_selective_log_softmax(new_hidden_states_chunk, completion_ids, temperature)1029 #This is needed to avoid race conditions with GPT OSS offload_embbed=True1030 #However, it seems that this line does not slow down or disrupt models.1031 device_synchronize()1032 all_logprobs_list.append(logprobs_chunk)1033 1034 new_logprobs = torch.cat(all_logprobs_list, dim=0)1035 1036 with autocaster:1037 loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1 = UnslothEfficientGRPO.apply(1038 new_logprobs,1039 old_logps,1040 ref_logps,1041 sampling_per_token_logps,1042 lm_head,1043 completion_input_ids,1044 completion_mask,1045 advantages,1046 trainer.beta,1047 trainer.accelerator.scaler,1048 1,1049 kwargs1050 )1051 1052 # Must force not returning hidden states but logits otherwise gibberish1053 os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0"1054 1055 return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, completion_mask1056 # Old non efficient code path1057 new_logits = torch.matmul(new_hidden_states, lm_head.t())1058 new_logits = new_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred1059 old_logits = torch.matmul(old_hidden_states, lm_head.t())1060 old_logits = old_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred1061 loss, completion_length, mean_kl = grpo_compute_loss(1062 old_logits,1063 new_logits,1064 completion_input_ids,1065 completion_mask,1066 trainer.beta,1067 advantages,1068 )1069 return loss, completion_length, mean_kl1070 pass1071 1072@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options)1073def grpo_compute_loss_slow(1074 ref,1075 new,1076 old,1077 sampling_per_token_logps,1078 input_ids,1079 mask,1080 beta,1081 advantages,1082 **kwargs1083):1084 # All Unsloth Zoo code licensed under AGPL31085 # Set defaults for optional arguments1086 loss_type = kwargs.get("loss_type", "grpo")1087 epsilon_low = kwargs.get("epsilon_low", 0.2)1088 epsilon_high = kwargs.get("epsilon_high", 0.2)1089 max_completion_length = kwargs.get("max_completion_length", 8192)1090 delta = kwargs.get("delta", None)1091 importance_sampling_level = kwargs.get("importance_sampling_level", "token")1092 num_items_in_batch = kwargs.get("num_items_in_batch", None)1093 current_gradient_accumulation_steps = kwargs.get("current_gradient_accumulation_steps", 1)1094 num_processes = kwargs.get("num_processes", 1)1095 use_vllm = kwargs.get("use_vllm", False)1096 vllm_importance_sampling_cap = kwargs.get("vllm_importance_sampling_cap", 2.0)1097 get_sapo_token_loss = kwargs.get("get_sapo_token_loss", None)1098 sapo_temperature_pos = kwargs.get("sapo_temperature_pos", 1.0)1099 sapo_temperature_neg = kwargs.get("sapo_temperature_neg", 1.05)1100 get_off_policy_mask = kwargs.get("get_off_policy_mask", None)1101 off_policy_mask_threshold = kwargs.get("off_policy_mask_threshold", None)1102 input_ids = input_ids.unsqueeze(-1)1103 1104 if advantages.dim() == 1:1105 advantages = advantages.unsqueeze(1)1106 1107 if off_policy_mask_threshold is not None:1108 off_policy_mask = get_off_policy_mask(1109 advantages=advantages,1110 per_token_logps=new,1111 old_per_token_logps=old,1112 mask=mask,1113 off_policy_threshold=off_policy_mask_threshold,1114 )1115 1116 with torch.no_grad():1117 if use_vllm and sampling_per_token_logps is not None:1118 #must filter out extra prompt tokens in begining after making input_ids left padded1119 importance_sampling_ratio = torch.exp((old * mask) - sampling_per_token_logps)1120 importance_sampling_ratio = torch.clamp(1121 importance_sampling_ratio, max=vllm_importance_sampling_cap1122 )1123 pass1124 1125 # Must detach - otherwise gradients are not propagated correctly!1126 # exp(x - x) == 11127 # loss_i = torch.exp(new - new.detach()) * advantages.unsqueeze(1)1128 if old is not None:1129 log_ratio = new - old1130 else:1131 log_ratio = new - new.detach()1132 1133 if importance_sampling_level == "token":1134 log_importance_weights = log_ratio1135 elif importance_sampling_level == "sequence":1136 log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)1137 log_importance_weights = log_importance_weights.unsqueeze(-1)1138 else:1139 raise ValueError(1140 f"Unknown importance sampling level: {importance_sampling_level}. Possible values are 'token' "1141 "and 'sequence'."1142 )1143 1144 coef_1 = torch.exp(log_importance_weights)1145 1146 # Reverse KL1147 # Note that this is a low variance low bias estimator for the KL divergence as used in GRPO paper1148 if beta != 0.0:1149 kl_i = torch.exp(ref - new) - (ref - new) - 1.01150 1151 else:1152 # set kl_i to a tensor of zeros with the correct shape1153 if importance_sampling_level == "sequence":1154 kl_i = new.new_zeros(new.size(0), 1)1155 else:1156 kl_i = torch.zeros_like(new)1157 # Full correct reverse KL divergence?? Missing term maybe?1158 # kl_i = torch.exp(new) * kl_i1159 1160 # Below is forward KL (normal KL)1161 # kl_i = torch.exp(old) * (old - new)1162 if loss_type == "cispo":1163 clamped_ratios = torch.clamp(coef_1, max=epsilon_high).detach()1164 loss_i = -clamped_ratios * advantages * new1165 #breakpoint()1166 elif loss_type in ["grpo", "bnpo", "dr_grpo", "dapo"]:1167 coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high)1168 1169 if delta is not None:1170 loss_1 = torch.clamp(coef_1, max=delta) * advantages1171 else:1172 loss_1 = coef_1 * advantages1173 pass1174 loss_2 = coef_2 * advantages1175 loss_i = -torch.min(loss_1, loss_2)1176 elif loss_type == "sapo":1177 if get_sapo_token_loss is None:1178 raise Exception(f"sapo is only available in TRL 0.26.0+")1179 loss_i = torch.empty_like(coef_1)1180 positive_advantages_mask = advantages.repeat([1, coef_1.shape[1]]) > 01181 #since we have n_chunks some tensors may error if they dont have elements in them1182 if coef_1[positive_advantages_mask].numel() != 0:1183 loss_i[positive_advantages_mask] = get_sapo_token_loss(1184 coef_1[positive_advantages_mask], sapo_temperature_pos1185 )1186 if coef_1[~positive_advantages_mask].numel() != 0:1187 loss_i[~positive_advantages_mask] = get_sapo_token_loss(1188 coef_1[~positive_advantages_mask], sapo_temperature_neg1189 )1190 loss_i = -loss_i * advantages1191 else:1192 raise ValueError(f"Unknown loss type: {loss_type}")1193 1194 if off_policy_mask_threshold is not None:1195 loss_i = loss_i * off_policy_mask1196 1197 if use_vllm and sampling_per_token_logps is not None:1198 loss_i = loss_i * importance_sampling_ratio1199 #delta for metric1200 with torch.no_grad():