Finetune on these tools button using LORA adapters

needle/cli.py

@@ -122,6 +122,10 @@ def main():
     p.add_argument("--batch-size", type=int, default=64)
     p.add_argument("--lr", type=float, default=3e-5)
     p.add_argument("--muon-lr", type=float, default=0.02)
+    p.add_argument("--lora-rank", type=int, default=0,
+                   help="Enable LoRA adaptation with this rank (default: 0 = disabled)")
+    p.add_argument("--lora-alpha", type=int, default=16,
+                   help="LoRA scaling alpha (default: 16)")
     p.add_argument("--d-model", type=int, default=512)
     p.add_argument("--num-heads", type=int, default=8)
     p.add_argument("--num-kv-heads", type=int, default=4)
@@ -237,6 +241,10 @@ def main():
     p.add_argument("--cache-dir", type=str, default=None)
     p.add_argument("--max-enc-len", type=int, default=None)
     p.add_argument("--max-dec-len", type=int, default=None)
+    p.add_argument("--lora-rank", type=int, default=0,
+                   help="Enable LoRA adaptation with this rank (default: 0 = disabled)")
+    p.add_argument("--lora-alpha", type=int, default=16,
+                   help="LoRA scaling alpha (default: 16)")
 
     p = sub.add_parser("playground", add_help=False)
     p.add_argument("--checkpoint", type=str, default=None)

needle/model/architecture.py

@@ -11,6 +11,10 @@ def default_init():
     return jinit.normal(stddev=0.02)
 
 
+def lora_init(shape, dtype=jnp.float32):
+    return jinit.zeros(shape, dtype)
+
+
 def residual_init(num_layers):
     return jinit.normal(stddev=0.02 / math.sqrt(2 * num_layers))
 
@@ -48,6 +52,8 @@ class TransformerConfig:
     dropout_rate: float = 0.1
     contrastive_dim: int = 128
     no_feedforward: bool = True
+    lora_rank: int = 0
+    lora_alpha: int = 16
 
     def __init__(self, **kwargs):
         valid = {f.name for f in self.__dataclass_fields__.values()}
@@ -81,23 +87,48 @@ def apply_rope(x, cos, sin):
     return jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1)
 
 
+class LoRADense(nn.Module):
+    features: int
+    dtype: jnp.dtype = jnp.bfloat16
+    use_bias: bool = False
+    kernel_init: callable = default_init
+    lora_rank: int = 0
+    lora_alpha: int = 16
+
+    @nn.compact
+    def __call__(self, inputs):
+        kernel = self.param("kernel", self.kernel_init, (inputs.shape[-1], self.features))
+        out = jnp.asarray(jnp.dot(inputs, kernel), dtype=self.dtype)
+        if self.lora_rank > 0:
+            lora_a = self.param("lora_A", lora_init, (inputs.shape[-1], self.lora_rank))
+            lora_b = self.param("lora_B", lora_init, (self.lora_rank, self.features))
+            scaling = self.lora_alpha / max(self.lora_rank, 1)
+            out = out + jnp.asarray(jnp.dot(jnp.dot(inputs, lora_a), lora_b) * scaling, dtype=self.dtype)
+        if self.use_bias:
+            bias = self.param("bias", jinit.zeros, (self.features,))
+            out = out + bias
+        return out
+
+
 class MultiHeadAttention(nn.Module):
     num_heads: int
     num_kv_heads: int
     d_model: int
     num_layers: int
     dtype: jnp.dtype = jnp.bfloat16
     rope_keys_only: bool = False
+    lora_rank: int = 0
+    lora_alpha: int = 16
 
     @nn.compact
     def __call__(self, q_input, kv_input, mask=None, rope=None):
         head_dim = self.d_model // self.num_heads
         kv_dim = self.num_kv_heads * head_dim
         B = q_input.shape[0]
 
