From f99cb5f4c6cba6af043a6a711b1e8907c5b6a7c1 Mon Sep 17 00:00:00 2001 From: casuallkk <2977808892@qq.com> Date: Tue, 1 Sep 2026 23:32:48 +0800 Subject: [PATCH] feat(gkd): add agentic monitoring metrics --- swift/megatron/trainers/gkd_trainer.py | 55 ++++++++++-- swift/megatron/trainers/rollout_mixin.py | 8 +- swift/ray/megatron/gkd_trainer.py | 35 ++++++-- swift/ray/megatron/loss/gkd.py | 45 +++++++++- swift/ray/megatron/megatron_worker.py | 4 +- swift/rlhf_trainers/gkd_helpers.py | 35 +++++++- swift/rlhf_trainers/gkd_loss.py | 103 ++++++++++++++++++++++- swift/rlhf_trainers/gkd_trainer.py | 52 ++++++++++-- swift/rlhf_trainers/rollout_mixin.py | 8 +- swift/rlhf_trainers/utils.py | 19 +++-- tests/train/test_gkd_monitoring.py | 101 ++++++++++++++++++++++ 11 files changed, 425 insertions(+), 40 deletions(-) create mode 100644 tests/train/test_gkd_monitoring.py diff --git a/swift/megatron/trainers/gkd_trainer.py b/swift/megatron/trainers/gkd_trainer.py index 1d558ed76f..2fdb09fb40 100644 --- a/swift/megatron/trainers/gkd_trainer.py +++ b/swift/megatron/trainers/gkd_trainer.py @@ -15,7 +15,7 @@ from swift.rl_core.resample import resample_encode_failed_inputs from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_output, build_opsd_samples, build_teacher_requests, encode_gkd_samples, fetch_teacher_parsed_by_routing) -from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss +from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss, gkd_monitoring_stats from swift.template import Template from swift.utils import get_logger, to_device from ..utils import forward_step_helper, get_padding_to @@ -259,11 +259,23 @@ def _compute_teacher_logits_local(self, encoded_batches: List[Dict], vp_stage: O if teacher_logits is not None: teacher_logits = teacher_logits.detach() + target_logprobs = None + if (teacher_logits is not None and teacher_labels is not None + and encoded_batch.get('data_source') == DataSource.STUDENT): + safe_labels = teacher_labels.masked_fill(teacher_labels == -100, 0).long() + teacher_logprobs = vocab_parallel_log_softmax(teacher_logits.float()) + target_logprobs = tp_gather_topk(teacher_logprobs, safe_labels.unsqueeze(-1)).squeeze(-1) + target_logprobs = target_logprobs.masked_fill(teacher_labels == -100, float('nan')) + if topk is not None and teacher_logits is not None: topk_logits, topk_indices = vocab_parallel_topk(teacher_logits, k=topk) - teacher_out = TeacherOutput(topk_logprobs=topk_logits, topk_indices=topk_indices) + teacher_out = TeacherOutput( + topk_logprobs=topk_logits, + topk_indices=topk_indices, + target_logprobs=target_logprobs, + ) else: - teacher_out = TeacherOutput(full_logits=teacher_logits) + teacher_out = TeacherOutput(full_logits=teacher_logits, target_logprobs=target_logprobs) teacher_out.labels = teacher_labels encoded_batch['teacher_output'] = teacher_out @@ -286,6 +298,11 @@ def _generate_and_score_completions(self, inputs: List[Dict]) -> List[Dict]: samples = self._gather_rollout_results(local_batch) self._log_completions_from_samples(samples) + num_turns_mean = None + if data_source == DataSource.STUDENT and samples and all(s.rollout_infos and 'num_turns' in s.rollout_infos + for s in samples): + num_turns_mean = sum(float(s.rollout_infos['num_turns']) for s in samples) / len(samples) + # Teacher API: build requests from samples, fetch logprobs local_parsed = None if self.use_teacher_api: @@ -299,7 +316,7 @@ def _generate_and_score_completions(self, inputs: List[Dict]) -> List[Dict]: self.teacher_clients, gather_fn=self._gather_teacher_requests, infer_fn=lambda handle, client: self._infer_teacher_requests( - handle, topk=self.gkd_logits_topk, teacher_client=client), + handle, topk=self.gkd_logits_topk, teacher_client=client, include_sampled=True), scatter_fn=self._scatter_teacher_parsed, is_main_process=self.is_main_process, tag_key=self.args.teacher_tag_key) @@ -314,6 +331,8 @@ def _generate_and_score_completions(self, inputs: List[Dict]) -> List[Dict]: sample_slice = samples[start_idx:end_idx] encoded_batch = self._encode_samples(sample_slice) encoded_batch['data_source'] = data_source + if num_turns_mean is not None: + encoded_batch['num_turns'] = num_turns_mean if local_parsed is not None: encoded_batch['_teacher_parsed'] = local_parsed[start_idx:end_idx] all_encoded_batches.append(encoded_batch) @@ -346,7 +365,8 @@ def loss_func(self, *, labels: torch.Tensor, teacher_output: TeacherOutput, - data_source: DataSource = DataSource.DATASET): + data_source: DataSource = DataSource.DATASET, + num_turns: Optional[float] = None): """Compute GKD loss (JSD + optional SFT loss).""" student_logits = output_tensor @@ -388,6 +408,29 @@ def loss_func(self, loss = loss + self.sft_alpha * sft_loss metric = {'loss': loss.detach().clone()} + if num_turns is not None: + metric['num_turns'] = loss.new_tensor(num_turns) + if data_source == DataSource.STUDENT: + monitor = gkd_monitoring_stats( + student_logits, + teacher_output, + labels, + full_vocab_topk=self.gkd_logits_topk or 16, + student_topk_fn=vocab_parallel_topk, + teacher_topk_fn=vocab_parallel_topk, + gather_fn=tp_gather_topk, + target_logprob_fn=lambda logits, target_ids: tp_gather_topk( + vocab_parallel_log_softmax(logits.float()), target_ids.unsqueeze(-1)).squeeze(-1)) + packed = torch.stack([ + monitor['topk_overlap_sum'], monitor['topk_overlap_count'], monitor['teacher_student_gap_sum'], + monitor['teacher_student_gap_count'] + ]) + if self.args.context_parallel_size > 1: + torch.distributed.all_reduce( + packed, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group()) + torch.distributed.all_reduce(packed, op=torch.distributed.ReduceOp.SUM, group=mpu.get_data_parallel_group()) + metric['gkd/topk_overlap'] = packed[0] / packed[1].clamp(min=1) + metric['gkd/teacher_student_gap'] = packed[2] / packed[3].clamp(min=1) if sft_loss is not None: metric['jsd_loss'] = jsd_loss_val.detach().clone() metric['sft_loss'] = sft_loss.detach().clone() @@ -408,6 +451,7 @@ def forward_step(self, data_iterator, model): data = next(data_iterator) data_source = data.pop('data_source', DataSource.DATASET) + num_turns = data.pop('num_turns', None) teacher_output = data.pop('teacher_output') data.pop('teacher_model_inputs', None) # consumed by _compute_teacher_logits; not needed for student forward data = self._prepare_batch(data, vp_stage) @@ -424,4 +468,5 @@ def forward_step(self, data_iterator, model): labels=labels, teacher_output=teacher_output, data_source=data_source, + num_turns=num_turns, ) diff --git a/swift/megatron/trainers/rollout_mixin.py b/swift/megatron/trainers/rollout_mixin.py index f821d00299..bf157a5e68 100644 --- a/swift/megatron/trainers/rollout_mixin.py +++ b/swift/megatron/trainers/rollout_mixin.py @@ -277,7 +277,11 @@ def _gather_teacher_requests(self, requests: List[RolloutInferRequest]) -> Dict[ flat_global = [req for dp in dp_ranks_sorted for req in segments_by_dp[dp]] return {'flat_global': flat_global, 'offset': offset, 'n_local': len(requests)} - def _infer_teacher_requests(self, handle: Dict[str, Any], topk: int, teacher_client: Optional[Any] = None): + def _infer_teacher_requests(self, + handle: Dict[str, Any], + topk: int, + teacher_client: Optional[Any] = None, + include_sampled: bool = False): """Phase 2 (main process only, no collective): run the teacher HTTP infer. Safe to call concurrently across teachers (distinct clients, no collective inside). @@ -287,7 +291,7 @@ def _infer_teacher_requests(self, handle: Dict[str, Any], topk: int, teacher_cli client = teacher_client if teacher_client is not None else self.teacher_clients[0] request_config = RequestConfig(prompt_logprobs=topk, max_tokens=1, temperature=0.0) responses = client.infer(handle['flat_global'], request_config=request_config, use_tqdm=False) - return [parse_prompt_logprobs(r, topk=topk) for r in responses] + return [parse_prompt_logprobs(r, topk=topk, include_sampled=include_sampled) for r in responses] def _scatter_teacher_parsed(self, handle: Dict[str, Any], parsed_global): """Phase 3 (all ranks, collective): broadcast the parsed result and slice this rank's part.""" diff --git a/swift/ray/megatron/gkd_trainer.py b/swift/ray/megatron/gkd_trainer.py index 2d55c5fc74..50df3be129 100644 --- a/swift/ray/megatron/gkd_trainer.py +++ b/swift/ray/megatron/gkd_trainer.py @@ -132,6 +132,10 @@ def _train_loop(self, tg, train_iters, iteration): chunk = source_items[step_idx * chunk_size:(step_idx + 1) * chunk_size] if not chunk: break + num_turns_mean = None + if data_source == DataSource.STUDENT and all(s.rollout_infos and 'num_turns' in s.rollout_infos + for s in chunk): + num_turns_mean = sum(float(s.rollout_infos['num_turns']) for s in chunk) / len(chunk) samples = self._encode_rollout_batch(chunk) use_colocated_teacher = self._teacher_use_disable_adapter or (self._teacher_model_dir @@ -144,7 +148,8 @@ def _train_loop(self, tg, train_iters, iteration): # Driver collates the student (and, for the colocated path, the teacher view) # micro-batches; the worker only runs prepare_batch (PP/CP slice) + forward. - dispatch = self._collate_for_workers_gkd(tg, samples, data_source, with_teacher=use_colocated_teacher) + dispatch = self._collate_for_workers_gkd( + tg, samples, data_source, with_teacher=use_colocated_teacher, num_turns=num_turns_mean) if use_colocated_teacher: # Teacher forwards on the worker (CP slicing keeps each rank's shard # aligned) and caches per-micro-batch; train_step attaches the cache. @@ -249,7 +254,7 @@ def _encode_rollout_batch(self, samples: List[GKDSample]): result.append(payload) return result - def _collate_for_workers_gkd(self, tg, samples: List[dict], data_source, *, with_teacher: bool): + def _collate_for_workers_gkd(self, tg, samples: List[dict], data_source, *, with_teacher: bool, num_turns=None): """Driver-side GKD collate: ``List[payload-dict]`` -> ``{dp_rank: [model_inputs]}``. Mirrors the non-Ray GKD ``_encode_samples`` (data_collator on the rank, teacher @@ -281,6 +286,8 @@ def _collate_for_workers_gkd(self, tg, samples: List[dict], data_source, *, with chunk = shard[i:i + mbs] model_inputs = template.data_collator([s['encoded'] for s in chunk], padding_to=padding_to) model_inputs['data_source'] = data_source + if num_turns is not None: + model_inputs['num_turns'] = num_turns if with_teacher: has_opsd = chunk[0].get('teacher_encoded') is not None key = 'teacher_encoded' if has_opsd else 'encoded' @@ -344,7 +351,7 @@ def _fetch_teacher_from_replicas(self, gkd_samples: List[GKDSample], samples): responses.extend(p) for sample, response, t_encoded in zip(samples, responses, teacher_encodeds): - parsed = parse_prompt_logprobs(response, topk=topk) + parsed = parse_prompt_logprobs(response, topk=topk, include_sampled=True) encoded = t_encoded if t_encoded is not None else sample['encoded'] teacher_labels = t_encoded.get('labels') if t_encoded is not None else None sample['teacher_output'] = self._build_per_sample_teacher_output(parsed, encoded, topk, teacher_labels) @@ -367,15 +374,27 @@ def _build_per_sample_teacher_output(parsed, encoded, topk, labels=None): parsed_len = len(lps) topk_logprobs = torch.full((seq_len, topk), float('-inf'), dtype=torch.float32) topk_indices = torch.zeros(seq_len, topk, dtype=torch.long) + target_logprobs = torch.full((seq_len, ), float('nan'), dtype=torch.float32) length = min(parsed_len, seq_len) if length > 0: - topk_logprobs[:length] = torch.tensor(lps[:length], dtype=torch.float32) - topk_indices[:length] = torch.tensor(ixs[:length], dtype=torch.long) - - kwargs = dict(topk_logprobs=topk_logprobs.unsqueeze(0), topk_indices=topk_indices.unsqueeze(0)) + topk_logprobs[:length] = torch.tensor([row[:topk] for row in lps[:length]], dtype=torch.float32) + topk_indices[:length] = torch.tensor([row[:topk] for row in ixs[:length]], dtype=torch.long) + flat_input_ids = input_ids if isinstance(input_ids, list) else input_ids.reshape(-1).tolist() + for pos in range(min(length, seq_len - 1)): + target_id = int(flat_input_ids[pos + 1]) + for lp, token_id in zip(lps[pos], ixs[pos]): + if int(token_id) == target_id: + target_logprobs[pos] = float(lp) + break + + kwargs = dict( + topk_logprobs=topk_logprobs.unsqueeze(0), + topk_indices=topk_indices.unsqueeze(0), + target_logprobs=target_logprobs.unsqueeze(0)) if labels is not None: t_labels = labels if not isinstance(t_labels, torch.Tensor): t_labels = torch.tensor(t_labels, dtype=torch.long) - kwargs['labels'] = t_labels.unsqueeze(0) if t_labels.dim() == 1 else t_labels + t_labels = t_labels.unsqueeze(0) if t_labels.dim() == 1 else t_labels + kwargs['labels'] = torch.roll(t_labels, shifts=-1, dims=-1) return TeacherOutput(**kwargs) diff --git a/swift/ray/megatron/loss/gkd.py b/swift/ray/megatron/loss/gkd.py index c0254a875e..e0cb59a4c0 100644 --- a/swift/ray/megatron/loss/gkd.py +++ b/swift/ray/megatron/loss/gkd.py @@ -11,7 +11,7 @@ from swift.megatron.trainers.utils import prepare_batch from swift.megatron.trainers.vocab_parallel_utils import vocab_parallel_kl_div, vocab_parallel_log_softmax from swift.megatron.utils import forward_step_helper -from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss +from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss, gkd_monitoring_stats from swift.utils import get_current_device, to_device from .base import Loss @@ -29,6 +29,7 @@ def forward_step(self, data_iterator, model): data = next(data_iterator) teacher_output = data.pop('teacher_output', TeacherOutput()) data_source = data.pop('data_source', None) + num_turns = data.pop('num_turns', None) data.pop('grpo_batch', None) # RL signals packed in GRPOBatch (not used by GKD loss) data = prepare_batch(self.args, data) @@ -43,6 +44,7 @@ def forward_step(self, data_iterator, model): labels=labels, teacher_output=teacher_output, data_source=data_source, + num_turns=num_turns, model=model, ) @@ -73,11 +75,23 @@ def compute_teacher_logits( outputs.append(TeacherOutput()) continue teacher_logits = teacher_logits.detach() + target_logprobs = None + if labels is not None: + safe_labels = labels.masked_fill(labels == -100, 0).long() + teacher_logprobs = vocab_parallel_log_softmax(teacher_logits.float()) + target_logprobs = tp_gather_topk(teacher_logprobs, safe_labels.unsqueeze(-1)).squeeze(-1) + target_logprobs = target_logprobs.masked_fill(labels == -100, float('nan')) if gkd_logits_topk is not None: topk_logits, topk_indices = vocab_parallel_topk(teacher_logits, k=gkd_logits_topk) - outputs.append(TeacherOutput(topk_logprobs=topk_logits, topk_indices=topk_indices, labels=labels)) + outputs.append( + TeacherOutput( + topk_logprobs=topk_logits, + topk_indices=topk_indices, + target_logprobs=target_logprobs, + labels=labels)) else: - outputs.append(TeacherOutput(full_logits=teacher_logits, labels=labels)) + outputs.append( + TeacherOutput(full_logits=teacher_logits, target_logprobs=target_logprobs, labels=labels)) del collated return outputs @@ -85,7 +99,7 @@ def compute_teacher_logits( # Loss computation # ------------------------------------------------------------------ - def loss_func(self, output_tensor, *, labels, teacher_output, data_source=None, model=None): + def loss_func(self, output_tensor, *, labels, teacher_output, data_source=None, num_turns=None, model=None): args = self.args student_logits = output_tensor @@ -128,6 +142,29 @@ def loss_func(self, output_tensor, *, labels, teacher_output, data_source=None, loss = loss + self.sft_alpha * sft_loss metric = {'loss': loss.detach().clone()} + if num_turns is not None: + metric['num_turns'] = loss.new_tensor(num_turns) + if data_source == DataSource.STUDENT: + monitor = gkd_monitoring_stats( + student_logits, + teacher_output, + labels, + full_vocab_topk=getattr(args, 'gkd_logits_topk', None) or 16, + student_topk_fn=vocab_parallel_topk, + teacher_topk_fn=vocab_parallel_topk, + gather_fn=tp_gather_topk, + target_logprob_fn=lambda logits, target_ids: tp_gather_topk( + vocab_parallel_log_softmax(logits.float()), target_ids.unsqueeze(-1)).squeeze(-1)) + packed = torch.stack([ + monitor['topk_overlap_sum'], monitor['topk_overlap_count'], monitor['teacher_student_gap_sum'], + monitor['teacher_student_gap_count'] + ]) + if args.context_parallel_size > 1: + torch.distributed.all_reduce( + packed, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group()) + torch.distributed.all_reduce(packed, op=torch.distributed.ReduceOp.SUM, group=mpu.get_data_parallel_group()) + metric['gkd/topk_overlap'] = packed[0] / packed[1].clamp(min=1) + metric['gkd/teacher_student_gap'] = packed[2] / packed[3].clamp(min=1) if sft_loss is not None: metric['jsd_loss'] = jsd_loss_val.detach().clone() metric['sft_loss'] = sft_loss.detach().clone() diff --git a/swift/ray/megatron/megatron_worker.py b/swift/ray/megatron/megatron_worker.py index 55318c8e88..39ef7b7c14 100644 --- a/swift/ray/megatron/megatron_worker.py +++ b/swift/ray/megatron/megatron_worker.py @@ -596,8 +596,8 @@ def _collate_teacher_outputs( """ from swift.rlhf_trainers.gkd_loss import TeacherOutput effective_target = None if is_opsd else target_seq_len - pad_vals = {'topk_logprobs': float('-inf'), 'labels': -100} - fields = ('full_logits', 'topk_logprobs', 'topk_indices', 'labels') + pad_vals = {'topk_logprobs': float('-inf'), 'target_logprobs': float('nan'), 'labels': -100} + fields = ('full_logits', 'topk_logprobs', 'topk_indices', 'target_logprobs', 'labels') kwargs = {} for field in fields: tensors = [getattr(t, field) for t in teacher_outputs] diff --git a/swift/rlhf_trainers/gkd_helpers.py b/swift/rlhf_trainers/gkd_helpers.py index 0ad0a6a24d..afe883aea9 100644 --- a/swift/rlhf_trainers/gkd_helpers.py +++ b/swift/rlhf_trainers/gkd_helpers.py @@ -176,7 +176,40 @@ def assemble_teacher_output( offsets=offsets, ) - teacher_out = TeacherOutput(topk_logprobs=topk_logprobs, topk_indices=topk_indices) + # ``prompt_logprobs=K`` may return the observed prompt token as an extra + # K+1 entry when it falls outside the teacher top-k. Preserve that value for + # the exact teacher-student log-probability gap monitor. + target_logprobs = torch.full((batch_size, seq_len), float('nan'), dtype=torch.float32) + input_ids_cpu = input_ids.detach().cpu() + + def _fill_target(row: int, pos: int, lps_row, ids_row) -> None: + target_pos = pos + 1 + if target_pos >= seq_len: + return + target_id = int(input_ids_cpu[row, target_pos]) + for lp, token_id in zip(lps_row, ids_row): + if int(token_id) == target_id: + target_logprobs[row, pos] = float(lp) + return + + if template_padding_free: + for i, (lps, ixs) in enumerate(parsed): + start, end = cu_seqlens[i], cu_seqlens[i + 1] + length = min(len(lps), end - start - 1) + for j in range(max(length, 0)): + _fill_target(0, start + j, lps[j], ixs[j]) + else: + for i, (lps, ixs) in enumerate(parsed): + start = offsets[i] if offsets is not None else 0 + length = min(len(lps), seq_len - start - 1) + for j in range(max(length, 0)): + _fill_target(i, start + j, lps[j], ixs[j]) + + teacher_out = TeacherOutput( + topk_logprobs=topk_logprobs, + topk_indices=topk_indices, + target_logprobs=target_logprobs.to(device), + ) if 'labels' in teacher_model_inputs: teacher_out.labels = teacher_model_inputs['labels'] return teacher_out diff --git a/swift/rlhf_trainers/gkd_loss.py b/swift/rlhf_trainers/gkd_loss.py index 82ae2f25d0..85ecb039a7 100644 --- a/swift/rlhf_trainers/gkd_loss.py +++ b/swift/rlhf_trainers/gkd_loss.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from dataclasses import dataclass from enum import Enum -from typing import Callable, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple # --------------------------------------------------------------------------- # Data types — shared across all backends @@ -26,6 +26,10 @@ class TeacherOutput: full_logits: Optional[torch.Tensor] = None topk_logprobs: Optional[torch.Tensor] = None topk_indices: Optional[torch.Tensor] = None + # Log-probability assigned by the teacher to the actually observed response + # token at each position. This keeps the teacher-student gap exact even when + # the distillation loss itself only retains the teacher's top-k support. + target_logprobs: Optional[torch.Tensor] = None labels: Optional[torch.Tensor] = None @property @@ -35,7 +39,7 @@ def is_topk_mode(self) -> bool: def to_device(self, device) -> 'TeacherOutput': """Move all tensor fields to ``device`` in place (Ray: teacher_output is collated on the CPU driver, moved to the GPU worker before forward).""" - for name in ('full_logits', 'topk_logprobs', 'topk_indices', 'labels'): + for name in ('full_logits', 'topk_logprobs', 'topk_indices', 'target_logprobs', 'labels'): v = getattr(self, name) if isinstance(v, torch.Tensor): setattr(self, name, v.to(device)) @@ -52,6 +56,7 @@ def select(self, mask: torch.Tensor) -> 'TeacherOutput': full_logits=self.full_logits[mask] if self.full_logits is not None else None, topk_logprobs=self.topk_logprobs[mask] if self.topk_logprobs is not None else None, topk_indices=self.topk_indices[mask] if self.topk_indices is not None else None, + target_logprobs=self.target_logprobs[mask] if self.target_logprobs is not None else None, labels=self.labels[mask] if self.labels is not None else None, ) @@ -64,6 +69,7 @@ def to_topk(self, k: int, topk_fn=None) -> 'TeacherOutput': return TeacherOutput( topk_logprobs=vals, topk_indices=ids, + target_logprobs=self.target_logprobs, labels=self.labels, ) @@ -216,6 +222,99 @@ def extract_active( return s_active, t_active, torch.tensor(int(s_active.shape[0]), device=labels.device) +@torch.no_grad() +def gkd_monitoring_stats( + student_logits: torch.Tensor, + teacher_output: TeacherOutput, + labels: torch.Tensor, + *, + full_vocab_topk: int = 16, + student_topk_fn: Callable = torch.topk, + teacher_topk_fn: Callable = torch.topk, + gather_fn: Callable = default_gather, + target_logprob_fn: Optional[Callable] = None, +) -> Dict[str, torch.Tensor]: + """Return additive GKD diagnostics over active response tokens. + + ``topk_overlap`` follows the standard token-level definition + ``|TopK(student) intersect TopK(teacher)| / K``. In a top-k teacher path, + K is the retained teacher width; for full-vocabulary distillation it defaults + to 16. + + ``teacher_student_gap`` is computed on the observed response token exactly as + ``log p_teacher(y_t) - log p_student(y_t)``. The returned values are sums and + counts so callers can aggregate them correctly across DP/CP ranks. + """ + s_active, t_active, num_valid = extract_active(student_logits, teacher_output, labels) + zero = student_logits.new_zeros((), dtype=torch.float32) + if int(num_valid.item()) == 0: + return { + 'topk_overlap_sum': zero, + 'topk_overlap_count': zero, + 'teacher_student_gap_sum': zero, + 'teacher_student_gap_count': zero, + } + + if t_active.is_topk_mode: + k = min(t_active.topk_indices.shape[-1], s_active.shape[-1]) + teacher_topk_ids = t_active.topk_indices[..., :k] + else: + k = min(full_vocab_topk, s_active.shape[-1], t_active.full_logits.shape[-1]) + _, teacher_topk_ids = teacher_topk_fn(t_active.full_logits, k) + + _, student_topk_ids = student_topk_fn(s_active, k) + overlap_count = (teacher_topk_ids.unsqueeze(-1) == student_topk_ids.unsqueeze(-2)).any(dim=-1).sum(dim=-1) + overlap_sum = (overlap_count.float() / k).sum() + + if t_active.labels is not None: + active_target_ids = t_active.labels.long() + else: + active_target_ids = labels[labels != -100].long() + if active_target_ids.numel() != s_active.shape[0]: + # A legacy top-k path can omit entire uncovered rows without carrying + # teacher labels. The overlap metric is still valid, but an exact gap + # cannot be aligned to the remaining observed tokens. + return { + 'topk_overlap_sum': overlap_sum.float(), + 'topk_overlap_count': num_valid.float(), + 'teacher_student_gap_sum': zero, + 'teacher_student_gap_count': zero, + } + if target_logprob_fn is None: + # Avoid materializing a second full-vocabulary log-probability tensor in + # the common non-TP path. + student_target_logits = gather_fn(s_active.float(), active_target_ids.unsqueeze(-1)).squeeze(-1) + student_target_logprobs = student_target_logits - torch.logsumexp(s_active.float(), dim=-1) + else: + student_target_logprobs = target_logprob_fn(s_active, active_target_ids) + + if t_active.target_logprobs is not None: + teacher_target_logprobs = t_active.target_logprobs.float() + gap_mask = torch.isfinite(teacher_target_logprobs) + elif t_active.full_logits is not None: + if target_logprob_fn is None: + teacher_target_logits = gather_fn(t_active.full_logits.float(), active_target_ids.unsqueeze(-1)).squeeze(-1) + teacher_target_logprobs = teacher_target_logits - torch.logsumexp(t_active.full_logits.float(), dim=-1) + else: + teacher_target_logprobs = target_logprob_fn(t_active.full_logits, active_target_ids) + gap_mask = torch.ones_like(teacher_target_logprobs, dtype=torch.bool) + else: + # Compatibility fallback for top-k tensors produced by older paths. It + # is exact only where the observed token is present in the retained set. + matches = t_active.topk_indices == active_target_ids.unsqueeze(-1) + gap_mask = matches.any(dim=-1) + match_pos = matches.float().argmax(dim=-1, keepdim=True) + teacher_target_logprobs = torch.gather(t_active.topk_logprobs.float(), -1, match_pos).squeeze(-1) + + gap = teacher_target_logprobs - student_target_logprobs + return { + 'topk_overlap_sum': overlap_sum.float(), + 'topk_overlap_count': num_valid.float(), + 'teacher_student_gap_sum': gap.masked_fill(~gap_mask, 0).sum().float(), + 'teacher_student_gap_count': gap_mask.sum().float(), + } + + # --------------------------------------------------------------------------- # gkd_loss — full pipeline: mask → prepare → jsd # --------------------------------------------------------------------------- diff --git a/swift/rlhf_trainers/gkd_trainer.py b/swift/rlhf_trainers/gkd_trainer.py index f997bfc7cf..393ec8937e 100644 --- a/swift/rlhf_trainers/gkd_trainer.py +++ b/swift/rlhf_trainers/gkd_trainer.py @@ -18,7 +18,7 @@ from swift.rl_core.data import GKDBatch, GKDSample from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_output, build_teacher_requests, encode_gkd_samples, fetch_teacher_parsed_by_routing) -from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss +from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss, gkd_monitoring_stats from swift.template import TemplateInputs from swift.trainers import SwiftMixin, disable_gradient_checkpointing from swift.utils import (JsonlWriter, get_logger, is_swanlab_available, is_wandb_available, remove_response, @@ -102,10 +102,12 @@ def get_train_dataloader(self): is_training=True, ) - def _compute_jsd_loss(self, student_logits, teacher_output: TeacherOutput, labels): + def _compute_jsd_loss(self, student_logits, teacher_output: TeacherOutput, labels, *, record_metrics=False): """Compute JSD loss. teacher_output.labels is always set (equals student labels when non-OPSD).""" shifted_labels = torch.roll(labels, shifts=-1, dims=1) teacher_output.labels = torch.roll(teacher_output.labels, shifts=-1, dims=1) + if record_metrics: + self._record_gkd_monitoring(student_logits, teacher_output, shifted_labels) if self.gkd_logits_topk is not None: teacher_output = teacher_output.to_topk(self.gkd_logits_topk) total, num_valid = gkd_loss(student_logits, teacher_output, shifted_labels, self.beta, self.temperature) @@ -113,6 +115,19 @@ def _compute_jsd_loss(self, student_logits, teacher_output: TeacherOutput, label return total * 0 return total / num_valid + def _record_gkd_monitoring(self, student_logits, teacher_output: TeacherOutput, labels) -> None: + stats = gkd_monitoring_stats(student_logits, teacher_output, labels, full_vocab_topk=self.gkd_logits_topk or 16) + packed = torch.stack([ + stats['topk_overlap_sum'], stats['topk_overlap_count'], stats['teacher_student_gap_sum'], + stats['teacher_student_gap_count'] + ]) + packed = self.accelerator.reduce(packed, reduction='sum') + mode = 'train' if self.model.training else 'eval' + if packed[1].item() > 0: + self._metrics[mode]['gkd/topk_overlap'].append((packed[0] / packed[1]).item()) + if packed[3].item() > 0: + self._metrics[mode]['gkd/teacher_student_gap'].append((packed[2] / packed[3]).item()) + @profiling_decorator def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): model_inputs = inputs['model_inputs'] @@ -239,7 +254,11 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N outputs_teacher = self.teacher_model(**t_fwd) teacher_out = TeacherOutput(full_logits=outputs_teacher.logits, labels=teacher_labels) - loss = self._compute_jsd_loss(outputs_student.logits, teacher_out, model_inputs['labels']) + loss = self._compute_jsd_loss( + outputs_student.logits, + teacher_out, + model_inputs['labels'], + record_metrics=data_source == DataSource.STUDENT) if self.args.sft_alpha > 0 and data_source != DataSource.STUDENT: loss = loss + self.args.sft_alpha * outputs_student.loss @@ -342,7 +361,16 @@ def _postprocess_batch(self, samples: List[GKDSample], batch_encoded_inputs: Lis self._fetch_and_assemble_teacher_logprobs(batch_encoded_inputs) def _log_rollout(self, samples: List[GKDSample]) -> None: - """Student completions are logged in ``_rollout_samples``; nothing extra here.""" + """Log multi-turn trajectory metadata for on-policy GKD rollouts.""" + if not samples or not all(s.rollout_infos and 'num_turns' in s.rollout_infos for s in samples): + return + num_turns = self._gather_and_flatten([s.rollout_infos['num_turns'] for s in samples], flatten_level=0) + mode = 'train' if self.model.training else 'eval' + self._metrics[mode]['num_turns'].append(sum(num_turns) / len(num_turns)) + if self.log_completions: + if 'num_turns' not in self._logs: + self._logs['num_turns'] = deque() + self._logs['num_turns'].extend(num_turns) @profiling_decorator def _prepare_inputs(self, inputs: DataType) -> Dict[str, torch.Tensor]: @@ -378,7 +406,7 @@ def _fetch_and_assemble_teacher_logprobs(self, chunks): self.teacher_clients, gather_fn=self._gather_teacher_requests, infer_fn=lambda handle, client: self._infer_teacher_requests( - handle, topk=self.gkd_logits_topk, teacher_client=client), + handle, topk=self.gkd_logits_topk, teacher_client=client, include_sampled=True), scatter_fn=self._scatter_teacher_parsed, is_main_process=self.accelerator.is_main_process, tag_key=self.args.teacher_tag_key) @@ -493,6 +521,12 @@ def _apply_chat_template_to_messages_list(self, messages_list: DataType): def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None: """Override log method to include completion table logging (aligned with GRPO).""" + mode = 'train' if self.model.training else 'eval' + metrics = {key: sum(values) / len(values) for key, values in self._metrics[mode].items() if values} + if mode == 'eval': + metrics = {f'eval_{key}': value for key, value in metrics.items()} + logs.update(metrics) + # Call parent log method import transformers from packaging import version @@ -500,6 +534,7 @@ def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> Non super().log(logs, start_time) else: super().log(logs) + self._metrics[mode].clear() # Log completions table if we have data (only for on-policy generations) if self.accelerator.is_main_process and self.log_completions and len(self._logs['prompt']) > 0: @@ -509,12 +544,15 @@ def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> Non 'prompt': list(self._logs['prompt'])[:seen_nums], 'completion': list(self._logs['completion'])[:seen_nums], } + for key, value in self._logs.items(): + if key not in table: + table[key] = list(value)[:seen_nums] # Write to jsonl self.jsonl_writer.append(table) - self._logs['prompt'].clear() - self._logs['completion'].clear() + for value in self._logs.values(): + value.clear() # Log to wandb if enabled report_to_wandb = self.args.report_to and 'wandb' in self.args.report_to and wandb.run is not None if report_to_wandb: diff --git a/swift/rlhf_trainers/rollout_mixin.py b/swift/rlhf_trainers/rollout_mixin.py index 0e2f47e439..5586488068 100644 --- a/swift/rlhf_trainers/rollout_mixin.py +++ b/swift/rlhf_trainers/rollout_mixin.py @@ -275,7 +275,11 @@ def _gather_teacher_requests(self, requests: List[Any]) -> Dict[str, Any]: all_counts = gather_object([n_local]) # per-rank counts in rank order return {'all_requests': all_requests, 'all_counts': all_counts, 'n_local': n_local} - def _infer_teacher_requests(self, handle: Dict[str, Any], topk: int, teacher_client: Optional[Any] = None): + def _infer_teacher_requests(self, + handle: Dict[str, Any], + topk: int, + teacher_client: Optional[Any] = None, + include_sampled: bool = False): """Phase 2 (main process only, no collective): run the teacher HTTP infer. Safe to call concurrently across teachers (distinct clients, no collective inside). @@ -286,7 +290,7 @@ def _infer_teacher_requests(self, handle: Dict[str, Any], topk: int, teacher_cli client = teacher_client if teacher_client is not None else self.teacher_clients[0] request_config = RequestConfig(prompt_logprobs=topk, max_tokens=1, temperature=0.0) responses = client.infer(handle['all_requests'], request_config=request_config, use_tqdm=False) - return [parse_prompt_logprobs(r, topk=topk) for r in responses] + return [parse_prompt_logprobs(r, topk=topk, include_sampled=include_sampled) for r in responses] def _scatter_teacher_parsed(self, handle: Dict[str, Any], parsed_global): """Phase 3 (all ranks, collective): broadcast the parsed result and slice this rank's part.""" diff --git a/swift/rlhf_trainers/utils.py b/swift/rlhf_trainers/utils.py index 329187f7c9..cf0b758a27 100644 --- a/swift/rlhf_trainers/utils.py +++ b/swift/rlhf_trainers/utils.py @@ -836,7 +836,9 @@ def replace_assistant_response_with_ids(messages: 'Messages', return messages -def parse_prompt_logprobs(response, topk: int) -> Tuple[List[List[float]], List[List[int]]]: +def parse_prompt_logprobs(response, + topk: int, + include_sampled: bool = False) -> Tuple[List[List[float]], List[List[int]]]: """Parse vLLM prompt_logprobs into per-position (logprobs, token_ids). vLLM's ``prompt_logprobs[i]`` is ``{token_id: {logprob, rank, ...}}`` for predicting @@ -847,7 +849,8 @@ def parse_prompt_logprobs(response, topk: int) -> Tuple[List[List[float]], List[ token — take that single entry. This is token-in-token-out, NOT the top-1: the sampled token may have any rank. - ``topk > 0`` (top-k, GKD): the ``topk`` highest-probability tokens, ordered by logprob - (== rank order). The sampled token, if returned as an extra ``k+1``-th entry, is dropped. + (== rank order). The sampled token, if returned as an extra ``k+1``-th entry, is dropped + unless ``include_sampled`` is true. """ raw = response.prompt_logprobs or [] lps: List[List[float]] = [] @@ -859,7 +862,9 @@ def parse_prompt_logprobs(response, topk: int) -> Tuple[List[List[float]], List[ lps.append([info['logprob']]) ixs.append([int(tid)]) else: - items = sorted(pos_lp.items(), key=lambda x: -x[1]['logprob'])[:topk] + items = sorted(pos_lp.items(), key=lambda x: -x[1]['logprob']) + if not include_sampled: + items = items[:topk] lps.append([info['logprob'] for _, info in items]) ixs.append([int(tid) for tid, _ in items]) return lps, ixs @@ -888,8 +893,8 @@ def assemble_teacher_topk_logprobs( length = min(len(lps), end - start) if length <= 0: continue - out_lp[start:start + length] = torch.tensor(lps[:length], dtype=torch.float32) - out_ix[start:start + length] = torch.tensor(ixs[:length], dtype=torch.long) + out_lp[start:start + length] = torch.tensor([row[:topk] for row in lps[:length]], dtype=torch.float32) + out_ix[start:start + length] = torch.tensor([row[:topk] for row in ixs[:length]], dtype=torch.long) return out_lp.unsqueeze(0).to(device), out_ix.unsqueeze(0).to(device) out_lp = torch.full((batch_size, seq_len, topk), float('-inf'), dtype=torch.float32) @@ -902,8 +907,8 @@ def assemble_teacher_topk_logprobs( length = min(P, seq_len - start) if length <= 0: continue - out_lp[idx, start:start + length] = torch.tensor(lps[:length], dtype=torch.float32) - out_ix[idx, start:start + length] = torch.tensor(ixs[:length], dtype=torch.long) + out_lp[idx, start:start + length] = torch.tensor([row[:topk] for row in lps[:length]], dtype=torch.float32) + out_ix[idx, start:start + length] = torch.tensor([row[:topk] for row in ixs[:length]], dtype=torch.long) return out_lp.to(device), out_ix.to(device) diff --git a/tests/train/test_gkd_monitoring.py b/tests/train/test_gkd_monitoring.py new file mode 100644 index 0000000000..3d702e119f --- /dev/null +++ b/tests/train/test_gkd_monitoring.py @@ -0,0 +1,101 @@ +import math +import torch +from collections import defaultdict, deque +from types import SimpleNamespace + +from swift.rl_core.data import GKDSample +from swift.rlhf_trainers.gkd_helpers import assemble_teacher_output +from swift.rlhf_trainers.gkd_loss import TeacherOutput, gkd_monitoring_stats +from swift.rlhf_trainers.gkd_trainer import GKDTrainer +from swift.rlhf_trainers.utils import parse_prompt_logprobs + + +def test_gkd_monitoring_full_logits_overlap_and_exact_gap(): + student = torch.tensor([[[-9., -9., -9., -9.], [4., 3., 1., 0.], [4., 0., 3., 2.]]]) + teacher = torch.tensor([[[-9., -9., -9., -9.], [4., 3., 0., 1.], [0., 1., 4., 3.]]]) + labels = torch.tensor([[-100, 1, 2]]) + teacher_out = TeacherOutput(full_logits=teacher, labels=labels) + + stats = gkd_monitoring_stats(student, teacher_out, labels, full_vocab_topk=2) + assert torch.isclose(stats['topk_overlap_sum'] / stats['topk_overlap_count'], torch.tensor(0.75)) + + expected = (torch.log_softmax(teacher[0, 1], -1)[1] - torch.log_softmax(student[0, 1], -1)[1] + + torch.log_softmax(teacher[0, 2], -1)[2] - torch.log_softmax(student[0, 2], -1)[2]) / 2 + actual = stats['teacher_student_gap_sum'] / stats['teacher_student_gap_count'] + torch.testing.assert_close(actual, expected) + + +def test_gkd_monitoring_topk_uses_observed_token_logprob_outside_topk(): + student = torch.tensor([[[3., 2., 1., 0.]]]) + labels = torch.tensor([[3]]) # observed token is outside the teacher's retained top-2 + teacher_target_lp = torch.tensor([[-0.25]]) + teacher_out = TeacherOutput( + topk_logprobs=torch.tensor([[[-0.1, -0.2]]]), + topk_indices=torch.tensor([[[0, 1]]]), + target_logprobs=teacher_target_lp, + labels=labels, + ) + + stats = gkd_monitoring_stats(student, teacher_out, labels) + expected_gap = teacher_target_lp.item() - torch.log_softmax(student[0, 0], -1)[3].item() + assert stats['teacher_student_gap_count'].item() == 1 + assert math.isclose( + (stats['teacher_student_gap_sum'] / stats['teacher_student_gap_count']).item(), expected_gap, rel_tol=1e-6) + + +def test_teacher_api_keeps_observed_token_outside_topk_for_gap(): + response = SimpleNamespace(prompt_logprobs=[ + None, + { + 0: { + 'logprob': -0.1 + }, + 1: { + 'logprob': -0.2 + }, + 9: { + 'logprob': -2.5 + }, + }, + { + 2: { + 'logprob': -0.3 + }, + 0: { + 'logprob': -0.4 + }, + 3: { + 'logprob': -3.5 + }, + }, + ]) + parsed = [parse_prompt_logprobs(response, topk=2, include_sampled=True)] + inputs = { + 'input_ids': torch.tensor([[5, 9, 3]]), + 'labels': torch.tensor([[-100, 9, 3]]), + 'attention_mask': torch.ones(1, 3, dtype=torch.long), + } + + teacher_out = assemble_teacher_output(parsed, inputs, topk=2, template_padding_free=False, device='cpu') + + torch.testing.assert_close(teacher_out.target_logprobs[0, :2], torch.tensor([-2.5, -3.5])) + assert teacher_out.topk_indices[0, 0].tolist() == [0, 1] + assert teacher_out.topk_indices[0, 1].tolist() == [2, 0] + + +def test_gkd_rollout_logs_num_turns(): + trainer = GKDTrainer.__new__(GKDTrainer) + trainer.model = SimpleNamespace(training=True) + trainer.log_completions = True + trainer._metrics = {'train': defaultdict(list), 'eval': defaultdict(list)} + trainer._logs = {'prompt': deque(), 'completion': deque()} + trainer._gather_and_flatten = lambda values, **_: values + samples = [ + GKDSample(messages=[], rollout_infos={'num_turns': 2}), + GKDSample(messages=[], rollout_infos={'num_turns': 4}), + ] + + trainer._log_rollout(samples) + + assert trainer._metrics['train']['num_turns'] == [3.0] + assert list(trainer._logs['num_turns']) == [2, 4]