From b60c138131abc5947835d9d68f66d5b6ff58cb57 Mon Sep 17 00:00:00 2001 From: cm2435 Date: Mon, 22 Jan 2024 20:34:12 +0000 Subject: [PATCH] formatting and typo fix --- unsloth/kernels/layernorm.py | 3 +++ unsloth/models/llama.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py index 6c9c7ce22d..4d83b607ce 100644 --- a/unsloth/kernels/layernorm.py +++ b/unsloth/kernels/layernorm.py @@ -50,6 +50,7 @@ def _layer_norm_fwd_fused( a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) _mean += a mean = tl.sum(_mean, axis=0) / N + # Compute variance _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for off in range(0, N, BLOCK_SIZE): @@ -59,9 +60,11 @@ def _layer_norm_fwd_fused( _var += x * x var = tl.sum(_var, axis=0) / N rstd = 1 / tl.sqrt(var + eps) + # Write mean / rstd tl.store(Mean + row, mean) tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation for off in range(0, N, BLOCK_SIZE): cols = off + tl.arange(0, BLOCK_SIZE) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 309625d962..41ff844481 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -785,7 +785,7 @@ class FastLlamaModel: layers = model.model.layers # Torch.compile fails on embedding matrix?? - # Workaround randomnly fixes it for torch versions < 2.2 + # Workaround randomly fixes it for torch versions < 2.2 model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight) model.config.update({"unsloth_version" : __version__})