sclarkson 1feb711f46 Fix compilation with clang on ARM64 (#1285) | 1 hafta önce | |
---|---|---|
.. | ||
README.md | 1 yıl önce | |
interface.cpp | 2 yıl önce | |
setup.py | 9 ay önce | |
xentropy_kernel.cu | 1 hafta önce |
This CUDA extension implements optimized cross-entropy loss, adapted from Apex's Xentropy. We make it work for bfloat16 and support in-place backward to save memory.
It has only been tested on A100s.
cd csrc/xentropy && pip install .
As of 2023-09-15, this extension is no longer used in the FlashAttention repo. We've instead switched to a Triton-based implementation. See the CrossEntropyLoss module for more details.