constexpr

This commit is contained in:
Daniel Han 2024-11-04 00:03:24 -08:00
commit d2725e7910
2 changed files with 9 additions and 9 deletions

View file

@ -38,9 +38,9 @@ def _cross_entropy_forward(
labels_ptr ,
VOCAB_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP : tl.constexpr,
DO_LOGIT_SCALING : tl.constexpr,
DO_LOGIT_SCALING ,
LOGIT_SCALE : tl.constexpr,
):
"""
@ -112,9 +112,9 @@ def _chunked_cross_entropy_forward(
VOCAB_SIZE : tl.constexpr,
N_CHUNKS : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP : tl.constexpr,
DO_LOGIT_SCALING : tl.constexpr,
DO_LOGIT_SCALING ,
LOGIT_SCALE : tl.constexpr,
):
"""
@ -195,9 +195,9 @@ def _cross_entropy_backward(
labels_ptr ,
VOCAB_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP : tl.constexpr,
DO_LOGIT_SCALING : tl.constexpr,
DO_LOGIT_SCALING ,
LOGIT_SCALE : tl.constexpr,
):
"""

View file

@ -366,8 +366,8 @@ pass
# =============================================
# Torch compile settings
UNSLOTH_COMPILE_DEBUG = True
UNSLOTH_COMPILE_DEBUG = "UNSLOTH_COMPILE_DEBUG" in os.environ
UNSLOTH_COMPILE_MAXIMUM = "UNSLOTH_COMPILE_MAXIMUM" in os.environ
# Just remove max_autotune_gemm warning
import functools
@functools.lru_cache(None)
@ -379,7 +379,7 @@ def is_big_gpu(index):
return True
import torch._inductor.utils
torch._inductor.utils.is_big_gpu = is_big_gpu
patch_torch_compile(debug = UNSLOTH_COMPILE_DEBUG, O3 = False)
patch_torch_compile(debug = UNSLOTH_COMPILE_DEBUG, O3 = UNSLOTH_COMPILE_MAXIMUM)
torch_compile_options = {
"epilogue_fusion" : True,