Update modeling_zhinao.py

modeling_zhinao.py

@@ -748,6 +748,17 @@ class ZhinaoForCausalLM(ZhinaoPreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
+        if config.use_flash_attn == "auto":
+            if flash_attn_varlen_func:
+                if config.bf16 or config.fp16:
+                    logger.warn("Try importing flash-attention.")
+                    config.use_flash_attn = True
+                else:
+                    config.use_flash_attn = False
+                    logger.warn("Flash attention will be disabled because it does NOT support fp32.")
+            else:
+                config.use_flash_attn = False
+                logger.warn("Please install FlashAttention first, " "e.g., with pip install flash-attn")   
         self.model = ZhinaoModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
@@ -761,19 +772,7 @@ class ZhinaoForCausalLM(ZhinaoPreTrainedModel):
         if config.fp16:
             self.model.half()
             self.lm_head.half()
-            self.linear.half()
-
-        if config.use_flash_attn == "auto":
-            if flash_attn_varlen_func:
-                if config.bf16 or config.fp16:
-                    logger.warn("Try importing flash-attention.")
-                    config.use_flash_attn = True
-                else:
-                    config.use_flash_attn = False
-                    logger.warn("Flash attention will be disabled because it does NOT support fp32.")
-            else:
-                config.use_flash_attn = False
-                logger.warn("Please install FlashAttention first, " "e.g., with pip install flash-attn")                   
+            self.linear.half()                
 
         self.post_init()