diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index 9901ebc4553..30b3bc627ee 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1676,15 +1676,17 @@ def generate_with_one_unscorable_token(prompts): trainer._generate = generate_with_one_unscorable_token - # Snapshot the divergence metric as soon as it is produced. Reading `_metrics` after training is useless, - # because the dict is cleared on every log. + # Snapshot the divergence metric as soon as it is produced. Reading the accumulators after training is + # useless, because they are cleared on every log. original_score = trainer._generate_and_score_completions recorded_metrics = [] def record_metrics(inputs): outputs = original_score(inputs) - for key in ["sampling/sampling_logp_difference/mean", "sampling/sampling_logp_difference/max"]: - recorded_metrics.extend((key, value) for value in trainer._metrics["train"][key]) + total, count = trainer._metric_stats["train"]["sampling/sampling_logp_difference/mean"] + recorded_metrics.append(("sampling/sampling_logp_difference/mean", (total / count).item())) + max_delta = trainer._metric_maxs["train"]["sampling/sampling_logp_difference/max"] + recorded_metrics.append(("sampling/sampling_logp_difference/max", max_delta.item())) return outputs trainer._generate_and_score_completions = record_metrics diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py index 6482537b9f0..b7f4bc49c8a 100644 --- a/tests/test_sft_trainer.py +++ b/tests/test_sft_trainer.py @@ -1607,6 +1607,22 @@ def convert_to_json(example): new_param = trainer.model.get_parameter(n) assert not torch.equal(param, new_param), f"Parameter {n} has not changed." + def test_log_averages_over_the_window_weighted_by_count(self): + dataset = load_dataset("trl-internal-testing/zen", "standard_language_modeling", split="train") + training_args = SFTConfig(output_dir=self.tmp_dir, report_to="none") + trainer = SFTTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args=training_args, train_dataset=dataset + ) + + # Two steps in one logging window: 1 token with entropy 3.0, then 9 tokens with entropy 1.0. The logged + # value must weight them by their token counts (12 / 10), not average the per-step ratios ((3 + 1) / 2). + trainer.model.train() + trainer._metric_stats["train"]["entropy"] = torch.tensor([3.0 + 9.0, 1.0 + 9.0], device=torch_device) + logs = {} + trainer.log(logs) + + assert logs["entropy"] == pytest.approx(1.2) + def test_train_with_eval(self): dataset = load_dataset("trl-internal-testing/zen", "standard_language_modeling") diff --git a/trl/trainer/distillation_trainer.py b/trl/trainer/distillation_trainer.py index 07292656f9a..391bbd3e2eb 100644 --- a/trl/trainer/distillation_trainer.py +++ b/trl/trainer/distillation_trainer.py @@ -826,6 +826,10 @@ def __init__( # Metrics & Logging self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + # Each entry is a running `(total, count)` pair of on-device tensors, summed in `compute_loss` with no + # collective and no host sync. `log()` reduces the pairs across ranks in a single collective and divides, + # so every metric is weighted by whatever its count counts (tokens, batches). + self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)} self._total_train_tokens = 0 self._current_train_step_time = 0.0 self.log_completions = args.log_completions @@ -1852,15 +1856,14 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N model, unwrapped_student, self._compute_loss, unwrapped_student, inputs, num_items_in_batch ) - # Log the mean per-token student entropy (in nats). The reduction runs here, after `_forward_redirection` - # returns, so the `gather_for_metrics` collective does not run inside the DDP/FSDP-wrapped forward (a hang/ - # ordering risk). The Liger path produces no entropy, so it logs none. Mirrors `SFTTrainer.compute_loss`. + # Log the mean per-token student entropy (in nats). Accumulated as per-rank running sums; the cross-rank + # reduction happens in `log()`, once per logging window, so no collective runs inside the DDP/FSDP-wrapped + # forward. The Liger path produces no entropy, so it logs none. Mirrors `SFTTrainer.compute_loss`. if entropy_sum is not None: mode = "train" if self.model.training else "eval" - num_valid_tokens = self.accelerator.gather_for_metrics(num_valid_tokens).sum() - entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() - entropy = (entropy_sum / num_valid_tokens).item() if num_valid_tokens > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) + self._metric_stats[mode]["entropy"] += torch.stack( + [entropy_sum.detach(), num_valid_tokens.detach().float()] + ) return (loss, None) if return_outputs else loss @@ -2019,6 +2022,15 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: valid = [v for v in val if not math.isnan(v)] metrics[key] = sum(valid) / len(valid) if valid else None + # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that + # all ranks stack them in the same order. + stats = self._metric_stats[mode] + if stats: + keys = sorted(stats) + values = torch.stack([stats[key].double() for key in keys]) + totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True)) + metrics.update({key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}) + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. if mode == "eval": @@ -2027,6 +2039,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: logs.update(metrics) super().log(logs, start_time) self._metrics[mode].clear() + self._metric_stats[mode].clear() if self.accelerator.is_main_process and self.log_completions: if is_rich_available(): diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index a43f3764163..73e104bff6f 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -942,7 +942,10 @@ def __init__( disable_dropout_in_model(self.ref_model) # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + # Each entry is a running `(total, count)` pair of on-device tensors, summed in `compute_loss` with no + # collective and no host sync. `log()` reduces the pairs across ranks in a single collective and divides, + # so every metric is weighted by whatever its count counts (tokens, pairs, batches). + self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)} self._total_train_tokens = 0 # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the @@ -1328,31 +1331,26 @@ def _compute_loss_liger(self, model, inputs, return_outputs): rejected_rewards, ) = metrics + stats = self._metric_stats[mode] if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + num_tokens_in_batch = inputs["attention_mask"].sum() + stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)]) - avg_chosen_logits = self.accelerator.gather_for_metrics(chosen_logits_mean).mean().item() - avg_rejected_logits = self.accelerator.gather_for_metrics(rejected_logits_mean).mean().item() - self._metrics[mode]["logits/chosen"].append(avg_chosen_logits) - self._metrics[mode]["logits/rejected"].append(avg_rejected_logits) + stats["logits/chosen"] += torch.stack([chosen_logits_mean.detach(), torch.ones_like(chosen_logits_mean)]) + stats["logits/rejected"] += torch.stack([rejected_logits_mean.detach(), torch.ones_like(rejected_logits_mean)]) - agg_chosen_rewards = self.accelerator.gather(chosen_rewards) - agg_rejected_rewards = self.accelerator.gather(rejected_rewards) - self._metrics[mode]["rewards/chosen"].append(agg_chosen_rewards.mean().item()) - self._metrics[mode]["rewards/rejected"].append(agg_rejected_rewards.mean().item()) + num_pairs = torch.tensor(float(chosen_rewards.numel()), device=chosen_rewards.device) + stats["rewards/chosen"] += torch.stack([chosen_rewards.detach().sum(), num_pairs]) + stats["rewards/rejected"] += torch.stack([rejected_rewards.detach().sum(), num_pairs]) reward_accuracies = (chosen_rewards > rejected_rewards).float() - agg_reward_accuracies = self.accelerator.gather(reward_accuracies) - self._metrics[mode]["rewards/accuracies"].append(agg_reward_accuracies.mean().item()) + stats["rewards/accuracies"] += torch.stack([reward_accuracies.sum(), num_pairs]) margins = chosen_rewards - rejected_rewards - agg_margins = self.accelerator.gather(margins) - self._metrics[mode]["rewards/margins"].append(agg_margins.mean().item()) + stats["rewards/margins"] += torch.stack([margins.detach().sum(), num_pairs]) - self._metrics[mode]["logps/chosen"].append(self.accelerator.gather(chosen_logps).mean().item()) - self._metrics[mode]["logps/rejected"].append(self.accelerator.gather(rejected_logps).mean().item()) + stats["logps/chosen"] += torch.stack([chosen_logps.detach().sum(), num_pairs]) + stats["logps/rejected"] += torch.stack([rejected_logps.detach().sum(), num_pairs]) return loss @@ -1610,77 +1608,56 @@ def _compute_loss(self, model, inputs, return_outputs): if self.aux_loss_enabled: aux_loss = outputs.aux_loss loss = loss + self.router_aux_loss_coef * aux_loss - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) + self._metric_stats[mode]["aux_loss"] += torch.stack([aux_loss.detach(), torch.ones_like(aux_loss)]) # Log the metrics + stats = self._metric_stats[mode] + # Entropy per_token_entropy = entropy_from_logits(shift_logits.detach()) mask = shift_completion_mask - entropy_sum = (per_token_entropy * mask).sum() - total_tokens = mask.sum() - - # Gather counts across ranks and weight-average - entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() - total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum() - entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) + stats["entropy"] += torch.stack([(per_token_entropy * mask).sum(), mask.sum().float()]) # Number of tokens if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + num_tokens_in_batch = inputs["attention_mask"].sum() + stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)]) # Average logits for chosen and rejected completions chosen_logits, rejected_logits = shift_logits.detach().chunk(2, dim=0) chosen_mask, rejected_mask = shift_completion_mask.chunk(2, dim=0) - total_chosen_logits = chosen_logits[chosen_mask.bool()].mean(-1).sum() - total_chosen_tokens = chosen_mask.sum() - total_rejected_logits = rejected_logits[rejected_mask.bool()].mean(-1).sum() - total_rejected_tokens = rejected_mask.sum() - total_chosen_logits = self.accelerator.gather_for_metrics(total_chosen_logits).sum().item() - total_chosen_tokens = self.accelerator.gather_for_metrics(total_chosen_tokens).sum().item() - total_rejected_logits = self.accelerator.gather_for_metrics(total_rejected_logits).sum().item() - total_rejected_tokens = self.accelerator.gather_for_metrics(total_rejected_tokens).sum().item() - avg_chosen_logits = total_chosen_logits / total_chosen_tokens if total_chosen_tokens > 0 else 0.0 - avg_rejected_logits = total_rejected_logits / total_rejected_tokens if total_rejected_tokens > 0 else 0.0 - self._metrics[mode]["logits/chosen"].append(avg_chosen_logits) - self._metrics[mode]["logits/rejected"].append(avg_rejected_logits) + stats["logits/chosen"] += torch.stack( + [chosen_logits[chosen_mask.bool()].mean(-1).sum(), chosen_mask.sum().float()] + ) + stats["logits/rejected"] += torch.stack( + [rejected_logits[rejected_mask.bool()].mean(-1).sum(), rejected_mask.sum().float()] + ) # Token accuracy for the chosen completions predictions = chosen_logits.argmax(dim=-1) chosen_mask = shift_completion_mask[: len(shift_completion_mask) // 2].bool() chosen_labels = shift_labels[: len(shift_labels) // 2] correct_predictions = (predictions == chosen_labels) & chosen_mask - total_tokens = chosen_mask.sum() - correct_tokens = correct_predictions.sum() - correct_tokens = self.accelerator.gather_for_metrics(correct_tokens) - total_tokens = self.accelerator.gather_for_metrics(total_tokens) - total_sum = total_tokens.sum() - accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0 - self._metrics[mode]["mean_token_accuracy"].append(accuracy) + stats["mean_token_accuracy"] += torch.stack([correct_predictions.sum(), chosen_mask.sum()]) # Rewards for chosen and rejected completions chosen_rewards = self.beta * chosen_logratios.detach() rejected_rewards = self.beta * rejected_logratios.detach() - agg_chosen_rewards = self.accelerator.gather(chosen_rewards) - agg_rejected_rewards = self.accelerator.gather(rejected_rewards) - self._metrics[mode]["rewards/chosen"].append(agg_chosen_rewards.mean().item()) - self._metrics[mode]["rewards/rejected"].append(agg_rejected_rewards.mean().item()) + num_pairs = torch.tensor(float(chosen_rewards.numel()), device=chosen_rewards.device) + stats["rewards/chosen"] += torch.stack([chosen_rewards.sum(), num_pairs]) + stats["rewards/rejected"] += torch.stack([rejected_rewards.sum(), num_pairs]) # Reward accuracy reward_accuracies = (chosen_rewards > rejected_rewards).float() - agg_reward_accuracies = self.accelerator.gather(reward_accuracies) - self._metrics[mode]["rewards/accuracies"].append(agg_reward_accuracies.mean().item()) + stats["rewards/accuracies"] += torch.stack([reward_accuracies.sum(), num_pairs]) # Reward margins margins = chosen_rewards - rejected_rewards - agg_margins = self.accelerator.gather(margins) - self._metrics[mode]["rewards/margins"].append(agg_margins.mean().item()) + stats["rewards/margins"] += torch.stack([margins.sum(), num_pairs]) # Average log probabilities for chosen and rejected completions - self._metrics[mode]["logps/chosen"].append(self.accelerator.gather(chosen_logps).mean().item()) - self._metrics[mode]["logps/rejected"].append(self.accelerator.gather(rejected_logps).mean().item()) + stats["logps/chosen"] += torch.stack([chosen_logps.detach().sum(), num_pairs]) + stats["logps/rejected"] += torch.stack([rejected_logps.detach().sum(), num_pairs]) return (loss, outputs) if return_outputs else loss @@ -1783,14 +1760,29 @@ def training_step(self, *args, **kwargs): def log(self, logs: dict[str, float], start_time: float | None = None) -> None: mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that + # all ranks stack them in the same order. + metrics = {} + stats = self._metric_stats[mode] + if stats: + keys = sorted(stats) + values = torch.stack([stats[key].double() for key in keys]) + totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True)) + metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()} + # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances + # on a train-mode log, so an eval log in between can lag by up to one logging window. + if mode == "train" and "num_tokens" in totals: + self._total_train_tokens += int(totals["num_tokens"][0]) + metrics["num_tokens"] = self._total_train_tokens + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. if mode == "eval": metrics = {f"eval_{key}": val for key, val in metrics.items()} logs.update(metrics) super().log(logs, start_time) - self._metrics[mode].clear() + self._metric_stats[mode].clear() # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index de126792d2f..a8240e04bb6 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -84,8 +84,6 @@ get_config_model_id, identity, maybe_gather_lm_head_ctx, - nanmax, - nanmin, nanstd, pad, print_prompt_completions_sample, @@ -1058,6 +1056,13 @@ def cast_outputs_to_original_dtype(module, args, output): # Initialize the metrics self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + # `_metric_stats` entries are running `(total, count)` pairs and `_metric_mins`/`_metric_maxs` are running + # extrema, all on-device: they only receive local tensors (no collective, no host sync), and `log()` + # aggregates them across ranks in one collective per kind. `_metrics` above keeps plain floats for values + # that are already identical on every rank. + self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)} + self._metric_mins = {"train": {}, "eval": {}} + self._metric_maxs = {"train": {}, "eval": {}} self._total_train_tokens = 0 self._current_train_step_time = 0.0 self.log_completions = args.log_completions @@ -2865,56 +2870,52 @@ def _generate_and_score_completions( self._logs["extra"][column].extend(gather_object(self._pending_extra_logs[column])) self._pending_extra_logs.clear() - # Flush user-logged metrics (from log_metric), averaging across processes. - # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values - # get mis-attributed across metrics (dict insertion order may differ between processes). - for name in sorted(self._pending_metrics): + # Flush user-logged metrics (from log_metric), accumulated locally and averaged across processes at `log()` + # time. Every rank must log the same metric names, otherwise the log-time aggregation mismatches ranks + # (the same requirement the previous per-step gather had). + for name in self._pending_metrics: values = self._pending_metrics[name] local_mean = sum(values) / len(values) - global_mean = self.accelerator.gather(torch.tensor(local_mean, device=device)).mean().item() - self._metrics[mode][name].append(global_mean) + local_mean = torch.tensor(local_mean, device=device) + self._metric_stats[mode][name] += torch.stack([local_mean, torch.ones_like(local_mean)]) self._pending_metrics.clear() if images is not None and self.log_multimodal: self._logs["images"].extend(gather_object(images)) if self.use_vllm and self.vllm_importance_sampling_correction: + stats = self._metric_stats[mode] + mins, maxs = self._metric_mins[mode], self._metric_maxs[mode] delta = torch.abs(old_per_token_logps - sampling_per_token_logps) mask = completion_mask.bool() if tool_mask is None else (completion_mask * tool_mask).bool() # Tokens vLLM could not score carry NaN, so exclude them rather than let them turn the reported # divergence into NaN. Counting them as zero instead would understate the divergence. - delta = delta[mask & ~torch.isnan(delta)] - mean_delta = torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) - max_delta = torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) - self._metrics[mode]["sampling/sampling_logp_difference/mean"].append( - self.accelerator.gather(mean_delta).mean().item() - ) - self._metrics[mode]["sampling/sampling_logp_difference/max"].append( - self.accelerator.gather(max_delta).max().item() + delta_valid = mask & ~torch.isnan(delta) + stats["sampling/sampling_logp_difference/mean"] += torch.stack( + [torch.where(delta_valid, delta, 0.0).sum(), delta_valid.sum().float()] ) + max_delta = torch.where(delta_valid, delta, -torch.inf).max() + key = "sampling/sampling_logp_difference/max" + maxs[key] = torch.maximum(maxs.get(key, max_delta), max_delta) if sequence_level_is: flat_is_ratio = vllm_importance_sampling_ratio.flatten() else: flat_is_ratio = vllm_importance_sampling_ratio[mask] - min_importance_sampling_ratio = ( - torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) - ) - mean_importance_sampling_ratio = ( - torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) - ) - max_importance_sampling_ratio = ( - torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) - ) - self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( - nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() - ) - self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( - self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() - ) - self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( - nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + is_ratio_valid = ~torch.isnan(flat_is_ratio) + stats["sampling/importance_sampling_ratio/mean"] += torch.stack( + [torch.where(is_ratio_valid, flat_is_ratio, 0.0).sum(), is_ratio_valid.sum().float()] ) + if flat_is_ratio.numel() > 0: + min_is_ratio = torch.where(is_ratio_valid, flat_is_ratio, torch.inf).min() + max_is_ratio = torch.where(is_ratio_valid, flat_is_ratio, -torch.inf).max() + else: + min_is_ratio = torch.tensor(torch.inf, device=device) + max_is_ratio = torch.tensor(-torch.inf, device=device) + key = "sampling/importance_sampling_ratio/min" + mins[key] = torch.minimum(mins.get(key, min_is_ratio), min_is_ratio) + key = "sampling/importance_sampling_ratio/max" + maxs[key] = torch.maximum(maxs.get(key, max_is_ratio), max_is_ratio) output = { "prompt_ids": prompt_ids, @@ -3006,8 +3007,10 @@ def compute_liger_loss(self, unwrapped_model, inputs): mode = "train" if self.model.training else "eval" if self.beta != 0.0: - self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).mean().item()) - self._metrics[mode]["clip_ratio"].append(self.accelerator.gather(clip_ratio).mean().item()) + mean_kl = mean_kl.detach() + self._metric_stats[mode]["kl"] += torch.stack([mean_kl, torch.ones_like(mean_kl)]) + clip_ratio = clip_ratio.detach() + self._metric_stats[mode]["clip_ratio"] += torch.stack([clip_ratio, torch.ones_like(clip_ratio)]) # DAPO/CISPO/VESPO normalize by num_items_in_batch / num_processes (applied internally by # the Liger loss), then need a `current_gradient_accumulation_steps / steps_per_generation` # rescale to land on the per-window token-mean — matching the non-Liger path @@ -3299,7 +3302,10 @@ def _compute_loss(self, model, inputs): loss = loss - apply_coef * entropy_loss - self._metrics[mode]["policy_loss"].append(self.accelerator.gather(policy_loss).nanmean().item()) + policy_loss_valid = ~torch.isnan(policy_loss) + self._metric_stats[mode]["policy_loss"] += torch.stack( + [torch.where(policy_loss_valid, policy_loss, 0.0).detach(), policy_loss_valid.float()] + ) # Adaptive update. Gated on train mode so evaluation cannot mutate the entropy controller state. if self.use_adaptive_entropy and mode == "train": @@ -3337,7 +3343,10 @@ def _compute_loss(self, model, inputs): if self.aux_loss_enabled: normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 loss = loss + self.router_aux_loss_coef * aux_loss / normalizer - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) + detached_aux_loss = aux_loss.detach() + self._metric_stats[mode]["aux_loss"] += torch.stack( + [detached_aux_loss, torch.ones_like(detached_aux_loss)] + ) # Log the metrics def masked_seq_mean(x): @@ -3345,36 +3354,39 @@ def masked_seq_mean(x): return x.squeeze(1) return (x * mask).sum(-1) / mask.sum(-1) - def global_masked_mean(x): + def accumulate_masked_mean(name, x): + x = x.detach() if x.shape[1] == 1: # when importance_sampling_level == "sequence": one value per sequence local_sum, local_count = x.sum(), torch.tensor(float(x.shape[0]), device=x.device) else: local_sum, local_count = (x * mask).sum(), mask.sum().float() - totals = self.accelerator.reduce(torch.stack([local_sum, local_count]), reduction="sum") - return (totals[0] / totals[1].clamp(min=1.0)).item() + self._metric_stats[mode][name] += torch.stack([local_sum, local_count]) if self.beta != 0.0: - self._metrics[mode]["kl"].append(global_masked_mean(per_token_kl)) + accumulate_masked_mean("kl", per_token_kl) - self._metrics[mode]["entropy"].append(global_masked_mean(entropies)) + accumulate_masked_mean("entropy", entropies) if self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]: # Compute the clipped probability ratios is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0) is_region_clipped = is_low_clipped | is_high_clipped - self._metrics[mode]["clip_ratio/low_mean"].append(global_masked_mean(is_low_clipped.float())) - self._metrics[mode]["clip_ratio/high_mean"].append(global_masked_mean(is_high_clipped.float())) - self._metrics[mode]["clip_ratio/region_mean"].append(global_masked_mean(is_region_clipped.float())) - gathered_low_clip = self.accelerator.gather(masked_seq_mean(is_low_clipped.float())) - self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) - gathered_high_clip = self.accelerator.gather(masked_seq_mean(is_high_clipped.float())) - self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + accumulate_masked_mean("clip_ratio/low_mean", is_low_clipped.float()) + accumulate_masked_mean("clip_ratio/high_mean", is_high_clipped.float()) + accumulate_masked_mean("clip_ratio/region_mean", is_region_clipped.float()) + mins, maxs = self._metric_mins[mode], self._metric_maxs[mode] + low_clip_seq = masked_seq_mean(is_low_clipped.float()) + low_min = torch.where(low_clip_seq.isnan(), torch.inf, low_clip_seq).min() + mins["clip_ratio/low_min"] = torch.minimum(mins.get("clip_ratio/low_min", low_min), low_min) + high_clip_seq = masked_seq_mean(is_high_clipped.float()) + high_max = torch.where(high_clip_seq.isnan(), -torch.inf, high_clip_seq).max() + maxs["clip_ratio/high_max"] = torch.maximum(maxs.get("clip_ratio/high_max", high_max), high_max) elif self.loss_type == "cispo": is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0) - self._metrics[mode]["cispo_clip_ratio"].append(global_masked_mean(is_cispo_clipped.float())) + accumulate_masked_mean("cispo_clip_ratio", is_cispo_clipped.float()) elif self.loss_type == "vespo": - self._metrics[mode]["vespo/phi_seq_mean"].append(global_masked_mean(phi_seq)) + accumulate_masked_mean("vespo/phi_seq_mean", phi_seq) return loss @@ -3400,6 +3412,24 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: valid = [v for v in val if not math.isnan(v)] metrics[key] = sum(valid) / len(valid) if valid else None + # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that + # all ranks stack them in the same order. + stats = self._metric_stats[mode] + if stats: + keys = sorted(stats) + values = torch.stack([stats[key].double() for key in keys]) + totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True)) + metrics.update({key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}) + # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A + # rank with no data contributed +/-inf sentinels; a window with no data at all logs None. + for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)): + if extrema: + keys = sorted(extrema) + values = torch.stack([sign * extrema[key].double() for key in keys]) + reduced = sign * self.accelerator.reduce(values, reduction="max") + for key, value in zip(keys, reduced.tolist(), strict=True): + metrics[key] = value if math.isfinite(value) else None + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. if mode == "eval": @@ -3408,6 +3438,9 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: logs.update(metrics) super().log(logs, start_time) self._metrics[mode].clear() + self._metric_stats[mode].clear() + self._metric_mins[mode].clear() + self._metric_maxs[mode].clear() if self.accelerator.is_main_process and self.log_completions: if is_rich_available(): diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py index 632cea2423d..519a0acf6ff 100644 --- a/trl/trainer/kto_trainer.py +++ b/trl/trainer/kto_trainer.py @@ -945,7 +945,10 @@ def __init__( disable_dropout_in_model(self.ref_model) # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + # Each entry is a running `(total, count)` pair of on-device tensors, summed in `compute_loss` with no + # collective and no host sync. `log()` reduces the pairs across ranks in a single collective and divides, + # so every metric is weighted by whatever its count counts (tokens, pairs, batches). + self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)} self._total_train_tokens = 0 # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the @@ -1449,43 +1452,23 @@ def _compute_loss_liger(self, model, inputs, return_outputs): kl=kl, ) - self._metrics[mode]["kl"].append(kl.item()) + stats = self._metric_stats[mode] + # `kl` is already global here (the gather above is required by the loss), so only the window average is left + # for `log()` to compute. + kl_total = kl.detach().sum() + stats["kl"] += torch.stack([kl_total, torch.ones_like(kl_total)]) # Number of tokens if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(batch["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + num_tokens_in_batch = batch["attention_mask"].sum() + stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)]) - all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() - all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() - - if all_num_chosen > 0: - self._metrics[mode]["rewards/chosen"].append( - self.accelerator.gather_for_metrics(chosen_rewards_sum.nansum()).nansum().item() / all_num_chosen - ) - self._metrics[mode]["logps/chosen"].append( - self.accelerator.gather_for_metrics(chosen_logps_sum.nansum()).nansum().item() / all_num_chosen - ) - self._metrics[mode]["logits/chosen"].append( - self.accelerator.gather_for_metrics(chosen_logits_sum.nansum()).nansum().item() / all_num_chosen - ) - - if all_num_rejected > 0: - self._metrics[mode]["rewards/rejected"].append( - self.accelerator.gather_for_metrics(rejected_rewards_sum.nansum()).nansum().item() / all_num_rejected - ) - self._metrics[mode]["logps/rejected"].append( - self.accelerator.gather_for_metrics(rejected_logps_sum.nansum()).nansum().item() / all_num_rejected - ) - self._metrics[mode]["logits/rejected"].append( - self.accelerator.gather_for_metrics(rejected_logits_sum.nansum()).nansum().item() / all_num_rejected - ) - - if all_num_chosen > 0 and all_num_rejected > 0: - self._metrics[mode]["rewards/margins"].append( - self._metrics[mode]["rewards/chosen"][-1] - self._metrics[mode]["rewards/rejected"][-1] - ) + stats["rewards/chosen"] += torch.stack([chosen_rewards_sum.nansum().detach(), num_chosen.float()]) + stats["logps/chosen"] += torch.stack([chosen_logps_sum.nansum().detach(), num_chosen.float()]) + stats["logits/chosen"] += torch.stack([chosen_logits_sum.nansum().detach(), num_chosen.float()]) + stats["rewards/rejected"] += torch.stack([rejected_rewards_sum.nansum().detach(), num_rejected.float()]) + stats["logps/rejected"] += torch.stack([rejected_logps_sum.nansum().detach(), num_rejected.float()]) + stats["logits/rejected"] += torch.stack([rejected_logits_sum.nansum().detach(), num_rejected.float()]) return loss @@ -1605,25 +1588,21 @@ def _compute_loss(self, model, inputs, return_outputs): 0, ) - self._metrics[mode]["kl"].append(kl.item()) + stats = self._metric_stats[mode] + # `kl` is already global here (the gather above is required by the loss), so only the window average is left + # for `log()` to compute. + kl_total = kl.detach().sum() + stats["kl"] += torch.stack([kl_total, torch.ones_like(kl_total)]) # Entropy per_token_entropy = entropy_from_logits(shift_logits.detach()) mask = batch["completion_mask"][:, 1:] - entropy_sum = (per_token_entropy * mask).sum() - total_tokens = mask.sum() - - # Gather counts across ranks and weight-average - entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() - total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum() - entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) + stats["entropy"] += torch.stack([(per_token_entropy * mask).sum(), mask.sum().float()]) # Number of tokens if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(batch["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + num_tokens_in_batch = batch["attention_mask"].sum() + stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)]) # Average logits for chosen and rejected completions shift_completion_mask = batch["completion_mask"][:, 1:] @@ -1631,48 +1610,23 @@ def _compute_loss(self, model, inputs, return_outputs): rejected_logits = shift_logits.detach().index_select(0, rejected_idx) chosen_mask = shift_completion_mask.index_select(0, chosen_idx) rejected_mask = shift_completion_mask.index_select(0, rejected_idx) - total_chosen_logits = chosen_logits[chosen_mask.bool()].mean(-1).sum() - total_chosen_tokens = chosen_mask.sum() - total_rejected_logits = rejected_logits[rejected_mask.bool()].mean(-1).sum() - total_rejected_tokens = rejected_mask.sum() - total_chosen_logits = self.accelerator.gather_for_metrics(total_chosen_logits).sum().item() - total_chosen_tokens = self.accelerator.gather_for_metrics(total_chosen_tokens).sum().item() - total_rejected_logits = self.accelerator.gather_for_metrics(total_rejected_logits).sum().item() - total_rejected_tokens = self.accelerator.gather_for_metrics(total_rejected_tokens).sum().item() - if total_chosen_tokens > 0: - self._metrics[mode]["logits/chosen"].append(total_chosen_logits / total_chosen_tokens) - if total_rejected_tokens > 0: - self._metrics[mode]["logits/rejected"].append(total_rejected_logits / total_rejected_tokens) - - all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() - all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() - - if all_num_chosen > 0: - self._metrics[mode]["rewards/chosen"].append( - self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() / all_num_chosen - ) - self._metrics[mode]["logps/chosen"].append( - self.accelerator.gather_for_metrics(chosen_logps.nansum()).nansum().item() / all_num_chosen - ) - - if all_num_rejected > 0: - self._metrics[mode]["rewards/rejected"].append( - self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() / all_num_rejected - ) - self._metrics[mode]["logps/rejected"].append( - self.accelerator.gather_for_metrics(rejected_logps.nansum()).nansum().item() / all_num_rejected - ) + stats["logits/chosen"] += torch.stack( + [chosen_logits[chosen_mask.bool()].mean(-1).sum(), chosen_mask.sum().float()] + ) + stats["logits/rejected"] += torch.stack( + [rejected_logits[rejected_mask.bool()].mean(-1).sum(), rejected_mask.sum().float()] + ) - if all_num_chosen > 0 and all_num_rejected > 0: - self._metrics[mode]["rewards/margins"].append( - self._metrics[mode]["rewards/chosen"][-1] - self._metrics[mode]["rewards/rejected"][-1] - ) + stats["rewards/chosen"] += torch.stack([chosen_rewards.nansum(), num_chosen.float()]) + stats["logps/chosen"] += torch.stack([chosen_logps.detach().nansum(), num_chosen.float()]) + stats["rewards/rejected"] += torch.stack([rejected_rewards.nansum(), num_rejected.float()]) + stats["logps/rejected"] += torch.stack([rejected_logps.detach().nansum(), num_rejected.float()]) loss = losses.nanmean() if self.aux_loss_enabled: aux_loss = outputs.aux_loss loss = loss + self.router_aux_loss_coef * aux_loss - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) + stats["aux_loss"] += torch.stack([aux_loss.detach(), torch.ones_like(aux_loss)]) return (loss, outputs) if return_outputs else loss @@ -1775,14 +1729,33 @@ def training_step(self, *args, **kwargs): def log(self, logs: dict[str, float], start_time: float | None = None) -> None: mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that + # all ranks stack them in the same order. + metrics = {} + stats = self._metric_stats[mode] + if stats: + keys = sorted(stats) + values = torch.stack([stats[key].double() for key in keys]) + totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True)) + metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()} + # KTO sees chosen and rejected completions in separate, independently sized groups, so the margin is a + # difference of the two window means rather than a metric with a count of its own. + if totals["rewards/chosen"][1] > 0 and totals["rewards/rejected"][1] > 0: + metrics["rewards/margins"] = metrics["rewards/chosen"] - metrics["rewards/rejected"] + # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances + # on a train-mode log, so an eval log in between can lag by up to one logging window. + if mode == "train" and "num_tokens" in totals: + self._total_train_tokens += int(totals["num_tokens"][0]) + metrics["num_tokens"] = self._total_train_tokens + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. if mode == "eval": metrics = {f"eval_{key}": val for key, val in metrics.items()} logs.update(metrics) super().log(logs, start_time) - self._metrics[mode].clear() + self._metric_stats[mode].clear() # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py index c1ac2977f3d..31b3d12dc4c 100644 --- a/trl/trainer/reward_trainer.py +++ b/trl/trainer/reward_trainer.py @@ -15,6 +15,7 @@ import contextlib import json import logging +import math import os import re import warnings @@ -601,8 +602,12 @@ def __init__( else: self.maybe_activation_offload_context = contextlib.nullcontext() - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + # `_metric_stats` entries are running `(total, count)` pairs and `_metric_mins`/`_metric_maxs` are running + # extrema, all on-device: `compute_loss` only adds local tensors here (no collective, no host sync), and + # `log()` aggregates them across ranks in one collective per kind. + self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)} + self._metric_mins = {"train": {}, "eval": {}} + self._metric_maxs = {"train": {}, "eval": {}} self._total_train_tokens = 0 # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the @@ -699,8 +704,10 @@ def tokenize_fn(example, processing_class): if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` map_kwargs["desc"] = f"Filtering {dataset_name} >{args.max_length} tokens" dataset = dataset.filter( - lambda example: len(example["chosen_ids"]) <= args.max_length - and len(example["rejected_ids"]) <= args.max_length, + lambda example: ( + len(example["chosen_ids"]) <= args.max_length + and len(example["rejected_ids"]) <= args.max_length + ), **map_kwargs, ) @@ -760,24 +767,28 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2) if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + num_tokens_in_batch = inputs["attention_mask"].sum() + self._metric_stats[mode]["num_tokens"] += torch.stack( + [num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)] + ) # Compute min, mean, max, accuracy and margin with torch.no_grad(): - all_rewards = self.accelerator.gather(outputs.logits) - self._metrics[mode]["min_reward"].append(all_rewards.min().item()) - self._metrics[mode]["mean_reward"].append(all_rewards.mean().item()) - self._metrics[mode]["max_reward"].append(all_rewards.max().item()) + stats = self._metric_stats[mode] + mins, maxs = self._metric_mins[mode], self._metric_maxs[mode] + rewards = outputs.logits.detach() + num_pairs = torch.tensor(float(rewards_chosen.numel()), device=rewards.device) + stats["mean_reward"] += torch.stack( + [rewards.sum(), torch.tensor(float(rewards.numel()), device=rewards.device)] + ) + min_reward = rewards.min() + mins["min_reward"] = torch.minimum(mins.get("min_reward", min_reward), min_reward) + max_reward = rewards.max() + maxs["max_reward"] = torch.maximum(maxs.get("max_reward", max_reward), max_reward) - mean_accuracy = (rewards_chosen > rewards_rejected).float().mean() - mean_accuracy = self.accelerator.gather_for_metrics(mean_accuracy).mean().item() - self._metrics[mode]["accuracy"].append(mean_accuracy) + stats["accuracy"] += torch.stack([(rewards_chosen > rewards_rejected).float().sum(), num_pairs]) - mean_margin = (rewards_chosen - rewards_rejected).mean() - mean_margin = self.accelerator.gather_for_metrics(mean_margin).mean() - self._metrics[mode]["margin"].append(mean_margin.item()) + stats["margin"] += torch.stack([(rewards_chosen - rewards_rejected).sum(), num_pairs]) return (loss, outputs) if return_outputs else loss @@ -788,7 +799,30 @@ def training_step(self, *args, **kwargs): def log(self, logs: dict[str, float], start_time: float | None = None) -> None: mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that + # all ranks stack them in the same order. + metrics = {} + stats = self._metric_stats[mode] + if stats: + keys = sorted(stats) + values = torch.stack([stats[key].double() for key in keys]) + totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True)) + metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()} + # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances + # on a train-mode log, so an eval log in between can lag by up to one logging window. + if mode == "train" and "num_tokens" in totals: + self._total_train_tokens += int(totals["num_tokens"][0]) + metrics["num_tokens"] = self._total_train_tokens + # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A + # rank with no data contributed +/-inf sentinels; a window with no data at all logs None. + for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)): + if extrema: + keys = sorted(extrema) + values = torch.stack([sign * extrema[key].double() for key in keys]) + reduced = sign * self.accelerator.reduce(values, reduction="max") + for key, value in zip(keys, reduced.tolist(), strict=True): + metrics[key] = value if math.isfinite(value) else None # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. @@ -797,7 +831,9 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: logs.update(metrics) super().log(logs, start_time) - self._metrics[mode].clear() + self._metric_stats[mode].clear() + self._metric_mins[mode].clear() + self._metric_maxs[mode].clear() # Ensure the model card is saved along with the checkpoint def _save_checkpoint(self, model, trial): diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py index 3c9de79d254..a532d697495 100644 --- a/trl/trainer/rloo_trainer.py +++ b/trl/trainer/rloo_trainer.py @@ -67,8 +67,6 @@ get_callable_name, get_config_model_id, identity, - nanmax, - nanmin, nanstd, pad, print_prompt_completions_sample, @@ -626,6 +624,13 @@ def __init__( # Initialize the metrics self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + # `_metric_stats` entries are running `(total, count)` pairs and `_metric_mins`/`_metric_maxs` are running + # extrema, all on-device: they only receive local tensors (no collective, no host sync), and `log()` + # aggregates them across ranks in one collective per kind. `_metrics` above keeps plain floats for values + # that are already identical on every rank. + self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)} + self._metric_mins = {"train": {}, "eval": {}} + self._metric_maxs = {"train": {}, "eval": {}} self._total_train_tokens = 0 self._current_train_step_time = 0.0 self.log_completions = args.log_completions @@ -1617,10 +1622,9 @@ def _generate_and_score_completions( # Calculate and log the mean KL divergence between current and reference model if self.beta != 0.0: - kl_stats = self.accelerator.reduce( - torch.stack([(per_token_kl * completion_mask).sum(), completion_mask.sum().float()]), reduction="sum" + self._metric_stats[mode]["kl"] += torch.stack( + [(per_token_kl * completion_mask).sum().detach(), completion_mask.sum().float()] ) - self._metrics[mode]["kl"].append((kl_stats[0] / kl_stats[1].clamp(min=1.0)).item()) # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) for i, reward_func_name in enumerate(self.reward_func_names): @@ -1648,14 +1652,14 @@ def _generate_and_score_completions( self._logs["extra"][column].extend(gather_object(self._pending_extra_logs[column])) self._pending_extra_logs.clear() - # Flush user-logged metrics (from log_metric), averaging across processes. - # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values - # get mis-attributed across metrics (dict insertion order may differ between processes). - for name in sorted(self._pending_metrics): + # Flush user-logged metrics (from log_metric), accumulated locally and averaged across processes at `log()` + # time. Every rank must log the same metric names, otherwise the log-time aggregation mismatches ranks + # (the same requirement the previous per-step gather had). + for name in self._pending_metrics: values = self._pending_metrics[name] local_mean = sum(values) / len(values) - global_mean = self.accelerator.gather(torch.tensor(local_mean, device=device)).mean().item() - self._metrics[mode][name].append(global_mean) + local_mean = torch.tensor(local_mean, device=device) + self._metric_stats[mode][name] += torch.stack([local_mean, torch.ones_like(local_mean)]) self._pending_metrics.clear() if images is not None and self.log_multimodal: @@ -1744,26 +1748,30 @@ def _compute_loss(self, model, inputs): # RLOO returns an unscaled loss (the HF Trainer divides by gradient accumulation), so add the aux term unscaled if self.aux_loss_enabled: loss = loss + self.router_aux_loss_coef * aux_loss - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) + detached_aux_loss = aux_loss.detach() + self._metric_stats[mode]["aux_loss"] += torch.stack( + [detached_aux_loss, torch.ones_like(detached_aux_loss)] + ) # Entropy - entropy_stats = self.accelerator.reduce( - torch.stack([(entropies * completion_mask).sum(), completion_mask.sum().float()]), reduction="sum" + self._metric_stats[mode]["entropy"] += torch.stack( + [(entropies * completion_mask).sum().detach(), completion_mask.sum().float()] ) - self._metrics[mode]["entropy"].append((entropy_stats[0] / entropy_stats[1].clamp(min=1.0)).item()) # Compute the clipped probability ratios + stats = self._metric_stats[mode] + mins, maxs = self._metric_mins[mode], self._metric_maxs[mode] is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0) is_region_clipped = is_low_clipped | is_high_clipped - gathered_low_clip = self.accelerator.gather(is_low_clipped.float()) - self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) - self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) - gathered_high_clip = self.accelerator.gather(is_high_clipped.float()) - self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) - self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) - gathered_clip_ratio = self.accelerator.gather(is_region_clipped.float()) - self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + clip_count = torch.tensor(float(is_low_clipped.numel()), device=is_low_clipped.device) + stats["clip_ratio/low_mean"] += torch.stack([is_low_clipped.float().sum(), clip_count]) + low_min = is_low_clipped.float().min() + mins["clip_ratio/low_min"] = torch.minimum(mins.get("clip_ratio/low_min", low_min), low_min) + stats["clip_ratio/high_mean"] += torch.stack([is_high_clipped.float().sum(), clip_count]) + high_max = is_high_clipped.float().max() + maxs["clip_ratio/high_max"] = torch.maximum(maxs.get("clip_ratio/high_max", high_max), high_max) + stats["clip_ratio/region_mean"] += torch.stack([is_region_clipped.float().sum(), clip_count]) return loss # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and @@ -1788,6 +1796,24 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: valid = [v for v in val if not math.isnan(v)] metrics[key] = sum(valid) / len(valid) if valid else None + # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that + # all ranks stack them in the same order. + stats = self._metric_stats[mode] + if stats: + keys = sorted(stats) + values = torch.stack([stats[key].double() for key in keys]) + totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True)) + metrics.update({key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}) + # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A + # rank with no data contributed +/-inf sentinels; a window with no data at all logs None. + for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)): + if extrema: + keys = sorted(extrema) + values = torch.stack([sign * extrema[key].double() for key in keys]) + reduced = sign * self.accelerator.reduce(values, reduction="max") + for key, value in zip(keys, reduced.tolist(), strict=True): + metrics[key] = value if math.isfinite(value) else None + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. if mode == "eval": @@ -1796,6 +1822,9 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: logs.update(metrics) super().log(logs, start_time) self._metrics[mode].clear() + self._metric_stats[mode].clear() + self._metric_mins[mode].clear() + self._metric_maxs[mode].clear() if self.accelerator.is_main_process and self.log_completions: if is_rich_available(): diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py index cab3f550c79..0770bbf076a 100644 --- a/trl/trainer/sft_trainer.py +++ b/trl/trainer/sft_trainer.py @@ -1414,8 +1414,10 @@ def __init__( text_config.output_router_logits = self.aux_loss_enabled text_config.router_aux_loss_coef = self.args.router_aux_loss_coef - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + # Initialize the metrics. Each entry is a running `(total, count)` pair of on-device tensors, summed in + # `compute_loss` with no collective and no host sync. `log()` reduces the pairs across ranks in a single + # collective and divides, so every metric is weighted by whatever its count counts (tokens, batches). + self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)} self._total_train_tokens = 0 # Add tags to the model @@ -1804,10 +1806,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N # PEFT (PromptTuning, P-Tuning) prepends `-100`-padded virtual tokens before delegating into the patched # forward, so the valid-token count over the padded labels can differ from the un-padded `labels[..., 1:]` # count by up to one per sequence; using the patched output keeps numerator and denominator aligned. - num_valid = self.accelerator.gather_for_metrics(outputs.num_valid_tokens).sum() - entropy_sum = self.accelerator.gather_for_metrics(outputs.entropy_sum).sum() - entropy = (entropy_sum / num_valid).item() if num_valid > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) + num_valid = outputs.num_valid_tokens.detach().float() + self._metric_stats[mode]["entropy"] += torch.stack([outputs.entropy_sum.detach(), num_valid]) elif not self.args.use_liger_kernel: # liger doesn't return logits with torch.no_grad(): if "shift_labels" in inputs: @@ -1833,39 +1833,38 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N total_tokens = mask.sum() correct_predictions = (predictions == shift_labels) & mask correct_tokens = correct_predictions.sum() - - # Gather counts across ranks and weight-average - entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() - total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum() - correct_tokens = self.accelerator.gather_for_metrics(correct_tokens) - entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0 - - total_sum = total_tokens.sum() - accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) - self._metrics[mode]["mean_token_accuracy"].append(accuracy) + self._metric_stats[mode]["entropy"] += torch.stack([entropy_sum, total_tokens.float()]) + self._metric_stats[mode]["mean_token_accuracy"] += torch.stack( + [correct_tokens.float(), total_tokens.float()] + ) if mode == "train": # When using padding-free, the attention_mask is not present in the inputs, instead we have cu_seq_lens_q, # cu_seq_lens_k, and max_length_k, max_length_q and position_ids. if "attention_mask" in inputs: - num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() + num_tokens_in_batch = inputs["attention_mask"].sum() elif "position_ids" in inputs: - local_num_tokens = torch.tensor(inputs["position_ids"].size(1), device=inputs["position_ids"].device) - num_tokens_in_batch = self.accelerator.gather_for_metrics(local_num_tokens).sum().item() + num_tokens_in_batch = torch.tensor( + inputs["position_ids"].size(1), device=inputs["position_ids"].device + ) else: raise ValueError("Expected 'attention_mask' or 'position_ids' in inputs.") - self._total_train_tokens += num_tokens_in_batch - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + # `num_tokens` is a running total rather than a window average, so `log()` uses the pair's total and + # ignores its count. + self._metric_stats[mode]["num_tokens"] += torch.stack( + [num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)] + ) if self.args.loss_type == "chunked_nll": - correct = self.accelerator.gather_for_metrics(outputs.num_correct_tokens).sum() - accuracy = (correct / num_valid).item() if num_valid > 0 else 0.0 - self._metrics[mode]["mean_token_accuracy"].append(accuracy) + self._metric_stats[mode]["mean_token_accuracy"] += torch.stack( + [outputs.num_correct_tokens.detach(), num_valid] + ) elif self.args.use_liger_kernel: if hasattr(outputs, "token_accuracy") and outputs.token_accuracy is not None: - token_accuracy = self.accelerator.gather_for_metrics(outputs.token_accuracy).mean().item() - self._metrics[mode]["mean_token_accuracy"].append(token_accuracy) + token_accuracy = outputs.token_accuracy.detach() + self._metric_stats[mode]["mean_token_accuracy"] += torch.stack( + [token_accuracy, torch.ones_like(token_accuracy)] + ) else: warnings.warn( "liger-kernel did not return token_accuracy when requested. The mean_token_accuracy metric will " @@ -1874,9 +1873,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N ) # Log auxiliary loss if enabled (applies to both Liger and non-Liger) if self.aux_loss_enabled: - aux_loss = outputs.aux_loss - aux_loss = self.accelerator.gather_for_metrics(aux_loss).mean().item() - self._metrics[mode]["aux_loss"].append(aux_loss) + aux_loss = outputs.aux_loss.detach() + self._metric_stats[mode]["aux_loss"] += torch.stack([aux_loss, torch.ones_like(aux_loss)]) return (loss, outputs) if return_outputs else loss @@ -1892,7 +1890,21 @@ def training_step(self, *args, **kwargs): def log(self, logs: dict[str, float], start_time: float | None = None) -> None: mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that + # all ranks stack them in the same order. + metrics = {} + stats = self._metric_stats[mode] + if stats: + keys = sorted(stats) + values = torch.stack([stats[key].double() for key in keys]) + totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True)) + metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()} + # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances + # on a train-mode log, so an eval log in between can lag by up to one logging window. + if mode == "train" and "num_tokens" in totals: + self._total_train_tokens += int(totals["num_tokens"][0]) + metrics["num_tokens"] = self._total_train_tokens # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. @@ -1901,7 +1913,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: logs.update(metrics) super().log(logs, start_time) - self._metrics[mode].clear() + self._metric_stats[mode].clear() # Ensure the model card is saved along with the checkpoint def _save_checkpoint(self, model, trial):