sequence parallel for uneven heads

[inkcherry]

In sequence_parallel (Ulysses), the sequence parallel size is constrained by the requirement to be divisible by the number of heads, which prevents some models/workloads from setting a specific sequence parallel size. This PR implements uneven all-to-all heads splitting.

• both support batch first (b,s,...) and seq_len first(s,b..) layout.
• Added unit tests with numerical checks. Locally also tested with 7 heads with sp=4 and 20 heads with sp=8, and it passed.

[delock]

@inkcherry thanks, I have no further questions. Hi @tjruwase @loadams this PR is to enable sequence parallel for model with number of heads not power of two, which is requested from customer. Can this PR be reviewed? Thanks!

[loadams]

@inkcherry thanks, I have no further questions. Hi @tjruwase @loadams this PR is to enable sequence parallel for model with number of heads not power of two, which is requested from customer. Can this PR be reviewed? Thanks!

Thanks @delock, things look good now, but we will just need to get a review.