Skip to content

Commit

Permalink
fix device mapping issue of llama gptq (#2101)
Browse files Browse the repository at this point in the history
Signed-off-by: Xin He <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
xin3he and pre-commit-ci[bot] authored Dec 27, 2024
1 parent b2d019f commit 9bddd52
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion neural_compressor/torch/algorithms/weight_only/gptq.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,8 @@ def trace_gptq_target_blocks(module, module_types=[torch.nn.ModuleList, torch.nn
gptq_related_blocks["transformers"] = m
find_transformers = True
# return gptq_related_blocks
elif is_leaf(m) and not find_transformers:
elif (is_leaf(m) and not find_transformers) or "Embedding" in type(m).__name__:
# "Embedding" in type(m).__name__ to resolve 'LlamaRotaryEmbedding'
gptq_related_blocks["embeddings"][n] = m
elif n.find(gptq_related_blocks["transformers_name"]) == -1 and find_transformers:
# no longer belong to transformers
Expand Down

0 comments on commit 9bddd52

Please sign in to comment.