# Copyright (c) ModelScope Contributors. All rights reserved. import inspect import os import random import torch import torch.nn as nn import trl from accelerate.utils import gather_object, is_peft_model from collections import defaultdict, deque from contextlib import contextmanager, nullcontext from copy import deepcopy from packaging import version from transformers import PreTrainedModel, Trainer from trl import SFTTrainer as HFSFTTrainer from trl.trainer.utils import RepeatSampler from typing import Any, Dict, List, Optional, Union 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.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, swanlab_get_run, to_device) from .rollout_mixin import DataType, RolloutTrainerMixin from .utils import get_gather_if_zero3_context, identity_data_collator, profiling_decorator try: from liger_kernel.chunked_loss import LigerFusedLinearJSDLoss _liger_kernel_available = True except ImportError: _liger_kernel_available = False if version.parse(trl.__version__) >= version.parse('0.26.0'): from trl.experimental.gkd import GKDTrainer as HFGKDTrainer else: from trl import GKDTrainer as HFGKDTrainer del HFGKDTrainer.__init__ del HFSFTTrainer.__init__ logger = get_logger() if is_wandb_available(): import wandb if is_swanlab_available(): import swanlab class GKDTrainer(RolloutTrainerMixin, SwiftMixin, HFGKDTrainer): sample_cls = GKDSample def __init__(self, model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, *_args, **kwargs): self.vllm_client = kwargs.pop('vllm_client', None) self.gkd_logits_topk = kwargs.pop('gkd_logits_topk', None) self._pop_teacher_kwargs(kwargs) # teacher kwargs (shared with GRPO via RolloutTrainerMixin) super().__init__(model, None, *_args, **kwargs) args = kwargs['args'] self.lmbda = args.lmbda self.temperature = args.temperature self._metrics = {'train': defaultdict(list), 'eval': defaultdict(list)} self._total_train_tokens = 0 # Initialize logging components self._prepare_logging() # Initialize liger loss if enabled self._prepare_liger_loss() self._setup_teacher() # Initialize rollout infrastructure for vLLM support self.prepare_rollout() # Initialize resample data iterator for truncation_strategy 'raise'('delete') if self.template.truncation_strategy == 'raise': self._prepare_resample_data_iterator() self._step = 0 self._buffered_inputs = None def _get_data_collator(self, args, template): return identity_data_collator def _get_train_sampler(self, train_dataset=None): return RepeatSampler( data_source=train_dataset or self.train_dataset, mini_repeat_count=1, batch_size=self.args.generation_batch_size, repeat_count=self.args.steps_per_generation, shuffle=getattr(self, 'shuffle_dataset', True), seed=self.args.seed, ) def get_train_dataloader(self): return self._get_dataloader( dataset=self.train_dataset, description='Training', batch_size=self._train_batch_size * self.args.steps_per_generation, sampler_fn=self._get_train_sampler, is_training=True, ) def _compute_jsd_loss(self, student_logits, teacher_output: TeacherOutput, labels): """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 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) if num_valid == 0: return total * 0 return total / num_valid @profiling_decorator def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): model_inputs = inputs['model_inputs'] gkd_batch: GKDBatch = inputs['gkd_batch'] data_source = gkd_batch.data_source teacher_model_inputs = inputs['teacher_model_inputs'] # logits_to_keep trims each forward to its completion region. student and teacher # are masked independently in the loss via their own labels (extract_active), so the # two sides can be trimmed independently even when OPSD makes their lengths differ. use_logits_to_keep = self.get_use_logits_to_keep(self.template.sequence_parallel_size == 1) if use_logits_to_keep and not self.use_liger_gkd_loss: self.prepare_logits_to_keep(model_inputs) # Trim the teacher only when it runs a forward. Teacher API does not (top-k is # pre-assembled on the full sequence). For a local teacher forward, non-OPSD reuses # the student encoding (separate dict) and OPSD has its own encoding — both are trimmed # here so the teacher forward also benefits, while extract_active aligns each side via # its own labels. if not self.use_teacher_api: self.prepare_logits_to_keep(teacher_model_inputs) # Teacher labels come from the teacher-side encoding (== student encoding when non-OPSD). # Must be read AFTER prepare_logits_to_keep which may truncate labels. teacher_labels = teacher_model_inputs['labels'] # model_inputs is clean from template.encode; exclude 'labels' for model forward. forward_inputs = {k: v for k, v in model_inputs.items() if k != 'labels'} teacher_fwd_inputs = {k: v for k, v in teacher_model_inputs.items() if k != 'labels'} if self.use_liger_gkd_loss: # Liger fused JSD loss for memory efficiency # Get base models (exclude lm_head to save memory) unwrapped_student = self.accelerator.unwrap_model(model) if is_peft_model(unwrapped_student): unwrapped_student = unwrapped_student.base_model.model base_student = getattr(unwrapped_student, getattr(unwrapped_student, 'base_model_prefix', 'model'), unwrapped_student) unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model) base_teacher = getattr(unwrapped_teacher, getattr(unwrapped_teacher, 'base_model_prefix', 'model'), unwrapped_teacher) # Forward through base models student_outputs = base_student(**forward_inputs, use_cache=False) load_context = self.load_teacher_model_context() if self.args.offload_teacher_model else nullcontext() with load_context: with torch.no_grad(), disable_gradient_checkpointing(self.teacher_model, self.args.gradient_checkpointing_kwargs): teacher_outputs = base_teacher(**forward_inputs, use_cache=False) # Get hidden states (shifted) student_hidden = student_outputs.last_hidden_state[:, :-1] teacher_hidden = teacher_outputs.last_hidden_state[:, :-1] # Release full outputs to free memory del student_outputs, teacher_outputs # Prepare labels (shifted) labels_mask = model_inputs['labels'] != -100 masked_input_ids = torch.where(labels_mask, model_inputs['input_ids'], torch.full_like(model_inputs['input_ids'], -100)) true_labels = masked_input_ids[:, 1:].contiguous() # Release intermediate tensors del labels_mask, masked_input_ids # Get output heads student_head = unwrapped_student.get_output_embeddings() teacher_head = unwrapped_teacher.get_output_embeddings() # Prepare context managers for gathering parameters in zero3 teacher_context = get_gather_if_zero3_context(self, is_zero3=self.is_teacher_ds3)(teacher_head.weight) student_context = get_gather_if_zero3_context(self)(student_head.weight) with teacher_context, student_context: # Compute liger fused JSD loss loss = self.liger_jsd_loss( student_input=student_hidden, student_weight=student_head.weight, teacher_input=teacher_hidden, teacher_weight=teacher_head.weight, true_labels=true_labels, student_bias=getattr(student_head, 'bias', None), teacher_bias=getattr(teacher_head, 'bias', None), ) # Release hidden states after loss computation del student_hidden, teacher_hidden, true_labels outputs_student = None # Non-liger path: student forward + teacher output construction + JSD loss else: if self.args.sft_alpha > 0: forward_inputs['labels'] = model_inputs['labels'] outputs_student = model(**forward_inputs) # Construct teacher_out — the only part that differs across modes if self.use_teacher_api: teacher_topk_logprobs = gkd_batch.teacher_topk_logprobs teacher_topk_indices = gkd_batch.teacher_topk_indices assert teacher_topk_logprobs is not None and teacher_topk_indices is not None, ( 'teacher_topk_logprobs/teacher_topk_indices missing on gkd_batch; ' 'ensure _prepare_inputs (or prediction_step) populated them.') teacher_out = TeacherOutput( topk_logprobs=teacher_topk_logprobs, topk_indices=teacher_topk_indices, labels=teacher_labels, ) else: t_fwd = teacher_fwd_inputs if self._is_self_distillation: adapter_ctx = ( self.accelerator.unwrap_model(model).disable_adapter() if self._teacher_use_disable_adapter else nullcontext()) with torch.no_grad(), adapter_ctx, \ disable_gradient_checkpointing(model, self.args.gradient_checkpointing_kwargs): outputs_teacher = model(**t_fwd) else: assert self.teacher_model is not None load_context = ( self.load_teacher_model_context() if self.args.offload_teacher_model else nullcontext()) with torch.no_grad(), load_context, disable_gradient_checkpointing( self.teacher_model, self.args.gradient_checkpointing_kwargs): 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']) if self.args.sft_alpha > 0 and data_source != DataSource.STUDENT: loss = loss + self.args.sft_alpha * outputs_student.loss # Return loss if return_outputs: return (loss, outputs_student) else: return loss def _log_student_completions(self, generated_samples) -> None: """Log student completions from generated GKDSample list.""" if not self.log_completions: return messages = [s.messages[:-1] for s in generated_samples] completions = [deepcopy(s.messages[-1]['content']) for s in generated_samples] valid_messages = gather_object(messages) valid_completions = gather_object(completions) self._logs['prompt'].extend(self._apply_chat_template_to_messages_list(valid_messages)) self._logs['completion'].extend(valid_completions) def _encode_samples(self, samples: List[GKDSample]): """Encode student + teacher (OPSD-aware). Mirrors Megatron GKD _encode_samples. Args: samples: Samples to encode (may be generated samples for STUDENT mode). Returns: (encoded_inputs, teacher_encoded_or_None) — ``teacher_encoded`` is non-None only for OPSD; non-OPSD returns None (teacher reuses the student encoding). """ template = self.template with self._template_context(template): student_encoded_list, teacher_encoded_list, has_opsd = encode_gkd_samples(samples, template) encoded_inputs = to_device(template.data_collator(student_encoded_list), self.model.device) if has_opsd: with self._template_context(template): teacher_encoded = to_device(template.data_collator(teacher_encoded_list), self.model.device) else: teacher_encoded = None return encoded_inputs, teacher_encoded def _build_teacher_requests(self, samples: List[GKDSample]): """Build teacher API requests from samples (mirrors Megatron GKD).""" return build_teacher_requests(samples, self.template) def _rollout_samples(self, inputs: DataType) -> List[GKDSample]: """Pick the data source, convert rows to samples, and generate for the STUDENT branch.""" if self._get_random_num() <= self.lmbda: self._data_source = DataSource.STUDENT else: self._data_source = DataSource.DATASET if self.template.truncation_strategy == 'raise': inputs = self.resample_encode_failed_inputs(inputs, strip_response=self._data_source != DataSource.DATASET) samples = self.to_samples(inputs) if self._data_source == DataSource.STUDENT: # Generation at full-batch level (uses _generate_completions from mixin) samples = self._generate_completions(samples) self._log_student_completions(samples) return samples def _score_completions(self, samples: List[GKDSample]) -> List[GKDSample]: """GKD has no reward scoring; teacher signals are filled in ``_postprocess_batch``.""" return samples @profiling_decorator def _prepare_batch_inputs(self, samples: List[GKDSample]) -> List[Dict[str, Any]]: """Encode + collate samples into per-micro-batch ``{model_inputs, gkd_batch}`` dicts. Mirrors GRPO ``_prepare_batch_inputs``: takes all samples and splits internally. """ sample_chunks = self.split_by_mini_batches(samples) ga_batch_encoded_inputs: List[Dict[str, Any]] = [] for chunk_samples in sample_chunks: encoded_inputs, teacher_encoded = self._encode_samples(chunk_samples) # teacher_encoded is non-None only for OPSD. Non-OPSD reuses the student encoding # for the teacher, but as a separate dict so each side can run prepare_logits_to_keep # independently without aliasing. teacher_model_inputs = teacher_encoded or dict(encoded_inputs) result: Dict[str, Any] = { 'model_inputs': encoded_inputs, 'gkd_batch': GKDBatch(data_source=self._data_source), 'teacher_model_inputs': teacher_model_inputs, } if self.use_teacher_api: result['_teacher_requests'] = self._build_teacher_requests(chunk_samples) result['_chunk_samples'] = chunk_samples ga_batch_encoded_inputs.append(result) return ga_batch_encoded_inputs def _postprocess_batch(self, samples: List[GKDSample], batch_encoded_inputs: List[Dict[str, Any]]) -> None: """Fetch teacher logprobs via API and fill them into each ``gkd_batch``.""" if self.use_teacher_api: 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.""" @profiling_decorator def _prepare_inputs(self, inputs: DataType) -> Dict[str, torch.Tensor]: mode = 'train' if self.model.training else 'eval' steps_per_generation = self.args.steps_per_generation if mode == 'train': if self._step % steps_per_generation == 0 or self._buffered_inputs is None: self._buffered_inputs = self._generate_and_score_completions(inputs) inputs = self._buffered_inputs[self._step % steps_per_generation] self._step += 1 else: inputs = self._generate_and_score_completions(inputs)[0] return inputs def _fetch_and_assemble_teacher_logprobs(self, chunks): """Fetch teacher logprobs via API using sample-based _teacher_requests. Each sample routes to exactly one teacher by tag (single teacher = all samples). Routing runs once over the flattened chunks, so a single teacher is one DP gather → infer → slice. The parsed results are sliced back per chunk and fed to ``assemble_teacher_output``. """ flat_samples, flat_requests, chunk_sizes = [], [], [] for c in chunks: reqs = c.pop('_teacher_requests', []) flat_samples.extend(c.pop('_chunk_samples', [])) flat_requests.extend(reqs) chunk_sizes.append(len(reqs)) parsed = fetch_teacher_parsed_by_routing( flat_samples, flat_requests, self.teacher_configs, 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), scatter_fn=self._scatter_teacher_parsed, is_main_process=self.accelerator.is_main_process, tag_key=self.args.teacher_tag_key) per_chunk_parsed, offset = [], 0 for cs in chunk_sizes: per_chunk_parsed.append(parsed[offset:offset + cs]) offset += cs for c, chunk_parsed in zip(chunks, per_chunk_parsed): gkd_batch: GKDBatch = c['gkd_batch'] target = c['teacher_model_inputs'] teacher_out = assemble_teacher_output( chunk_parsed, teacher_model_inputs=target, topk=self.gkd_logits_topk, template_padding_free=self.template.padding_free, device=target['input_ids'].device, ) gkd_batch.teacher_topk_logprobs = teacher_out.topk_logprobs gkd_batch.teacher_topk_indices = teacher_out.topk_indices @profiling_decorator def training_step(self, model: nn.Module, inputs: DataType, num_items_in_batch: Optional[int] = None) -> torch.Tensor: return HFSFTTrainer.training_step(self, model, inputs, num_items_in_batch) def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): inputs = self._prepare_inputs(inputs) with torch.no_grad(): with self.compute_loss_context_manager(): loss = self.compute_loss(model, inputs) loss = loss.mean().detach() return loss, None, None @contextmanager def offload_context(self): """Offload student model and optimizer to CPU during vLLM on-policy generation.""" if self.args.offload_model: self.offload_model(self.accelerator.unwrap_model(self.model)) if getattr(self, 'optimizer', None) and self.args.offload_optimizer: self.offload_optimizer() try: yield finally: if self.args.offload_model: self.load_model(self.accelerator.unwrap_model(self.model)) if getattr(self, 'optimizer', None) and self.args.offload_optimizer: self.load_optimizer() def _get_random_num(self) -> float: """ Generate a deterministic random number. Uses an isolated Random instance to avoid interfering with the global random state, ensuring thread-safety and consistent behavior across processes. Returns: float: A random number in the range [0.0, 1.0). """ seed = int(getattr(self.args, 'seed', 0)) seed += int(self.state.global_step) rng = random.Random(seed) return rng.random() def _prepare_liger_loss(self): """Initialize liger loss if enabled.""" args = self.args self.use_liger_gkd_loss = False if getattr(args, 'use_liger_kernel', False): if not _liger_kernel_available: raise ImportError( 'Liger kernel is not installed. Please install liger-kernel by running: pip install liger-kernel') assert self.args.sft_alpha == 0, 'SFT loss is not supported with liger loss' assert self.gkd_logits_topk is None, 'Top-k mode is not supported with liger loss' self.liger_jsd_loss = LigerFusedLinearJSDLoss( beta=self.beta, ignore_index=-100, temperature=self.temperature, compiled=False, ) self.use_liger_gkd_loss = True def _prepare_logging(self): """Initialize logging components for on-policy rollout tracking.""" args = self.args self.log_completions = args.log_completions self.wandb_log_unique_prompts = getattr(args, 'wandb_log_unique_prompts', False) self.jsonl_writer = JsonlWriter(os.path.join(self.args.output_dir, 'completions.jsonl')) # Initialize logs deque for storing rollout data (aligned with GRPO) self._logs = { 'prompt': deque(), 'completion': deque(), } def _apply_chat_template_to_messages_list(self, messages_list: DataType): """Convert messages list to prompt text list using template (aligned with GRPO).""" prompts_text = [] for messages in messages_list: remove_response(messages) template_inputs = TemplateInputs.from_dict({'messages': messages}) res = self.template.encode(template_inputs) prompts_text.append(self.template.safe_decode(res['input_ids'])) return prompts_text def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None: """Override log method to include completion table logging (aligned with GRPO).""" # Call parent log method import transformers from packaging import version if version.parse(transformers.__version__) >= version.parse('4.47.0.dev0'): super().log(logs, start_time) else: super().log(logs) # 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: seen_nums = len(self._logs['prompt']) table = { 'step': [str(self.state.global_step)] * seen_nums, 'prompt': list(self._logs['prompt'])[:seen_nums], 'completion': list(self._logs['completion'])[:seen_nums], } # Write to jsonl self.jsonl_writer.append(table) self._logs['prompt'].clear() self._logs['completion'].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: wandb_table = table.copy() import pandas as pd df = pd.DataFrame(wandb_table) if self.wandb_log_unique_prompts: df = df.drop_duplicates(subset=['prompt']) wandb.log({'completions': wandb.Table(dataframe=df)}) # Log to swanlab if enabled report_to_swanlab = self.args.report_to and 'swanlab' in self.args.report_to and swanlab_get_run( ) is not None if report_to_swanlab: headers = list(table.keys()) rows = [] for i in range(len(table['step'])): row = [table[header][i] for header in headers] rows.append(row) swanlab.log({'completions': swanlab.echarts.Table().add(headers, rows)})