always precompute freqs_cis

pytorch · Oct 30, 2024 · 3c23bdb · 3c23bdb
1 parent 5cd6d15
commit 3c23bdb
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/scripts/convert_llama_to_dcp.py b/scripts/convert_llama_to_dcp.py
@@ -124,12 +124,12 @@ def convert_llama_weights(input_dir, output_dir, max_seq_len: int):
         for i in range(len(shards)):
             del shards[i]["output.weight"]
 
-        # NOTE: precompute freqs_cis because must be persisted by default in torchtitan
-        state_dict["freqs_cis"] = precompute_freqs_cis(
-            dims_per_head,
-            max_seq_len,
-            params.get("rope_theta", 500000),
-        )
+    # NOTE: precompute freqs_cis because must be persisted by default in torchtitan
+    state_dict["freqs_cis"] = precompute_freqs_cis(
+        dims_per_head,
+        max_seq_len,
+        params.get("rope_theta", 500000),
+    )
 
     logger.info(f"Writing to DCP at '{output_dir}'")
     output_dir.mkdir(parents=True, exist_ok=True)