PP microbatching + block-causal FlexAttention composability issue

[tianyu-l]

Today init_attention_mask is called before PP does microbatch split
https://github.com/pytorch/torchtitan/blob/main/torchtitan/train.py#L421

I haven't tested, but this will likely cause wrong block mask being applied to non-first microbatches.
E.g. consider a local batch

batch = [
  b0,
  b1,
  b2,
  b3,
]

Block mask mask is created for batch, but after (say, size-1) microbatching mask will be used for 4 different smaller batches [b0], [b1], [b2], [b3]. For [b0] it might be fine, but for the others the mask is wrong.

The solution could be either of

1. do init_attention_mask after microbatching.
2. when PP does microbatching, modify the block mask as well.

cc @fegin @H-Huang @drisspg

[fegin]

3. If we are going to make CP register a model forward hook, the hook can modify the BlockMask, provided that CP knows PP's information.

The above only work when CP is enabled but this issue can happen even if CP is not enabled.

[ezyang]

Just a question for my curiosity: the way that init_attention_mask works seems a bit... ugly to me?

torchtitan/torchtitan/models/attention.py

Line 186 in 60645bc

FlexAttention.block_masks[mask_key] = block_mask

Is the reason the block masks member is a class variable because you want to share them across all attention blocks? The general pattern smacks of wanting some sort of laziness, where you init it the first time you get a batch for real (so at that point you get a microbatch) (although this is not easily compilable!)

[drisspg]

The pattern that we have recommended is to have blockmask be an input to your model. Torchtune had set this up and made a thin type that could either be a blockmask or tensor or none so that it could be used for sdpa or flex

[ezyang]

Yes, and then it would be obvious that you also have to microbatch it when doing PP

[tianyu-l]

One concern we had in the past is that
Different layers of TransformerBlock have different attention masks, e.g. some layers use sliding window, some layers use full causal, etc. So it will become a bit cumbersome to pass all these masks around and let PP microbatch everything.

Another option is to simply pass the eos_id around (which is the only thing affecting the masks from iteration to iteration), which PP microbatching doesn't need to do anything.

This is further complicated by CP inputs sharding with load balancing. See discussions in #1728 (comment)

[fegin]

The pattern that we have recommended is to have blockmask be an input to your model. Torchtune had set this up and made a thin type that could either be a blockmask or tensor or none so that it could be used for sdpa or flex

We are thinking to do this as well. But as @tianyu-l described, the main concern was different layers masks and how does it compose

But PP cannot micro-batch BlockMask atm as it doesn't understand BlockMask. I'm wondering if we should let PP understands BlockMask, just like PP understands torch.Tensor as BlockMask is from torch core. cc., @H-Huang