Upload 2 files

README.md

@@ -129,16 +129,29 @@ config = GenerationConfig.from_pretrained("Qwen/Qwen-7B-Chat-Int4", trust_remote
 response, history = model.chat(tokenizer, "Hi", history=None, generation_config=config)
 ```
 
+### 效果评测
+
+我们对BF16和Int4模型在基准评测上做了测试，发现量化模型效果损失较小，结果如下所示：
+
+We illustrate the model performance of both BF16 and Int4 models on the benchmark, and we find that the quantized model does not suffer from significant performance degradation. Results are shown below:
+
+|  Quantization |   MMLU     |  CEval (val) |  GSM8K |  Humaneval |
+| ------------- | :--------: | :----------: | :----: | :--------: |
+| BF16          |    53.9    |     54.2     |  41.1  |    24.4    |
+| Int4          |    52.6    |     52.9     |  38.1  |    23.8    |
+
+
+
 ### 推理速度 (Inference Speed)
 
-我们测算了BF16和Int4模型生成8192个token的平均推理速度。如图所示：
+我们测算了BF16和Int4模型生成2048和8192个token的平均推理速度。如图所示：
 
-We measured the average inference speed of generating 8192 tokens under BF16 precision and Int4 quantization level, respectively.
+We measured the average inference speed of generating 2048 and 8192 tokens under BF16 precision and Int4 quantization level, respectively.
 
-| Quantization Level | Inference Speed (tokens/s) |
-| ------------------ | :-------------------------:|
-|        BF16        | 22.59                      |
-|        Int4        | 24.91                      |
+|  Quantization | Speed (2048 tokens) | Speed (8192 tokens) |
+| ------------- | :------------------:| :------------------:|
+|      BF16     | 30.53               | 28.51               |
+|      Int4     | 45.60               | 33.83               |
 
 具体而言，我们记录在长度为1的上下文的条件下生成8192个token的性能。评测运行于单张A100-SXM4-80G GPU，使用PyTorch 2.0.1和CUDA 11.4。推理速度是生成8192个token的速度均值。
 
@@ -159,8 +172,6 @@ We also profile the peak GPU memory usage for encoding 2048 tokens as context (a
 
 The above speed and memory profiling are conducted using [this script](https://qianwen-res.oss-cn-beijing.aliyuncs.com/profile.py).
 
-
-
 ## Tokenizer
 
 > 注：作为术语的“tokenization”在中文中尚无共识的概念对应，本文档采用英文表达以利说明。
@@ -345,7 +356,6 @@ Qwen-7B-Chat also has the capability to be used as a [HuggingFace Agent](https:/
 | StarCoder-15.5B |      87.04      |    87.96    |   68.89   |
 | **Qwen-7B**     |      90.74      |    92.59    |   74.07   |
 
-
 ## FAQ
 
 如遇到问题，敬请查阅[FAQ](https://github.com/QwenLM/Qwen-7B/blob/main/FAQ_zh.md)以及issue区，如仍无法解决再提交issue。
@@ -364,4 +374,3 @@ Our code and checkpoints are open to research purpose, and they are allowed for
 
 If you are interested to leave a message to either our research team or product team, feel free to send an email to [email protected].
 
-

modeling_qwen.py

@@ -131,6 +131,7 @@ class FlashSelfAttention(torch.nn.Module):
         assert all((i.is_cuda for i in (q, k, v)))
         batch_size, seqlen_q = q.shape[0], q.shape[1]
         seqlen_k = k.shape[1]
+
         q, k, v = [rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]]
         cu_seqlens_q = torch.arange(
             0,
@@ -155,6 +156,7 @@ class FlashSelfAttention(torch.nn.Module):
                 device=q.device,
             )
             self.dropout_p = 0
+
         output = flash_attn_unpadded_func(
             q,
             k,
@@ -168,7 +170,8 @@ class FlashSelfAttention(torch.nn.Module):
             causal=is_causal,
         )
 
-        output = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
+        new_shape = (batch_size, output.shape[0] // batch_size) + output.shape[1:]
+        output = output.view(new_shape)
         return output
 
 
@@ -220,19 +223,6 @@ class QWenAttention(nn.Module):
 
         self.bf16 = config.bf16
 
-        if config.rotary_pct == 1.0:
-            self.rotary_ndims = None
-        else:
-            assert config.rotary_pct < 1
-            self.rotary_ndims = int(
-                self.hidden_size_per_attention_head * config.rotary_pct
-            )
-        dim = (
-            self.rotary_ndims
-            if self.rotary_ndims is not None
-            else self.hidden_size_per_attention_head
-        )
-        self.rotary_emb = RotaryEmbedding(dim, base=config.rotary_emb_base)
 
         self.use_dynamic_ntk = config.use_dynamic_ntk
         self.use_logn_attn = config.use_logn_attn
@@ -242,7 +232,6 @@ class QWenAttention(nn.Module):
             for i in range(1, 32768)
         ]
         self.logn_tensor = torch.tensor(logn_list)[None, :, None, None]
-        self._ntk_cached = 1.0
 
         self.attn_dropout = nn.Dropout(config.attn_dropout_prob)
 
@@ -351,6 +340,7 @@ class QWenAttention(nn.Module):
     def forward(
         self,
         hidden_states: Optional[Tuple[torch.FloatTensor]],
+        rotary_pos_emb: Optional[List[torch.Tensor]] = None,
         layer_past: Optional[Tuple[torch.Tensor]] = None,
         attention_mask: Optional[torch.FloatTensor] = None,
         head_mask: Optional[torch.FloatTensor] = None,
@@ -361,43 +351,19 @@ class QWenAttention(nn.Module):
     ):
 
         mixed_x_layer = self.c_attn(hidden_states)
+
         query, key, value = mixed_x_layer.split(self.split_size, dim=2)
 
         query = self._split_heads(query, self.num_heads, self.head_dim)
         key = self._split_heads(key, self.num_heads, self.head_dim)
         value = self._split_heads(value, self.num_heads, self.head_dim)
 
-        kv_seq_len = hidden_states.size()[1]
-        if layer_past:
-            # layer past[0] shape: bs * seq_len * head_num * dim
-            kv_seq_len += layer_past[0].shape[1]
-        if (
-            self.use_dynamic_ntk
-            and kv_seq_len == hidden_states.size()[1]
-            and not self.training
-        ):
-            context_value = math.log(kv_seq_len / self.seq_length, 2) + 1
-            ntk_alpha = 2 ** math.ceil(context_value) - 1
-            ntk_alpha = max(ntk_alpha, 1)
-            self._ntk_cached = ntk_alpha
-        else:
-            ntk_alpha = self._ntk_cached
-        rotary_pos_emb = self.rotary_emb(kv_seq_len, ntk_alpha=ntk_alpha).to(
-            hidden_states.device
-        )
-
-        if rotary_pos_emb is not None:
-            if isinstance(rotary_pos_emb, tuple):
-                rotary_pos_emb = rotary_pos_emb
-            else:
-                rotary_pos_emb = (rotary_pos_emb,) * 2
-
         if rotary_pos_emb is not None:
+            cur_len = query.shape[1]
+            rotary_pos_emb = [i[:, -cur_len:, :, :] for i in rotary_pos_emb]
+            rotary_pos_emb = (rotary_pos_emb,) * 2
             q_pos_emb, k_pos_emb = rotary_pos_emb
             # Slice the pos emb for current inference
-            cur_len = query.shape[1]
-            q_pos_emb = q_pos_emb[:, -cur_len:, :, :]
-            k_pos_emb = k_pos_emb[:, -cur_len:, :, :]
             query = apply_rotary_pos_emb(query, q_pos_emb)
             key = apply_rotary_pos_emb(key, k_pos_emb)
 
@@ -428,9 +394,9 @@ class QWenAttention(nn.Module):
             q, k, v = query, key, value
             context_layer = self.core_attention_flash(q, k, v)
 
-            context_layer = rearrange(
-                context_layer, "b s h d -> b s (h d)"
-            ).contiguous()
+            # b s h d -> b s (h d)
+            context_layer = context_layer.flatten(2,3).contiguous()
+
         else:
             query = query.permute(0, 2, 1, 3)
             key = key.permute(0, 2, 1, 3)
@@ -443,6 +409,7 @@ class QWenAttention(nn.Module):
             )
 
         attn_output = self.c_proj(context_layer)
+
         outputs = (attn_output, present)
         if output_attentions:
             if (
@@ -476,7 +443,6 @@ class QWenMLP(nn.Module):
         output = self.c_proj(intermediate_parallel)
         return output
 
-
 class QWenBlock(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -498,6 +464,7 @@ class QWenBlock(nn.Module):
     def forward(
         self,
         hidden_states: Optional[Tuple[torch.FloatTensor]],
+        rotary_pos_emb: Optional[List[torch.Tensor]] = None,
         layer_past: Optional[Tuple[torch.Tensor]] = None,
         attention_mask: Optional[torch.FloatTensor] = None,
         head_mask: Optional[torch.FloatTensor] = None,
@@ -510,6 +477,7 @@ class QWenBlock(nn.Module):
 
         attn_outputs = self.attn(
             layernorm_output,
+            rotary_pos_emb,
             layer_past=layer_past,
             attention_mask=attention_mask,
             head_mask=head_mask,
@@ -585,10 +553,28 @@ class QWenModel(QWenPreTrainedModel):
         self.embed_dim = config.hidden_size
 
         self.gradient_checkpointing = False
+        self.use_dynamic_ntk = config.use_dynamic_ntk
+        self.seq_length = config.seq_length
 
         self.wte = nn.Embedding(self.vocab_size, self.embed_dim)
 
         self.drop = nn.Dropout(config.emb_dropout_prob)
+
+
+        if config.rotary_pct == 1.0:
+            self.rotary_ndims = None
+        else:
+            assert config.rotary_pct < 1
+            self.rotary_ndims = int(
+                config.kv_channels * config.rotary_pct
+            )
+        dim = (
+            self.rotary_ndims
+            if self.rotary_ndims is not None
+            else config.kv_channels
+        )
+        self.rotary_emb = RotaryEmbedding(dim, base=config.rotary_emb_base)
+
         self.h = nn.ModuleList(
             [
                 QWenBlock(
@@ -692,6 +678,25 @@ class QWenModel(QWenPreTrainedModel):
             inputs_embeds = self.wte(input_ids)
         hidden_states = inputs_embeds
 
+        kv_seq_len = hidden_states.size()[1]
+        if past_key_values[0] is not None:
+            # past key values[0][0] shape: bs * seq_len * head_num * dim
+            kv_seq_len += past_key_values[0][0].shape[1]
+        if (
+            self.use_dynamic_ntk
+            and kv_seq_len == hidden_states.size()[1]
+            and not self.training
+        ):
+            context_value = math.log(kv_seq_len / self.seq_length, 2) + 1
+            ntk_alpha = 2 ** math.ceil(context_value) - 1
+            ntk_alpha = max(ntk_alpha, 1)
+        else:
+            ntk_alpha = self.rotary_emb._ntk_alpha_cached
+
+        rotary_pos_emb = self.rotary_emb(kv_seq_len, ntk_alpha=ntk_alpha)
+        for idx in range(len(rotary_pos_emb)):
+            rotary_pos_emb[idx] = rotary_pos_emb[idx].to(hidden_states.device)
+
         hidden_states = self.drop(hidden_states)
         output_shape = input_shape + (hidden_states.size(-1),)
 
@@ -722,6 +727,7 @@ class QWenModel(QWenPreTrainedModel):
                 outputs = torch.utils.checkpoint.checkpoint(
                     create_custom_forward(block),
                     hidden_states,
+                    rotary_pos_emb,
                     None,
                     attention_mask,
                     head_mask[i],
@@ -732,6 +738,7 @@ class QWenModel(QWenPreTrainedModel):
                 outputs = block(
                     hidden_states,
                     layer_past=layer_past,
+                    rotary_pos_emb=rotary_pos_emb,
                     attention_mask=attention_mask,
                     head_mask=head_mask[i],
                     encoder_hidden_states=encoder_hidden_states,
@@ -1142,14 +1149,19 @@ class RotaryEmbedding(torch.nn.Module):
             self._ntk_alpha_cached = ntk_alpha
             seq = torch.arange(self._seq_len_cached, device=self.inv_freq.device)
             freqs = torch.outer(seq.type_as(self.inv_freq), self.inv_freq)
+            
             emb = torch.cat((freqs, freqs), dim=-1)
             from einops import rearrange
 
-            self._rotary_pos_emb_cache = rearrange(emb, "n d -> 1 n 1 d")
+            emb = rearrange(emb, "n d -> 1 n 1 d")
+
+            cos, sin = emb.cos(), emb.sin()
+            self._rotary_pos_emb_cache = [cos, sin]
 
     def forward(self, max_seq_len, offset=0, ntk_alpha=1.0):
         self.update_rotary_pos_emb_cache(max_seq_len, offset, ntk_alpha)
-        return self._rotary_pos_emb_cache[:, offset : offset + max_seq_len]
+        cos, sin = self._rotary_pos_emb_cache
+        return [cos[:, offset : offset + max_seq_len], sin[:, offset : offset + max_seq_len]]
 
 
 def _rotate_half(x):
@@ -1161,19 +1173,20 @@ def _rotate_half(x):
 
 
 def apply_rotary_pos_emb(t, freqs):
+    cos, sin = freqs
     if apply_rotary_emb_func is not None and t.is_cuda:
         t_ = t.float()
-        freqs = freqs.squeeze(0).squeeze(1)
-        cos = freqs[:, : freqs.shape[-1] // 2].cos()
-        sin = freqs[:, : freqs.shape[-1] // 2].sin()
+        cos = cos.squeeze(0).squeeze(1)[:, : cos.shape[-1] // 2]
+        sin = sin.squeeze(0).squeeze(1)[:, : sin.shape[-1] // 2]
         output = apply_rotary_emb_func(t_, cos, sin).type_as(t)
         return output
     else:
-        rot_dim = freqs.shape[-1]
+        rot_dim = freqs[0].shape[-1]
+        cos, sin = freqs
         t_, t_pass_ = t[..., :rot_dim], t[..., rot_dim:]
         t_ = t_.float()
         t_pass_ = t_pass_.float()
-        t_ = (t_ * freqs.cos()) + (_rotate_half(t_) * freqs.sin())
+        t_ = (t_ * cos) + (_rotate_half(t_) * sin)
         return torch.cat((t_, t_pass_), dim=-1).type_as(t)