minor fix on llama checkpoint conversion

PiperOrigin-RevId: 572349696 Change-Id: Icda9d683e5b181b4d0bdd4e06eb39674729042be
google · Oct 10, 2023 · 79aae7d · 79aae7d
1 parent 291d708
commit 79aae7d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/saxml/tools/convert_llama_ckpt.py b/saxml/tools/convert_llama_ckpt.py
@@ -94,7 +94,7 @@ def convert(base_model_path, pax_model_path, model_size):
                   }
               },
           'final_ln': {
-              'scale': pytorch_vars[0]['norm.weight'].numpy()
+              'scale': pytorch_vars[0]['norm.weight'].type(torch.float16).numpy()
               },
           'transformer': {}
           }
@@ -165,7 +165,7 @@ def convert(base_model_path, pax_model_path, model_size):
 
   print(f'Saving the pax model to {pax_model_path}')
   jax_states = train_states.TrainState(
-      step=np.zeros(1),
+      step=0,
       mdl_vars={'params': jax_weights},
       opt_states={})