frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[AOTInductor] Option to not include weight in .so (#141997)

Browse Source 
Summary: Add an option in config to not include weights in .so

Test Plan: `test/inductor:test_aot_inductor -- -r test_so_without_weight_cuda`

Reviewed By: desertfire

Differential Revision: D65968885

Pull Request resolved: https://github.com/pytorch/pytorch/pull/141997
Approved by: https://github.com/desertfire
This commit is contained in:
Mu-Chu Lee
2024-12-05 03:35:54 +00:00
committed by PyTorch MergeBot
parent 51cbac4e6a
commit b08bc07cd7
6 changed files with 105 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
torch/_inductor/codecache.py
View File

@ -1637,11 +1637,15 @@ class AotCodeCompiler:
raw_bytes = bytes(raw_array.contents)
return raw_bytes if all_cuda else _pad_to_alignment(raw_bytes)
serialized_weights = b"".join(
_to_bytes(graph.get_original_value_of_constant(name), all_cuda)
for name in graph.constants.keys()
if name not in graph.folded_constants
)
if config.aot_inductor.package_constants_in_so:
serialized_weights = b"".join(
_to_bytes(graph.get_original_value_of_constant(name), all_cuda)
for name in graph.constants.keys()
if name not in graph.folded_constants
)
else:
serialized_weights = b""
consts_size = len(serialized_weights)
# TODO: Fix mmap weights with cuda
@ -1711,6 +1715,7 @@ class AotCodeCompiler:
aot_mode=graph.aot_mode,
use_absolute_path=use_absolute_path,
)
so_builder = CppBuilder(
name=output_name,
sources=[output_o, consts_o, kernels_o],