fix llama

Signed-off-by: ftgreat <[email protected]>
arfmatos · Jun 4, 2023 · d48fa78 · d48fa78
1 parent befa604
commit d48fa78
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/flagai/model/blocks/llama_block.py b/flagai/model/blocks/llama_block.py
@@ -40,9 +40,9 @@ def __init__(self, layer_id, config ):
 
         self.layer_id = layer_id
         if config.flash_atten_llama_style:
-            from flash_attn.ops.rms_norm import RMSNorm
-            self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
-            self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+            import flash_attn
+            self.attention_norm = flash_attn.ops.rms_norm.RMSNorm(config.dim, eps=config.norm_eps)
+            self.ffn_norm = flash_attn.ops.rms_norm.RMSNorm(config.dim, eps=config.norm_eps)
         else:
             self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
             self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)

diff --git a/flagai/model/layers/attentions.py b/flagai/model/layers/attentions.py
@@ -171,11 +171,6 @@ def forward(
 
             xq, xk = apply_rotary_pos_emb(xq, xk, freqs_cis=freqs_cis)
 
-            xq = xq.view(bsz, seqlen, 1, self.n_local_heads, self.head_dim)
-            keys = keys.view(bsz, seqlen, 1, self.n_local_heads, self.head_dim)
-            values = values.view(bsz, seqlen, 1, self.n_local_heads, self.head_dim)
-            qkv = torch.concat([xq, keys, values], dim=2)
-
             if use_cache:
                 self.cache_k = self.cache_k.to(xq)
                 self.cache_v = self.cache_v.to(xq)
@@ -189,6 +184,11 @@ def forward(
                 keys = xk
                 values = xv 
 
+            xq = xq.view(bsz, seqlen, 1, self.n_local_heads, self.head_dim)
+            keys = keys.view(bsz, seqlen, 1, self.n_local_heads, self.head_dim)
+            values = values.view(bsz, seqlen, 1, self.n_local_heads, self.head_dim)
+            qkv = torch.concat([xq, keys, values], dim=2)
+
         if self.config.flash_atten or (self.config.flash_atten_llama_style and not self.training):
             qkv = einops.rearrange(qkv, 'b s ... -> (b s) ...')
 

diff --git a/flagai/model/llama_model.py b/flagai/model/llama_model.py
@@ -159,8 +159,8 @@ def __init__(self, config, **kwargs):
             self.layers.append(LLAMABlock(layer_id, config))
 
         if config.flash_atten_llama_style:
-            from flash_attn.ops.rms_norm import RMSNorm
-            self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+            import flash_attn
+            self.norm = flash_attn.ops.rms_norm.RMSNorm(config.dim, eps=config.norm_eps)
         else:
             self.norm = RMSNorm(config.dim, eps=config.norm_eps)
         if os.getenv("ENV_TYPE") == "deepspeed+mpu":