Reduction Op Softmax, Cross Entropy Loss, and LogSoftmax Test Failed While Parsing JITFunc in Triton

[DuanYaQi]

Issue

Reduction Op Softmax, Cross Entropy Loss, and LogSoftmax Test Failed While Parsing JITFunc in Triton

Env

FlagGems Commit: 3c62c9c
Python 3.8.18
Torch: 2.3.1
Triton: 2.3.1
Pytest: 8.2.2

Command

python -m pytest -svvv tests/test_reduction_ops.py::test_accuracy_softmax'[0-dtype0-shape0]'
python -m pytest -svvv tests/test_reduction_ops.py::test_accuracy_cross_entropy_loss'[1-dtype0-shape0-mean-None-None]'
python -m pytest -svvv tests/test_reduction_ops.py::test_accuracy_log_softmax'[dtype0-shape0]'

Softmax Result

BT

[iclementine]

This is a bug with python's builtin module inspect, which is used in triton's frontend to parse the source of triton jit function. In short, inspect fails to parse function correctly when some form of lambda is in the function decorator.

TL;DR upgrading python interpreter to 3.10.10+ can solve this.

Also, as a workaround, maybe you can try modifying the lambda in the decorator. We will try it, too.

Related issues
python/cpython#65416
triton-lang/triton#1589
python/cpython#102647

The solution in cpython interpreter
python/cpython#99654

[DuanYaQi]

This is a bug with python's builtin module inspect, which is used in triton's frontend to parse the source of triton jit function. In short, inspect fails to parse function correctly when some form of lambda is in the function decorator.

TL;DR upgrading python interpreter to 3.10.2+ can solve this.

Also, as a workaround, maybe you can try modifying the lambda in the decorator. We will try it, too.

Related issues python/cpython#65416 triton-lang/triton#1589 python/cpython#102647

thanks for your reply. It looks like it can only handle lambda functions with simple logic.