Merge branch 'main' of https://github.com/hpcaitech/ColossalAI

applications/Chat/coati/models/llama/llama_critic.py

@@ -13,19 +13,18 @@ class LlamaCritic(Critic):
     Args:
         pretrained (str): Pretrained model name or path.
         config (LlamaConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
         lora_rank (int): LoRA rank.
         lora_train_bias (str): LoRA bias training mode.
     """
 
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[LlamaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none',
-                 **kwargs) -> None:
-
+    def __init__(
+        self,
+        pretrained: Optional[str] = None,
+        config: Optional[LlamaConfig] = None,
+        lora_rank: int = 0,
+        lora_train_bias: str = "none",
+        **kwargs,
+    ) -> None:
         if pretrained is not None:
             model = LlamaModel.from_pretrained(pretrained)
         elif config is not None:
@@ -33,9 +32,5 @@ class LlamaCritic(Critic):
         else:
             model = LlamaModel(LlamaConfig())
 
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-
         value_head = nn.Linear(model.config.hidden_size, 1)
-
         super().__init__(model, value_head, lora_rank, lora_train_bias, **kwargs)

applications/Chat/coati/models/llama/llama_rm.py

 from typing import Optional
 
 import torch.nn as nn
-from transformers import LlamaConfig, LlamaForCausalLM, LlamaModel
+from transformers import LlamaConfig, LlamaModel
 
 from ..base import RewardModel
 
@@ -13,18 +13,17 @@ class LlamaRM(RewardModel):
     Args:
         pretrained (str): Pretrained model name or path.
         config (LlamaConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
         lora_rank (int): LoRA rank.
         lora_train_bias (str): LoRA bias training mode.
     """
 
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[LlamaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
-
+    def __init__(
+        self,
+        pretrained: Optional[str] = None,
+        config: Optional[LlamaConfig] = None,
+        lora_rank: int = 0,
+        lora_train_bias: str = "none",
+    ) -> None:
         if pretrained is not None:
             model = LlamaModel.from_pretrained(pretrained)
         elif config is not None:
@@ -32,8 +31,6 @@ class LlamaRM(RewardModel):
         else:
             model = LlamaModel(LlamaConfig())
 
-        if checkpoint:
-            model.gradient_checkpointing_enable()
         value_head = nn.Linear(model.config.hidden_size, 1)
         value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.hidden_size + 1))
 

applications/Chat/coati/models/lora.py

+import dataclasses
 import math
+import warnings
 from typing import Optional
 
 import loralib as lora
@@ -7,9 +9,16 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
+@dataclasses.dataclass
+class LoRAManager:
+    merge_weights: bool = False
+
+
+LORA_MANAGER = LoRAManager()
+
+
 class LoraLinear(lora.LoRALayer, nn.Module):
-    """Replace in-place ops to out-of-place ops to fit gemini. Convert a torch.nn.Linear to LoraLinear.
-    """
+    """Replace in-place ops to out-of-place ops to fit gemini. Convert a torch.nn.Linear to LoraLinear."""
 
     def __init__(
         self,
@@ -17,16 +26,12 @@ class LoraLinear(lora.LoRALayer, nn.Module):
         bias: Optional[nn.Parameter],
         r: int = 0,
         lora_alpha: int = 1,
-        lora_dropout: float = 0.,
-        fan_in_fan_out: bool = False,    # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
-        merge_weights: bool = True,
+        lora_dropout: float = 0.0,
+        # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+        fan_in_fan_out: bool = False,
     ):
         nn.Module.__init__(self)
-        lora.LoRALayer.__init__(self,
-                                r=r,
-                                lora_alpha=lora_alpha,
-                                lora_dropout=lora_dropout,
-                                merge_weights=merge_weights)
+        lora.LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False)
         self.weight = weight
         self.bias = bias
 
@@ -47,39 +52,42 @@ class LoraLinear(lora.LoRALayer, nn.Module):
             self.weight.data = self.weight.data.T
 
     def reset_parameters(self):
-        if hasattr(self, 'lora_A'):
-            # initialize A the same way as the default for nn.Linear and B to zero
+        if hasattr(self, "lora_A"):
+            # Initialize A with the default values for nn.Linear and set B to zero.
             nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
             nn.init.zeros_(self.lora_B)
 
     def train(self, mode: bool = True):
-
-        def T(w):
-            return w.T if self.fan_in_fan_out else w
-
-        nn.Module.train(self, mode)
-        if self.merge_weights and self.merged:
-            # Make sure that the weights are not merged
-            if self.r > 0:
-                self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
-            self.merged = False
-
-    def eval(self):
-
         def T(w):
             return w.T if self.fan_in_fan_out else w
 
