cehrgpt 0.0.2__py3-none-any.whl → 0.1.0__py3-none-any.whl

cehrgpt/rl_finetune/cehrgpt_ppo_trainer.py

@@ -1,464 +0,0 @@
-import time
-from typing import List, Optional
-
-import numpy as np
-import torch
-from torch.nn.utils.rnn import pad_sequence
-from trl.core import (
-    WANDB_PADDING,
-    PPODecorators,
-    convert_to_scalar,
-    logprobs_from_logits,
-    stack_dicts,
-    stats_to_np,
-)
-from trl.trainer import PPOTrainer
-
-from cehrgpt.models.tokenization_hf_cehrgpt import CehrGptTokenizer
-
-
-class CehrGptPPODataCollator:
-    def __init__(self, tokenizer: CehrGptTokenizer, max_length: int):
-        self.tokenizer = tokenizer
-        self.max_length = max_length
-
-    def __call__(self, examples):
-
-        batch = {}
-
-        # Pad sequences to the max length in the batch
-        batch["input_ids"] = pad_sequence(
-            [example["input_ids"] for example in examples],
-            batch_first=True,
-            padding_value=self.tokenizer.pad_token_id,
-        ).to(torch.int64)
-
-        batch["attention_mask"] = pad_sequence(
-            [example["attention_mask"] for example in examples],
-            batch_first=True,
-            padding_value=0.0,
-        )
-
-        assert (
-            batch["input_ids"].shape[1] <= self.max_length
-        ), f"Invalid input_ids length: {batch['input_ids'].shape[1]}"
-
-        if "value_indicators" in examples[0]:
-            batch["value_indicators"] = pad_sequence(
-                [example["value_indicators"] for example in examples],
-                batch_first=True,
-                padding_value=False,
-            )
-
-        if "values" in examples[0]:
-            batch["values"] = pad_sequence(
-                [example["values"] for example in examples],
-                batch_first=True,
-                padding_value=self.tokenizer.pad_value_token_id,
-            )
-            assert batch["value_indicators"].shape[1] <= self.max_length
-            assert batch["values"].shape[1] <= self.max_length
-
-        return batch
-
-
-class CehrGptPPOTrainer(PPOTrainer):
-    def _step_safety_checker(
-        self,
-        batch_size: int,
-        queries: List[torch.LongTensor],
-        responses: List[torch.LongTensor],
-        scores: List[torch.FloatTensor],
-        values: List[torch.Tensor] = None,
-        value_indicators: List[torch.BoolTensor] = None,
-        masks: Optional[List[torch.LongTensor]] = None,
-    ):
-        """
-        Check if the input data is valid for training.
-
-        Args:
-            batch_size (int):
-                Batch size from the config file.
-            queries (List[`torch.LongTensor`]):
-                List of tensors containing the encoded queries of shape (`query_length`)
-            responses (List[`torch.LongTensor`]):
-                List of tensors containing the encoded responses of shape (`response_length`)
-            scores (List[`torch.FloatTensor`]):
-                List of tensors containing the scores.
-            masks (List[`torch.LongTensor`], *optional*):
-                list of optional tensors containing the masks of shape (`response_length`)
-
-        Returns:
-            `tuple`: The input processed data.
-        """
-        for name, tensor_list in zip(
-            ["queries", "responses", "scores", "values", "value_indicators"],
-            [queries, responses, scores, values, value_indicators],
-        ):
-            if not isinstance(tensor_list, list):
-                raise ValueError(
-                    f"{name} must be a list of tensors - got {type(tensor_list)}"
-                )
-            if not isinstance(tensor_list[0], torch.Tensor):
-                raise ValueError(
-                    f"Elements in {name} must be tensors - got {type(tensor_list[0])}"
-                )
-            if batch_size is not None and len(tensor_list) != batch_size:
-                raise ValueError(
-                    f"Batch size ({batch_size}) does not match number of examples - but got {len(tensor_list)} for: {name}"
-                )
-
-        # add queries, scores and responses on the correct device
-        queries = [tensor.to(self.current_device) for tensor in queries]
-        responses = [tensor.to(self.current_device) for tensor in responses]
-        scores = [tensor.to(self.current_device) for tensor in scores]
-        masks = (
-            [tensor.to(self.current_device) for tensor in masks]
-            if masks is not None
-            else None
-        )
-        values = (
-            [tensor.to(self.current_device) for tensor in values]
-            if values is not None
-            else None
-        )
-        value_indicators = (
-            [tensor.to(self.current_device) for tensor in value_indicators]
-            if value_indicators is not None
-            else None
-        )
-
-        # squeeze scores if needed
-        for i, score in enumerate(scores):
-            if score.dim() > 1:
-                raise ValueError(
-                    f"Scores must be 1-dimensional - got {score.dim()} for {score}"
-                )
-            elif score.dim() == 1:
-                scores[i] = score.squeeze()
-
-        return queries, responses, scores, values, value_indicators, masks
-
-    @PPODecorators.empty_device_cache()
-    def step(
-        self,
-        queries: List[torch.LongTensor],
-        responses: List[torch.LongTensor],
-        scores: List[torch.FloatTensor],
-        values: List[torch.Tensor] = None,
-        value_indicators: List[torch.BoolTensor] = None,
-        response_masks: Optional[List[torch.LongTensor]] = None,
-    ):
-
-        bs = self.config.batch_size
-
-        queries, responses, scores, values, value_indicators, response_masks = (
-            self._step_safety_checker(
-                bs, queries, responses, scores, values, value_indicators, response_masks
-            )
-        )
-        scores = torch.tensor(scores, device=self.current_device)
-        if self.config.use_score_scaling:
-            # Score scaling
-            scores_mean, scores_std = self.running.update(scores)
-            score_scaling_factor = scores_std + torch.finfo(scores.dtype).eps
-            if self.config.use_score_norm:
-                scores = (scores - scores_mean) / score_scaling_factor
-            else:
-                scores /= score_scaling_factor
-
-        if self.config.score_clip is not None:
-            # Score clipping
-            scores_dtype = scores.dtype
-            scores = torch.clip(
-                scores.float(), -self.config.score_clip, self.config.score_clip
-            ).to(dtype=scores_dtype)
-
-        # if we want to push best model to the hub
-        if hasattr(self, "highest_reward"):
-            if self.compare_step % self.config.compare_steps == 0:
-                curr_mean_reward = scores.mean()
-                # if the best reward ever seen
-                if curr_mean_reward > self.highest_reward:
-                    self.highest_reward = curr_mean_reward
-                    # push model to hub
-                    self.push_to_hub(**self.push_to_hub_kwargs)
-            self.compare_step += 1
-
-        timing = dict()
-        t0 = time.time()
-
-        t = time.time()
-
-        model_inputs = self.prepare_model_inputs(
-            queries, responses, values, value_indicators
-        )
-
-        if self.is_distributed:
-            pad_first = self.tokenizer.padding_side == "left"
-
-            model_inputs["input_ids"] = self.accelerator.pad_across_processes(
-                model_inputs["input_ids"],
-                dim=1,
-                pad_index=self.tokenizer.pad_token_id,
-                pad_first=pad_first,
-            )
-            model_inputs["attention_mask"] = self.accelerator.pad_across_processes(
-                model_inputs["attention_mask"], dim=1, pad_index=0, pad_first=pad_first
-            )
-            if values is not None:
-                model_inputs["values"] = self.accelerator.pad_across_processes(
-                    model_inputs["values"],
-                    dim=1,
-                    pad_index=self.tokenizer.pad_value_token_id,
-                    pad_first=pad_first,
-                )
-            if value_indicators is not None:
-                model_inputs["value_indicators"] = (
-                    self.accelerator.pad_across_processes(
-                        model_inputs["value_indicators"],
-                        dim=1,
-                        pad_index=False,
-                        pad_first=pad_first,
-                    )
-                )
-            if self.is_encoder_decoder:
-                model_inputs["decoder_input_ids"] = (
-                    self.accelerator.pad_across_processes(
-                        model_inputs["decoder_input_ids"],
-                        dim=1,
-                        pad_index=self.tokenizer.pad_token_id,
-                        pad_first=pad_first,
-                    )
-                )
-                model_inputs["decoder_attention_mask"] = (
-                    self.accelerator.pad_across_processes(
-                        model_inputs["decoder_attention_mask"],
-                        dim=1,
-                        pad_index=0,
-                        pad_first=pad_first,
-                    )
-                )
-
-        model_inputs_names = list(model_inputs.keys())
-
-        full_kl_penalty = self.config.kl_penalty == "full"
-
-        with torch.no_grad():
-            all_logprobs, logits_or_none, states_values, masks = (
-                self.batched_forward_pass(
-                    self.model,
-                    queries,
-                    responses,
-                    model_inputs,
-                    response_masks=response_masks,
-                    return_logits=full_kl_penalty,
-                )
-            )
-            with self.optional_peft_ctx():
-                ref_logprobs, ref_logits_or_none, _, _ = self.batched_forward_pass(
-                    self.model if self.is_peft_model else self.ref_model,
-                    queries,
-                    responses,
-                    model_inputs,
-                    return_logits=full_kl_penalty,
-                )
-
-        timing["time/ppo/forward_pass"] = time.time() - t
-
-        with torch.no_grad():
-            t = time.time()
-            if full_kl_penalty:
-                active_full_logprobs = logprobs_from_logits(
-                    logits_or_none, None, gather=False
-                )
-                ref_full_logprobs = logprobs_from_logits(
-                    ref_logits_or_none, None, gather=False
-                )
-
-                rewards, non_score_reward, kls = self.compute_rewards(
-                    scores, active_full_logprobs, ref_full_logprobs, masks
-                )
-            else:
-                rewards, non_score_reward, kls = self.compute_rewards(
-                    scores, all_logprobs, ref_logprobs, masks
-                )
-            timing["time/ppo/compute_rewards"] = time.time() - t
-
-            t = time.time()
-            states_values, advantages, returns = self.compute_advantages(
-                states_values, rewards, masks
-            )
-            timing["time/ppo/compute_advantages"] = time.time() - t
-
-        # upcast to float32 to avoid dataset issues
-        batch_dict = {
-            "queries": queries,
-            "responses": responses,
-            "logprobs": all_logprobs.to(torch.float32),
-            "states_values": states_values.to(torch.float32),
-            "masks": masks,
-            "advantages": advantages,
-            "returns": returns,
-        }
-        batch_dict.update(model_inputs)
-
-        t = time.time()
-        all_stats = []
-        early_stop = False
-        for _ in range(self.config.ppo_epochs):
-            if early_stop:
-                break
-            b_inds = np.random.permutation(bs)
-            for backward_batch_start in range(0, bs, self.config.backward_batch_size):
-                backward_batch_end = (
-                    backward_batch_start + self.config.backward_batch_size
-                )
-                backward_batch_inds = b_inds[backward_batch_start:backward_batch_end]
-
-                for mini_batch_start in range(
-                    0, self.config.backward_batch_size, self.config.mini_batch_size
-                ):
-                    mini_batch_end = mini_batch_start + self.config.mini_batch_size
-                    mini_batch_inds = backward_batch_inds[
-                        mini_batch_start:mini_batch_end
-                    ]
-                    mini_batch_dict = {
-                        "logprobs": batch_dict["logprobs"][mini_batch_inds],
-                        "states_values": batch_dict["states_values"][mini_batch_inds],
-                        "masks": batch_dict["masks"][mini_batch_inds],
-                        # hacks: the queries and responses are ragged.
-                        "queries": [batch_dict["queries"][i] for i in mini_batch_inds],
-                        "responses": [
-                            batch_dict["responses"][i] for i in mini_batch_inds
-                        ],
-                        "advantages": batch_dict["advantages"][mini_batch_inds],
-                        "returns": batch_dict["returns"][mini_batch_inds],
-                    }
-                    for k in model_inputs_names:
-                        mini_batch_dict[k] = batch_dict[k][mini_batch_inds]
-                    with self.accelerator.accumulate(self.model):
-                        model_inputs = {
-                            k: mini_batch_dict[k] for k in model_inputs_names
-                        }
-
-                        logprobs, logits, vpreds, _ = self.batched_forward_pass(
-                            self.model,
-                            mini_batch_dict["queries"],
-                            mini_batch_dict["responses"],
-                            model_inputs,
-                            return_logits=True,
-                        )
-                        train_stats = self.train_minibatch(
-                            mini_batch_dict["logprobs"],
-                            mini_batch_dict["states_values"],
-                            logprobs,
-                            logits,
-                            vpreds,
-                            mini_batch_dict["masks"],
-                            mini_batch_dict["advantages"],
-                            mini_batch_dict["returns"],
-                        )
-                        all_stats.append(train_stats)
-
-            # typically, early stopping is done at the epoch level
-            if self.config.early_stopping:
-                policykl = train_stats["policy/policykl"]
-                early_stop = self._early_stop(policykl)
-                if early_stop:
-                    break
-
-        timing["time/ppo/optimize_step"] = time.time() - t
-
-        t = time.time()
-        train_stats = stack_dicts(all_stats)
-
-        # reshape advantages/ratios such that they are not averaged.
-        train_stats["policy/advantages"] = torch.flatten(
-            train_stats["policy/advantages"]
-        ).unsqueeze(0)
-        train_stats["policy/advantages"] = torch.nan_to_num(
-            train_stats["policy/advantages"], WANDB_PADDING
-        )
-        train_stats["policy/ratio"] = torch.flatten(
-            train_stats["policy/ratio"]
-        ).unsqueeze(0)
-
-        stats = self.record_step_stats(
-            scores=scores,
-            logprobs=all_logprobs,
-            ref_logprobs=ref_logprobs,
-            non_score_reward=non_score_reward,
-            train_stats=train_stats,
-            kl_coef=self.kl_ctl.value,
-            masks=masks,
-            queries=queries,
-            responses=responses,
-            kls=kls,
-        )
-        # Gather/Reduce stats from all processes
-        if self.is_distributed:
-            stats = self.gather_stats(stats)
-        stats = stats_to_np(stats)
-        timing["time/ppo/calc_stats"] = time.time() - t
-        stats["ppo/learning_rate"] = self.optimizer.param_groups[0]["lr"]
-
-        # Update the KL control - multiply the batch_size by the number of processes
-        self.kl_ctl.update(
-            stats["objective/kl"],
-            self.config.batch_size * self.accelerator.num_processes,
-        )
-
-        # Log the total ppo time
-        timing["time/ppo/total"] = time.time() - t0
-        stats.update(timing)
-
-        # post-process stats for tensorboard and other loggers
-        if self.config.log_with != "wandb":
-            stats = convert_to_scalar(stats)
-
-        if self.lr_scheduler is not None:
-            self.lr_scheduler.step()
-
-        return stats
-
-    def prepare_model_inputs(
-        self,
-        queries: torch.Tensor,
-        responses: torch.Tensor,
-        values: torch.Tensor,
-        value_indicators: torch.Tensor,
-    ):
-        if self.is_encoder_decoder:
-            input_data = self.data_collator(
-                [
-                    {"input_ids": q, "attention_mask": torch.ones_like(q)}
-                    for q in queries
-                ]
-            ).to(self.current_device)
-
-            decoder_inputs = self.data_collator(
-                [
-                    {"input_ids": r, "attention_mask": torch.ones_like(r)}
-                    for r in responses
-                ]
-            ).to(self.current_device)
-            input_data["decoder_input_ids"] = decoder_inputs["input_ids"]
-            input_data["decoder_attention_mask"] = decoder_inputs["attention_mask"]
-        else:
-            input_ids = [torch.cat([q, r]) for q, r in zip(queries, responses)]
-            input_data = self.data_collator(
-                [
-                    {
-                        "input_ids": ids,
-                        "attention_mask": torch.ones_like(ids),
-                        "values": v_s,
-                        "value_indicators": v_indicators,
-                    }
-                    for ids, v_s, v_indicators in zip(
-                        input_ids, values, value_indicators
-                    )
-                ]
-            )
-        input_data.pop("labels", None)  # we don't want to compute LM losses
-        return input_data