Memory efficient GRPO, DPO etc (#1716)
* Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Faster inference? * Update llama.py * Update llama.py * Update utils.py * Update llama.py * Update llama.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update mapper.py * Fast Inference via vLLM * Update llama.py * Update llama.py * Update utils.py * Create rl.py * PatchRL * Update rl.py * Update rl.py * Update rl.py * PatchRLStatistics * Update rl.py * Update rl.py * Update rl.py * Update utils.py * Update utils.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * RL metrics * Update rl.py * RL metrics * Update __init__.py * Update rl.py * Update rl.py * Update rl.py * Update chat_templates.py * Update mapper.py * Fp8 cache * Update llama.py * Update llama.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update __init__.py * Update loader.py * Update rl.py * Update rl.py * Update _utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Better TRL handling * Update rl.py * Update tokenizer_utils.py * Auto patching * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update rl.py * Update tokenizer_utils.py * Update rl.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update tokenizer_utils.py * Update rl.py * Update rl.py * Update rl.py * max seq length * Update rl.py * Update rl.py * Patching * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * NEFTune * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Extra replacements * Update rl_replacements.py * Update rl.py * extra RL replacements * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update _utils.py * Update loader_utils.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * autocast * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update rl_replacements.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update pyproject.toml * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update _utils.py * Update llama.py * Update _utils.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update rl_replacements.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * GRPO optimized * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Selective Log softmax * Fix GRPO bsz * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Fix TRL
This commit is contained in:
Daniel Han
GitHub
parent
c51f8b41db
commit
045e94033f
4 changed files with 167 additions and 38 deletions
|
|
@ -39,7 +39,7 @@ triton = [
|
|||
"triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
|
||||
]
|
||||
huggingface = [
|
||||
"unsloth_zoo>=2025.2.2",
|
||||
"unsloth_zoo>=2025.2.5",
|
||||
"packaging",
|
||||
"tyro",
|
||||
"transformers>=4.46.1,!=4.47.0",
|
||||
|
|
@ -50,7 +50,7 @@ huggingface = [
|
|||
"wheel>=0.42.0",
|
||||
"numpy",
|
||||
"accelerate>=0.34.1",
|
||||
"trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,<0.15.0",
|
||||
"trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3",
|
||||
"peft>=0.7.1,!=0.11.0",
|
||||
"protobuf<4.0.0",
|
||||
"huggingface_hub",
|
||||
|
|
@ -176,26 +176,26 @@ cu124onlytorch251 = [
|
|||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
|
||||
]
|
||||
cu118onlytorch260 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
]
|
||||
cu124onlytorch260 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp310-cp310-win_amd64.whl ; python_version=='3.10' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
|
||||
]
|
||||
cu126onlytorch260 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
]
|
||||
cu118 = [
|
||||
"unsloth[huggingface]",
|
||||
|
|
@ -344,7 +344,7 @@ colab-ampere-torch220 = [
|
|||
"flash-attn>=2.6.3",
|
||||
]
|
||||
colab-new = [
|
||||
"unsloth_zoo>=2025.2.2",
|
||||
"unsloth_zoo>=2025.2.5",
|
||||
"packaging",
|
||||
"tyro",
|
||||
"transformers>=4.46.1,!=4.47.0",
|
||||
|
|
@ -362,7 +362,7 @@ colab-new = [
|
|||
]
|
||||
colab-no-deps = [
|
||||
"accelerate>=0.34.1",
|
||||
"trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,<0.15.0",
|
||||
"trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3",
|
||||
"peft>=0.7.1",
|
||||
"xformers",
|
||||
"bitsandbytes>=0.46.1",
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__version__ = "2025.2.9"
|
||||
__version__ = "2025.2.10"
|
||||
|
||||
__all__ = [
|
||||
"SUPPORTS_BFLOAT16",
|
||||
|
|
|
|||
|
|
@ -21,12 +21,25 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
|
|||
import inspect
|
||||
import os
|
||||
import re
|
||||
import torch
|
||||
from unsloth_zoo.compiler import create_new_function
|
||||
from unsloth_zoo.logging_utils import PatchRLStatistics
|
||||
from unsloth_zoo.rl_replacements import RL_REPLACEMENTS
|
||||
from .rl_replacements import (
|
||||
RL_EXTRA_ARGS,
|
||||
RL_FUNCTIONS,
|
||||
RL_PRE_ITEMS,
|
||||
RL_CONFIG_CHANGES,
|
||||
)
|
||||
selective_log_softmax = RL_REPLACEMENTS["selective_log_softmax"]
|
||||
|
||||
torch_compile_options = {
|
||||
"epilogue_fusion" : True,
|
||||
"max_autotune" : True,
|
||||
"shape_padding" : True,
|
||||
"trace.enabled" : False,
|
||||
"triton.cudagraphs" : False,
|
||||
}
|
||||
|
||||
def PatchRL(FastLanguageModel):
|
||||
|
||||
|
|
@ -81,6 +94,17 @@ from dataclasses import dataclass, field
|
|||
from packaging.version import Version
|
||||
import torch
|
||||
from contextlib import nullcontext
|
||||
from torch.nn import functional as F
|
||||
torch_compile_options = {{
|
||||
"epilogue_fusion" : True,
|
||||
"max_autotune" : True,
|
||||
"shape_padding" : True,
|
||||
"trace.enabled" : False,
|
||||
"triton.cudagraphs" : False,
|
||||
}}
|
||||
|
||||
{selective_log_softmax_code}
|
||||
{RL_pre}
|
||||
|
||||
@dataclass
|
||||
class Unsloth{RLConfig_name}({RLConfig_name}):
|
||||
|
|
@ -142,8 +166,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
if RLTrainer.__name__.startswith("Unsloth"): return
|
||||
if RLConfig .__name__.startswith("Unsloth"): return
|
||||
|
||||
# Get old source
|
||||
old_RLTrainer_source = inspect.getsource(RLTrainer)
|
||||
old_RLConfig_source = inspect.getsource(RLConfig)
|
||||
|
||||
all_imports = dir(trainer)
|
||||
imports = [x for x in all_imports if not x.startswith("_")]
|
||||
# Fix _deprecate_arguments not getting imported so stop __ but not _
|
||||
imports = [x for x in all_imports if not x.startswith("__")]
|
||||
|
||||
# Get default arguments
|
||||
EMPTY = inspect.Parameter.empty
|
||||
|
|
@ -358,6 +387,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
extra_args += num_proc_check
|
||||
pass
|
||||
|
||||
# Edit config with anything extra
|
||||
if trainer_file in RL_CONFIG_CHANGES:
|
||||
process_extra_args = RL_CONFIG_CHANGES[trainer_file]
|
||||
for process_extra_arg in process_extra_args:
|
||||
extra_args += process_extra_arg(old_RLTrainer_source, old_RLConfig_source)
|
||||
pass
|
||||
|
||||
# Edit report_to and default it to nothing if max_steps is like 60
|
||||
|
||||
# Create RLConfig args
|
||||
|
|
@ -377,6 +413,17 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
__RLTrainer_doc__ = eval(f"trl.trainer.{RLTrainer_name}").__doc__
|
||||
__RLConfig_doc__ = eval(f"trl.trainer.{RLConfig_name}") .__doc__
|
||||
|
||||
# Get all pre-modules
|
||||
if trainer_file in RL_PRE_ITEMS:
|
||||
RL_pre = "\n".join(RL_PRE_ITEMS[trainer_file])
|
||||
else:
|
||||
RL_pre = ""
|
||||
pass
|
||||
|
||||
# Selective log softmax
|
||||
selective_log_softmax_code = inspect.getsource(selective_log_softmax)
|
||||
|
||||
# Get final source code
|
||||
RLTrainer_source = RLTrainer_replacement.format(
|
||||
RLTrainer_name = RLTrainer_name,
|
||||
__RLTrainer_doc__ = __RLTrainer_doc__,
|
||||
|
|
@ -394,6 +441,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
|
||||
RLTrainer_extras = RLTrainer_extras,
|
||||
RLTrainer_post = RLTrainer_post,
|
||||
RL_pre = RL_pre,
|
||||
|
||||
selective_log_softmax_code = selective_log_softmax_code,
|
||||
)
|
||||
|
||||
# Create new function
|
||||
|
|
@ -402,7 +452,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
RLTrainer_source,
|
||||
f"trl.trainer.{trainer_file}",
|
||||
imports,
|
||||
overwrite = False,
|
||||
overwrite = True,
|
||||
)
|
||||
|
||||
# Patch Trainer
|
||||
|
|
|
|||
|
|
@ -15,14 +15,27 @@
|
|||
__all__ = [
|
||||
"RL_EXTRA_ARGS",
|
||||
"RL_FUNCTIONS",
|
||||
"RL_PRE_ITEMS",
|
||||
"RL_CONFIG_CHANGES",
|
||||
]
|
||||
|
||||
import re
|
||||
import torch
|
||||
import inspect
|
||||
from collections import defaultdict
|
||||
RL_EXTRA_ARGS = defaultdict(list)
|
||||
RL_FUNCTIONS = defaultdict(list)
|
||||
from unsloth_zoo.rl_replacements import RL_REPLACEMENTS
|
||||
RL_EXTRA_ARGS = defaultdict(list)
|
||||
RL_FUNCTIONS = defaultdict(list)
|
||||
RL_PRE_ITEMS = defaultdict(list)
|
||||
RL_CONFIG_CHANGES = defaultdict(list)
|
||||
|
||||
torch_compile_options = {
|
||||
"epilogue_fusion" : True,
|
||||
"max_autotune" : True,
|
||||
"shape_padding" : True,
|
||||
"trace.enabled" : False,
|
||||
"triton.cudagraphs" : False,
|
||||
}
|
||||
|
||||
# Check untrained tokens
|
||||
def sft_trainer_fix_untraiend_tokens(call_args, extra_args):
|
||||
|
|
@ -161,23 +174,89 @@ RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__move_model_to_vllm)
|
|||
def grpo_trainer__get_per_token_logps(function_name, function):
|
||||
if function_name != "_get_per_token_logps": return function
|
||||
|
||||
# Edit model to autocast it
|
||||
# .*? matches first match. .+? matches final match.
|
||||
original = re.findall(
|
||||
r"\n([ ]{4,})(logits = model\(.*?\))",
|
||||
function,
|
||||
flags = re.MULTILINE | re.DOTALL,
|
||||
)
|
||||
if len(original) != 0:
|
||||
spaces, original = original[0]
|
||||
spaces = len(spaces)
|
||||
replacer = \
|
||||
"if not hasattr(self, '_autocast_dtype'):\n" + \
|
||||
" "*(spaces + 4) + "self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16\n" + \
|
||||
" "*(spaces + 0) + "with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype):\n" + \
|
||||
" "*(spaces + 4) + original
|
||||
function = function.replace(original, replacer)
|
||||
def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep):
|
||||
if not hasattr(self, '_autocast_dtype'):
|
||||
self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16
|
||||
with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype):
|
||||
# We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
|
||||
logits = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits
|
||||
logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
|
||||
|
||||
input_ids = input_ids[:, -logits_to_keep:]
|
||||
# For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
|
||||
# See https://github.com/huggingface/trl/issues/2770
|
||||
logits = logits[:, -logits_to_keep:]
|
||||
return logits
|
||||
# return selective_log_softmax(logits, input_ids) # compute logprobs for the input tokens
|
||||
pass
|
||||
pass
|
||||
|
||||
function = inspect.getsource(_get_per_token_logps)
|
||||
return function
|
||||
pass
|
||||
RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps)
|
||||
|
||||
grpo_compute_loss = RL_REPLACEMENTS["grpo_compute_loss"]
|
||||
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_compute_loss))
|
||||
|
||||
# Edit _get_per_token_logps to handle mixed precision
|
||||
def grpo_trainer_compute_loss(function_name, function):
|
||||
if function_name != "compute_loss": return function
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch = None):
|
||||
if return_outputs:
|
||||
raise ValueError("The GRPOTrainer does not support returning outputs")
|
||||
# Compute the per-token log probabilities for the model
|
||||
|
||||
prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
|
||||
completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"]
|
||||
input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
|
||||
# attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
|
||||
attention_mask = None
|
||||
logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens
|
||||
|
||||
per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
|
||||
|
||||
# Compute the KL divergence between the model and the reference model
|
||||
ref_per_token_logps = inputs["ref_per_token_logps"]
|
||||
# per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
|
||||
|
||||
# x - x.detach() allows for preserving gradients from x
|
||||
advantages = inputs["advantages"]
|
||||
# per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
|
||||
# per_token_loss = -(per_token_loss - self.beta * per_token_kl)
|
||||
# loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
|
||||
input_ids = input_ids[:, -logits_to_keep:]
|
||||
loss, completion_length, mean_kl = grpo_compute_loss(
|
||||
ref_per_token_logps, per_token_logps, input_ids, completion_mask, self.beta, advantages,
|
||||
)
|
||||
# Log the metrics
|
||||
# completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
|
||||
self._metrics["completion_length"].append(completion_length.item())
|
||||
|
||||
# mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
|
||||
# self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
|
||||
self._metrics["kl"].append(mean_kl.item())
|
||||
return loss
|
||||
pass
|
||||
|
||||
function = inspect.getsource(compute_loss)
|
||||
return function
|
||||
pass
|
||||
RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_compute_loss)
|
||||
|
||||
# https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356
|
||||
# TRL warns if batch size is not a multiple of num_generations -> fix this.
|
||||
def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source):
|
||||
if "divisible by the number of generations" not in RLTrainer_source: return ""
|
||||
if "num_generations" not in RLConfig_source: return ""
|
||||
|
||||
check_batch_size = \
|
||||
"div = per_device_train_batch_size // num_generations\n"\
|
||||
"if div * num_generations != per_device_train_batch_size:\n"\
|
||||
" print('Unsloth: We know expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n"\
|
||||
"We will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations))\n"\
|
||||
" per_device_train_batch_size = num_generations\n"
|
||||
return check_batch_size
|
||||
pass
|
||||
RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue