Fix

Lightning-AI · Mar 21, 2024 · 9ecbd38 · 9ecbd38
1 parent 9e3b2d6
commit 9ecbd38
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/thunder/tests/litgpt_model.py b/thunder/tests/litgpt_model.py
@@ -18,9 +18,9 @@
         rotary_percentage=1.0,
         parallel_residual=False,
         bias=False,
-        _norm_class="RMSNorm",
+        norm_class_str="RMSNorm",
         norm_eps=1e-6,
-        _mlp_class="LLaMAMLP",
+        mlp_class_str="LLaMAMLP",
         intermediate_size=1376,
     ),
     dict(
@@ -34,8 +34,8 @@
         rotary_percentage=1.0,
         parallel_residual=False,
         bias=False,
-        _norm_class="RMSNorm",
-        _mlp_class="LLaMAMLP",
+        norm_class_str="RMSNorm",
+        mlp_class_str="LLaMAMLP",
         intermediate_size=11008,
         rope_condense_ratio=4,
     ),
@@ -49,8 +49,8 @@
         rotary_percentage=1.0,
         parallel_residual=False,
         bias=False,
-        _norm_class="RMSNorm",
-        _mlp_class="LLaMAMLP",
+        norm_class_str="RMSNorm",
+        mlp_class_str="LLaMAMLP",
         intermediate_size=1376,
     ),
     dict(
@@ -87,9 +87,9 @@
         rotary_percentage=1.0,
         parallel_residual=False,
         bias=False,
-        _norm_class="RMSNorm",
+        norm_class_str="RMSNorm",
         norm_eps=1e-05,
-        _mlp_class="LLaMAMLP",
+        mlp_class_str="LLaMAMLP",
         intermediate_size=1376,
         rope_base=1000000,
     ),
@@ -104,9 +104,9 @@
         n_query_groups=8,
         parallel_residual=False,
         bias=False,
-        _norm_class="RMSNorm",
+        norm_class_str="RMSNorm",
         norm_eps=1e-05,
-        _mlp_class="LLaMAMoE",
+        mlp_class_str="LLaMAMoE",
         intermediate_size=224,
         rope_base=1000000,
         n_expert=8,