Repository: hpcaitech/ColossalAI

Issue:
[BUG]:   Coati train_rm use naive stra
### 🐛 Describe the bug

![image](https://user-images.githubusercontent.com/45619430/229027112-d09019b2-5295-4719-9e97-a64702ab7b6e.png)
Description:
When attempting to save a model using the save_model() method in the coati/trainer/rm.py module, an error is encountered with an unexpected keyword argument 'tokenizer'. The error traceback suggests that the error occurs in line 135 of the coati/trainer/rm.py file.

Steps to Reproduce:

Run the train_rm.sh script.
Call the save_model() method with the 'tokenizer' keyword argument.
Expected Behavior:
The save_model() method should save the model without any errors or issues.

Actual Behavior:
The save_model() method raises a TypeError with the message "save_model() got an unexpected keyword argument 'tokenizer'".

Possible Solution:
It seems that the save_model() method in the coati/trainer/rm.py module does not expect a 'tokenizer' keyword argument. To resolve the issue, either remove the 'tokenizer' argument when calling the save_model() method, or modify the method to accept the 'tokenizer' argument if it is required for the intended functionality.

Additional Information:

Version of Coati used: 1.0.0
Python version: 3.7
Operating system: Linux

### Environment

_No response_
[tensor] fix some unittests

[tensor] fix some unittests



Discussion:
 
> Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿      
----  
**Title:** [BUG]:    
      



Resolving patch:
diff --git a/applications/Chat/coati/trainer/strategies/ddp.py b/applications/Chat/coati/trainer/strategies/ddp.py
--- a/applications/Chat/coati/trainer/strategies/ddp.py
+++ b/applications/Chat/coati/trainer/strategies/ddp.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 import os
 import random
 
@@ -5,12 +7,13 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from coati.models.base import Actor
+from coati.models.base import LM, Actor, RewardModel
 from coati.models.lora import LoraLinear
 from coati.replay_buffer import ReplayBuffer
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from .base import Strategy
 from .naive import NaiveStrategy
@@ -72,17 +75,32 @@ def _unwrap_actor(actor: Actor) -> nn.Module:
         model: DDP = Strategy._unwrap_actor(actor)
         return model.module
 
-    def save_model(self, model: nn.Module, path: str, only_rank0: bool = False) -> None:
+    def save_model(self, model: nn.Module, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
+        if only_rank0 and dist.get_rank() != 0:
+            return None
+        
         for module in model.modules():
             if isinstance(module, LoraLinear):
                 module.merge_weights = True
                 module.eval()
-
-        if only_rank0 and dist.get_rank() != 0:
-            return
-        model = model.model.module
-        state_dict = model.state_dict()
-        torch.save(state_dict, path)
+        
+        if isinstance(model, RewardModel):
+            state_dict = model.state_dict()
+            if only_rank0 and dist.get_rank() != 0:
+                return
+            torch.save(state_dict, path)
+        else:
+            try:
+                if isinstance(model, LM):
+                    model = model.model
+                model.save_pretrained(path)
+                if tokenizer is not None:
+                    tokenizer.save_pretrained(path)
+            except AttributeError:
+                state_dict = model.state_dict()
+                if only_rank0 and dist.get_rank() != 0:
+                    return
+                torch.save(state_dict, path)
 
     def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None:
         if only_rank0 and dist.get_rank() != 0:
diff --git a/applications/Chat/coati/trainer/strategies/naive.py b/applications/Chat/coati/trainer/strategies/naive.py
--- a/applications/Chat/coati/trainer/strategies/naive.py
+++ b/applications/Chat/coati/trainer/strategies/naive.py
@@ -1,11 +1,14 @@
-from typing import Any
+from typing import Any, Optional
 
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from coati.replay_buffer import ReplayBuffer
+from coati.models.base import LM, RewardModel
+from coati.models.lora import LoraLinear
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from .base import Strategy
 
@@ -38,9 +41,25 @@ def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False
                           pin_memory=pin_memory,
                           collate_fn=replay_buffer.collate_fn)
 
-    def save_model(self, model: nn.Module, path: str, only_rank0: bool = False) -> None:
-        unwrapped_model = self._unwrap_model(model)
-        torch.save(unwrapped_model.state_dict(), path)
+    def save_model(self, model: nn.Module, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
+        for module in model.modules():
+            if isinstance(module, LoraLinear):
+                module.merge_weights = True
+                module.eval()
+        
+        if isinstance(model, RewardModel):
+            state_dict = model.state_dict()
+            torch.save(state_dict, path)
+        else:
+            try:
+                if isinstance(model, LM):
+                    model = model.model
+                model.save_pretrained(path)
+                if tokenizer is not None:
+                    tokenizer.save_pretrained(path)
+            except AttributeError:
+                state_dict = model.state_dict()
+                torch.save(state_dict, path)
 
     def load_model(self, model: nn.Module, path: str, map_location: Any = None, strict: bool = True) -> None:
         unwrapped_model = self._unwrap_model(model)


Test patch: