| |
|
|
| from typing import List, Union |
| from torch.utils.data import DataLoader |
| from torch.utils.tensorboard import SummaryWriter |
|
|
| from mindspeed_rl.utils.tokenizer import BaseTokenizer |
| from mindspeed_rl.workers.rule_reward import RuleReward |
| from mindspeed_rl.trainer.utils.compute_utils import FixedKLController, AdaptiveKLController |
| from mindspeed_rl.workers.scheduler.launcher import RayActorGroup |
| from mindspeed_rl.utils.loggers import WandbLogger |
|
|
|
|
| class RayBaseTrainer(object): |
| """ |
| Note that this trainer runs on the driver process on a single CPU/GPU node. |
| """ |
|
|
| def __init__(self, |
| actor_worker: RayActorGroup = None, |
| ref_worker: RayActorGroup = None, |
| vit_worker: RayActorGroup = None, |
| reward_list: List[Union[RayActorGroup, RuleReward]] = None, |
| train_iters: int = 1, |
| save_interval: int = 1, |
| gamma: float = 1.0, |
| lam: float = 0.95, |
| adv_estimator: str = "group_norm", |
| missing_eos_penalty: float = 1.0, |
| kl_penalty: str = 'low_var_kl', |
| kl_ctrl_type: str = 'fixed', |
| kl_horizon: int = 1000, |
| kl_target: float = 100.0, |
| init_kl_coef: float = 0.001, |
| global_batch_size: int = 32, |
| micro_batch_size: int = 1, |
| n_samples_per_prompt: int = 1, |
| tokenizer: BaseTokenizer = None, |
| dataset_additional_keys: List[str] = None, |
| blocking: bool = False, |
| guarantee_order: bool = False, |
| num_cpus_for_local_task: float = 0.1, |
| **kwargs): |
|
|
| self.actor_worker = actor_worker |
| self.ref_worker = ref_worker |
| self.vit_worker = vit_worker |
| self.reward_list = reward_list |
| self.train_iters = train_iters |
| self.save_interval = save_interval |
| self.gamma = gamma |
| self.lam = lam |
| self.adv_estimator = adv_estimator |
| self.missing_eos_penalty = missing_eos_penalty |
| self.kl_penalty = kl_penalty |
| self.kl_ctrl_type = kl_ctrl_type |
| self.kl_horizon = kl_horizon |
| self.kl_target = kl_target |
| self.init_kl_coef = init_kl_coef |
| self.global_batch_size = global_batch_size |
| self.micro_batch_size = micro_batch_size |
| self.n_samples_per_prompt = n_samples_per_prompt |
| self.tokenizer = tokenizer |
| self.dataset_additional_keys = dataset_additional_keys |
| self.blocking = blocking |
| self.guarantee_order = guarantee_order |
| self.num_cpus_for_local_task = num_cpus_for_local_task |
| self.kwargs = kwargs |
|
|
| |
| if kl_ctrl_type == 'fixed': |
| self.kl_ctrl = FixedKLController(init_kl_coef=self.init_kl_coef) |
| elif kl_ctrl_type == 'adaptive': |
| if self.kl_horizon <= 0: |
| raise ValueError(f'horizon must be larger than 0. Got {self.kl_horizon}') |
| self.kl_ctrl = AdaptiveKLController(init_kl_coef=init_kl_coef, |
| target_kl=kl_target, |
| horizon=kl_horizon) |
| else: |
| raise NotImplementedError |
|
|
| self.wandb = None |
| self.tensorboard = None |
| if kwargs.get("use_wandb", ""): |
| self.wandb = WandbLogger(kwargs) |
| if kwargs.get("use_tensorboard", "") and self.wandb is None: |
| self.tensorboard = SummaryWriter() |
|
|
| def transfer_dock_init(self): |
| pass |
|
|
| def fit(self, data_loader: DataLoader): |
| """ |
| The utils loop of xx |
| """ |
| pass |
|
|
| def save_checkpoint(self, iteration): |
| pass |
|
|