-        q = nn.Dense(self.d_model, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="q_proj")(q_input)
-        k = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="k_proj")(kv_input)
-        v = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="v_proj")(kv_input)
+        q = LoRADense(self.d_model, dtype=self.dtype, use_bias=False, kernel_init=default_init(), lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, name="q_proj")(q_input)
+        k = LoRADense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, name="k_proj")(kv_input)
+        v = LoRADense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, name="v_proj")(kv_input)
 
         q = q.reshape(B, -1, self.num_heads, head_dim).transpose(0, 2, 1, 3)
         k = k.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
@@ -127,7 +158,7 @@ def __call__(self, q_input, kv_input, mask=None, rope=None):
 
         out = jnp.matmul(attn_weights, v)
         out = out.transpose(0, 2, 1, 3).reshape(B, -1, self.d_model)
-        return nn.Dense(self.d_model, dtype=self.dtype, use_bias=False, kernel_init=residual_init(self.num_layers), name="out_proj")(out)
+        return LoRADense(self.d_model, dtype=self.dtype, use_bias=False, kernel_init=residual_init(self.num_layers), lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, name="out_proj")(out)
 
 
 class FeedForward(nn.Module):
@@ -136,11 +167,13 @@ class FeedForward(nn.Module):
     num_layers: int
     dtype: jnp.dtype = jnp.bfloat16
     activation: str = "drelu"
+    lora_rank: int = 0
+    lora_alpha: int = 16
 
     @nn.compact
     def __call__(self, x, ffn_mask=None):
-        gate = nn.Dense(self.d_ff, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="gate_proj")(x)
-        up = nn.Dense(self.d_ff, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="up_proj")(x)
+        gate = LoRADense(self.d_ff, dtype=self.dtype, use_bias=False, kernel_init=default_init(), lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, name="gate_proj")(x)
+        up = LoRADense(self.d_ff, dtype=self.dtype, use_bias=False, kernel_init=default_init(), lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, name="up_proj")(x)
         if self.activation == "swiglu":
             h = nn.silu(gate) * up
         elif self.activation == "geglu":
@@ -149,7 +182,7 @@ def __call__(self, x, ffn_mask=None):
             h = nn.relu(gate) * nn.relu(up)
         if ffn_mask is not None:
             h = h * ffn_mask[:, None, :] 
-        return nn.Dense(self.d_model, dtype=self.dtype, use_bias=False, kernel_init=residual_init(self.num_layers), name="down_proj")(h)
+        return LoRADense(self.d_model, dtype=self.dtype, use_bias=False, kernel_init=residual_init(self.num_layers), lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, name="down_proj")(h)
 
 
 class EncoderBlock(nn.Module):
@@ -163,22 +196,28 @@ class EncoderBlock(nn.Module):
     activation: str = "drelu"
     dropout_rate: float = 0.0
     no_feedforward: bool = True
+    lora_rank: int = 0
+    lora_alpha: int = 16
 
     @nn.compact
     def __call__(self, x, mask=None, rope=None, ffn_mask=None, deterministic=True):
         gate = nn.sigmoid(self.param("attn_gate", jinit.zeros, ())).astype(self.dtype)
         residual = x
         x = ZCRMSNorm(dtype=self.dtype)(x)
-        x = MultiHeadAttention(self.num_heads, self.num_kv_heads, self.d_model, self.num_layers, self.dtype, name="self_attn")(
-            x, x, mask=mask, rope=rope
-        )
+        x = MultiHeadAttention(
+            self.num_heads, self.num_kv_heads, self.d_model, self.num_layers,
+            self.dtype, name="self_attn", lora_rank=self.lora_rank, lora_alpha=self.lora_alpha,
+        )(x, x, mask=mask, rope=rope)
         x = residual + gate * nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
 
         if not self.no_feedforward:
             ffn_gate = nn.sigmoid(self.param("ffn_gate", jinit.zeros, ())).astype(self.dtype)
             residual = x
             x = ZCRMSNorm(dtype=self.dtype)(x)
