linkedin / linkedin/Liger-Kernel
cross entropy kernel: inplace modification of input
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
Hello,
Regarding inplace modification of PyTorch tensors, there are already multiple (#254, #262, #272) issues. I would also like to point out that according to PyTorch docs for [mark_dirty()](https://pytorch.org/docs/stable/generated/torch.autograd.function.FunctionCtx.mark_dirty.html#torch-autograd-function-functionctx-mark-dirty):
```
Every tensor that’s been modified in-place in a call to forward() should be given to this function, to ensure correctness of our checks
```
So can you please look into it and come up with a resolution? :-)
Now I noticed another thing.
Since currently gradient is being calculated during forward pass, during backward pass, if grad of loss w.r.t. cross entropy output is 1.0, no calculation is done and only the grad calculated during the forward pass is returned. However, just doing that seems to make backward pass take almost the same time as the forward pass. I do not understand why that would be so.
I profiled both forward and backward passes using Triton's proton profiler with the following script:
```
import torch
# from torch.nn import CrossEntropyLoss as TorchCrossEntropyLoss
from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss
import triton.profiler as proton
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
elif torch.xpu.is_available():
device = 'xpu'
else:
print('No accelerator available, running on CPU')
def run():
# Create Tensors to hold input and outputs.
with proton.scope("0-init"):
V = 131072
B = 8
T = 2048
_input = torch.randn(B * T, V, requires_grad=True, device=device)
target = torch.randint(V, (B * T, 1), device=device).squeeze(1)
# layer = TorchCrossEntropyLoss()
layer = LigerCrossEntropyLoss()
with proton.scope("1-warmup"):
for _ in range(10):
y = layer(_input, target)
y.backward(retain_graph=True)
for _ in range(100):
# Forward pass:
with proton.scope("2-forward"):
y = layer(_input, target)
# Backward pass:
with proton.scope("3-backward"):
y.backward(retain_graph=True)
func = proton.profile(run, name="cross_entropy", context='shadow')
func()
# Write out the profile
# Visualize using `proton-viewer -m time/ms ./cross_entropy.hatchet`
proton.finalize()
```
Below is the time breakdown among different kernels:
```
3361.863 ROOT
├─ 8.295 0-init
│ ├─ 8.292 _ZN2at6native54_GLOBAL__N__d8ceb000_21_DistributionNormal_cu_0c5b6e8543distribution_elementwise_grid_stride_kernelIfLi4EZNS0_9templates4cuda20normal_and_transformIffPNS_17CUDAGeneratorImplEZZZNS4_13normal_kernelIS7_EEvRKNS_10TensorBaseEddT_ENKUlvE_clEvENKUlvE0_clEvEUlfE_EEvRNS_18TensorIteratorBaseET1_T2_EUlP24curandStatePhilox4_32_10E0_ZNS1_27distribution_nullary_kernelIff6float4S7_SM_SF_EEvSH_SJ_RKT3_T4_EUlifE_EEviNS_15PhiloxCudaStateESI_SJ_
│ └─ 0.003 _ZN2at6native60_GLOBAL__N__2b2da9da_27_DistributionRandomKernel_cu_f88cee4443distribution_elementwise_grid_stride_kernelIjLi4EZZZNS0_9templates4cuda21random_from_to_kernelIPNS_17CUDAGeneratorImplEEEvRNS_18TensorIteratorBaseEmlT_ENKUlvE_clEvENKUlvE2_clEvEUlP24curandStatePhilox4_32_10E0_ZNS1_27distribution_nullary_kernelIlj5uint4S7_SF_ZZZS5_IS7_EvS9_mlSA_ENKSB_clEvENKSC_clEvEUljE_EEvS9_T2_RKT3_T4_EUlijE_EEviNS_15PhiloxCudaStateET1_SJ_
├─ 291.116 1-warmup
│ ├─ 0.087 _ZN2at6native13reduce_kernelILi512ELi1ENS0_8ReduceOpIfNS0_14func_wrapper_tIfZNS0_11sum_functorIfffEclERNS_14TensorIteratorEEUlffE_EEjfLi4EEEEEvT1_
│ ├─ 0.058 _ZN2at6native13reduce_kernelILi512ELi1ENS0_8ReduceOpIlNS0_14func_wrapper_tIlZNS0_11sum_functorIlllEclERNS_14TensorIteratorEEUlllE_EEjlLi4EEEEEvT1_
│ ├─ 0.041 _ZN2at6native27unrolled_elementwise_kernelIZZZNS0_23direct_copy_kernel_cudaERNS_18TensorIteratorBaseEENKUlvE1_clEvENKUlvE2_clEvEUllE_NS_6detail5ArrayIPcLi2EEE23TrivialOffsetCalculatorILi1EjESC_NS0_6memory12LoadWithCastILi1EEENSD_13StoreWithCastILi1EEEEEviT_T0_T1_T2_T3_T4_
│ ├─ 0.039 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_11FillFunctorIfEENS_6detail5ArrayIPcLi1EEEEEviT0_T1_
│ ├─ 0.027 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_13AUnaryFunctorIllbNS0_51_GLOBAL__N__28ce311f_18_CompareEQKernel_cu_d8008c9616CompareEqFunctorIlEEEENS_6detail5ArrayIPcLi2EEEEEviT0_T1_
│ ├─ 0.015 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_13BinaryFunctorIffbNS0_51_GLOBAL__N__28ce311f_18_CompareEQKernel_cu_d8008c9616CompareEqFunctorIfEEEENS_6detail5ArrayIPcLi3EEEEEviT0_T1_
│ ├─ 136.551 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_15CUDAFunctor_addIfEENS_6detail5ArrayIPcLi3EEEEEviT0_T1_
│ └─ 154.297 liger_cross_entropy_kernel
├─ 1545.133 2-forward
│ ├─ 0.878 _ZN2at6native13reduce_kernelILi512ELi1ENS0_8ReduceOpIfNS0_14func_wrapper_tIfZNS0_11sum_functorIfffEclERNS_14TensorIteratorEEUlffE_EEjfLi4EEEEEvT1_
│ ├─ 0.569 _ZN2at6native13reduce_kernelILi512ELi1ENS0_8ReduceOpIlNS0_14func_wrapper_tIlZNS0_11sum_functorIlllEclERNS_14TensorIteratorEEUlllE_EEjlLi4EEEEEvT1_
│ ├─ 0.401 _ZN2at6native27unrolled_elementwise_kernelIZZZNS0_23direct_copy_kernel_cudaERNS_18TensorIteratorBaseEENKUlvE1_clEvENKUlvE2_clEvEUllE_NS_6detail5ArrayIPcLi2EEE23TrivialOffsetCalculatorILi1EjESC_NS0_6memory12LoadWithCastILi1EEENSD_13StoreWithCastILi1EEEEEviT_T0_T1_T2_T3_T4_
│ ├─ 0.207 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_11FillFunctorIfEENS_6detail5ArrayIPcLi1EEEEEviT0_T1_
│ ├─ 0.275 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_13AUnaryFunctorIllbNS0_51_GLOBAL__N__28ce311f_18_CompareEQKernel_cu_d8008c9616CompareEqFunctorIlEEEENS_6detail5ArrayIPcLi2EEEEEviT0_T1_
│ └─ 1542.804 liger_cross_entropy_kernel
└─ 1517.319 3-backward
├─ 0.171 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_11FillFunctorIfEENS_6detail5ArrayIPcLi1EEEEEviT0_T1_
├─ 0.148 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_13BinaryFunctorIffbNS0_51_GLOBAL__N__28ce311f_18_CompareEQKernel_cu_d8008c9616CompareEqFunctorIfEEEENS_6detail5ArrayIPcLi3EEEEEviT0_T1_
└─ 1517.000 _ZN2at6native29vectorized_elementwise_kernelILi4ENS0_15CUDAFunctor_addIfEENS_6detail5ArrayIPcLi3EEEEEviT0_T1_
```
Notice that during backward, almost 100% time is spent in the last kernel, namely `_ZN2at6native29vectorized_elementwise_kernelILi4ENS0_15CUDAFunctor_addIfEENS_6detail5ArrayIPcLi3EEEEEviT0_T1_`. I have checked this kernel takes similar amount of time in the PyTorch/huggingface version of the calculation.
So are we gaining much from the inplace calculation? It already seems illegal and there is hardly any backward pass speedup even with it
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the LigerCrossEntropyLoss entry point in liger_kernel.transformers.cross_entropy and reproduce the supplied Triton profiler script. Check the in-place tensor behavior against PyTorch's mark_dirty() documentation, then compare the forward and backward kernel timings; done means the correctness concern and the reported backward-pass cost have a documented resolution.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100