/swift/megatron/trainers/gkd_trainer.py

# Copyright (c) ModelScope Contributors. All rights reserved.
import os
import random
import torch
import torch.nn.functional as F
from contextlib import contextmanager
from functools import partial
from mcore_bridge import set_random_seed
from megatron.core import mpu
from megatron.core.rerun_state_machine import RerunDataIterator
from transformers.utils import ContextManagers
from typing import Dict, List, Optional

from swift.megatron.arguments import MegatronArguments
from swift.rl_core.data import GKDSample
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.template import Template
from swift.utils import get_logger, to_device
from ..utils import forward_step_helper, get_padding_to
from .gkd_utils import cp_reduce, tp_gather_topk, vocab_parallel_topk
from .rlhf_mixin import MegatronRLHFTrainer
from .rollout_mixin import MegatronRolloutMixin
from .utils import gather_object
from .vocab_parallel_utils import vocab_parallel_kl_div, vocab_parallel_log_softmax

logger = get_logger()


class MegatronGKDTrainer(MegatronRolloutMixin, MegatronRLHFTrainer):

    sample_cls = GKDSample

    def __init__(self, args: MegatronArguments, template, **kwargs):
        self.vllm_client = kwargs.pop('vllm_client', None)

        # GKD-specific parameters
        self.beta = args.beta  # JSD interpolation coefficient
        self.temperature = args.temperature
        self.lmbda = args.lmbda  # On-policy probability
        self.args = args
        self._setup_teacher()
        if self._teacher_use_disable_adapter:
            logger.info('Self-distillation mode: using disable_adapter() for fixed teacher (no extra model)')
        self.sft_alpha = getattr(args, 'sft_alpha', 0.0)  # Weight for SFT loss

        # GKD top-k logits configuration
        self.gkd_logits_topk = getattr(args, 'gkd_logits_topk', None)

        self.use_vllm = getattr(args, 'use_vllm', False)
        self.steps_per_generation = args.steps_per_generation
        self.generation_batch_size = args.generation_batch_size
        super().__init__(args, template)

        if self.use_teacher_api:
            logger.info(f'Using teacher model API for logprobs, top_logprobs={self.gkd_logits_topk}')

        # Get device for data processing
        self.device = torch.cuda.current_device()

        # Initialize vLLM rollout engine if on-policy generation is enabled
        self._init_rollout_engine()

        # Truncation strategy for handling sequences that exceed max_length
        self.truncation_strategy = args.truncation_strategy
        self.max_completion_length = args.max_completion_length

        self.resample_data_iterator = None
        self._buffered_inputs = None

        self._prepare_logging()

    def train(self, train_dataset, val_dataset):
        if self.truncation_strategy == 'delete':
            self.resample_data_iterator = self._init_resample_data_iterator(train_dataset)
        super().train(train_dataset, val_dataset)

    def prepare_model(self):
        super().prepare_model()
        if self.use_teacher_api:
            logger.info('Skipping local teacher model loading - using external API for teacher logprobs')
        elif self._is_self_distillation:
            logger.info('Self-distillation mode: using student model as teacher (no separate teacher loaded)')
        self._load_teacher_model()

    @contextmanager
    def _template_context(self, template: Template, max_length: Optional[int] = None):
        """Context manager to temporarily modify max_length constraint from template."""
        original_max_length = template.max_length
        template.max_length = max_length
        try:
            yield
        finally:
            template.max_length = original_max_length

    def _build_teacher_requests(self, samples: List[GKDSample]):
        if not self.use_teacher_api:
            return []
        return build_teacher_requests(samples, self.template)

    def _encode_samples(self, samples: List[GKDSample]) -> Dict[str, torch.Tensor]:
        template = self.template
        args = self.args

        with self._template_context(template):
            student_encoded_list, teacher_encoded_list, has_opsd = encode_gkd_samples(samples, template)

        padding_to = get_padding_to(args)
        encoded_batch = to_device(template.data_collator(student_encoded_list, padding_to=padding_to), self.device)
        if has_opsd:
            teacher_model_inputs = to_device(
                template.data_collator(teacher_encoded_list, padding_to=padding_to), self.device)
        else:
            teacher_model_inputs = encoded_batch.copy()
        encoded_batch['teacher_model_inputs'] = teacher_model_inputs
        return encoded_batch

    def _get_random_num(self) -> float:
        """Generate a deterministic random number consistent across all processes.

        Uses an isolated Random instance with seed based on args.seed + step counter

        Returns:
            float: A random number in the range [0.0, 1.0).
        """
        seed = int(getattr(self.args, 'seed', 0))
        seed += int(self._step)
        rng = random.Random(seed)
        return rng.random()

    def _determine_data_source(self) -> DataSource:
        """Determine data source for current step based on GKD algorithm.

        GKD training mode selection logic:
        1. With probability lmbda: On-Policy (student generates)
        2. Otherwise: Off-Policy (use dataset responses)

        Returns:
            DataSource enum indicating which source to use.
        """
        random_num = self._get_random_num()

        if random_num < self.lmbda:
            # Mode 1: On-Policy learning, student model generates responses
            if self.use_vllm:
                return DataSource.STUDENT
            else:
                # If vLLM not enabled, fall back to dataset
                logger.warning_once('On-policy mode triggered but use_vllm=False. '
                                    'Falling back to dataset responses. Enable vLLM for on-policy generation.')
                return DataSource.DATASET
        else:
            # Mode 2: Off-Policy learning, use dataset responses
            return DataSource.DATASET

    def _init_resample_data_iterator(self, train_dataset):
        """Initialize an independent data iterator for resampling.

        Uses a different seed (args.seed + 1) to avoid overlapping with training samples.

        Args:
            train_dataset: The training dataset to create the resample iterator from.

        Returns:
            The resample data iterator (first element of the iterator tuple).
        """
        args = self.args
        resample_seed = getattr(args, 'seed', 42) + 1
        try:
            set_random_seed(
                resample_seed,
                args.data_parallel_random_init,
                args.te_rng_tracker,
            )
            resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0]
        finally:
            set_random_seed(
                args.seed,
                args.data_parallel_random_init,
                args.te_rng_tracker,
            )
        return resample_data_iterator

    def resample_encode_failed_inputs(self, inputs: List[Dict], max_resample_rounds: int = 10) -> List[Dict]:
        """Attempt to encode each input. If encoding fails, resample until we have enough valid samples.

        Args:
            inputs: List of input data samples
            max_resample_rounds: Maximum number of resample rounds

        Returns:
            List of successfully encoded input samples with the same length as inputs
        """
        return resample_encode_failed_inputs(
            self.template,
            self.resample_data_iterator,
            inputs,
            max_resample_rounds=max_resample_rounds,
            strip_response=False,
        )

    def _assemble_teacher_outputs(self, encoded_batches: List[Dict]) -> None:
        for encoded_batch in encoded_batches:
            parsed = encoded_batch.pop('_teacher_parsed')
            teacher_model_inputs = encoded_batch['teacher_model_inputs']
            teacher_out = assemble_teacher_output(
                parsed,
                teacher_model_inputs=teacher_model_inputs,
                topk=self.gkd_logits_topk,
                template_padding_free=self.template.padding_free,
                device=self.device,
            )
            if teacher_out.labels is not None:
                teacher_out.labels = torch.roll(teacher_out.labels, shifts=-1, dims=-1)
            encoded_batch['teacher_output'] = teacher_out

    def _compute_teacher_logits(self, encoded_batches: List[Dict], vp_stage: Optional[int] = None) -> None:
        if self.use_teacher_api:
            self._assemble_teacher_outputs(encoded_batches)
            return
        if self._is_self_distillation:
            # Self-distillation teacher == current student weights. Computing it here (at batch
            # preparation, once per steps_per_generation cycle) would reuse stale weights across the
            # cycle's train steps. Defer to _replace_data_iterator so each train step recomputes the
            # teacher with up-to-date student weights (weights are constant within a train step).
            return
        self._compute_teacher_logits_local(encoded_batches, vp_stage)

    def _compute_teacher_logits_local(self, encoded_batches: List[Dict], vp_stage: Optional[int] = None) -> None:
        """Compute teacher_output for each micro-batch via a local forward.

        Handles both a separate fixed teacher and self-distillation (teacher == current student
        weights). For self-distillation the caller is responsible for invoking this per train step
        so the weights are current.
        """
        topk = self.gkd_logits_topk
        if self._is_self_distillation:
            teacher_model = self.unwrapped_models[vp_stage or 0]
            adapter_contexts = []
            if self._teacher_use_disable_adapter:
                adapter_contexts = [m.disable_adapter() for m in self.peft_models]
            outer_context = ContextManagers(adapter_contexts)
        else:
            teacher_model = self.teacher_models[vp_stage or 0]
            outer_context = self.load_teacher_model_context()

        with torch.no_grad(), outer_context:
            for encoded_batch in encoded_batches:
                teacher_model_inputs = encoded_batch['teacher_model_inputs']
                teacher_batch = {
                    k: v.clone() if isinstance(v, torch.Tensor) else v
                    for k, v in teacher_model_inputs.items()
                }
                teacher_data = self._prepare_batch(teacher_batch, vp_stage)
                teacher_data.pop('loss_scale', None)
                teacher_labels = teacher_data.pop('labels', None)
                teacher_logits = forward_step_helper(teacher_model, teacher_data)
                if teacher_logits is not None:
                    teacher_logits = teacher_logits.detach()

                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)
                else:
                    teacher_out = TeacherOutput(full_logits=teacher_logits)

                teacher_out.labels = teacher_labels
                encoded_batch['teacher_output'] = teacher_out

    def _generate_and_score_completions(self, inputs: List[Dict]) -> List[Dict]:
        """Unified rollout → teacher → encode pipeline (mirrors Megatron GRPO).

        Stages: determine data source → to_samples → (student) generate → teacher
        requests/logprobs → encode micro-batches → teacher logits. Returns the flat
        list of encoded micro-batches (length == total microbatches).
        """
        data_source = self._determine_data_source()

        # Convert to samples (resample operates on dict, to_samples after)
        samples = self.to_samples(inputs)

        if data_source == DataSource.STUDENT:
            local_batch = self._get_local_rollout_batch(samples)
            local_batch = self._generate_completions(local_batch)
            samples = self._gather_rollout_results(local_batch)
            self._log_completions_from_samples(samples)

        # Teacher API: build requests from samples, fetch logprobs
        local_parsed = None
        if self.use_teacher_api:
            build_opsd_samples(samples)
            teacher_requests = self._build_teacher_requests(samples)
            if teacher_requests:
                local_parsed = fetch_teacher_parsed_by_routing(
                    samples,
                    teacher_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.is_main_process,
                    tag_key=self.args.teacher_tag_key)

        # Encode micro-batches
        total_microbatches = self.args.num_microbatches * self.steps_per_generation
        micro_batch_size = len(samples) // total_microbatches
        assert micro_batch_size == self.args.micro_batch_size
        all_encoded_batches = []
        for i in range(total_microbatches):
            start_idx = i * micro_batch_size
            end_idx = start_idx + micro_batch_size
            sample_slice = samples[start_idx:end_idx]
            encoded_batch = self._encode_samples(sample_slice)
            encoded_batch['data_source'] = data_source
            if local_parsed is not None:
                encoded_batch['_teacher_parsed'] = local_parsed[start_idx:end_idx]
            all_encoded_batches.append(encoded_batch)
        self._compute_teacher_logits(all_encoded_batches)
        return all_encoded_batches

    def _replace_data_iterator(self, data_iterator):
        num_microbatches = self.args.num_microbatches
        steps_per_generation = self.steps_per_generation

        if self._step % steps_per_generation == 0:
            total_microbatches = num_microbatches * steps_per_generation
            global_batch = []
            for _ in range(total_microbatches):
                raw_batch = next(data_iterator)
                if self.truncation_strategy == 'delete' and self.resample_data_iterator is not None:
                    raw_batch = self.resample_encode_failed_inputs(raw_batch)
                global_batch.extend(raw_batch)

            all_encoded_batches = self._generate_and_score_completions(global_batch)
            self._buffered_inputs = [
                all_encoded_batches[i * num_microbatches:(i + 1) * num_microbatches]
                for i in range(steps_per_generation)
            ]

        step_idx = self._step % steps_per_generation
        encoded_batches = self._buffered_inputs[step_idx]

        # Self-distillation teacher == current student weights. Recompute per train step (weights are
        # constant within a step) instead of once per generation cycle, so it tracks student updates
        # across steps_per_generation. Runs outside the pipeline schedule, so PP > 1 is supported.
        if self._is_self_distillation:
            self._compute_teacher_logits_local(encoded_batches)

        self._step += 1

        return RerunDataIterator(iter(encoded_batches))

    def loss_func(self,
                  output_tensor: torch.Tensor,
                  *,
                  labels: torch.Tensor,
                  teacher_output: TeacherOutput,
                  data_source: DataSource = DataSource.DATASET):
        """Compute GKD loss (JSD + optional SFT loss)."""
        student_logits = output_tensor

        jsd_total, jsd_num_valid = gkd_loss(
            student_logits,
            teacher_output,
            labels,
            self.beta,
            self.temperature,
            gather_fn=tp_gather_topk,
            log_softmax_fn=vocab_parallel_log_softmax,
            kl_div_fn=vocab_parallel_kl_div)
        jsd_loss_val = cp_reduce(jsd_total, jsd_num_valid, cp_size=self.args.context_parallel_size)

        loss = jsd_loss_val

        # Add SFT loss if enabled (skip for student-generated responses)
        sft_loss = None
        if self.sft_alpha > 0 and data_source != DataSource.STUDENT:
            args = self.args
            logits_sbv = student_logits.transpose(0, 1).contiguous()
            model = self.unwrapped_models[0]
            if hasattr(model, 'language_model'):
                model = model.language_model
            per_token_loss = model.compute_language_model_loss(labels, logits_sbv)
            loss_mask = labels != -100
            sft_loss_sum = (per_token_loss * loss_mask).sum()
            sft_loss_count = loss_mask.sum().float()

            # All-reduce across CP group for correct averaging
            if args.context_parallel_size > 1:
                sft_stats = torch.stack([sft_loss_sum, sft_loss_count])
                torch.distributed.all_reduce(
                    sft_stats, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group())
                sft_loss_sum, sft_loss_count = sft_stats[0], sft_stats[1]

            sft_loss = sft_loss_sum / sft_loss_count

            loss = loss + self.sft_alpha * sft_loss

        metric = {'loss': loss.detach().clone()}
        if sft_loss is not None:
            metric['jsd_loss'] = jsd_loss_val.detach().clone()
            metric['sft_loss'] = sft_loss.detach().clone()
        metric = self._all_reduce_metric(metric)

        loss = loss / mpu.get_context_parallel_world_size()

        # Flush completion logs at generation cycle boundaries.
        if (self._step - 1) % self.steps_per_generation == 0:
            self._flush_log_completions()

        return loss, metric

    def forward_step(self, data_iterator, model):
        unwrapped_model = model.module.module
        input_tensor = unwrapped_model.get_input_tensor()
        vp_stage = unwrapped_model.vp_stage

        data = next(data_iterator)
        data_source = data.pop('data_source', DataSource.DATASET)
        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)

        data.pop('loss_scale', None)
        labels = data.pop('labels', None)

        if input_tensor is not None:
            unwrapped_model.set_input_tensor(input_tensor)
        student_output = model(**data)

        return student_output, partial(
            self.loss_func,
            labels=labels,
            teacher_output=teacher_output,
            data_source=data_source,
        )