From 86f073edd6032b56f4e5f9bdea9220d9adac7083 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 1 Apr 2024 02:02:13 +0000 Subject: [PATCH] Add reference --- vllm/model_executor/models/tpu/gemma.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/model_executor/models/tpu/gemma.py b/vllm/model_executor/models/tpu/gemma.py index 8242f29551f1e..56c0701b864f6 100644 --- a/vllm/model_executor/models/tpu/gemma.py +++ b/vllm/model_executor/models/tpu/gemma.py @@ -1,5 +1,8 @@ """Inference-only Gemma model compatible with HF weights. +Adapted from +https://github.com/google/gemma_pytorch/blob/main/gemma/model_xla.py + NOTE(woosuk): This is a temporary workaround to run the Gemma model using PyTorch XLA. This should be merged into the main Gemma model implementation once the custom ops are refactored and the model becomes torch.compile-able.