wehub-resource-sync
381 lines
19 KiB
Python
381 lines
19 KiB
Python
# 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)
|