# Copyright (c) ModelScope Contributors. All rights reserved. """Driver-side GKD trainer for Ray-based Megatron training.""" from __future__ import annotations import copy import random import ray import torch from contextlib import contextmanager from typing import List from swift.infer_engine.protocol import RequestConfig, RolloutOutput from swift.rl_core.data import GKDSample from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput from swift.rlhf_trainers.utils import parse_prompt_logprobs from swift.rollout import MultiTurnScheduler, invoke_async_hook, multi_turns, run_multi_turn from swift.utils import get_logger, remove_response from .base_trainer import BaseRayTrainer from .driver_utils import extract_iteration from .worker_group import DPDispatchedDict logger = get_logger() class GKDTrainer(BaseRayTrainer): def _prepare_state(self) -> None: super()._prepare_state() args = self.args self.sft_alpha = args.sft_alpha self.gkd_logits_topk = args.gkd_logits_topk # GKD on-policy schedule: each step is on-policy (student generates) with # probability ``lmbda``; otherwise off-policy (distill on dataset responses). self.lmbda = args.lmbda self._data_source_rng = random.Random(getattr(args, 'seed', 42)) # GKD generates exactly one completion per prompt (on-policy student generation), # so num_generations is always 1 here regardless of the (GRPO-oriented) default. self._data_info['num_generations'] = 1 self._teacher_model_dir = getattr(args, 'teacher_model_dir', None) or args.teacher_model self._teacher_model_server = args.teacher_model_server self._teacher_use_disable_adapter = args._teacher_use_disable_adapter if self._teacher_use_disable_adapter: self._teacher_model_dir = None if self._teacher_model_server and not self.teacher_replicas: raise NotImplementedError('teacher_model_server is not yet supported in the Ray pipeline. ' 'Use teacher_model (colocated) or teacher replicas (teacher.gpus > 0) instead.') vp_size = getattr(args, 'virtual_pipeline_model_parallel_size', None) assert vp_size is None or vp_size == 1, \ 'Ray GKD does not support VPP (virtual_pipeline_model_parallel_size > 1).' # truncation_strategy='delete': resample prompts whose encode fails (over max_length). self.truncation_strategy = args.truncation_strategy self.max_completion_length = args.max_completion_length self._max_resample_rounds = getattr(args, 'max_resample_times', 3) self._needs_resample_iterator = self.truncation_strategy == 'delete' self._prepare_multi_turn() def _prepare_multi_turn(self) -> None: args = self.args self._multi_turn_scheduler: MultiTurnScheduler | None = None self._max_turns: int | None = getattr(args, 'max_turns', None) self._enable_server_multi_turn = False scheduler_cfg = getattr(args, 'multi_turn_scheduler', None) if not scheduler_cfg: return if isinstance(scheduler_cfg, str): if scheduler_cfg not in multi_turns: raise ValueError(f'Unknown multi_turn_scheduler: {scheduler_cfg!r}; ' f'available: {list(multi_turns)}') scheduler_kwargs = {'max_turns': self._max_turns} tokenizer = getattr(getattr(self, 'template', None), 'tokenizer', None) if tokenizer is not None: scheduler_kwargs['tokenizer'] = tokenizer gym_env = getattr(args, 'gym_env', None) if gym_env is not None: scheduler_kwargs['gym_env'] = gym_env self._multi_turn_scheduler = multi_turns[scheduler_cfg](**scheduler_kwargs) else: assert isinstance(scheduler_cfg, MultiTurnScheduler) self._multi_turn_scheduler = scheduler_cfg def _train_loop(self, tg, train_iters, iteration): ckpt = self.ckpt_manager spg = self._steps_per_generation # Initialize colocated teacher if configured (skip for self-distillation via disable_adapter) if self._teacher_model_dir and not self._teacher_model_server and not self._teacher_use_disable_adapter: tg.execute('init_teacher_model', self._teacher_model_dir) logger.info('Colocated teacher model initialized from %s', self._teacher_model_dir) while iteration < train_iters: # One generation (a single data_source) feeds ``spg`` training steps. prompt_batch = next(self._data_iter) if self.truncation_strategy == 'delete': prompt_batch = self._resample_failed_prompts(prompt_batch, strip_response=False) data_source = self._determine_data_source() if data_source == DataSource.STUDENT: # On-policy: sync the latest weights to the rollout engine and generate. ckpt.sync_weights(merge_and_sync=True) with self._generation_context(tg, ckpt): rollout_batch = self._expand_for_generation(prompt_batch) completions = self._generate(rollout_batch) gkd_samples = self._postprocess_rollout(rollout_batch, completions) source_items = gkd_samples else: # Off-policy (lmbda<1): distill on the dataset's ground-truth responses, # no generation and no weight sync to the rollout engine. gkd_samples = [GKDSample.from_row(item) for item in prompt_batch] source_items = gkd_samples self._maybe_log_completions(gkd_samples if data_source == DataSource.STUDENT else None, gen_step=iteration) # Split one generation into ``spg`` chunks; each chunk is one training step # (same data_source). spg=1 degenerates to a single chunk == the whole batch. # n == global_batch_size * spg: the driver dataloader uses drop_last=True + a # cyclic iterator (see _setup_dataloader) and GKD uses num_generations=1, so n is # always an exact multiple of spg and the spg chunks tile source_items with no # remainder. (max(., 1) only guards the impossible spg > n case.) n = len(source_items) chunk_size = max(n // spg, 1) for step_idx in range(spg): if iteration >= train_iters: break chunk = source_items[step_idx * chunk_size:(step_idx + 1) * chunk_size] if not chunk: break samples = self._encode_rollout_batch(chunk) use_colocated_teacher = self._teacher_use_disable_adapter or (self._teacher_model_dir and not self._teacher_model_server) if self.teacher_replicas: if data_source != DataSource.STUDENT: raise NotImplementedError('Teacher replicas currently require on-policy generation (lmbda=1). ' 'Use a colocated teacher_model for lmbda<1 (off-policy) training.') self._fetch_teacher_from_replicas(chunk, samples) # 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) 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. tg.compute_teacher_logits(dispatch) results = tg.train_step(dispatch) iteration = extract_iteration(results) return iteration def _determine_data_source(self): """Pick the data source for this step (GKD on/off-policy schedule). With probability ``lmbda`` the step is on-policy (the student generates the response); otherwise it is off-policy and we distill on the dataset's ground-truth response. """ if self._data_source_rng.random() < self.lmbda: return DataSource.STUDENT return DataSource.DATASET def _expand_for_generation(self, prompt_batch): """Convert prompt dicts to GKDSample and strip response for generation.""" samples = [GKDSample.from_row(item) for item in prompt_batch] for s in samples: remove_response(s.messages) return samples def _generate(self, batch: List[GKDSample]) -> List[RolloutOutput]: args = self.args request_config = RequestConfig( n=1, max_tokens=args.max_completion_length, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, stop=args.stop_words or None, return_details=True, ) # Convert samples to RolloutInferRequest at the engine boundary # (same pattern as GRPO Ray trainer). requests = [s.to_infer_request() for s in batch] if self._multi_turn_scheduler is not None and not self._enable_server_multi_turn: # Mode A: driver-side trainer loop. run_multi_turn mutates `messages` # in place on RolloutInferRequest objects. invoke_async_hook(self._multi_turn_scheduler.on_trajectory_start(requests)) first_turn = [ RolloutOutput(response=resp) for resp in self._distribute_to_replicas(requests, request_config) ] return run_multi_turn( requests=requests, first_turn_outputs=first_turn, scheduler=self._multi_turn_scheduler, rollout_fn=lambda reqs, cfg: [RolloutOutput(response=resp) for resp in self._distribute_to_replicas(reqs, cfg)], request_config=request_config, max_turns=self._max_turns, ) completions = self._distribute_to_replicas(requests, request_config) return [RolloutOutput(response=resp) for resp in completions] def _postprocess_rollout(self, samples, outputs): """Merge rollout outputs back onto GKDSample (deepcopy to match HF path).""" results = [] for sample, output in zip(samples, outputs): if output is not None: sample = copy.deepcopy(sample) sample.apply_rollout_output(rollout_output=output) results.append(sample) return results @contextmanager def _extended_max_length(self): """Temporarily extend template.max_length by max_completion_length so the prompt+response (token-in-token-out) encodes without truncation. """ template = self.template original = template.max_length template.max_length = original + self.args.max_completion_length try: yield finally: template.max_length = original def _encode_rollout_batch(self, samples: List[GKDSample]): """Encode samples into per-sample worker payloads. Uses the shared ``encode_gkd_samples`` helper (same as HF / Megatron GKD) so OPSD logic is fully encapsulated. Returns per-sample payloads with ``encoded`` (student) and optionally ``teacher_encoded`` (OPSD teacher view). """ from swift.rlhf_trainers.gkd_helpers import encode_gkd_samples template = self.template with self._extended_max_length(): student_list, teacher_list, has_opsd = encode_gkd_samples(samples, template) result = [] for i in range(len(samples)): payload = {'encoded': student_list[i]} if has_opsd: payload['teacher_encoded'] = teacher_list[i] result.append(payload) return result def _collate_for_workers_gkd(self, tg, samples: List[dict], data_source, *, with_teacher: bool): """Driver-side GKD collate: ``List[payload-dict]`` -> ``{dp_rank: [model_inputs]}``. Mirrors the non-Ray GKD ``_encode_samples`` (data_collator on the rank, teacher forward later via prepare_batch). Each micro-batch ``model_inputs`` carries: - student forward tensors (``template.data_collator`` of ``encoded``), - ``data_source`` (SFT gating in loss_func), - ``teacher_model_inputs`` (colocated path): the collated teacher VIEW for the worker teacher forward — OPSD uses ``teacher_encoded``, else ``encoded``, - ``teacher_output`` (teacher-replicas path): the per-sample TeacherOutputs (already on each sample) collated into one batched TeacherOutput. """ from swift.megatron.utils import get_padding_to from .megatron_worker import MegatronWorker template = self.template padding_to = self._padding_to if self._padding_to is not None else get_padding_to(self.args) dp_size = tg.dp_size mbs = int(self.args.micro_batch_size) n = len(samples) if n % dp_size != 0: raise ValueError(f'_collate_for_workers_gkd: batch size {n} not divisible by dp_size {dp_size}.') shard_size = n // dp_size dispatch = DPDispatchedDict() for dp_rank in range(dp_size): shard = samples[dp_rank * shard_size:(dp_rank + 1) * shard_size] micro_batches = [] for i in range(0, len(shard), mbs): 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 with_teacher: has_opsd = chunk[0].get('teacher_encoded') is not None key = 'teacher_encoded' if has_opsd else 'encoded' model_inputs['teacher_model_inputs'] = template.data_collator( batch=[s[key] for s in chunk], padding_to=padding_to) elif chunk[0].get('teacher_output') is not None: # Teacher-replicas path: per-sample TeacherOutputs collated on the driver # (pure tensor ops). The teacher seq length differs from the student under # OPSD, so align by mask (is_opsd) rather than padding to the student length. if getattr(self.args, 'context_parallel_size', 1) > 1: raise ValueError('Standalone teacher replicas (teacher.gpus > 0) do not support ' 'context_parallel_size > 1: per-sample teacher token-logprobs are built ' 'from raw sequence lengths and cannot be CP-sharded to align with the ' 'student. Use a colocated teacher_model for CP>1.') has_opsd = any(s.get('teacher_encoded') is not None for s in chunk) model_inputs['teacher_output'] = MegatronWorker._collate_teacher_outputs( [s['teacher_output'] for s in chunk], self.device, padding_free=template.padding_free, target_seq_len=model_inputs['labels'].shape[-1], is_opsd=has_opsd) micro_batches.append(model_inputs) dispatch[dp_rank] = micro_batches return dispatch def _fetch_teacher_from_replicas(self, gkd_samples: List[GKDSample], samples): """Fetch teacher logprobs from Ray teacher replicas. Uses to_infer_request() + teacher_messages replacement (OPSD) to build unified RolloutInferRequest objects, matching HF GKD's _build_teacher_requests. """ topk = self.gkd_logits_topk assert topk is not None, 'gkd_logits_topk must be set when using teacher replicas' requests = [] teacher_encodeds = [] # teacher-side encoded (OPSD) or None (non-OPSD) for s, sample in zip(gkd_samples, samples): req = s.to_infer_request() teacher_encoded = sample.get('teacher_encoded') if s.teacher_messages: req.messages = s.teacher_messages teacher_encodeds.append(teacher_encoded) else: teacher_encodeds.append(None) requests.append(req) request_config = RequestConfig(prompt_logprobs=topk, max_tokens=1, temperature=0.0) replicas = self.teacher_replicas n = len(replicas) chunk_size = (len(requests) + n - 1) // n refs = [] for i, replica in enumerate(replicas): shard = requests[i * chunk_size:(i + 1) * chunk_size] if not shard: continue refs.append(replica.generate(shard, request_config)) parts = ray.get(refs) responses = [] for p in parts: responses.extend(p) for sample, response, t_encoded in zip(samples, responses, teacher_encodeds): parsed = parse_prompt_logprobs(response, topk=topk) 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) @staticmethod def _build_per_sample_teacher_output(parsed, encoded, topk, labels=None): """Build a per-sample TeacherOutput from parsed prompt logprobs. For OPSD, ``encoded`` is the teacher-side encoding and ``labels`` are its labels; together they let ``extract_active`` mask-align the shared response. For non-OPSD, teacher and student share the same encoding, so we fall back to ``encoded['labels']`` when ``labels`` is not provided. """ if labels is None: labels = encoded.get('labels') lps, ixs = parsed input_ids = encoded['input_ids'] seq_len = len(input_ids) if isinstance(input_ids, list) else input_ids.shape[-1] 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) 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)) 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 return TeacherOutput(**kwargs)