Skip to content

Qualcomm AI Engine Direct - Fixed the order of the transforms for llama #5221

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

Merged
merged 2 commits into from
Sep 10, 2024
Merged

Qualcomm AI Engine Direct - Fixed the order of the transforms for llama #5221

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
0e23b45
Qualcomm AI Engine Direct - Fixed the order of the transforms for llama
shewu-quic Sep 10, 2024
dfd80b0
fixed ci
shewu-quic Sep 10, 2024
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
7 changes: 6 additions & 1 deletion examples/models/llama2/eval_llama_lib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,13 +41,15 @@ def __init__(
tokenizer: Union[SentencePieceTokenizer, Tiktoken],
max_seq_length: Optional[int] = None,
use_kv_cache: bool = False,
generate_full_logits: bool = False,
enable_dynamic_shape: bool = True,
):
super().__init__(
model=model, tokenizer=tokenizer, max_seq_length=max_seq_length
)
self._model = model.to(self.device)
self._use_kv_cache = use_kv_cache
self._generate_full_logits = generate_full_logits
self._enable_dynamic_shape = enable_dynamic_shape

def _model_call(self, inps):
Expand All @@ -60,7 +62,10 @@ def _model_call(self, inps):
pos_tensor = torch.tensor([pos], dtype=torch.int64)
logits = self._model(inps[:, pos : pos + 1], pos_tensor)
result_logits.append(logits)
return torch.cat(result_logits, dim=1)
if self._generate_full_logits:
return torch.cat(result_logits, dim=1)
else:
return torch.stack(result_logits, dim=1)
Comment on lines +65 to +68
Copy link
Contributor

@cccclai cccclai Sep 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm what's the difference between these?

Copy link
Collaborator Author

@shewu-quic shewu-quic Sep 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because the shape of the function output should be (batch, seq, vocab_size).
If _generate_full_logits, the shape of each result logit in result_logits are (batch, seq, vocab_size)
We could just use cat by dim=1.
If not _generate_full_logits, the shape of each result logit in result_logits are (batch, vocab_size).
We will need use stack to get one more dimension (batch, seq, vocab_size)

else:
pos_tensor = torch.tensor([0], dtype=torch.int64, device=self.device)
# Batch process the whole sequence.
Expand Down
9 changes: 5 additions & 4 deletions examples/models/llama2/export_llama_lib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,7 @@ def build_args_parser() -> argparse.ArgumentParser:
"--optimized_rotation_path",
default=None,
required=False,
help="[QNN Backend] Optimized rotation checkpoint path. Just apply R1/R2 here."
help="[QNN backend] Optimized rotation checkpoint path. Just apply R1/R2 here."
"You can download the optimized rotation matrices from https://github.com/facebookresearch/SpinQuant/tree/main",
)
parser.add_argument(
Expand Down Expand Up @@ -440,6 +440,9 @@ def _prepare_for_llama_export(modelname: str, args) -> LLMEdgeManager:
transforms.append(replace_sdpa_with_flex_sdpa)
transforms.append(replace_causal_mask)
transforms.append(replace_rms_norm_with_native_rms_norm)
if args.optimized_rotation_path:
transforms.append(fuse_layer_norms)
transforms.append(get_model_with_r1_r2(args.optimized_rotation_path))
transforms.append(convert_linear_to_conv2d)

elif args.coreml or args.mps:
Expand All @@ -448,9 +451,6 @@ def _prepare_for_llama_export(modelname: str, args) -> LLMEdgeManager:
transforms.append(replace_sdpa_with_simple_sdpa)
transforms.append(replace_causal_mask)

if args.optimized_rotation_path:
transforms.append(fuse_layer_norms)
transforms.append(get_model_with_r1_r2(args.optimized_rotation_path))
return (
_load_llama_model(
modelname=modelname,
Expand Down Expand Up @@ -744,6 +744,7 @@ def _load_llama_model(
max_seq_len=model.params.max_seq_len,
dtype=dtype,
use_kv_cache=use_kv_cache,
generate_full_logits=generate_full_logits,
example_inputs=example_inputs,
enable_dynamic_shape=enable_dynamic_shape,
calibration_tasks=calibration_tasks,
Expand Down
10 changes: 9 additions & 1 deletion extension/llm/export/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ def __init__(
example_inputs,
args: Optional[Any] = None,
enable_dynamic_shape: bool = False,
generate_full_logits: bool = False,
calibration_tasks: Optional[List[str]] = None,
calibration_limit: Optional[int] = None,
calibration_seq_length: Optional[int] = None,
Expand All @@ -86,6 +87,7 @@ def __init__(
self.dtype = dtype
self.example_inputs = example_inputs
self.use_kv_cache = use_kv_cache
self.generate_full_logits = generate_full_logits
self.enable_dynamic_shape = enable_dynamic_shape
self.verbose = verbose
self.metadata = metadata
Expand Down Expand Up @@ -229,7 +231,12 @@ def calibrate_template(
)
pos += 1
if pos >= len(token_list):
token_list.append(torch.argmax(logits[:], dim=-1).item())
if self.generate_full_logits:
token_list.append(
torch.argmax(logits[:, -1], dim=-1).item()
)
else:
token_list.append(torch.argmax(logits[:], dim=-1).item())

calibrate_template(
module=prepared_module,
Expand All @@ -243,6 +250,7 @@ def calibrate_template(
tokenizer=tokenizer,
max_seq_length=calibration_seq_length,
use_kv_cache=self.use_kv_cache,
generate_full_logits=self.generate_full_logits,
enable_dynamic_shape=self.enable_dynamic_shape,
)
eval_results = evaluate_model(
Expand Down
9 changes: 1 addition & 8 deletions extension/llm/export/partitioner_lib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,16 +139,9 @@ def get_qnn_partitioner(
if pt2e_quantize is not None:
use_fp16 = False

soc_chip_table = {
"SM8650": QcomChipset.SM8650,
"SM8550": QcomChipset.SM8550,
"SM8475": QcomChipset.SM8475,
"SM8450": QcomChipset.SM8450,
}

return QnnPartitioner( # pyre-fixme[16]
generate_qnn_executorch_compiler_spec( # pyre-fixme[16]
soc_model=soc_chip_table[soc_model], # pyre-fixme[16]
soc_model=getattr(QcomChipset, soc_model), # pyre-fixme[16]
# pyre-fixme[16]
backend_options=generate_htp_compiler_spec(
use_fp16=use_fp16,
Expand Down