Skip to content

Commit 41c78df

Browse files
authored
[NPU]Fix OOM in benchmark_tvd on NPU (#1103)
## Summary `benchmark_tvd.py` can fail with an OOM error on Atlas 800I A2 (32GB) because of exceeding the available NPU global memory.This change addresses the issue and prevents the benchmark from exceeding device memory capacity. ## Testing Done Atlas 800I A2(32GB) - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence
1 parent 7644a0f commit 41c78df

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

benchmark/scripts/benchmark_tvd.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,10 @@ def full():
114114
# We know that the full test will require 66GBs for vocab size 2^17
115115
if gpu_memory_gbs >= 66:
116116
x_max = 17
117-
else:
117+
elif gpu_memory_gbs >= 32:
118118
x_max = 16
119+
else:
120+
x_max = 15
119121
common_args = {
120122
"kernel_name": "tvd",
121123
"x_name": "V",

0 commit comments

Comments
 (0)