HuggingFaceNmtEngine does not use SDPA attention implementation

[ddaspit]

The HuggingFaceNmEngine class sets output_attentions=True when calling the generate method on the model. This is not compatible with the SDPA attention implementation, so Huggingface reverts to using the default eager implementation. The class outputs the attentions so that they can be used to infer the alignment. A config option should be added to the class that allows the caller to disable inferring alignment from attention. This will allow Huggingface to use the SDPA attention implementation when it is configured.

[Enkidu93]

The HuggingFaceNmEngine class sets output_attentions=True when calling the generate method on the model. This is not compatible with the SDPA attention implementation, so Huggingface reverts to using the default eager implementation. The class outputs the attentions so that they can be used to infer the alignment. A config option should be added to the class that allows the caller to disable inferring alignment from attention. This will allow Huggingface to use the SDPA attention implementation when it is configured.

Is this causing the slow-downs observed in Serval jobs? I wondered about SDPA being the cause of the differences and searched for it in past issues only to find that you created this issue 😆.

[ddaspit]

I think this only affects inferencing, so it wouldn't explain slower training.

[Enkidu93]

I think this only affects inferencing, so it wouldn't explain slower training.

Oh, right - good point

[pmachapman]

@ddaspit Before I write unit tests and create a PR, I want to be sure I understood this task correctly. I have created a branch main...config_output_attentions with the change where I make output_attentions a value that can be set to False in the args.

Setting it to False will cause _try_translate_n_batch to crash, as output["token_attentions"] will be None. (I couldn't figure a way around this, as the code in this function needed the attention weights which were not in the output).

Is this the correct way to go about this task?

[ddaspit]

We just need to handle the case where output["token_attentions"] is None gracefully. Basically, this means that the word alignment will not be computed and will be blank.