update baichuan.py

vllm/model_executor/models/baichuan.py

@@ -399,6 +399,12 @@ class BaiChuanModel(nn.Module):
             
             for layername in loaded_params:
                 weight = params_dict[layername]
+                if "lm_head.weight" in layername and weight.shape[1] >= 4096:
+                    lay_key_words.append("lm_head.weight")
+                    combined_words = "|".join(lay_key_words)
+                    os.environ['LM_NN'] = '1'  
+                else:
+                    os.environ['LM_NN'] = '0' 
                 matches = re.findall(combined_words, layername)
                 if matches:      
                     # if self.use_gemm_pad and gemm_bank_conf(weight.data.shape[0]):