Weight the TiledLoss backward pass the way the forward does

[vineethsaivs]

TiledLoss returns a loss whose gradient is not the derivative of that loss.

forward takes a token weighted average over the shards and skips fully masked ones. backward runs autograd on each shard's own mean and divides by the shard count, so masked shards stay in the denominator and live ones lose their weights. Carry the per shard token counts through to backward, scale each shard's incoming gradient by its share of the labels, and drop the division by the shard count.

Test: TORCHDYNAMO_DISABLE=1 DS_ACCELERATOR=cpu PYTHONPATH=. python -m pytest tests/unit/ulysses_alst/test_tiled_loss.py -q. Three regressions fail before, the gradient off by 0.375x to 1.5x while the forward matches. Twelve CPU tests pass after, covering a fully masked shard, uneven token counts and more shards than live tokens, against token mean cross entropy in float64. Changed-file pre-commit passes. Executed on Apple M2 Pro CPU; no CUDA or distributed training run.

[vineethsaivs]

Could you approve the CI run? Three regressions fail before, the gradient off by 0.375x to 1.5x while the forward matches. Twelve CPU tests pass after, covering a fully masked shard, uneven token counts and more shards than live tokens, against token mean cross entropy in float64. Changed-file pre-commit passes.

[vineethsaivs]

Blast radius, since this touches the gradient every Ulysses SP run takes.

The even case is provably unchanged. When every shard holds the same number of live labels, good_i / total_good is 1 / shards for each of them, which is exactly the logits_grad /= shards this removes. Measured on a sequence of 8 with nothing masked, at 2 and 4 shards: the gradient is bit identical before and after, max gap 0.000e+00, and identical to plain token mean cross entropy in both.

What moves is only the case the forward already treats specially: shards with different numbers of non -100 labels, or a shard that is entirely masked. There the old backward was off by 0.375x to 1.5x on the shapes in the new test.

The forward was already correct and is untouched; it agrees with plain token mean cross entropy to 1.19e-07, which is float32 eps from the loss function's own .float() cast.

[stas00]

Thank you for the PR, @vineethsaivs

why are you trying to use TiledLoss? I just left it there in case someone was already using it - I probably should delete everything after

###########################################################
### below are older versions that some might still want ###
###########################################################

or perhaps mark those deprecated

Any reason why you're not using one of these 3 as documented here?

https://github.com/vineethsaivs/DeepSpeed/blob/666720b44a98cc79fe7e478c0101f0c9dadecb43/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L19-L21

That is:

• SequenceTiledCompute - generic autograd function to perform compute after tiling on the sequence dimension
• TiledMLP - a specific autograd function to perform tiled MLP (it's much easier to understand before trying to grok SequenceTiledCompute)
• TiledFusedLogitsLoss - a specific autograd function to perform loss computation without manifesting the full logits tensor and instead computing loss on shards of logits.

[vineethsaivs]

Honest answer: I'm not using it. I was reading through the sequence parallel loss code checking that each backward actually matches its own forward, and TiledLoss was the one that disagreed. TiledFusedLogitsLoss is fine on that count: its forward takes a plain mean and its backward pre-divides to match, and the docstring already warns callers to use the "sum" reduction when shards hold different numbers of unmasked tokens.

So if TiledLoss is legacy, deleting it is a better fix than mine and I'm happy to close this. One thing to watch if you do: UlyssesSPFwdLossBwdWithLogits sits in the same older-versions block and still calls it, so it goes along with it.

Want me to send a deprecation or removal PR instead?

[stas00]

sure, that would be very kind of you to delete them all after that line including the line, @vineethsaivs - thank you. you can re-use this PR.

I guess they were premature prototypes as I was developing this feature, so weren't well tested and you discovered proved to be buggy.

[alanhuangyoo]

Ran your tests on an H20 since the description says CPU only: same three gradient cases fail on master, all 12 pass here.

Small one, non-blocking: good_items is a device tensor, so ctx.good_items_per_shard keeps one per shard alive until backward and the weight multiply runs on device. int(...) at the point you append would make it a plain number and the masked-shard entry would stop being the odd one out type-wise.

[stas00]

@alanhuangyoo, if you read my comment above this PR is going to change to drop this class altogether - so nothing needs to be fixed.

[alanhuangyoo]

Got it, no action needed then. Deleting it is the better fix.