Repository: pytorch
Issue: [DTensor] loss_parallel not worked with shifted logits and labels
Labels: oncall: distributed

### 🐛 Describe the bug

Hi, I just found a bug of `loss_parallel` with shifted `logits`. 

Basically, I follows the tutorial in the [link](https://pytorch.org/tutorials/intermediate/TP_tutorial.html#apply-loss-parallel). I made changes by slicing `logits` and `labels` so we are predicting next label. The code is as follows:

```
class ParallelStyleTest(DTensorTestBase):

    @property
    def world_size(self) -> int:
        return 4

    @with_comms
    def test_slice_backward(self):
        device_mesh = self.build_device_mesh()
        logits = torch.rand((8, 15), device='cuda').requires_grad_()
        labels = torch.randint(0, 15, (8,), dtype=torch.long, device='cuda')
        loss = CrossEntropyLoss()(logits[:-1, :], labels[1:])
        print(f'{loss = }', flush=True)
        loss.backward()

        d_logits = DTensor.from_local(
            logits, device_mesh, (Replicate(),)
        ).redistribute(device_mesh, (Shard(-1),))
        d_labels = DTensor.from_local(labels, device_mesh, (Replicate(),))
        with loss_parallel():
            loss = CrossEntropyLoss()(d_logits[:-1, :], d_labels[1:])
            print(f'{loss = }', flush=True)
            loss.backward()


if __name__ == '__main__':
    run_tests()
```

Then an error is thrown. I found the error happens when `slice_backward` is called. But slice seems to work well without `loss_parallel`.

```
Traceback (most recent call last):
  File "/.../torch24/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 657, in run_test
    getattr(self, test_name)()
  File "/.../torch24/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 539, in wrapper
    fn()
  File "/.../torch24/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 2744, in wrapper
    method(*args, **kwargs)
  File "/.../torch24/lib/python3.10/site-packages/torch/testing/_internal/distributed/_tensor/common_dtensor.py", line 369, in wrapper
    func(self, *args, **kwargs)  # type: ignore[misc]
  File "/mnt/workspace/jiqi/Code/mtp_pai/mtp/tests/auto_parallel/test_parallel_style.py", line 42, in test_slice_backward
    loss.backward()
  File "/.../torch24/lib/python3.10/site-packages/torch/_tensor.py", line 521, in backward
    torch.autograd.backward(
  File "/.../torch24/lib/python3.10/site-packages/torch/autograd/__init__.py", line 289, in backward
    _engine_run_backward(
  File "/.../torch24/lib/python3.10/site-packages/torch/autograd/graph.py", line 768, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/mnt/workspace/jiqi/Code/mtp_pai/mtp/tests/auto_parallel/zoos/utils.py", line 14, in __torch_dispatch__
    return DTensor._op_dispatcher.dispatch(
  File "/.../torch24/lib/python3.10/site-packages/torch/distributed/_tensor/_dispatch.py", line 205, in dispatch
    local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
  File "/.../torch24/lib/python3.10/site-packages/torch/_ops.py", line 667, in __call__
    return self_._op(*args, **kwargs)
RuntimeError: The size of tensor a (15) must match the size of tensor b (3) at non-singleton dimension 1
```

### Versions

PyTorch 2.4.0

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o