Update fuji.py

apple · Nov 6, 2024 · a0142c1 · a0142c1
1 parent 27e92a3
commit a0142c1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py
@@ -300,7 +300,8 @@ def get_trainer_kwargs(
         trainer_kwargs = dict(
             model_kwargs=dict(
                 num_layers=126,
-                hidden_dim=53248,
+                hidden_dim=16384,
+                ffn_dim=53248,
                 num_heads=128,
                 # No GQA support in V1 models, so num_kv_heads is the same as num_heads.
                 num_kv_heads=None if version == Version.V1 else 8,