Reshape instead of view in TiledFusedLogitsLoss (#8362)
Follow-up to #8348, in the same file.
That PR moved `TiledMLP.backward` off `view` because the flatten of
batch and sequence into one axis needs a copy when a caller hands in a
non-contiguous activation. `TiledFusedLogitsLoss.forward` does the same
flatten, under the same comment, and is still on `view`:
```python
# flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1
x = x.view(-1, *x.shape[2:])
y = y.view(-1, *y.shape[2:])
if mask is not None:
mask = mask.view(-1)
```
A transposed activation — one of the two layouts #8348's test covers —
fails there:
```
RuntimeError: view size is not compatible with input tensor's size and stride
(at least one dimension spans across two contiguous subspaces)
```
### What is and is not reachable
I ran the four input shapes through `TiledFusedLogitsLoss.apply` on
master:
| input | contiguous | master |
| --- | --- | --- |
| `x` contiguous | yes | ok |
| **`x` transposed** | no | **RuntimeError: view size is not compatible
…** |
| `x` channel slice | no | ok |
| `y` / `mask` strided | no | ok |
Only the transposed case is reachable today. A hidden-dimension slice
keeps the wider row stride, and that still admits merging batch into
sequence, so it survives — same for the strided `y` and `mask` I tried.
`y` and `mask` go through the same flatten and are changed with it
rather than left on a spelling that happens to hold.
The unflatten at the end of `forward` stays a `view`: `x_grad` comes
from `zeros_like()` of the already-flattened `x`, so it is contiguous by
construction. Same reasoning for the `x_grad.view(x_shape_orig)` that
#8348 left alone in `TiledMLP.backward`.
`TiledLoss` has no such flatten, so nothing to do there.
### Test
`TestTiledFusedLogitsLossInputLayout` mirrors `TestTiledMLPInputLayout`
and checks the loss against the same input made contiguous, so it pins
the value and not just the absence of a throw.
On master:
```
FAILED tests/unit/ulysses_alst/test_tiled_compute.py::TestTiledFusedLogitsLossInputLayout::test_transposed_input_matches_a_contiguous_copy[2]
FAILED ...[4]
RuntimeError: view size is not compatible with input tensor's size and stride
2 failed
```
With this PR:
```
tests/unit/ulysses_alst/test_tiled_compute.py 14 passed
yapf --diff / flake8 clean
```
---------
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