-        nn.Module.eval(self)
-        if self.merge_weights and not self.merged:
-            # Merge the weights and mark it
-            if self.r > 0:
-                self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
-                delattr(self, 'lora_A')
-                delattr(self, 'lora_B')
-            self.merged = True
+        self.training = mode
+        if LORA_MANAGER.merge_weights:
+            if mode and self.merged:
+                warnings.warn("Invoke module.train() would unmerge LoRA weights.")
+                raise NotImplementedError("LoRA unmerge is not tested.")
+                # Make sure that the weights are not merged
+                if self.r > 0:
+                    if not hasattr(self, "lora_A") or not hasattr(self, "lora_B"):
+                        # FIXME(csric): temporary fix
+                        self.lora_A = nn.Parameter(self.weight.new_empty((self.r, self.in_features)))
+                        self.lora_B = nn.Parameter(self.weight.new_empty((self.out_features, self.r)))
+                        self.reset_parameters()
+                    else:
+                        self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
+                self.merged = False
+            elif not mode and not self.merged:
+                warnings.warn("Invoke module.eval() would merge LoRA weights.")
+                # Merge the weights and mark it
+                if self.r > 0:
+                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
+                    delattr(self, "lora_A")
+                    delattr(self, "lora_B")
+                self.merged = True
+
+        return self
 
     def forward(self, x: torch.Tensor):
-
         def T(w):
             return w.T if self.fan_in_fan_out else w
 
@@ -92,21 +100,23 @@ class LoraLinear(lora.LoRALayer, nn.Module):
             return F.linear(x, T(self.weight), bias=self.bias)
 
 
-def lora_linear_wrapper(linear: nn.Linear, lora_rank: int) -> LoraLinear:
-    assert lora_rank <= linear.in_features, f'LoRA rank ({lora_rank}) must be less than or equal to in features ({linear.in_features})'
-    lora_linear = LoraLinear(linear.weight, linear.bias, r=lora_rank, merge_weights=False)
+def _lora_linear_wrapper(linear: nn.Linear, lora_rank: int) -> LoraLinear:
+    assert (
+        lora_rank <= linear.in_features
+    ), f"LoRA rank ({lora_rank}) must be less than or equal to in features ({linear.in_features})"
+    lora_linear = LoraLinear(linear.weight, linear.bias, r=lora_rank)
     return lora_linear
 
 
-def convert_to_lora_recursively(module: nn.Module, lora_rank: int) -> None:
+def _convert_to_lora_recursively(module: nn.Module, lora_rank: int) -> None:
     for name, child in module.named_children():
         if isinstance(child, nn.Linear):
-            setattr(module, name, lora_linear_wrapper(child, lora_rank))
+            setattr(module, name, _lora_linear_wrapper(child, lora_rank))
         else:
-            convert_to_lora_recursively(child, lora_rank)
+            _convert_to_lora_recursively(child, lora_rank)
 
 
-def convert_to_lora_module(module: nn.Module, lora_rank: int, lora_train_bias: str = 'none') -> nn.Module:
+def convert_to_lora_module(module: nn.Module, lora_rank: int, lora_train_bias: str = "none") -> nn.Module:
     """Convert a torch.nn.Module to a LoRA module.
 
     Args:
@@ -118,7 +128,7 @@ def convert_to_lora_module(module: nn.Module, lora_rank: int, lora_train_bias: s
     """
     if lora_rank <= 0:
         return module
-    convert_to_lora_recursively(module, lora_rank)
+    _convert_to_lora_recursively(module, lora_rank)
     lora.mark_only_lora_as_trainable(module, lora_train_bias)
     return module
 
