diff --git a/pyproject.toml b/pyproject.toml index 7d955082d7..05c4191974 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ huggingface = [ "psutil", "wheel>=0.42.0", "numpy", + "triton", ] cu118only = [ "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.22.post7%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'", diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index cae9d30b21..8ad8bf20be 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -52,8 +52,12 @@ def _rms_layernorm_forward( normed = normed.to(W_row.dtype) # Exact copy from HF # For Gemma - cannot do += 1 since float16 - maybe use FMADD - if not ADD_ONE: output = normed * W_row - else: output = normed * (W_row + 1.0) + if not ADD_ONE: + output = normed * W_row + else: + # Error analysis shows we need to do +1 in float32 then downcast to float16 + output = normed * (W_row.to(tl.float32) + 1.0).to(W_row.dtype) + pass tl.store(Y + col_offsets, output, mask = mask) pass @@ -93,8 +97,11 @@ def _rms_layernorm_backward( normed = X_row * inv_var # For Gemma - cannot do += 1 since float16 - maybe use FMADD - if not ADD_ONE: dY_W = dY_row * W_row - else: dY_W = dY_row * (W_row + 1.0) + if not ADD_ONE: + dY_W = dY_row * W_row + else: + dY_W = dY_row * (W_row + 1.0) + pass rowsum_dY_normed = tl.sum(dY_W * normed, axis = 0) output = inv_var/n_cols * (n_cols*dY_W - normed*rowsum_dY_normed)