From a7ae1fb3cccd44960ad3fec3cef8a5b8cfe87e21 Mon Sep 17 00:00:00 2001 From: leslie-fang-intel Date: Thu, 13 Jun 2024 10:50:36 -0700 Subject: [PATCH] Change the llama max_batch_size larger than default eval batch size (#2283) Summary: Fix the PyTorch issue: https://github.com/pytorch/pytorch/issues/106110 for llama dynamic batch test case. The root-cause and discussion is as described in https://github.com/pytorch/pytorch/issues/106110#issuecomment-1950964863. Pull Request resolved: https://github.com/pytorch/benchmark/pull/2283 Reviewed By: aaronenyeshi Differential Revision: D58532244 Pulled By: xuzhao9 fbshipit-source-id: ef2a51974c77e8dcf7e041db10e71869768126a3 --- torchbenchmark/models/llama/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchbenchmark/models/llama/model.py b/torchbenchmark/models/llama/model.py index a01f4cae6a..760179e020 100644 --- a/torchbenchmark/models/llama/model.py +++ b/torchbenchmark/models/llama/model.py @@ -19,7 +19,7 @@ class ModelArgs: multiple_of: int = 256 # make SwiGLU hidden layer size multiple of large power of 2 norm_eps: float = 1e-5 - max_batch_size: int = 32 # From the paper they use a batch size of 4M for training + max_batch_size: int = 64 # From the paper they use a batch size of 4M for training max_seq_len: int = 1024 device: Optional[str] = None