repeat only shorter ctxs

configs/fw_qa.yaml

@@ -0,0 +1,44 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00002
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 5
+weight_decay: 0.01
+warmup_ratio: 0.1
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 16
+# LoRA
+lora_r: 8
+lora_dropout: 0.05
+target_modules:
+  - down_proj
+  - up_proj
+
+# data
+train_ds_names:
+- fw_qa
+
+val_ds_names:
+- fw_qa

configs/fw_qa_ctx_qa.yaml

@@ -0,0 +1,50 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00002
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 5
+weight_decay: 0.01
+warmup_ratio: 0.1
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 16
+# LoRA
+lora_r: 8
+lora_dropout: 0.05
+target_modules:
+  - down_proj
+  - up_proj
+
+# data
+train_ds_names:
+- fw_qa
+- ctx_qa
+- pwc
+- hotpot_qa
+
+val_ds_names:
+- fw_qa
+- ctx_qa
+- pwc
+- hotpot_qa

configs/fw_qa_pwc_hotpot.yaml

@@ -0,0 +1,51 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00002
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 5
+weight_decay: 0.01
+warmup_ratio: 0.1
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 16
+# LoRA
+lora_r: 8
+lora_dropout: 0.05
+target_modules:
+  - down_proj
+  - up_proj
+
+# data
+train_ds_names:
+- fw_qa
+- pwc
+- hotpot_qa
+
+val_ds_names:
+- fw_qa
+- pwc
+- hotpot_qa
+
+load_best_model_at_end: true
+metric_for_best_model: eval_pwc_loss

configs/fw_qa_tiny.yaml

@@ -16,14 +16,14 @@ label_names: ["labels"]
 # w/o this the trainer stores logits of all sample in memory...
 # batch_eval_metrics: true
 
-per_device_train_batch_size: 16
-per_device_eval_batch_size: 16
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
 max_val_samples_per_ds: 1000
 # optim: schedule_free_adamw
 
 learning_rate: 0.00002
 # lr_scheduler_type: "constant_with_warmup"
-neftune_noise_alpha: 1
+neftune_noise_alpha: 5
 weight_decay: 0.01
 warmup_ratio: 0.1
 
@@ -35,7 +35,6 @@ lora_dropout: 0.05
 target_modules:
   - down_proj
   - up_proj
-  - gate_proj
 
 # data
 train_ds_names:

hyperlora/data_utils.py

@@ -99,8 +99,6 @@ def filter_long_samples(samples):
 
 
 def add_repeat_prompt_fn(samples):
-    # TODO: maybe repeat only short context (<5K chars)
-    # TODO: also check batch size and mem util
     unique_contexts = set()
     ctxs, prompts, responses = [], [], []
     for ctx, prompt, response in zip(
@@ -109,8 +107,10 @@ def add_repeat_prompt_fn(samples):
         # Only process if the context is not already in the set
         if ctx in unique_contexts:
             continue
-
         unique_contexts.add(ctx)
+        # remove too long contexts
+        if len(ctx) > 3000:
+            continue
         ctxs.append(ctx)
         responses.append(ctx)
         prompts.append("Repeat the text above.")
@@ -143,11 +143,11 @@ def add_negative_prompt_fn(samples):
     ):
         if ctx in unique_contexts:
             continue
+        unique_contexts.add(ctx)
         if any(keyword in prompt for keyword in keywords):
             # Skip samples where the prompt contains any of the specified keywords
             continue
 
-        unique_contexts.add(ctx)
         ctxs.append(ctx)
         prompts.append(prompt)
         responses.append(response)

hyperlora/intx_sft.py

@@ -80,7 +80,7 @@ LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
 def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
     indices = torch.where(valid_masks)
     acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
-    return {"perplexity": acc}
+    return {"per_token_acc": acc}
 
 
 def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
@@ -105,7 +105,7 @@ def compute_perplexity(shift_logits, shift_labels, valid_masks):
     loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
     loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1)
     perplexity = torch.exp(loss).mean().item()
-    return {"per_token_ppl": perplexity}
+    return {"perplexity": perplexity}
 
 
 class Evaluator:
@@ -369,8 +369,8 @@ def main():
         probabilities=[l / total_len for l in train_ds_len],
     )
 
-    val_train_indices = np.random.permutation(len(train_ds))[:500]
-    val_ds["train"] = train_ds.select(val_train_indices)
+    # val_train_indices = np.random.permutation(len(train_ds))[:500]
+    # val_ds["train"] = train_ds.select(val_train_indices)
 
     test_ds = dict()
     if "test" in tokenized_ds: