diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 000000000..176a458f9 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +* text=auto diff --git a/README.md b/README.md index c82f1c15c..41745a0de 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,7 @@ sh INSTALL_MEGATRON.sh | pp/tp/cp MoE finetuning | megatron | [Script](cookbook/megatron/tp_moe.py) | | Multimodal FSDP finetuning | transformers | [Script](cookbook/mm/fsdp2.py) | | GRPO RL training | megatron | [Script](cookbook/rl/grpo/grpo.py) | +| PPO RL training | transformers | [Script](cookbook/rl/ppo/ppo.py) | | GRPO Multimodal RL training | megatron | [Script](cookbook/rl/grpo/grpo_mm.py) | | GRPO Math RL training | megatron | [Script](cookbook/rl/grpo/short_math_grpo.py) | | DPO full-parameter training | transformers | [Script](cookbook/rl/dpo/dpo_full.py) | diff --git a/README_ZH.md b/README_ZH.md index 31cab0f93..2b6a0ba61 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -87,6 +87,7 @@ sh INSTALL_MEGATRON.sh | pp/tp/cp MoE 微调 | megatron | [脚本](cookbook/megatron/tp_moe.py) | | 多模态 FSDP 微调 | transformers | [脚本](cookbook/mm/fsdp2.py) | | GRPO 强化学习训练 | megatron | [脚本](cookbook/rl/grpo/grpo.py) | +| PPO 强化学习训练 | transformers | [脚本](cookbook/rl/ppo/ppo.py) | | GRPO 多模态强化学习训练 | megatron | [脚本](cookbook/rl/grpo/grpo_mm.py) | | GRPO 数学强化学习训练 | megatron | [脚本](cookbook/rl/grpo/short_math_grpo.py) | | DPO 全参数训练 | transformers | [脚本](cookbook/rl/dpo/dpo_full.py) | diff --git a/cookbook/rl/ppo/ppo.py b/cookbook/rl/ppo/ppo.py new file mode 100644 index 000000000..aedf59dcb --- /dev/null +++ b/cookbook/rl/ppo/ppo.py @@ -0,0 +1,240 @@ +"""Standard PPO training on GSM8K with a LoRA policy and full-parameter critic. + +The first implementation supports the Transformers/Accelerate-FSDP backend. Policy, +critic, and vLLM sampler use separate GPU groups. The frozen policy base model is +used as the reference policy. +""" +import random +from typing import Any, Dict, List, Tuple + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GAEAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.cli import CLI +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.metric import CompletionRewardMetric, PPOMetric, PPOValueMetric +from twinkle.model import TransformersModel, TransformersValueModel +from twinkle.processor import InputProcessor +from twinkle.preprocessor.llm import GSM8KProcessor +from twinkle.reward import GSM8KAccuracyReward, GSM8KFormatReward +from twinkle.sampler import vLLMSampler + +logger = get_logger() +args = CLI.from_args() + +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.5-4B' +POLICY_GPUS = args.infra.model_gpus or 4 +CRITIC_GPUS = args.infra.critic_model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = POLICY_GPUS + CRITIC_GPUS + SAMPLER_GPUS +NUM_GENERATIONS = args.rl.num_generations or 4 +MAX_NEW_TOKENS = args.sampling.max_tokens or 1024 +POLICY_LR = args.optimizer.learning_rate or 1e-5 +CRITIC_LR = args.rl.critic_learning_rate +MAX_STEPS = args.training.max_steps or 200 +BATCH_SIZE = args.training.batch_size or 4 +MINI_BATCH_SIZE = args.training.mini_batch_size or 4 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 1 +# Number of policy/value updates over each rollout batch. Reuse the common +# training argument whaohile preserving PPO's historical default. +PPO_EPOCHS = args.training.num_train_epochs if args.training.num_train_epochs is not None else 4 +SAVE_STEPS = args.training.save_steps or 50 +ADAPTER_NAME = args.lora.adapter_name or 'default' + + +def create_gsm8k_dataset(): + dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=400) + dataset.map(GSM8KProcessor()) + dataset.encode(add_generation_prompt=True) + return dataset + + +def compute_rewards(trajectories: List[Dict[str, Any]]) -> Tuple[List[float], List[float], List[float]]: + accuracy = GSM8KAccuracyReward()(trajectories) + formatting = GSM8KFormatReward()(trajectories) + return [a + f for a, f in zip(accuracy, formatting)], formatting, accuracy + + +def response_rows(full_values, trajectories) -> List[List[float]]: + """Extract response-token rows from collected model outputs.""" + import torch + + value_rows = [] + tensors = full_values if isinstance(full_values, list) else [full_values] + for tensor in tensors: + if tensor is None: + continue + tensor = torch.as_tensor(tensor) + if tensor.dim() == 1: + tensor = tensor.unsqueeze(0) + value_rows.extend(tensor) + if len(value_rows) != len(trajectories): + raise ValueError(f'model output batch mismatch: {len(value_rows)} rows for {len(trajectories)} trajectories') + + rows = [] + for value_row, trajectory in zip(value_rows, trajectories): + mask = torch.as_tensor(trajectory['labels'], device=value_row.device) != -100 + rows.append(value_row[:mask.numel()][mask].detach().float().cpu().tolist()) + return rows + + +def main(): + critic_start = POLICY_GPUS + sampler_start = POLICY_GPUS + CRITIC_GPUS + groups = [ + DeviceGroup(name='policy', ranks=list(range(POLICY_GPUS)), device_type='GPU'), + DeviceGroup(name='critic', ranks=list(range(critic_start, sampler_start)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(sampler_start, NUM_GPUS)), device_type='GPU'), + ] + policy_mesh = DeviceMesh.from_sizes(world_size=POLICY_GPUS, fsdp_size=POLICY_GPUS) + critic_mesh = DeviceMesh.from_sizes(world_size=CRITIC_GPUS, fsdp_size=CRITIC_GPUS) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=groups, lazy_collect=False) + + policy = TransformersModel( + model_id=MODEL_ID, device_mesh=policy_mesh, remote_group='policy') + lora_config = LoraConfig( + target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'], + r=32, + lora_alpha=64, + lora_dropout=0.05, + ) + policy.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=1) + policy.set_optimizer('AdamW', lr=POLICY_LR) + policy.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + policy.set_loss( + 'PPOLoss', + epsilon=args.loss.epsilon, + entropy_coef=args.loss.entropy_coef, + loss_agg_mode='token-mean', + ) + policy.add_metric(PPOMetric, epsilon=args.loss.epsilon) + policy.set_processor(InputProcessor) + policy.set_template('Qwen3_5Template', model_id=MODEL_ID) + + critic = TransformersValueModel( + model_id=MODEL_ID, device_mesh=critic_mesh, remote_group='critic') + critic.set_optimizer('AdamW', lr=CRITIC_LR) + critic.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + critic.set_loss('PPOValueLoss', epsilon=args.loss.value_clip) + critic.add_metric(PPOValueMetric, epsilon=args.loss.value_clip) + critic.set_processor(InputProcessor) + critic.set_template('Qwen3_5Template', model_id=MODEL_ID) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 400 + MAX_NEW_TOKENS, + 'max_lora_rank': 32, + 'enable_lora': True, + 'tensor_parallel_size': 1, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID) + checkpoint_manager = CheckpointEngineManager(model=policy, sampler=sampler) + dataloader = DataLoader( + dataset=create_gsm8k_dataset, + batch_size=BATCH_SIZE, + min_batch_size=BATCH_SIZE, + device_mesh=policy_mesh, + remote_group='policy', + ) + gae = GAEAdvantage(args.rl.gamma, args.rl.gae_lambda, args.rl.normalize_advantages) + reward_metric = CompletionRewardMetric() + sampling_params = SamplingParams(max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1) + + optim_step = 0 + rollout_step = 0 + logger.info(get_device_placement()) + while optim_step < MAX_STEPS: + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + reward_metric.reset() + prompts = batch if isinstance(batch, list) else [batch] + checkpoint_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + expanded = [prompt for prompt in prompts for _ in range(NUM_GENERATIONS)] + samples = sampler.sample(expanded, sampling_params) + + trajectories, old_logps, lengths = [], [], [] + for response in samples: + for sequence in response.sequences: + trajectories.append(sequence.new_input_feature) + old_logps.append([entry[0][1] for entry in sequence.logprobs]) + lengths.append(len(sequence.tokens)) + rewards, format_rewards, accuracy_rewards = compute_rewards(trajectories) + reward_metric.accumulate( + completion_lengths=lengths, + rewards={'total': rewards, 'format': format_rewards, 'accuracy': accuracy_rewards}, + ) + + reference = policy.forward_only(inputs=trajectories, disable_lora=True) + ref_logps = response_rows(reference['logps'], trajectories) + critic_outputs = critic.forward_only(inputs=trajectories) + old_values = response_rows(critic_outputs['values'], trajectories) + token_rewards = gae.build_token_rewards( + rewards, lengths, old_logps=old_logps, ref_logps=ref_logps, kl_coef=args.rl.kl_coef) + max_len = max(lengths) + padded_rewards = [row + [0.0] * (max_len - len(row)) for row in token_rewards] + padded_values = [row + [0.0] * (max_len - len(row)) for row in old_values] + masks = [[True] * length + [False] * (max_len - length) for length in lengths] + advantages, returns = gae(padded_rewards, padded_values, masks=masks) + advantages = [advantages[i, :length].tolist() for i, length in enumerate(lengths)] + returns = [returns[i, :length].tolist() for i, length in enumerate(lengths)] + + indices = list(range(len(trajectories))) + for _ in range(PPO_EPOCHS): + random.shuffle(indices) + for start in range(0, len(indices), MINI_BATCH_SIZE): + chosen = indices[start:start + MINI_BATCH_SIZE] + mb_inputs = [trajectories[i] for i in chosen] + mb_old_logps = [old_logps[i] for i in chosen] + mb_old_values = [old_values[i] for i in chosen] + mb_advantages = [advantages[i] for i in chosen] + mb_returns = [returns[i] for i in chosen] + policy.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + advantages=mb_advantages, + micro_batch_size=MICRO_BATCH_SIZE, + ) + policy.clip_grad_and_step() + critic.forward_backward( + inputs=mb_inputs, + old_values=mb_old_values, + returns=mb_returns, + advantages=mb_advantages, + micro_batch_size=MICRO_BATCH_SIZE, + ) + critic.clip_grad_and_step() + optim_step += 1 + if optim_step % SAVE_STEPS == 0: + policy.save(f'ppo-policy-checkpoint-{optim_step}') + critic.save(f'ppo-critic-checkpoint-{optim_step}') + if optim_step >= MAX_STEPS: + break + if optim_step >= MAX_STEPS: + break + + logs = reward_metric.calculate() + logs.update(policy.calculate_metric(is_training=True)) + logs.update(critic.calculate_metric(is_training=True)) + rollout_step += 1 + logger.info(f'[Rollout {rollout_step}, optim step {optim_step}/{MAX_STEPS}] {logs}') + + policy.save('ppo-policy-final') + critic.save('ppo-critic-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/ppo/ppo.sh b/cookbook/rl/ppo/ppo.sh new file mode 100644 index 000000000..015cb5ac4 --- /dev/null +++ b/cookbook/rl/ppo/ppo.sh @@ -0,0 +1,29 @@ +#!/bin/sh +set -eu + +# Standard PPO on GSM8K via Ray. +# Transformers/Accelerate-FSDP: 4 policy + 4 full-parameter critic + 4 sampler GPUs. +# Override any option after the defaults, for example: +# sh ppo.sh --max-steps 20 --num-train-epochs 1 + +python ppo.py \ + --model-id ms://Qwen/Qwen3.5-4B \ + --model-gpus 4 \ + --critic-model-gpus 4 \ + --sampler-gpus 4 \ + --num-generations 2 \ + --max-tokens 1024 \ + --batch-size 4 \ + --mini-batch-size 4 \ + --micro-batch-size 1 \ + --num-train-epochs 4 \ + --gamma 1.0 \ + --gae-lambda 0.95 \ + --kl-coef 0.01 \ + --value-clip 0.2 \ + --lr 1e-5 \ + --critic-learning-rate 1e-5 \ + --max-steps 200 \ + --save-steps 50 \ + --adapter-name default \ + "$@" diff --git a/src/twinkle/advantage/__init__.py b/src/twinkle/advantage/__init__.py index cbf5565d1..519dd1fe4 100644 --- a/src/twinkle/advantage/__init__.py +++ b/src/twinkle/advantage/__init__.py @@ -1,10 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .base import Advantage +from .gae import GAEAdvantage from .grpo import GRPOAdvantage from .rloo import RLOOAdvantage __all__ = [ 'Advantage', + 'GAEAdvantage', 'GRPOAdvantage', 'RLOOAdvantage', ] diff --git a/src/twinkle/advantage/gae.py b/src/twinkle/advantage/gae.py new file mode 100644 index 000000000..747bcb429 --- /dev/null +++ b/src/twinkle/advantage/gae.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING, List, Optional, Tuple, Union + +from .base import Advantage + +if TYPE_CHECKING: + import torch + + +class GAEAdvantage(Advantage): + """Token-level generalized advantage estimation for terminal completions.""" + + def __init__(self, gamma: float = 1.0, gae_lambda: float = 0.95, normalize: bool = True): + if not 0.0 <= gamma <= 1.0: + raise ValueError('gamma must be in [0, 1]') + if not 0.0 <= gae_lambda <= 1.0: + raise ValueError('gae_lambda must be in [0, 1]') + self.gamma = gamma + self.gae_lambda = gae_lambda + self.normalize = normalize + + @staticmethod + def build_token_rewards( + rewards: Union['torch.Tensor', List[float]], + lengths: List[int], + *, + old_logps: Optional[List[List[float]]] = None, + ref_logps: Optional[List[List[float]]] = None, + kl_coef: float = 0.0, + ) -> List[List[float]]: + import torch + + rewards = torch.as_tensor(rewards, dtype=torch.float32).flatten().tolist() + if len(rewards) != len(lengths): + raise ValueError('rewards and lengths must have the same batch size') + if (old_logps is None) != (ref_logps is None): + raise ValueError('old_logps and ref_logps must be provided together') + + token_rewards = [] + for i, (reward, length) in enumerate(zip(rewards, lengths)): + if length <= 0: + raise ValueError('completion lengths must be positive') + values = [0.0] * length + if old_logps is not None: + if len(old_logps[i]) != length or len(ref_logps[i]) != length: + raise ValueError(f'log-prob length mismatch at sample {i}') + values = [-kl_coef * (float(old) - float(ref)) for old, ref in zip(old_logps[i], ref_logps[i])] + values[-1] += float(reward) + token_rewards.append(values) + return token_rewards + + def __call__( + self, + rewards: Union['torch.Tensor', List[List[float]]], + values: Union['torch.Tensor', List[List[float]]], + *, + masks: Optional[Union['torch.Tensor', List[List[bool]]]] = None, + normalize: Optional[bool] = None, + **kwargs, + ) -> Tuple['torch.Tensor', 'torch.Tensor']: + import torch + + rewards = torch.as_tensor(rewards, dtype=torch.float32) + values = torch.as_tensor(values, dtype=torch.float32, device=rewards.device) + if rewards.dim() == 1: + rewards = rewards.unsqueeze(0) + if values.dim() == 1: + values = values.unsqueeze(0) + if rewards.shape != values.shape: + raise ValueError(f'rewards and values must have identical shapes, got {rewards.shape} and {values.shape}') + + if masks is None: + masks = torch.ones_like(rewards, dtype=torch.bool) + else: + masks = torch.as_tensor(masks, dtype=torch.bool, device=rewards.device) + if masks.shape != rewards.shape: + raise ValueError('masks must have the same shape as rewards') + + advantages = torch.zeros_like(rewards) + for batch_idx in range(rewards.shape[0]): + valid = masks[batch_idx].nonzero(as_tuple=True)[0] + last_gae = rewards.new_zeros(()) + for j in range(len(valid) - 1, -1, -1): + pos = valid[j] + if j + 1 < len(valid): + next_value = values[batch_idx, valid[j + 1]] + else: + next_value = rewards.new_zeros(()) + delta = rewards[batch_idx, pos] + self.gamma * next_value - values[batch_idx, pos] + last_gae = delta + self.gamma * self.gae_lambda * last_gae + advantages[batch_idx, pos] = last_gae + + returns = advantages + values + should_normalize = self.normalize if normalize is None else normalize + if should_normalize: + valid_advantages = advantages[masks] + if valid_advantages.numel() > 1: + mean = valid_advantages.mean() + std = valid_advantages.std(unbiased=False) + advantages = torch.where(masks, (advantages - mean) / (std + 1e-8), advantages) + advantages = advantages.masked_fill(~masks, 0.0) + returns = returns.masked_fill(~masks, 0.0) + return advantages, returns diff --git a/src/twinkle/cli/cli.py b/src/twinkle/cli/cli.py index a12594d00..5bc6aed0e 100644 --- a/src/twinkle/cli/cli.py +++ b/src/twinkle/cli/cli.py @@ -134,6 +134,7 @@ class LossArgs: beta: float = 0.1 sft_weight: float = 1.0 entropy_coef: float = 0.0 + value_clip: float = 0.2 ignore_index: int = -100 @@ -169,6 +170,7 @@ class InfraArgs: model_gpus: int | None = None sampler_gpus: int | None = None ref_model_gpus: int | None = None + critic_model_gpus: int | None = None world_size: int | None = None dp_size: int | None = None fsdp_size: int | None = None @@ -205,6 +207,11 @@ class RLArgs: gkd_temperature: float = 1.0 gkd_topk: int = 64 router_replay_mode: Literal['disabled', 'R2', 'R3'] = 'disabled' + gamma: float = 1.0 + gae_lambda: float = 0.95 + kl_coef: float = 0.0 + normalize_advantages: bool = True + critic_learning_rate: float = 1e-5 @dataclass diff --git a/src/twinkle/data_format/output.py b/src/twinkle/data_format/output.py index 6a248649b..62a9e93eb 100644 --- a/src/twinkle/data_format/output.py +++ b/src/twinkle/data_format/output.py @@ -20,11 +20,13 @@ class ModelOutput(TypedDict, total=False): loss: The loss calculated by the model. logps: The log-probabilities of correct tokens by the model. num_tokens: The token denominator associated with ``loss``. + values: Per-token scalar value estimates used by critic models. embeddings: The embeddings output by the model, used be embedding task. """ logits: Optional[OutputType] loss: Optional[OutputType] logps: Optional[OutputType] + values: Optional[OutputType] num_tokens: Optional[OutputType] embeddings: Optional[OutputType] routed_experts: Optional[OutputType] diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 29485e052..6eb991383 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -16,6 +16,31 @@ class RayHelper: _remote_components: Dict[str, Any] = {} + @staticmethod + def _get_ray_custom_resources(device_groups: List[DeviceGroup]) -> Dict[str, float]: + """Return custom accelerator resources needed for local Ray startup. + + Ray discovers CUDA GPUs itself. Other accelerators, such as Ascend + NPUs, need to be registered as Ray custom resources so that placement + groups can request them. + """ + device_types = {group.device_type.upper() for group in device_groups} - {'CPU'} + + # ResourceManager supports one accelerator type per run. Only NPU + # currently needs an explicit Ray custom-resource registration. + if device_types != {'NPU'}: + return {} + + try: + import torch + + npu = getattr(torch, 'npu', None) + npu_count = npu.device_count() if npu is not None and npu.is_available() else 0 + except (ImportError, AttributeError, RuntimeError): + return {} + + return {'NPU': float(npu_count)} if npu_count > 0 else {} + @staticmethod def init_registry(): if RayHelper._registry is not None: @@ -71,7 +96,8 @@ def initialize(nproc_per_node: int, ncpu_proc_per_node: int, device_groups: List import ray RayHelper.device_groups = device_groups if not RayHelper.ray_inited(): - ray.init(ignore_reinit_error=True) + resources = RayHelper._get_ray_custom_resources(device_groups) + ray.init(ignore_reinit_error=True, resources=resources or None) if RayHelper.resource_manager is None: # Resource manager initializes only once in the pipeline process. diff --git a/src/twinkle/loss/__init__.py b/src/twinkle/loss/__init__.py index d062a3fab..18b93f5ae 100644 --- a/src/twinkle/loss/__init__.py +++ b/src/twinkle/loss/__init__.py @@ -4,10 +4,11 @@ from .cross_entropy import CrossEntropyLoss from .dpo import CPOLoss, DPOLoss, ORPOLoss, SimPOLoss from .gkd import GKDLoss -from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, SAPOLoss +from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, PPOLoss, SAPOLoss from .infonce import InfonceLoss from .liger_fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss from .mse import MSELoss +from .value import PPOValueLoss torch_loss_mapping = { 'mse': MSELoss, @@ -18,6 +19,8 @@ 'gkd': GKDLoss, # RL losses 'grpo': GRPOLoss, + 'ppo': PPOLoss, + 'ppo_value': PPOValueLoss, 'gspo': GSPOLoss, 'sapo': SAPOLoss, 'cispo': CISPOLoss, diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py index 5fd046ae7..c7f112c42 100644 --- a/src/twinkle/loss/base.py +++ b/src/twinkle/loss/base.py @@ -7,6 +7,7 @@ class Loss: require_logits = False require_entropy = False require_logps = True + require_values = False def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput: ... diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 781b22060..81e0b9208 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -1,6 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import numpy as np -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union from twinkle.data_format import LossOutput from twinkle.loss.base import Loss @@ -304,6 +304,39 @@ def __call__( return LossOutput(loss=loss, num_tokens=0) +class PPOLoss(GRPOLoss): + """PPO clipped policy loss. + + PPO and GRPO share the same clipped policy objective. The algorithms differ + in how advantages are produced, so this class provides the PPO-facing name + while using PPO-specific loss aggregation. + + Args: + loss_agg_mode: ``'token-mean'`` averages over every valid response + token. ``'seq-mean-token-mean'`` averages tokens within each + sequence and then averages the sequences. + """ + + _LOSS_AGG_MODES = {'token-mean', 'seq-mean-token-mean'} + + def __init__(self, loss_agg_mode: Literal['token-mean', 'seq-mean-token-mean'] = 'token-mean', **kwargs): + super().__init__(**kwargs) + if loss_agg_mode not in self._LOSS_AGG_MODES: + raise ValueError(f'Unsupported PPO loss_agg_mode: {loss_agg_mode}. ' + f'Expected one of {sorted(self._LOSS_AGG_MODES)}.') + self.loss_agg_mode = loss_agg_mode + + def _aggregate_loss( + self, + per_token_loss: 'torch.Tensor', + loss_mask: 'torch.Tensor', + **kwargs, + ) -> 'torch.Tensor': + if self.loss_agg_mode == 'token-mean': + return (per_token_loss * loss_mask).sum() / loss_mask.sum().clamp(min=1.0) + return super()._aggregate_loss(per_token_loss, loss_mask, **kwargs) + + class GSPOLoss(GRPOLoss): """ GRPO with sequence-level importance sampling. diff --git a/src/twinkle/loss/value.py b/src/twinkle/loss/value.py new file mode 100644 index 000000000..563cdaef5 --- /dev/null +++ b/src/twinkle/loss/value.py @@ -0,0 +1,62 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +from twinkle.data_format import LossOutput +from twinkle.loss.base import Loss +from twinkle.loss.grpo import GRPOLoss + +if TYPE_CHECKING: + import torch + + +class PPOValueLoss(Loss): + """Clipped PPO value-function loss over response tokens.""" + + require_logps = False + require_values = True + + def __init__(self, epsilon: float = 0.2, ignore_index: int = -100, **kwargs): + self.epsilon = epsilon + self.ignore_index = ignore_index + self._aligner = GRPOLoss(ignore_index=ignore_index) + + def __call__( + self, + inputs: Dict, + outputs: Dict, + *, + old_values: Optional[Union['torch.Tensor', List, np.ndarray]] = None, + returns: Optional[Union['torch.Tensor', List, np.ndarray]] = None, + **kwargs, + ) -> LossOutput: + import torch + + labels = inputs.get('labels') + assert labels is not None, "inputs must contain 'labels'" + labels = torch.as_tensor(labels) + if labels.dim() == 1: + labels = labels.unsqueeze(0) + mask = (labels != self.ignore_index).bool() + + values = outputs.get('values') + assert values is not None, "outputs must contain 'values'" + if values.dim() == 3 and values.shape[-1] == 1: + values = values.squeeze(-1) + if values.dim() == 1: + values = values.unsqueeze(0) + if values.shape != mask.shape: + raise AssertionError(f'values/mask shape mismatch: values={tuple(values.shape)} mask={tuple(mask.shape)}') + assert old_values is not None, 'old_values are required for PPO value clipping' + assert returns is not None, 'returns are required for PPO value loss' + + old_values = self._aligner._pad_and_align_to_batch(old_values, mask, values.device, values.dtype) + returns = self._aligner._pad_and_align_to_batch(returns, mask, values.device, values.dtype) + + clipped_values = old_values + torch.clamp(values - old_values, -self.epsilon, self.epsilon) + loss_unclipped = (values - returns).square() + loss_clipped = (clipped_values - returns).square() + per_token_loss = 0.5 * torch.maximum(loss_unclipped, loss_clipped) + mask_f = mask.to(values.dtype) + loss = (per_token_loss * mask_f).sum() / mask_f.sum().clamp(min=1.0) + return LossOutput(loss=loss, num_tokens=0) diff --git a/src/twinkle/metric/__init__.py b/src/twinkle/metric/__init__.py index baeb6c1c9..f6ac5120d 100644 --- a/src/twinkle/metric/__init__.py +++ b/src/twinkle/metric/__init__.py @@ -4,6 +4,7 @@ from .completion_and_reward import CompletionRewardMetric from .dpo import DPOMetric from .embedding import EmbeddingMetric -from .grpo import CISPOMetric, GRPOMetric, GSPOMetric +from .grpo import CISPOMetric, GRPOMetric, GSPOMetric, PPOMetric from .loss import LossMetric +from .ppo import PPOValueMetric from .train_metric import TrainMetric diff --git a/src/twinkle/metric/grpo.py b/src/twinkle/metric/grpo.py index bd85aab67..e3eaacd20 100644 --- a/src/twinkle/metric/grpo.py +++ b/src/twinkle/metric/grpo.py @@ -373,6 +373,10 @@ def calculate(self) -> Dict[str, Any]: return results +class PPOMetric(GRPOMetric): + """PPO policy metric; shares token-level ratio and clipping statistics with GRPO.""" + + class GSPOMetric(GRPOMetric): """GRPOMetric variant for GSPO: clip applies to per-sequence geometric-mean ratio.""" diff --git a/src/twinkle/metric/ppo.py b/src/twinkle/metric/ppo.py new file mode 100644 index 000000000..ebe90d3fd --- /dev/null +++ b/src/twinkle/metric/ppo.py @@ -0,0 +1,84 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import Any, Dict, List, Union + +from twinkle.data_format import InputFeature, ModelOutput +from .base import Metric + + +class PPOValueMetric(Metric): + """Aggregate PPO critic statistics over valid response tokens.""" + + def __init__(self, device_mesh=None, process_group=None, epsilon: float = 0.2, ignore_index: int = -100, **kwargs): + super().__init__(device_mesh, process_group, **kwargs) + self.epsilon = epsilon + self.ignore_index = ignore_index + self.reset() + + def reset(self): + self.records = [] + + def accumulate(self, + inputs: Union[InputFeature, List[InputFeature]], + outputs: ModelOutput, + *, + old_values=None, + returns=None, + advantages=None, + **kwargs): + import torch + + from twinkle.utils.transformers_utils import align_logps_to_mask + + if outputs is None or outputs.get('values') is None or old_values is None or returns is None: + return + inputs_list = inputs if isinstance(inputs, list) else [inputs] + values_list = outputs['values'] if isinstance(outputs['values'], list) else [outputs['values']] + cursor = 0 + for mb_input, values in zip(inputs_list, values_list): + labels = torch.as_tensor(mb_input['labels'], device=values.device) + if labels.dim() == 1: + labels = labels.unsqueeze(0) + if values.dim() == 1: + values = values.unsqueeze(0) + mask = labels != self.ignore_index + batch_size = labels.shape[0] + old = align_logps_to_mask(old_values[cursor:cursor + batch_size], mask, values.dtype) + target = align_logps_to_mask(returns[cursor:cursor + batch_size], mask, values.dtype) + adv = ( + align_logps_to_mask(advantages[cursor:cursor + + batch_size], mask, values.dtype) if advantages is not None else None) + cursor += batch_size + if old is None or target is None or not mask.any(): + continue + clipped = old + (values - old).clamp(-self.epsilon, self.epsilon) + self.records.append({ + 'values': values[mask].detach().float().cpu().tolist(), + 'returns': target[mask].detach().float().cpu().tolist(), + 'advantages': adv[mask].detach().float().cpu().tolist() if adv is not None else [], + 'clipped': ((values - old).abs() > self.epsilon)[mask].detach().float().cpu().tolist(), + 'clipped_values': clipped[mask].detach().float().cpu().tolist(), + }) + + def calculate(self) -> Dict[str, Any]: + import torch + + records = self.gather_results(self.records) + self.reset() + if not records: + return {} + values = torch.tensor([v for record in records for v in record['values']]) + returns = torch.tensor([v for record in records for v in record['returns']]) + clipped = torch.tensor([v for record in records for v in record['clipped']]) + advantages = torch.tensor([v for record in records for v in record['advantages']]) + return_var = returns.var(unbiased=False) + explained_variance = 1.0 - (returns - values).var(unbiased=False) / return_var.clamp(min=1e-8) + result = { + 'train/value_mean': values.mean().item(), + 'train/return_mean': returns.mean().item(), + 'train/value_clip_ratio': clipped.mean().item(), + 'train/explained_variance': explained_variance.item(), + } + if advantages.numel(): + result['train/advantage_mean'] = advantages.mean().item() + result['train/advantage_std'] = advantages.std(unbiased=False).item() + return result diff --git a/src/twinkle/model/__init__.py b/src/twinkle/model/__init__.py index 88f544d6c..91401a085 100644 --- a/src/twinkle/model/__init__.py +++ b/src/twinkle/model/__init__.py @@ -6,12 +6,12 @@ if TYPE_CHECKING: from .base import TwinkleModel from .megatron import MegatronModel, MultiLoraMegatronModel - from .transformers import MultiLoraTransformersModel, TransformersModel + from .transformers import MultiLoraTransformersModel, TransformersModel, TransformersValueModel else: _import_structure = { 'base': ['TwinkleModel'], - 'transformers': ['TransformersModel', 'MultiLoraTransformersModel'], + 'transformers': ['TransformersModel', 'MultiLoraTransformersModel', 'TransformersValueModel'], 'megatron': ['MegatronModel', 'MultiLoraMegatronModel'], } diff --git a/src/twinkle/model/transformers/__init__.py b/src/twinkle/model/transformers/__init__.py index 9ffe98661..cd3775ace 100644 --- a/src/twinkle/model/transformers/__init__.py +++ b/src/twinkle/model/transformers/__init__.py @@ -1,3 +1,4 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .multi_lora_transformers import MultiLoraTransformersModel from .transformers import TransformersModel +from .value_model import TransformersValueModel diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 017a515b3..ffbcebccd 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -466,6 +466,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec loss_require_logits = getattr(loss_instance, 'require_logits', False) loss_require_entropy = getattr(loss_instance, 'require_entropy', False) loss_require_logps = getattr(loss_instance, 'require_logps', True) + loss_require_values = getattr(loss_instance, 'require_values', False) assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' inputs: Dict[str, Any] = processor( inputs, @@ -504,6 +505,9 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec else: outputs['logps'] = selective_log_softmax(logits, masked_labels) del logits + if loss_require_values: + values = outputs['logits'] + outputs['values'] = values.squeeze(-1) if values.shape[-1] == 1 else values outputs['past_key_values'] = None if not (return_logits or loss_require_logits): outputs['logits'] = None @@ -557,6 +561,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T loss_require_logits = getattr(loss_instance, 'require_logits', False) loss_require_entropy = getattr(loss_instance, 'require_entropy', False) loss_require_logps = getattr(loss_instance, 'require_logps', True) + loss_require_values = getattr(loss_instance, 'require_values', False) inputs: Dict[str, Any] = processor( inputs, sp_strategy=self.sp_strategy, @@ -598,6 +603,9 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T else: outputs['logps'] = selective_log_softmax(logits, masked_labels) del logits + if loss_require_values: + values = outputs['logits'] + outputs['values'] = values.squeeze(-1) if values.shape[-1] == 1 else values outputs['past_key_values'] = None if not (return_logits or loss_require_logits): outputs['logits'] = None diff --git a/src/twinkle/model/transformers/value_model.py b/src/twinkle/model/transformers/value_model.py new file mode 100644 index 000000000..110650f6d --- /dev/null +++ b/src/twinkle/model/transformers/value_model.py @@ -0,0 +1,39 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from torch import nn +from typing import Optional + +from twinkle import DeviceMesh, remote_class +from .transformers import TransformersModel + + +@remote_class() +class TransformersValueModel(TransformersModel): + """Transformers causal-LM backbone with a scalar value head.""" + + def __init__(self, *args, device_mesh: Optional[DeviceMesh] = None, **kwargs): + super().__init__(*args, device_mesh=device_mesh, **kwargs) + output_head = self.model.get_output_embeddings() + if output_head is None or not hasattr(output_head, 'in_features'): + raise ValueError('The model must expose a linear output embedding to construct a value head') + value_head = nn.Linear( + output_head.in_features, 1, bias=True, device=output_head.weight.device, dtype=output_head.weight.dtype) + nn.init.zeros_(value_head.weight) + nn.init.zeros_(value_head.bias) + self.model.set_output_embeddings(value_head) + self.model.config.tie_word_embeddings = False + + def add_adapter_to_model(self, *args, **kwargs): + raise NotImplementedError('PPO critic LoRA is not supported; train the critic as a full-parameter model') + + def forward(self, *, inputs, **kwargs): + task = kwargs.get('task', 'causal_lm') + if task != 'causal_lm': + raise ValueError('TransformersValueModel only supports task="causal_lm"') + return super().forward(inputs=inputs, **kwargs) + + def forward_only(self, *, inputs, **kwargs): + task = kwargs.get('task', 'causal_lm') + if task != 'causal_lm': + raise ValueError('TransformersValueModel only supports task="causal_lm"') + return super().forward_only(inputs=inputs, **kwargs) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 0a819e1ca..483fb821d 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -624,7 +624,7 @@ def unpack_packed_sequences( # Collect output keys to unpack: (key, pad_value) output_keys = [] - for key, pad_val in [('logps', 0), ('entropies', 0), ('logits', 0)]: + for key, pad_val in [('logps', 0), ('values', 0), ('entropies', 0), ('logits', 0)]: if outputs and outputs.get(key) is not None: output_keys.append((key, pad_val)) diff --git a/src/twinkle/utils/nccl_safe.py b/src/twinkle/utils/nccl_safe.py index b22b10137..f1e6d4095 100644 --- a/src/twinkle/utils/nccl_safe.py +++ b/src/twinkle/utils/nccl_safe.py @@ -78,6 +78,7 @@ def __init__(self, loss_instance): self.require_logps = getattr(loss_instance, 'require_logps', True) self.require_entropy = getattr(loss_instance, 'require_entropy', False) self.require_logits = getattr(loss_instance, 'require_logits', False) + self.require_values = getattr(loss_instance, 'require_values', False) self.reduction = getattr(loss_instance, 'reduction', 'mean') self._nccl_safe_wrapped = True @@ -102,7 +103,7 @@ def _zero_loss(outputs) -> 'LossOutput': """ import torch if isinstance(outputs, dict): - for key in ('logps', 'logits', 'loss'): + for key in ('logps', 'values', 'logits', 'loss'): t = outputs.get(key) if t is not None and isinstance(t, torch.Tensor) and t.requires_grad: return LossOutput(loss=(t.flatten()[:1] * 0).sum(), num_tokens=0) @@ -232,7 +233,7 @@ def _force_zero_backward(model, og, adapter_name, kwargs): # Find a graph-connected tensor for zero loss zero_loss = None if outputs is not None and isinstance(outputs, dict): - for key in ('logps', 'logits', 'loss'): + for key in ('logps', 'values', 'logits', 'loss'): t = outputs.get(key) if t is not None and isinstance(t, torch.Tensor) and t.requires_grad: zero_loss = (t.flatten()[:1] * 0).sum() diff --git a/tests/advantage/test_gae.py b/tests/advantage/test_gae.py new file mode 100644 index 000000000..0e706359b --- /dev/null +++ b/tests/advantage/test_gae.py @@ -0,0 +1,44 @@ +import pytest +import torch + +from twinkle.advantage import GAEAdvantage + + +def test_single_terminal_token(): + advantages, returns = GAEAdvantage(gamma=1.0, gae_lambda=1.0, normalize=False)([[2.0]], [[0.5]]) + torch.testing.assert_close(advantages, torch.tensor([[1.5]])) + torch.testing.assert_close(returns, torch.tensor([[2.0]])) + + +def test_multi_token_terminal_gae(): + advantages, returns = GAEAdvantage(gamma=1.0, gae_lambda=1.0, normalize=False)( + [[0.0, 0.0, 1.0]], [[0.2, 0.3, 0.4]]) + torch.testing.assert_close(advantages, torch.tensor([[0.8, 0.7, 0.6]])) + torch.testing.assert_close(returns, torch.ones(1, 3)) + + +def test_padding_is_ignored(): + advantages, returns = GAEAdvantage(normalize=False)( + [[0.0, 1.0, 99.0]], [[0.2, 0.3, 42.0]], masks=[[True, True, False]]) + assert advantages[0, 2] == 0 + assert returns[0, 2] == 0 + + +def test_advantage_normalization(): + advantages, _ = GAEAdvantage(gamma=0.0, gae_lambda=0.0)( + [[1.0, 2.0], [3.0, 4.0]], torch.zeros(2, 2)) + assert advantages.mean().item() == pytest.approx(0.0, abs=1e-6) + assert advantages.std(unbiased=False).item() == pytest.approx(1.0, abs=1e-6) + + +def test_terminal_reward_and_kl_shaping(): + rewards = GAEAdvantage.build_token_rewards( + [2.0], [2], old_logps=[[-1.0, -2.0]], ref_logps=[[-1.5, -1.5]], kl_coef=0.1) + assert rewards[0] == pytest.approx([-0.05, 2.05]) + + +def test_invalid_hyperparameters(): + with pytest.raises(ValueError): + GAEAdvantage(gamma=1.1) + with pytest.raises(ValueError): + GAEAdvantage(gae_lambda=-0.1) diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 6147fa13e..3b8847d3c 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -307,6 +307,33 @@ def test_cli_new_infra_fields(self): assert args.infra.sequence_parallel is True assert args.infra.world_size == 32 + def test_cli_ppo_fields(self): + args = CLI.from_args(argv=[ + '--critic_model_gpus', + '2', + '--num_train_epochs', + '3', + '--gamma', + '0.99', + '--gae_lambda', + '0.9', + '--kl_coef', + '0.01', + '--no_normalize_advantages', + '--critic_learning_rate', + '2e-5', + '--value_clip', + '0.1', + ]) + assert args.infra.critic_model_gpus == 2 + assert args.training.num_train_epochs == 3 + assert args.rl.gamma == pytest.approx(0.99) + assert args.rl.gae_lambda == pytest.approx(0.9) + assert args.rl.kl_coef == pytest.approx(0.01) + assert args.rl.normalize_advantages is False + assert args.rl.critic_learning_rate == pytest.approx(2e-5) + assert args.loss.value_clip == pytest.approx(0.1) + def test_use_megatron_true_flips_strategy(self): args = CLI.from_args(argv=['--use_megatron', 'true']) assert args.model.strategy == 'native_fsdp' diff --git a/tests/loss/test_ppo.py b/tests/loss/test_ppo.py new file mode 100644 index 000000000..882ea5797 --- /dev/null +++ b/tests/loss/test_ppo.py @@ -0,0 +1,106 @@ +import torch + +from twinkle.advantage import GAEAdvantage +from twinkle.loss import GRPOLoss, PPOLoss, PPOValueLoss + + +def test_ppo_policy_loss_reuses_grpo_objective(): + labels = torch.tensor([[1, 2, -100]]) + logps = torch.tensor([[-0.8, -1.2, 0.0]], requires_grad=True) + old_logps = [[-1.0, -1.0]] + advantages = [[1.0, -1.0]] + inputs = {'labels': labels} + outputs = {'logps': logps} + ppo = PPOLoss(epsilon=0.2)(inputs, outputs, old_logps=old_logps, advantages=advantages)['loss'] + grpo = GRPOLoss(epsilon=0.2)(inputs, outputs, old_logps=old_logps, advantages=advantages)['loss'] + torch.testing.assert_close(ppo, grpo) + + +def test_ppo_loss_aggregation_modes(): + per_token_loss = torch.tensor([[1.0, 3.0, 0.0], [2.0, 4.0, 6.0]]) + loss_mask = torch.tensor([[True, True, False], [True, True, True]]) + + token_mean = PPOLoss()._aggregate_loss(per_token_loss, loss_mask) + sequence_mean = PPOLoss(loss_agg_mode='seq-mean-token-mean')._aggregate_loss(per_token_loss, loss_mask) + + # token-mean: (1 + 3 + 2 + 4 + 6) / 5 + torch.testing.assert_close(token_mean, torch.tensor(3.2)) + # seq-mean-token-mean: ((1 + 3) / 2 + (2 + 4 + 6) / 3) / 2 + torch.testing.assert_close(sequence_mean, torch.tensor(3.0)) + + +def test_value_loss_without_clipping(): + values = torch.tensor([[0.0, 1.0, 10.0]], requires_grad=True) + result = PPOValueLoss(epsilon=0.2)( + {'labels': torch.tensor([[1, 2, -100]])}, + {'values': values}, + old_values=[[0.0, 1.0]], + returns=[[1.0, 1.0]], + ) + assert result['loss'].item() == 0.25 + result['loss'].backward() + assert torch.isfinite(values.grad).all() + assert values.grad[0, 2] == 0 + + +def test_value_loss_uses_clipped_maximum(): + values = torch.tensor([[2.0]], requires_grad=True) + result = PPOValueLoss(epsilon=0.2)( + {'labels': torch.tensor([[1]])}, + {'values': values}, + old_values=[[0.0]], + returns=[[1.0]], + ) + assert result['loss'].item() == 0.5 + + +def test_value_loss_accepts_trailing_singleton(): + values = torch.tensor([[[0.5], [0.0]]], requires_grad=True) + result = PPOValueLoss()( + {'labels': torch.tensor([[1, -100]])}, + {'values': values}, + old_values=[0.5], + returns=[1.0], + ) + assert torch.isfinite(result['loss']) + + +def test_ppo_rollout_updates_actor_and_critic(): + labels = torch.tensor([[-100, 1, 2, 3], [-100, 4, 5, 6]]) + old_logps = [[-0.9, -1.1, -1.0], [-1.2, -0.8, -1.1]] + old_values = [[0.1, 0.2, 0.3], [0.2, 0.1, 0.4]] + token_rewards = GAEAdvantage.build_token_rewards([1.0, -0.5], [3, 3]) + advantages, returns = GAEAdvantage(gamma=1.0, gae_lambda=0.95, normalize=True)( + token_rewards, old_values) + + actor_logps = torch.nn.Parameter(torch.tensor([ + [0.0, -0.7, -1.3, -0.8], + [0.0, -1.0, -0.7, -1.4], + ])) + critic_values = torch.nn.Parameter(torch.tensor([ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ])) + optimizer = torch.optim.SGD([actor_logps, critic_values], lr=0.1) + actor_before = actor_logps.detach().clone() + critic_before = critic_values.detach().clone() + + actor_loss = PPOLoss(epsilon=0.2)( + {'labels': labels}, + {'logps': actor_logps}, + old_logps=old_logps, + advantages=advantages, + )['loss'] + critic_loss = PPOValueLoss(epsilon=0.2)( + {'labels': labels}, + {'values': critic_values}, + old_values=old_values, + returns=returns, + )['loss'] + (actor_loss + critic_loss).backward() + optimizer.step() + + assert torch.isfinite(actor_loss) + assert torch.isfinite(critic_loss) + assert not torch.equal(actor_logps.detach(), actor_before) + assert not torch.equal(critic_values.detach(), critic_before) diff --git a/tests/model/test_value_model.py b/tests/model/test_value_model.py new file mode 100644 index 000000000..c4ee7c14d --- /dev/null +++ b/tests/model/test_value_model.py @@ -0,0 +1,76 @@ +import inspect +import tempfile +from unittest.mock import patch + +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast + +from twinkle.loss import PPOValueLoss +from twinkle.model import TransformersValueModel + + +def _tiny_model_dir(): + path = tempfile.mkdtemp() + config = GPT2Config( + vocab_size=8, + n_positions=16, + n_embd=8, + n_layer=1, + n_head=2, + bos_token_id=1, + eos_token_id=2, + ) + GPT2LMHeadModel(config).save_pretrained(path) + tokenizer = Tokenizer(WordLevel({ + '[UNK]': 0, + 'a': 1, + 'b': 2, + 'c': 3, + 'd': 4, + 'e': 5, + 'f': 6, + 'g': 7, + }, unk_token='[UNK]')) + PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token='[UNK]', pad_token='[UNK]').save_pretrained(path) + return path + + +def test_value_model_constructor_exposes_device_mesh(): + """Keep ``device_mesh`` visible to ``@remote_class`` dispatch.""" + parameter = inspect.signature(TransformersValueModel.__init__).parameters['device_mesh'] + assert parameter.default is None + + +def test_value_model_forward_and_backward(): + model = TransformersValueModel(model_id=_tiny_model_dir(), mixed_precision='no') + model.set_loss(PPOValueLoss()) + model._lazy_wrap_model() + model_device = str(next(model.model.parameters()).device) + with patch('twinkle.processor.base.Platform.get_local_device', return_value=model_device): + outputs = model.forward_backward( + inputs=[{'input_ids': [1, 2, 3], 'labels': [-100, 2, 3]}], + old_values=[[0.0, 0.0]], + returns=[[1.0, 1.0]], + ) + assert outputs['values'].shape == (1, 3) + assert outputs.get('logps') is None + head = model.model.get_output_embeddings() + assert head.out_features == 1 + assert head.weight.grad is not None + assert torch.isfinite(head.weight.grad).all() + + +def test_value_model_forward_only_returns_values(): + model = TransformersValueModel(model_id=_tiny_model_dir(), mixed_precision='no') + model.set_loss(PPOValueLoss()) + model._lazy_wrap_model() + model_device = str(next(model.model.parameters()).device) + with patch('twinkle.processor.base.Platform.get_local_device', return_value=model_device): + outputs = model.forward_only( + inputs=[{'input_ids': [1, 2, 3], 'labels': [-100, 2, 3]}], + ) + assert outputs['values'].shape == (1, 3) + assert outputs.get('logps') is None + assert not outputs['values'].requires_grad