Split from #8457.
## Problem
`prepare_autosp_inputs()` takes a `seq_dim` argument, but only passes it
to `mark_dynamic`. The AutoSP graph passes ignore it:
`pass_shard_seq_dim` and `shard_tensor_node` always read the sequence
length from `shape[1]` and slice dimension 1.
Sequence-first `[seq, batch]` inputs (`seq_dim=0`) are therefore
mishandled:
- If the batch dimension is static, compilation fails with `expected
sequence dimension to be of type SymInt but found int`.
- If the batch dimension is dynamic, AutoSP shards the batch instead of
the sequence, and attention runs over the wrong tokens.
## Fix
Dynamo copies only the `tag` attribute of an input tensor into graph
metadata (`_extract_tensor_dict`), so `seq_dim` has to travel in the
tag:
- `prepare_autosp_inputs()` tags inputs as `(key, seq_dim)`.
- `find_node_by_tag()` accepts both the tuple and a bare key. A bare
key, as set manually in existing code and tests, keeps the previous
default of dimension 1.
- New `get_autosp_seq_dim()` reads the recorded dimension.
`pass_shard_seq_dim` and `shard_tensor_node` use it to find the sequence
length, and `shard_tensor_node` builds its slice for the tensor's full
rank.
- `label_id` is now bounds-checked against `seq_dim`, like the other
prepared inputs, since it is also sharded along that dimension.
- Docs: explain that `seq_dim` applies to every prepared tensor, and
that `seq_dim=0` is for sequence-first inputs.
## Tests
`TestAutoSPSeqDim` (2 ranks, `sp_size=2`) runs a small causal SDPA model
with the sequence on `seq_dim` ∈ {0, 1}. For each of three batches, the
mean of the per-rank AutoSP losses must match a single-device run on the
full sequence (`rtol=1e-4`). Causal attention makes the loss depend on
which dimension is sharded, so slicing the batch instead of the sequence
would not match.
- `seq_dim=0` fails on master and passes with this PR.
- `seq_dim=1` passes on both, which covers the default path with the new
tuple tags.
The test compares forward losses only. AutoSP currently sums gradients
over the SP group, while each rank's loss is a local mean, so gradients
are `sp_size`× a full-sequence reference. That is part of the
causal-loss aggregation follow-up from #8457, not this change.
Results on 4× NVIDIA RTX 5090 D, torch 2.13.0, transformers 5.17.0:
- `tests/unit/v1/compile/test_compile_autosp.py -k "not
TestAutoSPCompile"`: 38 passed.
- `TestAutoSPCompile` fails identically on unmodified master in this
environment. The error is a Dynamo `InternalTorchDynamoError` in the
Ulysses baseline's `torch.compile` forward, before AutoSP runs, so I
excluded it.
- `pre-commit` passes on the changed files.
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: pengdurice <pengduhit@gmail.com>
