update files

config.json

@@ -13,7 +13,7 @@
   "hidden_size": 5120,
   "initializer_range": 0.02,
   "intermediate_size": 13824,
-  "max_position_embeddings": 262144,
+  "max_position_embeddings": 32768,
   "max_tokenizer_truncation": 262144,
   "model_type": "xverse",
   "num_attention_heads": 40,

modeling_xverse.py

@@ -33,6 +33,7 @@ from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
 from transformers.generation.utils import GenerationConfig
 from .configuration_xverse import XverseConfig
+from xformers import ops as xops
 
 
 logger = logging.get_logger(__name__)
@@ -48,13 +49,15 @@ def _make_causal_mask(
     Make causal mask used for bi-directional self-attention.
     """
     bsz, tgt_len = input_ids_shape
-    mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
+    mask = torch.full((tgt_len, tgt_len), torch.tensor(
+        torch.finfo(dtype).min, device=device), device=device)
     mask_cond = torch.arange(mask.size(-1), device=device)
     mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
     mask = mask.to(dtype)
 
     if past_key_values_length > 0:
-        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+        mask = torch.cat([torch.zeros(
+            tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
     return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
 
 
@@ -66,7 +69,8 @@ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int]
     bsz, src_len = mask.size()
     tgt_len = tgt_len if tgt_len is not None else src_len
 
-    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+    expanded_mask = mask[:, None, None, :].expand(
+        bsz, 1, tgt_len, src_len).to(dtype)
 
     inverted_mask = 1.0 - expanded_mask
 
@@ -84,8 +88,10 @@ class XverseRMSNorm(nn.Module):
 
     def forward(self, hidden_states):
         input_dtype = hidden_states.dtype
-        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        variance = hidden_states.to(torch.float32).pow(
+            2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * \
+            torch.rsqrt(variance + self.variance_epsilon)
 
         return (self.weight * hidden_states).to(input_dtype)
 
@@ -93,29 +99,47 @@ class XverseRMSNorm(nn.Module):
 class XverseRotaryEmbedding(torch.nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=500000, device=None):
         super().__init__()
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+        self.base = base
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        inv_freq = 1.0 / \
+            (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
         self.register_buffer("inv_freq", inv_freq)
 
         # Build here to make `torch.jit.trace` work.
         self.max_seq_len_cached = max_position_embeddings
-        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+
+        t = torch.arange(self.max_seq_len_cached,
+                         device=self.inv_freq.device, dtype=self.inv_freq.dtype)
         freqs = torch.einsum("i,j->ij", t, self.inv_freq)
         # Different from paper, but it uses a different permutation in order to obtain the same calculation
         emb = torch.cat((freqs, freqs), dim=-1)
-        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
-        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
+        self.register_buffer("cos_cached", emb.cos()[
+                             None, None, :, :], persistent=False)
+        self.register_buffer("sin_cached", emb.sin()[
+                             None, None, :, :], persistent=False)
 
     def forward(self, x, seq_len=None):
         # x: [bs, num_attention_heads, seq_len, head_size]
         # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
         if seq_len > self.max_seq_len_cached:
-            self.max_seq_len_cached = seq_len
-            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
-            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
-            # Different from paper, but it uses a different permutation in order to obtain the same calculation
+
+            t = torch.arange(seq_len, device=x.device, dtype=torch.float32)
+            dim = self.dim
+            alpha = (seq_len / (self.max_position_embeddings/2) - 1)
+            base = self.base * alpha ** (dim / (dim-2))
+            ntk_inv_freq = 1.0 / \
+                (base ** (torch.arange(0, dim, 2).float().to(x.device) / dim))
+
+            freqs = torch.einsum("i,j->ij", t, ntk_inv_freq)
             emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
-            self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
-            self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
+            cos_cached = emb.cos()[None, None, :, :]
+            sin_cached = emb.sin()[None, None, :, :]
+            return (
+                cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+                sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype)
+            )
+
         return (
             self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
             self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
@@ -125,7 +149,7 @@ class XverseRotaryEmbedding(torch.nn.Module):
 def rotate_half(x):
     """Rotates half the hidden dims of the input."""
     x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
+    x2 = x[..., x.shape[-1] // 2:]
     return torch.cat((-x2, x1), dim=-1)
 
 
@@ -173,11 +197,16 @@ class XverseAttention(nn.Module):
                 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                 f" and `num_heads`: {self.num_heads})."
             )
-        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
-        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
-        self.rotary_emb = XverseRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.rotary_emb = XverseRotaryEmbedding(
+            self.head_dim, max_position_embeddings=self.max_position_embeddings)
 
     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
         return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
@@ -190,64 +219,50 @@ class XverseAttention(nn.Module):
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        dropout: Optional[float] = 0.1,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        bsz, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        kv_seq_len = key_states.shape[-2]
-        if past_key_value is not None:
-            kv_seq_len += past_key_value[0].shape[-2]
-        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
-        # [bsz, nh, t, hd]
 
-        if past_key_value is not None:
-            # reuse k, v, self_attention
-            key_states = torch.cat([past_key_value[0], key_states], dim=2)
-            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        bsz, q_len, _ = hidden_states.size()  # [bsz, q_len, hidden_size]
 
-        past_key_value = (key_states, value_states) if use_cache else None
+        query_states = self.q_proj(hidden_states).view(
+            bsz, q_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(
+            bsz, q_len, self.num_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(
+            bsz, q_len, self.num_heads, self.head_dim)
+        # [bsz, q_len, nh, hd]
 
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        # [bsz, nh, q_len, hd]
 
-        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
-            raise ValueError(
-                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
-                f" {attn_weights.size()}"
-            )
+        kv_seq_len = key_states.shape[-2]  # q_len
+        n_head = key_states.shape[-3]  # nh
 
-        if attention_mask is not None:
-            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
-                raise ValueError(
-                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
-                )
-            attn_weights = attn_weights + attention_mask
-            attn_weights = torch.max(
-                attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min, device=attn_weights.device)
-            )
+        assert past_key_value is None, "past_key_value is not supported"
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids)
 
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-        attn_output = torch.matmul(attn_weights, value_states)
+        assert not output_attentions, "output_attentions is not supported"
+        assert not use_cache, "use_cache is not supported"
 
-        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
+        """
+        Input tensors must be in format ``[B, M, H, K]``, where B is the batch size, M 
+        the sequence length, H the number of heads, and K the embeding size per head
+        """
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
 
-        attn_output = attn_output.transpose(1, 2)
+        attn_output = xops.memory_efficient_attention(
+            query_states, key_states, value_states, attn_bias=xops.LowerTriangularMask(), p=dropout
+        )
+        
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-
         attn_output = self.o_proj(attn_output)
-
-        if not output_attentions:
-            attn_weights = None
-
-        return attn_output, attn_weights, past_key_value
+        return attn_output, None, None
 
 
 class XverseDecoderLayer(nn.Module):
@@ -260,8 +275,10 @@ class XverseDecoderLayer(nn.Module):
             intermediate_size=config.intermediate_size,
             hidden_act=config.hidden_act,
         )
-        self.input_layernorm = XverseRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = XverseRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.input_layernorm = XverseRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = XverseRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
         self,
@@ -426,6 +443,7 @@ XVERSE_INPUTS_DOCSTRING = r"""
             Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
 """
 
+
 @add_start_docstrings(
     "The bare Xverse Model outputting raw hidden-states without any specific head on top.",
     XVERSE_START_DOCSTRING,
@@ -443,8 +461,10 @@ class XverseModel(XversePreTrainedModel):
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
 
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        self.layers = nn.ModuleList([XverseDecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [XverseDecoderLayer(config) for _ in range(config.num_hidden_layers)])
         self.norm = XverseRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
         self.gradient_checkpointing = False
@@ -476,7 +496,8 @@ class XverseModel(XversePreTrainedModel):
                 inputs_embeds.device
             )
             combined_attention_mask = (
-                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask +
+                combined_attention_mask
             )
 
         return combined_attention_mask
@@ -504,13 +525,15 @@ class XverseModel(XversePreTrainedModel):
 
         # retrieve input_ids and inputs_embeds
         if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+            raise ValueError(
+                "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
         elif input_ids is not None:
             batch_size, seq_length = input_ids.shape
         elif inputs_embeds is not None:
             batch_size, seq_length, _ = inputs_embeds.shape
         else:
-            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+            raise ValueError(
+                "You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
         seq_length_with_past = seq_length
         past_key_values_length = 0
@@ -536,7 +559,8 @@ class XverseModel(XversePreTrainedModel):
                 (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
             )
         attention_mask = self._prepare_decoder_attention_mask(
-            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            attention_mask, (batch_size,
+                             seq_length), inputs_embeds, past_key_values_length
         )
 
         hidden_states = inputs_embeds
@@ -588,7 +612,8 @@ class XverseModel(XversePreTrainedModel):
             hidden_states = layer_outputs[0]
 
             if use_cache:
-                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+                next_decoder_cache += (
+                    layer_outputs[2 if output_attentions else 1],)
 
             if output_attentions:
                 all_self_attns += (layer_outputs[1],)
@@ -615,7 +640,8 @@ class XverseForCausalLM(XversePreTrainedModel):
         super().__init__(config)
         self.model = XverseModel(config)
 
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head = nn.Linear(
+            config.hidden_size, config.vocab_size, bias=False)
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -726,25 +752,32 @@ class XverseForCausalLM(XversePreTrainedModel):
             attentions=outputs.attentions,
         )
 
-    def _build_chat_input(self, tokenizer, messages: List[dict], max_new_tokens: int=2048):
+    def _build_chat_input(self, tokenizer, messages: List[dict], max_new_tokens: int = 2048):
         max_new_tokens = max_new_tokens or self.generation_config.max_new_tokens
         max_input_tokens = self.config.max_position_embeddings - max_new_tokens
-        max_input_tokens = max(self.config.max_position_embeddings // 2, max_input_tokens)
-        max_input_tokens = min(self.config.max_tokenizer_truncation, max_input_tokens)
+        max_input_tokens = max(
+            self.config.max_position_embeddings // 2, max_input_tokens)
+        max_input_tokens = min(
+            self.config.max_tokenizer_truncation, max_input_tokens)
 
         total_input, round_input = [], []
-        user_prompt_tokens = tokenizer.encode("Human: ", return_token_type_ids=False)
-        exec_prompt_tokens = tokenizer.encode("Exec: ", return_token_type_ids=False)
-        assist_prompt_tokens = tokenizer.encode("Assistant: ", return_token_type_ids=False)
+        user_prompt_tokens = tokenizer.encode(
+            "Human: ", return_token_type_ids=False)
+        exec_prompt_tokens = tokenizer.encode(
+            "Exec: ", return_token_type_ids=False)
+        assist_prompt_tokens = tokenizer.encode(
+            "Assistant: ", return_token_type_ids=False)
         assist_prompt_len = len(assist_prompt_tokens)
 
         for i, message in enumerate(messages[::-1]):
             if message['role'] == 'user' or message['role'] == 'exec':
                 user_content = f"{message['content']}\n\n"
                 content_tokens = user_prompt_tokens + tokenizer.encode(user_content, return_token_type_ids=False) if message['role'] == 'user' else \
-                    exec_prompt_tokens + tokenizer.encode(user_content, return_token_type_ids=False)
+                    exec_prompt_tokens + \
+                    tokenizer.encode(user_content, return_token_type_ids=False)
                 if i == 0:
-                    content_tokens = content_tokens[:max_input_tokens-assist_prompt_len]
+                    content_tokens = content_tokens[:max_input_tokens -
+                                                    assist_prompt_len]
                     content_tokens += assist_prompt_tokens
                 round_input = content_tokens + round_input
 
@@ -760,27 +793,33 @@ class XverseForCausalLM(XversePreTrainedModel):
                 round_input = []
             elif message['role'] == 'assistant':
                 assist_content = f"{message['content']}"
-                content_tokens = assist_prompt_tokens + tokenizer.encode(assist_content, return_token_type_ids=False)
-                round_input = content_tokens + [self.generation_config.eos_token_id] + round_input
+                content_tokens = assist_prompt_tokens + \
+                    tokenizer.encode(
+                        assist_content, return_token_type_ids=False)
+                round_input = content_tokens + \
+                    [self.generation_config.eos_token_id] + round_input
             elif message['role'] == 'system':
                 assert i == len(messages) - 1
                 user_content = f"{message['content']}\n"
-                content_tokens = tokenizer.encode(user_content, return_token_type_ids=False)
+                content_tokens = tokenizer.encode(
+                    user_content, return_token_type_ids=False)
                 round_input = user_prompt_tokens + content_tokens + round_input
                 if len(total_input) + len(round_input) > max_input_tokens:
                     break
                 else:
                     total_input = round_input + total_input
             else:
-                raise ValueError(f"message role not supported yet: {message['role']}")
+                raise ValueError(
+                    f"message role not supported yet: {message['role']}")
         total_input = torch.LongTensor([total_input]).to(self.device)
         return total_input
 
     @torch.no_grad()
     def chat(self, tokenizer, messages: List[dict], stream=False,
-             generation_config: Optional[GenerationConfig]=None):
+             generation_config: Optional[GenerationConfig] = None):
         generation_config = generation_config or self.generation_config
-        input_ids = self._build_chat_input(tokenizer, messages, generation_config.max_new_tokens)
+        input_ids = self._build_chat_input(
+            tokenizer, messages, generation_config.max_new_tokens)
         if stream:
             from transformers import TextIteratorStreamer
             from threading import Thread
@@ -788,7 +827,8 @@ class XverseForCausalLM(XversePreTrainedModel):
             self.__class__.generate = PreTrainedModel.generate
 
             def stream_generator():
-                generation_kwargs = dict(inputs=input_ids, generation_config=generation_config, streamer=streamer)
+                generation_kwargs = dict(
+                    inputs=input_ids, generation_config=generation_config, streamer=streamer)
                 thread = Thread(target=self.generate, kwargs=generation_kwargs)
                 thread.start()
                 for next_text in streamer:
@@ -797,8 +837,10 @@ class XverseForCausalLM(XversePreTrainedModel):
             return stream_generator()
         else:
             self.__class__.generate = PreTrainedModel.generate  # disable stream
-            outputs = self.generate(input_ids, generation_config=generation_config)
-            response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
+            outputs = self.generate(
+                input_ids, generation_config=generation_config)
+            response = tokenizer.decode(
+                outputs[0][len(input_ids[0]):], skip_special_tokens=True)
             return response
 
     def prepare_inputs_for_generation(
@@ -835,7 +877,8 @@ class XverseForCausalLM(XversePreTrainedModel):
     def _reorder_cache(past_key_values, beam_idx):
         reordered_past = ()
         for layer_past in past_key_values:
-            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+            reordered_past += (tuple(past_state.index_select(0, beam_idx)
+                               for past_state in layer_past),)
         return reordered_past
 
     def quantize(self, bit_length: int):
@@ -844,37 +887,44 @@ class XverseForCausalLM(XversePreTrainedModel):
         for layer in self.model.layers:
             layer.self_attn.q_proj = QuantizationLinear(
                 bit_length=bit_length,
-                weight=layer.self_attn.q_proj.weight.to(torch.cuda.current_device()),
+                weight=layer.self_attn.q_proj.weight.to(
+                    torch.cuda.current_device()),
                 device=layer.self_attn.q_proj.weight.device,
             )
             layer.self_attn.k_proj = QuantizationLinear(
                 bit_length=bit_length,
-                weight=layer.self_attn.k_proj.weight.to(torch.cuda.current_device()),
+                weight=layer.self_attn.k_proj.weight.to(
+                    torch.cuda.current_device()),
                 device=layer.self_attn.k_proj.weight.device
             )
             layer.self_attn.v_proj = QuantizationLinear(
                 bit_length=bit_length,
-                weight=layer.self_attn.v_proj.weight.to(torch.cuda.current_device()),
+                weight=layer.self_attn.v_proj.weight.to(
+                    torch.cuda.current_device()),
                 device=layer.self_attn.v_proj.weight.device
             )
             layer.self_attn.o_proj = QuantizationLinear(
                 bit_length=bit_length,
-                weight=layer.self_attn.o_proj.weight.to(torch.cuda.current_device()),
+                weight=layer.self_attn.o_proj.weight.to(
+                    torch.cuda.current_device()),
                 device=layer.self_attn.o_proj.weight.device
             )
             layer.mlp.gate_proj = QuantizationLinear(
                 bit_length=bit_length,
-                weight=layer.mlp.gate_proj.weight.to(torch.cuda.current_device()),
+                weight=layer.mlp.gate_proj.weight.to(
+                    torch.cuda.current_device()),
                 device=layer.mlp.gate_proj.weight.device
             )
             layer.mlp.down_proj = QuantizationLinear(
                 bit_length=bit_length,
-                weight=layer.mlp.down_proj.weight.to(torch.cuda.current_device()),
+                weight=layer.mlp.down_proj.weight.to(
+                    torch.cuda.current_device()),
                 device=layer.mlp.down_proj.weight.device
             )
             layer.mlp.up_proj = QuantizationLinear(
                 bit_length=bit_length,
-                weight=layer.mlp.up_proj.weight.to(torch.cuda.current_device()),
+                weight=layer.mlp.up_proj.weight.to(
+                    torch.cuda.current_device()),
                 device=layer.mlp.up_proj.weight.device
             )
 

ms_wrapper.py

@@ -0,0 +1,95 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import torch
+
+from modelscope.models.base import TorchModel
+from modelscope.preprocessors.base import Preprocessor
+from modelscope.pipelines.base import Model, Pipeline
+from modelscope.utils.config import Config
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.preprocessors.builder import PREPROCESSORS
+from modelscope.models.builder import MODELS
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+@MODELS.register_module('text-generation', module_name='XVERSE-13B')
+class XVERSE13BTextGeneration(TorchModel):
+
+    def __init__(self, model_dir, *args, **kwargs):
+        super().__init__(model_dir, *args, **kwargs)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
+        self.model = self.model.eval()
+
+    def forward(self, inputs, **forward_params):
+        inputs = self.tokenizer(inputs, return_tensors='pt').input_ids
+        inputs = inputs.cuda()
+        generated_ids = self.model.generate(inputs, eos_token_id=self.tokenizer.eos_token_id, **forward_params)
+        return self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+
+
+@PIPELINES.register_module('text-generation', module_name='XVERSE-13B-pipeline')
+class XVERSE13BTextGenerationPipeline(Pipeline):
+    """ Give simple introduction to this pipeline.
+
+    Examples:
+
+    >>> from modelscope.pipelines import pipeline
+    >>> input = "Hello, ModelScope!"
+    >>> my_pipeline = pipeline('text-generation', 'xverse/XVERSE-13B')
+    >>> result = my_pipeline(input)
+
+    """
+
+    def __init__(self, model, **kwargs):
+        """
+        use `model` and `preprocessor` to create a custom pipeline for prediction
+        Args:
+            model: model id on modelscope hub.
+            preprocessor: the class of method be init_preprocessor
+        """
+        assert isinstance(model, str) or isinstance(model, Model), \
+            'model must be a single str or Model'
+        if isinstance(model, str):
+            pipe_model = Model.from_pretrained(model)
+        elif isinstance(model, Model):
+            pipe_model = model
+        else:
+            raise NotImplementedError
+
+        super().__init__(model=pipe_model, **kwargs)
+
+    def _sanitize_parameters(self, **pipeline_parameters):
+        """
+        this method should sanitize the keyword args to preprocessor params,
+        forward params and postprocess params on '__call__' or '_process_single' method
+        considered to be a normal classmethod with default implementation / output
+
+        Default Returns:
+            Dict[str, str]:  preprocess_params = {}
+            Dict[str, str]:  forward_params = {}
+            Dict[str, str]:  postprocess_params = pipeline_parameters
+        """
+        return {}, pipeline_parameters, {}
+
+    def preprocess(self, inputs, **preprocess_params):
+        return inputs
+
+    def forward(self, inputs, **forward_params):
+        """ Provide default implementation using self.model and user can reimplement it
+        """ 
+        output = super().forward(inputs, **forward_params)
+        return {'text': output}
+
+    def postprocess(self, inputs):
+        """ If current pipeline support model reuse, common postprocess
+            code should be write here.
+
+        Args:
+            inputs:  input data
+
+        Return:
+            dict of results:  a dict containing outputs of model, each
+                output should have the standard output name.
+        """
+        return inputs
+