@@ -134,7 +144,7 @@ class LoRAModule(nn.Module):
             Defaults to 'none'.
     """
 
-    def __init__(self, lora_rank: int = 0, lora_train_bias: str = 'none') -> None:
+    def __init__(self, lora_rank: int = 0, lora_train_bias: str = "none") -> None:
         super().__init__()
         self.lora_rank = lora_rank
         self.lora_train_bias = lora_train_bias

applications/Chat/coati/models/loss.py

@@ -13,6 +13,7 @@ class GPTLMLoss(nn.Module):
 
     def __init__(self):
         super().__init__()
+        # NOTE: default ignore_index is -100, which is equal to IGNORE_INDEX in sft_dataset.py
         self.loss = nn.CrossEntropyLoss()
 
     def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
@@ -31,11 +32,13 @@ class PolicyLoss(nn.Module):
         super().__init__()
         self.clip_eps = clip_eps
 
-    def forward(self,
-                log_probs: torch.Tensor,
-                old_log_probs: torch.Tensor,
-                advantages: torch.Tensor,
-                action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def forward(
+        self,
+        log_probs: torch.Tensor,
+        old_log_probs: torch.Tensor,
+        advantages: torch.Tensor,
+        action_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         ratio = (log_probs - old_log_probs).exp()
         surr1 = ratio * advantages
         surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages
@@ -55,44 +58,21 @@ class ValueLoss(nn.Module):
         super().__init__()
         self.clip_eps = clip_eps
 
-    def forward(self,
-                values: torch.Tensor,
-                old_values: torch.Tensor,
-                reward: torch.Tensor,
-                action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def forward(
+        self,
+        values: torch.Tensor,
+        old_values: torch.Tensor,
+        reward: torch.Tensor,
+        action_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         values_clipped = old_values + (values - old_values).clamp(-self.clip_eps, self.clip_eps)
-        surr1 = (values_clipped - reward)**2
-        surr2 = (values - reward)**2
+        surr1 = (values_clipped - reward) ** 2
+        surr2 = (values - reward) ** 2
         loss = torch.max(surr1, surr2)
         loss = loss.mean()
         return 0.5 * loss
 
 
-class PPOPtxActorLoss(nn.Module):
-    """
-    To Do:
-
-    PPO-ptx Actor Loss
-    """
-
-    def __init__(self, policy_clip_eps: float = 0.2, pretrain_coef: float = 0.0, pretrain_loss_fn=GPTLMLoss()) -> None:
-        super().__init__()
-        self.pretrain_coef = pretrain_coef
-        self.policy_loss_fn = PolicyLoss(clip_eps=policy_clip_eps)
-        self.pretrain_loss_fn = pretrain_loss_fn
-
-    def forward(self,
-                log_probs: torch.Tensor,
-                old_log_probs: torch.Tensor,
-                advantages: torch.Tensor,
-                lm_logits: torch.Tensor,
-                lm_input_ids: torch.Tensor,
-                action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
-        policy_loss = self.policy_loss_fn(log_probs, old_log_probs, advantages, action_mask=action_mask)
-        lm_loss = self.pretrain_loss_fn(lm_logits, lm_input_ids)
-        return policy_loss + self.pretrain_coef * lm_loss
-
-
 class LogSigLoss(nn.Module):
     """
     Pairwise Loss for Reward Model

applications/Chat/coati/models/opt/__init__.py

@@ -2,4 +2,4 @@ from .opt_actor import OPTActor
 from .opt_critic import OPTCritic
 from .opt_rm import OPTRM
 
-__all__ = ['OPTActor', 'OPTCritic', 'OPTRM']
+__all__ = ["OPTActor", "OPTCritic", "OPTRM"]

applications/Chat/coati/models/opt/opt_actor.py

