diff --git a/chat_templates/Qwen/Qwen3-4B-Instruct-2507.jinja b/chat_templates/Qwen/Qwen3-4B-Instruct-2507.jinja new file mode 100644 index 0000000..09575ea --- /dev/null +++ b/chat_templates/Qwen/Qwen3-4B-Instruct-2507.jinja @@ -0,0 +1,18 @@ +{%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} +{%- endif %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{ content + '<|im_end|>' }}{% endgeneration %}{{ '\n' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} +{%- endif %} \ No newline at end of file diff --git a/chat_templates/mistralai/Mistral-7B-Instruct-v0.2.jinja b/chat_templates/mistralai/Mistral-7B-Instruct-v0.2.jinja new file mode 100644 index 0000000..fa2ff94 --- /dev/null +++ b/chat_templates/mistralai/Mistral-7B-Instruct-v0.2.jinja @@ -0,0 +1,24 @@ +{%- if messages[0]['role'] == 'system' %} + {%- set system_message = messages[0]['content'] %} + {%- set loop_messages = messages[1:] %} +{%- else %} + {%- set loop_messages = messages %} +{%- endif %} + +{{- bos_token }} +{%- for message in loop_messages %} + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %} + {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }} + {%- endif %} + {%- if message['role'] == 'user' %} + {%- if loop.first and system_message is defined %} + {{- ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST] ' }} + {%- else %} + {{- ' [INST] ' + message['content'] + ' [/INST] ' }} + {%- endif %} + {%- elif message['role'] == 'assistant' %} + {% generation %}{{ message['content'] + eos_token }}{% endgeneration %} + {%- else %} + {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }} + {%- endif %} +{%- endfor %} \ No newline at end of file diff --git a/scripts/niah/train_mistral_7b.sh b/scripts/niah/train_mistral_7b.sh new file mode 100644 index 0000000..3d2eaab --- /dev/null +++ b/scripts/niah/train_mistral_7b.sh @@ -0,0 +1,39 @@ +WANDB_MODE=disabled run uv run train.py configs/niah_exp/ctx_magic_number_32_256.yaml \ +--model_name_or_path=mistralai/Mistral-7B-Instruct-v0.2 \ +--num_train_epochs=1 \ +--per_device_train_batch_size=-1 \ +--gradient_accumulation_steps=64 \ +--per_device_eval_batch_size=16 \ +--exp_setup=hyper_lora \ +--aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_blocks=8 \ +--num_self_attn_per_block=0 \ +--num_pre_head_layers=1 \ +--lora_r=8 \ +--eval_steps=100 \ +--logging_steps=10 \ +--save_steps=1000 \ +--learning_rate=4e-5 \ +--lora_dropout=0.0 \ +--neftune_noise_alpha=0 \ +--per_rank_gen=True \ +--per_layer_processing=True \ +--gen_lora_l1_reg_coef=2.0 \ +--use_sequence_packing=True \ +--max_packed_inp_len=1024 \ +--max_packed_ctx_len=1024 \ +--dataloader_num_workers=0 \ +--dataloader_prefetch_factor=None \ +--eval_on_start=False \ +--ctx_encoder_type=early_exit \ +--n_latent_queries=208 \ +--use_kl_loss=False \ +--eval_on_start=True \ +--max_ctx_chunk_len=512 \ +--min_ctx_chunk_len=25 \ +--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \ +--max_val_samples_per_ds=100 \ +--seed=1 \ +--use_per_ctx_average_loss=True \ +--torch_empty_cache_steps=10 diff --git a/scripts/niah/train_qwen3_4b.sh b/scripts/niah/train_qwen3_4b.sh new file mode 100644 index 0000000..78c32be --- /dev/null +++ b/scripts/niah/train_qwen3_4b.sh @@ -0,0 +1,39 @@ +WANDB_MODE=disabled run uv run train.py configs/niah_exp/ctx_magic_number_32_256.yaml \ +--model_name_or_path=Qwen/Qwen3-4B-Instruct-2507 \ +--num_train_epochs=1 \ +--per_device_train_batch_size=-1 \ +--gradient_accumulation_steps=32 \ +--per_device_eval_batch_size=16 \ +--exp_setup=hyper_lora \ +--aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_blocks=8 \ +--num_self_attn_per_block=0 \ +--num_pre_head_layers=1 \ +--lora_r=8 \ +--eval_steps=100 \ +--logging_steps=10 \ +--save_steps=1000 \ +--learning_rate=4e-5 \ +--lora_dropout=0.0 \ +--neftune_noise_alpha=0 \ +--per_rank_gen=True \ +--per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.5 \ +--use_sequence_packing=True \ +--max_packed_inp_len=2048 \ +--max_packed_ctx_len=2048 \ +--dataloader_num_workers=0 \ +--dataloader_prefetch_factor=None \ +--eval_on_start=False \ +--ctx_encoder_type=early_exit \ +--n_latent_queries=208 \ +--use_kl_loss=False \ +--eval_on_start=True \ +--max_ctx_chunk_len=512 \ +--min_ctx_chunk_len=25 \ +--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \ +--max_val_samples_per_ds=100 \ +--seed=1 \ +--use_per_ctx_average_loss=True \ +--torch_empty_cache_steps=10 diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index 1484aa8..65c4032 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/src/ctx_to_lora/data/definitions.py @@ -10,6 +10,16 @@ CTX_AFFIXES = { "prefix": [2, 106, 1645, 110], # user\n\n\n "suffix": [107, 108, 106, 2516, 108], # \nmodel\n }, + "mistralai/Mistral-7B-Instruct-v0.2": { + "prefix": [1, 733, 16289, 28793, 28705, 13, 13], # ` [INST] \n\n` + "suffix": [733, 28748, 16289, 28793], # ` [/INST] ` + }, + "Qwen/Qwen3-4B-Instruct-2507": { + # `<|im_start|>system\n<|im_end|>\n<|im_start|>user\n` + "prefix": [151644, 8948, 198, 151645, 198, 151644, 872, 198], + # `<|im_end|>\n<|im_start|>assistant\n` + "suffix": [151645, 198, 151644, 77091, 198], + }, }