From 462ac33d7b2b44806464b2a4d94c3ccfe61b0146 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 17 Jun 2025 11:34:47 +0000 Subject: [PATCH] per layer lightweightlora --- .../gemma_self_gen_qa_rank_16.sh | 24 +++ .../gemma_self_gen_qa_rank_16_lightweight.sh | 26 ++++ .../gemma_self_gen_qa_repeat_rank_16.sh | 24 +++ ..._self_gen_qa_repeat_rank_16_lightweight.sh | 26 ++++ src/ctx_to_lora/modeling/hypernet.py | 144 ++++++++++++------ 5 files changed, 194 insertions(+), 50 deletions(-) create mode 100644 scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16.sh create mode 100644 scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16_lightweight.sh create mode 100644 scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16.sh create mode 100644 scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16_lightweight.sh diff --git a/scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16.sh b/scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16.sh new file mode 100644 index 0000000..3871e79 --- /dev/null +++ b/scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16.sh @@ -0,0 +1,24 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora_medium +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=8 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +uv run accelerate launch --num_processes=8 --gradient_accumulation_steps=16 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29561 intx_sft.py configs/self_gen_3_small.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=2.2 --per_device_train_batch_size=-1 \ +--gradient_accumulation_steps=16 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=1 \ +--lora_r=16 \ +--eval_steps=1000000000000000 --do_eval=no --do_predict=no --eval_on_start=False \ +--save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=8192 \ +--per_rank_gen=True \ +--per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.1 \ diff --git a/scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16_lightweight.sh b/scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16_lightweight.sh new file mode 100644 index 0000000..6df3fb5 --- /dev/null +++ b/scripts/gemma_rank_exp/gemma_self_gen_qa_rank_16_lightweight.sh @@ -0,0 +1,26 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora_medium +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=32 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29563 intx_sft.py configs/self_gen_3_small.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 \ +--gradient_accumulation_steps=32 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=1 \ +--lora_r=16 \ +--eval_steps=1000000000000000 --do_eval=no --do_predict=no --eval_on_start=False \ +--save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 \ +--num_pre_head_layers=1 \ +--latent_size=256 \ +--neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 \ +--add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=8192 \ +--per_rank_gen=True \ +--per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.1 \ diff --git a/scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16.sh b/scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16.sh new file mode 100644 index 0000000..7515897 --- /dev/null +++ b/scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16.sh @@ -0,0 +1,24 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora_medium +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=8 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +uv run accelerate launch --num_processes=8 --gradient_accumulation_steps=16 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29562 intx_sft.py configs/self_gen_3_small.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=2.2 --per_device_train_batch_size=-1 \ +--gradient_accumulation_steps=16 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=1 \ +--lora_r=16 \ +--eval_steps=1000000000000000 --do_eval=no --do_predict=no --eval_on_start=False \ +--save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--add_negative_prompt=False \ +--add_repeat_prompt=True --repeat_prob=0.1 \ +--use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=8192 \ +--per_rank_gen=True \ +--per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.1 \ diff --git a/scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16_lightweight.sh b/scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16_lightweight.sh new file mode 100644 index 0000000..ad39662 --- /dev/null +++ b/scripts/gemma_rank_exp/gemma_self_gen_qa_repeat_rank_16_lightweight.sh @@ -0,0 +1,26 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora_medium +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=32 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29564 intx_sft.py configs/self_gen_3_small.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 \ +--gradient_accumulation_steps=32 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=1 \ +--lora_r=16 \ +--eval_steps=1000000000000000 --do_eval=no --do_predict=no --eval_on_start=False \ +--save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 \ +--num_pre_head_layers=1 \ +--latent_size=256 \ +--neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 \ +--add_negative_prompt=False \ +--add_repeat_prompt=True --repeat_prob=0.1 \ +--use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=8192 \ +--per_rank_gen=True \ +--per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.1 \ diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 62c6539..1c02ce1 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -6,7 +6,7 @@ from dataclasses import dataclass from typing import Any import torch -from einops import einsum, rearrange, unpack +from einops import rearrange, unpack from einops.layers.torch import EinMix as Mix from jaxtyping import Float, Integer from peft import ( @@ -59,6 +59,7 @@ class HypernetConfig: per_rank_gen: bool per_layer_processing: bool use_token_mixing: bool + num_pre_head_layers: int dropout_rate: float lora_config: LoraConfig @@ -240,7 +241,7 @@ class HyperLoRA(nn.Module): self.d_latent * 4, self.d_latent, ) - for _ in range(4) + for _ in range(self.config.num_pre_head_layers) ] else: layers = [ @@ -250,49 +251,86 @@ class HyperLoRA(nn.Module): output_size=self.config.latent_size, dropout_rate=getattr(self.config, "dropout_rate", 0), ) - for _ in range(4) + for _ in range(self.config.num_pre_head_layers) ] self.layers = nn.Sequential(*layers) d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) - # if self.config.use_light_weight_lora: - # # light-weight lora projection (per module) + if self.config.use_light_weight_lora: + # light-weight lora projection (per layer, per module) + # self.pre_lora_projection = nn.ParameterDict( + # { + # k: nn.Parameter( + # torch.randn( + # self.d_in[k], self.config.light_weight_latent_size + # ) + # ) + # for k in self.target_modules + # } + # ) + self.pre_lora_projection = nn.ParameterDict( + { + k: Mix( + "bs n_layers r d_latent -> bs n_layers r d_in", + weight_shape="n_layers d_latent d_in", + bias_shape="n_layers d_in", + n_layers=len(self.layer_indices), + d_latent=self.config.light_weight_latent_size, + d_in=self.d_in[k], + ) + for k in self.target_modules + } + ) + with torch.no_grad(): + for param in self.pre_lora_projection.values(): + for i in range(param.weight.shape[0]): + nn.init.orthogonal_( + param.weight[i], + # 1 / self.config.light_weight_latent_size**0.5, + ) + param.bias.data[:] = 0 - # self.pre_lora_projection = nn.ParameterDict( - # { - # k: nn.Parameter( - # torch.randn( - # self.d_in[k], self.config.light_weight_latent_size - # ) - # ) - # for k in self.target_modules - # } - # ) - # for param in self.pre_lora_projection.values(): - # nn.init.orthogonal_(param) + # self.post_lora_projection = nn.ParameterDict( + # { + # k: nn.Parameter( + # torch.randn( + # self.d_out[k], self.config.light_weight_latent_size + # ) + # ) + # for k in self.target_modules + # } + # ) + self.post_lora_projection = nn.ParameterDict( + { + k: Mix( + "bs n_layers d_latent r -> bs n_layers d_out r", + weight_shape="n_layers d_latent d_out", + bias_shape="n_layers d_out", + n_layers=len(self.layer_indices), + d_latent=self.config.light_weight_latent_size, + d_out=self.d_out[k], + ) + for k in self.target_modules + } + ) + with torch.no_grad(): + for param in self.post_lora_projection.values(): + for i in range(param.weight.shape[0]): + nn.init.orthogonal_( + param.weight[i], + # 1 / self.config.light_weight_latent_size**0.5, + ) + param.bias.data[:] = 0 - # self.post_lora_projection = nn.ParameterDict( - # { - # k: nn.Parameter( - # torch.randn( - # self.d_out[k], self.config.light_weight_latent_size - # ) - # ) - # for k in self.target_modules - # } - # ) - # for param in self.post_lora_projection.values(): - # nn.init.orthogonal_(param) + d_lora = self.config.light_weight_latent_size * 2 + self.d_in = {k: self.config.light_weight_latent_size for k in self.d_in} + self.d_out = { + k: self.config.light_weight_latent_size for k in self.d_out + } - # d_lora = self.config.light_weight_latent_size * 2 - # self.d_in = {k: self.config.light_weight_latent_size for k in self.d_in} - # self.d_out = { - # k: self.config.light_weight_latent_size for k in self.d_out - # } - - # logger.info(f"Using light-weight LoRA with d_lora = {d_lora // 2}") + logger.info(f"Using light-weight LoRA with d_lora = {d_lora // 2}") n_modules = len(self.target_modules) # have to do this otherwise doesnt work with adamw_torch_fused @@ -404,7 +442,7 @@ class HyperLoRA(nn.Module): # # output_size=self.config.latent_size, # # dropout_rate=getattr(self.config, "dropout_rate", 0), # # ) - # # for _ in range(4) + # # for _ in range(self.config.num_pre_head_layers) # # ] # # ) # self.extra_layers = ResMLPBlock( @@ -464,16 +502,18 @@ class HyperLoRA(nn.Module): # transpose B B = rearrange(B, "bs n_layers r d_out -> bs n_layers d_out r") if self.config.use_light_weight_lora: - A = einsum( - self.pre_lora_projection[module], - A, - "d_in d_latent, bs n_layers r d_latent -> bs n_layers r d_in", - ) - B = einsum( - self.post_lora_projection[module], - B, - "d_out d_latent, bs n_layers d_latent r -> bs n_layers d_out r", - ) + # A = einsum( + # self.pre_lora_projection[module], + # A, + # "d_in d_latent, bs n_layers r d_latent -> bs n_layers r d_in", + # ) + # B = einsum( + # self.post_lora_projection[module], + # B, + # "d_out d_latent, bs n_layers d_latent r -> bs n_layers d_out r", + # ) + A = self.pre_lora_projection[module](A) + B = self.post_lora_projection[module](B) lora_dict[module] = dict(A=A, B=B) @@ -573,6 +613,8 @@ class ModulatedPretrainedModel(nn.Module): use_flash_attn=use_flash_attn, ) hypernet_config = state_dict["hypernet_config"] + if getattr(hypernet_config, "num_pre_head_layers", None) is None: + hypernet_config.num_pre_head_layers = 4 ctx_encoder_args = state_dict["ctx_encoder_args"] model = cls(base_model, hypernet_config, ctx_encoder_args) model.load_state_dict(state_dict) @@ -631,13 +673,15 @@ class ModulatedPretrainedModel(nn.Module): A = A[0:1] B = B[:, 0:1] biases = [A, B.T] - # TODO: scale the biases by the output size? - # e.g., rank=16 gives bigger (2x?) gradient magnitudes at initialization - # compared to rank=8 # bias-hyperinit # init weights to zeros and bias to the base weights bias_cat = torch.cat(biases, dim=1) + # scale the biases by the output size + # e.g., rank=16 gives bigger (2x?) gradient magnitudes at initialization + # compared to rank=8 + r = self.hypernet_config.lora_config.r + bias_cat = bias_cat / r self.hypernet.head.bias.data[..., i, :, : bias_cat.shape[1]] = bias_cat def state_dict(self, *args, **kwargs):