Qualcomm AI Engine Direct - support embedding op

[haowhsu-quic]

summary:

• support embedding op with int32 index input
• make mobilebert / llama2 be fully delegated
• add requantize passes for mixed precision
• bug fixes

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2057

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit f223b65 with merge base 81b3232 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[cccclai]

Two main comments:

• Is there a way to repro the error on our side? maybe it's an edge case we'd need to fix
• Can we leave a todo (maybe an expected failing unit test)?

[haowhsu-quic]

Two main comments:

• Is there a way to repro the error on our side? maybe it's an edge case we'd need to fix
• Can we leave a todo (maybe an expected failing unit test)?

Thanks for reviewing! I've added TODO items for next action.

It could be reproduced with following patch:

[examples/qualcomm/scripts/mobilebert_fine_tune.py]
@@ -58,10 +58,18 @@ def accuracy_per_class(preds, goldens, labels):
def get_dataset(data_val):
    # prepare input data
    inputs, input_list = [], ""
-   # max_position_embeddings defaults to 512
-   position_ids = torch.arange(512).expand((1, -1)).to(torch.int32)
    for index, data in enumerate(data_val):
        data = [d.to(torch.int32) for d in data]
+       # input_ids, attention_mask, token_type_ids
+       inputs.append((*data[:2], torch.zeros(data[0].size(), dtype=torch.int32)))
-       # input_ids, attention_mask, token_type_ids, position_ids
-       inputs.append(
-           (
-               *data[:2],
-               torch.zeros(data[0].size(), dtype=torch.int32),
-               position_ids[:, : data[0].shape[1]],
-           )
-       )
        input_text = " ".join(
            [f"input_{index}_{i}.raw" for i in range(len(inputs[-1]))]
        )
@@ -204,9 +212,6 @@ def get_fine_tuned_mobilebert(artifacts_dir, pretrained_weight, batch_size):
            map_location=torch.device("cpu"),
        ),
    )
+   # hack for changing dtype of "position_ids" from int64 to int32
+   sub_module = model.mobilebert.embeddings
+   sub_module.position_ids = sub_module.position_ids.to(torch.int32)

    return model.eval(), dataloader_val, labels

[cccclai]

Thanks for the update. Seems like some llama related CI jobs start failign. Those changes look legit but Is it okay to remove changes in examples/models/llama2/model.py to get CI green? We can do a seperate pr for this.

[haowhsu-quic]

Thank you, I moved the datatype casting from examples/models/llama2/model.py to our own script as mobilebert does.

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@cccclai merged this pull request in 57e192b.