Reshape instead of view in TiledFusedLogitsLoss

[alanhuangyoo]

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:

# 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

[ebarkhordar]

The new test compares the two losses but never the gradient, and the flatten you changed feeds the gradient path as well: x_grad is a zeros_like of the flattened x, it is scattered per shard through x_grad.narrow(0, shard_offset, shard_step).view_as(x_shard) at ulysses_sp.py:1190, and unflattened at :1208. TestTiledMLPInputLayout, which your docstring names as the same caller contract, does assert x_tiled.grad against the reference.

The test also never calls backward(), so TiledFusedLogitsLoss.backward and that saved x_grad are not exercised at all.

Adding both passes at 307ff15, in a clean python:3.12-slim container with torch 2.14.0+cpu:

shards=2  loss_equal=True  x_grad_equal=True  max|dx|=0.000e+00
shards=4  loss_equal=True  x_grad_equal=True  max|dx|=0.000e+00

Concretely that is loss.backward() inside the loop, then comparing strided.grad with contiguous.grad after it. The parameter gradients match too, if you want the third assertion the MLP test makes.

[alanhuangyoo]

Both added in ac714ff — loss.backward() in the loop, strided.grad against contiguous.grad, and the parameter gradients as the third comparison, matching TestTiledMLPInputLayout.

Same numbers you measured, on torch 2.13.0:

shards=2  loss_equal=True  x_grad_equal=True  max|dx|=0.000e+00  max|dW|=0.000e+00
shards=4  loss_equal=True  x_grad_equal=True  max|dx|=0.000e+00  max|dW|=0.000e+00

One thing I checked while doing this, since it changes what the assertion is worth: reverting the reshape and running the strengthened test fails at the forward flatten (ulysses_sp.py:1150, view size is not compatible ...) before backward is ever reached. So the gradient assertion is not what catches this particular bug — the forward already did.

What it does cover is the part your comment named that the forward does not. reshape on a non-contiguous input returns a copy, so x_grad is a zeros_like of that copy rather than of the tensor the caller handed in, and the shard scatter at :1190 and the unflatten at :1208 have to put those values back on the original layout. Nothing was checking that they did. max|dx| = 0.000e+00 against the contiguous reference is what says the copy does not scramble it, and grad.abs().sum() > 0 says it is not passing by both sides being empty.

Thanks — the test was asserting the shallower half of its own docstring.

[ebarkhordar]

Good catch on the revert check, and you are right that it changes what the assertion is worth. The forward flatten fails first, so the gradient comparison is not the guard against this particular regression and I should not have implied it was.

What it does pin is the part nothing was covering: the scatter at :1190 and the unflatten at :1208 putting values back on the caller's layout after reshape hands back a copy. grad.abs().sum() > 0 alongside it is the right guard against both sides being empty.

[alanhuangyoo]

@stas00 — this is the sibling of your #8348 in the same file, and it has had no reviewer for a week, so flagging rather than letting it sit.

#8348 moved TiledMLP.backward off view because the batch+sequence flatten needs a copy when the caller hands in a non-contiguous activation. TiledFusedLogitsLoss.forward does the same flatten under the same comment and was left on view, so it raises on exactly the inputs #8348 was written for.

If you would rather fold it into something else you have in flight, that is fine by me — I mostly want to make sure it did not get skipped for looking like a duplicate of #8348.

[stas00]

Thank you for covering this class as well and the heads up, @alanhuangyoo!