Fix attention mask handling in the Hybrid Engine Bloom flow

[deepcharm]

The Bloom flow in Hybrid Engine applies the same transformation of the input mask which is already performed earlier by the transformers BloomModel::forward.

This results in the non-convergence of scores, specifically in Deepspeed Chat on different accelerators, including CUDA and HPU.

The fix removes redundant mask transformation and application, producing correct convergence.

[lekurile]

Hello @deepcharm,

Thank you for the contribution. I've studied the code in BloomModel::forward and can see that the call to _prepare_4d_causal_attention_mask is made here. I believe here is where the referenced 1-mask operation is happening in the _prepare_4d_causal_attention_mask function.

I think in this case the change makes sense since we don't want to do the redundant mask processing operation, however, I just want to be careful that we don't change ds_attention.py masking behavior fundamentally and account for model support beyond BLOOM, since I'm not sure if all transformer model implementations do this mask processing outside the transformers block.

One option is to add an optional config parameter in our inference config.py that will enable skipping of this operation. We can set this parameter to False by default and to True only in the BLOOM container.

If more models need this behavior, we can enable this in the corresponding model-specific container.

Please let me know if you have feedback or questions.

Thanks,
Lev

[deepcharm]

Hi @lekurile

Thanks for your feedback. I will implement the option that you have described and submit another patch.

Max

[tjruwase]

@deepcharm, thank for improving the PR. Please ping us again when it is ready for review.

[deepcharm]

Hi @lekurile, @tjruwase

As advised, I've added an optional config parameter invert_mask into DeepSpeedInferenceConfig (True by default),
which enables skipping the invert operation for some transformer implementations, such as BLOOM.

Kindly review the change. Thanks.

[lekurile]

Hi @lekurile, @tjruwase

As advised, I've added an optional config parameter invert_mask into DeepSpeedInferenceConfig (True by default), which enables skipping the invert operation for some transformer implementations, such as BLOOM.

Kindly review the change. Thanks.

@deepcharm Thank you! Looks good to me. Approved and running checks.