Incorrect‌ `attention dropout` and fake `flash_conformer`

[lzqlzzq]

Thanks for the great model! But I found some problems in the code.

Incorrect attention_dropout

In src/muq/muq/modules/flash_conformer.py, the dropout remains even the .eval() mode is on, would cause inconsistent output between eval() runs.

MuQ/src/muq/muq/modules/flash_conformer.py

Line 696 in 28847ea

hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=self.dropout_p, is_causal=self.is_causal)

It is ought to be something like:

    hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=self.dropout_p if self.training else 0.0, is_causal=self.is_causal) 

Fake flash_conformer

In src/muq/muq/models/muq_model.py, becase of the path structure, in a normal working directory, python interpreter would never resolve the modules.flash_conformer. So is_flash=True would never work.

I proposed a fix in: https://github.com/lzqlzzq/flashMuQ
This fix use AttentionInterface in 🤗transformers to dispatch actual attention implementation, as most modern models do.
I benchmarked the implementation with small batch. It has no numerical difference and wall time reduced to 0.86x, moreover 12x vram saving:

muq_large_attention_benchmark device=cuda batch=1 sequence_length=750 dtype=float16
muq_large_attention_backend configured_backend=eager actual_backend=eager allocated_before=607302144 allocated_after=608838144 allocated_delta=1536000 peak_delta=68911104 driver_before=None driver_after=None median_forward_ms=13.592
muq_large_attention_backend configured_backend=sdpa actual_backend=flash_attention allocated_before=607732224 allocated_after=609268224 allocated_delta=1536000 peak_delta=18874368 driver_before=None driver_after=None median_forward_ms=11.646
muq_large_attention_comparison sdpa_over_eager_allocated_delta=1.000000 sdpa_over_eager_forward_time=0.856852