Skip to content

Commit 7588843

Browse files
committed
Update on "Save foundation weights separately"
This diff: 1. Introduces SerializationConfig to llm_config. Currently, this allows user to save the foundation weights in a separate file; majorly useful for lora case. 2. Adds a pass to tag foundation (non-lora) weights. This is at the top-level (export_llama_lib). The tags are preserved through run_decomps/other passes, and do not affect functionality. 3. Tags are read when placing constants into the named_data_store. 4. Tagged weights are serialized to a separate file. Notes 1. Adding tags to node.meta['custom']['blah'] means that they will not be discarded by run_decompositions 2. Adding tags to the lifted model (ep.graph_module) requires the EP to check is_param_node for xnnpack constants. Instead, add tags to the unlifted model (ep.module()), so we do not need to go through a re-export to get the EP. 3. Not an issue for this diff as llama doesn't have any higher order ops. Adding tags to models with higher-order ops is problematic due to nested submodules. Differential Revision: [D79181064](https://our.internmc.facebook.com/intern/diff/D79181064/) [ghstack-poisoned]
2 parents 7ceacde + 62f7eb5 commit 7588843

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

exir/program/_program.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1892,7 +1892,9 @@ def write_tensor_data_to_file(self, outdir) -> None:
18921892
"""
18931893
assert self._tensor_data is not None
18941894
for filename, cord in self._tensor_data.items():
1895-
with open(os.path.join(outdir, f"{filename}.ptd"), "wb") as f:
1895+
if not filename.endswith(".ptd"):
1896+
filename += ".ptd"
1897+
with open(os.path.join(outdir, f"{filename}"), "wb") as f:
18961898
logging.info(f"Writing data file to {filename}")
18971899
cord.write_to_file(f)
18981900

0 commit comments

Comments
 (0)