-            x = FeedForward(self.d_model, self.d_ff, self.num_layers, self.dtype, self.activation)(x, ffn_mask=ffn_mask)
+            x = FeedForward(
+                self.d_model, self.d_ff, self.num_layers, self.dtype,
+                self.activation, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha,
+            )(x, ffn_mask=ffn_mask)
             x = residual + ffn_gate * nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
 
         return x
@@ -195,6 +234,8 @@ class _EncoderScanBody(nn.Module):
     activation: str = "drelu"
     dropout_rate: float = 0.0
     no_feedforward: bool = True
+    lora_rank: int = 0
+    lora_alpha: int = 16
     deterministic: bool = True
 
     @nn.compact
@@ -203,7 +244,7 @@ def __call__(self, carry, _):
         x = EncoderBlock(
             self.num_heads, self.num_kv_heads, self.d_model, self.d_ff,
             self.num_layers, self.dtype, self.activation, self.dropout_rate,
-            self.no_feedforward,
+            self.no_feedforward, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha,
         )(x, mask, rope, ffn_mask, self.deterministic)
         return (x, mask, rope, ffn_mask), None
 
@@ -227,7 +268,7 @@ def __call__(self, x, mask=None, rope=None, ffn_mask=None, deterministic=True):
         (x, _, _, _), _ = ScanBlock(
             cfg.num_heads, cfg.num_kv_heads, cfg.d_model, cfg.d_ff,
             cfg.total_layers, dt, cfg.activation, cfg.dropout_rate,
-            cfg.no_feedforward, deterministic, name="layers",
+            cfg.no_feedforward, cfg.lora_rank, cfg.lora_alpha, deterministic, name="layers",
         )((x, mask, rope, ffn_mask), None)
 
         x = ZCRMSNorm(dtype=dt, name="final_norm")(x)
@@ -244,30 +285,37 @@ class DecoderBlock(nn.Module):
     activation: str = "drelu"
     dropout_rate: float = 0.0
     no_feedforward: bool = True
+    lora_rank: int = 0
+    lora_alpha: int = 16
 
     @nn.compact
     def __call__(self, x, encoder_out, self_mask=None, cross_mask=None, rope=None, ffn_mask=None, deterministic=True):
         self_gate = nn.sigmoid(self.param("self_attn_gate", jinit.zeros, ())).astype(self.dtype)
         residual = x
         x = ZCRMSNorm(dtype=self.dtype)(x)
-        x = MultiHeadAttention(self.num_heads, self.num_kv_heads, self.d_model, self.num_layers, self.dtype, name="self_attn")(
-            x, x, mask=self_mask, rope=rope
-        )
+        x = MultiHeadAttention(
+            self.num_heads, self.num_kv_heads, self.d_model, self.num_layers,
+            self.dtype, name="self_attn", lora_rank=self.lora_rank, lora_alpha=self.lora_alpha,
+        )(x, x, mask=self_mask, rope=rope)
         x = residual + self_gate * nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
 
         cross_gate = nn.sigmoid(self.param("cross_attn_gate", jinit.zeros, ())).astype(self.dtype)
         residual = x
         x = ZCRMSNorm(dtype=self.dtype)(x)
-        x = MultiHeadAttention(self.num_heads, self.num_kv_heads, self.d_model, self.num_layers, self.dtype, name="cross_attn")(
-            x, encoder_out, mask=cross_mask
-        )
+        x = MultiHeadAttention(
+            self.num_heads, self.num_kv_heads, self.d_model, self.num_layers,
+            self.dtype, name="cross_attn", lora_rank=self.lora_rank, lora_alpha=self.lora_alpha,
+        )(x, encoder_out, mask=cross_mask)
         x = residual + cross_gate * nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
 
         if not self.no_feedforward:
             ffn_gate = nn.sigmoid(self.param("ffn_gate", jinit.zeros, ())).astype(self.dtype)
             residual = x
             x = ZCRMSNorm(dtype=self.dtype)(x)
-            x = FeedForward(self.d_model, self.d_ff, self.num_layers, self.dtype, self.activation)(x, ffn_mask=ffn_mask)
+            x = FeedForward(
+                self.d_model, self.d_ff, self.num_layers, self.dtype,
+                self.activation, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha,
+            )(x, ffn_mask=ffn_mask)
             x = residual + ffn_gate * nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
 
         return x
@@ -285,6 +333,8 @@ class _DecoderScanBody(nn.Module):
     activation: str = "drelu"
     dropout_rate: float = 0.0
     no_feedforward: bool = True
+    lora_rank: int = 0
+    lora_alpha: int = 16
     deterministic: bool = True
 
     @nn.compact
@@ -293,7 +343,7 @@ def __call__(self, carry, _):
         x = DecoderBlock(
             self.num_heads, self.num_kv_heads, self.d_model, self.d_ff,
             self.num_layers, self.dtype, self.activation, self.dropout_rate,
-            self.no_feedforward,
+            self.no_feedforward, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha,
         )(x, encoder_out, self_mask, cross_mask, rope, ffn_mask, self.deterministic)
         return (x, encoder_out, self_mask, cross_mask, rope, ffn_mask), None
 
@@ -316,7 +366,7 @@ def __call__(self, x, encoder_out, self_mask=None, cross_mask=None, rope=None, f
         (x, _, _, _, _, _), _ = ScanBlock(
             cfg.num_heads, cfg.num_kv_heads, cfg.d_model, cfg.d_ff,
             cfg.total_layers, dt, cfg.activation, cfg.dropout_rate,
-            cfg.no_feedforward, deterministic, name="layers",
+            cfg.no_feedforward, cfg.lora_rank, cfg.lora_alpha, deterministic, name="layers",
         )((x, encoder_out, self_mask, cross_mask, rope, ffn_mask), None)
 
         x = ZCRMSNorm(dtype=dt)(x)

needle/training/finetune.py

@@ -311,27 +311,48 @@ def finetune_local(args):
         print("Starting training...")
         approx_steps = max(1, (train_kept // args.batch_size) * args.epochs)
         cfg = ckpt_config
-        train_args = argparse.Namespace(
-            name=experiment_name, checkpoint=args.checkpoint, init_from=None,
-            epochs=args.epochs, batch_size=args.batch_size, lr=3e-5, muon_lr=0.02,
+        train_args_options = dict(
+            name=experiment_name,
+            checkpoint=args.checkpoint,
+            init_from=None,
+            epochs=args.epochs,
+            batch_size=args.batch_size,
+            lr=3e-5,
+            muon_lr=0.02,
             d_model=cfg["d_model"],
             num_heads=cfg["num_heads"],
             num_kv_heads=cfg.get("num_kv_heads", cfg["num_heads"]),
             num_layers=cfg["num_encoder_layers"],
             num_dec_layers=cfg["num_decoder_layers"],
             d_ff=cfg.get("d_ff", cfg["d_model"] * 4),
-            max_enc_len=max_enc_len, max_dec_len=max_dec_len, max_samples=None,
-            warmup_ratio=0.05, decay_ratio=0.05, wandb=False,
+            max_enc_len=max_enc_len,
+            max_dec_len=max_dec_len,
+            max_samples=None,
+            warmup_ratio=0.05,
+            decay_ratio=0.05,
+            wandb=False,
             dtype=cfg.get("dtype", "bfloat16"),
-            checkpoint_dir=args.checkpoint_dir, seed=42,
-            eval_every=max(1, approx_steps), max_eval_samples=min(val_kept, 50),
+            checkpoint_dir=args.checkpoint_dir,
+            seed=42,
+            eval_every=max(1, approx_steps),
+            max_eval_samples=min(val_kept, 50),
             contrastive_weight=0.1,
             contrastive_dim=cfg.get("contrastive_dim", 128),
             num_memory_slots=cfg.get("num_memory_slots", 64),
-            w_name=2.0, w_value=4.0, w_key=1.5,
+            w_name=2.0,
+            w_value=4.0,
+            w_key=1.5,
             val_ds=val_ds,
+            lora_rank=getattr(args, "lora_rank", 0),
+            lora_alpha=getattr(args, "lora_alpha", 16),
         )
 
+        if getattr(args, "lora_rank", 0) > 0:
+            train_args_options["init_from"] = args.checkpoint
+            train_args_options["checkpoint"] = None
+
+        train_args = argparse.Namespace(**train_args_options)
+
         from ..training.train import train
         train(train_args)
         best_path = _ensure_best_checkpoint(args.checkpoint_dir, run_id)

needle/training/optim.py

@@ -56,12 +56,26 @@ def ortho(g):
 
 
 def _param_labels(params):
-    """Label each param: 'muon' for Dense kernels, 'adam' for the rest."""
+    """Label each param: 'muon' for Dense kernels, 'adam' for the rest.
+
+    If LoRA weights are present, freeze base LoRADense kernels and only train
+    the adapter parameters.
+    """
+
+    def _has_lora_siblings(path):
+        subtree = params
+        for node in path[:-1]:
+            key = node.key if hasattr(node, "key") else str(node)
+            subtree = subtree[key]
+        return hasattr(subtree, "__getitem__") and ("lora_A" in subtree or "lora_B" in subtree)
 
     def _label(path, leaf):
         name = path[-1].key if hasattr(path[-1], "key") else str(path[-1])
-        if name == "kernel" and leaf.ndim in (2, 3):
-            return "muon"
+        if name == "kernel":
+            if _has_lora_siblings(path):
+                return "frozen"
+            if leaf.ndim in (2, 3):
+                return "muon"
         return "adam"
 
     return jax.tree_util.tree_map_with_path(_label, params)
@@ -106,10 +120,11 @@ def create_train_state(rng, config, learning_rate, muon_lr, total_steps, warmup_
         optax.adamw(adam_schedule, b2=0.95, weight_decay=0.0),
     )
 
+    frozen_opt = optax.set_to_zero()
     tx = optax.chain(
         optax.clip_by_global_norm(1.0),
         optax.multi_transform(
-            {"muon": muon_opt, "adam": adam_opt},
+            {"muon": muon_opt, "adam": adam_opt, "frozen": frozen_opt},
             _param_labels,
         ),
     )

needle/training/train.py

@@ -217,6 +217,13 @@ def train(args):
         with open(resume_checkpoint, "rb") as f:
             ckpt_data = pickle.load(f)
         config = TransformerConfig(**ckpt_data["config"])
+        if getattr(args, "lora_rank", 0) > 0:
+            config.lora_rank = args.lora_rank
+            config.lora_alpha = args.lora_alpha
+            if getattr(args, "init_from", None) is None and config.lora_rank > 0 and ckpt_data["config"].get("lora_rank", 0) == 0:
+                print(f"  Switching from resume checkpoint to init_from for LoRA training (rank={config.lora_rank})")
+                args.init_from = resume_checkpoint
+                resume_checkpoint = None
         ckpt_params = jax.tree.map(jnp.array, ckpt_data["params"])
         print(f"  Config: d={config.d_model}, heads={config.num_heads}, layers={config.num_encoder_layers}/{config.num_decoder_layers}")
     else:
@@ -231,6 +238,8 @@ def train(args):
             dtype=args.dtype,
             num_memory_slots=getattr(args, "num_memory_slots", 64),
             contrastive_dim=getattr(args, "contrastive_dim", 128),
+            lora_rank=getattr(args, "lora_rank", 0),
+            lora_alpha=getattr(args, "lora_alpha", 16),
         )
 
     global _PRECISION, _CONTRASTIVE_WEIGHT, _LOSS_WEIGHT_MAP