We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG]: The indent of save_model in ppo.py will cause save_model error
### 🐛 Describe the bug


error indent:
![image](https://user-images.githubusercontent.com/233871/229757061-afedc23c-3cf2-40ad-aab9-1da52caa33b5.png)
correct indent:
![image](https://user-images.githubusercontent.com/233871/229757194-68bae50e-c2da-4328-9faa-a5e28439a7cf.png)

error when save model
│ /home/bmb/projects/ColossalAI/applications/Chat/examples/train_prompts.py:168 in main            │
│                                                                                                  │
│   165 │   │   │   │   update_timesteps=args.update_timesteps)                                    │
│   166 │                                                                                          │
│   167 │   # save model checkpoint after fitting                                                  │
│ ❱ 168 │   trainer.save_model(args.save_path, only_rank0=True, tokenizer=tokenizer)               │
│   169 │   # save optimizer checkpoint on all ranks                                               │
│   170 │   if args.need_optim_ckpt:                                                               │
│   171 │   │   strategy.save_optimizer(actor_optim,                                               │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
AttributeError: 'PPOTrainer' object has no attribute 'save_model'


### Environment

_No response_
[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `applications/Chat/coati/trainer/ppo.py`
Content:
```
1 from typing import Any, Callable, Dict, List, Optional
2 
3 import torch
4 import torch.nn as nn
5 from coati.experience_maker import Experience, NaiveExperienceMaker
6 from coati.models.base import Actor, Critic
7 from coati.models.generation_utils import update_model_kwargs_fn
8 from coati.models.loss import PolicyLoss, ValueLoss
9 from coati.replay_buffer import NaiveReplayBuffer
10 from torch.optim import Optimizer
11 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
12 
13 from .base import Trainer
14 from .callbacks import Callback
15 from .strategies import Strategy
16 
17 
18 class PPOTrainer(Trainer):
19     """
20         Trainer for PPO algorithm.
21 
22     Args:
23         strategy (Strategy): the strategy to use for training
24         actor (Actor): the actor model in ppo algorithm
25         critic (Critic): the critic model in ppo algorithm
26         reward_model (nn.Module): the reward model in rlhf algorithm to make reward of sentences
27         initial_model (Actor): the initial model in rlhf algorithm to generate reference logits to limit the update of actor
28         actor_optim (Optimizer): the optimizer to use for actor model
29         critic_optim (Optimizer): the optimizer to use for critic model
30         kl_coef (float, defaults to 0.1): the coefficient of kl divergence loss
31         train_batch_size (int, defaults to 8): the batch size to use for training
32         buffer_limit (int, defaults to 0): the max_size limitaiton of replay buffer
33         buffer_cpu_offload (bool, defaults to True): whether to offload replay buffer to cpu
34         eps_clip (float, defaults to 0.2): the clip coefficient of policy loss
35         value_clip (float, defaults to 0.4): the clip coefficient of value loss
36         experience_batch_size (int, defaults to 8): the batch size to use for experience generation
37         max_epochs (int, defaults to 1): the number of epochs of training process
38         tokenier (Callable, optional): the tokenizer to use for tokenizing the input
39         sample_replay_buffer (bool, defaults to False): whether to sample from replay buffer
40         dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader
41         callbacks (List[Callback], defaults to []): the callbacks to call during training process
42         generate_kwargs (dict, optional): the kwargs to use while model generating
43     """
44 
45     def __init__(self,
46                  strategy: Strategy,
47                  actor: Actor,
48                  critic: Critic,
49                  reward_model: nn.Module,
50                  initial_model: Actor,
51                  actor_optim: Optimizer,
52                  critic_optim: Optimizer,
53                  kl_coef: float = 0.1,
54                  ptx_coef: float = 0.9,
55                  train_batch_size: int = 8,
56                  buffer_limit: int = 0,
57                  buffer_cpu_offload: bool = True,
58                  eps_clip: float = 0.2,
59                  value_clip: float = 0.4,
60                  experience_batch_size: int = 8,
61                  max_epochs: int = 1,
62                  tokenizer: Optional[Callable[[Any], dict]] = None,
63                  sample_replay_buffer: bool = False,
64                  dataloader_pin_memory: bool = True,
65                  callbacks: List[Callback] = [],
66                  **generate_kwargs) -> None:
67         experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, kl_coef)
68         replay_buffer = NaiveReplayBuffer(train_batch_size, buffer_limit, buffer_cpu_offload)
69         generate_kwargs = _set_default_generate_kwargs(strategy, generate_kwargs, actor)
70         super().__init__(strategy, experience_maker, replay_buffer, experience_batch_size, max_epochs, tokenizer,
71                          sample_replay_buffer, dataloader_pin_memory, callbacks, **generate_kwargs)
72         self.actor = actor
73         self.critic = critic
74 
75         self.actor_loss_fn = PolicyLoss(eps_clip)
76         self.critic_loss_fn = ValueLoss(value_clip)
77         self.ptx_loss_fn = nn.CrossEntropyLoss(ignore_index=-100)
78         self.ptx_coef = ptx_coef
79         self.actor_optim = actor_optim
80         self.critic_optim = critic_optim
81 
82     def training_step(self, experience: Experience) -> Dict[str, float]:
83         self.actor.train()
84         self.critic.train()
85         # policy loss
86         num_actions = experience.action_mask.size(1)
87         action_log_probs = self.actor(experience.sequences, num_actions, attention_mask=experience.attention_mask)
88         actor_loss = self.actor_loss_fn(action_log_probs,
89                                         experience.action_log_probs,
90                                         experience.advantages,
91                                         action_mask=experience.action_mask)
92 
93         # ptx loss
94         if self.ptx_coef != 0:
95             ptx = next(iter(self.pretrain_dataloader))['input_ids'].to(torch.cuda.current_device())
96             label = next(iter(self.pretrain_dataloader))['labels'].to(torch.cuda.current_device())[:, 1:]
97             attention_mask = next(iter(self.pretrain_dataloader))['attention_mask'].to(torch.cuda.current_device())
98             ptx_log_probs = self.actor.get_base_model()(ptx, attention_mask=attention_mask)['logits'][..., :-1, :]
99             ptx_loss = self.ptx_loss_fn(ptx_log_probs.view(-1, ptx_log_probs.size(-1)), label.view(-1))
100             actor_loss = ptx_loss * self.ptx_coef + actor_loss * (1 - self.ptx_coef)
101 
102         self.strategy.backward(actor_loss, self.actor, self.actor_optim)
103         self.strategy.optimizer_step(self.actor_optim)
104         self.actor_optim.zero_grad()
105 
106         # value loss
107         values = self.critic(experience.sequences,
108                              action_mask=experience.action_mask,
109                              attention_mask=experience.attention_mask)
110         critic_loss = self.critic_loss_fn(values,
111                                           experience.values,
112                                           experience.reward,
113                                           action_mask=experience.action_mask)
114         self.strategy.backward(critic_loss, self.critic, self.critic_optim)
115         self.strategy.optimizer_step(self.critic_optim)
116         self.critic_optim.zero_grad()
117 
118         return {'reward': experience.reward.mean().item()}
119 
120 
121 def _set_default_generate_kwargs(strategy: Strategy, generate_kwargs: dict, actor: Actor) -> None:
122     origin_model = strategy._unwrap_actor(actor)
123     new_kwargs = {**generate_kwargs}
124     # use huggingface models method directly
125     if 'prepare_inputs_fn' not in generate_kwargs and hasattr(origin_model, 'prepare_inputs_for_generation'):
126         new_kwargs['prepare_inputs_fn'] = origin_model.prepare_inputs_for_generation
127 
128     if 'update_model_kwargs_fn' not in generate_kwargs:
129         new_kwargs['update_model_kwargs_fn'] = update_model_kwargs_fn
130 
131     return new_kwargs
132 
133 
134 def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
135     self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
136
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/applications/Chat/coati/trainer/ppo.py b/applications/Chat/coati/trainer/ppo.py
--- a/applications/Chat/coati/trainer/ppo.py
+++ b/applications/Chat/coati/trainer/ppo.py
@@ -117,6 +117,9 @@
 
         return {'reward': experience.reward.mean().item()}
 
+    def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
+        self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
+
 
 def _set_default_generate_kwargs(strategy: Strategy, generate_kwargs: dict, actor: Actor) -> None:
     origin_model = strategy._unwrap_actor(actor)
@@ -129,7 +132,3 @@
         new_kwargs['update_model_kwargs_fn'] = update_model_kwargs_fn
 
     return new_kwargs
-
-
-def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
-    self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)

{"golden_diff": "diff --git a/applications/Chat/coati/trainer/ppo.py b/applications/Chat/coati/trainer/ppo.py\n--- a/applications/Chat/coati/trainer/ppo.py\n+++ b/applications/Chat/coati/trainer/ppo.py\n@@ -117,6 +117,9 @@\n \n         return {'reward': experience.reward.mean().item()}\n \n+    def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:\n+        self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)\n+\n \n def _set_default_generate_kwargs(strategy: Strategy, generate_kwargs: dict, actor: Actor) -> None:\n     origin_model = strategy._unwrap_actor(actor)\n@@ -129,7 +132,3 @@\n         new_kwargs['update_model_kwargs_fn'] = update_model_kwargs_fn\n \n     return new_kwargs\n-\n-\n-def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:\n-    self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)\n", "issue": "[BUG]: The indent of save_model in ppo.py will cause save_model error\n### \ud83d\udc1b Describe the bug\n\n\r\nerror indent:\r\n![image](https://user-images.githubusercontent.com/233871/229757061-afedc23c-3cf2-40ad-aab9-1da52caa33b5.png)\r\ncorrect indent:\r\n![image](https://user-images.githubusercontent.com/233871/229757194-68bae50e-c2da-4328-9faa-a5e28439a7cf.png)\r\n\r\nerror when save model\r\n\u2502 /home/bmb/projects/ColossalAI/applications/Chat/examples/train_prompts.py:168 in main            \u2502\r\n\u2502                                                                                                  \u2502\r\n\u2502   165 \u2502   \u2502   \u2502   \u2502   update_timesteps=args.update_timesteps)                                    \u2502\r\n\u2502   166 \u2502                                                                                          \u2502\r\n\u2502   167 \u2502   # save model checkpoint after fitting                                                  \u2502\r\n\u2502 \u2771 168 \u2502   trainer.save_model(args.save_path, only_rank0=True, tokenizer=tokenizer)               \u2502\r\n\u2502   169 \u2502   # save optimizer checkpoint on all ranks                                               \u2502\r\n\u2502   170 \u2502   if args.need_optim_ckpt:                                                               \u2502\r\n\u2502   171 \u2502   \u2502   strategy.save_optimizer(actor_optim,                                               \u2502\r\n\u2570\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u256f\r\nAttributeError: 'PPOTrainer' object has no attribute 'save_model'\r\n\n\n### Environment\n\n_No response_\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "from typing import Any, Callable, Dict, List, Optional\n\nimport torch\nimport torch.nn as nn\nfrom coati.experience_maker import Experience, NaiveExperienceMaker\nfrom coati.models.base import Actor, Critic\nfrom coati.models.generation_utils import update_model_kwargs_fn\nfrom coati.models.loss import PolicyLoss, ValueLoss\nfrom coati.replay_buffer import NaiveReplayBuffer\nfrom torch.optim import Optimizer\nfrom transformers.tokenization_utils_base import PreTrainedTokenizerBase\n\nfrom .base import Trainer\nfrom .callbacks import Callback\nfrom .strategies import Strategy\n\n\nclass PPOTrainer(Trainer):\n    \"\"\"\n        Trainer for PPO algorithm.\n\n    Args:\n        strategy (Strategy): the strategy to use for training\n        actor (Actor): the actor model in ppo algorithm\n        critic (Critic): the critic model in ppo algorithm\n        reward_model (nn.Module): the reward model in rlhf algorithm to make reward of sentences\n        initial_model (Actor): the initial model in rlhf algorithm to generate reference logits to limit the update of actor\n        actor_optim (Optimizer): the optimizer to use for actor model\n        critic_optim (Optimizer): the optimizer to use for critic model\n        kl_coef (float, defaults to 0.1): the coefficient of kl divergence loss\n        train_batch_size (int, defaults to 8): the batch size to use for training\n        buffer_limit (int, defaults to 0): the max_size limitaiton of replay buffer\n        buffer_cpu_offload (bool, defaults to True): whether to offload replay buffer to cpu\n        eps_clip (float, defaults to 0.2): the clip coefficient of policy loss\n        value_clip (float, defaults to 0.4): the clip coefficient of value loss\n        experience_batch_size (int, defaults to 8): the batch size to use for experience generation\n        max_epochs (int, defaults to 1): the number of epochs of training process\n        tokenier (Callable, optional): the tokenizer to use for tokenizing the input\n        sample_replay_buffer (bool, defaults to False): whether to sample from replay buffer\n        dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader\n        callbacks (List[Callback], defaults to []): the callbacks to call during training process\n        generate_kwargs (dict, optional): the kwargs to use while model generating\n    \"\"\"\n\n    def __init__(self,\n                 strategy: Strategy,\n                 actor: Actor,\n                 critic: Critic,\n                 reward_model: nn.Module,\n                 initial_model: Actor,\n                 actor_optim: Optimizer,\n                 critic_optim: Optimizer,\n                 kl_coef: float = 0.1,\n                 ptx_coef: float = 0.9,\n                 train_batch_size: int = 8,\n                 buffer_limit: int = 0,\n                 buffer_cpu_offload: bool = True,\n                 eps_clip: float = 0.2,\n                 value_clip: float = 0.4,\n                 experience_batch_size: int = 8,\n                 max_epochs: int = 1,\n                 tokenizer: Optional[Callable[[Any], dict]] = None,\n                 sample_replay_buffer: bool = False,\n                 dataloader_pin_memory: bool = True,\n                 callbacks: List[Callback] = [],\n                 **generate_kwargs) -> None:\n        experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, kl_coef)\n        replay_buffer = NaiveReplayBuffer(train_batch_size, buffer_limit, buffer_cpu_offload)\n        generate_kwargs = _set_default_generate_kwargs(strategy, generate_kwargs, actor)\n        super().__init__(strategy, experience_maker, replay_buffer, experience_batch_size, max_epochs, tokenizer,\n                         sample_replay_buffer, dataloader_pin_memory, callbacks, **generate_kwargs)\n        self.actor = actor\n        self.critic = critic\n\n        self.actor_loss_fn = PolicyLoss(eps_clip)\n        self.critic_loss_fn = ValueLoss(value_clip)\n        self.ptx_loss_fn = nn.CrossEntropyLoss(ignore_index=-100)\n        self.ptx_coef = ptx_coef\n        self.actor_optim = actor_optim\n        self.critic_optim = critic_optim\n\n    def training_step(self, experience: Experience) -> Dict[str, float]:\n        self.actor.train()\n        self.critic.train()\n        # policy loss\n        num_actions = experience.action_mask.size(1)\n        action_log_probs = self.actor(experience.sequences, num_actions, attention_mask=experience.attention_mask)\n        actor_loss = self.actor_loss_fn(action_log_probs,\n                                        experience.action_log_probs,\n                                        experience.advantages,\n                                        action_mask=experience.action_mask)\n\n        # ptx loss\n        if self.ptx_coef != 0:\n            ptx = next(iter(self.pretrain_dataloader))['input_ids'].to(torch.cuda.current_device())\n            label = next(iter(self.pretrain_dataloader))['labels'].to(torch.cuda.current_device())[:, 1:]\n            attention_mask = next(iter(self.pretrain_dataloader))['attention_mask'].to(torch.cuda.current_device())\n            ptx_log_probs = self.actor.get_base_model()(ptx, attention_mask=attention_mask)['logits'][..., :-1, :]\n            ptx_loss = self.ptx_loss_fn(ptx_log_probs.view(-1, ptx_log_probs.size(-1)), label.view(-1))\n            actor_loss = ptx_loss * self.ptx_coef + actor_loss * (1 - self.ptx_coef)\n\n        self.strategy.backward(actor_loss, self.actor, self.actor_optim)\n        self.strategy.optimizer_step(self.actor_optim)\n        self.actor_optim.zero_grad()\n\n        # value loss\n        values = self.critic(experience.sequences,\n                             action_mask=experience.action_mask,\n                             attention_mask=experience.attention_mask)\n        critic_loss = self.critic_loss_fn(values,\n                                          experience.values,\n                                          experience.reward,\n                                          action_mask=experience.action_mask)\n        self.strategy.backward(critic_loss, self.critic, self.critic_optim)\n        self.strategy.optimizer_step(self.critic_optim)\n        self.critic_optim.zero_grad()\n\n        return {'reward': experience.reward.mean().item()}\n\n\ndef _set_default_generate_kwargs(strategy: Strategy, generate_kwargs: dict, actor: Actor) -> None:\n    origin_model = strategy._unwrap_actor(actor)\n    new_kwargs = {**generate_kwargs}\n    # use huggingface models method directly\n    if 'prepare_inputs_fn' not in generate_kwargs and hasattr(origin_model, 'prepare_inputs_for_generation'):\n        new_kwargs['prepare_inputs_fn'] = origin_model.prepare_inputs_for_generation\n\n    if 'update_model_kwargs_fn' not in generate_kwargs:\n        new_kwargs['update_model_kwargs_fn'] = update_model_kwargs_fn\n\n    return new_kwargs\n\n\ndef save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:\n    self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)\n", "path": "applications/Chat/coati/trainer/ppo.py"}], "after_files": [{"content": "from typing import Any, Callable, Dict, List, Optional\n\nimport torch\nimport torch.nn as nn\nfrom coati.experience_maker import Experience, NaiveExperienceMaker\nfrom coati.models.base import Actor, Critic\nfrom coati.models.generation_utils import update_model_kwargs_fn\nfrom coati.models.loss import PolicyLoss, ValueLoss\nfrom coati.replay_buffer import NaiveReplayBuffer\nfrom torch.optim import Optimizer\nfrom transformers.tokenization_utils_base import PreTrainedTokenizerBase\n\nfrom .base import Trainer\nfrom .callbacks import Callback\nfrom .strategies import Strategy\n\n\nclass PPOTrainer(Trainer):\n    \"\"\"\n        Trainer for PPO algorithm.\n\n    Args:\n        strategy (Strategy): the strategy to use for training\n        actor (Actor): the actor model in ppo algorithm\n        critic (Critic): the critic model in ppo algorithm\n        reward_model (nn.Module): the reward model in rlhf algorithm to make reward of sentences\n        initial_model (Actor): the initial model in rlhf algorithm to generate reference logits to limit the update of actor\n        actor_optim (Optimizer): the optimizer to use for actor model\n        critic_optim (Optimizer): the optimizer to use for critic model\n        kl_coef (float, defaults to 0.1): the coefficient of kl divergence loss\n        train_batch_size (int, defaults to 8): the batch size to use for training\n        buffer_limit (int, defaults to 0): the max_size limitaiton of replay buffer\n        buffer_cpu_offload (bool, defaults to True): whether to offload replay buffer to cpu\n        eps_clip (float, defaults to 0.2): the clip coefficient of policy loss\n        value_clip (float, defaults to 0.4): the clip coefficient of value loss\n        experience_batch_size (int, defaults to 8): the batch size to use for experience generation\n        max_epochs (int, defaults to 1): the number of epochs of training process\n        tokenier (Callable, optional): the tokenizer to use for tokenizing the input\n        sample_replay_buffer (bool, defaults to False): whether to sample from replay buffer\n        dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader\n        callbacks (List[Callback], defaults to []): the callbacks to call during training process\n        generate_kwargs (dict, optional): the kwargs to use while model generating\n    \"\"\"\n\n    def __init__(self,\n                 strategy: Strategy,\n                 actor: Actor,\n                 critic: Critic,\n                 reward_model: nn.Module,\n                 initial_model: Actor,\n                 actor_optim: Optimizer,\n                 critic_optim: Optimizer,\n                 kl_coef: float = 0.1,\n                 ptx_coef: float = 0.9,\n                 train_batch_size: int = 8,\n                 buffer_limit: int = 0,\n                 buffer_cpu_offload: bool = True,\n                 eps_clip: float = 0.2,\n                 value_clip: float = 0.4,\n                 experience_batch_size: int = 8,\n                 max_epochs: int = 1,\n                 tokenizer: Optional[Callable[[Any], dict]] = None,\n                 sample_replay_buffer: bool = False,\n                 dataloader_pin_memory: bool = True,\n                 callbacks: List[Callback] = [],\n                 **generate_kwargs) -> None:\n        experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, kl_coef)\n        replay_buffer = NaiveReplayBuffer(train_batch_size, buffer_limit, buffer_cpu_offload)\n        generate_kwargs = _set_default_generate_kwargs(strategy, generate_kwargs, actor)\n        super().__init__(strategy, experience_maker, replay_buffer, experience_batch_size, max_epochs, tokenizer,\n                         sample_replay_buffer, dataloader_pin_memory, callbacks, **generate_kwargs)\n        self.actor = actor\n        self.critic = critic\n\n        self.actor_loss_fn = PolicyLoss(eps_clip)\n        self.critic_loss_fn = ValueLoss(value_clip)\n        self.ptx_loss_fn = nn.CrossEntropyLoss(ignore_index=-100)\n        self.ptx_coef = ptx_coef\n        self.actor_optim = actor_optim\n        self.critic_optim = critic_optim\n\n    def training_step(self, experience: Experience) -> Dict[str, float]:\n        self.actor.train()\n        self.critic.train()\n        # policy loss\n        num_actions = experience.action_mask.size(1)\n        action_log_probs = self.actor(experience.sequences, num_actions, attention_mask=experience.attention_mask)\n        actor_loss = self.actor_loss_fn(action_log_probs,\n                                        experience.action_log_probs,\n                                        experience.advantages,\n                                        action_mask=experience.action_mask)\n\n        # ptx loss\n        if self.ptx_coef != 0:\n            ptx = next(iter(self.pretrain_dataloader))['input_ids'].to(torch.cuda.current_device())\n            label = next(iter(self.pretrain_dataloader))['labels'].to(torch.cuda.current_device())[:, 1:]\n            attention_mask = next(iter(self.pretrain_dataloader))['attention_mask'].to(torch.cuda.current_device())\n            ptx_log_probs = self.actor.get_base_model()(ptx, attention_mask=attention_mask)['logits'][..., :-1, :]\n            ptx_loss = self.ptx_loss_fn(ptx_log_probs.view(-1, ptx_log_probs.size(-1)), label.view(-1))\n            actor_loss = ptx_loss * self.ptx_coef + actor_loss * (1 - self.ptx_coef)\n\n        self.strategy.backward(actor_loss, self.actor, self.actor_optim)\n        self.strategy.optimizer_step(self.actor_optim)\n        self.actor_optim.zero_grad()\n\n        # value loss\n        values = self.critic(experience.sequences,\n                             action_mask=experience.action_mask,\n                             attention_mask=experience.attention_mask)\n        critic_loss = self.critic_loss_fn(values,\n                                          experience.values,\n                                          experience.reward,\n                                          action_mask=experience.action_mask)\n        self.strategy.backward(critic_loss, self.critic, self.critic_optim)\n        self.strategy.optimizer_step(self.critic_optim)\n        self.critic_optim.zero_grad()\n\n        return {'reward': experience.reward.mean().item()}\n\n    def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:\n        self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)\n\n\ndef _set_default_generate_kwargs(strategy: Strategy, generate_kwargs: dict, actor: Actor) -> None:\n    origin_model = strategy._unwrap_actor(actor)\n    new_kwargs = {**generate_kwargs}\n    # use huggingface models method directly\n    if 'prepare_inputs_fn' not in generate_kwargs and hasattr(origin_model, 'prepare_inputs_for_generation'):\n        new_kwargs['prepare_inputs_fn'] = origin_model.prepare_inputs_for_generation\n\n    if 'update_model_kwargs_fn' not in generate_kwargs:\n        new_kwargs['update_model_kwargs_fn'] = update_model_kwargs_fn\n\n    return new_kwargs\n", "path": "applications/Chat/coati/trainer/ppo.py"}]}