constexpr

This commit is contained in:
Daniel Han 2024-11-04 00:05:51 -08:00
commit 7d5111a40c
2 changed files with 7 additions and 7 deletions

View file

@ -39,9 +39,9 @@ def _cross_entropy_forward(
VOCAB_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP : tl.constexpr,
SOFTCAP ,
DO_LOGIT_SCALING ,
LOGIT_SCALE : tl.constexpr,
LOGIT_SCALE ,
):
"""
Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
@ -113,9 +113,9 @@ def _chunked_cross_entropy_forward(
N_CHUNKS : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP : tl.constexpr,
SOFTCAP ,
DO_LOGIT_SCALING ,
LOGIT_SCALE : tl.constexpr,
LOGIT_SCALE ,
):
"""
256K vocab divided in 4 chunks
@ -196,9 +196,9 @@ def _cross_entropy_backward(
VOCAB_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP : tl.constexpr,
SOFTCAP ,
DO_LOGIT_SCALING ,
LOGIT_SCALE : tl.constexpr,
LOGIT_SCALE ,
):
"""
CE_i = -y log(P) = y * (log[sum(exp(x))] - x)

View file

@ -366,7 +366,7 @@ pass
# =============================================
# Torch compile settings
UNSLOTH_COMPILE_DEBUG = "UNSLOTH_COMPILE_DEBUG" in os.environ
UNSLOTH_COMPILE_DEBUG = True #"UNSLOTH_COMPILE_DEBUG" in os.environ
UNSLOTH_COMPILE_MAXIMUM = "UNSLOTH_COMPILE_MAXIMUM" in os.environ
# Just remove max_autotune_gemm warning
import functools