From 371bf8ad88e52c6d9de035ca12f105bb069db832 Mon Sep 17 00:00:00 2001
From: Jack Zhang <32371937+jackzhxng@users.noreply.github.com>
Date: Tue, 27 May 2025 20:02:10 -0700
Subject: [PATCH] Reduce memory requirement on export_llama tests with no
 params

[ghstack-poisoned]
---
 examples/models/llama/model_args.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/models/llama/model_args.py b/examples/models/llama/model_args.py
index 94dbb5a0651..5734cd66ef7 100644
--- a/examples/models/llama/model_args.py
+++ b/examples/models/llama/model_args.py
@@ -5,7 +5,7 @@
 @dataclass
 class ModelArgs:
     dim: int = 4096
-    n_layers: int = 8
+    n_layers: int = 1
     n_heads: int = 32
     n_kv_heads: Optional[int] = None
     vocab_size: int = 512  # Arbitrary value, should be defined later by tokenizer.