Mllama fixes #39182
Mllama fixes #39182
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: mllama |
| class MllamaVisionModel(MllamaPreTrainedModel): | ||
| config_class = MllamaVisionConfig | ||
| base_model_prefix = "vision_model" | ||
| _supports_flash_attn_2 = False # the vision model always adds a 4D attn mask which is not supported by FA2 |
There was a problem hiding this comment.
Nice catch! IMO we can still run FA2 with vision module but we need to prepare the mask correctly. In text models usually for FA2, we keep the 2D mask and don't expand it to 4D.
We can do similar thing in Mllama and skip the Reshape to 2D and create 4D attention mask part in case of FA2. We might need to check cross attention as well, which also uses 4D mask
| class MllamaModel(MllamaPreTrainedModel): | ||
| _checkpoint_conversion_mapping = {"language_model.model": "language_model"} | ||
| _supports_quantized_cache = False # quant cache not supported in encoder-decoder setting | ||
| _supports_flash_attn_2 = False # the vision model does not support FA2 |
There was a problem hiding this comment.
_supports_flash_attn_2 = False should be defined only on the module that doesn't support it, i.e. on MllamaVisionModel. I guess it's the issue with tests is that they don't check all submodules for _supports_flash_attn_2
There was a problem hiding this comment.
Ok, I was wondering about that. I will check it out and revert the change
3 participants
This PR adds the
is_causalattribute to some Attention modules in mllama and disables FA2 for theMllamaVisionModel.Some tests used to fail when
is_causalwas missing:MllamaForCausalLMModelTest::test_flash_attn_2_fp32_ln,MllamaForConditionalGenerationModelTest::test_eager_matches_fa2_generate, ...and when it was added FA2 failed for these tests, on both Mi355 and A100. This is because the vision model uses a 4D attn mask which is not supported by FA2.Not an expert in VLM so can you chek
is_causalis right @zucchini-nlp please?I also added Expectations for AMD MI355. After these changes, we go from
15 failed, 185 passed, 94 skippedto196 passed, 98 skippedon AMD MI355.