import math import os import sys from typing import Dict, List import ray import tinker import torch from tinker import types from trinity.algorithm import ALGORITHM_TYPE from trinity.algorithm.advantage_fn import ADVANTAGE_FN from trinity.algorithm.entropy_loss_fn import ENTROPY_LOSS_FN from trinity.algorithm.entropy_loss_fn.entropy_loss_fn import DummyEntropyLossFn from trinity.algorithm.kl_fn import KL_FN from trinity.algorithm.policy_loss_fn import POLICY_LOSS_FN from trinity.algorithm.utils import prefix_metrics from trinity.common.config import Config from trinity.common.experience import Experience from trinity.manager.synchronizer import Synchronizer from trinity.trainer.tinker.utils import ( compute_data_metrics, compute_throughout_metrics, compute_timing_metrics, to_tinker_input, ) from trinity.trainer.trainer import TrainEngineWrapper from trinity.utils.log import get_logger from trinity.utils.timer import Timer class TinkerTrainerWrapper(TrainEngineWrapper): def __init__(self, config: Config): self.config = config self.logger = get_logger("tinker_trainer") self._init_algorithm() self.synchronizer = Synchronizer.get_actor(namespace=self.config.synchronizer.ray_namespace) def _init_algorithm(self): self.algorithm = ALGORITHM_TYPE.get(self.config.algorithm.algorithm_type) self.algorithm_config = algorithm_config = self.config.algorithm if self.algorithm.compute_advantage_in_trainer: self.advantage_fn = ADVANTAGE_FN.get(algorithm_config.advantage_fn)( **algorithm_config.advantage_fn_args ) self.kl_fn = KL_FN.get(algorithm_config.kl_penalty_fn)( **algorithm_config.kl_penalty_fn_args ) # TODO raise NotImplementedError( "`compute_advantage_in_trainer` is not implemented yet in tinker" ) self.loss_agg_mode = algorithm_config.loss_agg_mode self.policy_loss_fn = POLICY_LOSS_FN.get(algorithm_config.policy_loss_fn)( backend="tinker", **algorithm_config.policy_loss_fn_args ) self.kl_loss_fn = KL_FN.get(algorithm_config.kl_loss_fn)(**algorithm_config.kl_loss_fn_args) self.entropy_loss_fn = ENTROPY_LOSS_FN.get(algorithm_config.entropy_loss_fn)( **algorithm_config.entropy_loss_fn_args ) # EXPERIMENTAL: apply loss scale fix self.do_fix_actor_microbatch_loss_scale = ( self.config.trainer.fix_actor_microbatch_loss_scale and (self.loss_agg_mode == "token-mean") ) self.lr_scheduler_type = algorithm_config.optimizer.lr_scheduler_type self.total_steps = self.config.trainer.total_steps or sys.maxsize self.num_warmup_steps = algorithm_config.optimizer.lr_warmup_steps if self.num_warmup_steps < 0: self.num_warmup_steps = int( algorithm_config.optimizer.lr_warmup_steps_ratio * self.total_steps ) self.min_lr_ratio = algorithm_config.optimizer.min_lr_ratio assert 0.0 <= self.min_lr_ratio <= 1.0 self.logger.info( f"Total steps: {self.total_steps}, num_warmup_steps: {self.num_warmup_steps}" ) if self.lr_scheduler_type not in {"constant", "cosine"}: raise NotImplementedError( f"LR scheduler type {self.lr_scheduler_type} is not supported" ) @property def _current_lr_factor(self): train_step_num = self._train_step_num # warmup if train_step_num < self.num_warmup_steps: factor = float(train_step_num) / float(max(1.0, self.num_warmup_steps)) factor = self.min_lr_ratio + (1.0 - self.min_lr_ratio) * factor return factor # decay if train_step_num >= self.total_steps: progress = 1.0 else: progress = float(train_step_num - self.num_warmup_steps) / float( max(1.0, self.total_steps - self.num_warmup_steps) ) if self.lr_scheduler_type == "constant": factor = 1.0 elif self.lr_scheduler_type == "cosine": num_cycles = 0.5 # TODO: may add to config factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) factor = self.min_lr_ratio + (1.0 - self.min_lr_ratio) * factor return max(self.min_lr_ratio, factor) @property def current_learning_rate(self): return self._current_lr_factor * self.algorithm_config.optimizer.lr @property def adam_params(self): return types.AdamParams( learning_rate=self.current_learning_rate, beta1=self.algorithm_config.optimizer.betas[0], beta2=self.algorithm_config.optimizer.betas[1], # eps is currently not in config weight_decay=self.algorithm_config.optimizer.weight_decay, grad_clip_norm=self.config.trainer.grad_clip, ) async def prepare(self): self.service_client = tinker.ServiceClient() name_prefix_list = [self.config.project, self.config.group, self.config.name] self.tinker_checkpoint_name_prefix = "-".join( [prefix for prefix in name_prefix_list if prefix] ) self.default_local_dir = self.config.checkpoint_job_dir self.local_latest_checkpointed_iteration = os.path.join( self.config.checkpoint_job_dir, "latest_checkpointed_iteration.txt" ) self.local_latest_state_dict_iteration = os.path.join( self.config.checkpoint_job_dir, "latest_state_dict_iteration.txt" ) if os.path.exists(self.local_latest_checkpointed_iteration): with open(self.local_latest_checkpointed_iteration, "r") as f: self._train_step_num = self.latest_remote_checkpoint_step = int(f.read().strip()) checkpoint_file_path = os.path.join( self.default_local_dir, f"global_step_{self._train_step_num}", "remote_checkpoint_path.txt", ) with open(checkpoint_file_path, "r") as f: self.latest_remote_checkpoint_path = f.read().strip() self.actor_client = ( await self.service_client.create_training_client_from_state_with_optimizer_async( path=self.latest_remote_checkpoint_path, ) ) else: self.actor_client = await self.service_client.create_lora_training_client_async( base_model=self.config.model.model_path, rank=self.config.model.tinker.rank, seed=self.config.model.tinker.seed, train_mlp=self.config.model.tinker.train_mlp, train_attn=self.config.model.tinker.train_attn, train_unembed=self.config.model.tinker.train_unembed, ) self.latest_remote_checkpoint_step = 0 self.latest_remote_checkpoint_path = None self._train_step_num = 0 if os.path.exists(self.local_latest_state_dict_iteration): with open(self.local_latest_state_dict_iteration, "r") as f: self.latest_remote_sampler_step = int(f.read().strip()) sampler_file_path = os.path.join( self.default_local_dir, f"global_step_{self.latest_remote_sampler_step}", "remote_sampler_path.txt", ) with open(sampler_file_path, "r") as f: self.latest_remote_sampler_path = f.read().strip() else: self.latest_remote_sampler_step = 0 self.latest_remote_sampler_path = None self.ref_client = await self.service_client.create_sampling_client_async( base_model=self.config.model.model_path, ) @property def train_step_num(self) -> int: """Get the current training step number.""" return self._train_step_num def _loss_func( self, batch: list[types.Datum], logprobs: list[torch.Tensor] ) -> tuple[torch.Tensor, dict[str, float]]: total_loss = 0.0 metrics = {} assert len(self.model_inputs_list) == len( logprobs ), "len(self.model_inputs_list) must equal to len(logprobs)" for model_inputs, logprob in zip(self.model_inputs_list, logprobs): micro_batch_metrics = {} response_mask = model_inputs["action_mask"] logprob = logprob[-response_mask.shape[0] :] pg_loss, pg_loss_metrics = self.policy_loss_fn(logprob=logprob, **model_inputs) prefix_metrics( src_metrics=pg_loss_metrics, prefix="actor", dst_metrics=micro_batch_metrics ) if self.entropy_loss_fn != DummyEntropyLossFn: entropy = -(logprob * logprob.exp()) else: entropy = None # compute entropy loss from entropy entropy_loss, entropy_loss_metrics = self.entropy_loss_fn( # type: ignore entropy=entropy, **model_inputs, loss_agg_mode=self.loss_agg_mode, ) prefix_metrics( src_metrics=entropy_loss_metrics, prefix="actor", dst_metrics=micro_batch_metrics, ) # compute kl loss kl_loss, kl_loss_metrics = self.kl_loss_fn.calculate_kl_loss( logprob=logprob, ref_logprob=model_inputs["ref_logprob"], response_mask=response_mask, loss_agg_mode=self.loss_agg_mode, old_logprob=model_inputs["old_logprob"], ) prefix_metrics( src_metrics=kl_loss_metrics, prefix="actor", dst_metrics=micro_batch_metrics, ) # compute policy loss policy_loss = pg_loss - entropy_loss + kl_loss loss_scale = 1.0 if not self.do_fix_actor_microbatch_loss_scale: loss_scale /= len(logprobs) loss = policy_loss * loss_scale total_loss = total_loss + loss micro_batch_metrics["actor/final_loss"] = loss.detach().item() # update metrics for key, val in micro_batch_metrics.items(): if key not in metrics: metrics[key] = [] metrics[key].append(val) avg_metrics = {k: sum(v) / len(v) for k, v in metrics.items()} return total_loss, avg_metrics async def train_step(self, batch_exps: List[Experience]) -> Dict: """Training one step. Args: batch (List[Experience]): A batch of experiences to train. Returns: Dict: Metrics of the training step. """ batch, batch_input_tokens, model_inputs_list = to_tinker_input(batch_exps, self.logger) self.model_inputs_list = model_inputs_list timing_raw = {} metrics = {} self._train_step_num += 1 with Timer(timing_raw, "step"): if self.algorithm.use_reference: # ref_logprob may not be used import asyncio ref_logprobs = await asyncio.gather( *[ self.ref_client.compute_logprobs_async(input_tokens) for input_tokens in batch_input_tokens ] ) for model_inputs, ref_logprob in zip(model_inputs_list, ref_logprobs): response_length = model_inputs["action_mask"].shape[0] model_inputs["ref_logprob"] = torch.tensor(ref_logprob[-response_length:]) if self.algorithm.compute_advantage_in_trainer: # TODO: following is verl format, which is not compatible with tinker raise NotImplementedError( "`compute_advantage_in_trainer` is not implemented yet in tinker" ) else: # skip token_level_scores for sft/dpo for model_inputs in model_inputs_list: if "token_level_scores" in model_inputs: assert "token_level_rewards" not in model_inputs model_inputs["token_level_rewards"] = model_inputs["token_level_scores"] # update actor with Timer(timing_raw, "update_actor"): fwdbwd_future = await self.actor_client.forward_backward_custom_async( batch, self._loss_func ) optim_future = await self.actor_client.optim_step_async(self.adam_params) fwdbwd_result = await fwdbwd_future optim_result = await optim_future metrics.update(fwdbwd_result.metrics) if optim_result.metrics: metrics.update(optim_result.metrics) # collect metrics metrics.update(compute_data_metrics(batch=self.model_inputs_list)) timing_metrics = compute_timing_metrics(batch=self.model_inputs_list, timing_raw=timing_raw) metrics.update({k.replace("timing_s/", "time/"): v for k, v in timing_metrics.items()}) metrics.update( compute_throughout_metrics(batch=self.model_inputs_list, timing_raw=timing_raw) ) return metrics def save_checkpoint(self, block_until_saved: bool = False, save_as_hf: bool = False) -> None: """Save the checkpoint.""" if self.train_step_num == self.latest_remote_checkpoint_step: return self.latest_remote_checkpoint_step = self.train_step_num checkpoint_name = f"{self.tinker_checkpoint_name_prefix}-state-{self.train_step_num}" self.latest_remote_checkpoint_path = ( self.actor_client.save_state(checkpoint_name).result().path ) local_path = os.path.join( self.default_local_dir, f"global_step_{self.train_step_num}", ) os.makedirs(local_path, exist_ok=True) # save a flag to indicate this is a full checkpoint dir # make sure this flag is created before notifying the synchronizer # to avoid the synchronizer recognizing it as a state_dict-only checkpoint # TODO: use a better way to indicate full checkpoint flag_path = os.path.join(local_path, ".full_checkpoint") with open(flag_path, "w") as f: f.write("") remote_checkpoint_path = os.path.join(local_path, "remote_checkpoint_path.txt") with open(remote_checkpoint_path, "w") as f: f.write(self.latest_remote_checkpoint_path) with open(self.local_latest_checkpointed_iteration, "w") as f: f.write(str(self.train_step_num)) def sync_weight(self) -> None: """Sync the model weight.""" raise NotImplementedError("Tinker trainer does not support NCCL sync") def upload_state_dict(self) -> None: """Upload the state dict to Synchronizer.""" self.save_state_dict() ray.get( self.synchronizer.set_model_state_dict.remote( self.latest_remote_sampler_path, self.train_step_num ) ) def save_state_dict(self) -> None: """Only save the model state dict for Synchronizer.""" if self.train_step_num == self.latest_remote_sampler_step: return self.latest_remote_sampler_step = self.train_step_num checkpoint_name = f"{self.tinker_checkpoint_name_prefix}-sampler-{self.train_step_num}" self.latest_remote_sampler_path = ( self.actor_client.save_weights_for_sampler(checkpoint_name).result().path ) local_path = os.path.join( self.default_local_dir, f"global_step_{self.train_step_num}", ) os.makedirs(local_path, exist_ok=True) remote_sampler_path = os.path.join(local_path, "remote_sampler_path.txt") with open(remote_sampler_path, "w") as f: f.write(self.latest_remote_sampler_path) with open(self.local_latest_state_dict_iteration, "w") as f: f.write(str(self.train_step_num))