diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2.sh new file mode 100644 index 0000000..0ab70bf --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2.sh @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29560 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=2 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_neft1.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_neft1.sh new file mode 100644 index 0000000..653d6ff --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_neft1.sh @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29560 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=2 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=1 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4.sh new file mode 100644 index 0000000..9db5102 --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4.sh @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29563 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=4 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing.sh new file mode 100644 index 0000000..9252e4b --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing.sh @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29562 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=4 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.1 diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_decoder_2.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_decoder_2.sh new file mode 100644 index 0000000..db57cf1 --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_decoder_2.sh @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29561 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=4 \ +--decoder_depth=2 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_no_reg.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_no_reg.sh new file mode 100644 index 0000000..e93bda7 --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_no_reg.sh @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29562 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=4 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.0 diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_token_mixing.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_token_mixing.sh new file mode 100644 index 0000000..d68d68d --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac4_per_layer_processing_token_mixing.sh @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29562 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_self_attends_per_block=8 --num_latent_factor=4 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--neftune_noise_alpha=5 --use_light_weight_lora=False \ +--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True \ +--use_token_mixing=True diff --git a/src/ctx_to_lora/modeling_idefics2.py b/src/ctx_to_lora/modeling_idefics2.py index 04a9e48..4f60c38 100644 --- a/src/ctx_to_lora/modeling_idefics2.py +++ b/src/ctx_to_lora/modeling_idefics2.py @@ -80,7 +80,7 @@ class Idefics2PerceiverConfig(PretrainedConfig): self, input_size: int, intermediate_size_factor: int = 1, - hidden_act="gelu_pytorch_tanh", + hidden_act="silu", hidden_size=4096, rms_norm_eps=1e-06, resampler_n_latents=64, diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 50881cd..c53debf 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -288,7 +288,7 @@ class Mixer(nn.Module): self.gate_proj = nn.Linear(input_size, intermediate_emb_size, bias=False) self.up_proj = nn.Linear(input_size, intermediate_emb_size, bias=False) self.down_proj = nn.Linear(intermediate_emb_size, output_size, bias=False) - self.act_fn = nn.GELU(approximate="tanh") + self.act_fn = nn.SiLU() def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) @@ -309,14 +309,14 @@ class MLPResidualBlock(nn.Module): # if pre_layer_norm: # layers.append(nn.LayerNorm(input_size)) layers = [ - Gemma3RMSNorm(input_size), + nn.LayerNorm(input_size), nn.Dropout(dropout_rate), nn.Linear(input_size, hidden_size), - nn.GELU(approximate="tanh"), + nn.SiLU(), nn.Dropout(dropout_rate), nn.Linear(hidden_size, output_size), - Gemma3RMSNorm(output_size), - # nn.GELU(approximate="tanh"), + nn.LayerNorm(output_size), + # nn.SiLU(), ] # if post_dropout: # layers.append(nn.Dropout(dropout_rate)) @@ -576,26 +576,64 @@ class MixerPerLayer(nn.Module): self.up_proj = UpMixPerLayer(n_layers, n_modules, r, d_in, d_in * 4) self.gate_proj = UpMixPerLayer(n_layers, n_modules, r, d_in, d_in * 4) self.down_proj = DownMixPerLayer(n_layers, n_modules, r, d_in * 4, d_in) - self.act_fn = nn.GELU(approximate="tanh") + self.act_fn = nn.SiLU() def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) +# class ResMLPBlockPerLayer(nn.Module): +# def __init__(self, n_layers: int, n_modules: int, r: int, d_in: int): +# super().__init__() +# # input shape: [bs, n_layers, n_modules, feature_dim] +# self.pre_norm = nn.LayerNorm(d_in) +# self.post_norm = nn.LayerNorm(d_in) +# self.mixer = MixerPerLayer(n_layers, n_modules, r, d_in) + +# def forward(self, x): +# inp = x +# x = self.pre_norm(x) +# x = self.mixer(x) +# x = self.post_norm(x) +# return x + inp + + class ResMLPBlockPerLayer(nn.Module): - def __init__(self, n_layers: int, n_modules: int, r: int, d_in: int): + def __init__( + self, + n_layers: int, + input_size: int, + hidden_size: int, + output_size: int, + ): super().__init__() - # input shape: [bs, n_layers, n_modules, feature_dim] - self.pre_norm = Gemma3RMSNorm(d_in) - self.post_norm = Gemma3RMSNorm(d_in) - self.mixer = MixerPerLayer(n_layers, n_modules, r, d_in) + layers = [ + nn.LayerNorm(input_size), + # nn.Linear(input_size, hidden_size), + Mix( + "bs n_layers n_modules r d_in -> bs n_layers n_modules r d_hid", + weight_shape="n_layers d_in d_hid", + bias_shape="n_layers d_hid", + n_layers=n_layers, + d_in=input_size, + d_hid=hidden_size, + ), + nn.SiLU(), + # nn.Linear(hidden_size, output_size), + Mix( + "bs n_layers n_modules r d_hid -> bs n_layers n_modules r d_out", + weight_shape="n_layers d_hid d_out", + bias_shape="n_layers d_out", + n_layers=n_layers, + d_hid=hidden_size, + d_out=output_size, + ), + nn.LayerNorm(output_size), + ] + self.layers = nn.Sequential(*layers) def forward(self, x): - inp = x - x = self.pre_norm(x) - x = self.mixer(x) - x = self.post_norm(x) - return x + inp + return x + self.layers(x) class ResMLPTokenMixingPerLayer(nn.Module): @@ -603,8 +641,8 @@ class ResMLPTokenMixingPerLayer(nn.Module): super().__init__() # input shape: [bs, n_layers, n_modules, feature_dim] d_hid = d_in // 4 - self.pre_norm = Gemma3RMSNorm(d_in) - self.post_norm = Gemma3RMSNorm(d_in) + self.pre_norm = nn.LayerNorm(d_in) + self.post_norm = nn.LayerNorm(d_in) # self.proj_down = Mix( # "bs n_layers n_modules r d_in -> bs n_layers n_modules r d_hid", # weight_shape="n_layers d_in d_hid", @@ -624,7 +662,7 @@ class ResMLPTokenMixingPerLayer(nn.Module): # d_hid=d_hid, # d_in=d_in, # ) - # act_fn = nn.GELU(approximate="tanh") + # act_fn = nn.SiLU() # out_proj = Mix( # "bs n_layers n_modules r d_in -> bs n_layers n_modules r0 d_in0", # weight_shape="n_layers r d_in r0 d_in0", @@ -652,7 +690,7 @@ class ResMLPTokenMixingPerLayer(nn.Module): r=r, r0=r * n_token_factors, ) - act_fn = nn.GELU(approximate="tanh") + act_fn = nn.SiLU() token_mix_down = Mix( "bs n_layers n_modules r0 d_in -> bs n_layers n_modules r d_in", weight_shape="n_layers r0 r", @@ -717,17 +755,29 @@ class HyperLoRA(nn.Module): # ) if self.config.per_layer_processing: layers = [ - nn.Linear(self.d_latent, self.d_latent), + # nn.Linear(self.d_latent, self.d_latent), + # Mix( + # "bs n_layers n_modules r d -> bs n_layers n_modules r d0", + # weight_shape="n_layers d d0", + # bias_shape="n_layers d0", + # n_layers=self.n_layers, + # d=self.d_latent, + # d0=self.d_latent, + # ), ResMLPBlockPerLayer( self.n_layers, - self.num_modules, - self.lora_config.r, self.d_latent, + self.d_latent * 4, + self.d_latent, + # self.num_modules, + # self.lora_config.r, + # self.d_latent, ), ] else: layers = [ - nn.Linear(self.d_latent, self.d_latent), + # nn.Linear(self.d_latent, self.d_latent), + # nn.LayerNorm(self.d_latent), MLPResidualBlock( input_size=self.config.latent_size, hidden_size=self.config.latent_size * 4, @@ -819,9 +869,9 @@ class HyperLoRA(nn.Module): if self.config.per_layer_processing: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", - weight_shape="n_layers n_modules r d_latent d_lora", + weight_shape="n_layers n_modules d_latent d_lora", # bias_shape=None, # no bias - bias_shape="n_layers n_modules r d_lora", + bias_shape="n_layers n_modules d_lora", n_layers=len(self.layer_indices), n_modules=n_modules, d_latent=self.config.latent_size, @@ -831,9 +881,9 @@ class HyperLoRA(nn.Module): else: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", - weight_shape="n_modules r d_latent d_lora", + weight_shape="n_modules d_latent d_lora", # bias_shape=None, # no bias - bias_shape="n_modules r d_lora", + bias_shape="n_modules d_lora", n_layers=len(self.layer_indices), n_modules=n_modules, d_latent=self.config.latent_size, @@ -852,15 +902,16 @@ class HyperLoRA(nn.Module): r=self.config.lora_config.r, d_lora=d_lora, ) - self.head = Mix( - "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", - weight_shape="d_latent r d_lora", - # bias_shape=None, # no bias - bias_shape="r d_lora", - d_latent=self.config.latent_size, - r=self.config.lora_config.r, - d_lora=d_lora, - ) + else: + self.head = Mix( + "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", + weight_shape="d_latent r d_lora", + # bias_shape=None, # no bias + bias_shape="r d_lora", + d_latent=self.config.latent_size, + r=self.config.lora_config.r, + d_lora=d_lora, + ) else: # each module processes d -> r d_out independently if self.config.per_layer_processing: diff --git a/src/ctx_to_lora/training_utils.py b/src/ctx_to_lora/training_utils.py index 23016bb..2153d04 100644 --- a/src/ctx_to_lora/training_utils.py +++ b/src/ctx_to_lora/training_utils.py @@ -114,7 +114,7 @@ def train_model( logger.info(f"Resuming from the checkpoint: {checkpoint}") is_modulated_model = isinstance(model, ModulatedPretrainedModel) - trainer_cls = Trainer if not is_modulated_model else ModulatedModelTrainer + trainer_cls = Trainer # if not is_modulated_model else ModulatedModelTrainer trainer_kwargs = dict( model=model, args=training_args, @@ -123,10 +123,10 @@ def train_model( data_collator=train_collator, compute_metrics=compute_metrics, ) - if is_modulated_model: - logger.info(f"Training with modulated model. Using CustomTrainer.") - trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef - del training_args.gen_lora_l1_reg_coef + # if is_modulated_model: + # logger.info(f"Training with modulated model. Using CustomTrainer.") + # trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef + # del training_args.gen_lora_l1_reg_coef trainer = trainer_cls(**trainer_kwargs)