diff --git a/needle/cli.py b/needle/cli.py index 705ba92..2259e12 100644 --- a/needle/cli.py +++ b/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) diff --git a/needle/model/architecture.py b/needle/model/architecture.py index c62cc4d..8af850a 100644 --- a/needle/model/architecture.py +++ b/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,6 +87,29 @@ 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 @@ -88,6 +117,8 @@ class MultiHeadAttention(nn.Module): 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): @@ -95,9 +126,9 @@ def __call__(self, q_input, kv_input, mask=None, rope=None): 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) diff --git a/needle/training/finetune.py b/needle/training/finetune.py index e494308..08cf6e6 100644 --- a/needle/training/finetune.py +++ b/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) diff --git a/needle/training/optim.py b/needle/training/optim.py index be8890b..78aad48 100644 --- a/needle/training/optim.py +++ b/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, ), ) diff --git a/needle/training/train.py b/needle/training/train.py index 657559f..6bc5176 100644 --- a/needle/training/train.py +++ b/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