diff --git a/swift/rlhf_trainers/grpo_trainer.py b/swift/rlhf_trainers/grpo_trainer.py index 42113fdbce..4e3f84d392 100644 --- a/swift/rlhf_trainers/grpo_trainer.py +++ b/swift/rlhf_trainers/grpo_trainer.py @@ -1186,7 +1186,13 @@ def masked_batch_mean(x): # Add rollout correction metrics if rollout_correction_metrics: - metrics_data['rollout_correction'] = rollout_correction_metrics + # Gather one row per rank to preserve gather_for_metrics' dataloader remainder handling. + values = torch.stack(list(rollout_correction_metrics.values())).detach().unsqueeze(0) + gathered = self.accelerator.gather_for_metrics(values) + metrics_data['rollout_correction'] = dict(zip(rollout_correction_metrics, gathered.nanmean(0).unbind())) + for key, reduce in [('log_ppl_diff_max', torch.max), ('log_ppl_diff_min', torch.min)]: + index = list(rollout_correction_metrics).index(key) + metrics_data['rollout_correction'][key] = reduce(gathered[:, index]) # Compute the clipped probability ratios if self.loss_type == 'cispo': @@ -1844,7 +1850,14 @@ def offload_context(self): def log(self, logs: Dict[str, float], start_time: Optional[float] = 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 + rollout_metrics = { + key: val + for key, val in self._metrics[mode].items() if key.startswith('rollout_correction/') + } + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items() if key not in rollout_metrics} + if rollout_metrics: + values = torch.stack([torch.stack(val) for val in rollout_metrics.values()]) + metrics.update(zip(rollout_metrics, values.cpu().double().mean(1).tolist())) # 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. @@ -2414,7 +2427,7 @@ def _compute_rollout_offpolicy_metrics( per_token_logps: torch.Tensor, rollout_per_token_logps: torch.Tensor, completion_mask: torch.Tensor, - ) -> Dict[str, float]: + ) -> Dict[str, torch.Tensor]: """ Compute off-policy diagnostic metrics (always computed for monitoring). @@ -2454,11 +2467,10 @@ def masked_mean(x, mask, axis=None): # Formula: exp(-1/|T| * Σ log π_training(y_t|y_ Dict[str, float]: + ) -> Dict[str, torch.Tensor]: """ Compute importance sampling correction metrics (ess, clipped_frac, is_weight_mean). Only called when rollout_importance_sampling_mode is enabled. @@ -2565,7 +2576,7 @@ def masked_mean(x, mask): # 1. IS weight statistics mean_is_weight = masked_mean(is_weights, completion_mask) - metrics['is_weight_mean'] = self.accelerator.gather_for_metrics(mean_is_weight).nanmean().item() + metrics['is_weight_mean'] = mean_is_weight # 2. Compute Effective Sample Size (ESS) for IS weights # ESS = 1 / E[(w_i / E[w_i])²] (using clamped weights for stability) @@ -2574,7 +2585,7 @@ def masked_mean(x, mask): mean_for_ess = masked_mean(weights_for_ess, completion_mask) is_weights_normalized = weights_for_ess / (mean_for_ess + 1e-8) # Avoid division by zero ess = 1.0 / masked_mean(is_weights_normalized.square(), completion_mask).clamp(min=1e-10) - metrics['ess'] = self.accelerator.gather_for_metrics(ess).nanmean().item() + metrics['ess'] = ess # 3. Fraction of clipped/masked samples if self.rollout_importance_sampling_mode in ['token_truncate', 'token_mask']: @@ -2583,12 +2594,12 @@ def masked_mean(x, mask): clipped_frac = masked_mean((is_ratio > threshold).float(), completion_mask) else: # token_mask clipped_frac = masked_mean((is_weights == 0).float(), completion_mask) - metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item() + metrics['clipped_frac'] = clipped_frac else: # Sequence-level (both truncate and mask) seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask) clipped_frac = (seq_ratios > threshold).float().mean() - metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item() + metrics['clipped_frac'] = clipped_frac return metrics