Upgrade trl fix (#2544)
* Update llama.py making set and reset functions in order to properly use autoSequenceClassification * Update fast_lora.py, added mixed precising pytorch autocasting * Update llama.py did not included rotary embeddings in the reset functions correctly * Update rl.py: correct get reward model added as well as the eval step stuff * Update rl.py removed function that did not need to be patched * Update llama.py: kept reset functions and made their names generic * Update fast_lora.py * Update rl.py, try except * Update fast_lora.py, removing downcasting stuff * Update llama.py removed depircate LLamaLinearScalingRotaryEmbedding * Update rl.py for VLLM RLOO and PPO * Update rl.py reverted * Update rl.py with peft cahnges * Update rl.py, disabling adapters screws inference up * Update rl.py getting PPO support * Update rl.py cleanup * Update rl.py cleaned up not useful commented code * Update llama.py, enabled new flag, keep padding * Upgrade trl fix Signed-off-by: Dattu Sharma <venkatadattasainimmaturi@gmail.com> * Update rl.py made changes relative to the review * Revert accidental patch block for non grpo Signed-off-by: Dattu Sharma <venkatadattasainimmaturi@gmail.com> * Fixup sampling params issue * Fix rl.py regex Signed-off-by: Dattu Sharma <venkatadattasainimmaturi@gmail.com> * loss type: grpo, drgrpo and bnpo Signed-off-by: Dattu Sharma <venkatadattasainimmaturi@gmail.com> * Add trl version check for vllm colocate mode for RL trainers * Update rl.py For TRL 0.18.0 (Main branch of TRL at the time because its on 0.17.0) , the SFT trainer for some reason deletes the labels column and unsloth internal loss funcitons need that column for hte claculations so I add it back in like this. * Update llama.py, merge it to be dattas llama version * Update rl.py, sft changes to get 0.18.0 to be working * Update rl_replacements.py, added hidden state stuff * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py, rechanged the accumlated loss * Fixup num_iterations>1 for grpo Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com> * Update rl_replacements.py * no unnecessary logits upcast. fix naming Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com> * Update rl_replacements.py returned hidden states from logprobs * Update rl_replacements.py removed debug logic * Update rl_replacements.py, should be fine now * Update rl_replacements.py, should take new args for GRPO trainer * Update rl_replacements.py, made it compatible with trl 0.15.2 * Update rl_replacements.py, fixed typo in per tokne-Logps --------- Signed-off-by: Dattu Sharma <venkatadattasainimmaturi@gmail.com> Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com> Co-authored-by: pluesclues <136766175+pluesclues@users.noreply.github.com>
This commit is contained in:
Datta Nimmaturi
GitHub
parent
7a3df703a0
commit
16a007a283
3 changed files with 101 additions and 26 deletions
|
|
@ -43,6 +43,7 @@ torch_compile_options = {
|
|||
"triton.cudagraphs" : False,
|
||||
}
|
||||
|
||||
from trl import __version__ as trl_version
|
||||
|
||||
def vLLMSamplingParams(**kwargs):
|
||||
from vllm import SamplingParams
|
||||
|
|
@ -545,7 +546,12 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
|
||||
selective_log_softmax_code = selective_log_softmax_code,
|
||||
)
|
||||
|
||||
|
||||
if RLTrainer_name == "SFTTrainer":
|
||||
original_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask"]'
|
||||
new_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask","labels"]'
|
||||
RLTrainer_source = RLTrainer_source.replace(original_text, new_text)
|
||||
|
||||
# Remove multiple doc strings
|
||||
if __RLConfig_doc__ != "" and RLTrainer_source.count(__RLTrainer_doc__) == 2:
|
||||
RLTrainer_source = RLTrainer_source.replace(__RLTrainer_doc__, "", 1)
|
||||
|
|
@ -597,9 +603,15 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
|
|||
if len(replacer) != 0:
|
||||
replacer = replacer[0]
|
||||
vllm_setter = "\n" + " "*8 + \
|
||||
"if hasattr(model, 'vllm_engine') and "\
|
||||
"hasattr(args, 'use_vllm') and (getattr(args, 'use_vllm', False) == False): "\
|
||||
"args.use_vllm = True\n"
|
||||
"if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'):\n" + \
|
||||
" " * 12 + "if (getattr(args, 'use_vllm', False) == False):\n" + \
|
||||
" " * 16 + "args.use_vllm = True\n"
|
||||
|
||||
if "grpo" in trainer_file and trl_version >= "0.18":
|
||||
# If model has vllm_engine, then use vllm in colocate mode. Donot wait for server
|
||||
vllm_setter += \
|
||||
" " * 12 + "args.vllm_mode='colocate'\n"
|
||||
|
||||
init = init.replace(replacer, replacer + vllm_setter)
|
||||
pass
|
||||
pass
|
||||
|
|
@ -615,7 +627,8 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
|
|||
if len(vllm_part) == 1:
|
||||
vllm_part, args = vllm_part[0][0], vllm_part[0][1]
|
||||
# Strip all comments
|
||||
new_vllm_part = re.sub(r"\#[^\n]{1,}\n", "", vllm_part)
|
||||
new_vllm_part = re.sub(r"^\s*\#[^\n]*\n?", "", vllm_part, flags=re.MULTILINE) # to also remove whole comment line instead of just starting at #
|
||||
new_vllm_part = re.sub(r"\s*\#.*$", "", new_vllm_part, flags=re.MULTILINE) # remove comments that occur after code
|
||||
|
||||
# Get SamplingParams
|
||||
sampling_params = re.findall(
|
||||
|
|
@ -624,9 +637,9 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
|
|||
new_vllm_part,
|
||||
flags = re.MULTILINE | re.DOTALL,
|
||||
)
|
||||
|
||||
if len(sampling_params) == 1:
|
||||
sampling_params = sampling_params[0]
|
||||
|
||||
# Fix guided_decoding
|
||||
sampling_params = sampling_params.replace(
|
||||
"guided_decoding=guided_decoding,",
|
||||
|
|
@ -638,11 +651,18 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
|
|||
sampling_params = \
|
||||
" "*12 + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + \
|
||||
sampling_params # Add spaces
|
||||
|
||||
# count the indentation of last line of sampling_params.
|
||||
last_line = sampling_params.split("\n")[-1]
|
||||
last_prev_line = sampling_params.split("\n")[-2]
|
||||
last_prev_indentation = len(last_prev_line) - len(last_prev_line.lstrip())
|
||||
last_indentation = len(last_line) - len(last_line.lstrip())
|
||||
|
||||
|
||||
# Add extra arguments to SamplingParams
|
||||
extra = "**getattr(getattr(args, 'vllm_sampling_params', vLLMSamplingParams()), '_set_kwargs', {})"
|
||||
# Backwards replace
|
||||
to_replace = "," + extra + "," + ")"
|
||||
to_replace = ",\n" + " "*last_prev_indentation + extra + ",\n" + " "*last_indentation + ")"
|
||||
sampling_params = to_replace.join(sampling_params.rsplit(")", 1))
|
||||
# Strip multiple commas
|
||||
sampling_params = re.sub(r"[\,][\s]{0,}\,", ",", sampling_params)
|
||||
|
|
@ -650,9 +670,21 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
|
|||
new_vllm_part = \
|
||||
f"\n{' '*8}if {args}.use_vllm:\n{sampling_params}"\
|
||||
f"\n{' '*8}else:\n"
|
||||
|
||||
init = init.replace(vllm_part, new_vllm_part)
|
||||
pass
|
||||
|
||||
if trl_version >= "0.18":
|
||||
# Replace LLM init with already existing vLLM engine for colocate mode
|
||||
vllm_llm_init_pattern = r"self\.llm\s*=\s*LLM\([^)]*\)*\)"
|
||||
vllm_llm_replacement = "self.llm = model.vllm_engine\n"
|
||||
new_vllm_part = re.sub(
|
||||
vllm_llm_init_pattern,
|
||||
vllm_llm_replacement,
|
||||
new_vllm_part,
|
||||
flags=re.DOTALL # Ensure . matches newlines [[5]]
|
||||
)
|
||||
|
||||
init = init.replace(vllm_part, new_vllm_part)
|
||||
|
||||
pass
|
||||
|
||||
# Search for vLLM calling in all child functions
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ __all__ = [
|
|||
"RL_METRICS_CHANGES",
|
||||
]
|
||||
|
||||
import os
|
||||
import re
|
||||
import torch
|
||||
import inspect
|
||||
|
|
@ -207,24 +208,34 @@ 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
|
||||
|
||||
def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep):
|
||||
if os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0':
|
||||
def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep, calc_logprob_flag = None):
|
||||
if os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0' and not calc_logprob_flag:
|
||||
return None # Unsloth efficient GRPO
|
||||
# Otherwise, calculate normally:
|
||||
if not hasattr(self, '_autocast_dtype'):
|
||||
self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16
|
||||
if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': self._autocast_dtype = torch.float16
|
||||
|
||||
os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
|
||||
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:]
|
||||
hidden_states = 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
|
||||
return hidden_states
|
||||
# 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
|
||||
# logits = logits[:, -logits_to_keep:]
|
||||
# return logits
|
||||
# logps = selective_log_softmax(logits, input_ids)
|
||||
|
||||
# row_indices, col_indices = torch.where(logps < -20)
|
||||
|
||||
# # Method 1: Check if tensors have elements
|
||||
# if len(row_indices) > 0 and len(col_indices) > 0:
|
||||
# breakpoint() # Breakpoint triggered here
|
||||
# print("Found high values!")
|
||||
# return logps # compute logprobs for the input tokens
|
||||
pass
|
||||
pass
|
||||
|
||||
|
|
@ -264,7 +275,13 @@ def grpo_trainer_compute_loss(function_name, function):
|
|||
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"]
|
||||
# _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves.
|
||||
# https://github.com/huggingface/trl/blob/05bc43e960396581e458195b8388efe6b82cae1f/trl/trainer/grpo_trainer.py#L1328
|
||||
if self.beta != 0.0:
|
||||
with torch.inference_mode(), model.disable_adapter():
|
||||
ref_per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
|
||||
else:
|
||||
ref_per_token_logps = None
|
||||
# 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
|
||||
|
|
@ -272,16 +289,35 @@ def grpo_trainer_compute_loss(function_name, function):
|
|||
# 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()
|
||||
if "old_per_token_logps" in inputs.keys():
|
||||
old_hidden_states = inputs["old_per_token_logps"]
|
||||
else:
|
||||
old_hidden_states = None
|
||||
input_ids = input_ids[:, -logits_to_keep:]
|
||||
if per_token_logps is not None:
|
||||
loss, completion_length, mean_kl = grpo_compute_loss_slow(
|
||||
ref_per_token_logps, per_token_logps, input_ids, completion_mask, self.beta, advantages,
|
||||
ref_per_token_logps, per_token_logps, old_hidden_states, input_ids, completion_mask, self.beta, advantages,
|
||||
loss_type = self.args.loss_type,
|
||||
epsilon_low = self.epsilon_low, epsilon_high = self.epsilon_high,
|
||||
max_completion_length = self.args.max_completion_length,
|
||||
delta = self.args.delta,
|
||||
)
|
||||
else:
|
||||
loss, completion_length, mean_kl = grpo_accumulated_loss(
|
||||
self, _input_ids, logits_to_keep, completion_mask, advantages,
|
||||
n_chunks = self.args.unsloth_num_chunks,
|
||||
)
|
||||
if hasattr(self.args, "loss_type"):
|
||||
loss, completion_length, mean_kl = grpo_accumulated_loss(
|
||||
self, _input_ids, logits_to_keep, completion_mask, advantages, old_hidden_states,
|
||||
n_chunks = self.args.unsloth_num_chunks,
|
||||
loss_type = self.args.loss_type,
|
||||
epsilon_low = self.epsilon_low, epsilon_high = self.epsilon_high,
|
||||
max_completion_length = self.args.max_completion_length,
|
||||
delta = self.args.delta,
|
||||
)
|
||||
else:
|
||||
# to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17
|
||||
loss, completion_length, mean_kl = grpo_accumulated_loss(
|
||||
self, _input_ids, logits_to_keep, completion_mask, advantages, old_hidden_states,
|
||||
n_chunks = self.args.unsloth_num_chunks,
|
||||
)
|
||||
|
||||
# Log the metrics
|
||||
# completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
|
||||
|
|
|
|||
|
|
@ -194,8 +194,15 @@ def _backwards_compatible_trainer(trainer_class, config_class):
|
|||
config_dict.update(additional_config_kwargs)
|
||||
|
||||
# Create Config with all the collected parameters
|
||||
config = config_class(**config_dict)
|
||||
|
||||
# Reinitialising config class with parameters (that were none initially but populated on first init)
|
||||
# causes the 2nd init to fail as there are mutual exclusive checks on pairs of parameters.
|
||||
# Refer: https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_config.py#L499-L502 for example
|
||||
# So we only create config class if the previous init was not TrainingArguments
|
||||
if not isinstance(training_args, TrainingArguments):
|
||||
config = config_class(**config_dict)
|
||||
else:
|
||||
config = training_args
|
||||
|
||||
# Reconstruct kwargs for Trainer
|
||||
kwargs = trainer_kwargs
|
||||
kwargs["args"] = config
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue