FlexAttention `BlockMask` breaks Pipeline Parallelism with split-backward (zero-bubble / multi) schedules

https://github.com/pytorch/torchtitan/pull/3571 discovered this issue; the problematic CI tests are temporarily disabled. ## Summary When a language model uses **FlexAttention** (the default attention backend after SDPA was deprecated for language models) together with **Pipeline Parallelism**, training fails on any schedule that computes the input gradient separately from the weight gradient (zero-bubble and CSV/multi schedules): ``` AttributeError: 'BlockMask' object has no attribute 'requires_grad' ``` Full-backward schedules (`1F1B`, `GPipe`, `Interleaved1F1B`) are **not** affected. ## Root cause The trainer builds a FlexAttention `BlockMask` and forwards it to **every** PP stage as a keyword argument through the pipeline schedule: ```python # torchtitan/trainer.py (forward_backward_step) self.pp_schedule.step( inputs, **extra_inputs, **extra_kwargs, # contains attention_masks (a BlockMask) and positions ... ) ``` Split-backward schedules call `torch.distributed.pipelining`'s `stage_backward_input`, which walks **all** recorded stage inputs to collect their `grad_fn`s: ```python # torch/distributed/pipelining/_backward.py def _get_grad_fn_or_grad_acc(t): if t.requires_grad and t.grad_fn is None: # <-- t is a BlockMask here ... ``` `positions` survives this because it is a plain `Tensor` (`requires_grad=False`). The `BlockMask`, however, is not a `Tensor` and has no `requires_grad` attribute, so the call raises. Full-backward schedules never reach `stage_backward_input` (they use `stage_backward`), which is why they don't hit this. This is the same class of problem already noted in `torchtitan` for `_skip_lm_head` (PP backward calling `.requires_grad` on non-tensor forward inputs), i.e. **`stage_backward_input` should skip non-tensor stage inputs.** ## Abbreviated traceback ``` File "torchtitan/trainer.py", in forward_backward_step self.pp_schedule.step(...) File "torch/distributed/pipelining/schedules.py", in _perform_action stage.backward_one_chunk(...) File "torch/distributed/pipelining/stage.py", in backward_maybe_with_nosync result = perform_backward(backward_type)() # backward_type == "input" File "torch/distributed/pipelining/_backward.py", in stage_backward_input stage_input_grad_fns = list(map(_get_grad_fn_or_grad_acc, ...)) File "torch/distributed/pipelining/_backward.py", in _get_grad_fn_or_grad_acc if t.requires_grad and t.grad_fn is None: AttributeError: 'BlockMask' object has no attribute 'requires_grad' ``` ## Affected vs. unaffected schedules | Schedule | Splits backward (calls `stage_backward_input`)? | Status | | --- | --- | --- | | `InterleavedZeroBubble` | yes | ❌ fails | | `ZBVZeroBubble` | yes | ❌ fails | | `PipelineScheduleMulti` (custom CSV with separate I/W) | yes | ❌ fails | | `1F1B` | no | ✅ works | | `GPipe` | no | ✅ works | | `Interleaved1F1B` | no | ✅ works | ## Reproduction Any language model on the FlexAttention backend (e.g. llama3 debug model, which now defaults to flex) with a split-backward PP schedule: ```bash NGPU=4 ./run_train.sh \ --parallelism.pipeline_parallel_degree 4 \ --parallelism.pipeline_parallel_schedule InterleavedZeroBubble \ --activation_checkpoint.mode full ``` ## Impact - In `tests/integration_tests/features.py`, the `pp_looped_zero_bubble`, `pp_zbv`, and `pp_custom_csv` flavors fail. They are currently marked `disabled=True` with a TODO referencing this issue. - FlexAttention + zero-bubble/multi PP is effectively unsupported until fixed. ## Possible fixes 1. **Upstream (preferred):** make `stage_backward_input` / `_get_grad_fn_or_grad_acc` skip non-tensor stage inputs (treat them as non-differentiable), matching how full backward already tolerates them. 2. **torchtitan workaround:** stop forwarding the `BlockMask` as a chunked PP input and instead reconstruct it inside the model's `forward` from the `positions` tensor (which is already forwarded and is PP-safe). This avoids putting a non-tensor object into the PP stage inputs entirely, but is a larger change to mask flow and interacts with Context Parallel (which shards the trainer-built mask). @tianyu-l 's comment for 2: There's tricky composability issue with CP. Today we are doing - create inputs and attention masks, without CP sharding, without PP microbatching - do CP sharding of inputs and attention masks (non-trivial!) outside the model forward, before PP sees the inputs - PP sees the input and do microbatching, which includes adjusting the block masks Fix 2 would invalidate the order, thus requiring other global changes. ## Notes - `positions` (Tensor) is forwarded the same way and works fine; only the non-tensor `BlockMask` is the problem. - Previously this path was never exercised because the language-model default was SDPA (no `BlockMask`); it surfaced when flex became the default.