[Bugfix] use float32 precision in samplers/test_logprobs.py for comparing with HF (#6409)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
This commit is contained in:
Thomas Parnell
2024-07-15 19:14:49 +02:00
committed by GitHub
parent eaec4b9153
commit 4ef95b0f06
2 changed files with 8 additions and 1 deletions

View File

@@ -687,6 +687,12 @@ if triton.__version__ >= "2.1.0":
cap = current_platform.get_device_capability()
BLOCK = 128 if cap[0] >= 8 else 64
# need to reduce num. blocks when using fp32
# due to increased use of GPU shared memory
if q.dtype is torch.float32:
BLOCK = BLOCK // 2
# shape constraints
Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
assert Lq == Lk and Lk == Lv