@@ -18,12 +18,14 @@ class OPTActor(Actor):
         lora_train_bias (str): LoRA bias training mode.
     """
 
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[OPTConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
+    def __init__(
+        self,
+        pretrained: Optional[str] = None,
+        config: Optional[OPTConfig] = None,
+        checkpoint: bool = False,
+        lora_rank: int = 0,
+        lora_train_bias: str = "none",
+    ) -> None:
         if pretrained is not None:
             model = OPTForCausalLM.from_pretrained(pretrained)
         elif config is not None:

applications/Chat/coati/models/opt/opt_critic.py

@@ -14,25 +14,24 @@ class OPTCritic(Critic):
     Args:
         pretrained (str): Pretrained model name or path.
         config (OPTConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
         lora_rank (int): Rank of the low-rank approximation.
         lora_train_bias (str): LoRA bias training mode.
     """
 
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[OPTConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none',
-                 **kwargs) -> None:
+    def __init__(
+        self,
+        pretrained: Optional[str] = None,
+        config: Optional[OPTConfig] = None,
+        lora_rank: int = 0,
+        lora_train_bias: str = "none",
+        **kwargs,
+    ) -> None:
         if pretrained is not None:
             model = OPTModel.from_pretrained(pretrained)
         elif config is not None:
             model = OPTModel(config)
         else:
             model = OPTModel(OPTConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
+
         value_head = nn.Linear(model.config.word_embed_proj_dim, 1)
         super().__init__(model, value_head, lora_rank, lora_train_bias, **kwargs)

applications/Chat/coati/models/opt/opt_rm.py

@@ -13,25 +13,23 @@ class OPTRM(RewardModel):
     Args:
         pretrained (str): Pretrained model name or path.
         config (OPTConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
         lora_rank (int): Rank of the low-rank approximation.
         lora_train_bias (str): LoRA bias training mode.
     """
 
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[OPTConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
+    def __init__(
+        self,
+        pretrained: Optional[str] = None,
+        config: Optional[OPTConfig] = None,
+        lora_rank: int = 0,
+        lora_train_bias: str = "none",
+    ) -> None:
         if pretrained is not None:
             model = OPTModel.from_pretrained(pretrained)
         elif config is not None:
             model = OPTModel(config)
         else:
             model = OPTModel(OPTConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
 
         value_head = nn.Linear(model.config.word_embed_proj_dim, 1)
         value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.word_embed_proj_dim + 1))

applications/Chat/coati/models/roberta/__init__.py

-from .roberta_actor import RoBERTaActor
-from .roberta_critic import RoBERTaCritic
-from .roberta_rm import RoBERTaRM
-
-__all__ = ['RoBERTaActor', 'RoBERTaCritic', 'RoBERTaRM']
\ No newline at end of file

applications/Chat/coati/models/roberta/roberta_actor.py

-from typing import Optional
-
-from transformers.models.roberta.configuration_roberta import RobertaConfig
-from transformers.models.roberta.modeling_roberta import RobertaForCausalLM
-
-from ..base import Actor
-
-class RoBERTaActor(Actor):
-    """
-    RoBERTa Actor model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (RoBERTaConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the low-rank approximation.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[RobertaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
-        if pretrained is not None:
-            model = RobertaForCausalLM.from_pretrained(pretrained)
-        elif config is not None:
-            model = RobertaForCausalLM(config)
-        else:
-            model = RobertaForCausalLM(RobertaConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-        super().__init__(model, lora_rank, lora_train_bias)

applications/Chat/coati/models/roberta/roberta_critic.py

-from typing import Optional
-
-import torch.nn as nn
-from transformers.models.roberta.configuration_roberta import RobertaConfig
-from transformers.models.roberta.modeling_roberta import RobertaModel
-
-from ..base import Critic
-
-
-class RoBERTaCritic(Critic):
-    """
-    RoBERTa Critic model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (RoBERTa Config): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the low-rank approximation.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[RobertaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none',
-                 **kwargs) -> None:
-        if pretrained is not None:
-            model = RobertaModel.from_pretrained(pretrained, add_pooling_layer=False)
-        elif config is not None:
-            model = RobertaModel(config)
-        else:
-            model = RobertaModel(RobertaConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-        value_head = nn.Linear(model.config.hidden_size, 1)
-        super().__init__(model, value_head, lora_rank, lora_train_bias, **kwargs)

applications/Chat/coati/models/roberta/roberta_rm.py

-from typing import Optional
-
-import torch.nn as nn
-from transformers import RobertaConfig, RobertaModel
-
-
-from ..base import RewardModel
-
-
-class RoBERTaRM(RewardModel):
-    """
-    RoBERTa Reward model.
-
-    Args:
-        pretrained (str): Pretrained model name or path.
-        config (RoBERTaConfig): Model config.
-        checkpoint (bool): Enable gradient checkpointing.
-        lora_rank (int): Rank of the low-rank approximation.
-        lora_train_bias (str): LoRA bias training mode.
-    """
-
-    def __init__(self,
-                 pretrained: Optional[str] = None,
-                 config: Optional[RobertaConfig] = None,
-                 checkpoint: bool = False,
-                 lora_rank: int = 0,
-                 lora_train_bias: str = 'none') -> None:
-        if pretrained is not None:
-            model = RobertaModel.from_pretrained(pretrained, add_pooling_layer=False)
-        elif config is not None:
-            model = RobertaModel(config)
-        else:
-            model = RobertaModel(RobertaConfig())
-        if checkpoint:
-            model.gradient_checkpointing_enable()
-
-        value_head = nn.Linear(model.config.hidden_size, 1)
-        value_head.weight.data.normal_(mean=0.0, std=1/(model.config.hidden_size + 1))
-        super().__init__(model, value_head, lora_rank, lora_train_bias)
\ No newline at end of file

applications/Chat/coati/models/utils.py

 from typing import Optional, Union
 
-import loralib as lora
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 
 
-def compute_approx_kl(log_probs: torch.Tensor,
-                      log_probs_base: torch.Tensor,
-                      action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+def _compute_approx_kl(
+    log_probs: torch.Tensor, log_probs_base: torch.Tensor, action_mask: Optional[torch.Tensor] = None
+) -> torch.Tensor:
     """
     Compute the approximate KL divergence between two distributions.
     Schulman blog: http://joschu.net/blog/kl-approx.html
@@ -19,7 +17,7 @@ def compute_approx_kl(log_probs: torch.Tensor,
         action_mask: Mask for actions.
     """
 
-    log_ratio = log_probs - log_probs_base
+    log_ratio = log_probs_base - log_probs
     approx_kl = (log_ratio.exp() - 1) - log_ratio
     if action_mask is not None:
         approx_kl = masked_mean(approx_kl, action_mask, dim=1)
@@ -28,65 +26,44 @@ def compute_approx_kl(log_probs: torch.Tensor,
     return approx_kl
 
 
-def compute_reward(r: Union[torch.Tensor, float],
-                   kl_coef: float,
-                   log_probs: torch.Tensor,
-                   log_probs_base: torch.Tensor,
-                   action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+def compute_reward(
+    r: Union[torch.Tensor, float],
+    kl_coef: float,
+    log_probs: torch.Tensor,
+    log_probs_base: torch.Tensor,
+    action_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
     if kl_coef <= 0.0:
         return r
-    kl = compute_approx_kl(log_probs, log_probs_base, action_mask=action_mask)
+    kl = _compute_approx_kl(log_probs, log_probs_base, action_mask=action_mask)
     reward = r - kl_coef * kl
     return reward
 
 
-def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+def _log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
     log_probs = F.log_softmax(logits, dim=-1)
     log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
     return log_probs_labels.squeeze(-1)
 
 
+def calc_action_log_probs(logits: torch.Tensor, sequences: torch.LongTensor, num_actions: int) -> torch.Tensor:
+    """Calculate action log probs.
+
+    Args:
+        output (torch.Tensor): Output tensor of Actor.forward.logits.
+        sequences (torch.LongTensor): Input sequences.
+        num_actions (int): Number of actions.
+
+    Returns:
+        torch.Tensor: Action log probs.
+    """
+    log_probs = _log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:])
+    return log_probs[:, -num_actions:]
+
+
 def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:
     tensor = tensor * mask
     tensor = tensor.sum(dim=dim)
     mask_sum = mask.sum(dim=dim)
     mean = tensor / (mask_sum + 1e-8)
     return mean
-
-
-def masked_normalize(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1, eps: float = 1e-8) -> torch.Tensor:
-    tensor = tensor * mask
-    mean = masked_mean(tensor, mask, dim=dim)
-    mean_centered = tensor - mean
-    var = masked_mean(mean_centered**2, mask, dim=dim)
-    return mean_centered * var.clamp(min=eps).rsqrt()
-
-
-def normalize(tensor: torch.Tensor, dim: int = 0, eps: float = 1e-8) -> torch.Tensor:
-    mean = tensor.mean(dim)
-    mean_centered = tensor - mean
-    var = (mean_centered**2).mean(dim)
-    norm = mean_centered * var.clamp(min=eps).rsqrt()
-    return norm
-
-
-def convert_to_lora(model: nn.Module,
-                    input_size: int,
-                    output_size: int,
-                    lora_rank: int = 16,
-                    lora_alpha: int = 1,
-                    lora_dropout: float = 0.,
-                    fan_in_fan_out: bool = False,
-                    merge_weights: bool = True):
-    if lora_rank > min(input_size, output_size):
-        raise ValueError(f"LoRA rank {lora_rank} must be less or equal than {min(input_size, output_size)}")
-
-    for name, module in model.named_modules():
-        if isinstance(module, nn.Linear):
-            module._modules[name] = lora.Linear(input_size,
-                                                output_size,
-                                                r=lora_rank,
-                                                lora_alpha=lora_alpha,
-                                                lora_dropout=lora_dropout,
-                                                fan_in_fan_out=fan_in_fan_out,
-                                                merge_weights=merge_weights)

applications/Chat/coati/quant/__init__.py

+from .llama_gptq import load_quant as llama_load_quant
+from .utils import low_resource_init
+
+__all__ = [
+    "llama_load_quant",
+    "low_resource_init",
+]

applications/Chat/inference/llama_gptq/__init__.py → applications/Chat/coati/quant/llama_gptq/__init__.py

 from .loader import load_quant
 
 __all__ = [
-    'load_quant',
+    "load_quant",
 ]

applications/Chat/inference/llama_gptq/loader.py → applications/Chat/coati/quant/llama_gptq/loader.py

 import torch
 import torch.nn as nn
-import transformers
-from transformers import LlamaConfig, LlamaForCausalLM
 
 from .model_utils import find_layers
 from .quant import make_quant
 
 
-def load_quant(pretrained: str, checkpoint: str, wbits: int, groupsize: int):
-    config = LlamaConfig.from_pretrained(pretrained)
-
-    def noop(*args, **kwargs):
-        pass
-
-    torch.nn.init.kaiming_uniform_ = noop
-    torch.nn.init.uniform_ = noop
-    torch.nn.init.normal_ = noop
-
-    torch.set_default_dtype(torch.half)
-    transformers.modeling_utils._init_weights = False
-    torch.set_default_dtype(torch.half)
-    model = LlamaForCausalLM(config)
-    torch.set_default_dtype(torch.float)
+def load_quant(model: nn.Module, checkpoint: str, wbits: int, groupsize: int):
     model = model.eval()
     layers = find_layers(model)
-    for name in ['lm_head']:
+
+    # ignore lm head
+    layers = find_layers(model)
+    for name in ["lm_head"]:
         if name in layers:
             del layers[name]
+
     make_quant(model, layers, wbits, groupsize)
 
-    print(f'Loading model with {wbits} bits...')
-    if checkpoint.endswith('.safetensors'):
+    if checkpoint.endswith(".safetensors"):
         from safetensors.torch import load_file as safe_load
+
         model.load_state_dict(safe_load(checkpoint))
     else:
         model.load_state_dict(torch.load(checkpoint))
-    model.seqlen = 2048
-    print('Done.')
 
     return model

applications/Chat/inference/llama_gptq/model_utils.py → applications/Chat/coati/quant/llama_gptq/model_utils.py

 # copied from https://github.com/qwopqwop200/GPTQ-for-LLaMa/blob/past/modelutils.py
 
-import torch
 import torch.nn as nn
 
 
-def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''):
+def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=""):
     if type(module) in layers:
         return {name: module}
     res = {}
     for name1, child in module.named_children():
-        res.update(find_layers(child, layers=layers, name=name + '.' + name1 if name != '' else name1))
+        res.update(find_layers(child, layers=layers, name=name + "." + name1 if name != "" else name1))
     return res

applications/Chat/inference/llama_gptq/quant.py → applications/Chat/coati/quant/llama_gptq/quant.py

@@ -13,14 +13,13 @@ def quantize(x, scale, zero, maxq):
 
 
 class Quantizer(nn.Module):
-
     def __init__(self, shape=1):
         super(Quantizer, self).__init__()
-        self.register_buffer('maxq', torch.tensor(0))
-        self.register_buffer('scale', torch.zeros(shape))
-        self.register_buffer('zero', torch.zeros(shape))
+        self.register_buffer("maxq", torch.tensor(0))
+        self.register_buffer("scale", torch.zeros(shape))
+        self.register_buffer("zero", torch.zeros(shape))
 
-    def configure(self, bits, perchannel=False, sym=True, mse=False, norm=2.4, grid=100, maxshrink=.8):
+    def configure(self, bits, perchannel=False, sym=True, mse=False, norm=2.4, grid=100, maxshrink=0.8):
         self.maxq = torch.tensor(2**bits - 1)
         self.perchannel = perchannel
         self.sym = sym
@@ -68,7 +67,7 @@ class Quantizer(nn.Module):
             self.zero = torch.round(-xmin / self.scale)
 
         if self.mse:
-            best = torch.full([x.shape[0]], float('inf'), device=dev)
+            best = torch.full([x.shape[0]], float("inf"), device=dev)
             for i in range(int(self.maxshrink * self.grid)):
                 p = 1 - i / self.grid
                 xmin1 = p * xmin
@@ -123,13 +122,12 @@ class Quantizer(nn.Module):
 try:
     import quant_cuda
 except:
-    print('CUDA extension not installed.')
+    print("CUDA extension not installed.")
 
 # Assumes layer is perfectly divisible into 256 * 256 blocks
 
 
 class QuantLinear(nn.Module):
-
     def __init__(self, bits, groupsize, infeatures, outfeatures):
         super().__init__()
         if bits not in [2, 3, 4, 8]:
@@ -142,11 +140,11 @@ class QuantLinear(nn.Module):
         groupsize = groupsize if groupsize != -1 else infeatures
         self.groupsize = groupsize
         self.register_buffer(
-            'qzeros', torch.zeros((math.ceil(infeatures / groupsize), outfeatures // 256 * (bits * 8)),
-                                  dtype=torch.int))
-        self.register_buffer('scales', torch.zeros((math.ceil(infeatures / groupsize), outfeatures)))
-        self.register_buffer('bias', torch.zeros(outfeatures))
-        self.register_buffer('qweight', torch.zeros((infeatures // 256 * (bits * 8), outfeatures), dtype=torch.int))
+            "qzeros", torch.zeros((math.ceil(infeatures / groupsize), outfeatures // 256 * (bits * 8)), dtype=torch.int)
+        )
+        self.register_buffer("scales", torch.zeros((math.ceil(infeatures / groupsize), outfeatures)))
+        self.register_buffer("bias", torch.zeros(outfeatures))
+        self.register_buffer("qweight", torch.zeros((infeatures // 256 * (bits * 8), outfeatures), dtype=torch.int))
         self._initialized_quant_state = False
 
     def pack(self, linear, scales, zeros):
@@ -161,8 +159,10 @@ class QuantLinear(nn.Module):
         for idx in range(self.infeatures):
             g_idx = idx // self.groupsize
             intweight.append(
-                torch.round((linear.weight.data[:, idx] + scale_zeros[g_idx]) / self.scales[g_idx]).to(torch.int)[:,
-                                                                                                                  None])
+                torch.round((linear.weight.data[:, idx] + scale_zeros[g_idx]) / self.scales[g_idx]).to(torch.int)[
+                    :, None
+                ]
+            )
         intweight = torch.cat(intweight, dim=1)
         intweight = intweight.t().contiguous()
         intweight = intweight.numpy().astype(np.uint32)
@@ -271,13 +271,13 @@ class QuantLinear(nn.Module):
         return y.reshape(outshape)
 
 
-def make_quant(module, names, bits, groupsize, name=''):
+def make_quant(module, names, bits, groupsize, name=""):
     if isinstance(module, QuantLinear):
         return
     for attr in dir(module):
         tmp = getattr(module, attr)
-        name1 = name + '.' + attr if name != '' else attr
+        name1 = name + "." + attr if name != "" else attr
         if name1 in names:
             setattr(module, attr, QuantLinear(bits, groupsize, tmp.in_features, tmp.out_features))
     for name1, child in module.named_children():
-        make_quant(child, names, bits, groupsize, name + '.' + name1 if name != '' else name1)
+        make_quant(child, names, bits, groupsize, name + "." + name1 if name != "" else name1)

applications/Chat/coati/quant/utils.py

+from contextlib import contextmanager
+
+import torch
+
+
+def _noop(*args, **kwargs):
+    pass
+
+
+@contextmanager
+def low_resource_init():
+    """This context manager disables weight initialization and sets the default float dtype to half."""
+    old_kaiming_uniform_ = torch.nn.init.kaiming_uniform_
+    old_uniform_ = torch.nn.init.uniform_
+    old_normal_ = torch.nn.init.normal_
+    dtype = torch.get_default_dtype()
+    try:
+        torch.nn.init.kaiming_uniform_ = _noop
+        torch.nn.init.uniform_ = _noop
+        torch.nn.init.normal_ = _noop
+        torch.set_default_dtype(torch.half)
+        yield
+    finally:
+        torch.nn.init.kaiming_uniform_ = old_kaiming_uniform_
+        torch.nn.init.uniform_ = old_uniform_
+        torch.nn.init.normal_ = old_normal_
+        torch.set_default_dtype(dtype)

applications/Chat/coati/ray/README.md

+:warning: **This content may be outdated since the major update of Colossal Chat. We will update this content soon.**
+
+# Distributed PPO Training on Stage 3
+
+## Detach Experience Makers and Trainers
+
+We can completely separate the trainers and makers.
+
+<p align="center">
+<img src="https://github.com/hpcaitech/public_assets/blob/main/applications/chat/basic_structure.png?raw=true" width=600/>
+</p>
+
+- The experience maker performs inference, produces experience, and remotely delivers it to the trainer (1).
+- The trainer consumes experience to train models, and periodically transmits new model parameters to the maker (2.1, 2.2).
+- Using an experience buffer to overlap transmission and computing.
+
+In this manner, each node will work continuously without model idle time, and different optimization strategies can be applied for inference and training to meet the needs of speed or storage. It is also helpful for scalability.
+
+`DetachedPPOTrainer` and `ExperienceMakerHolder` are Ray Actors (distinguished from Actor Model), representing Trainer and Experience Maker on the graph above, respectively.
+
+[More about Ray Core](https://docs.ray.io/en/latest/ray-core/walkthrough.html)
+
+## Usage
+
+See examples at `ColossalAI/application/Chat/examples/ray`
+
+### Setup Makers
+
+- define makers' environment variables :
+
+  ```python
+  env_info_makers = [{
+      'local_rank': '0',
+      'rank': str(rank),
+      'world_size': str(num_makers),
+      'master_port': maker_port,
+      'master_addr': master_addr
+  } for rank in range(num_makers)]
+
+  ```
+
+- define maker models :
+
+  ```python
+  def model_fn():
+      actor = get_actor_from_args(...)
+      critic = get_critic_from_args(...)
+      reward_model = get_reward_model_from_args(...)
+      initial_model = get_actor_from_args(...)
+      return actor, critic, reward_model, initial_model
+
+  ```
+
+- set experience_holder_refs :
+
+  ```python
+  experience_holder_refs = [
+      ExperienceMakerHolder.options(
+          name=f"maker_{i}",
+          num_gpus=1,
+          max_concurrency=2
+      ).remote(
+          detached_trainer_name_list=[f"trainer_{x}" for x in target_trainers(...)],
+          model_fn=model_fn,
+          ...)
+      for i, env_info_maker in enumerate(env_info_makers)
+  ]
+  ```
+
+  The names in the `detached_trainer_name_list` refer to the target trainers that the maker should send experience to.
+  We set a trainer's name the same as a maker, by `.options(name="str")`. See below.
+
+### Setup Trainers
+
+- define trainers' environment variables :
+  ```python
+  env_info_trainers = [{
+      'local_rank': '0',
+      'rank': str(rank),
+      'world_size': str(num_trainers),
+      'master_port': trainer_port,
+      'master_addr': master_addr
+  } for rank in range(num_trainers)]
+  ```
+- define trainer models :
+
+  ```python
+  def trainer_model_fn():
+      actor = get_actor_from_args(...)
+      critic = get_critic_from_args(...)
+      return actor, critic
+  ```
+
+- set trainer_refs :
+  ```python
+  trainer_refs = [
+      DetachedPPOTrainer.options(
+          name=f"trainer{i}",
+          num_gpus=1,
+          max_concurrency=2
+      ).remote(
+          experience_maker_holder_name_list=[f"maker{x}" for x in target_makers(...)],
+          model_fn = trainer_model_fn(),
+          ...)
+      for i, env_info_trainer in enumerate(env_info_trainers)
+  ]
+  ```
+  The names in `experience_maker_holder_name_list` refer to the target makers that the trainer should send updated models to.
+  By setting `detached_trainer_name_list` and `experience_maker_holder_name_list`, we can customize the transmission graph.
+
+### Launch Jobs
+
+- define data_loader :
+
+  ```python
+  def data_loader_fn():
+      return = torch.utils.data.DataLoader(dataset=dataset)
+
+  ```
+
+- launch makers :
+
+  ```python
+  wait_tasks = []
+  for experience_holder_ref in experience_holder_refs:
+      wait_tasks.append(
+          experience_holder_ref.workingloop.remote(data_loader_fn(),
+                                                   num_steps=experience_steps))
+
+  ```
+
+- launch trainers :
+
+  ```python
+  for trainer_ref in trainer_refs:
+      wait_tasks.append(trainer_ref.fit.remote(total_steps, update_steps, train_epochs))
+  ```
+
+- wait for done :
+  ```python
+  ray.get(wait_tasks)
+  ```
+
+## Flexible Structure
+
+We can deploy different strategies to makers and trainers. Here are some notions.
+
+### 2 Makers 1 Trainer
+
+<p align="center">
+<img src="https://github.com/hpcaitech/public_assets/blob/main/applications/chat/2m1t.png?raw=true" width=600/>
+</p>
+
+### 2 Makers 2 Trainer
+
+<p align="center">
+<img src="https://github.com/hpcaitech/public_assets/blob/main/applications/chat/2m2t.png?raw=true" width=600/>
+</p>
+
+### Maker Inference Quantization
+
+<p align="center">
+<img src="https://github.com/hpcaitech/public_assets/blob/main/applications/chat/2m2t_quantize.png?raw=true" width=600/>
+</p>
+
+### Tensor Parallel
+
+<p align="center">
+<img src="https://github.com/hpcaitech/public_assets/blob/main/applications/chat/tp_ddp_hybrid.png?raw=true" width=600/>
+</p>
+
+## TODO
+
+- [ ] Support LoRA
+- [ ] Support TP & PP