Repository: pytorch
Issue: `CrossEntropyLoss` much slower with class probabilities vs. class indices target
Labels: module: nn,module: loss,triaged

### 🐛 Describe the bug

Hi, I noticed this when training ViT-S/16 on the ImageNet-1k dataset, with and without [MixUp](https://arxiv.org/abs/1710.09412). In short: torchvision's [`v2.MixUp`](https://pytorch.org/vision/main/generated/torchvision.transforms.v2.MixUp.html), just like other common implementations, converts the class indices to one-hot class probabilities before linearly interpolate them, so the same call to `CrossEntropyLoss()` goes to different kernels with vs. without MixUp. In a less optimized setting ([specific commit](https://github.com/EIFY/mup-vit/tree/51bcdc77ea6e26ceef049e96f59702a9d9ef15d1)), I found that with vs. without MixUp makes the difference between **0.198s vs. 0.120s** per step (excluding data loading) with batch size 1024 on a 8x A100-SXM4-40GB [Lambda](https://lambdalabs.com/) instance. In further iterations, I included `CrossEntropyLoss` in the model so they compile together. I then tested a workaround in which the custom MixUp returns both targets and the weight scalar lam(bda):

```python
            lam = self._dist.sample()
            images = images.roll(1, 0).mul_(1.0 - lam).add_(images, alpha=lam)
            return images, lam, labels, labels.roll(1, 0)
```

and the model calculates the cross entropy loss by using the class indices twice:

```python
        logprob = F.log_softmax(out, dim=1)
        return lam * F.nll_loss(logprob, target1) + (1.0 - lam) * F.nll_loss(logprob, target2)
```

See https://github.com/EIFY/mup-vit/compare/normal-prefetch...loss-optim. I found that this is faster (**0.184s** per step) than just calling `CrossEntropyLoss()` with class probabilities even with `torch.compile(...,mode="max-autotune-no-cudagraphs")` (**0.191s** per step), excluding data loading. Is this expected? Intuitively I can see that the sparsity helps forward & backward pass in loss calculation, but backward pass beyond that should be the same and it's hard to understand the observed second/step difference with a 12-layer transformer.

To replicate the timing result, you can run the code with fake data:
```shell
MUPVIT_MAIN=~/Downloads/mup-vit/main.py
PYTHON=torchrun
N_WORKERS=80
N_THREADS=124
EPOCH=90

NUMEXPR_MAX_THREADS=$N_THREADS $PYTHON $MUPVIT_MAIN --fake-data --workers $N_WORKERS --multiprocessing-distributed --epochs $EPOCH --batch-size 1024 --mlp-head --report-to wandb --name cross-entropy-loss-class-prob

NUMEXPR_MAX_THREADS=$N_THREADS $PYTHON $MUPVIT_MAIN --fake-data --workers $N_WORKERS --multiprocessing-distributed --epochs $EPOCH --batch-size 1024 --mlp-head --mixup-alpha 0.0 --report-to wandb --name cross-entropy-loss-class-index
```
with either the [normal-prefetch](https://github.com/EIFY/mup-vit/tree/normal-prefetch) or the [loss-optim](https://github.com/EIFY/mup-vit/tree/loss-optim) branch.

### Versions

The output below shows PyTorch version: 2.3.1 and torchvision==0.18.1, but I am actually running
```
$ pip freeze
(...)
torch==2.5.1
torchaudio==2.5.1
torchvision==0.20.1
(...)
```
I don't know why `python3 -mpip list --format=freeze` finds the old packages.

```
Collecting environment information...
PyTorch version: 2.3.1
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.8.0-47-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 550.90.12
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               124
On-line CPU(s) list:                  0-123
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7542 32-Core Processor
CPU family:                           23
Model:                                49
Thread(s) per core:                   1
Core(s) per socket:                   1
Socket(s):                            124
Stepping:                             0
BogoMIPS:                             5800.00
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr wbnoinvd arat npt nrip_save umip rdpid arch_capabilities
Virtualization:                       AMD-V
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            7.8 MiB (124 instances)
L1i cache:                            7.8 MiB (124 instances)
L2 cache:                             62 MiB (124 instances)
L3 cache:                             1.9 GiB (124 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-31,64-95
NUMA node1 CPU(s):                    32-63,96-123
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; untrained return thunk; SMT disabled
Vulnerability Spec rstack overflow:   Vulnerable: Safe RET, no microcode
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] flake8==4.0.1
[pip3] numpy==1.21.5
[pip3] optree==0.12.1
[pip3] torch==2.3.1
[pip3] torchvision==0.18.1
[pip3] triton==2.3.1
[conda] Could not collect
```

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki