constexpr
This commit is contained in:
parent
006bafb437
commit
d2725e7910
2 changed files with 9 additions and 9 deletions
|
|
@ -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,
|
||||
):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue