Update on "Add LoRA linear definition"

^
Add lora linear definition. Pull out linears from attention, and allow custom linear (eg. lora linear) to be passed in.

If none, construct linear (current behaviour).

Differential Revision: [D73953776](https://our.internmc.facebook.com/intern/diff/D73953776/)

[ghstack-poisoned]

Author: lucylq

examples/models/llama/attention.py

@@ -5,8 +5,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from executorch.examples.models.llama.model_args import ModelArgs
 from executorch.examples.models.llama.lora import LoRALinear
+from executorch.examples.models.llama.model_args import ModelArgs
 from executorch.examples.models.llama.norm import RMSNorm
 from executorch.examples.models.llama.rope import Rope
 
@@ -373,8 +373,7 @@ def __init__(
                 dropout=0.0,
                 use_bias=args.attention_qkv_bias,
             )
-            if args.target_modules is not None
-            and "q_proj" in args.target_modules
+            if args.target_modules is not None and "q_proj" in args.target_modules
             else nn.Linear(
                 self.dim, self.n_heads * self.head_dim, bias=self.attention_qkv_bias
             )
@@ -388,8 +387,7 @@ def __init__(
                 dropout=0.0,
                 use_bias=args.attention_qkv_bias,
             )
-            if args.target_modules is not None
-            and "k_proj" in args.target_modules
+            if args.target_modules is not None and "k_proj" in args.target_modules
             else nn.Linear(
                 self.dim, self.n_kv_heads * self.head_dim, bias=self.attention_qkv_bias
             )
@@ -403,8 +401,7 @@ def __init__(
                 dropout=0.0,
                 use_bias=args.attention_qkv_bias,
             )
-            if args.target_modules is not None
-            and "v_proj" in args.target_modules
+            if args.target_modules is not None and "v_proj" in args.target_modules
             else nn.Linear(
                 self.dim, self.n_kv_heads * self.head_dim, bias=self.attention_qkv_bias
             )
@@ -418,8 +415,7 @@ def __init__(
                 dropout=0.0,
                 use_bias=args.attention_qkv_bias,
             )
-            if args.target_modules is not None
-            and "output_proj" in args.target_modules
+            if args.target_modules is not None and "output_proj" in args.target_modules
             else nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False)
         )