From 616e97ef165215527aa448b08907373982269fef Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Sep 2024 17:03:19 -0700 Subject: [PATCH] Update layernorm.py --- unsloth/kernels/layernorm.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py index 7939613804..c0ff4d7440 100644 --- a/unsloth/kernels/layernorm.py +++ b/unsloth/kernels/layernorm.py @@ -188,16 +188,18 @@ pass def testing_suite_layernorm(): for dim in [512, 1024, 2048]: for dtype in [torch.float16, torch.bfloat16]: - for seqlen in [3341, 2048, 349]: - for random_state in [3407, 42]: - test_layernorm( - dim = dim, - eps = 1e-5, - dtype = dtype, - bsz = 21, - random_state = random_state, - seqlen = seqlen, - ) + with torch.autocast(device_type = "cuda", dtype = dtype): + for seqlen in [3341, 2048, 349]: + for random_state in [3407, 42]: + test_layernorm( + dim = dim, + eps = 1e-5, + dtype = dtype, + bsz = 21, + random_state = random_state, + seqlen = seqlen, + ) + pass pass pass pass