diff --git a/examples/models/llama/feed_forward.py b/examples/models/llama/feed_forward.py index 60d58c973ea..b23429b89f7 100644 --- a/examples/models/llama/feed_forward.py +++ b/examples/models/llama/feed_forward.py @@ -1,6 +1,6 @@ import torch.nn.functional as F -from executorch.examples.models.llama.lora import LoRALinear +from executorch.examples.models.llama.lora import lora_call, LoRALinear from executorch.examples.models.llama.model_args import ModelArgs from torch import nn @@ -65,15 +65,8 @@ def __init__(self, dim: int, hidden_dim: int, args: ModelArgs): ) def forward(self, x, lora_blob=None): - # CoreML LoRA-as-IO Path-2: when `lora_blob` is provided, route per- - # projection slices to LoRALinear instances tagged with `_lora_key`. - # Default behavior (lora_blob=None) is unchanged. - def _call(linear, x_in): - if lora_blob is not None: - key = getattr(linear, "_lora_key", None) - if key is not None and key in lora_blob: - a, b = lora_blob[key] - return linear(x_in, a, b) - return linear(x_in) - - return _call(self.w2, F.silu(_call(self.w1, x)) * _call(self.w3, x)) + return lora_call( + self.w2, + F.silu(lora_call(self.w1, x, lora_blob)) * lora_call(self.w3, x, lora_blob), + lora_blob, + )