Relax checks for flash-attn

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84

[cyanguwa]

I had a few problems running NeMo+TE with flash-attn on. Specifically, NVTE_FLASH_ATTN=1 python nemo/examples/nlp/language_modeling/megatron_gpt_pretraining.py did not run through the flash-attn path. I had to make a few changes as listed in the PR to make it work.

• attention_softmax_in_fp32: NeMo sets it to False by default, but this could be a NeMo usage problem. Or we can relax our check here in TE. We can think about this.
• apply_query_key_layer_scaling: do we expect users to always remember to add model.apply_query_key_layer_scaling=False to NeMo run command when they want to turn flash-attn on? Seems a bit too much work for an average user.
• the flash-attn path gets turned off if attention_mask is passed in. Could we make it still work, but just provide a warning instead? Otherwise, it's just too easy for the code to not go through the flash-attn path.

Thanks!

[ksivaman]

Hmm, I see the issue here of flash-attn being used only in a specific config of arguments. But I think what we have to ensure is that the FA path is used with TE's default settings, which currently is the case. For some of the args you specifically list:

• apply_query_key_layer_scaling: This should really be switched off by default in NeMo/MLM since its an old numeric hack to get LLM training in fp16, which isn't the default training precision anyway.
• attention_mask: If the user has passed an attention mask, I think we must assume that it's a custom attention mask that they want to use and not the redundant causal mask that we can discard. I did have a warning about this in the original PR but we got rid of it during review.
• attention_softmax_in_fp32: Maybe you're right about this one, we have an extra assert on L223 that we'll have to remove. My only gripe in that case is that we're ignoring this parameter entirely. We could set it to True explicitly in case of FA. We already do something similar here for the QK layer scaling.

[cyanguwa]

Maybe @ptrendx can chime in here?

• apply_query_key_layer_scaling: what if some users are still training with fp16? Could be some public users. I guess we still need to make sure fp16 converges?
• attention_mask: then we need to make sure that we don't pass in attention_mask, which again could be a user problem here, i.e. NeMo.
• attention_softmax_in_fp32: I agree.

[ptrendx]

• apply_query_key_layer_scaling: It should not apply to FlashAttention as the BMM1->softmax path is inside the kernel and performed in FP32 anyway
• attention_softmax_in_fp32: it also should not matter for FA, since softmax inside it uses FP32 math
• attention_mask: this is a bigger problem I think, since the FA kernel does not take the mask. We could potentially do the masked_fill to get the thrown away values to -10000, although this would eat into the performance gains.

[ksivaman]

Closing in favor of #84