How to use flash attention when setting "--reset-position-ids"?

[t90tank]

It seems that Megatron do not pass attention mask to flash attention, but when setting "--reset-position-ids", different inputs have different attention masks.
Does Megatron support this case?

Also, can we open sequence parallel when setting "--reset-position-ids"?

[torshie]

It seems --reset-position-ids doesn't work with flash attention. I checked flash attention's API, it supports causal attention mask only, not custom attention masks.

[mayank31398]

Yes, this won't work
See this PR: bigcode-project#53 in the BigCode fork

[t90tank]

Thanks, I will try other ways.