From 81a87cc262c08c7dedbf8abe3caf27889b57709a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 Nov 2024 00:09:11 -0800 Subject: [PATCH] types --- unsloth/kernels/rms_layernorm.py | 2 +- unsloth/kernels/rope_embedding.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 4b22f8c3e5..10b435dd57 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -191,7 +191,7 @@ pass # [TODO] Unsure why RMS Layernorm is not torch.compiling properly -@torch.compiler.disable +# @torch.compiler.disable def fast_rms_layernorm(layernorm, X : torch.Tensor, gemma : bool = False): W : torch.Tensor = layernorm.weight eps : float = layernorm.variance_epsilon if \ diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index a8173d2453..246055dc7b 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -16,7 +16,7 @@ import triton import triton.language as tl import torch from .utils import calculate_settings -ROPE_GROUP_SIZE = 4 +ROPE_GROUP_SIZE : int = 4 @triton.heuristics({"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),}) @triton.jit