ExecuTorch

ExecuTorch 是一个平台，使 PyTorch 训练和推理程序能够在移动和边缘设备上运行。它由 torch.compile 和 torch.export 提供支持，以实现性能和部署。

您可以将 ExecuTorch 与 Transformers 和 torch.export 一起使用。convert_and_export_with_cache() 方法将 PreTrainedModel 转换为可导出的模块。在底层，它使用 torch.export 导出模型，确保与 ExecuTorch 的兼容性。

import torch
from transformers import LlamaForCausalLM, AutoTokenizer, GenerationConfig
from transformers.integrations.executorch import(
    TorchExportableModuleWithStaticCache,
    convert_and_export_with_cache
)

generation_config = GenerationConfig(
    use_cache=True,
    cache_implementation="static",
    cache_config={
        "batch_size": 1,
        "max_cache_len": 20,
    }
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B", pad_token="</s>", padding_side="right")
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="sdpa", generation_config=generation_config)

exported_program = convert_and_export_with_cache(model)

导出的 PyTorch 模型现在可以与 ExecuTorch 一起使用了。使用 TorchExportableModuleWithStaticCache 包装模型以生成文本。

prompts = ["Simply put, the theory of relativity states that "]
prompt_tokens = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
prompt_token_ids = prompt_tokens["input_ids"]

generated_ids = TorchExportableModuleWithStaticCache.generate(
    exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=20,
)
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
print(generated_text)
['Simply put, the theory of relativity states that 1) the speed of light is the']

< > 更新在 GitHub 上