fix default value of xentropy

modeling_flash_llama.py

@@ -44,11 +44,19 @@ try:
     from flash_attn.bert_padding import unpad_input, pad_input
     flash_attn_v2_installed = True
     print('>>>> Flash Attention installed')
-    from flash_attn.losses.cross_entropy import CrossEntropyLoss as xCrossEntropyLoss
 except ImportError:
     flash_attn_v2_installed = False
     raise ImportError('Please install Flash Attention: `pip install flash-attn --no-build-isolation`')
 
+try:
+    from flash_attn.losses.cross_entropy import CrossEntropyLoss as xCrossEntropyLoss
+    flash_xentropy_installed = True
+    print('>>>> xentropy installed')
+except ImportError:
+    flash_xentropy_installed = False
+    raise ImportError('Please install xentropy kernels: `pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/xentropy`')
+
+
 try:
     from flash_attn.layers.rotary import apply_rotary_emb_func
     flash_rope_installed = True
@@ -774,7 +782,7 @@ class LlamaForCausalLM(LlamaPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         only_last_logit: Optional[bool] = None,
-        xentropy: Optional[bool] = None,
+        xentropy: Optional[bool] = False,
         is_padded_inputs: Optional[bool] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
@@ -869,7 +877,7 @@ class LlamaForCausalLM(LlamaPreTrainedModel):
 
     def prepare_inputs_for_generation(
         self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, only_last_logit=False,
-        xentropy=False, **kwargs
+        xentropy=True, **kwargs
     ):
         if past_key_values:
             input_ids = input_ids[:, -1:]