Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Free params after converting in param loader #112

Merged
merged 2 commits into from
Dec 13, 2023
Merged

Free params after converting in param loader #112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2e40ced
Free params after converting in param loader
vinx13 Dec 13, 2023
9f9b69b
check mixtral
vinx13 Dec 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion mlc_llm/relax_model/llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1735,7 +1735,8 @@ def get_model(args, hf_config):
if args.max_seq_len != -1:
config.max_sequence_length = args.max_seq_len

param_manager = ParamManager()
keep_params_after_load = isinstance(config, MixtralConfig) and args.quantization.name == "q4f16_ft"
param_manager = ParamManager(keep_params_after_load)
bb = relax.BlockBuilder()

bb.add_func(get_scatter_func(dtype), "scatter")
Expand Down
10 changes: 7 additions & 3 deletions mlc_llm/relax_model/param_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,7 @@ class ParamManager:
pidx2pname: Dict[int, str]
torch_pname2binname: Dict[str, str]

def __init__(self) -> None:
def __init__(self, keep_params_after_load=False) -> None:
self.params = {}
self.param_names = []
self.params_in_func = {}
Expand All @@ -219,6 +219,9 @@ def __init__(self) -> None:

self.qspec_updater_classes = []

# this is a workaround only needed for mixtral q4f16_ft
self.keep_params_after_load = keep_params_after_load

def register_params(
self,
model: nn.Module,
Expand Down Expand Up @@ -606,8 +609,9 @@ def get_item(i):
relax_pname,
[cached_torch_params[torch_pname] for torch_pname in torch_pnames],
)
# for torch_pname in torch_pnames:
# del cached_torch_params[torch_pname]
if not self.keep_params_after_load:
for torch_pname in torch_pnames:
del cached_torch_params[torch_pname]

assert i in cached_relax_params
assert i not in loaded_idx_set
Expand Down