diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 0fd0d9f5d28..5d41071f320 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -174,3 +174,12 @@ class GRPOConfig(TrainingArguments): default=0.04, metadata={"help": "KL coefficient."}, ) + + per_device_micro_batch_size: int = field( + default=8, + metadata={ + "help": "Micro batch size per GPU/TPU/MPS/NPU core/CPU for computing loss terms. " + "These microbatches will be accumulated over, resulting in the same effective " + "batch size as `per_device_train_batch_size*num_generations`, but with lower mem footprint." + }, + ) diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index f130d5155f0..5e8064fb21d 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -37,7 +37,7 @@ is_wandb_available, ) from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled -from transformers.utils import is_peft_available +from transformers.utils import is_peft_available, logging from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template from ..import_utils import is_vllm_available @@ -55,6 +55,9 @@ if is_wandb_available(): import wandb + +logger = logging.get_logger("GRPOTrainer") + # What we call a reward function is a callable that takes a list of prompts and completions and returns a list of # rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. RewardFunc = Union[str, PreTrainedModel, Callable[[list, list], list[float]]] @@ -145,6 +148,7 @@ class GRPOTrainer(Trainer): """ _tag_names = ["trl", "grpo"] + args: GRPOConfig # helps with type hinting def __init__( self, @@ -383,6 +387,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -self.max_prompt_length :] prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -self.max_prompt_length :] + logger.info("Starting generation") # Generate completions using either vLLM or regular generation if self.args.use_vllm: # First, have main process load weights if needed @@ -419,80 +424,122 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N else: # Regular generation path with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model: + unwrapped_model.eval() # Needed to make sure use_cache works with gradient_checkpointing prompt_completion_ids = unwrapped_model.generate( **prompt_inputs, generation_config=self.generation_config ) + model.train() + logger.info("Finished generation") + bsz = prompt_completion_ids.size(0) + micro_bsz = self.args.per_device_micro_batch_size or bsz # default to full batch if not set prompt_length = prompt_inputs["input_ids"].size(1) - completion_ids = prompt_completion_ids[:, prompt_length:] - - # Get the per-token log probabilities for the completions for the model and the reference model - def get_per_token_logps(model, input_ids, logits_to_keep): - # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded - logits = model(input_ids, logits_to_keep=logits_to_keep + 1).logits # (B, L, V) - logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred - - # Compute the log probabilities for the input tokens. Use a loop to reduce memory peak. - per_token_logps = [] - for logits_row, input_ids_row in zip(logits, input_ids[:, -logits_to_keep:]): - log_probs = logits_row.log_softmax(dim=-1) - token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1) - per_token_logps.append(token_log_prob) - return torch.stack(per_token_logps) - - logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens - per_token_logps = get_per_token_logps(model, prompt_completion_ids, logits_to_keep) - - with torch.inference_mode(): - if self.ref_model is not None: - ref_per_token_logps = get_per_token_logps(self.ref_model, prompt_completion_ids, logits_to_keep) - else: - with self.accelerator.unwrap_model(model).disable_adapter(): - ref_per_token_logps = get_per_token_logps(model, prompt_completion_ids, logits_to_keep) - - # Compute the KL divergence between the model and the reference model - per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 - - # Mask everything after the first EOS token - is_eos = completion_ids == self.processing_class.eos_token_id - eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) - eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] - sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) - completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + prompts = [prompt for prompt in prompts for _ in range(self.num_generations)] - # Decode the generated completions - completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - if is_conversational(inputs[0]): - completions = [[{"role": "assistant", "content": completion}] for completion in completions] + # Prepare reward kwargs before looping through microbatches + if any(not isinstance(reward_func, PreTrainedModel) for reward_func in self.reward_funcs): + # Repeat all input columns (but "prompt" and "completion") to match the number of generations + all_reward_kwargs = {key: [] for key in inputs[0].keys() if key not in ["prompt", "completion"]} + for key in all_reward_kwargs: + for example in inputs: + all_reward_kwargs[key].extend([example[key]] * self.num_generations) + + # iterate through bsz in loss_bsz chunks and accumulate: + # - rewards + # - per-token log probabilities + # - per-token KL divergences + # - masks of completion tokens + all_rewards_per_func = [] + all_per_token_logps = [] + all_per_token_kl = [] + all_completion_mask = [] + for i in range(0, bsz, micro_bsz): + current_batch_span = slice(i, i + micro_bsz) + + micro_prompt_completion_ids = prompt_completion_ids[current_batch_span] + micro_completion_ids = micro_prompt_completion_ids[:, prompt_length:] + current_bsz = micro_completion_ids.size(0) # last one may be