Batch upload part 2

.gitattributes

@@ -51,3 +51,5 @@ assets/ominicontrol_art/DistractedBoyfriend.webp filter=lfs diff=lfs merge=lfs -
 assets/ominicontrol_art/PulpFiction.jpg filter=lfs diff=lfs merge=lfs -text
 assets/ominicontrol_art/breakingbad.jpg filter=lfs diff=lfs merge=lfs -text
 assets/ominicontrol_art/oiiai.png filter=lfs diff=lfs merge=lfs -text
+nl_tasks/data/MetaMathQA-40K/MetaMathQA-40K.json filter=lfs diff=lfs merge=lfs -text
+nl_tasks/data/MetaMathQA/MetaMathQA-395K.json filter=lfs diff=lfs merge=lfs -text

examples/subject.ipynb

@@ -0,0 +1,216 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.chdir(\"..\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from diffusers.pipelines import FluxPipeline\n",
+    "from PIL import Image\n",
+    "\n",
+    "from omini.pipeline.flux_omini import Condition, generate, seed_everything"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe = FluxPipeline.from_pretrained(\n",
+    "    \"black-forest-labs/FLUX.1-schnell\", torch_dtype=torch.bfloat16\n",
+    ")\n",
+    "pipe = pipe.to(\"cuda\")\n",
+    "pipe.load_lora_weights(\n",
+    "    \"Yuanshi/OminiControl\",\n",
+    "    weight_name=f\"omini/subject_512.safetensors\",\n",
+    "    adapter_name=\"subject\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/penguin.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "# For this model, the position_delta is (0, 32).\n",
+    "# For more details of position_delta, please refer to:\n",
+    "# https://github.com/Yuanshi9815/OminiControl/issues/89#issuecomment-2827080344\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, 32))\n",
+    "\n",
+    "prompt = \"On Christmas evening, on a crowded sidewalk, this item sits on the road, covered in snow and wearing a Christmas hat.\"\n",
+    "\n",
+    "\n",
+    "seed_everything(0)\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=512,\n",
+    "    width=512,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024, 512))\n",
+    "concat_image.paste(image, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/tshirt.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, 32))\n",
+    "\n",
+    "prompt = \"On the beach, a lady sits under a beach umbrella. She's wearing this shirt and has a big smile on her face, with her surfboard hehind her. The sun is setting in the background. The sky is a beautiful shade of orange and purple.\"\n",
+    "\n",
+    "\n",
+    "seed_everything()\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=512,\n",
+    "    width=512,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024, 512))\n",
+    "concat_image.paste(condition.condition, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/rc_car.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, 32))\n",
+    "\n",
+    "prompt = \"A film style shot. On the moon, this item drives across the moon surface. The background is that Earth looms large in the foreground.\"\n",
+    "\n",
+    "seed_everything()\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=512,\n",
+    "    width=512,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024, 512))\n",
+    "concat_image.paste(condition.condition, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/clock.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, 32))\n",
+    "\n",
+    "prompt = \"In a Bauhaus style room, this item is placed on a shiny glass table, with a vase of flowers next to it. In the afternoon sun, the shadows of the blinds are cast on the wall.\"\n",
+    "\n",
+    "seed_everything()\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=512,\n",
+    "    width=512,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024, 512))\n",
+    "concat_image.paste(condition.condition, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/oranges.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, 32))\n",
+    "\n",
+    "prompt = \"A very close up view of this item. It is placed on a wooden table. The background is a dark room, the TV is on, and the screen is showing a cooking show.\"\n",
+    "\n",
+    "seed_everything()\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=512,\n",
+    "    width=512,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024, 512))\n",
+    "concat_image.paste(condition.condition, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/subject_1024.ipynb

@@ -0,0 +1,216 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.chdir(\"..\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from diffusers.pipelines import FluxPipeline\n",
+    "from PIL import Image\n",
+    "\n",
+    "from omini.pipeline.flux_omini import Condition, generate, seed_everything"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe = FluxPipeline.from_pretrained(\n",
+    "    \"black-forest-labs/FLUX.1-schnell\", torch_dtype=torch.bfloat16\n",
+    ")\n",
+    "pipe = pipe.to(\"cuda\")\n",
+    "pipe.load_lora_weights(\n",
+    "    \"Yuanshi/OminiControl\",\n",
+    "    weight_name=f\"omini/subject_1024_beta.safetensors\",\n",
+    "    adapter_name=\"subject\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/penguin.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "# For this model, the position_delta is (0, -32).\n",
+    "# For more details of position_delta, please refer to:\n",
+    "# https://github.com/Yuanshi9815/OminiControl/issues/89#issuecomment-2827080344\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, -32))\n",
+    "\n",
+    "prompt = \"On Christmas evening, on a crowded sidewalk, this item sits on the road, covered in snow and wearing a Christmas hat.\"\n",
+    "\n",
+    "\n",
+    "seed_everything(0)\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=1024,\n",
+    "    width=1024,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024 + 512, 1024))\n",
+    "concat_image.paste(image, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/tshirt.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, -32))\n",
+    "\n",
+    "prompt = \"On the beach, a lady sits under a beach umbrella. She's wearing this shirt and has a big smile on her face, with her surfboard hehind her. The sun is setting in the background. The sky is a beautiful shade of orange and purple.\"\n",
+    "\n",
+    "\n",
+    "seed_everything(0)\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=1024,\n",
+    "    width=1024,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024 + 512, 1024))\n",
+    "concat_image.paste(image, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/rc_car.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, -32))\n",
+    "\n",
+    "prompt = \"A film style shot. On the moon, this item drives across the moon surface. The background is that Earth looms large in the foreground.\"\n",
+    "\n",
+    "seed_everything()\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=1024,\n",
+    "    width=1024,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024 + 512, 1024))\n",
+    "concat_image.paste(image, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/clock.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, -32))\n",
+    "\n",
+    "prompt = \"In a Bauhaus style room, this item is placed on a shiny glass table, with a vase of flowers next to it. In the afternoon sun, the shadows of the blinds are cast on the wall.\"\n",
+    "\n",
+    "seed_everything(0)\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=1024,\n",
+    "    width=1024,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024 + 512, 1024))\n",
+    "concat_image.paste(image, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = Image.open(\"assets/oranges.jpg\").convert(\"RGB\").resize((512, 512))\n",
+    "\n",
+    "condition = Condition(image, \"subject\", position_delta=(0, -32))\n",
+    "\n",
+    "prompt = \"A very close up view of this item. It is placed on a wooden table. The background is a dark room, the TV is on, and the screen is showing a cooking show.\"\n",
+    "\n",
+    "seed_everything()\n",
+    "\n",
+    "result_img = generate(\n",
+    "    pipe,\n",
+    "    prompt=prompt,\n",
+    "    conditions=[condition],\n",
+    "    num_inference_steps=8,\n",
+    "    height=1024,\n",
+    "    width=1024,\n",
+    ").images[0]\n",
+    "\n",
+    "concat_image = Image.new(\"RGB\", (1024 + 512, 1024))\n",
+    "concat_image.paste(image, (0, 0))\n",
+    "concat_image.paste(result_img, (512, 0))\n",
+    "concat_image"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

nl_tasks/README.md

@@ -0,0 +1,45 @@
+Dynamo
+export TORCH_COMPILE_DISABLE=1
+unset TORCH_COMPILE_DISABLE
+echo $TORCH_COMPILE_DISABLE
+
+
+Untracked files:
+  (use "git add <file>..." to include in what will be committed)
+        nl_tasks/README.md
+        nl_tasks/config/commonsense_opt.yaml
+        nl_tasks/config/glue.yaml
+        nl_tasks/config/math395.yaml
+        nl_tasks/data/MetaMathQA/
+        nl_tasks/data/gsm8k_infer.py
+        nl_tasks/exp100/
+        nl_tasks/exp395/
+        nl_tasks/exp_init/
+        nl_tasks/expsBOFT/
+        nl_tasks/expsOFT/
+        nl_tasks/repro.sh
+        nl_tasks/run_all/
+        nl_tasks/run_exps/
+        nl_tasks/scripts/copy train_cms_reasoning.sh
+        nl_tasks/scripts/inference.sh
+        nl_tasks/scripts/merge_100k.sh
+        nl_tasks/scripts/merge_math.sh
+        nl_tasks/scripts/peft_merge.sh
+        nl_tasks/scripts/train_100math.sh
+        nl_tasks/scripts/train_initn40k.sh
+        nl_tasks/scripts/train_math.sh
+        nl_tasks/src/ft_mathQ.py
+        nl_tasks/src/peft_merge.py
+        nl_tasks/src/testLlama.py
+        nl_tasks/testLlama.sh
+        nl_tasks/training_metrics_bs8.json
+        nlu/1Mgrid/
+        nlu/_scripts_/
+        nlu/glue22_exp/
+        nlu/glue_exp00/
+        nlu/glue_test/
+        nlu/seeds/
+        nlu/src/test.py
+        nlu/test.sh
+        nlu/training_metrics_bs8.json
+

nl_tasks/config/commonsense.yaml

@@ -0,0 +1,44 @@
+
+model:
+  model_name: meta-llama/Llama-2-7b-hf #facebook/opt-125m #meta-llama/Llama-2-7b-hf #"openai-community/gpt2" #EleutherAI/pythia-160m #Qwen/Qwen2.5-0.5B
+  # model_name: facebook/opt-125m
+  # adapter_path: "./run_all/exnr15/ft2"
+  # adapter_path: './exp_init/run_ex01/ft2'
+  data_collator_mode: 'dynamic'
+
+rotation_adapter_config:
+  r: 4
+  num_rotations: 4
+  # target_modules: ["q_proj", "v_proj", "v_proj", "o_proj", "gate_proj","up_proj","down_proj"]
+  target_modules: ["q_proj", "v_proj",]
+
+data:
+  dataset_name: 'math'
+  split_ratio: 0.025
+  # path: "./data/gsm8k_test.jsonl"
+  path: ./data/MetaMathQA-40K/MetaMathQA-40K.json
+  # path: ./data/MetaMathQA/MetaMathQA-395K.json
+  dataset_split: train
+  # dataset_field: [question, answer]
+  dataset_field: [query, response]
+
+trainer_args:
+  learning_rate: 2e-4
+  # eval_strategy: steps
+  per_device_train_batch_size: 32
+  per_device_eval_batch_size: 64
+  # accumulate_grad_batches: 1
+  # save_steps: 1000
+  gradient_checkpointing: False # (Turn off for faster training)
+  output_dir: "./run_exps"
+  # save_path: "runs"
+
+  report_to: wandb
+  logging_steps: 25
+  # eval_steps: 100
+  #dataloader_num_workers: 4
+
+  num_train_epochs: 2.0
+  # max_steps: -1
+
+# device: 'cuda'

nl_tasks/config/commonsense_opt.yaml

@@ -0,0 +1,32 @@
+
+model:
+  model_name: facebook/opt-125m #"openai-community/gpt2" #EleutherAI/pythia-160m #Qwen/Qwen2.5-0.5B
+  # adapter_path: "./nl_tasks/run_exps/ft2"
+  data_collator_mode: 'dynamic'
+
+rotation_adapter_config:
+  r: 4
+  num_rotations: 2
+  target_modules: ["q_proj", "v_proj"]
+
+data:
+  dataset_name: 'math'
+  # path: "./nl_tasks/data/MetaMathQA-40K" #MetaMathQA-40K.json"
+  path: "./data/gsm8k_test.jsonl"
+  dataset_split: train[:200]
+  dataset_field: [question, answer]
+
+trainer_args:
+  learning_rate: 2e-4
+  # accumulate_grad_batches: 1
+  # dataloader_workers: 5
+  # save_interval: 1000
+  # sample_interval: 100
+  # max_steps: -1
+  gradient_checkpointing: False # (Turn off for faster training)
+  output_dir: "./run_exps"
+  # save_path: "runs"
+
+  max_steps: 40
+
+# device: 'cuda'

nl_tasks/config/glue.yaml

@@ -0,0 +1,48 @@
+
+model:
+  model_name: microsoft/deberta-v3-base  #facebook/opt-125m #meta-llama/Llama-2-7b-hf #"openai-community/gpt2" #EleutherAI/pythia-160m #Qwen/Qwen2.5-0.5B
+  # model_name: facebook/opt-125m
+  # adapter_path: "./run_all/exnr15/ft2"
+  # adapter_path: './run_all/run_exps9/ft2'
+  # adapter_path: "./exp395/run_ex07/ft2"
+  data_collator_mode: 'dynamic'
+
+rotation_adapter_config:
+  r: 5
+  num_rotations: 1
+  # target_modules: ["q_proj", "v_proj", "v_proj", "o_proj", "gate_proj","up_proj","down_proj"]
+  target_modules: ["query_proj", "value_proj", "key_proj", 'attention.output.dense', 'intermediate.dense', 'output.dense']
+  task_type: "SEQ_CLS"
+
+data:
+  dataset_name: 'math'
+  split_ratio: 0.00258
+  # path: "./data/gsm8k_test.jsonl"
+  # path: ./data/MetaMathQA-40K/MetaMathQA-40K.json
+  path: ./data/MetaMathQA/MetaMathQA-395K.json
+  dataset_split: train[:100000]
+  # dataset_field: [question, answer]
+  dataset_field: [query, response]
+
+trainer_args:
+  learning_rate: 2e-4
+  warmup_ratio: 0.01
+  # eval_strategy: steps
+  per_device_train_batch_size: 32
+  per_device_eval_batch_size: 64
+  # accumulate_grad_batches: 1
+  # save_steps: 1000
+  gradient_checkpointing: False # (Turn off for faster training)
+  output_dir: "./exps/run_exps"
+  # save_path: "runs"
+
+  # report_to: wandb
+  logging_steps: 200
+  # eval_steps: 1000
+  #dataloader_num_workers: 4
+
+  num_train_epochs: 2.0
+  # max_steps: 21
+  # torch_compile: False
+
+# device: 'cuda'

nl_tasks/config/math395.yaml

@@ -0,0 +1,46 @@
+
+model:
+  model_name: meta-llama/Llama-2-7b-hf  #facebook/opt-125m #meta-llama/Llama-2-7b-hf #"openai-community/gpt2" #EleutherAI/pythia-160m #Qwen/Qwen2.5-0.5B
+  # model_name: facebook/opt-125m
+  # adapter_path: "./run_all/exnr15/ft2"
+  # adapter_path: './run_all/run_exps9/ft2'
+  # adapter_path: "./exp395/run_ex07/ft2"
+  data_collator_mode: 'dynamic'
+
+rotation_adapter_config:
+  r: 16
+  num_rotations: 1
+  # target_modules: ["q_proj", "v_proj", "v_proj", "o_proj", "gate_proj","up_proj","down_proj"]
+  target_modules: ["q_proj", "v_proj",]
+
+data:
+  dataset_name: 'math'
+  split_ratio: 0.00258
+  # path: "./data/gsm8k_test.jsonl"
+  # path: ./data/MetaMathQA-40K/MetaMathQA-40K.json
+  path: ./data/MetaMathQA/MetaMathQA-395K.json
+  dataset_split: train[:100000]
+  # dataset_field: [question, answer]
+  dataset_field: [query, response]
+
+trainer_args:
+  learning_rate: 2e-4
+  warmup_ratio: 0.01
+  # eval_strategy: steps
+  per_device_train_batch_size: 32
+  per_device_eval_batch_size: 64
+  # accumulate_grad_batches: 1
+  # save_steps: 1000
+  gradient_checkpointing: False # (Turn off for faster training)
+  output_dir: "./exps/run_exps"
+  # save_path: "runs"
+
+  report_to: wandb
+  logging_steps: 200
+  # eval_steps: 1000
+  #dataloader_num_workers: 4
+
+  num_train_epochs: 2.0
+  # max_steps: 21
+
+# device: 'cuda'

nl_tasks/data/MetaMathQA-40K/MetaMathQA-40K.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c884f10e8aa1229a6e73a6bba2c9134ee0c7b7de92a02a7b8c9459085a59e117
+size 31076207

nl_tasks/data/MetaMathQA/MetaMathQA-395K.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb39a5d8c05c042ece92eae37dfd5ea414a5979df2bf3ad3b86411bef8205725
+size 395626321

nl_tasks/environment.yaml

@@ -0,0 +1,55 @@
+# environment.yml
+name: allm # The name of the environment
+
+channels: # The conda channels to search for packages
+  # - pytorch
+  - conda-forge
+  # - dnachun
+  # - anaconda
+channel_priority: strict
+
+dependencies:
+  # Packages to install with conda
+  - python=3.11.3
+  #- pytorch-cuda=12.4
+  #- pytorch
+  # - numpy
+  - transformers>=4.55
+  - einops
+  - jaxtyping
+
+  - tensorboard
+  - omegaconf
+  - accelerate
+  - peft
+
+  - wandb
+
+  - scipy
+  - pandas
+  - matplotlib
+  - scikit-image
+  - scikit-learn
+  - joblib
+  - pillow
+  - datasets
+  ## NO - huggingface_hub
+  - tqdm
+  - nltk
+  - future
+  
+
+  - defusedxml
+  - ipdb
+  - torchinfo
+
+
+
+  - timm
+  - graphviz #anaconda::graphviz
+  - dnachun::torchviz
+  - pip:
+  #  - draccus
+     - fraction
+     - vllm
+    

nl_tasks/exps/run_ex01/trainer_state.json

@@ -0,0 +1,914 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 2.0,
+  "eval_steps": 100,
+  "global_step": 2438,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.020508613617719443,
+      "grad_norm": 0.06690643727779388,
+      "learning_rate": 4.918032786885246e-06,
+      "loss": 0.751,
+      "step": 25
+    },
+    {
+      "epoch": 0.04101722723543889,
+      "grad_norm": 0.23132523894309998,
+      "learning_rate": 1.0040983606557377e-05,
+      "loss": 0.7344,
+      "step": 50
+    },
+    {
+      "epoch": 0.06152584085315833,
+      "grad_norm": 0.2116735428571701,
+      "learning_rate": 1.5163934426229509e-05,
+      "loss": 0.6404,
+      "step": 75
+    },
+    {
+      "epoch": 0.08203445447087777,
+      "grad_norm": 0.1672675907611847,
+      "learning_rate": 2.028688524590164e-05,
+      "loss": 0.486,
+      "step": 100
+    },
+    {
+      "epoch": 0.08203445447087777,
+      "eval_loss": 0.4427674114704132,
+      "eval_runtime": 19.6288,
+      "eval_samples_per_second": 50.945,
+      "eval_steps_per_second": 0.815,
+      "step": 100
+    },
+    {
+      "epoch": 0.10254306808859721,
+      "grad_norm": 0.16888108849525452,
+      "learning_rate": 2.540983606557377e-05,
+      "loss": 0.4407,
+      "step": 125
+    },
+    {
+      "epoch": 0.12305168170631665,
+      "grad_norm": 0.17033565044403076,
+      "learning_rate": 3.05327868852459e-05,
+      "loss": 0.4031,
+      "step": 150
+    },
+    {
+      "epoch": 0.1435602953240361,
+      "grad_norm": 0.194916769862175,
+      "learning_rate": 3.5655737704918037e-05,
+      "loss": 0.3787,
+      "step": 175
+    },
+    {
+      "epoch": 0.16406890894175555,
+      "grad_norm": 0.29443657398223877,
+      "learning_rate": 4.077868852459016e-05,
+      "loss": 0.3769,
+      "step": 200
+    },
+    {
+      "epoch": 0.16406890894175555,
+      "eval_loss": 0.3537510335445404,
+      "eval_runtime": 19.4681,
+      "eval_samples_per_second": 51.366,
+      "eval_steps_per_second": 0.822,
+      "step": 200
+    },
+    {
+      "epoch": 0.184577522559475,
+      "grad_norm": 0.2323056161403656,
+      "learning_rate": 4.59016393442623e-05,
+      "loss": 0.3658,
+      "step": 225
+    },
+    {
+      "epoch": 0.20508613617719443,
+      "grad_norm": 0.239767923951149,
+      "learning_rate": 4.999935927058032e-05,
+      "loss": 0.3402,
+      "step": 250
+    },
+    {
+      "epoch": 0.22559474979491387,
+      "grad_norm": 0.21633633971214294,
+      "learning_rate": 4.997693718919013e-05,
+      "loss": 0.3342,
+      "step": 275
+    },
+    {
+      "epoch": 0.2461033634126333,
+      "grad_norm": 0.23770427703857422,
+      "learning_rate": 4.992251147198466e-05,
+      "loss": 0.3366,
+      "step": 300
+    },
+    {
+      "epoch": 0.2461033634126333,
+      "eval_loss": 0.3253124952316284,
+      "eval_runtime": 19.4632,
+      "eval_samples_per_second": 51.379,
+      "eval_steps_per_second": 0.822,
+      "step": 300
+    },
+    {
+      "epoch": 0.2666119770303528,
+      "grad_norm": 0.2381717562675476,
+      "learning_rate": 4.98361518561306e-05,
+      "loss": 0.3379,
+      "step": 325
+    },
+    {
+      "epoch": 0.2871205906480722,
+      "grad_norm": 0.2646787166595459,
+      "learning_rate": 4.971796899657632e-05,
+      "loss": 0.329,
+      "step": 350
+    },
+    {
+      "epoch": 0.30762920426579166,
+      "grad_norm": 0.3449915647506714,
+      "learning_rate": 4.9568114324266624e-05,
+      "loss": 0.3406,
+      "step": 375
+    },
+    {
+      "epoch": 0.3281378178835111,
+      "grad_norm": 0.3168472647666931,
+      "learning_rate": 4.938677985211011e-05,
+      "loss": 0.3227,
+      "step": 400
+    },
+    {
+      "epoch": 0.3281378178835111,
+      "eval_loss": 0.3124091625213623,
+      "eval_runtime": 19.4632,
+      "eval_samples_per_second": 51.379,
+      "eval_steps_per_second": 0.822,
+      "step": 400
+    },
+    {
+      "epoch": 0.34864643150123054,
+      "grad_norm": 0.24941708147525787,
+      "learning_rate": 4.9174197928947795e-05,
+      "loss": 0.3301,
+      "step": 425
+    },
+    {
+      "epoch": 0.36915504511895,
+      "grad_norm": 0.25401854515075684,
+      "learning_rate": 4.8930640941838104e-05,
+      "loss": 0.3267,
+      "step": 450
+    },
+    {
+      "epoch": 0.3896636587366694,
+      "grad_norm": 0.2830125689506531,
+      "learning_rate": 4.86564209670399e-05,
+      "loss": 0.3159,
+      "step": 475
+    },
+    {
+      "epoch": 0.41017227235438886,
+      "grad_norm": 0.26643845438957214,
+      "learning_rate": 4.835188937014059e-05,
+      "loss": 0.3006,
+      "step": 500
+    },
+    {
+      "epoch": 0.41017227235438886,
+      "eval_loss": 0.30365315079689026,
+      "eval_runtime": 19.5051,
+      "eval_samples_per_second": 51.269,
+      "eval_steps_per_second": 0.82,
+      "step": 500
+    },
+    {
+      "epoch": 0.4306808859721083,
+      "grad_norm": 0.2782154977321625,
+      "learning_rate": 4.801743635584168e-05,
+      "loss": 0.3015,
+      "step": 525
+    },
+    {
+      "epoch": 0.45118949958982774,
+      "grad_norm": 0.23164071142673492,
+      "learning_rate": 4.7653490467978906e-05,
+      "loss": 0.3095,
+      "step": 550
+    },
+    {
+      "epoch": 0.4716981132075472,
+      "grad_norm": 0.282122403383255,
+      "learning_rate": 4.726051804041709e-05,
+      "loss": 0.3049,
+      "step": 575
+    },
+    {
+      "epoch": 0.4922067268252666,
+      "grad_norm": 0.2644311487674713,
+      "learning_rate": 4.683902259952387e-05,
+      "loss": 0.3213,
+      "step": 600
+    },
+    {
+      "epoch": 0.4922067268252666,
+      "eval_loss": 0.29857054352760315,
+      "eval_runtime": 19.4613,
+      "eval_samples_per_second": 51.384,
+      "eval_steps_per_second": 0.822,
+      "step": 600
+    },
+    {
+      "epoch": 0.5127153404429861,
+      "grad_norm": 0.2529069185256958,
+      "learning_rate": 4.638954421898746e-05,
+      "loss": 0.3001,
+      "step": 625
+    },
+    {
+      "epoch": 0.5332239540607056,
+      "grad_norm": 0.2736290395259857,
+      "learning_rate": 4.5912658827805425e-05,
+      "loss": 0.2916,
+      "step": 650
+    },
+    {
+      "epoch": 0.5537325676784249,
+      "grad_norm": 0.2555182874202728,
+      "learning_rate": 4.5408977472331005e-05,
+      "loss": 0.3052,
+      "step": 675
+    },
+    {
+      "epoch": 0.5742411812961444,
+      "grad_norm": 0.30381742119789124,
+      "learning_rate": 4.48791455333227e-05,
+      "loss": 0.3036,
+      "step": 700
+    },
+    {
+      "epoch": 0.5742411812961444,
+      "eval_loss": 0.292624831199646,
+      "eval_runtime": 19.4487,
+      "eval_samples_per_second": 51.417,
+      "eval_steps_per_second": 0.823,
+      "step": 700
+    },
+    {
+      "epoch": 0.5947497949138638,
+      "grad_norm": 0.26833590865135193,
+      "learning_rate": 4.432384189900008e-05,
+      "loss": 0.3023,
+      "step": 725
+    },
+    {
+      "epoch": 0.6152584085315833,
+      "grad_norm": 0.263784259557724,
+      "learning_rate": 4.3743778095165764e-05,
+      "loss": 0.3016,
+      "step": 750
+    },
+    {
+      "epoch": 0.6357670221493027,
+      "grad_norm": 0.30153971910476685,
+      "learning_rate": 4.313969737350775e-05,
+      "loss": 0.2984,
+      "step": 775
+    },
+    {
+      "epoch": 0.6562756357670222,
+      "grad_norm": 0.27196648716926575,
+      "learning_rate": 4.251237375925071e-05,
+      "loss": 0.3034,
+      "step": 800
+    },
+    {
+      "epoch": 0.6562756357670222,
+      "eval_loss": 0.28956037759780884,
+      "eval_runtime": 19.4528,
+      "eval_samples_per_second": 51.407,
+      "eval_steps_per_second": 0.823,
+      "step": 800
+    },
+    {
+      "epoch": 0.6767842493847416,
+      "grad_norm": 0.25102221965789795,
+      "learning_rate": 4.186261105937612e-05,
+      "loss": 0.2961,
+      "step": 825
+    },
+    {
+      "epoch": 0.6972928630024611,
+      "grad_norm": 0.30812397599220276,
+      "learning_rate": 4.1191241832682364e-05,
+      "loss": 0.2995,
+      "step": 850
+    },
+    {
+      "epoch": 0.7178014766201805,
+      "grad_norm": 0.2563766837120056,
+      "learning_rate": 4.049912632300421e-05,
+      "loss": 0.2878,
+      "step": 875
+    },
+    {
+      "epoch": 0.7383100902379,
+      "grad_norm": 0.3218742311000824,
+      "learning_rate": 3.978715135695881e-05,
+      "loss": 0.296,
+      "step": 900
+    },
+    {
+      "epoch": 0.7383100902379,
+      "eval_loss": 0.2854253053665161,
+      "eval_runtime": 19.4687,
+      "eval_samples_per_second": 51.365,
+      "eval_steps_per_second": 0.822,
+      "step": 900
+    },
+    {
+      "epoch": 0.7588187038556193,
+      "grad_norm": 0.27102115750312805,
+      "learning_rate": 3.905622920763031e-05,
+      "loss": 0.2944,
+      "step": 925
+    },
+    {
+      "epoch": 0.7793273174733388,
+      "grad_norm": 0.276564359664917,
+      "learning_rate": 3.83072964256494e-05,
+      "loss": 0.2847,
+      "step": 950
+    },
+    {
+      "epoch": 0.7998359310910582,
+      "grad_norm": 0.3700260519981384,
+      "learning_rate": 3.7541312639165145e-05,
+      "loss": 0.2877,
+      "step": 975
+    },
+    {
+      "epoch": 0.8203445447087777,
+      "grad_norm": 0.2763277292251587,
+      "learning_rate": 3.675925932424715e-05,
+      "loss": 0.2819,
+      "step": 1000
+    },
+    {
+      "epoch": 0.8203445447087777,
+      "eval_loss": 0.2817630469799042,
+      "eval_runtime": 19.4695,
+      "eval_samples_per_second": 51.362,
+      "eval_steps_per_second": 0.822,
+      "step": 1000
+    },
+    {
+      "epoch": 0.8408531583264971,
+      "grad_norm": 0.3116626739501953,
+      "learning_rate": 3.596213854729328e-05,
+      "loss": 0.2855,
+      "step": 1025
+    },
+    {
+      "epoch": 0.8613617719442166,
+      "grad_norm": 0.27529171109199524,
+      "learning_rate": 3.515097168105444e-05,
+      "loss": 0.2847,
+      "step": 1050
+    },
+    {
+      "epoch": 0.881870385561936,
+      "grad_norm": 0.3135339021682739,
+      "learning_rate": 3.4326798095921656e-05,
+      "loss": 0.2875,
+      "step": 1075
+    },
+    {
+      "epoch": 0.9023789991796555,
+      "grad_norm": 0.33244284987449646,
+      "learning_rate": 3.349067382815217e-05,
+      "loss": 0.2885,
+      "step": 1100
+    },
+    {
+      "epoch": 0.9023789991796555,
+      "eval_loss": 0.27884459495544434,
+      "eval_runtime": 19.4401,
+      "eval_samples_per_second": 51.44,
+      "eval_steps_per_second": 0.823,
+      "step": 1100
+    },
+    {
+      "epoch": 0.9228876127973749,
+      "grad_norm": 0.2761545479297638,
+      "learning_rate": 3.264367022674124e-05,
+      "loss": 0.2857,
+      "step": 1125
+    },
+    {
+      "epoch": 0.9433962264150944,
+      "grad_norm": 0.24953879415988922,
+      "learning_rate": 3.1786872580673214e-05,
+      "loss": 0.2832,
+      "step": 1150
+    },
+    {
+      "epoch": 0.9639048400328137,
+      "grad_norm": 0.32240793108940125,
+      "learning_rate": 3.09213787283109e-05,
+      "loss": 0.2909,
+      "step": 1175
+    },
+    {
+      "epoch": 0.9844134536505332,
+      "grad_norm": 0.2800053358078003,
+      "learning_rate": 3.004829765070516e-05,
+      "loss": 0.297,
+      "step": 1200
+    },
+    {
+      "epoch": 0.9844134536505332,
+      "eval_loss": 0.2768399119377136,
+      "eval_runtime": 19.4646,
+      "eval_samples_per_second": 51.375,
+      "eval_steps_per_second": 0.822,
+      "step": 1200
+    },
+    {
+      "epoch": 1.0049220672682526,
+      "grad_norm": 0.27479347586631775,
+      "learning_rate": 2.916874805062701e-05,
+      "loss": 0.2817,
+      "step": 1225
+    },
+    {
+      "epoch": 1.0254306808859721,
+      "grad_norm": 0.28057900071144104,
+      "learning_rate": 2.828385691914301e-05,
+      "loss": 0.2821,
+      "step": 1250
+    },
+    {
+      "epoch": 1.0459392945036916,
+      "grad_norm": 0.2822401225566864,
+      "learning_rate": 2.7394758091570664e-05,
+      "loss": 0.2831,
+      "step": 1275
+    },
+    {
+      "epoch": 1.066447908121411,
+      "grad_norm": 0.2816685438156128,
+      "learning_rate": 2.6502590794664073e-05,
+      "loss": 0.285,
+      "step": 1300
+    },
+    {
+      "epoch": 1.066447908121411,
+      "eval_loss": 0.2752681076526642,
+      "eval_runtime": 19.4336,
+      "eval_samples_per_second": 51.457,
+      "eval_steps_per_second": 0.823,
+      "step": 1300
+    },
+    {
+      "epoch": 1.0869565217391304,
+      "grad_norm": 0.2750077545642853,
+      "learning_rate": 2.560849818689141e-05,
+      "loss": 0.2829,
+      "step": 1325
+    },
+    {
+      "epoch": 1.1074651353568499,
+      "grad_norm": 0.29034796357154846,
+      "learning_rate": 2.471362589367452e-05,
+      "loss": 0.2727,
+      "step": 1350
+    },
+    {
+      "epoch": 1.1279737489745694,
+      "grad_norm": 0.33246326446533203,
+      "learning_rate": 2.3819120539467663e-05,
+      "loss": 0.2806,
+      "step": 1375
+    },
+    {
+      "epoch": 1.1484823625922886,
+      "grad_norm": 0.26562532782554626,
+      "learning_rate": 2.2926128278556052e-05,
+      "loss": 0.2666,
+      "step": 1400
+    },
+    {
+      "epoch": 1.1484823625922886,
+      "eval_loss": 0.27384424209594727,
+      "eval_runtime": 19.4305,
+      "eval_samples_per_second": 51.466,
+      "eval_steps_per_second": 0.823,
+      "step": 1400
+    },
+    {
+      "epoch": 1.1689909762100081,
+      "grad_norm": 0.3040076494216919,
+      "learning_rate": 2.2035793326456883e-05,
+      "loss": 0.2799,
+      "step": 1425
+    },
+    {
+      "epoch": 1.1894995898277276,
+      "grad_norm": 0.38336554169654846,
+      "learning_rate": 2.1149256493804576e-05,
+      "loss": 0.2858,
+      "step": 1450
+    },
+    {
+      "epoch": 1.2100082034454471,
+      "grad_norm": 0.28841930627822876,
+      "learning_rate": 2.0267653724598747e-05,
+      "loss": 0.2777,
+      "step": 1475
+    },
+    {
+      "epoch": 1.2305168170631666,
+      "grad_norm": 0.3314826786518097,
+      "learning_rate": 1.9392114640687985e-05,
+      "loss": 0.2884,
+      "step": 1500
+    },
+    {
+      "epoch": 1.2305168170631666,
+      "eval_loss": 0.27183249592781067,
+      "eval_runtime": 19.4549,
+      "eval_samples_per_second": 51.401,
+      "eval_steps_per_second": 0.822,
+      "step": 1500
+    },
+    {
+      "epoch": 1.251025430680886,
+      "grad_norm": 0.35494905710220337,
+      "learning_rate": 1.8523761094354304e-05,
+      "loss": 0.2833,
+      "step": 1525
+    },
+    {
+      "epoch": 1.2715340442986054,
+      "grad_norm": 0.30966347455978394,
+      "learning_rate": 1.7663705730853012e-05,
+      "loss": 0.276,
+      "step": 1550
+    },
+    {
+      "epoch": 1.2920426579163249,
+      "grad_norm": 0.2722432613372803,
+      "learning_rate": 1.6813050562749778e-05,
+      "loss": 0.2775,
+      "step": 1575
+    },
+    {
+      "epoch": 1.3125512715340442,
+      "grad_norm": 0.2977288067340851,
+      "learning_rate": 1.5972885557881666e-05,
+      "loss": 0.269,
+      "step": 1600
+    },
+    {
+      "epoch": 1.3125512715340442,
+      "eval_loss": 0.27014729380607605,
+      "eval_runtime": 19.4752,
+      "eval_samples_per_second": 51.347,
+      "eval_steps_per_second": 0.822,
+      "step": 1600
+    },
+    {
+      "epoch": 1.3330598851517639,
+      "grad_norm": 0.2877044081687927,
+      "learning_rate": 1.5144287242751378e-05,
+      "loss": 0.2727,
+      "step": 1625
+    },
+    {
+      "epoch": 1.3535684987694832,
+      "grad_norm": 0.3222461938858032,
+      "learning_rate": 1.4328317323144284e-05,
+      "loss": 0.2742,
+      "step": 1650
+    },
+    {
+      "epoch": 1.3740771123872026,
+      "grad_norm": 0.3473854959011078,
+      "learning_rate": 1.3526021323735626e-05,
+      "loss": 0.2724,
+      "step": 1675
+    },
+    {
+      "epoch": 1.3945857260049221,
+      "grad_norm": 0.3066307306289673,
+      "learning_rate": 1.2738427248431028e-05,
+      "loss": 0.2696,
+      "step": 1700
+    },
+    {
+      "epoch": 1.3945857260049221,
+      "eval_loss": 0.26889586448669434,
+      "eval_runtime": 19.4809,
+      "eval_samples_per_second": 51.332,
+      "eval_steps_per_second": 0.821,
+      "step": 1700
+    },
+    {
+      "epoch": 1.4150943396226414,
+      "grad_norm": 0.2990802228450775,
+      "learning_rate": 1.1966544263156865e-05,
+      "loss": 0.269,
+      "step": 1725
+    },
+    {
+      "epoch": 1.435602953240361,
+      "grad_norm": 0.29301875829696655,
+      "learning_rate": 1.1211361402788226e-05,
+      "loss": 0.2681,
+      "step": 1750
+    },
+    {
+      "epoch": 1.4561115668580804,
+      "grad_norm": 0.32831257581710815,
+      "learning_rate": 1.047384630387131e-05,
+      "loss": 0.2771,
+      "step": 1775
+    },
+    {
+      "epoch": 1.4766201804758,
+      "grad_norm": 0.3254742920398712,
+      "learning_rate": 9.75494396476423e-06,
+      "loss": 0.2644,
+      "step": 1800
+    },
+    {
+      "epoch": 1.4766201804758,
+      "eval_loss": 0.2682200074195862,
+      "eval_runtime": 19.4711,
+      "eval_samples_per_second": 51.358,
+      "eval_steps_per_second": 0.822,
+      "step": 1800
+    },
+    {
+      "epoch": 1.4971287940935194,
+      "grad_norm": 0.29671165347099304,
+      "learning_rate": 9.05557553478459e-06,
+      "loss": 0.2658,
+      "step": 1825
+    },
+    {
+      "epoch": 1.5176374077112387,
+      "grad_norm": 0.29359179735183716,
+      "learning_rate": 8.376637133915558e-06,
+      "loss": 0.2676,
+      "step": 1850
+    },
+    {
+      "epoch": 1.5381460213289582,
+      "grad_norm": 0.3013196587562561,
+      "learning_rate": 7.718998704582739e-06,
+      "loss": 0.2708,
+      "step": 1875
+    },
+    {
+      "epoch": 1.5586546349466777,
+      "grad_norm": 0.33183860778808594,
+      "learning_rate": 7.0835028969730185e-06,
+      "loss": 0.2727,
+      "step": 1900
+    },
+    {
+      "epoch": 1.5586546349466777,
+      "eval_loss": 0.26749491691589355,
+      "eval_runtime": 19.4495,
+      "eval_samples_per_second": 51.415,
+      "eval_steps_per_second": 0.823,
+      "step": 1900
+    },
+    {
+      "epoch": 1.579163248564397,
+      "grad_norm": 0.3297445774078369,
+      "learning_rate": 6.470963989323764e-06,
+      "loss": 0.2792,
+      "step": 1925
+    },
+    {
+      "epoch": 1.5996718621821167,
+      "grad_norm": 0.33076536655426025,
+      "learning_rate": 5.8821668445656924e-06,
+      "loss": 0.2749,
+      "step": 1950
+    },
+    {
+      "epoch": 1.620180475799836,
+      "grad_norm": 0.290751188993454,
+      "learning_rate": 5.317865904656497e-06,
+      "loss": 0.2653,
+      "step": 1975
+    },
+    {
+      "epoch": 1.6406890894175554,
+      "grad_norm": 0.30254530906677246,
+      "learning_rate": 4.778784223893601e-06,
+      "loss": 0.2767,
+      "step": 2000
+    },
+    {
+      "epoch": 1.6406890894175554,
+      "eval_loss": 0.26705402135849,
+      "eval_runtime": 19.451,
+      "eval_samples_per_second": 51.411,
+      "eval_steps_per_second": 0.823,
+      "step": 2000
+    },
+    {
+      "epoch": 1.661197703035275,
+      "grad_norm": 0.30978214740753174,
+      "learning_rate": 4.265612542444827e-06,
+      "loss": 0.2661,
+      "step": 2025
+    },
+    {
+      "epoch": 1.6817063166529942,
+      "grad_norm": 0.3454751968383789,
+      "learning_rate": 3.7790084012840453e-06,
+      "loss": 0.2717,
+      "step": 2050
+    },
+    {
+      "epoch": 1.7022149302707137,
+      "grad_norm": 0.3721317946910858,
+      "learning_rate": 3.319595299665873e-06,
+      "loss": 0.2767,
+      "step": 2075
+    },
+    {
+      "epoch": 1.7227235438884332,
+      "grad_norm": 0.33564668893814087,
+      "learning_rate": 2.8879618962189326e-06,
+      "loss": 0.2614,
+      "step": 2100
+    },
+    {
+      "epoch": 1.7227235438884332,
+      "eval_loss": 0.26674506068229675,
+      "eval_runtime": 19.4651,
+      "eval_samples_per_second": 51.374,
+      "eval_steps_per_second": 0.822,
+      "step": 2100
+    },
+    {
+      "epoch": 1.7432321575061525,
+      "grad_norm": 0.31175196170806885,
+      "learning_rate": 2.484661254681381e-06,
+      "loss": 0.2689,
+      "step": 2125
+    },
+    {
+      "epoch": 1.7637407711238722,
+      "grad_norm": 0.27521854639053345,
+      "learning_rate": 2.110210135245147e-06,
+      "loss": 0.266,
+      "step": 2150
+    },
+    {
+      "epoch": 1.7842493847415914,
+      "grad_norm": 0.301544189453125,
+      "learning_rate": 1.765088332416917e-06,
+      "loss": 0.2746,
+      "step": 2175
+    },
+    {
+      "epoch": 1.804757998359311,
+      "grad_norm": 0.29790034890174866,
+      "learning_rate": 1.4497380602442378e-06,
+      "loss": 0.2671,
+      "step": 2200
+    },
+    {
+      "epoch": 1.804757998359311,
+      "eval_loss": 0.2663249373435974,
+      "eval_runtime": 19.4746,
+      "eval_samples_per_second": 51.349,
+      "eval_steps_per_second": 0.822,
+      "step": 2200
+    },
+    {
+      "epoch": 1.8252666119770304,
+      "grad_norm": 0.3353317677974701,
+      "learning_rate": 1.1645633856944977e-06,
+      "loss": 0.2693,
+      "step": 2225
+    },
+    {
+      "epoch": 1.8457752255947497,
+      "grad_norm": 0.2950851023197174,
+      "learning_rate": 9.099297109128407e-07,
+      "loss": 0.2704,
+      "step": 2250
+    },
+    {
+      "epoch": 1.8662838392124692,
+      "grad_norm": 0.3311655521392822,
+      "learning_rate": 6.861633050223526e-07,
+      "loss": 0.2757,
+      "step": 2275
+    },
+    {
+      "epoch": 1.8867924528301887,
+      "grad_norm": 0.3037591576576233,
+      "learning_rate": 4.935508860664601e-07,
+      "loss": 0.2684,
+      "step": 2300
+    },
+    {
+      "epoch": 1.8867924528301887,
+      "eval_loss": 0.26620855927467346,
+      "eval_runtime": 19.477,
+      "eval_samples_per_second": 51.343,
+      "eval_steps_per_second": 0.821,
+      "step": 2300
+    },
+    {
+      "epoch": 1.907301066447908,
+      "grad_norm": 0.3440220057964325,
+      "learning_rate": 3.323392536292436e-07,
+      "loss": 0.2773,
+      "step": 2325
+    },
+    {
+      "epoch": 1.9278096800656277,
+      "grad_norm": 0.33441346883773804,
+      "learning_rate": 2.0273497260433204e-07,
+      "loss": 0.2582,
+      "step": 2350
+    },
+    {
+      "epoch": 1.948318293683347,
+      "grad_norm": 0.28934305906295776,
+      "learning_rate": 1.0490410851763943e-07,
+      "loss": 0.2653,
+      "step": 2375
+    },
+    {
+      "epoch": 1.9688269073010665,
+      "grad_norm": 0.304415225982666,
+      "learning_rate": 3.8972014743038356e-08,
+      "loss": 0.268,
+      "step": 2400
+    },
+    {
+      "epoch": 1.9688269073010665,
+      "eval_loss": 0.26615387201309204,
+      "eval_runtime": 19.4583,
+      "eval_samples_per_second": 51.392,
+      "eval_steps_per_second": 0.822,
+      "step": 2400
+    },
+    {
+      "epoch": 1.989335520918786,
+      "grad_norm": 0.2703385055065155,
+      "learning_rate": 5.023171883647426e-09,
+      "loss": 0.263,
+      "step": 2425
+    },
+    {
+      "epoch": 2.0,
+      "step": 2438,
+      "total_flos": 1.58523627405312e+18,
+      "train_loss": 0.3074496070120157,
+      "train_runtime": 3055.5986,
+      "train_samples_per_second": 25.527,
+      "train_steps_per_second": 0.798
+    }
+  ],
+  "logging_steps": 25,
+  "max_steps": 2438,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 2,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 1.58523627405312e+18,
+  "train_batch_size": 32,
+  "trial_name": null,
+  "trial_params": null
+}

nl_tasks/repro.sh

@@ -0,0 +1,87 @@
+export OMINI_CONFIG=./config/commonsense.yaml
+
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include"
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+export WANDB_PROJECT="Llama2_7B_FT_Math40k_2"
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+date +"%F %T"
+TEXT=("oft" "boft" "loco" "hra")
+
+# --run_text "$text"  --dynamo_backend no 
+export ACCELERATE_DYNAMO_BACKEND="no"
+#       --trainer_args.max_steps=81 \
+
+accelerate launch --dynamo_backend no --main_process_port 41353 -m src.testLlama \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./expsBOFT/seed44/" --trainer_args.learning_rate=8e-4 \
+      --run_text "boft" --trainer_args.per_device_train_batch_size 32 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.gradient_accumulation_steps 2 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+      --trainer_args.eval_strategy '"no"' \
+      --trainer_args.load_best_model_at_end False \
+      --trainer_args.save_strategy '"no"' \
+      --trainer_args.logging_step 50 \
+      --trainer_args.report_to none --trainer_args.warmup_steps 100 \
+      --seed 44
+date +"%F %T"
+
+accelerate launch --dynamo_backend no --main_process_port 41353 -m src.testLlama \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./expsBOFT/seed43/" --trainer_args.learning_rate=8e-4 \
+      --run_text "boft" --trainer_args.per_device_train_batch_size 32 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.gradient_accumulation_steps 2 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+      --trainer_args.eval_strategy '"no"' \
+      --trainer_args.load_best_model_at_end False \
+      --trainer_args.save_strategy '"no"' \
+      --trainer_args.logging_step 50 \
+      --trainer_args.report_to none --trainer_args.warmup_steps 100 \
+      --seed 43
+date +"%F %T"
+
+
+accelerate launch --main_process_port 41353 -m src.testLlama \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./expsOFT/seed43/" --trainer_args.learning_rate=8e-4 \
+      --run_text "oft" --trainer_args.per_device_train_batch_size 64 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.gradient_accumulation_steps 1 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+      --trainer_args.eval_strategy '"no"' \
+      --trainer_args.load_best_model_at_end False \
+      --trainer_args.save_strategy '"no"' \
+      --trainer_args.logging_step 50 \
+      --trainer_args.report_to none --trainer_args.warmup_steps 100 \
+      --seed 43
+
+date +"%F %T"
+
+accelerate launch --main_process_port 41353 -m src.testLlama \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./expsOFT/seed44/" --trainer_args.learning_rate=8e-4 \
+      --run_text "oft" --trainer_args.per_device_train_batch_size 64 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.gradient_accumulation_steps 1 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+      --trainer_args.eval_strategy '"no"' \
+      --trainer_args.load_best_model_at_end False \
+      --trainer_args.save_strategy '"no"' \
+      --trainer_args.logging_step 50 \
+      --trainer_args.report_to none --trainer_args.warmup_steps 100 \
+      --seed 44
+
+date +"%F %T"
+

nl_tasks/rpeft/__init__.py

@@ -0,0 +1,43 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all.
+
+# coding=utf-8
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__version__ = "0.0.1"
+
+from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PEFT_TYPE_TO_CONFIG_MAPPING,\
+      get_peft_config, get_peft_model #, PEFT_TYPE_TO_TUNER_MAPPING
+
+
+from .rotation import (
+    RotationConfig,
+    RotationTuner,
+)
+from .utils import (
+    TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING,
+    PeftConfig,
+    PeftType,
+    PromptLearningConfig,
+    TaskType,
+    bloom_model_postprocess_past_key_value,
+    get_peft_model_state_dict,
+    prepare_model_for_int8_training,
+    set_peft_model_state_dict,
+    shift_tokens_right,
+)
+
+from .peft_model import PeftModel

nl_tasks/rpeft/mapping.py

@@ -0,0 +1,273 @@
+# coding=utf-8
+# Original License:
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .peft_model import (
+    PeftModel,
+    PeftModelForCausalLM,
+    PeftModelForSeq2SeqLM,
+    PeftModelForSequenceClassification,
+    PeftModelForTokenClassification,
+)
+from .rotation import RotationConfig, RotationTuner
+from .utils import PromptLearningConfig
+
+from transformers import PreTrainedModel
+
+MODEL_TYPE_TO_PEFT_MODEL_MAPPING = {
+    "SEQ_CLS": PeftModelForSequenceClassification,
+    "SEQ_2_SEQ_LM": PeftModelForSeq2SeqLM,
+    "CAUSAL_LM": PeftModelForCausalLM,
+    "TOKEN_CLS": PeftModelForTokenClassification,
+}
+
+PEFT_TYPE_TO_CONFIG_MAPPING: dict = {
+    "ROTATION": RotationConfig,
+}
+
+PEFT_TYPE_TO_TUNER_MAPPING: dict = {
+    "ROTATION": RotationTuner
+}
+
+TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
+    "t5": ["q", "v"],
+    "mt5": ["q", "v"],
+    "bart": ["q_proj", "v_proj"],
+    "gpt2": ["c_attn"],
+    "bloom": ["query_key_value"],
+    "blip-2": ["q", "v", "q_proj", "v_proj"],
+    "opt": ["q_proj", "v_proj"],
+    "gptj": ["q_proj", "v_proj"],
+    "gpt_neox": ["query_key_value"],
+    "gpt_neo": ["q_proj", "v_proj"],
+    "bert": ["query", "value"],
+    "roberta": ["query", "value"],
+    "xlm-roberta": ["query", "value"],
+    "electra": ["query", "value"],
+    "deberta-v2": ["query_proj", "value_proj"],
+    "deberta": ["in_proj"],
+    "layoutlm": ["query", "value"],
+    "llama": ["q_proj", "v_proj"],
+    "chatglm": ["query_key_value"],
+    "gpt_bigcode": ["c_attn"],
+    "mpt": ["Wqkv"],
+    "RefinedWebModel": ["query_key_value"],
+    "RefinedWeb": ["query_key_value"],
+    "falcon": ["query_key_value"],
+    "btlm": ["c_proj", "c_attn"],
+    "codegen": ["qkv_proj"],
+    "mistral": ["q_proj", "v_proj"],
+    "mixtral": ["q_proj", "v_proj"],
+    "stablelm": ["q_proj", "v_proj"],
+    "phi": ["q_proj", "v_proj", "fc1", "fc2"],
+    "gemma": ["q_proj", "v_proj"],
+}
+
+
+
+def get_peft_config(config_dict):
+    """
+    Returns a Peft config object from a dictionary.
+
+    Args:
+        config_dict (`Dict[str, Any]`): Dictionary containing the configuration parameters.
+    """
+
+    return PEFT_TYPE_TO_CONFIG_MAPPING[config_dict["peft_type"]](**config_dict)
+
+
+def _prepare_prompt_learning_config(peft_config, model_config):
+    if peft_config.num_layers is None:
+        if "num_hidden_layers" in model_config:
+            num_layers = model_config["num_hidden_layers"]
+        elif "num_layers" in model_config:
+            num_layers = model_config["num_layers"]
+        elif "n_layer" in model_config:
+            num_layers = model_config["n_layer"]
+        else:
+            raise ValueError("Please specify `num_layers` in `peft_config`")
+        peft_config.num_layers = num_layers
+
+    if peft_config.token_dim is None:
+        if "hidden_size" in model_config:
+            token_dim = model_config["hidden_size"]
+        elif "n_embd" in model_config:
+            token_dim = model_config["n_embd"]
+        elif "d_model" in model_config:
+            token_dim = model_config["d_model"]
+        else:
+            raise ValueError("Please specify `token_dim` in `peft_config`")
+        peft_config.token_dim = token_dim
+
+    if peft_config.num_attention_heads is None:
+        if "num_attention_heads" in model_config:
+            num_attention_heads = model_config["num_attention_heads"]
+        elif "n_head" in model_config:
+            num_attention_heads = model_config["n_head"]
+        elif "num_heads" in model_config:
+            num_attention_heads = model_config["num_heads"]
+        elif "encoder_attention_heads" in model_config:
+            num_attention_heads = model_config["encoder_attention_heads"]
+        else:
+            raise ValueError("Please specify `num_attention_heads` in `peft_config`")
+        peft_config.num_attention_heads = num_attention_heads
+
+    if getattr(peft_config, "encoder_hidden_size", None) is None:
+        setattr(peft_config, "encoder_hidden_size", token_dim)
+
+    return peft_config
+
+
+def _prepare_lora_config(peft_config, model_config):
+    if peft_config.target_modules is None:
+        if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
+            raise ValueError("Please specify `target_modules` in `peft_config`")
+        peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]]
+    if len(peft_config.target_modules) == 1:
+        peft_config.fan_in_fan_out = True
+        peft_config.enable_lora = [True, False, True]
+    if peft_config.inference_mode:
+        peft_config.merge_weights = True
+    return peft_config
+    
+
+
+def get_peft_model(model, peft_config,
+                    adapter_name: str = "default"):
+    """
+    Returns a Peft model object from a model and a config.
+
+    Args:
+        model ([`transformers.PreTrainedModel`]): Model to be wrapped.
+        peft_config ([`PeftConfig`]): Configuration object containing the parameters of the Peft model.
+    """
+    model_config = model.config.to_dict()
+    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)
+    if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys():
+        if peft_config.peft_type == "LORA" or "QUANTA":
+            peft_config = _prepare_lora_config(peft_config, model_config)
+            return PeftModel(model, peft_config)
+    if not isinstance(peft_config, PromptLearningConfig):
+        if peft_config.peft_type == "LORA" or "QUANTA":
+            peft_config = _prepare_lora_config(peft_config, model_config)
+    else:
+        peft_config = _prepare_prompt_learning_config(peft_config, model_config)
+    # assert False
+    return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](
+        model,
+        peft_config,
+        adapter_name=adapter_name,
+    )
+
+
+# def get_peft_model(
+#     model: PreTrainedModel,
+#     peft_config,
+#     adapter_name: str = "default",
+#     mixed: bool = False,
+#     autocast_adapter_dtype: bool = True,
+#     revision: Optional[str] = None,
+#     low_cpu_mem_usage: bool = False,
+# ) -> PeftModel | PeftMixedModel:
+#     """
+#     Returns a Peft model object from a model and a config, where the model will be modified in-place.
+
+#     Args:
+#         model ([`transformers.PreTrainedModel`]):
+#             Model to be wrapped.
+#         peft_config ([`PeftConfig`]):
+#             Configuration object containing the parameters of the Peft model.
+#         adapter_name (`str`, `optional`, defaults to `"default"`):
+#             The name of the adapter to be injected, if not provided, the default adapter name is used ("default").
+#         mixed (`bool`, `optional`, defaults to `False`):
+#             Whether to allow mixing different (compatible) adapter types.
+#         autocast_adapter_dtype (`bool`, *optional*):
+#             Whether to autocast the adapter dtype. Defaults to `True`. Right now, this will only cast adapter weights
+#             using float16 or bfloat16 to float32, as this is typically required for stable training, and only affect
+#             select PEFT tuners.
+#         revision (`str`, `optional`, defaults to `main`):
+#             The revision of the base model. If this isn't set, the saved peft model will load the `main` revision for
+#             the base model
+#         low_cpu_mem_usage (`bool`, `optional`, defaults to `False`):
+#             Create empty adapter weights on meta device. Useful to speed up the loading process. Leave this setting as
+#             False if you intend on training the model, unless the adapter weights will be replaced by different weights
+#             before training starts.
+#     """
+#     model_config = BaseTuner.get_model_config(model)
+#     old_name = peft_config.base_model_name_or_path
+#     new_name = model.__dict__.get("name_or_path", None)
+#     peft_config.base_model_name_or_path = new_name
+
+#     # Especially in notebook environments there could be a case that a user wants to experiment with different
+#     # configuration values. However, it is likely that there won't be any changes for new configs on an already
+#     # initialized PEFT model. The best we can do is warn the user about it.
+#     if any(isinstance(module, BaseTunerLayer) for module in model.modules()):
+#         warnings.warn(
+#             "You are trying to modify a model with PEFT for a second time. If you want to reload the model with a "
+#             "different config, make sure to call `.unload()` before."
+#         )
+
+#     if (old_name is not None) and (old_name != new_name):
+#         warnings.warn(
+#             f"The PEFT config's `base_model_name_or_path` was renamed from '{old_name}' to '{new_name}'. "
+#             "Please ensure that the correct base model is loaded when loading this checkpoint."
+#         )
+
+#     if revision is not None:
+#         if peft_config.revision is not None and peft_config.revision != revision:
+#             warnings.warn(
+#                 f"peft config has already set base model revision to {peft_config.revision}, overwriting with revision {revision}"
+#             )
+#         peft_config.revision = revision
+
+#     if (
+#         (isinstance(peft_config, PEFT_TYPE_TO_CONFIG_MAPPING["LORA"]))
+#         and (peft_config.init_lora_weights == "eva")
+#         and not low_cpu_mem_usage
+#     ):
+#         warnings.warn(
+#             "lora with eva initialization used with low_cpu_mem_usage=False. "
+#             "Setting low_cpu_mem_usage=True can improve the maximum batch size possible for eva initialization."
+#         )
+
+#     prefix = PEFT_TYPE_TO_PREFIX_MAPPING.get(peft_config.peft_type)
+#     if prefix and adapter_name in prefix:
+#         warnings.warn(
+#             f"Adapter name '{adapter_name}' should not be contained in the prefix '{prefix}'. "
+#             "This may lead to reinitialization of the adapter weights during loading."
+#         )
+
+#     if mixed:
+#         # note: PeftMixedModel does not support autocast_adapter_dtype, so don't pass it
+#         return PeftMixedModel(model, peft_config, adapter_name=adapter_name)
+
+#     # We explicitly exclude prompt learning here since prompt learning is specific to the task and needs special
+#     # handling in the PEFT model's forward method.
+#     if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not peft_config.is_prompt_learning:
+#         return PeftModel(
+#             model,
+#             peft_config,
+#             adapter_name=adapter_name,
+#             autocast_adapter_dtype=autocast_adapter_dtype,
+#             low_cpu_mem_usage=low_cpu_mem_usage,
+#         )
+
+#     return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](
+#         model,
+#         peft_config,
+#         adapter_name=adapter_name,
+#         autocast_adapter_dtype=autocast_adapter_dtype,
+#         low_cpu_mem_usage=low_cpu_mem_usage,
+#     )

nl_tasks/rpeft/peft_model.py

@@ -0,0 +1,922 @@
+# coding=utf-8
+# Original License:
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import os
+import warnings
+from contextlib import contextmanager
+
+import torch
+from accelerate import dispatch_model, infer_auto_device_map
+from accelerate.hooks import AlignDevicesHook, add_hook_to_module, remove_hook_from_submodules
+from accelerate.utils import get_balanced_memory
+from huggingface_hub import hf_hub_download
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput
+from transformers.utils import PushToHubMixin
+
+import packaging.version
+import transformers
+from typing import Any, Literal, Optional, Union
+
+from .rotation import RotationTuner
+from .utils import (
+    TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING,
+    WEIGHTS_NAME,
+    PeftConfig,
+    PeftType,
+    PromptLearningConfig,
+    TaskType,
+    _set_trainable,
+    get_peft_model_state_dict,
+    set_peft_model_state_dict,
+    shift_tokens_right,
+)
+
+class PeftModel(PushToHubMixin, torch.nn.Module):
+    """
+
+    """
+
+    def __init__(self, model, peft_config: PeftConfig, adapter_name: str = "default"):
+        super().__init__()
+        self.peft_config = peft_config
+        self.base_model = model
+        self.config = self.base_model.config
+        self.modules_to_save = None
+        self.active_adapter = adapter_name
+        ##### Diff do nothing with active_adapter
+        if isinstance(self.peft_config, PromptLearningConfig):
+            self._setup_prompt_encoder()
+        else:
+            if self.peft_config.peft_type == PeftType.ROTATION:
+                self.base_model = RotationTuner(model, {adapter_name: peft_config}, adapter_name)
+        if getattr(self.peft_config, "modules_to_save", None) is not None:
+            self.modules_to_save = self.peft_config.modules_to_save
+            _set_trainable(self)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.base_model_torch_dtype = getattr(model, "dtype", None)
+
+    def save_pretrained(self, save_directory, **kwargs):
+        r"""
+        Args:
+        This function saves the adapter model and the adapter configuration files to a directory, so that it can be
+        re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub`
+        method.
+            save_directory (`str`):
+                Directory where the adapter model and configuration files will be saved (will be created if it does not
+                exist).
+            **kwargs:
+                Additional keyword arguments passed along to the `push_to_hub` method.
+        """
+        if os.path.isfile(save_directory):
+            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save only the trainable weights
+        output_state_dict = get_peft_model_state_dict(self, kwargs.get("state_dict", None))
+        torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
+
+        # save the config and change the inference mode to `True`
+        if self.peft_config.base_model_name_or_path is None:
+            self.peft_config.base_model_name_or_path = (
+                self.base_model.__dict__.get("name_or_path", None)
+                if isinstance(self.peft_config, PromptLearningConfig)
+                else self.base_model.model.__dict__.get("name_or_path", None)
+            )
+        inference_mode = self.peft_config.inference_mode
+        self.peft_config.inference_mode = True
+        self.peft_config.save_pretrained(save_directory)
+        self.peft_config.inference_mode = inference_mode
+
+    @classmethod
+    def from_pretrained(cls, model, model_id, is_trainable = False, **kwargs):
+        r"""
+        Args:
+        Instantiate a `LoraModel` from a pretrained Lora configuration and weights.
+            model (`transformers.PreTrainedModel`):
+                The model to be adapted. The model should be initialized with the `from_pretrained` method. from
+                `transformers` library.
+            model_id (`str`):
+                The name of the Lora configuration to use. Can be either:
+                    - A string, the `model id` of a Lora configuration hosted inside a model repo on
+                        huggingface Hub
+                    - A path to a directory containing a Lora configuration file saved using the
+                        `save_pretrained` method, e.g., ``./my_lora_config_directory/``.
+        """
+        from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PEFT_TYPE_TO_CONFIG_MAPPING
+        # load the config
+        config = PEFT_TYPE_TO_CONFIG_MAPPING[PeftConfig.from_pretrained(model_id).peft_type].from_pretrained(model_id)
+        config.inference_mode = not is_trainable
+
+        if getattr(model, "hf_device_map", None) is not None:
+            remove_hook_from_submodules(model)
+
+        if config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys():
+            model = cls(model, config)
+        else:
+            model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config)
+
+        # load weights if any
+        if os.path.exists(os.path.join(model_id, WEIGHTS_NAME)):
+            filename = os.path.join(model_id, WEIGHTS_NAME)
+        else:
+            try:
+                filename = hf_hub_download(model_id, WEIGHTS_NAME)
+            except:  # noqa
+                raise ValueError(
+                    f"Can't find weights for {model_id} in {model_id} or in the Hugging Face Hub. "
+                    f"Please check that the file {WEIGHTS_NAME} is present at {model_id}."
+                )
+
+        adapters_weights = torch.load(
+            filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        # load the weights into the model
+        model = set_peft_model_state_dict(model, adapters_weights)
+        if getattr(model, "hf_device_map", None) is not None:
+            device_map = kwargs.get("device_map", "auto")
+            max_memory = kwargs.get("max_memory", None)
+            no_split_module_classes = model._no_split_modules
+            if device_map != "sequential":
+                max_memory = get_balanced_memory(
+                    model,
+                    max_memory=max_memory,
+                    no_split_module_classes=no_split_module_classes,
+                    low_zero=(device_map == "balanced_low_0"),
+                )
+            if isinstance(device_map, str):
+                device_map = infer_auto_device_map(
+                    model, max_memory=max_memory, no_split_module_classes=no_split_module_classes
+                )
+            model = dispatch_model(model, device_map=device_map)
+            hook = AlignDevicesHook(io_same_device=True)
+            if model.peft_config.peft_type == PeftType.LORA or model.peft_config.peft_type == PeftType.BOTTLENECK \
+                or model.peft_config.peft_type == "ROTATION":
+                add_hook_to_module(model.base_model.model, hook)
+            else:
+                remove_hook_from_submodules(model.prompt_encoder)
+                add_hook_to_module(model.base_model, hook)
+            # if model.peft_config.is_prompt_learning:
+            #     remove_hook_from_submodules(model.prompt_encoder)
+            # add_hook_to_module(model.base_model, hook)
+        return model
+
+    def _setup_prompt_encoder(self):
+        transformer_backbone = None
+        for name, module in self.base_model.named_children():
+            for param in module.parameters():
+                param.requires_grad = False
+            if isinstance(module, PreTrainedModel):
+                # Make sure to freeze Tranformers model
+                if transformer_backbone is None:
+                    transformer_backbone = module
+                    self.transformer_backbone_name = name
+
+        if self.peft_config.num_transformer_submodules is None:
+            self.peft_config.num_transformer_submodules = (
+                2 if self.peft_config.task_type == TaskType.SEQ_2_SEQ_LM else 1
+            )
+
+        for named_param, value in list(transformer_backbone.named_parameters()):
+            if value.shape[0] == self.base_model.config.vocab_size:
+                self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", ""))
+                break
+
+        if self.peft_config.peft_type == PeftType.PROMPT_TUNING:
+            prompt_encoder = PromptEmbedding(self.peft_config, self.word_embeddings)
+        elif self.peft_config.peft_type == PeftType.P_TUNING:
+            prompt_encoder = PromptEncoder(self.peft_config)
+        elif self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            prompt_encoder = PrefixEncoder(self.peft_config)
+        else:
+            raise ValueError("Not supported")
+        self.prompt_encoder = prompt_encoder
+        self.prompt_tokens = torch.arange(
+            self.peft_config.num_virtual_tokens * self.peft_config.num_transformer_submodules
+        ).long()
+
+    def get_prompt_embedding_to_save(self):
+        """
+        Returns the prompt embedding to save when saving the model. Only applicable when `peft_config.peft_type !=
+        PeftType.LORA`.
+        """
+        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.device)
+        if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens]
+        prompt_embeddings = self.prompt_encoder(prompt_tokens)
+        return prompt_embeddings[0].detach().cpu()
+
+    def get_prompt(self, batch_size):
+        """
+        Returns the virtual prompts to use for Peft. Only applicable when `peft_config.peft_type != PeftType.LORA`.
+        """
+        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.device)
+        if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens]
+            if self.peft_config.inference_mode:
+                past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1)
+            else:
+                past_key_values = self.prompt_encoder(prompt_tokens)
+            past_key_values = past_key_values.view(
+                batch_size,
+                self.peft_config.num_virtual_tokens,
+                self.peft_config.num_layers * 2,
+                self.peft_config.num_attention_heads,
+                self.peft_config.token_dim // self.peft_config.num_attention_heads,
+            )
+            if self.peft_config.num_transformer_submodules == 2:
+                past_key_values = torch.cat([past_key_values, past_key_values], dim=2)
+            past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(
+                self.peft_config.num_transformer_submodules * 2
+            )
+            if TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING.get(self.config.model_type, None) is not None:
+                post_process_fn = TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING[self.config.model_type]
+                past_key_values = post_process_fn(past_key_values)
+            return past_key_values
+        else:
+            if self.peft_config.inference_mode:
+                prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1)
+            else:
+                prompts = self.prompt_encoder(prompt_tokens)
+            return prompts
+
+    def print_trainable_parameters(self):
+        """
+        Prints the number of trainable parameters in the model.
+        """
+        trainable_params = 0
+        all_param = 0
+        for _, param in self.named_parameters():
+            num_params = param.numel()
+            # if using DS Zero 3 and the weights are initialized empty
+            if num_params == 0 and hasattr(param, "ds_numel"):
+                num_params = param.ds_numel
+
+            all_param += num_params
+            if param.requires_grad:
+                trainable_params += num_params
+        print(
+            f"trainable params: {trainable_params:,} || all params: {all_param:,} || trainable: {100 * trainable_params / all_param}%"
+        )
+
+    def __getattr__(self, name: str):
+        """Forward missing attributes to the wrapped module."""
+        try:
+            return super().__getattr__(name)  # defer to nn.Module's logic
+        except AttributeError:
+            return getattr(self.base_model, name)
+
+    def forward(self, *args, **kwargs):
+        """
+        Forward pass of the model.
+        """
+        return self.get_base_model()(*args, **kwargs)
+
+    @contextmanager
+    def disable_adapter(self):
+        """
+        Disables the adapter module.
+        """
+        if isinstance(self.peft_config, PromptLearningConfig):
+            old_forward = self.forward
+            self.forward = self.base_model.forward
+        else:
+            self.base_model.disable_adapter_layers()
+        yield
+        if isinstance(self.peft_config, PromptLearningConfig):
+            self.forward = old_forward
+        else:
+            self.base_model.enable_adapter_layers()
+
+    def get_base_model(self):
+        """
+        Returns the base model.
+        """
+        return self.base_model if isinstance(self.peft_config, PromptLearningConfig) else self.base_model.model
+
+
+
+class PeftModelForSequenceClassification(PeftModel):
+    """
+    """
+
+    def __init__(self, model, peft_config: PeftConfig, adapter_name: str = "default"):
+        super().__init__(model, peft_config, adapter_name)
+        self.modules_to_save = ["classifier", "score", "pooler"]
+
+        # for name, _ in self.base_model.named_children():
+        #     if any(module_name in name for module_name in self.modules_to_save):
+        #         self.cls_layer_name = name
+        #         break
+        user_modules = getattr(peft_config, "modules_to_save", None) or []
+        default_modules = ["classifier", "score"]
+        self.modules_to_save = list(set(user_modules + default_modules))
+
+        #from .rotation import RotationTuner # Import để check type
+        if isinstance(self.base_model, RotationTuner):
+            real_model = self.base_model.model
+        else:
+            real_model = self.base_model
+
+        # 3. Tìm tên layer thực tế
+        for name, _ in real_model.named_children():
+            if any(module_name in name for module_name in self.modules_to_save):
+                self.cls_layer_name = name
+
+        # # to make sure classifier layer is trainable
+        _set_trainable(self)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        if "num_items_in_batch" in kwargs:
+            kwargs.pop("num_items_in_batch")
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not isinstance(self.peft_config, PromptLearningConfig):
+            return self.base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        batch_size = input_ids.shape[0]
+        if attention_mask is not None:
+            # concat prompt attention mask
+            prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device)
+            attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+        if kwargs.get("position_ids", None) is not None:
+            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
+            kwargs["position_ids"] = None
+        kwargs.update(
+            {
+                "attention_mask": attention_mask,
+                "labels": labels,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+            }
+        )
+
+        if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            return self._prefix_tuning_forward(input_ids=input_ids, **kwargs)
+        else:
+            if kwargs.get("token_type_ids", None) is not None:
+                kwargs["token_type_ids"] = torch.cat(
+                    (
+                        torch.zeros(batch_size, self.peft_config.num_virtual_tokens).to(self.device),
+                        kwargs["token_type_ids"],
+                    ),
+                    dim=1,
+                ).long()
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+            prompts = self.get_prompt(batch_size=batch_size)
+            prompts = prompts.to(inputs_embeds.dtype)
+            inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
+            return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
+
+    def _prefix_tuning_forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        batch_size = input_ids.shape[0]
+        past_key_values = self.get_prompt(batch_size)
+        fwd_params = list(inspect.signature(self.base_model.forward).parameters.keys())
+        kwargs.update(
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "inputs_embeds": inputs_embeds,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+                "past_key_values": past_key_values,
+            }
+        )
+        if "past_key_values" in fwd_params:
+            return self.base_model(labels=labels, **kwargs)
+        else:
+            transformer_backbone_name = self.base_model.get_submodule(self.transformer_backbone_name)
+            fwd_params = list(inspect.signature(transformer_backbone_name.forward).parameters.keys())
+            if "past_key_values" not in fwd_params:
+                raise ValueError("Model does not support past key values which are required for prefix tuning.")
+            outputs = transformer_backbone_name(**kwargs)
+            pooled_output = outputs[1] if len(outputs) > 1 else outputs[0]
+            if "dropout" in [name for name, _ in list(self.base_model.named_children())]:
+                pooled_output = self.base_model.dropout(pooled_output)
+            logits = self.base_model.get_submodule(self.cls_layer_name)(pooled_output)
+
+            loss = None
+            if labels is not None:
+                if self.config.problem_type is None:
+                    if self.base_model.num_labels == 1:
+                        self.config.problem_type = "regression"
+                    elif self.base_model.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                        self.config.problem_type = "single_label_classification"
+                    else:
+                        self.config.problem_type = "multi_label_classification"
+
+                if self.config.problem_type == "regression":
+                    loss_fct = MSELoss()
+                    if self.base_model.num_labels == 1:
+                        loss = loss_fct(logits.squeeze(), labels.squeeze())
+                    else:
+                        loss = loss_fct(logits, labels)
+                elif self.config.problem_type == "single_label_classification":
+                    loss_fct = CrossEntropyLoss()
+                    loss = loss_fct(logits.view(-1, self.base_model.num_labels), labels.view(-1))
+                elif self.config.problem_type == "multi_label_classification":
+                    loss_fct = BCEWithLogitsLoss()
+                    loss = loss_fct(logits, labels)
+            if not return_dict:
+                output = (logits,) + outputs[2:]
+                return ((loss,) + output) if loss is not None else output
+
+            return SequenceClassifierOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            )
+
+
+class PeftModelForCausalLM(PeftModel):
+    """
+    Peft model for Causal LM
+
+    Args:
+        model ([`PreTrainedModel`]): Base transformer model
+        peft_config ([`PeftConfig`]): Peft config.
+
+
+    Example::
+
+        >>> from transformers import AutoModelForCausalLM >>> from peft_local_tensor import PeftModelForCausalLM, get_peft_config
+        >>> config = {
+                'peft_type': 'PREFIX_TUNING', 'task_type': 'CAUSAL_LM', 'inference_mode': False, 'num_virtual_tokens':
+                20, 'token_dim': 1280, 'num_transformer_submodules': 1, 'num_attention_heads': 20, 'num_layers': 36,
+                'encoder_hidden_size': 1280, 'prefix_projection': False, 'postprocess_past_key_value_function': None
+            }
+        >>> peft_config = get_peft_config(config) >>> model = AutoModelForCausalLM.from_pretrained("gpt2-large") >>>
+        peft_model = PeftModelForCausalLM(model, peft_config) >>> peft_model.print_trainable_parameters() trainable
+        params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544
+    """
+
+    def __init__(self, model, peft_config: PeftConfig, adapter_name: str = "default"):
+        self.prompt_encoder = None #### don't know why
+        self.modules_to_save = None
+        super().__init__(model, peft_config, adapter_name)
+        self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        if not isinstance(self.peft_config, PromptLearningConfig):
+            return self.base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        batch_size = input_ids.shape[0]
+        if attention_mask is not None:
+            # concat prompt attention mask
+            prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device)
+            attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        if kwargs.get("position_ids", None) is not None:
+            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
+            kwargs["position_ids"] = None
+        if kwargs.get("token_type_ids", None) is not None:
+            warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids")
+            kwargs["token_type_ids"] = None
+        kwargs.update(
+            {
+                "attention_mask": attention_mask,
+                "labels": labels,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+            }
+        )
+
+        if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            past_key_values = self.get_prompt(batch_size)
+            return self.base_model(input_ids=input_ids, past_key_values=past_key_values, **kwargs)
+        else:
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+            # concat prompt labels
+            if labels is not None:
+                prefix_labels = torch.full((batch_size, self.peft_config.num_virtual_tokens), -100).to(self.device)
+                kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
+            prompts = self.get_prompt(batch_size=batch_size)
+            prompts = prompts.to(inputs_embeds.dtype)
+            inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
+            return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
+
+    def generate(self, **kwargs):
+        self.base_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation
+        try:
+            if not isinstance(self.peft_config, PromptLearningConfig):
+                outputs = self.base_model.generate(**kwargs)
+            else:
+                if "input_ids" not in kwargs:
+                    raise ValueError("input_ids must be provided for Peft model generation")
+                if kwargs.get("attention_mask", None) is not None:
+                    # concat prompt attention mask
+                    prefix_attention_mask = torch.ones(
+                        kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens
+                    ).to(kwargs["input_ids"].device)
+                    kwargs["attention_mask"] = torch.cat((prefix_attention_mask, kwargs["attention_mask"]), dim=1)
+
+                if kwargs.get("position_ids", None) is not None:
+                    warnings.warn(
+                        "Position ids are not supported for parameter efficient tuning. Ignoring position ids."
+                    )
+                    kwargs["position_ids"] = None
+                if kwargs.get("token_type_ids", None) is not None:
+                    warnings.warn(
+                        "Token type ids are not supported for parameter efficient tuning. Ignoring token type ids"
+                    )
+                    kwargs["token_type_ids"] = None
+
+                outputs = self.base_model.generate(**kwargs)
+        except:
+            self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation
+            raise
+        else:
+            self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation
+            return outputs
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs)
+        if isinstance(self.peft_config, PromptLearningConfig):
+            if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+                prefix_attention_mask = torch.ones(
+                    model_kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens
+                ).to(model_kwargs["input_ids"].device)
+                model_kwargs["attention_mask"] = torch.cat(
+                    (prefix_attention_mask, model_kwargs["attention_mask"]), dim=1
+                )
+
+            if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+                past_key_values = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0])
+                if self.base_model_torch_dtype is not None:
+                    # handle the case for Bloom where it outputs tuple of tuples
+                    if isinstance(past_key_values[0], tuple):
+                        past_key_values = tuple(
+                            tuple(
+                                past_key_value.to(self.base_model_torch_dtype)
+                                for past_key_value in past_key_value_tuple
+                            )
+                            for past_key_value_tuple in past_key_values
+                        )
+                    else:
+                        past_key_values = tuple(
+                            past_key_value.to(self.base_model_torch_dtype) for past_key_value in past_key_values
+                        )
+                        
+                model_kwargs["past_key_values"] = past_key_values
+            else:
+                if model_kwargs["past_key_values"] is None:
+                    inputs_embeds = self.word_embeddings(model_kwargs["input_ids"])
+                    prompts = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0])
+                    prompts = prompts.to(inputs_embeds.dtype)
+                    model_kwargs["inputs_embeds"] = torch.cat((prompts, inputs_embeds), dim=1)
+                    model_kwargs["input_ids"] = None
+
+        return model_kwargs
+
+
+class PeftModelForSeq2SeqLM(PeftModel):
+    """
+
+    """
+
+    def __init__(self, model, peft_config: PeftConfig):
+        super().__init__(model, peft_config)
+        self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation
+        self.base_model_prepare_encoder_decoder_kwargs_for_generation = (
+            self.base_model._prepare_encoder_decoder_kwargs_for_generation
+        )
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        decoder_input_ids=None,
+        decoder_attention_mask=None,
+        decoder_inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        if not isinstance(self.peft_config, PromptLearningConfig):
+            return self.base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                decoder_input_ids=decoder_input_ids,
+                decoder_attention_mask=decoder_attention_mask,
+                decoder_inputs_embeds=decoder_inputs_embeds,
+                labels=labels,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        batch_size = input_ids.shape[0]
+        if decoder_attention_mask is not None:
+            # concat prompt attention mask
+            prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device)
+            decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1)
+
+        if kwargs.get("position_ids", None) is not None:
+            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
+            kwargs["position_ids"] = None
+        if kwargs.get("token_type_ids", None) is not None:
+            warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids")
+            kwargs["token_type_ids"] = None
+        kwargs.update(
+            {
+                "attention_mask": attention_mask,
+                "decoder_attention_mask": decoder_attention_mask,
+                "labels": labels,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+            }
+        )
+
+        if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            past_key_values = self.get_prompt(batch_size)
+            return self.base_model(
+                input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs
+            )
+        else:
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+            if decoder_inputs_embeds is None and decoder_input_ids is None:
+                decoder_input_ids = shift_tokens_right(
+                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
+                )
+                decoder_inputs_embeds = self.word_embeddings(decoder_input_ids)
+
+            if attention_mask is not None:
+                # concat prompt attention mask
+                prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device)
+                kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+            # concat prompt labels
+            if labels is not None:
+                if self.peft_config.num_transformer_submodules == 1:
+                    kwargs["labels"] = labels
+                elif self.peft_config.num_transformer_submodules == 2:
+                    prefix_labels = torch.full((batch_size, self.peft_config.num_virtual_tokens), -100).to(self.device)
+                    kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
+            prompts = self.get_prompt(batch_size=batch_size)
+            prompts = prompts.to(inputs_embeds.dtype)
+            inputs_embeds = torch.cat((prompts[:, : self.peft_config.num_virtual_tokens], inputs_embeds), dim=1)
+            if self.peft_config.num_transformer_submodules == 1:
+                return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
+            elif self.peft_config.num_transformer_submodules == 2:
+                decoder_inputs_embeds = torch.cat(
+                    (prompts[:, self.peft_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1
+                )
+                return self.base_model(
+                    inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs
+                )
+
+    def generate(self, **kwargs):
+        self.base_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation
+        self.base_model._prepare_encoder_decoder_kwargs_for_generation = (
+            self._prepare_encoder_decoder_kwargs_for_generation
+        )
+        try:
+            if not isinstance(self.peft_config, PromptLearningConfig):
+                outputs = self.base_model.generate(**kwargs)
+            else:
+                if "input_ids" not in kwargs:
+                    raise ValueError("input_ids must be provided for Peft model generation")
+                if kwargs.get("position_ids", None) is not None:
+                    warnings.warn(
+                        "Position ids are not supported for parameter efficient tuning. Ignoring position ids."
+                    )
+                    kwargs["position_ids"] = None
+                if kwargs.get("token_type_ids", None) is not None:
+                    warnings.warn(
+                        "Token type ids are not supported for parameter efficient tuning. Ignoring token type ids"
+                    )
+                    kwargs["token_type_ids"] = None
+
+                if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+                    outputs = self.base_model.generate(**kwargs)
+                else:
+                    raise NotImplementedError
+        except:
+            self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation
+            self.base_model._prepare_encoder_decoder_kwargs_for_generation = (
+                self.base_model_prepare_encoder_decoder_kwargs_for_generation
+            )
+            raise
+        else:
+            self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation
+            self.base_model._prepare_encoder_decoder_kwargs_for_generation = (
+                self.base_model_prepare_encoder_decoder_kwargs_for_generation
+            )
+            return outputs
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs)
+        if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            batch_size = model_kwargs["decoder_input_ids"].shape[0]
+            past_key_values = self.get_prompt(batch_size)
+            model_kwargs["past_key_values"] = past_key_values
+        return model_kwargs
+
+
+class PeftModelForTokenClassification(PeftModel):
+    """
+
+    """
+
+    def __init__(self, model, peft_config: PeftConfig):
+        super().__init__(model, peft_config)
+        self.modules_to_save = ["classifier", "score"]
+
+        for name, _ in self.base_model.named_children():
+            if any(module_name in name for module_name in self.modules_to_save):
+                self.cls_layer_name = name
+                break
+
+        # to make sure classifier layer is trainable
+        _set_trainable(self)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not isinstance(self.peft_config, PromptLearningConfig):
+            return self.base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        batch_size = input_ids.shape[0]
+        if attention_mask is not None:
+            # concat prompt attention mask
+            prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device)
+            attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+        if kwargs.get("position_ids", None) is not None:
+            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
+            kwargs["position_ids"] = None
+        kwargs.update(
+            {
+                "attention_mask": attention_mask,
+                "labels": labels,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+            }
+        )
+
+        if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+            return self._prefix_tuning_forward(input_ids=input_ids, **kwargs)
+        else:
+            if kwargs.get("token_type_ids", None) is not None:
+                kwargs["token_type_ids"] = torch.cat(
+                    (
+                        torch.zeros(batch_size, self.peft_config.num_virtual_tokens).to(self.device),
+                        kwargs["token_type_ids"],
+                    ),
+                    dim=1,
+                ).long()
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+            prompts = self.get_prompt(batch_size=batch_size)
+            prompts = prompts.to(inputs_embeds.dtype)
+            inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
+            return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
+
+    def _prefix_tuning_forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        batch_size = input_ids.shape[0]
+        past_key_values = self.get_prompt(batch_size)
+        fwd_params = list(inspect.signature(self.base_model.forward).parameters.keys())
+        kwargs.update(
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "inputs_embeds": inputs_embeds,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+                "past_key_values": past_key_values,
+            }
+        )
+        if "past_key_values" in fwd_params:
+            return self.base_model(labels=labels, **kwargs)
+        else:
+            transformer_backbone_name = self.base_model.get_submodule(self.transformer_backbone_name)
+            fwd_params = list(inspect.signature(transformer_backbone_name.forward).parameters.keys())
+            if "past_key_values" not in fwd_params:
+                raise ValueError("Model does not support past key values which are required for prefix tuning.")
+            outputs = transformer_backbone_name(**kwargs)
+            sequence_output = outputs[0]
+            if "dropout" in [name for name, _ in list(self.base_model.named_children())]:
+                sequence_output = self.base_model.dropout(sequence_output)
+            logits = self.base_model.get_submodule(self.cls_layer_name)(sequence_output)
+
+            loss = None
+            loss = None
+            if labels is not None:
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+            if not return_dict:
+                output = (logits,) + outputs[2:]
+                return ((loss,) + output) if loss is not None else output
+
+            return TokenClassifierOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            )

nl_tasks/rpeft/rotation/__init__.py

@@ -0,0 +1,3 @@
+from .rotation_config import RotationConfig
+from .layer import RotationLayer
+from .model import RotationTuner

nl_tasks/rpeft/rotation/layer.py

@@ -0,0 +1,412 @@
+import torch
+import torch.nn as nn
+from typing import Optional, Set
+
+from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge
+import torch.nn.functional as F
+
+def inverse_2x2(matrices):
+
+    # Extract matrix elements
+    # matrices[..., 0, 0] corresponds to 'a' in [[a, b], [c, d]]
+    a = matrices[..., 0, 0]
+    b = matrices[..., 0, 1]
+    c = matrices[..., 1, 0]
+    d = matrices[..., 1, 1]
+    
+    # Compute determinant
+    det = a * d - b * c
+    
+    # Compute inverse using the formula:
+    # inv = (1/det) * [[d, -b], [-c, a]]
+    inv_det = 1.0 / det
+    
+    # Create output tensor
+    inv_matrices = torch.empty_like(matrices)
+    inv_matrices[..., 0, 0] = d * inv_det
+    inv_matrices[..., 0, 1] = -b * inv_det
+    inv_matrices[..., 1, 0] = -c * inv_det
+    inv_matrices[..., 1, 1] = a * inv_det
+    
+    return inv_matrices
+
+class Rotation(nn.Module):
+    """
+    Rotation layer based on Cayley transformation for parameter-efficient fine-tuning.
+    
+    This layer implements orthogonal fine-tuning through Cayley transformation:
+    h(x) = (I - A)^{-1} (I + A) x
+    
+    where A = XY^T with X = [U; -V] and Y = [V; U]
+    """
+    
+    def __init__(self, r, dim, T=1.0, num_rotations=4, drop_out=0.1):
+        super().__init__()
+        self.r = r
+        self.T = T
+        self.U = nn.Parameter(torch.randn(num_rotations, r, dim) * 0.002, requires_grad=True)
+        self.V = nn.Parameter(torch.randn(num_rotations, r, dim) * 0.0, requires_grad=True)
+        # self.U = nn.Parameter(torch.empty(num_rotations, r, dim), requires_grad=True)
+        # self.V = nn.Parameter(torch.empty(num_rotations, r, dim), requires_grad=True)
+        self.num_rotations = num_rotations
+        self.dropout = nn.Dropout(drop_out) if drop_out > 0 else nn.Identity()
+        self.dim = dim
+        
+        # elf._post_init()
+    # @property
+    # def U(self):
+    #     # Calculate U = [a, b] whenever self.U is accessed
+    #     # This function acts as the 'getter' for self.U
+    #     return torch.cat([self.a, self.b], dim=-1)
+    
+    # @property
+    # def V(self):
+    #     # Calculate V = [b, a] whenever self.V is accessed
+    #     # This function acts as the 'getter' for self.V
+    #     return torch.cat([self.b, self.a], dim=-1)
+
+    # def _post_init(self):
+    #     import torch.nn.init as init
+    #     import math
+    #     # init.kaiming_uniform_(self.U, a=math.sqrt(1), mode='fan_out')
+    #     init.normal_(self.U, 0, 1e-2)
+    #     with torch.no_grad():
+    #         self.V.data.copy_(self.U.data)
+
+    def forward(self, x):
+        """
+        Apply Cayley transformation to input x.
+        
+        A = XY^T where X = [U; -V], Y = [V; U]
+        Cayley transformation: h(x) = (I - A)^{-1} (I + A) x
+        
+        Uses Woodbury identity for efficient computation:
+        (I - XY^T)^{-1} = I + X (I - Y^T X)^{-1} Y^T
+        
+        Args:
+            x: Input tensor of shape (..., dim)
+            
+        Returns:
+            Transformed tensor of shape (..., dim)
+        """
+        x_dtype = x.dtype
+        
+        x = self.dropout(x) # NLU tasks do not use dropout
+        X = torch.cat([self.U, -self.V], dim=1)  # Shape: (num_rotations, 2r, dim)
+        Y = torch.cat([self.V, self.U], dim=1) * self.T   # Shape: (num_rotations, 2r, dim)
+
+        Y_T_X = torch.matmul(Y, X.transpose(1, 2))  # Shape: (num_rotations, 2r, 2r)
+        # I_2r = torch.eye(2 * self.r, device=x.device, dtype=x.dtype).repeat(self.num_rotations, 1, 1)
+        I_2r = torch.eye(2 * self.r, device=x.device, dtype=x.dtype).unsqueeze(0)
+        I_minus_YX = I_2r - Y_T_X
+        
+        if self.r == 1:
+            I_minus_YX_inv = inverse_2x2(I_minus_YX)
+        else:
+            # make it float32
+            I_minus_YX = I_minus_YX.to(torch.float32)
+            I_minus_YX_inv = torch.linalg.inv(I_minus_YX)  # Shape: (num_rotations, 2r, 2r)
+            I_minus_YX_inv = I_minus_YX_inv.to(x_dtype)
+        
+        # Yx = torch.einsum("...d,nrd->...nr", x, Y)   # Shape: (batch*seq_len, num_rotations, 2r)
+        
+        input_shape = x.shape
+        x_flat = x.reshape(-1, self.dim) # Shape: (B, dim)
+        Y_reshape = Y.reshape(-1, self.dim) # Shape: (nr, d)
+        Yx2_flat = F.linear(x_flat, Y_reshape)
+        Yx2 = Yx2_flat.view(-1, self.num_rotations, 2*self.r)
+        # is_close = torch.allclose(Yx.view(-1, self.num_rotations, 2*self.r), Yx2, atol=1e-5, rtol=1e-4)
+        # if is_close:
+        #     pass
+        #     # print("✅ SUCCESS 11: The optimized code produces identical results!")
+        # else:
+        #     print("❌ FAILURE: The results diverge.")
+
+        ###
+        # n of (r,r) @ n  of (r,1) -> n of r
+        # I_minus_YX_inv_Yx = torch.einsum("nrr,...nr->...nr", I_minus_YX_inv, Yx)
+        # I_minus_YX_inv_Yx = torch.einsum("...qr,...r->...q", I_minus_YX_inv, Yx)
+        
+        Yx2_expanded = Yx2.unsqueeze(-1)
+        I_minus_YX_inv_ex = I_minus_YX_inv.unsqueeze(0)
+        I_minus_YX_inv_Yx2 = I_minus_YX_inv_ex @ Yx2_expanded
+        I_minus_YX_inv_Yx2 = I_minus_YX_inv_Yx2.squeeze(-1)
+        # is_close = torch.allclose(I_minus_YX_inv_Yx.view(-1, self.num_rotations, 2*self.r), I_minus_YX_inv_Yx2, atol=1e-4, rtol=1e-3)
+        # if is_close:
+        #     pass
+        #     #print("✅ SUCCESS 22: The optimized code produces identical results!")
+        # else:
+        #     print("❌ FAILURE: The results diverge.")
+        #     exit()
+        
+        # n of (r,) @ n of (r,d)
+        # second_term = torch.einsum("...nr,nrd->...nd", I_minus_YX_inv_Yx, X)  # Shape: (batch*seq_len, num_rotations, dim)
+        # second_term = torch.einsum("...r, ...rd->...d", I_minus_YX_inv_Yx, X)
+        
+        
+        # I_minus_YX_inv_Yx_ex = I_minus_YX_inv_Yx2.unsqueeze(-2)
+        # X_ex = X.unsqueeze(0)
+        # second_term2 = I_minus_YX_inv_Yx_ex @ X_ex
+        # second_term2 = second_term2.squeeze(-2)
+        # is_close = torch.allclose(second_term, second_term2, atol=1e-5, rtol=1e-4)
+        # if is_close:
+        #     pass
+        #     # print("✅ SUCCESS 33: The optimized code produces identical results!")
+        # else:
+        #     print("❌ FAILURE: The results diverge.")
+
+        # second_term = second_term.sum(dim=-2)  # Sum over rotations
+        
+        coeffs_flat = I_minus_YX_inv_Yx2.reshape(-1, self.num_rotations * 2 * self.r) # (batch*len, 2n r)
+        X_flat = X.reshape(-1, self.dim) #(N*2r, dim)
+        second_term3 = torch.matmul(coeffs_flat, X_flat)
+        # is_close = torch.allclose(second_term.view(-1, self.dim), second_term3, atol=1e-5, rtol=1e-4)
+        # if is_close:
+        #     pass
+        #     # print("✅ SUCCESS 44: The optimized code produces identical results!")
+        # else:
+        #     print("❌ FAILURE: The results diverge.")
+        
+        # output = x + 2 * second_term  # Shape: (batch*seq_len, dim)
+        output = x_flat + 2 * second_term3 # (batch*seq_len, dim)
+        
+        # return output.to(x_dtype)
+        return output.view(*input_shape)
+    
+    def get_delta_weight(self):
+        """
+        Compute the delta weight matrix induced by the rotation layer.
+        
+        Returns:
+            Delta weight matrix of shape (dim, dim)
+        """
+
+        X = torch.cat([self.U, -self.V], dim=1)  # Shape: (num_rotations, 2r, dim)
+        Y = torch.cat([self.V, self.U], dim=1) * self.T   # Shape: (num_rotations, 2r, dim)
+
+        Y_T_X = torch.matmul(Y, X.transpose(1, 2))  # Shape: (num_rotations, 2r, 2r)
+        I_2r = torch.eye(2 * self.r, device=X.device, dtype=X.dtype).repeat(self.num_rotations, 1, 1)
+        I_minus_YX = I_2r - Y_T_X
+        
+        if self.r == 1:
+            I_minus_YX_inv = inverse_2x2(I_minus_YX)
+            I_minus_YX_inv_Y = torch.einsum("nRr,nrd->nRd", I_minus_YX_inv, Y) # Shape: (num_rotations, 2r, dim)
+            # I_minus_YX_inv_Y = torch.einsum("...rr,...rd->...rd", I_minus_YX_inv, Y) ## reproduce
+        else:
+            I_minus_YX_inv_Y = torch.linalg.solve(I_minus_YX.to(torch.float32), Y.to(torch.float32))  # Shape: (num_rotations, 2r, dim)
+
+            # I_minus_YX_inv = torch.linalg.inv(I_minus_YX)
+            # I_minus_YX_inv_Y = torch.einsum("...rr,...rd->...rd", I_minus_YX_inv, Y)  ## reproduce
+
+            I_minus_YX_inv_Y = I_minus_YX_inv_Y.to(X.dtype)
+        
+        # I_minus_YX_float = I_minus_YX.float()
+        # I_minus_YX_inv = torch.linalg.inv(I_minus_YX_float)  # Shape: (num_rotations, 2r, 2r)
+        # I_minus_YX_inv = I_minus_YX_inv.to(X.dtype)
+        
+        
+        # I_minus_YX_inv_Y = torch.einsum("nRr,nrd->nRd", I_minus_YX_inv, Y) # Shape: (num_rotations, 2r, dim)
+
+
+        second_term = torch.einsum("nrd,nrD->ndD", X, I_minus_YX_inv_Y)  # Shape: (num_rotations, dim, dim)
+        second_term = second_term.sum(dim=0)
+        total_delta_weight = 2 * second_term
+
+        return total_delta_weight
+    
+
+class RotationLayer(BaseTunerLayer):
+    """
+    Adapter-like wrapper that attaches Rotation modules to a base linear layer.
+    """
+
+    adapter_layer_names: tuple[str, ...] = ("rotation",)
+    other_param_names: tuple[str, ...] = ("r", "T", "num_rotations", "scaling")
+
+    def __init__(self, base_layer: nn.Module, **kwargs):
+        # Let BaseTunerLayer do its init (it usually subclasses nn.Module)
+        super().__init__()
+        # store base layer and adapter containers
+        self.base_layer = base_layer
+        self.rotation = nn.ModuleDict()  # mapping adapter_name -> Rotation module
+        self.scaling={}  # default scaling per adapter
+        self._adapter_config = {}  # store r, T, num_rotations per adapter
+
+        # flags (exposed in a simple way)
+        self._disable_adapters = False
+        self.merged_adapters: list[str] = []
+        self._cast_input_dtype_enabled = True
+        self.kwargs = kwargs
+
+        if isinstance(base_layer, nn.Linear):
+            self.in_features = base_layer.in_features
+            self.out_features = base_layer.out_features
+        else:
+            raise NotImplementedError("RotationLayer only supports nn.Linear base layers for now.")
+
+    @property
+    def _available_adapters(self) -> set[str]:
+        return set(self.rotation.keys())
+
+    @property
+    def disable_adapters(self) -> bool:
+        return self._disable_adapters
+
+    @property
+    def merged(self) -> bool:
+        return bool(self.merged_adapters)
+
+    @property
+    def active_adapters(self) -> list[str]:
+        # If some external mechanism sets active adapters, prefer it; else use all added adapters.
+        return getattr(self, "_active_adapters", list(self.rotation.keys()))
+
+    def get_base_layer(self) -> nn.Module:
+        return self.base_layer
+
+    def _cast_input_dtype(self, x: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        if not self._cast_input_dtype_enabled:
+            return x
+        return x.to(dtype)
+
+    def update_layer(
+        self,
+        adapter_name: str,
+        r: int,
+        T: float,
+        num_rotations: int,
+        drop_out: float,
+        **kwargs,
+    ):
+        """
+        Add / update a rotation adapter for this layer.
+        """
+        
+        if r <= 0:
+            raise ValueError(f"r must be positive, got {r}")
+        if num_rotations <= 0:
+            raise ValueError(f"num_rotations must be positive, got {num_rotations}")
+        
+        rot = Rotation(r=r, dim=self.in_features, T=T, num_rotations=num_rotations, drop_out=drop_out)
+        self.rotation[adapter_name] = rot
+        self.scaling[adapter_name] = 1.0
+        self._adapter_config[adapter_name] = {"r": r, "T": T, "num_rotations": num_rotations}
+
+    # (optional) helper to set currently active adapters externally
+    def set_active_adapters(self, adapters: Optional[list[str]]):
+        if adapters is None:
+            if hasattr(self, "_active_adapters"):
+                delattr(self, "_active_adapters")
+        else:
+            self._active_adapters = adapters
+            
+        
+class Linear(nn.Module, RotationLayer):
+    """
+    A linear layer with an integrated rotation layer for parameter-efficient fine-tuning.
+    """
+    
+    def __init__(self, 
+                 base_layer: nn.Linear,
+                 adapter_name: str,
+                 r: int,
+                 T: float,
+                 num_rotations: int,
+                 drop_out: float,
+                 **kwargs):
+        
+        super().__init__()
+        RotationLayer.__init__(self, base_layer=base_layer, **kwargs)
+        
+        self._active_adapter = adapter_name
+        
+        self.update_layer(
+            adapter_name=adapter_name,
+            r=r,
+            T=T,
+            num_rotations=num_rotations,
+            drop_out=drop_out,
+            **kwargs,
+        )
+        
+    def merge(self, safe_merge: bool = False, adapter_names: Optional[str] = None):
+        """
+        Merge the adapter effect into the base layer weights:
+            W_merged = W @ R
+        where R = I + delta (delta returned by get_delta_weight()).
+        """
+        adapter_names = check_adapters_to_merge(self, adapter_names)
+
+        if not adapter_names:
+            return
+        
+        base_layer = self.get_base_layer()
+        orig_dtype = base_layer.weight.dtype
+        # base_layer.weight shape: (out_features, in_features)
+        W = base_layer.weight.data  # (out, in)
+
+        for active_adapter in adapter_names:
+
+            if active_adapter not in self._available_adapters:
+                continue
+
+            delta_R = self.rotation[active_adapter].get_delta_weight()  # (in, in)
+            R = torch.eye(delta_R.size(0), device=delta_R.device, dtype=delta_R.dtype) + delta_R  # (in, in)
+            # merged W = W @ R
+            merged_W = W.to(R.dtype) @ R
+
+            if safe_merge and not torch.isfinite(merged_W).all():
+                raise ValueError("Merging resulted in non-finite weights. Aborting merge.")
+
+            base_layer.weight.data = merged_W.contiguous().to(orig_dtype)
+
+            # mark merged (so unmerge can restore by inverse)
+            self.merged_adapters.append(active_adapter)
+                
+                
+    def unmerge(self):
+        """
+        Reverse merges in LIFO order (pop merged adapters and invert R).
+        """
+        base_layer = self.get_base_layer()
+        orig_dtype = base_layer.weight.dtype
+
+        while self.merged_adapters:
+            active_adapter = self.merged_adapters.pop()
+            if active_adapter not in self._available_adapters:
+                continue
+            delta_R = self.rotation[active_adapter].get_delta_weight()  # (in, in)
+            R = torch.eye(delta_R.size(0), device=delta_R.device, dtype=delta_R.dtype) + delta_R
+            R_inv = torch.linalg.inv(R)
+            merged_W = base_layer.weight.data.to(R.dtype)
+            unmerged_W = merged_W @ R_inv
+            base_layer.weight.data = unmerged_W.contiguous().to(orig_dtype)
+            
+            
+    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        x_dtype = x.dtype
+        base_layer = self.get_base_layer()
+
+        if self.disable_adapters:
+            # if merged, unmerge to ensure base_layer produces original behavior
+            if self.merged:
+                self.unmerge()
+            return base_layer(x, *args, **kwargs).to(x_dtype)
+
+        if self.merged:
+            # if merged into base layer, just forward
+            return base_layer(x, *args, **kwargs).to(x_dtype)
+
+        # otherwise apply active adapters (transform inputs) then call base layer
+        for active_adapter in self.active_adapters:
+            if active_adapter not in self.rotation:
+                continue
+            rotation = self.rotation[active_adapter]
+            x = self._cast_input_dtype(x, rotation.U.dtype)
+            x = rotation(x)
+
+        return base_layer(x, *args, **kwargs).to(x_dtype)
+
+    def __repr__(self):
+        return f"rotation.{super().__repr__()}"
+        

nl_tasks/rpeft/rotation/layer_test.py

@@ -0,0 +1,296 @@
+import torch
+import torch.nn as nn
+from omini.rotation.layer import Linear, Rotation
+
+def test_rotation_merge():
+    """
+    Test that merging rotation adapter produces the same output as the unmerged version.
+    """
+    print("="*60)
+    print("Testing Rotation Layer Merge")
+    print("="*60)
+    
+    # Set random seed for reproducibility
+    torch.manual_seed(42)
+    
+    # Configuration
+    in_features = 512
+    out_features = 1024
+    r = 4
+    num_rotations = 4
+    T = 1.0
+    batch_size = 8
+    seq_len = 16
+    
+    # Create base linear layer
+    base_layer = nn.Linear(in_features, out_features, bias=True)
+    
+    # Create rotation layer
+    rotation_layer = Linear(
+        base_layer=base_layer,
+        adapter_name="default",
+        r=r,
+        T=T,
+        num_rotations=num_rotations
+    )
+    
+    # Create random input
+    x = torch.randn(batch_size, seq_len, in_features)
+    
+    # Test 1: Forward pass before merge
+    print("\n" + "-"*60)
+    print("Test 1: Computing output BEFORE merge")
+    print("-"*60)
+    rotation_layer.eval()
+    with torch.no_grad():
+        output_before = rotation_layer(x)
+    
+    print(f"Output shape: {output_before.shape}")
+    print(f"Output mean: {output_before.mean().item():.6f}")
+    print(f"Output std: {output_before.std().item():.6f}")
+    print(f"Output min: {output_before.min().item():.6f}")
+    print(f"Output max: {output_before.max().item():.6f}")
+    
+    # Save original weight for verification
+    original_weight = base_layer.weight.data.clone()
+    
+    # Test 2: Merge adapter
+    print("\n" + "-"*60)
+    print("Test 2: Merging adapter")
+    print("-"*60)
+    rotation_layer.merge(safe_merge=True, adapter_names=["default"])
+    print(f"✓ Adapter merged successfully")
+    print(f"✓ Merged adapters: {rotation_layer.merged_adapters}")
+    
+    # Check that weights have changed
+    weight_diff = (base_layer.weight.data - original_weight).abs().max().item()
+    print(f"Max weight change: {weight_diff:.6e}")
+    
+    # Test 3: Forward pass after merge
+    print("\n" + "-"*60)
+    print("Test 3: Computing output AFTER merge")
+    print("-"*60)
+    with torch.no_grad():
+        output_after = rotation_layer(x)
+    
+    print(f"Output shape: {output_after.shape}")
+    print(f"Output mean: {output_after.mean().item():.6f}")
+    print(f"Output std: {output_after.std().item():.6f}")
+    print(f"Output min: {output_after.min().item():.6f}")
+    print(f"Output max: {output_after.max().item():.6f}")
+    
+    # Test 4: Compare outputs
+    print("\n" + "-"*60)
+    print("Test 4: Comparing outputs")
+    print("-"*60)
+    
+    # Compute differences
+    abs_diff = (output_after - output_before).abs()
+    rel_diff = abs_diff / (output_before.abs() + 1e-8)
+    
+    max_abs_diff = abs_diff.max().item()
+    mean_abs_diff = abs_diff.mean().item()
+    max_rel_diff = rel_diff.max().item()
+    mean_rel_diff = rel_diff.mean().item()
+    
+    print(f"Max absolute difference: {max_abs_diff:.6e}")
+    print(f"Mean absolute difference: {mean_abs_diff:.6e}")
+    print(f"Max relative difference: {max_rel_diff:.6e}")
+    print(f"Mean relative difference: {mean_rel_diff:.6e}")
+    
+    # Check if outputs are close
+    atol = 1e-4  # Absolute tolerance
+    rtol = 1e-3  # Relative tolerance
+    
+    are_close = torch.allclose(output_before, output_after, atol=atol, rtol=rtol)
+    
+    if are_close:
+        print(f"\n✅ PASS: Outputs are identical (within atol={atol}, rtol={rtol})")
+    else:
+        print(f"\n❌ FAIL: Outputs differ significantly")
+        print(f"   Expected: atol < {atol}, rtol < {rtol}")
+        print(f"   Got: max_abs_diff = {max_abs_diff:.6e}, max_rel_diff = {max_rel_diff:.6e}")
+    
+    # Test 5: Unmerge and verify
+    print("\n" + "-"*60)
+    print("Test 5: Testing unmerge")
+    print("-"*60)
+    rotation_layer.unmerge()
+    print(f"✓ Adapter unmerged")
+    print(f"✓ Merged adapters: {rotation_layer.merged_adapters}")
+    
+    with torch.no_grad():
+        output_unmerged = rotation_layer(x)
+    
+    unmerge_diff = (output_unmerged - output_before).abs().max().item()
+    print(f"Max difference after unmerge: {unmerge_diff:.6e}")
+    
+    unmerge_close = torch.allclose(output_before, output_unmerged, atol=atol, rtol=rtol)
+    if unmerge_close:
+        print(f"✅ PASS: Unmerge restored original behavior")
+    else:
+        print(f"❌ FAIL: Unmerge did not restore original behavior")
+    
+    # Test 6: Verify weight restoration
+    weight_restored_diff = (base_layer.weight.data - original_weight).abs().max().item()
+    print(f"Max weight difference after unmerge: {weight_restored_diff:.6e}")
+    
+    weight_restored = torch.allclose(base_layer.weight.data, original_weight, atol=1e-5)
+    if weight_restored:
+        print(f"✅ PASS: Original weights restored")
+    else:
+        print(f"❌ FAIL: Original weights not fully restored")
+    
+    print("\n" + "="*60)
+    print("Test Summary")
+    print("="*60)
+    return are_close and unmerge_close and weight_restored
+
+
+def test_multiple_merges():
+    """
+    Test merging and unmerging multiple times.
+    """
+    print("\n" + "="*60)
+    print("Testing Multiple Merge/Unmerge Cycles")
+    print("="*60)
+    
+    torch.manual_seed(42)
+    
+    in_features = 256
+    out_features = 512
+    r = 4
+    num_rotations = 4
+    
+    base_layer = nn.Linear(in_features, out_features, bias=True)
+    rotation_layer = Linear(
+        base_layer=base_layer,
+        adapter_name="default",
+        r=r,
+        T=1.0,
+        num_rotations=num_rotations
+    )
+    
+    x = torch.randn(4, 8, in_features)
+    rotation_layer.eval()
+    
+    # Get original output
+    with torch.no_grad():
+        original_output = rotation_layer(x)
+    
+    # Test multiple cycles
+    all_passed = True
+    for cycle in range(3):
+        print(f"\nCycle {cycle + 1}:")
+        
+        # Merge
+        rotation_layer.merge(safe_merge=True)
+        with torch.no_grad():
+            merged_output = rotation_layer(x)
+        
+        merge_close = torch.allclose(original_output, merged_output, atol=1e-4, rtol=1e-3)
+        print(f"  Merge: {'✅ PASS' if merge_close else '❌ FAIL'}")
+        
+        # Unmerge
+        rotation_layer.unmerge()
+        with torch.no_grad():
+            unmerged_output = rotation_layer(x)
+        
+        unmerge_close = torch.allclose(original_output, unmerged_output, atol=1e-4, rtol=1e-3)
+        print(f"  Unmerge: {'✅ PASS' if unmerge_close else '❌ FAIL'}")
+        
+        all_passed = all_passed and merge_close and unmerge_close
+    
+    return all_passed
+
+
+def test_with_different_dtypes():
+    """
+    Test merging with different data types.
+    """
+    print("\n" + "="*60)
+    print("Testing Different Data Types")
+    print("="*60)
+    
+    torch.manual_seed(42)
+    
+    dtypes = [torch.float32, torch.float16, torch.bfloat16]
+    all_passed = True
+    
+    for dtype in dtypes:
+        print(f"\nTesting with dtype: {dtype}")
+        
+        in_features = 256
+        out_features = 512
+        r = 4
+        num_rotations = 4
+        
+        base_layer = nn.Linear(in_features, out_features, bias=True)
+        base_layer = base_layer.to(dtype)
+        
+        rotation_layer = Linear(
+            base_layer=base_layer,
+            adapter_name="default",
+            r=r,
+            T=1.0,
+            num_rotations=num_rotations
+        )
+        rotation_layer = rotation_layer.to(dtype)
+        
+        x = torch.randn(4, 8, in_features, dtype=dtype)
+        rotation_layer.eval()
+        
+        with torch.no_grad():
+            output_before = rotation_layer(x)
+            rotation_layer.merge(safe_merge=True)
+            output_after = rotation_layer(x)
+        
+        # Adjust tolerances based on dtype
+        if dtype == torch.float32:
+            atol, rtol = 1e-5, 1e-4
+        elif dtype == torch.float16:
+            atol, rtol = 1e-2, 1e-2
+        else:  # bfloat16
+            atol, rtol = 1e-2, 1e-2
+        
+        are_close = torch.allclose(output_before, output_after, atol=atol, rtol=rtol)
+        
+        if are_close:
+            print(f"  ✅ PASS")
+        else:
+            max_diff = (output_after - output_before).abs().max().item()
+            print(f"  ❌ FAIL (max diff: {max_diff:.6e})")
+        
+        all_passed = all_passed and are_close
+    
+    return all_passed
+
+
+if __name__ == "__main__":
+    print("\n" + "="*60)
+    print("ROTATION LAYER MERGE TEST SUITE")
+    print("="*60)
+    
+    results = {}
+    
+    # Run all tests
+    results["basic_merge"] = test_rotation_merge()
+    results["multiple_cycles"] = test_multiple_merges()
+    results["different_dtypes"] = test_with_different_dtypes()
+    
+    # Print summary
+    print("\n" + "="*60)
+    print("FINAL SUMMARY")
+    print("="*60)
+    
+    for test_name, passed in results.items():
+        status = "✅ PASS" if passed else "❌ FAIL"
+        print(f"{test_name}: {status}")
+    
+    all_passed = all(results.values())
+    print("\n" + "="*60)
+    if all_passed:
+        print("🎉 ALL TESTS PASSED!")
+    else:
+        print("⚠️  SOME TESTS FAILED")
+    print("="*60)

nl_tasks/rpeft/rotation/model.py

@@ -0,0 +1,392 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from enum import Enum
+from dataclasses import asdict
+from tqdm import tqdm
+
+
+from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists, onload_layer
+
+from peft.utils import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, ModulesToSaveWrapper, _get_submodules
+
+from .layer import RotationLayer, Linear
+
+TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING.copy()
+
+class RotationTuner(BaseTuner):
+    
+    prefix: str = "rotation"
+    tuner_layer_class = RotationLayer
+    target_module_mapping = TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING
+    
+    
+    @staticmethod
+    def _check_target_module_exists(rotation_config, key: str) -> bool:
+        return check_target_module_exists(rotation_config, key)
+    
+    def _create_and_replace(
+        self, 
+        rotation_config,
+        adapter_name: str, 
+        target: nn.Module, 
+        target_name: str, 
+        parent: nn.Module, 
+        current_key: str,
+        **optional_kwargs,
+    ) -> None:
+        """
+        Create and replace a target module with a rotation-augmented version.
+        
+        This method is called when an existing module is already a RotationLayer
+        and needs to have a new adapter added to it.
+        
+        Args:
+            rotation_config: Configuration for the rotation adapter
+            adapter_name: Name of the adapter to add
+            target: The target module to augment
+            target_name: Name of the target module
+            parent: Parent module containing the target
+            current_key: Full key path to the current module
+            **optional_kwargs: Additional optional arguments
+            
+        Raises:
+            ValueError: If current_key is not provided
+        """
+        
+        if current_key is None:
+            raise ValueError("current_key must be provided to create Rotation layer")
+        
+        # Check if target is already a RotationLayer
+        if isinstance(target, RotationLayer):
+            target.update_layer(
+                adapter_name=adapter_name,
+                r=rotation_config.r,
+                T=rotation_config.T,
+                num_rotations=rotation_config.num_rotations,
+            )
+        else:
+            # Create new rotation layer
+            new_module = self._create_new_module(
+                rotation_config=rotation_config,
+                adapter_name=adapter_name,
+                target=target,
+                **optional_kwargs,
+            )
+            if new_module is not None:
+                self._replace_module(parent, target_name, new_module, target)
+        
+    def _replace_module(self, parent, child_name, new_module, child):
+        
+        setattr(parent, child_name, new_module)
+        
+        # child layer wraps the original module, unpack it
+        if hasattr(child, "base_layer"):
+            child = child.base_layer
+        
+        meta = torch.device("meta")
+        # dispatch to correct device
+        for name, module in new_module.named_modules():
+            if (self.prefix in name) or ("ranknum" in name):
+                if hasattr(child, "qweight"):
+                    weight = child.qweight
+                elif hasattr(child, "W_q"):
+                    weight = child.W_q
+                elif hasattr(child, "weight"):
+                    weight = child.weight
+                elif getattr(child, "in_proj_weight", None) is not None:  # MHA
+                    weight = child.in_proj_weight
+                else:
+                    weight = next(child.parameters())
+                if not any(p.device == meta for p in module.parameters()):
+                    module.to(weight.device)
+    
+    def _mark_only_adapters_as_trainable(self, model):
+        
+        # First, freeze all parameters
+        for n, p in model.named_parameters():
+            # print(f'{n}, np  {p.requires_grad}')
+            if self.prefix not in n:
+                p.requires_grad = False
+            else:
+                p.requires_grad = True
+        
+        # Handle bias parameters based on config
+        for active_adapter in self.active_adapters:
+            bias_config = self.peft_config[active_adapter].bias
+            
+            if bias_config == "none":
+                continue
+            elif bias_config == "all":
+                # Enable all bias parameters
+                for n, p in model.named_parameters():
+                    if "bias" in n:
+                        p.requires_grad = True
+            elif bias_config == "rotation_only":
+                # Enable only bias in rotation layers
+                for name, m in model.named_modules():
+                    if isinstance(m, RotationLayer):
+                        if hasattr(m, "bias") and m.bias is not None:
+                            m.bias.requires_grad = True
+            else:
+                raise NotImplementedError(
+                    f"Requested bias configuration '{bias_config}' is not implemented. "
+                    f"Supported values: 'none', 'all', 'rotation_only'"
+                )
+    
+    @staticmethod
+    def _create_new_module(
+        rotation_config,
+        adapter_name: str,
+        target: nn.Module,
+        **kwargs,
+    ) -> Optional[nn.Module]:
+        """
+        Create a new rotation-augmented module.
+        
+        Args:
+            rotation_config: Configuration for the rotation adapter
+            adapter_name: Name of the adapter
+            target: Base module to augment
+            **kwargs: Additional arguments
+            
+        Returns:
+            New RotationLayer module wrapping the target, or None if unsupported
+        """
+        if isinstance(target, nn.Linear):
+            return Linear(
+                base_layer=target,
+                adapter_name=adapter_name,
+                r=rotation_config.r,
+                T=rotation_config.T,
+                num_rotations=rotation_config.num_rotations,
+                drop_out=rotation_config.drop_out,
+                **kwargs,
+            )
+        else:
+            # Unsupported layer type
+            print(
+                f"Rotation layer does not support {type(target).__name__} yet. "
+                f"Skipping this module."
+            )
+            return None
+    
+    
+    def __getattr__(self, name: str):
+        """Forward missing attributes to the wrapped module."""
+        try:
+            return super().__getattr__(name)  # defer to nn.Module's logic
+        except AttributeError:
+            if name == "model":  # see #1892: prevent infinite recursion if class is not initialized
+                raise
+            return getattr(self.model, name)
+    
+    def get_peft_config_as_dict(self, inference: bool = False):
+        config_dict = {}
+        for key, value in self.peft_config.items():
+            config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()}
+            if inference:
+                config["inference_mode"] = True
+        config_dict[key] = config
+        return config
+    
+    
+    def _set_adapter_layers(self, enabled=True):
+        for module in self.model.modules():
+            if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)):
+                module.enable_adapters(enabled)
+
+    def enable_adapter_layers(self) -> None:
+        """Enable all adapters.
+
+        Call this if you have previously disabled all adapters and want to re-enable them.
+        """
+        self._set_adapter_layers(enabled=True)
+
+    def disable_adapter_layers(self):
+        for active_adapter in self.active_adapters:
+            val = self.peft_config[active_adapter].bias
+            if val != "none":
+                msg = (
+                    f"Careful, disabling adapter layers with bias configured to be '{val}' does not produce the same "
+                    "output as the base model would without adaption."
+                )
+                print(msg)
+        self._set_adapter_layers(enabled=False)
+
+    def set_adapter(self, adapter_name, inference_mode):
+        """Set the active adapter(s).
+
+        Additionally, this function will set the specified adapters to trainable (i.e., requires_grad=True). If this is
+        not desired, use the following code.
+
+        ```py
+        >>> for name, param in model_peft.named_parameters():
+        ...     if ...:  # some check on name (ex. if 'lora' in name)
+        ...         param.requires_grad = False
+        ```
+
+        Args:
+            adapter_name (`str` or `list[str]`): Name of the adapter(s) to be activated.
+        """
+        for module in self.model.modules():
+            if isinstance(module, RotationLayer):
+                if module.merged:
+                    print("Adapter cannot be set when the model is merged. Unmerging the model first.")
+                    module.unmerge()
+                module.set_adapter(adapter_name, inference_mode)
+        self.active_adapter = adapter_name
+    
+    def merge_adapter(self, adapter_names: Optional[list[str]] = None) -> None:
+        """
+        Merge adapter weights into the base model weights.
+        
+        This can speed up inference by eliminating the need for runtime
+        rotation computations.
+        
+        Args:
+            adapter_names: List of adapter names to merge. If None, merges all
+                active adapters.
+        """
+        for module in self.model.modules():
+            if isinstance(module, RotationLayer):
+                module.merge(safe_merge=False, adapter_names=adapter_names)
+    
+    
+    def unmerge_adapter(self) -> None:
+        """
+        Unmerge adapter weights from the base model weights.
+        
+        This reverses the merge operation, restoring dynamic adapter behavior.
+        """
+        for module in self.model.modules():
+            if isinstance(module, RotationLayer):
+                module.unmerge()
+    
+    @staticmethod
+    def _prepare_adapter_config(peft_config, model_config):
+        
+        if peft_config.target_modules is None:
+            if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING:
+                raise ValueError("Please specify `target_modules` in `peft_config`")
+            peft_config.target_modules = set(
+                TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING[model_config["model_type"]]
+            )
+            
+        return peft_config
+    
+    
+    def _check_new_adapter_config(self, config) -> None:
+        """
+        Check the validity of a new adapter configuration.
+        
+        Args:
+            config: Configuration to validate
+            
+        Raises:
+            ValueError: If configuration is invalid
+        """
+        # Validate rank
+        if config.r <= 0:
+            raise ValueError(f"r must be positive, got {config.r}")
+        
+        # Validate num_rotations
+        if config.num_rotations <= 0:
+            raise ValueError(
+                f"num_rotations must be positive, got {config.num_rotations}"
+            )
+        
+        
+        # Validate bias configuration
+        valid_bias_configs = ["none", "all", "rotation_only"]
+        if hasattr(config, "bias") and config.bias not in valid_bias_configs:
+            raise ValueError(
+                f"Invalid bias configuration '{config.bias}'. "
+                f"Must be one of {valid_bias_configs}"
+            )
+            
+    
+    def _unload_and_optionally_merge(
+        self,
+        merge=True,
+        progressbar: bool = False,
+        safe_merge: bool = False,
+        adapter_names: Optional[list[str]] = None,
+    ):
+        if merge:
+            self._check_merge_allowed()
+
+        key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key]
+        desc = "Unloading " + ("and merging " if merge else "") + "model"
+        for key in tqdm(key_list, disable=not progressbar, desc=desc):
+            try:
+                parent, target, target_name = _get_submodules(self.model, key)
+            except AttributeError:
+                continue
+            with onload_layer(target):
+                if hasattr(target, "unload_and_optionally_merge_module"):
+                    # if layers have special unloading method, like MultiheadAttention, use that
+                    unloaded_module = target.unload_and_optionally_merge_module(
+                        merge=merge, safe_merge=safe_merge, adapter_names=adapter_names
+                    )
+                    self._replace_module(parent, target_name, unloaded_module, target)
+                elif hasattr(target, "base_layer"):
+                    if merge:
+                        target.merge(safe_merge=safe_merge, adapter_names=adapter_names)
+                    self._replace_module(parent, target_name, target.get_base_layer(), target)
+
+        return self.model
+
+    def delete_adapter(self, adapter_name: str) -> None:
+        """
+        Deletes an existing adapter.
+
+        Args:
+            adapter_name (str): Name of the adapter to be deleted.
+        """
+        if adapter_name not in list(self.peft_config.keys()):
+            raise ValueError(f"Adapter {adapter_name} does not exist")
+        del self.peft_config[adapter_name]
+
+        key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key]
+        new_adapter = None
+        for key in key_list:
+            _, target, _ = _get_submodules(self.model, key)
+            if isinstance(target, RotationLayer):
+                target.delete_adapter(adapter_name)
+                if new_adapter is None:
+                    new_adapter = target.active_adapters[:]
+
+        self.active_adapter = new_adapter or []
+        self._delete_auxiliary_adapter(adapter_name, new_active_adapters=new_adapter)
+
+    def merge_and_unload(
+        self, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[list[str]] = None
+    ) -> torch.nn.Module:
+        r"""
+        This method merges the OFT layers into the base model. This is needed if someone wants to use the base model as
+        a standalone model.
+
+        Args:
+            progressbar (`bool`):
+                whether to show a progressbar indicating the unload and merge process
+            safe_merge (`bool`):
+                whether to activate the safe merging check to check if there is any potential Nan in the adapter
+                weights
+            adapter_names (`List[str]`, *optional*):
+                The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults
+                to `None`.
+
+        """
+        return self._unload_and_optionally_merge(
+            progressbar=progressbar, safe_merge=safe_merge, adapter_names=adapter_names
+        )
+
+    def unload(self) -> torch.nn.Module:
+        """
+        Gets back the base model by removing all the oft modules without merging. This gives back the original base
+        model.
+        """
+        return self._unload_and_optionally_merge(merge=False)
+    

nl_tasks/rpeft/rotation/rotation_config.py

@@ -0,0 +1,89 @@
+from dataclasses import dataclass, field
+from typing import List, Optional
+# from peft.config import PeftConfig
+from rpeft.utils import PeftConfig
+from rpeft.utils import PeftType
+
+@dataclass
+class RotationConfig(PeftConfig):
+    """
+    Configuration class for Rotation-based Parameter-Efficient Fine-Tuning.
+    
+    This configuration stores all parameters needed to apply the Rotation method
+    (based on Cayley transformation) to a model's linear layers.
+    
+    Args:
+        r (`int`):
+            The rank parameter for the low-rank approximation in rotation matrices.
+        T (`float`, *optional*, defaults to 1.0):
+            Temperature parameter for the transformation.
+        num_rotations (`int`, *optional*, defaults to 4):
+            Number of rotation matrices to use in parallel.
+        target_modules (`Union[List[str], str]`):
+            Module names to apply rotation to (e.g., ["q_proj", "v_proj"]).
+        target_modules_to_skip (`Union[List[str], str]`, *optional*):
+            Module names to skip when applying rotation.
+        modules_to_save (`Union[List[str], str]`, *optional*):
+            Modules to save in addition to rotation parameters.
+        layers_to_transform (`Union[List[int], int]`, *optional*):
+            Layers to transform. If None, all layers matching target_modules are transformed.
+        apply_before (`bool`, *optional*, defaults to False):
+            If True, apply rotation before the base linear layer. If False, apply after.
+    """
+    
+    peft_type: str = field(default="ROTATION", init=False)  
+    target_modules: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "List of module names to apply rotation to (e.g., ['q_proj', 'v_proj', 'linear'])"
+        },
+    )
+    target_modules_to_skip: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": "List of module names to skip when applying rotation"},
+    )
+    modules_to_save: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": "List of modules to save in addition to rotation parameters"},
+    )
+    r: int = field(
+        default=8,
+        metadata={"help": "Rank parameter for low-rank approximation"},
+    )
+    T: float = field(
+        default=1.0,
+        metadata={"help": "Temperature parameter for Cayley transformation"},
+    )
+    num_rotations: int = field(
+        default=4,
+        metadata={"help": "Number of rotation matrices to use in parallel"},
+    )
+    
+    bias: str = field(
+        default="none",
+        metadata={
+            "help": "Bias training configuration. Options: 'none', 'all', 'rotation_only'"
+        }
+    )
+    layers_to_transform: Optional[List[int]] = field(
+        default=None,
+        metadata={"help": "Layers to transform. If None, all matching layers are transformed"},
+    )
+    drop_out: float = field(
+        default=0.0,
+        metadata={
+            'help': 'intput drop out rate'
+        }
+    )
+    
+    def __post_init__(self):
+        ##### Diff
+        self.peft_type = PeftType.ROTATION
+        self.target_modules = (
+            set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules
+        )
+        self.target_modules_to_skip = (
+            set(self.target_modules_to_skip) 
+            if isinstance(self.target_modules_to_skip, list) 
+            else self.target_modules_to_skip
+        )

nl_tasks/rpeft/utils/__init__.py

@@ -0,0 +1,29 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
+
+# coding=utf-8
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .adapters_utils import CONFIG_NAME, WEIGHTS_NAME
+from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType
+from .other import (
+    TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING,
+    _set_trainable,
+    bloom_model_postprocess_past_key_value,
+    prepare_model_for_int8_training,
+    shift_tokens_right,
+    transpose,
+)
+from .save_and_load import get_peft_model_state_dict, set_peft_model_state_dict

nl_tasks/rpeft/utils/adapters_utils.py

@@ -0,0 +1,19 @@
+# coding=utf-8
+# Original License:
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+WEIGHTS_NAME = "adapter_model.bin"
+CONFIG_NAME = "adapter_config.json"
+
+# TODO: add automapping and superclass here?

nl_tasks/rpeft/utils/config.py

@@ -0,0 +1,220 @@
+# coding=utf-8
+# Original License:
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import enum
+import json
+import os
+from dataclasses import asdict, dataclass, field
+from typing import Optional, Union
+
+from huggingface_hub import hf_hub_download
+from transformers.utils import PushToHubMixin, http_user_agent
+
+from .adapters_utils import CONFIG_NAME
+
+
+class PeftType(str, enum.Enum):
+    PROMPT_TUNING = "PROMPT_TUNING"
+    P_TUNING = "P_TUNING"
+    PREFIX_TUNING = "PREFIX_TUNING"
+    LORA = "LORA"
+    BOTTLENECK = "BOTTLENECK"
+    QUANTA = "QUANTA"
+
+    ROTATION = "ROTATION"
+
+
+
+class TaskType(str, enum.Enum):
+    SEQ_CLS = "SEQ_CLS"
+    SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM"
+    CAUSAL_LM = "CAUSAL_LM"
+    TOKEN_CLS = "TOKEN_CLS"
+
+
+@dataclass
+class PeftConfigMixin(PushToHubMixin):
+    r"""
+    This is the base configuration class for PEFT adapter models. It contains all the methods that are common to all
+    PEFT adapter models. This class inherits from `transformers.utils.PushToHubMixin` which contains the methods to
+    push your model to the Hub. The method `save_pretrained` will save the configuration of your adapter model in a
+    directory. The method `from_pretrained` will load the configuration of your adapter model from a directory.
+
+    Args:
+        peft_type (Union[[`~peft_local_tensor.utils.config.PeftType`], `str`]): The type of Peft method to use.
+    """
+    peft_type: Optional[PeftType] = field(default=None, metadata={"help": "The type of PEFT model."})
+
+    @property
+    def __dict__(self):
+        return asdict(self)
+
+    def to_dict(self):
+        return self.__dict__
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, subfolder: Optional[str] = None, **kwargs):
+        r"""
+        This method loads the configuration of your adapter model from a directory.
+
+        Args:
+            pretrained_model_name_or_path (`str`):
+                The directory or the Hub repository id where the configuration is saved.
+            kwargs (additional keyword arguments, *optional*):
+                Additional keyword arguments passed along to the child class initialization.
+        """
+        path = (
+            os.path.join(pretrained_model_name_or_path, subfolder)
+            if subfolder is not None
+            else pretrained_model_name_or_path
+        )
+
+        hf_hub_download_kwargs, class_kwargs, _ = cls._split_kwargs(kwargs)
+        if "user_agent" not in hf_hub_download_kwargs:
+            hf_hub_download_kwargs["user_agent"] = http_user_agent()
+
+        if os.path.isfile(os.path.join(path, CONFIG_NAME)):
+            config_file = os.path.join(path, CONFIG_NAME)
+        else:
+            try:
+                config_file = hf_hub_download(
+                    pretrained_model_name_or_path, CONFIG_NAME, subfolder=subfolder, **hf_hub_download_kwargs
+                )
+            except Exception as exc:
+                raise ValueError(f"Can't find '{CONFIG_NAME}' at '{pretrained_model_name_or_path}'") from exc
+
+        loaded_attributes = cls.from_json_file(config_file)
+        kwargs = {**class_kwargs, **loaded_attributes}
+        kwargs = cls.check_kwargs(**kwargs)
+        return cls.from_peft_type(**kwargs)
+    
+    def save_pretrained(self, save_directory, **kwargs):
+        r"""
+        This method saves the configuration of your adapter model in a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory where the configuration will be saved.
+            **kwargs:
+                Additional keyword arguments passed along to the `transformers.utils.PushToHubMixin.push_to_hub`
+                method.
+        """
+        if os.path.isfile(save_directory):
+            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
+
+        os.makedirs(save_directory, exist_ok=True)
+        auto_mapping_dict = kwargs.pop("auto_mapping_dict", None)
+
+        output_dict = self.to_dict()
+        # converting set type to list
+        for key, value in output_dict.items():
+            if isinstance(value, set):
+                output_dict[key] = list(value)
+
+        output_path = os.path.join(save_directory, CONFIG_NAME)
+
+        # Add auto mapping details for custom models.
+        if auto_mapping_dict is not None:
+            output_dict["auto_mapping"] = auto_mapping_dict
+
+        # save it
+        with open(output_path, "w") as writer:
+            writer.write(json.dumps(output_dict, indent=2, sort_keys=True))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        r"""
+        This method loads the configuration of your adapter model from a directory.
+
+        Args:
+            pretrained_model_name_or_path (`str`):
+                The directory or the hub-id where the configuration is saved.
+            **kwargs:
+                Additional keyword arguments passed along to the child class initialization.
+        """
+        if os.path.isfile(os.path.join(pretrained_model_name_or_path, CONFIG_NAME)):
+            config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+        else:
+            try:
+                config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME)
+            except Exception:
+                raise ValueError(f"Can't find config.json at '{pretrained_model_name_or_path}'")
+
+        loaded_attributes = cls.from_json_file(config_file)
+
+        config = cls(**kwargs)
+
+        for key, value in loaded_attributes.items():
+            if hasattr(config, key):
+                setattr(config, key, value)
+
+        return config
+
+    @classmethod
+    def from_json_file(cls, path_json_file, **kwargs):
+        r"""
+        Loads a configuration file from a json file.
+
+        Args:
+            path_json_file (`str`):
+                The path to the json file.
+        """
+        with open(path_json_file, "r") as file:
+            json_object = json.load(file)
+
+        return json_object
+
+
+@dataclass
+class PeftConfig(PeftConfigMixin):
+    """
+    This is the base configuration class to store the configuration of a :class:`~peft_local_tensor.PeftModel`.
+
+    Args:
+        peft_type (Union[[`~peft_local_tensor.utils.config.PeftType`], `str`]): The type of Peft method to use.
+        task_type (Union[[`~peft_local_tensor.utils.config.TaskType`], `str`]): The type of task to perform.
+        inference_mode (`bool`, defaults to `False`): Whether to use the Peft model in inference mode.
+    """
+
+    base_model_name_or_path: str = field(default=None, metadata={"help": "The name of the base model to use."})
+    revision: Optional[str] = field(default=None, metadata={"help": "The specific base model version to use."})
+    peft_type: Union[str, PeftType] = field(default=None, metadata={"help": "Peft type"})
+    task_type: Union[str, TaskType] = field(default=None, metadata={"help": "Task type"})
+    inference_mode: bool = field(default=False, metadata={"help": "Whether to use inference mode"})
+
+
+@dataclass
+class PromptLearningConfig(PeftConfig):
+    """
+    This is the base configuration class to store the configuration of a Union[[`~peft_local_tensor.PrefixTuning`],
+    [`~peft_local_tensor.PromptEncoder`], [`~peft_local_tensor.PromptTuning`]].
+
+    Args:
+        num_virtual_tokens (`int`): The number of virtual tokens to use.
+        token_dim (`int`): The hidden embedding dimension of the base transformer model.
+        num_transformer_submodules (`int`): The number of transformer submodules in the base transformer model.
+        num_attention_heads (`int`): The number of attention heads in the base transformer model.
+        num_layers (`int`): The number of layers in the base transformer model.
+    """
+
+    num_virtual_tokens: int = field(default=None, metadata={"help": "Number of virtual tokens"})
+    token_dim: int = field(
+        default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"}
+    )
+    num_transformer_submodules: Optional[int] = field(
+        default=None, metadata={"help": "Number of transformer submodules"}
+    )
+    num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"})
+    num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})

nl_tasks/rpeft/utils/other.py

@@ -0,0 +1,160 @@
+# coding=utf-8
+# Original License:
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+
+# needed for prefix-tuning of bloom model
+def bloom_model_postprocess_past_key_value(past_key_values):
+    past_key_values = torch.cat(past_key_values)
+    total_layers, batch_size, num_attention_heads, num_virtual_tokens, head_dim = past_key_values.shape
+    keys = past_key_values[: total_layers // 2]
+    keys = keys.transpose(2, 3).reshape(
+        total_layers // 2, batch_size * num_attention_heads, head_dim, num_virtual_tokens
+    )
+    values = past_key_values[total_layers // 2 :]
+    values = values.reshape(total_layers // 2, batch_size * num_attention_heads, num_virtual_tokens, head_dim)
+
+    return tuple(zip(keys, values))
+
+
+def prepare_model_for_int8_training(
+    model, output_embedding_layer_name="lm_head", use_gradient_checkpointing=True, layer_norm_names=["layer_norm"]
+):
+    r"""
+    This method wrapps the entire protocol for preparing a model before running a training. This includes:
+        1- Cast the layernorm in fp32 2- making output embedding layer require grads 3- Add the upcasting of the lm
+        head to fp32
+
+    Args:
+        model, (`transformers.PreTrainedModel`):
+            The loaded model from `transformers`
+    """
+    loaded_in_8bit = getattr(model, "is_loaded_in_8bit", False)
+
+    for name, param in model.named_parameters():
+        # freeze base model's layers
+        param.requires_grad = False
+
+        if loaded_in_8bit:
+            # cast layer norm in fp32 for stability for 8bit models
+            if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
+                param.data = param.data.to(torch.float32)
+
+    if loaded_in_8bit and use_gradient_checkpointing:
+        # For backward compatibility
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
+
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+
+            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+        # enable gradient checkpointing for memory efficiency
+        model.gradient_checkpointing_enable()
+
+    if hasattr(model, output_embedding_layer_name):
+        output_embedding_layer = getattr(model, output_embedding_layer_name)
+        input_dtype = output_embedding_layer.weight.dtype
+
+        class CastOutputToFloat(torch.nn.Sequential):
+            r"""
+            Manually cast to the expected dtype of the lm_head as sometimes there is a final layer norm that is casted
+            in fp32
+
+            """
+
+            def forward(self, x):
+                return super().forward(x.to(input_dtype)).to(torch.float32)
+
+        setattr(model, output_embedding_layer_name, CastOutputToFloat(output_embedding_layer))
+
+    return model
+
+
+TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = {
+    "bloom": bloom_model_postprocess_past_key_value,
+}
+
+
+# copied from transformers.models.bart.modeling_bart
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+    """
+    Shift input ids one token to the right.
+
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): input ids
+        pad_token_id (`int`): The id of the `padding` token.
+        decoder_start_token_id (`int`): The id of the `start` token.
+    """
+    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+    shifted_input_ids[:, 0] = decoder_start_token_id
+
+    if pad_token_id is None:
+        raise ValueError("self.model.config.pad_token_id has to be defined.")
+    # replace possible -100 values in labels by `pad_token_id`
+    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+    return shifted_input_ids
+
+
+def _set_trainable(model):
+    if model.modules_to_save is not None:
+        for name, param in model.named_parameters():
+            if any(module_name in name for module_name in model.modules_to_save):
+                param.requires_grad = True
+
+
+def fsdp_auto_wrap_policy(model):
+    import functools
+    import os
+
+    from accelerate import FullyShardedDataParallelPlugin
+    from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy
+
+    from ..tuners import PrefixEncoder, PromptEmbedding, PromptEncoder
+
+    def lambda_policy_fn(module):
+        if (
+            len(list(module.named_children())) == 0
+            and getattr(module, "weight", None) is not None
+            and module.weight.requires_grad
+        ):
+            return True
+        return False
+
+    lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn)
+    transformer_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls=(
+            PrefixEncoder,
+            PromptEncoder,
+            PromptEmbedding,
+            FullyShardedDataParallelPlugin.get_module_class_from_name(
+                model, os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "")
+            ),
+        ),
+    )
+
+    auto_wrap_policy = functools.partial(_or_policy, policies=[lambda_policy, transformer_wrap_policy])
+    return auto_wrap_policy
+
+
+def transpose(weight, fan_in_fan_out):
+    return weight.T if fan_in_fan_out else weight

nl_tasks/rpeft/utils/save_and_load.py

@@ -0,0 +1,166 @@
+# coding=utf-8
+# Original License:
+# Copyright 2023-present the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .config import PeftType
+import warnings
+import torch
+
+def _find_mismatched_keys(
+    model: torch.nn.Module, peft_model_state_dict: dict[str, torch.Tensor], ignore_mismatched_sizes: bool = False
+) -> tuple[dict[str, torch.Tensor], list[tuple[str, tuple[int, ...], tuple[int, ...]]]]:
+    if not ignore_mismatched_sizes:
+        return peft_model_state_dict, []
+
+    mismatched = []
+    state_dict = model.state_dict()
+    for key, tensor in peft_model_state_dict.items():
+        if key not in state_dict:
+            continue
+
+        # see https://github.com/huggingface/transformers/blob/09f9f566de83eef1f13ee83b5a1bbeebde5c80c1/src/transformers/modeling_utils.py#L3858-L3864
+        if (state_dict[key].shape[-1] == 1) and (state_dict[key].numel() * 2 == tensor.numel()):
+            # This skips size mismatches for 4-bit weights. Two 4-bit values share an 8-bit container, causing size
+            # differences. Without matching with module type or parameter type it seems like a practical way to detect
+            # valid 4bit weights.
+            continue
+
+        if state_dict[key].shape != tensor.shape:
+            mismatched.append((key, tensor.shape, state_dict[key].shape))
+
+    for key, _, _ in mismatched:
+        del peft_model_state_dict[key]
+
+    return peft_model_state_dict, mismatched
+
+def get_peft_model_state_dict(model, state_dict=None):
+    """
+    Get the state dict of the Peft model.
+
+    Args:
+        model ([`PeftModel`]): The Peft model. When using torch.nn.DistributedDataParallel, DeepSpeed or FSDP,
+        the model should be the underlying model/unwrapped model (i.e. model.module).
+        state_dict (`dict`, *optional*, defaults to `None`):
+            The state dict of the model. If not provided, the state dict of the model
+        will be used.
+    """
+    if state_dict is None:
+        state_dict = model.state_dict()
+    if model.peft_config.peft_type == PeftType.LORA:
+        # to_return = lora_state_dict(model, bias=model.peft_config.bias)
+        # adapted from `https://github.com/microsoft/LoRA/blob/main/loralib/utils.py`
+        # to directly with the state dict which is necessary when using DeepSpeed or FSDP
+        bias = model.peft_config.bias
+        if bias == "none":
+            to_return = {k: state_dict[k] for k in state_dict if "lora_" in k}
+        elif bias == "all":
+            to_return = {k: state_dict[k] for k in state_dict if "lora_" in k or "bias" in k}
+        elif bias == "lora_only":
+            to_return = {}
+            for k in state_dict:
+                if "lora_" in k:
+                    to_return[k] = state_dict[k]
+                    bias_name = k.split("lora_")[0] + "bias"
+                    if bias_name in state_dict:
+                        to_return[bias_name] = state_dict[bias_name]
+        else:
+            raise NotImplementedError
+    elif model.peft_config.peft_type == PeftType.BOTTLENECK:
+        # return the state dict of the model with Bottleneck adapters
+        bias = model.peft_config.bias
+        if bias == "none":
+            to_return = {k: state_dict[k] for k in state_dict if "adapter_" in k}
+        elif bias == "all":
+            to_return = {k: state_dict[k] for k in state_dict if "adapter_" in k or "bias" in k}
+        elif bias == "adapter_only":
+            to_return = {}
+            for k in state_dict:
+                if "adapter_" in k:
+                    to_return[k] = state_dict[k]
+                    bias_name = k.split("adapter_")[0] + "bias"
+                    if bias_name in state_dict:
+                        to_return[bias_name] = state_dict[bias_name]
+        else:
+            raise NotImplementedError
+        
+    elif model.peft_config.peft_type == PeftType.ROTATION:
+        bias = model.peft_config.bias
+        if bias == "none":
+            to_return = {k: state_dict[k] for k in state_dict if "rotation" in k}
+        elif bias == "all":
+            to_return = {k: state_dict[k] for k in state_dict if "rotation" in k or "bias" in k}
+        elif bias == "rotation_only":
+            to_return = {}
+            for k in state_dict:
+                if "rotation" in k:
+                    to_return[k] = state_dict[k]
+                    bias_name = k.split("rotation")[0] + "bias"
+                    if bias_name in state_dict:
+                        to_return[bias_name] = state_dict[bias_name]
+        else:
+            raise NotImplementedError
+        
+    elif model.peft_config.is_prompt_learning:
+        to_return = {}
+        if model.peft_config.inference_mode:
+            prompt_embeddings = model.prompt_encoder.embedding.weight
+        else:
+            prompt_embeddings = model.get_prompt_embedding_to_save()
+        to_return["prompt_embeddings"] = prompt_embeddings    
+    else:
+        raise NotImplementedError
+    if model.modules_to_save is not None:
+        for key, value in state_dict.items():
+            if any(module_name in key for module_name in model.modules_to_save):
+                to_return[key] = value
+    return to_return
+
+
+def set_peft_model_state_dict(model, peft_model_state_dict,
+                              adapter_name="default",
+                              ignore_mismatched_sizes: bool = False):
+    """
+    Set the state dict of the Peft model.
+
+    Args:
+        model ([`PeftModel`]): The Peft model.
+        peft_model_state_dict (`dict`): The state dict of the Peft model.
+        adapter_name (`str`, *optional*, defaults to `"default"`):
+            The name of the adapter whose state dict should be set.
+    """
+    peft_model_state_dict, mismatched_keys = _find_mismatched_keys(
+            model, peft_model_state_dict, ignore_mismatched_sizes=ignore_mismatched_sizes
+        )
+    if mismatched_keys:
+        # see https://github.com/huggingface/transformers/blob/09f9f566de83eef1f13ee83b5a1bbeebde5c80c1/src/transformers/modeling_utils.py#L4039
+        mismatched_warning = "\n".join(
+            [
+                f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
+                for key, shape1, shape2 in mismatched_keys
+            ]
+        )
+        msg = (
+            f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint "
+            f"and are being ignored because you passed `ignore_mismatched_sizes=True`: {mismatched_warning}."
+        )
+        warnings.warn(msg)
+
+    model.load_state_dict(peft_model_state_dict, strict=False)
+    if model.peft_config.peft_type != PeftType.LORA and model.peft_config.peft_type != PeftType.BOTTLENECK \
+        and model.peft_config.peft_type != PeftType.ROTATION:
+        model.prompt_encoder.embedding.load_state_dict(
+            {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True
+        )
+    return model

nl_tasks/scripts/.nfs80e7f26e00566c630000664a

@@ -0,0 +1,117 @@
+# export OMINI_CONFIG=./config/commonsense.yaml
+export OMINI_CONFIG=./config/commonsense.yaml
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include" 
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+#./exps/run_ex15_3ep
+# ./exp_init/run_ex02/ft2
+# ADAPTER = "--model.merge_adapter_path "./exps/run_ex12/ft2" --model.merge_output_path "./exps/run_ex12/merged"
+# export ADAPTER = "--model.merge_adapter_path ./exp395/run_ex01/ft2 --model.merge_output_path ./exp395/run_ex01/merged"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path "./exps/run_ex19_2ep/ft2" --model.merge_output_path "./exps/run_ex19_2ep/merged"
+
+# OUTPUT="./exps/run_ex19_2ep/merged"
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $OUTPUT
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $OUTPUT
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex24"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex25"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex26"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex27"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex28"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+MERGE_DIR="./exps/run_ex33"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+MERGE_DIR="./exps/run_ex34"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+
+
+

nl_tasks/scripts/.nfs80e7f26e0132942e00006649

@@ -0,0 +1,341 @@
+export OMINI_CONFIG=./config/commonsense.yaml
+
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include"
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+export WANDB_PROJECT="Llama2_7B_FT_Math40k_2"
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp_init/run_ex01" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text "init|kaim_out_u=v"
+
+# sleep 5 
+# echo "1st exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp_init/run_ex02" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text "init|kaim_out_u=v(ratio)"
+
+# sleep 5 
+# echo "2nd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex03" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+
+# sleep 5 
+# echo "3rd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex04" --trainer_args.learning_rate=2e-3 \
+#      --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+# sleep 5 
+# echo "4th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex05" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+# sleep 5 
+# echo "5th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex06" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+# sleep 5 
+# echo "6th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_exps7" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+# sleep 5 
+# echo "7th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex08" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+# sleep 5 
+# echo "8th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex09" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 16 --rotation_adapter_config.r 1
+
+# sleep 5 
+# echo "9th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex10" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 8 --rotation_adapter_config.r 2
+
+# sleep 5 
+# echo "10 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex11" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "11 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex12" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16  --run_text 'u=v,def'
+
+# sleep 5 
+# echo "12 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex13" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'u=vkaim'
+
+# sleep 5 
+# echo "13 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex14" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'a,b,def'
+
+# sleep 5 
+# echo "14 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex15" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex17" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex18" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "158exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex19" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "19 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex20" --trainer_args.learning_rate=8e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "20 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex21" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|2ep|1e3"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# back to official 40k
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex22" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|1e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex23" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|1e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex24" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|1e2|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex25" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|2e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex26" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|5e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex27" --trainer_args.learning_rate=8e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|8e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex28" --trainer_args.learning_rate=2e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|2e2|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex29" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex30" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex31" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex32" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex33" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex34" --trainer_args.learning_rate=2e-2 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+      --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+sleep 5 
+echo "29 exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+bash scripts/merge.sh

nl_tasks/scripts/copy train_cms_reasoning.sh

@@ -0,0 +1,133 @@
+export OMINI_CONFIG=./config/commonsense.yaml
+
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include"
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+export WANDB_PROJECT="Llama2_7B_FT_Math40k"
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps1" --trainer_args.learning_rate=1e-5
+
+# sleep 5 
+# echo "1st exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps2" --trainer_args.learning_rate=2e-5
+
+# sleep 5 
+# echo "2nd exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps3" --trainer_args.learning_rate=5e-5
+
+# sleep 5 
+# echo "3rd exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps4" --trainer_args.learning_rate=1e-4
+
+# sleep 5 
+# echo "4th exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps5" --trainer_args.learning_rate=2e-4
+
+# sleep 5 
+# echo "5th exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps6" --trainer_args.learning_rate=5e-4
+
+# sleep 5 
+# echo "6th exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps7" --trainer_args.learning_rate=8e-4
+
+# sleep 5 
+# echo "7th exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps8" --trainer_args.learning_rate=1e-3
+
+# sleep 5 
+# echo "8th exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exps9" --trainer_args.learning_rate=2e-3
+
+# sleep 5 
+# echo "9th exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exnr10" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 8 --rotation_adapter_config.r 2
+
+# sleep 5 
+# echo "10 exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exnr11" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+
+# sleep 5 
+# echo "11 exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exnr12" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 16 --rotation_adapter_config.r 1
+
+# sleep 5 
+# echo "12 exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_exnr13" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "13 exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_all/exnr14" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "14 exp finishes"
+# date +"%F %T"
+
+accelerate launch --main_process_port 41353 -m src.ft_mathR \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./run_all/exnr15" --trainer_args.learning_rate=1e-3 \
+      --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+
+# sleep 5 
+echo "15 exp finishes"
+date +"%F %T"

nl_tasks/scripts/down_math_train.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+DATASET_ID="meta-math/MetaMathQA"
+LOCAL_DIR="./data/MetaMathQA"
+
+# echo "Starting download for dataset: $DATASET_ID..."
+huggingface-cli download $DATASET_ID \
+    --repo-type dataset \
+    --local-dir $LOCAL_DIR \
+    --local-dir-use-symlinks False \
+    --resume-download \
+    --include "*.json"
+
+# echo "Download completed. Data is located at: $LOCAL_DIR"

nl_tasks/scripts/inference.sh

@@ -0,0 +1,14 @@
+# OUTPUT="./exps/run_ex12/merged"
+# OUTPUT="./exp395/run_ex02/merged"
+# OUTPUT="./exp_init/run_ex02/merged"
+OUTPUT="./exps/run_ex15_3ep/merged"
+
+date +"%F %T"
+
+echo 'test math'
+
+date +"%F %T"
+python inference/gsm8k_infer.py --model $OUTPUT
+date +"%F %T"
+python inference/MATH_infer.py --model $OUTPUT
+date +"%F %T"

nl_tasks/scripts/merge.sh

@@ -0,0 +1,137 @@
+# export OMINI_CONFIG=./config/commonsense.yaml
+export OMINI_CONFIG=./config/commonsense.yaml
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include" 
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+#./exps/run_ex15_3ep
+# ./exp_init/run_ex02/ft2
+# ADAPTER = "--model.merge_adapter_path "./exps/run_ex12/ft2" --model.merge_output_path "./exps/run_ex12/merged"
+# export ADAPTER = "--model.merge_adapter_path ./exp395/run_ex01/ft2 --model.merge_output_path ./exp395/run_ex01/merged"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path "./exps/run_ex19_2ep/ft2" --model.merge_output_path "./exps/run_ex19_2ep/merged"
+
+# OUTPUT="./exps/run_ex19_2ep/merged"
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $OUTPUT
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $OUTPUT
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex24"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex25"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex26"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex27"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex28"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+
+# MERGE_DIR="./exps/run_ex33"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+# python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+# date +"%F %T"
+
+# 140126
+MERGE_DIR="./exprep/run_ex30"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+MERGE_DIR="./exprep/run_ex31"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+MERGE_DIR="./exprep/run_ex32"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+

nl_tasks/scripts/merge_100k.sh

@@ -0,0 +1,100 @@
+# export OMINI_CONFIG=./config/commonsense.yaml
+export OMINI_CONFIG=./config/math395.yaml
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include" 
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+#./exps/run_ex15_3ep
+# ./exp_init/run_ex02/ft2
+# ADAPTER = "--model.merge_adapter_path "./exps/run_ex12/ft2" --model.merge_output_path "./exps/run_ex12/merged"
+# export ADAPTER = "--model.merge_adapter_path ./exp395/run_ex01/ft2 --model.merge_output_path ./exp395/run_ex01/merged"
+
+MERGE_DIR="./exp100/run_ex06"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+MERGE_DIR="./exp100/run_ex07"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+
+MERGE_DIR="./exp100/run_ex08"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+
+MERGE_DIR="./exp100/run_ex09"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+
+MERGE_DIR="./exp100/run_ex10"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+
+MERGE_DIR="./exp100/run_ex11"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+
+MERGE_DIR="./exp100/run_ex12"
+accelerate launch --main_process_port 41353 -m src.merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"

nl_tasks/scripts/merge_math.sh

@@ -0,0 +1,31 @@
+# export OMINI_CONFIG=./config/commonsense.yaml
+export OMINI_CONFIG=./config/math395.yaml
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include" 
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+#./exps/run_ex15_3ep
+# ./exp_init/run_ex02/ft2
+# ADAPTER = "--model.merge_adapter_path "./exps/run_ex12/ft2" --model.merge_output_path "./exps/run_ex12/merged"
+# export ADAPTER = "--model.merge_adapter_path ./exp395/run_ex01/ft2 --model.merge_output_path ./exp395/run_ex01/merged"
+
+MERGE_DIR="./exp395/run_ex10"
+# accelerate launch --main_process_port 41353 -m src.merge \
+#       --config_path $OMINI_CONFIG \
+#       --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+# # OUTPUT="./exp395/run_ex09/merged"
+
+# date +"%F %T"
+# python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"

nl_tasks/scripts/peft_merge.sh

@@ -0,0 +1,60 @@
+# export OMINI_CONFIG=./config/commonsense.yaml
+export OMINI_CONFIG=./config/commonsense.yaml
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include" 
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+
+
+MERGE_DIR="./expsBOFT/seed43"
+accelerate launch --main_process_port 41353 -m src.peft_merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+MERGE_DIR="./expsBOFT/seed44"
+accelerate launch --main_process_port 41353 -m src.peft_merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+
+MERGE_DIR="./expsOFT/seed43"
+accelerate launch --main_process_port 41353 -m src.peft_merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+
+MERGE_DIR="./expsOFT/seed44"
+accelerate launch --main_process_port 41353 -m src.peft_merge \
+      --config_path $OMINI_CONFIG \
+      --model.merge_adapter_path $MERGE_DIR/ft2/ --model.merge_output_path $MERGE_DIR/merged/
+date +"%F %T"
+python inference/MATH_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+python inference/gsm8k_infer.py --model $MERGE_DIR/merged/
+date +"%F %T"
+

nl_tasks/scripts/train_100math.sh

@@ -0,0 +1,184 @@
+export OMINI_CONFIG=./config/math395.yaml
+
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include"
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+export WANDB_PROJECT="Llama2_7B_FT_Math_395k"
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex01" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+#       --run_text 'def|o100k'
+
+# sleep 5 
+# echo "1st exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex02" --trainer_args.learning_rate=2e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+#       --run_text 'def|o100k'
+
+# sleep 5 
+# echo "2nd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# bash scripts/merge_100k.sh
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex03" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+#       --run_text 'def|o100k'
+
+# sleep 5 
+# echo "3rd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex04" --trainer_args.learning_rate=5e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+#       --run_text 'def|o100k'
+
+# sleep 5 
+# echo "4th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex05" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train[:100000] \
+#       --run_text 'def|o100k'
+# sleep 5 
+# echo "5th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+# bash scripts/merge_100math.sh
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex06" --trainer_args.learning_rate=1e-2 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+      --run_text 'def|o100k|b48' --trainer_args.per_device_train_batch_size 48
+sleep 5 
+echo "6th exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex07" --trainer_args.learning_rate=1e-2 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 3.0  --data.dataset_split train[:100000] \
+      --run_text 'def|o100k|b48' --trainer_args.per_device_train_batch_size 48
+sleep 5 
+echo "6th exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex08" --trainer_args.learning_rate=2e-2 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+      --run_text 'def|o100k|b48' --trainer_args.per_device_train_batch_size 48
+sleep 5 
+echo "8th exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex09" --trainer_args.learning_rate=2e-2 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 3.0  --data.dataset_split train[:100000] \
+      --run_text 'def|o100k|b48' --trainer_args.per_device_train_batch_size 48
+sleep 5 
+echo "9th exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex10" --trainer_args.learning_rate=3e-2 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+      --run_text 'def|o100k|b48' --trainer_args.per_device_train_batch_size 48
+sleep 5 
+echo "10th exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex11" --trainer_args.learning_rate=8e-3 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 2.0  --data.dataset_split train[:100000] \
+      --run_text 'def|o100k|b48' --trainer_args.per_device_train_batch_size 48
+
+sleep 5 
+echo "11 exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp100/run_ex12" --trainer_args.learning_rate=8e-3 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 3.0  --data.dataset_split train[:100000] \
+      --run_text 'def|o100k|b48' --trainer_args.per_device_train_batch_size 48
+
+sleep 5 
+echo "12 exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+
+bash ./scripts/merge_100k.sh
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex13" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'u=vkaim'
+
+# sleep 5 
+# echo "13 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex14" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'a,b,def'
+
+# sleep 5 
+# echo "14 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex15" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"

nl_tasks/scripts/train_cms_reasoning.sh

@@ -0,0 +1,260 @@
+export OMINI_CONFIG=./config/commonsense.yaml
+
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include"
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+export WANDB_PROJECT="Llama2_7B_FT_Math40k_2"
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex01" --trainer_args.learning_rate=5e-5 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+
+# sleep 5 
+# echo "1st exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex02" --trainer_args.learning_rate=5e-4 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+
+# sleep 5 
+# echo "2nd exp finishes"
+# date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex03" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+
+# sleep 5 
+# echo "3rd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex04" --trainer_args.learning_rate=2e-3 \
+#      --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+# sleep 5 
+# echo "4th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex05" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+# sleep 5 
+# echo "5th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex06" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+# sleep 5 
+# echo "6th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_exps7" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+# sleep 5 
+# echo "7th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex08" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+# sleep 5 
+# echo "8th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex09" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 16 --rotation_adapter_config.r 1
+
+# sleep 5 
+# echo "9th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex10" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 8 --rotation_adapter_config.r 2
+
+# sleep 5 
+# echo "10 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex11" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "11 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex12" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16  --run_text 'u=v,def'
+
+# sleep 5 
+# echo "12 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+### continue with 40k
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex13_3ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 200 
+
+# sleep 5 
+# echo "13 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex14_3ep" --trainer_args.learning_rate=2e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 200
+
+# sleep 5 
+# echo "14 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex15_3ep" --trainer_args.learning_rate=5e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 200
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex16_3ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.05' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 200
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex17_3ep" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.05' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 200
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex18_2ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 200
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex19_2ep" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 200
+
+# sleep 5 
+# echo "19 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+
+# 140126
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex20_2ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 11
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex21_2ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 10
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex24_3ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 100 --seed 10
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex25_3ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 100 --seed 12
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex22_2ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 12
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex23_3ep" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 100 --seed 11
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex26_2ep" --trainer_args.learning_rate=8e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 11
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex27_2ep" --trainer_args.learning_rate=8e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 10
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex28_2ep" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 11
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex29_2ep" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 10
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exprep/run_ex30" --trainer_args.learning_rate=8e-4 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+      --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 20
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exprep/run_ex31" --trainer_args.learning_rate=8e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 21
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exprep/run_ex32" --trainer_args.learning_rate=8e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'init=def|dr0.10' \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 100 --seed 22

nl_tasks/scripts/train_initn40k.sh

@@ -0,0 +1,341 @@
+export OMINI_CONFIG=./config/commonsense.yaml
+
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include"
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+export WANDB_PROJECT="Llama2_7B_FT_Math40k_2"
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp_init/run_ex01" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text "init|kaim_out_u=v"
+
+# sleep 5 
+# echo "1st exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp_init/run_ex02" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text "init|kaim_out_u=v(ratio)"
+
+# sleep 5 
+# echo "2nd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex03" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+
+# sleep 5 
+# echo "3rd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex04" --trainer_args.learning_rate=2e-3 \
+#      --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+# sleep 5 
+# echo "4th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex05" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 2 --rotation_adapter_config.r 8
+# sleep 5 
+# echo "5th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex06" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+# sleep 5 
+# echo "6th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_exps7" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+# sleep 5 
+# echo "7th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex08" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+# sleep 5 
+# echo "8th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex09" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 16 --rotation_adapter_config.r 1
+
+# sleep 5 
+# echo "9th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex10" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 8 --rotation_adapter_config.r 2
+
+# sleep 5 
+# echo "10 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex11" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "11 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex12" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16  --run_text 'u=v,def'
+
+# sleep 5 
+# echo "12 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex13" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'u=vkaim'
+
+# sleep 5 
+# echo "13 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex14" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'a,b,def'
+
+# sleep 5 
+# echo "14 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex15" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex17" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex18" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "158exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex19" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "19 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex20" --trainer_args.learning_rate=8e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|fix_token"
+# sleep 5 
+# echo "20 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex21" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 500  --data.dataset_split train[:41023] --data.split_ratio 0.02493 \
+#       --run_text "dropout|2ep|1e3"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# back to official 40k
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex22" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|1e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex23" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|1e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex24" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|1e2|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex25" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|2e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex26" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|5e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex27" --trainer_args.learning_rate=8e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|8e3|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex28" --trainer_args.learning_rate=2e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|2e2|40k"
+# sleep 5 
+# echo "21 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex29" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex30" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 2.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex31" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex32" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex33" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+#       --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+# sleep 5 
+# echo "29 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+
+accelerate launch --main_process_port 41353 -m src.ft_mathQ \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exps/run_ex34" --trainer_args.learning_rate=2e-2 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 3.0  --data.dataset_split train \
+      --run_text "drop0.1|2ep|initu=v=0.01|40k" --trainer_args.per_device_train_batch_size 48
+sleep 5 
+echo "29 exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+bash scripts/merge.sh

nl_tasks/scripts/train_math.sh

@@ -0,0 +1,162 @@
+export OMINI_CONFIG=./config/math395.yaml
+
+#echo $OMINI_CONFIG
+export TOKENIZERS_PARALLELISM=true
+
+# CUDA Include (/cuda.h)
+CUDA_INCLUDE_PATH="/home/work/miniconda3/envs/allm/include"
+
+# 3. Add into CPATH & CPLUS_INCLUDE_PATH (C/C++ compiler)
+export CPATH=$CPATH:$CUDA_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$CUDA_INCLUDE_PATH
+# echo "CPATH is set to: $CPATH"
+# echo "CPLUS_INCLUDE_PATH is set to: $CPLUS_INCLUDE_PATH"
+
+export WANDB_PROJECT="Llama2_7B_FT_Math_395k"
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+date +"%F %T"
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex01" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'def'
+
+# sleep 5 
+# echo "1st exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex02" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "2nd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex03" --trainer_args.learning_rate=2e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "3rd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex04" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "4rd exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex05" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500 --data.dataset_split train[:101011] --data.split_ratio 0.01
+# sleep 5 
+# echo "5th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex06" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500 --data.dataset_split train[:101011] --data.split_ratio 0.01
+# sleep 5 
+# echo "6th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex0s7" --trainer_args.learning_rate=5e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500 --data.dataset_split train[:101011] --data.split_ratio 0.01
+# sleep 5 
+# echo "7th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex08" --trainer_args.learning_rate=1e-4 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+#       --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500 --data.dataset_split train[:101011] --data.split_ratio 0.01 \
+#       --trainer_args.per_device_train_batch_size 32 --run_text 'u2e2,def'
+# sleep 5 
+# echo "8th exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+accelerate launch --main_process_port 41353 -m src.ft_mathR \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex09" --trainer_args.learning_rate=2e-3 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 3.0 --trainer_args.eval_steps 500 --data.dataset_split train[:101011] --data.split_ratio 0.01 \
+      --trainer_args.per_device_train_batch_size 32 --run_text 'init=def|fix_token'
+
+sleep 5 
+echo "9th exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+accelerate launch --main_process_port 41353 -m src.ft_mathR \
+      --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex10" --trainer_args.learning_rate=2e-3 \
+      --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 \
+      --trainer_args.num_train_epochs 2.0 --trainer_args.eval_steps 500 --data.dataset_split train[:101011] --data.split_ratio 0.01 \
+      --trainer_args.per_device_train_batch_size 32 --run_text "init=def|fix_token"
+
+sleep 5 
+echo "10 exp finishes"
+date +"%F %T"
+wandb sync wandb/latest-run
+
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex11" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16
+
+# sleep 5 
+# echo "11 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex12" --trainer_args.learning_rate=1e-2 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16  --run_text 'u=v,def'
+
+# sleep 5 
+# echo "12 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex13" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'u=vkaim'
+
+# sleep 5 
+# echo "13 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex14" --trainer_args.learning_rate=2e-3 \
+#       --rotation_adapter_config.num_rotations 1 --rotation_adapter_config.r 16 --run_text 'a,b,def'
+
+# sleep 5 
+# echo "14 exp finishes"
+# date +"%F %T"
+# wandb sync wandb/latest-run
+
+# accelerate launch --main_process_port 41353 -m src.ft_mathR \
+#       --config_path $OMINI_CONFIG --trainer_args.output_dir "./exp395/run_ex15" --trainer_args.learning_rate=1e-3 \
+#       --rotation_adapter_config.num_rotations 4 --rotation_adapter_config.r 4
+
+# sleep 5 
+# echo "15 exp finishes"
+# date +"%F %T"

nl_tasks/setup.py

@@ -0,0 +1,28 @@
+# Copyright [2024] [Zhuo Chen]
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import setuptools
+
+setuptools.setup(
+    name="rpeft",
+    version="0.0.2",
+    author="SDML",
+    packages=setuptools.find_packages(),
+    install_requires=[
+        'transformers>=4.0.0',
+        'torch>=2.0.0'
+    ],
+    python_requires='>=3.9',
+)

nl_tasks/src/config.py

@@ -0,0 +1,183 @@
+from dataclasses import dataclass, field, fields, asdict
+from typing import Optional, List, Literal, Dict, Any, Union
+from transformers import TrainingArguments, Trainer
+from omegaconf import OmegaConf
+import sys
+
+
+@dataclass
+class ModelConfig:
+    model_name: str = ""
+    dropout: float = 0.0
+    model_max_seq_length: int = field(default=512)
+    data_collator_mode: str=field(default='fixed', metadata={"help": "fixed or dynamic padding in DataCollator"})  
+    lambda_reg: float = field(default=1e-4, metadata={"help": "The control strength of regularity"})
+    adapter_path: Optional[str] = field(default=None)
+
+    merge_adapter_path: Optional[str] = field(default=None)
+    merge_output_path: Optional[str] = field(default=None)
+
+@dataclass
+class RotationConfig:
+    r: int = field(default=4)
+    num_rotations: int = field(default=4)
+    task_type: str = "CAUSAL_LM"
+    target_modules: List[str] = field(default_factory=lambda: ["q_proj",])
+
+@dataclass
+class DataConfig:
+    dataset_name: str = 'math'
+    split_ratio: float = field(default=0.01)
+    path: str = "./nl_tasks/data/MetaMathQA-40K/MetaMathQA-40K.json"
+    dataset_split: str = field(default="train[:1000]", metadata={"help": "(`['train', 'test', 'eval']`):"})
+    adapter_names: List[Optional[str]] = field(default_factory=lambda: ["default"])   ###
+    dataset_field: List[str] = field(default_factory=list, metadata={"help": "Fields of dataset input and output."})
+
+
+@dataclass
+class TrainingOverride:
+    optim: str=field(default="adamw_torch")   ##
+    eval_strategy: str=field(default='no')
+    per_device_train_batch_size: int=field(default=8) ##
+    per_device_eval_batch_size: int=field(default=8)  ##
+
+    learning_rate: float = field(default=1e-05)
+    lr_scheduler_type: str = field(default='cosine')
+    # warmup_ratio: float = field(default=0.1)
+    warmup_steps: int = field(default=0)
+    
+    gradient_checkpointing: bool = field(default=False)
+    gradient_accumulation_steps: int=field(default=1)
+    output_dir: str = field(default="runs")
+    save_steps: float = field(default=0)
+    save_strategy: str =field(default='no')
+    # save_total_limit: int=field(default=1) No need any more
+    bf16: bool=field(default=False)
+    bf16_full_eval: bool=field(default=False)
+    save_safetensors: bool=field(default=False)
+
+    report_to: Union[None, str, list[str]]=field(default="none")
+    logging_steps: int=field(default=25) # we use int only
+    # logging_first_step: bool=field(default=False)
+    eval_steps: Union[None,int]=field(default=None)  # we use int only f
+
+    dataloader_num_workers: int = field(default=1)
+    dataloader_pin_memory: bool = field(default=True)  ###
+    dataloader_persistent_workers: bool=field(default=True) ###
+    dataloader_prefetch_factor: int = field(default=1) ###
+
+    num_train_epochs: float = field(default=1.0)
+    max_steps: int=field(default=-1)
+    load_best_model_at_end: bool = field(default=True)
+
+@dataclass
+class GlueConfig:
+    task_name: str = field(default='mnli')
+    pad_to_max_length: bool = field(default=True)
+
+
+@dataclass
+class MainConfig:
+    model: ModelConfig = field(default_factory=ModelConfig)
+    rotation_adapter_config: RotationConfig = field(default_factory=RotationConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+    trainer_args: TrainingOverride = field(default_factory=TrainingOverride)
+
+    glue: GlueConfig = field(default_factory=GlueConfig)
+    project_name: str = "llm_rotation"
+    seed: int = 42
+    run_text: str=field(default='def')
+    # device: str = field(default='cpu')
+
+@dataclass
+class HFTrainingArguments(TrainingArguments):
+    extension: Optional[Dict[str, Any]] = field(
+        default=None, 
+        metadata={"help": "Serialized MainConfig excluding training args"}
+    )
+
+def convert_to_trainer_args(main_cfg: MainConfig) -> HFTrainingArguments:
+    """
+    Maps MainConfig to MyTrainingArguments.
+    Logic:
+    1. Extract 'training' fields -> Pass to TrainingArguments constructor.
+    2. Pack 'model', 'data', etc. -> Put into 'extension'.
+    """
+    KEY = "trainer_args"
+    # 1. Convert OmegaConf/Dataclass to pure Python dict
+    # resolve=True ensures variables like ${model.name} are interpolated
+    full_dict = asdict(main_cfg)
+    
+    # 2. Extract Training Arguments
+    # These will be unpack **kwargs to initialize the parent TrainingArguments
+    train_args_dict = full_dict.pop(KEY)
+    
+    # 3. The rest (model, data, seed) goes into extension
+    extension_payload = full_dict
+    
+    # 4. Initialize MyTrainingArguments
+    # Note: We must ensure train_args_dict keys match TrainingArguments fields.
+    try:
+        args = HFTrainingArguments(**train_args_dict)
+    except TypeError as e:
+        print(f"Error: Your 'training' config contains keys unknown to HF TrainingArguments: {e}")
+        sys.exit(1)
+        
+    # 5. Attach the extension
+    args.extension = extension_payload
+    
+    return args
+
+
+
+
+@dataclass
+class Training:
+    model_name_or_path: Optional[str] = field(default="huggyllama/llama-7b")
+    adapter_name_or_path: Optional[str] = field(default=None)
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+    dataset_split: str = field(
+        default="train[:100000]", metadata={"help": "(`['train', 'test', 'eval']`):"}
+    )
+    dataset_field: List[str] = field(
+        default=None, metadata={"help": "Fields of dataset input and output."}
+    )
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(default=512, metadata={
+        "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."}, )
+    hrft_r: int = field(default=8, metadata={
+        "help": "The rank of the adapter. When passing `None` and `adapter_name_or_path` is also `None`, full fine-tuning is used."})
+    init_a: float = field(default=1e-4, metadata={"help": "The initial weights"})
+    eps: float = field(default=1e-4, metadata={"help": "The control strength of COFT. The freedom of rotation."})
+    lamda: float = field(default=1e-4, metadata={"help": "The control strength of regularity"})
+    add_orth: str = field(default='none', metadata={"help": ""})
+    init_weights: Literal[True, "pissa"] = field(
+        default=True,
+        metadata={
+            "help": (
+                "Passing True (default) results in the LoRA initialization."
+                "Passing `pissa` results in PiSSA initialization."
+            ),
+        },
+    )
+    extension: Optional[Dict[str, Any]] = field(
+        default=None, 
+        metadata={"help": "Serialized MainConfig excluding training args"}
+    )
+
+    # target_modules: str = (
+    #     "(.*x_embedder"
+    #     "|.*(?<!single_)transformer_blocks\\.[0-9]+\\.norm1\\.linear"
+    #     "|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_k"
+    #     "|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_q"
+    #     "|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_v"
+    #     "|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_out\\.0"
+    #     "|.*(?<!single_)transformer_blocks\\.[0-9]+\\.ff\\.net\\.2"
+    #     "|.*single_transformer_blocks\\.[0-9]+\\.norm\\.linear"
+    #     "|.*single_transformer_blocks\\.[0-9]+\\.proj_mlp"
+    #     "|.*single_transformer_blocks\\.[0-9]+\\.proj_out"
+    #     "|.*single_transformer_blocks\\.[0-9]+\\.attn.to_k"
+    #     "|.*single_transformer_blocks\\.[0-9]+\\.attn.to_q"
+    #     "|.*single_transformer_blocks\\.[0-9]+\\.attn.to_v"
+    #     "|.*single_transformer_blocks\\.[0-9]+\\.attn.to_out)"
+    # )

nl_tasks/src/ft_mathQ.py

@@ -0,0 +1,702 @@
+#
+import sys
+#print('sys.path: ___ ', sys.path)
+#print(f"Current Python Executable: {sys.executable}")
+
+### dynamo warning
+import warnings
+
+# Ignore FutureWarning: prims_common.check, Online Softmax
+warnings.filterwarnings("ignore", category=FutureWarning, module='torch._inductor.lowering')
+warnings.filterwarnings("ignore", message=".*Online softmax is disabled on the fly.*", category=UserWarning)
+
+warnings.filterwarnings("ignore", message=".*Our suggested max number of worker in current system is 1.*", category=UserWarning)
+warnings.filterwarnings("ignore", message=".*will be initialized from a multivariate normal distribution.*")
+warnings.filterwarnings("ignore", message=".*that differ from the model config and generation config.*", category=UserWarning)
+warnings.filterwarnings("ignore", message=".*torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch..*", category=UserWarning)
+
+import torch
+torch.backends.cuda.matmul.fp32_precision = 'tf32'
+# import wandb
+import os
+torch.set_num_threads(1)
+os.environ["OMP_NUM_THREADS"]="1"
+os.environ["MKL_NUM_THREADS"]="1"
+import torch
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"PyTorch built with CUDA version: {torch.version.cuda}")
+
+import yaml
+#from peft import LoraConfig, get_peft_model_state_dict
+from torch.utils.data import DataLoader
+import time
+from datetime import datetime
+import math
+
+from typing import List, Tuple
+
+# import prodigyopt
+
+
+###
+import copy
+from dataclasses import field, dataclass, asdict
+from typing import Sequence, Literal, Dict
+
+import transformers
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+from transformers import Trainer
+from transformers.modeling_utils import *
+from transformers.trainer import _is_peft_model
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.data.data_collator import DataCollator
+
+from transformers.training_args import TrainingArguments
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.trainer_callback import TrainerCallback
+from transformers.trainer_utils import EvalPrediction
+from torch.utils.data import Dataset, IterableDataset
+from datasets import load_dataset
+##
+#from ..pipeline.flux_omini import transformer_forward, encode_images
+# from ...omini.rotation import RotationTuner, RotationConfig
+
+
+from rpeft.rotation import RotationTuner, RotationConfig
+from rpeft import get_peft_model, PeftModel
+from .config import MainConfig, convert_to_trainer_args
+import pyrallis
+from omegaconf import OmegaConf
+import torch.optim as optim
+import wandb
+from torch.nn.utils.rnn import pad_sequence
+
+IGNORE_INDEX = -100
+PROMPT = (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+)
+
+def get_rank():
+    try:
+        rank = int(os.environ.get("LOCAL_RANK"))
+    except:
+        rank = 0
+    return rank
+
+
+def get_config():
+    config_path = os.environ.get("OMINI_CONFIG")
+    assert config_path is not None, "Please set the OMINI_CONFIG environment variable"
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def init_wandb(wandb_config, run_name):
+    import wandb
+
+    try:
+        assert os.environ.get("WANDB_API_KEY") is not None
+        wandb.init(
+            project=wandb_config["project"],
+            name=run_name,
+            config={},
+        )
+    except Exception as e:
+        print("Failed to initialize WanDB:", e)
+
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+        
+
+def smart_tokenizer_and_embedding_resize(
+        special_tokens_dict: Dict,
+        tokenizer: transformers.PreTrainedTokenizer,
+        model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+def preprocess(
+        sources: Sequence[str],
+        targets: Sequence[str],
+        tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+
+# @dataclass
+# class DataCollatorForSupervisedDataset():
+#     """Collate examples for supervised fine-tuning."""
+
+#     tokenizer: transformers.PreTrainedTokenizer
+#     max_length: int = field(default=512)
+#     mode: str = field(default="fixed") # dynamic -> dynamo
+
+#     def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+#         if self.mode == 'dynamic':
+#             input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+#             input_ids = [torch.tensor(x) for x in input_ids]
+#             input_ids = torch.nn.utils.rnn.pad_sequence(
+#                 input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+#             )
+#             labels = [torch.tensor(x) for x in labels]
+#             labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+#             return dict(
+#                 input_ids=input_ids,
+#                 labels=labels,
+#                 attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+#             )
+#         elif self.mode == 'fixed':
+#             input_ids = [torch.tensor(x["input_ids"][:self.max_length]) for x in instances]
+#             input_ids = torch.stack([
+#                 torch.nn.functional.pad(x, (0, self.max_length - x.size(0)), value=self.tokenizer.pad_token_id)
+#                 for x in input_ids
+#             ])
+
+#             # Labels
+#             labels = [torch.tensor(x["labels"][:self.max_length]) for x in instances]
+#             labels = torch.stack([
+#                 torch.nn.functional.pad(x, (0, self.max_length - x.size(0)), value=IGNORE_INDEX)
+#                 for x in labels
+#             ])
+
+#             return dict(
+#                 input_ids=input_ids,
+#                 labels=labels,
+#                 attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+#             )
+#         else:
+#             raise NotImplementedError
+
+# @dataclass
+# class DataCollatorForSupervisedDataset(object):
+#     tokenizer: transformers.PreTrainedTokenizer
+#     max_length: int = field(default=512)
+#     mode: str = field(default="fixed")  # "dynamic" or "fixed"
+
+#     def _pad_to_length(self, tensors: Sequence[torch.Tensor], pad_value: int, target_len: int):
+#         """Pad a list of 1D tensors to target_len (int) and stack -> (B, target_len)."""
+#         batch_size = len(tensors)
+#         out = torch.full((batch_size, target_len), pad_value, dtype=tensors[0].dtype)
+#         for i, t in enumerate(tensors):
+#             L = min(t.size(0), target_len)
+#             out[i, :L] = t[:L]
+#         return out
+
+#     def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+#         # Collect raw sequences (lists or tensors)
+#         input_seqs = [torch.tensor(x["input_ids"], dtype=torch.long) for x in instances]
+#         label_seqs = [torch.tensor(x["labels"], dtype=torch.long) for x in instances]
+
+#         if self.mode == "dynamic":
+#             # pad to the max length present in this batch (<= self.max_length)
+#             batch_max_len = min(max([s.size(0) for s in input_seqs]), self.max_length)
+#             input_ids = self._pad_to_length(input_seqs, pad_value=self.tokenizer.pad_token_id, target_len=batch_max_len)
+#             labels = self._pad_to_length(label_seqs, pad_value=IGNORE_INDEX, target_len=batch_max_len)
+#         elif self.mode == "fixed":
+#             # always pad/truncate to self.max_length
+#             input_ids = self._pad_to_length(input_seqs, pad_value=self.tokenizer.pad_token_id, target_len=self.max_length)
+#             labels = self._pad_to_length(label_seqs, pad_value=IGNORE_INDEX, target_len=self.max_length)
+#         else:
+#             raise NotImplementedError(f"Unknown mode: {self.mode}")
+
+#         attention_mask = input_ids.ne(self.tokenizer.pad_token_id).long()
+
+#         return {
+#             "input_ids": input_ids,
+#             "labels": labels,
+#             "attention_mask": attention_mask
+#         }
+    
+@dataclass
+class DataCollatorForSupervisedDataset():
+    tokenizer: transformers.PreTrainedTokenizer
+    max_length: int = field(default=512)
+    mode: str = field(default="fixed")  # "dynamic" or "fixed"
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        # Extract inputs and labels
+        # Assuming instances is a list of dicts like {'input_ids': [...], 'labels': [...]}
+        input_ids_list = [torch.tensor(x["input_ids"], dtype=torch.long) for x in instances]
+        labels_list = [torch.tensor(x["labels"], dtype=torch.long) for x in instances]
+
+        # 1. Determine padding logic
+        if self.mode == "dynamic":
+            # Dynamic padding: pad to the longest sequence in the batch
+            # But cap it at self.max_length to prevent OOM
+            batch_max_len = max([len(x) for x in input_ids_list])
+            target_len = min(batch_max_len, self.max_length)
+        else:
+            # Fixed padding: always pad to max_length
+            target_len = self.max_length
+
+        # 2. Helper to pad and truncate
+        def pad_and_truncate(tensors, padding_value):
+            # First, pad everything using PyTorch's optimized utility (batch_first=True)
+            padded = pad_sequence(tensors, batch_first=True, padding_value=padding_value)
+            
+            # Handle truncation/extending to exact target_len
+            curr_len = padded.shape[1]
+            if curr_len > target_len:
+                # Truncate if too long (rare if filtered beforehand)
+                return padded[:, :target_len]
+            elif curr_len < target_len:
+                # Pad more if shorter than target_len (happens in fixed mode)
+                diff = target_len - curr_len
+                padding = torch.full((padded.shape[0], diff), padding_value, dtype=padded.dtype)
+                return torch.cat([padded, padding], dim=1)
+            else:
+                return padded
+
+        # 3. Apply padding
+        # Critical: tokenizer.pad_token_id must NOT be None here
+        if self.tokenizer.pad_token_id is None:
+            raise ValueError("Tokenizer.pad_token_id is None. Please set it to eos_token_id or unk_token_id.")
+            
+        input_ids = pad_and_truncate(input_ids_list, self.tokenizer.pad_token_id)
+        labels = pad_and_truncate(labels_list, IGNORE_INDEX)
+
+        # 4. Create Attention Mask explicitly
+        # .ne() creates Bools, .long() casts to 0s and 1s for compatibility
+        attention_mask = input_ids.ne(self.tokenizer.pad_token_id).long()
+
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": attention_mask
+        }
+        
+def train_tokenize_function(examples, tokenizer, query, response):
+    sources = [PROMPT.format_map(dict(instruction=instruction)) for instruction in examples[query]]
+    targets = [f"{output}{tokenizer.eos_token}" for output in examples[response]]
+    data_dict = preprocess(sources, targets, tokenizer)
+    return data_dict
+
+
+
+### Trainer
+def default_worker_init_fn(worker_id):
+    # mỗi worker chỉ 1 thread cho BLAS
+    try:
+        import numpy as _np
+    except Exception:
+        _np = None
+    torch.set_num_threads(1)
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
+    os.environ.setdefault("MKL_NUM_THREADS", "1")
+    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
+    # Optional: bind CPU affinity per worker to avoid contention (NUMA-aware)
+    try:
+        cpu_count = os.cpu_count() or 1
+        # chia đều CPU cho workers
+        num_workers = getattr(torch.utils.data, "_num_workers", None)
+        # fallback: if not available, compute from environment variable or pass externally
+        # We'll do a simple round-robin assignment using worker_id
+        # assign a small mask of cores to this worker (e.g., chunk size 4)
+        chunk = max(1, cpu_count // max(1, min(64, cpu_count)))
+        start = (worker_id * chunk) % cpu_count
+        end = start + chunk
+        mask = set(range(start, min(end, cpu_count)))
+        try:
+            os.sched_setaffinity(0, mask)
+        except Exception:
+            pass
+    except Exception:
+        pass
+
+def set_seed(seed: int):
+    # random.seed(seed)
+    # np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    transformers.set_seed(seed)
+
+
+@pyrallis.wrap()
+def main(mainCfg: MainConfig):
+    #mainCfg = get_config()
+    #print(mainCfg)
+    print('='*120)
+    # print(OmegaConf.to_yaml(mainCfg))
+    # print('-'*40)
+    # 
+    # print((training_args))
+    set_seed(mainCfg.seed)
+    training_args = convert_to_trainer_args(mainCfg)
+
+    # wandb 
+    ENTITY = "nvan-13-korea-university" 
+    PROJECT = os.environ.get("WANDB_PROJECT")
+    api = wandb.Api()
+    try:
+        runs_list = api.runs(f"{ENTITY}/{PROJECT}")
+        next_run_num = len(runs_list) + 1
+    except Exception as e:
+        next_run_num = 1
+
+    training_args.run_name = f'[{next_run_num}]lr={mainCfg.trainer_args.learning_rate:.1e},b={mainCfg.trainer_args.per_device_train_batch_size},'\
+                            f'n={mainCfg.rotation_adapter_config.num_rotations},r={mainCfg.rotation_adapter_config.r},'\
+                            f'init={mainCfg.run_text}'
+    # training_args.project = f'Rotation-Llama2-{mainCfg.data.dataset_name}'
+
+    # print('-'*40)
+    # print(training_args.to_json_string())
+    # exit()
+    
+    model = AutoModelForCausalLM.from_pretrained(mainCfg.model.model_name,
+                                                 device_map="auto", low_cpu_mem_usage=True,
+                                                 dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+                                                 attn_implementation="sdpa",
+                                                 )
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    print("DEVICE", DEVICE)
+    # for name, param in model.named_parameters():
+    #     if 'q_proj' in name and 'layers.5' in name:
+    #         print(f"Name: {name} | {param.shape} ")
+            # print(f"Name (pretrained): {name} | {param.shape} | {param.data[0:5,0:5]}")
+    # print('model', model)
+    # exit()
+
+    total_params_now = sum(p.numel() for p in model.parameters())
+    print(f'#params of the pretrained model, {total_params_now:,}')
+    # print(model)
+    if mainCfg.model.adapter_path is not None:
+        print('___ Loading from:  ', mainCfg.model.adapter_path)
+        model = PeftModel.from_pretrained(model, mainCfg.model.adapter_path, is_trainable = True)
+    elif mainCfg.rotation_adapter_config.r is not None:
+        rotation_adapter_config = asdict(mainCfg.rotation_adapter_config)
+        # rotation_adapter_config[peft_type]
+        
+        for adapter_name in mainCfg.data.adapter_names:
+            rotation_config = RotationConfig(**rotation_adapter_config)
+            model = get_peft_model(model, rotation_config, adapter_name=adapter_name)
+            # model.set_adapter(adapter_name)
+
+        # import peft
+        # from peft import OFTConfig
+        # oft_config = OFTConfig(
+        #     # r=16,
+        #     oft_block_size=4*mainCfg.rotation_adapter_config.r,
+        #     use_cayley_neumann=True,
+        #     target_modules=["q_proj", "v_proj",],
+        #     module_dropout=0.05, # mainCfg.rotation_adapter_config.drop_out,
+        #     # task_type="CAUSAL_LM",
+        #     bias="none",
+        # )
+
+        # for adapter_name in mainCfg.data.adapter_names:
+        #     model = peft.get_peft_model(model, oft_config, adapter_name=adapter_name)
+    else:
+        print("Full Parameter Fine-Tuning")
+    model = model.to(DEVICE)
+        
+    # print('model', model)
+    model.print_trainable_parameters()
+    exit()
+    # print("Program starts")
+    # time.sleep(300)
+    # exit()
+
+    # for name, param in model.named_parameters():
+    #     if 'q_proj' in name and 'rotation' in name and 'layers.5' in name:
+    #         print(f"Name: {name} | {param.shape} ")
+    #         print(f"Name (pretrained): {name} | {param.shape} ")
+    #         X = param.data
+    # print('model', type(model), X.shape)
+    # visualize_value_distribution(X)
+    # exit()
+
+    rotation_layers = filter(
+                lambda p: p.requires_grad, model.parameters()
+            )    
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        mainCfg.model.model_name,
+        model_max_length=mainCfg.model.model_max_seq_length,
+        padding_side="right",
+        use_fast=True,
+    )
+
+    if tokenizer.pad_token is None:
+        if tokenizer.unk_token_id is not None:
+            tokenizer.pad_token_id = tokenizer.unk_token_id
+            tokenizer.pad_token = tokenizer.unk_token
+            print("Set PAD token to UNK token.")
+        elif tokenizer.eos_token_id is not None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+            tokenizer.pad_token = tokenizer.eos_token
+            print("Set PAD token to EOS token.")
+
+        if model is not None:
+            model.config.pad_token_id = tokenizer.pad_token_id
+            if model.config.pad_token_id != tokenizer.pad_token_id:
+                raise ValueError("Failed to sync pad_token_id between tokenizer and model config")
+
+    # local MetaMathQA-40K
+    raw_datasets = load_dataset("json", data_files=mainCfg.data.path, split=mainCfg.data.dataset_split)
+    #raw_train_datasets = load_dataset("MetaMathQA-40K", split=mainCfg.data.dataset_split)
+    # print('raw', type(raw_train_datasets), len(raw_train_datasets))
+
+    # split a single set
+    # split_ratio = mainCfg.data.split_ratio
+    # split_data = raw_datasets.train_test_split(test_size=split_ratio, seed=42)
+    # raw_train_datasets = split_data['train']
+    # raw_valid_datasets = split_data['test']
+
+    train_dataset = raw_datasets.map(
+        train_tokenize_function,
+        batched=True,
+        batch_size=30000,
+        num_proc=32,
+        remove_columns=raw_datasets.column_names,
+        load_from_cache_file=True,
+        desc="Running tokenizer on train dataset",
+        fn_kwargs={"tokenizer": tokenizer, "query": mainCfg.data.dataset_field[0],
+                   "response": mainCfg.data.dataset_field[1]}
+    )
+
+    # valid_dataset = raw_valid_datasets.map(
+    #     train_tokenize_function,
+    #     batched=True,
+    #     batch_size=30000,
+    #     num_proc=32,
+    #     remove_columns=raw_train_datasets.column_names,
+    #     load_from_cache_file=True,
+    #     desc="Running tokenizer on train dataset",
+    #     fn_kwargs={"tokenizer": tokenizer, "query": mainCfg.data.dataset_field[0],
+    #                "response": mainCfg.data.dataset_field[1]}
+    # )
+    print('- dataset size: ', len(train_dataset))
+
+
+    # print('dataset', type(train_dataset))
+    # print('process', len(train_dataset))
+    # print(f"Sample features: {train_dataset.column_names}, {train_dataset.num_rows}")
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer, max_length=mainCfg.model.model_max_seq_length, 
+                                                     #mode=mainCfg.model.data_collator_mode,
+                                                     )
+    data_module = dict(train_dataset=train_dataset, data_collator=data_collator)
+
+    optimizer = optim.AdamW(
+        rotation_layers, 
+        lr=mainCfg.trainer_args.learning_rate, #
+        eps=1e-8
+    )
+    # print('model x', model)
+    start_time = datetime.now()
+    print('start time: ', start_time.strftime("%Y-%m-%d %H:%M:%S"))
+    trainer = MyTrainer(model=model, processing_class=tokenizer,
+                        lamda=mainCfg.model.lambda_reg,
+                        optimizers=(optimizer, None),
+                        args=training_args, **data_module)
+    model.config.use_cache = False
+
+
+    # now = time.time()
+    # for i in range(20):
+    #     next(iter(trainer.get_train_dataloader()))
+    # print('time', time.time()-now)
+    # now = time.time()
+    
+
+    # dl = trainer.get_train_dataloader()
+    # t0 = time.time()
+    # for i, batch in enumerate(dl):
+    #     if i==20: break
+    # print("time / 20 batches =", time.time() - t0)
+    # exit()
+
+
+    # model2 = model.merge_and_unload()
+    # results2 = trainer2.evaluate()
+    # print('results2: ', results2)
+    # exit()
+
+    trainer.train()
+    
+    end_time = datetime.now()
+    print('end time: ', end_time.strftime("%Y-%m-%d %H:%M:%S"), '| duration: ', end_time - start_time)
+    # Save Model (Includes Adapter weights & Config)
+    # trainer.save_model(os.path.join(training_args.output_dir, 'ft'))
+    # Save Tokenizer
+    tokenizer.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
+    # Save Training State (Metrics & Logs)
+    trainer.save_state()
+
+    # save peft_config. Or model.base_model.peft_config['default']
+    model.peft_config.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
+
+    # the easiest way
+    model.save_pretrained(os.path.join(training_args.output_dir, 'ft2'))
+    return
+
+
+
+class MyTrainer(Trainer):
+
+    def __init__(
+            self,
+            model: Union[PreTrainedModel, nn.Module] = None,
+            args: TrainingArguments = None,
+            data_collator: Optional[DataCollator] = None,
+            train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None,
+            eval_dataset: Optional[Union[Dataset, Dict[str, Dataset], "datasets.Dataset"]] = None,
+            processing_class: Optional[PreTrainedTokenizerBase] = None,
+            model_init: Optional[Callable[[], PreTrainedModel]] = None,
+            compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+            callbacks: Optional[List[TrainerCallback]] = None,
+            optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+            preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
+            #run_name: Optional[str] = None,
+            #report_to: Optional[Union[str, list[str]]] = None,
+            # project
+            lamda: float = 1e-4
+    ):
+        super().__init__(model=model, args=args, data_collator=data_collator,
+                         train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=processing_class,
+                         model_init=model_init, compute_metrics=compute_metrics, callbacks=callbacks,
+                         optimizers=optimizers, preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+                         #run_name=run_name, report_to=report_to
+                         )
+        self.lamda = lamda
+
+    # def compute_loss(self, model, inputs, return_outputs=False,
+    #                  num_items_in_batch: Optional[torch.Tensor] = None,):
+    #     """
+    #     How the loss is computed by Trainer. By default, all models return the loss in the first element.
+
+    #     Subclass and override for custom behavior.
+    #     """
+    #     if self.label_smoother is not None and "labels" in inputs:
+    #         labels = inputs.pop("labels")
+    #     else:
+    #         labels = None
+    #     if self.model_accepts_loss_kwargs:
+    #         kwargs = {}
+    #         if num_items_in_batch is not None:
+    #             kwargs["num_items_in_batch"] = num_items_in_batch
+    #         inputs = {**inputs, **kwargs}
+    #     outputs = model(**inputs)
+    #     # Save past state if it exists
+    #     # TODO: this needs to be fixed and made cleaner later.
+    #     if self.args.past_index >= 0:
+    #         self._past = outputs[self.args.past_index]
+
+    #     if labels is not None:
+    #         unwrapped_model = unwrap_model(model)
+    #         if _is_peft_model(unwrapped_model):
+    #             model_name = unwrapped_model.base_model.model._get_name()
+    #         else:
+    #             model_name = unwrapped_model._get_name()
+    #         if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
+    #             loss = self.label_smoother(outputs, labels, shift_labels=True)
+    #         else:
+    #             loss = self.label_smoother(outputs, labels)
+    #     else:
+    #         if isinstance(outputs, dict) and "loss" not in outputs:
+    #             raise ValueError(
+    #                 "The model did not return a loss from the inputs, only the following keys: "
+    #                 f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
+    #             )
+    #         # We don't use .loss here since the model may return tuples instead of ModelOutput.
+    #         loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
+    #     # ------------------------------------------------------------------------------
+
+    #     # for name, param in model.named_parameters():
+    #     #     if 'oft_r' in name:
+    #     #         device = param.device
+    #     #         householder_U_norm = param / param.norm(dim=0)
+    #     #         orth_loss = torch.norm(
+    #     #             torch.eye(householder_U_norm.size(1), device=device) - householder_U_norm.t() @ householder_U_norm)
+    #     #         print(self.lamda)
+    #     #         loss = loss + self.lamda * orth_loss.to(loss.device)
+
+    #     # ------------------------------------------------------------------------------
+
+    #     return (loss, outputs) if return_outputs else loss
+
+    def get_train_dataloader(self):
+        # get dataset & sampler from super
+        train_dataset = self.train_dataset
+        sampler = self._get_train_sampler()
+
+        # compute effective batch size per step (HF has some routines; we use per_device_train_batch_size)
+        batch_size = self.args.train_batch_size if hasattr(self.args, "train_batch_size") else self.args.per_device_train_batch_size
+
+        # recommended num_workers: start moderate (16), you can tune upward
+        num_workers = getattr(self.args, "dataloader_num_workers", 16)
+        pin_memory = getattr(self.args, "dataloader_pin_memory", True)
+        prefetch_factor = getattr(self.args, "dataloader_prefetch_factor", 2)
+        persistent_workers = getattr(self.args, "dataloader_persistent_workers", True)
+
+        return DataLoader(
+            train_dataset,
+            batch_size=batch_size,
+            sampler=sampler,
+            collate_fn=self.data_collator,
+            drop_last=self.args.dataloader_drop_last if hasattr(self.args, "dataloader_drop_last") else False,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            persistent_workers=persistent_workers,
+            prefetch_factor=prefetch_factor,
+            worker_init_fn=default_worker_init_fn,
+        )
+    
+if __name__ == "__main__":
+    main()

nl_tasks/src/ft_mathR.py

@@ -0,0 +1,689 @@
+#
+import sys
+#print('sys.path: ___ ', sys.path)
+#print(f"Current Python Executable: {sys.executable}")
+
+### dynamo warning
+import warnings
+
+# Ignore FutureWarning: prims_common.check, Online Softmax
+warnings.filterwarnings("ignore", category=FutureWarning, module='torch._inductor.lowering')
+warnings.filterwarnings("ignore", message=".*Online softmax is disabled on the fly.*", category=UserWarning)
+
+warnings.filterwarnings("ignore", message=".*Our suggested max number of worker in current system is 1.*", category=UserWarning)
+warnings.filterwarnings("ignore", message=".*will be initialized from a multivariate normal distribution.*")
+warnings.filterwarnings("ignore", message=".*that differ from the model config and generation config.*", category=UserWarning)
+warnings.filterwarnings("ignore", message=".*torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch..*", category=UserWarning)
+
+import torch
+torch.backends.cuda.matmul.fp32_precision = 'tf32'
+# import wandb
+import os
+torch.set_num_threads(1)
+os.environ["OMP_NUM_THREADS"]="1"
+os.environ["MKL_NUM_THREADS"]="1"
+import torch
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"PyTorch built with CUDA version: {torch.version.cuda}")
+
+import yaml
+#from peft import LoraConfig, get_peft_model_state_dict
+from torch.utils.data import DataLoader
+import time
+from datetime import datetime
+import math
+
+from typing import List, Tuple
+
+# import prodigyopt
+
+
+###
+import copy
+from dataclasses import field, dataclass, asdict
+from typing import Sequence, Literal, Dict
+
+import transformers
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+from transformers import Trainer
+from transformers.modeling_utils import *
+from transformers.trainer import _is_peft_model
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.data.data_collator import DataCollator
+
+from transformers.training_args import TrainingArguments
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.trainer_callback import TrainerCallback
+from transformers.trainer_utils import EvalPrediction
+from torch.utils.data import Dataset, IterableDataset
+from datasets import load_dataset
+##
+#from ..pipeline.flux_omini import transformer_forward, encode_images
+# from ...omini.rotation import RotationTuner, RotationConfig
+
+
+from rpeft.rotation import RotationTuner, RotationConfig
+from rpeft import get_peft_model, PeftModel
+from .config import MainConfig, convert_to_trainer_args
+import pyrallis
+from omegaconf import OmegaConf
+import torch.optim as optim
+import wandb
+from torch.nn.utils.rnn import pad_sequence
+
+IGNORE_INDEX = -100
+PROMPT = (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+)
+
+def get_rank():
+    try:
+        rank = int(os.environ.get("LOCAL_RANK"))
+    except:
+        rank = 0
+    return rank
+
+
+def get_config():
+    config_path = os.environ.get("OMINI_CONFIG")
+    assert config_path is not None, "Please set the OMINI_CONFIG environment variable"
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def init_wandb(wandb_config, run_name):
+    import wandb
+
+    try:
+        assert os.environ.get("WANDB_API_KEY") is not None
+        wandb.init(
+            project=wandb_config["project"],
+            name=run_name,
+            config={},
+        )
+    except Exception as e:
+        print("Failed to initialize WanDB:", e)
+
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+        
+
+def smart_tokenizer_and_embedding_resize(
+        special_tokens_dict: Dict,
+        tokenizer: transformers.PreTrainedTokenizer,
+        model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+def preprocess(
+        sources: Sequence[str],
+        targets: Sequence[str],
+        tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+
+# @dataclass
+# class DataCollatorForSupervisedDataset():
+#     """Collate examples for supervised fine-tuning."""
+
+#     tokenizer: transformers.PreTrainedTokenizer
+#     max_length: int = field(default=512)
+#     mode: str = field(default="fixed") # dynamic -> dynamo
+
+#     def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+#         if self.mode == 'dynamic':
+#             input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+#             input_ids = [torch.tensor(x) for x in input_ids]
+#             input_ids = torch.nn.utils.rnn.pad_sequence(
+#                 input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+#             )
+#             labels = [torch.tensor(x) for x in labels]
+#             labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+#             return dict(
+#                 input_ids=input_ids,
+#                 labels=labels,
+#                 attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+#             )
+#         elif self.mode == 'fixed':
+#             input_ids = [torch.tensor(x["input_ids"][:self.max_length]) for x in instances]
+#             input_ids = torch.stack([
+#                 torch.nn.functional.pad(x, (0, self.max_length - x.size(0)), value=self.tokenizer.pad_token_id)
+#                 for x in input_ids
+#             ])
+
+#             # Labels
+#             labels = [torch.tensor(x["labels"][:self.max_length]) for x in instances]
+#             labels = torch.stack([
+#                 torch.nn.functional.pad(x, (0, self.max_length - x.size(0)), value=IGNORE_INDEX)
+#                 for x in labels
+#             ])
+
+#             return dict(
+#                 input_ids=input_ids,
+#                 labels=labels,
+#                 attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+#             )
+#         else:
+#             raise NotImplementedError
+
+# @dataclass
+# class DataCollatorForSupervisedDataset(object):
+#     tokenizer: transformers.PreTrainedTokenizer
+#     max_length: int = field(default=512)
+#     mode: str = field(default="fixed")  # "dynamic" or "fixed"
+
+#     def _pad_to_length(self, tensors: Sequence[torch.Tensor], pad_value: int, target_len: int):
+#         """Pad a list of 1D tensors to target_len (int) and stack -> (B, target_len)."""
+#         batch_size = len(tensors)
+#         out = torch.full((batch_size, target_len), pad_value, dtype=tensors[0].dtype)
+#         for i, t in enumerate(tensors):
+#             L = min(t.size(0), target_len)
+#             out[i, :L] = t[:L]
+#         return out
+
+#     def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+#         # Collect raw sequences (lists or tensors)
+#         input_seqs = [torch.tensor(x["input_ids"], dtype=torch.long) for x in instances]
+#         label_seqs = [torch.tensor(x["labels"], dtype=torch.long) for x in instances]
+
+#         if self.mode == "dynamic":
+#             # pad to the max length present in this batch (<= self.max_length)
+#             batch_max_len = min(max([s.size(0) for s in input_seqs]), self.max_length)
+#             input_ids = self._pad_to_length(input_seqs, pad_value=self.tokenizer.pad_token_id, target_len=batch_max_len)
+#             labels = self._pad_to_length(label_seqs, pad_value=IGNORE_INDEX, target_len=batch_max_len)
+#         elif self.mode == "fixed":
+#             # always pad/truncate to self.max_length
+#             input_ids = self._pad_to_length(input_seqs, pad_value=self.tokenizer.pad_token_id, target_len=self.max_length)
+#             labels = self._pad_to_length(label_seqs, pad_value=IGNORE_INDEX, target_len=self.max_length)
+#         else:
+#             raise NotImplementedError(f"Unknown mode: {self.mode}")
+
+#         attention_mask = input_ids.ne(self.tokenizer.pad_token_id).long()
+
+#         return {
+#             "input_ids": input_ids,
+#             "labels": labels,
+#             "attention_mask": attention_mask
+#         }
+    
+@dataclass
+class DataCollatorForSupervisedDataset():
+    tokenizer: transformers.PreTrainedTokenizer
+    max_length: int = field(default=512)
+    mode: str = field(default="fixed")  # "dynamic" or "fixed"
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        # Extract inputs and labels
+        # Assuming instances is a list of dicts like {'input_ids': [...], 'labels': [...]}
+        input_ids_list = [torch.tensor(x["input_ids"], dtype=torch.long) for x in instances]
+        labels_list = [torch.tensor(x["labels"], dtype=torch.long) for x in instances]
+
+        # 1. Determine padding logic
+        if self.mode == "dynamic":
+            # Dynamic padding: pad to the longest sequence in the batch
+            # But cap it at self.max_length to prevent OOM
+            batch_max_len = max([len(x) for x in input_ids_list])
+            target_len = min(batch_max_len, self.max_length)
+        else:
+            # Fixed padding: always pad to max_length
+            target_len = self.max_length
+
+        # 2. Helper to pad and truncate
+        def pad_and_truncate(tensors, padding_value):
+            # First, pad everything using PyTorch's optimized utility (batch_first=True)
+            padded = pad_sequence(tensors, batch_first=True, padding_value=padding_value)
+            
+            # Handle truncation/extending to exact target_len
+            curr_len = padded.shape[1]
+            if curr_len > target_len:
+                # Truncate if too long (rare if filtered beforehand)
+                return padded[:, :target_len]
+            elif curr_len < target_len:
+                # Pad more if shorter than target_len (happens in fixed mode)
+                diff = target_len - curr_len
+                padding = torch.full((padded.shape[0], diff), padding_value, dtype=padded.dtype)
+                return torch.cat([padded, padding], dim=1)
+            else:
+                return padded
+
+        # 3. Apply padding
+        # Critical: tokenizer.pad_token_id must NOT be None here
+        if self.tokenizer.pad_token_id is None:
+            raise ValueError("Tokenizer.pad_token_id is None. Please set it to eos_token_id or unk_token_id.")
+            
+        input_ids = pad_and_truncate(input_ids_list, self.tokenizer.pad_token_id)
+        labels = pad_and_truncate(labels_list, IGNORE_INDEX)
+
+        # 4. Create Attention Mask explicitly
+        # .ne() creates Bools, .long() casts to 0s and 1s for compatibility
+        attention_mask = input_ids.ne(self.tokenizer.pad_token_id).long()
+
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": attention_mask
+        }
+        
+def train_tokenize_function(examples, tokenizer, query, response):
+    sources = [PROMPT.format_map(dict(instruction=instruction)) for instruction in examples[query]]
+    targets = [f"{output}{tokenizer.eos_token}" for output in examples[response]]
+    data_dict = preprocess(sources, targets, tokenizer)
+    return data_dict
+
+
+
+### Trainer
+def default_worker_init_fn(worker_id):
+    # mỗi worker chỉ 1 thread cho BLAS
+    try:
+        import numpy as _np
+    except Exception:
+        _np = None
+    torch.set_num_threads(1)
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
+    os.environ.setdefault("MKL_NUM_THREADS", "1")
+    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
+    # Optional: bind CPU affinity per worker to avoid contention (NUMA-aware)
+    try:
+        cpu_count = os.cpu_count() or 1
+        # chia đều CPU cho workers
+        num_workers = getattr(torch.utils.data, "_num_workers", None)
+        # fallback: if not available, compute from environment variable or pass externally
+        # We'll do a simple round-robin assignment using worker_id
+        # assign a small mask of cores to this worker (e.g., chunk size 4)
+        chunk = max(1, cpu_count // max(1, min(64, cpu_count)))
+        start = (worker_id * chunk) % cpu_count
+        end = start + chunk
+        mask = set(range(start, min(end, cpu_count)))
+        try:
+            os.sched_setaffinity(0, mask)
+        except Exception:
+            pass
+    except Exception:
+        pass
+
+def set_seed(seed: int):
+    # random.seed(seed)
+    # np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    transformers.set_seed(seed)
+
+
+@pyrallis.wrap()
+def main(mainCfg: MainConfig):
+    #mainCfg = get_config()
+    #print(mainCfg)
+    print('='*120)
+    # print(OmegaConf.to_yaml(mainCfg))
+    # print('-'*40)
+    # 
+    # print((training_args))
+    set_seed(mainCfg.seed)
+    training_args = convert_to_trainer_args(mainCfg)
+
+    # wandb 
+    ENTITY = "nvan-13-korea-university" 
+    PROJECT = os.environ.get("WANDB_PROJECT")
+    api = wandb.Api()
+    try:
+        runs_list = api.runs(f"{ENTITY}/{PROJECT}")
+        next_run_num = len(runs_list) + 1
+    except Exception as e:
+        next_run_num = 1
+
+    training_args.run_name = f'[{next_run_num}]lr={mainCfg.trainer_args.learning_rate:.1e},b={mainCfg.trainer_args.per_device_train_batch_size},'\
+                            f'n={mainCfg.rotation_adapter_config.num_rotations},r={mainCfg.rotation_adapter_config.r},'\
+                            f'init={mainCfg.run_text}'
+    # training_args.project = f'Rotation-Llama2-{mainCfg.data.dataset_name}'
+
+    # print('-'*40)
+    # print(training_args.to_json_string())
+    # exit()
+    
+    model = AutoModelForCausalLM.from_pretrained(mainCfg.model.model_name,
+                                                 device_map="auto", low_cpu_mem_usage=True,
+                                                 dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+                                                 attn_implementation="sdpa",
+                                                 )
+    
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    print("DEVICE", DEVICE)
+    # for name, param in model.named_parameters():
+    #     if 'q_proj' in name and 'layers.5' in name:
+    #         print(f"Name: {name} | {param.shape} ")
+            # print(f"Name (pretrained): {name} | {param.shape} | {param.data[0:5,0:5]}")
+    # print('model', model)
+    # exit()
+
+    total_params_now = sum(p.numel() for p in model.parameters())
+    print(f'#params of the pretrained model, {total_params_now:,}')
+    # print(model)
+    if mainCfg.model.adapter_path is not None:
+        print('___ Loading from:  ', mainCfg.model.adapter_path)
+        model = PeftModel.from_pretrained(model, mainCfg.model.adapter_path, is_trainable = True)
+    elif mainCfg.rotation_adapter_config.r is not None:
+        rotation_adapter_config = asdict(mainCfg.rotation_adapter_config)
+        # rotation_adapter_config[peft_type]
+        
+        for adapter_name in mainCfg.data.adapter_names:
+            rotation_config = RotationConfig(**rotation_adapter_config)
+            model = get_peft_model(model, rotation_config, adapter_name=adapter_name)
+            # model.set_adapter(adapter_name)
+
+    else:
+        print("Full Parameter Fine-Tuning")
+    model = model.to(DEVICE)
+        
+    # print('model', model)
+    model.print_trainable_parameters()
+    # print("Program starts")
+    # time.sleep(300)
+    # exit()
+
+    # for name, param in model.named_parameters():
+    #     if 'q_proj' in name and 'rotation' in name and 'layers.5' in name:
+    #         print(f"Name: {name} | {param.shape} ")
+    #         print(f"Name (pretrained): {name} | {param.shape} ")
+    #         X = param.data
+    # print('model', type(model), X.shape)
+    # visualize_value_distribution(X)
+    # exit()
+
+    rotation_layers = filter(
+                lambda p: p.requires_grad, model.parameters()
+            )    
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        mainCfg.model.model_name,
+        model_max_length=mainCfg.model.model_max_seq_length,
+        padding_side="right",
+        use_fast=True,
+    )
+
+    if tokenizer.pad_token is None:
+        if tokenizer.unk_token_id is not None:
+            tokenizer.pad_token_id = tokenizer.unk_token_id
+            tokenizer.pad_token = tokenizer.unk_token
+            print("Set PAD token to UNK token.")
+        elif tokenizer.eos_token_id is not None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+            tokenizer.pad_token = tokenizer.eos_token
+            print("Set PAD token to EOS token.")
+
+        if model is not None:
+            model.config.pad_token_id = tokenizer.pad_token_id
+            if model.config.pad_token_id != tokenizer.pad_token_id:
+                raise ValueError("Failed to sync pad_token_id between tokenizer and model config")
+
+    # local MetaMathQA-40K
+    raw_datasets = load_dataset("json", data_files=mainCfg.data.path, split=mainCfg.data.dataset_split)
+    #raw_train_datasets = load_dataset("MetaMathQA-40K", split=mainCfg.data.dataset_split)
+    # print('raw', type(raw_train_datasets), len(raw_train_datasets))
+
+    # split a single set
+    split_ratio = mainCfg.data.split_ratio
+    split_data = raw_datasets.train_test_split(test_size=split_ratio, seed=42)
+    raw_train_datasets = split_data['train']
+    raw_valid_datasets = split_data['test']
+
+    train_dataset = raw_train_datasets.map(
+        train_tokenize_function,
+        batched=True,
+        batch_size=30000,
+        num_proc=32,
+        remove_columns=raw_train_datasets.column_names,
+        load_from_cache_file=True,
+        desc="Running tokenizer on train dataset",
+        fn_kwargs={"tokenizer": tokenizer, "query": mainCfg.data.dataset_field[0],
+                   "response": mainCfg.data.dataset_field[1]}
+    )
+
+    valid_dataset = raw_valid_datasets.map(
+        train_tokenize_function,
+        batched=True,
+        batch_size=30000,
+        num_proc=32,
+        remove_columns=raw_train_datasets.column_names,
+        load_from_cache_file=True,
+        desc="Running tokenizer on train dataset",
+        fn_kwargs={"tokenizer": tokenizer, "query": mainCfg.data.dataset_field[0],
+                   "response": mainCfg.data.dataset_field[1]}
+    )
+    print('- dataset size: ', len(valid_dataset), len(train_dataset))
+
+
+    # print('dataset', type(train_dataset))
+    # print('process', len(train_dataset))
+    # print(f"Sample features: {train_dataset.column_names}, {train_dataset.num_rows}")
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer, max_length=mainCfg.model.model_max_seq_length, 
+                                                     #mode=mainCfg.model.data_collator_mode,
+                                                     )
+    data_module = dict(train_dataset=train_dataset, data_collator=data_collator, eval_dataset=valid_dataset)
+
+    optimizer = optim.AdamW(
+        rotation_layers, 
+        lr=mainCfg.trainer_args.learning_rate, #
+        eps=1e-8
+    )
+    # print('model x', model)
+    start_time = datetime.now()
+    print('start time: ', start_time.strftime("%Y-%m-%d %H:%M:%S"))
+    trainer = MyTrainer(model=model, processing_class=tokenizer,
+                        lamda=mainCfg.model.lambda_reg,
+                        optimizers=(optimizer, None),
+                        args=training_args, **data_module)
+    model.config.use_cache = False
+
+
+    # now = time.time()
+    # for i in range(20):
+    #     next(iter(trainer.get_train_dataloader()))
+    # print('time', time.time()-now)
+    # now = time.time()
+    
+
+    # dl = trainer.get_train_dataloader()
+    # t0 = time.time()
+    # for i, batch in enumerate(dl):
+    #     if i==20: break
+    # print("time / 20 batches =", time.time() - t0)
+    # exit()
+
+
+    # model2 = model.merge_and_unload()
+    # results2 = trainer2.evaluate()
+    # print('results2: ', results2)
+    # exit()
+
+    start_time = datetime.now()
+    trainer.train()
+    
+    end_time = datetime.now()
+    print('end time: ', end_time.strftime("%Y-%m-%d %H:%M:%S"), '| duration: ', end_time - start_time)
+    # Save Model (Includes Adapter weights & Config)
+    # trainer.save_model(os.path.join(training_args.output_dir, 'ft'))
+    # Save Tokenizer
+    tokenizer.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
+    # Save Training State (Metrics & Logs)
+    trainer.save_state()
+
+    # save peft_config. Or model.base_model.peft_config['default']
+    model.peft_config.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
+
+    # the easiest way
+    model.save_pretrained(os.path.join(training_args.output_dir, 'ft2'))
+    return
+
+
+
+class MyTrainer(Trainer):
+
+    def __init__(
+            self,
+            model: Union[PreTrainedModel, nn.Module] = None,
+            args: TrainingArguments = None,
+            data_collator: Optional[DataCollator] = None,
+            train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None,
+            eval_dataset: Optional[Union[Dataset, Dict[str, Dataset], "datasets.Dataset"]] = None,
+            processing_class: Optional[PreTrainedTokenizerBase] = None,
+            model_init: Optional[Callable[[], PreTrainedModel]] = None,
+            compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+            callbacks: Optional[List[TrainerCallback]] = None,
+            optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+            preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
+            #run_name: Optional[str] = None,
+            #report_to: Optional[Union[str, list[str]]] = None,
+            # project
+            lamda: float = 1e-4
+    ):
+        super().__init__(model=model, args=args, data_collator=data_collator,
+                         train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=processing_class,
+                         model_init=model_init, compute_metrics=compute_metrics, callbacks=callbacks,
+                         optimizers=optimizers, preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+                         #run_name=run_name, report_to=report_to
+                         )
+        self.lamda = lamda
+
+    # def compute_loss(self, model, inputs, return_outputs=False,
+    #                  num_items_in_batch: Optional[torch.Tensor] = None,):
+    #     """
+    #     How the loss is computed by Trainer. By default, all models return the loss in the first element.
+
+    #     Subclass and override for custom behavior.
+    #     """
+    #     if self.label_smoother is not None and "labels" in inputs:
+    #         labels = inputs.pop("labels")
+    #     else:
+    #         labels = None
+    #     if self.model_accepts_loss_kwargs:
+    #         kwargs = {}
+    #         if num_items_in_batch is not None:
+    #             kwargs["num_items_in_batch"] = num_items_in_batch
+    #         inputs = {**inputs, **kwargs}
+    #     outputs = model(**inputs)
+    #     # Save past state if it exists
+    #     # TODO: this needs to be fixed and made cleaner later.
+    #     if self.args.past_index >= 0:
+    #         self._past = outputs[self.args.past_index]
+
+    #     if labels is not None:
+    #         unwrapped_model = unwrap_model(model)
+    #         if _is_peft_model(unwrapped_model):
+    #             model_name = unwrapped_model.base_model.model._get_name()
+    #         else:
+    #             model_name = unwrapped_model._get_name()
+    #         if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
+    #             loss = self.label_smoother(outputs, labels, shift_labels=True)
+    #         else:
+    #             loss = self.label_smoother(outputs, labels)
+    #     else:
+    #         if isinstance(outputs, dict) and "loss" not in outputs:
+    #             raise ValueError(
+    #                 "The model did not return a loss from the inputs, only the following keys: "
+    #                 f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
+    #             )
+    #         # We don't use .loss here since the model may return tuples instead of ModelOutput.
+    #         loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
+    #     # ------------------------------------------------------------------------------
+
+    #     # for name, param in model.named_parameters():
+    #     #     if 'oft_r' in name:
+    #     #         device = param.device
+    #     #         householder_U_norm = param / param.norm(dim=0)
+    #     #         orth_loss = torch.norm(
+    #     #             torch.eye(householder_U_norm.size(1), device=device) - householder_U_norm.t() @ householder_U_norm)
+    #     #         print(self.lamda)
+    #     #         loss = loss + self.lamda * orth_loss.to(loss.device)
+
+    #     # ------------------------------------------------------------------------------
+
+    #     return (loss, outputs) if return_outputs else loss
+
+    def get_train_dataloader(self):
+        # get dataset & sampler from super
+        train_dataset = self.train_dataset
+        sampler = self._get_train_sampler()
+
+        # compute effective batch size per step (HF has some routines; we use per_device_train_batch_size)
+        batch_size = self.args.train_batch_size if hasattr(self.args, "train_batch_size") else self.args.per_device_train_batch_size
+
+        # recommended num_workers: start moderate (16), you can tune upward
+        num_workers = getattr(self.args, "dataloader_num_workers", 16)
+        pin_memory = getattr(self.args, "dataloader_pin_memory", True)
+        prefetch_factor = getattr(self.args, "dataloader_prefetch_factor", 2)
+        persistent_workers = getattr(self.args, "dataloader_persistent_workers", True)
+
+        return DataLoader(
+            train_dataset,
+            batch_size=batch_size,
+            sampler=sampler,
+            collate_fn=self.data_collator,
+            drop_last=self.args.dataloader_drop_last if hasattr(self.args, "dataloader_drop_last") else False,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            persistent_workers=persistent_workers,
+            prefetch_factor=prefetch_factor,
+            worker_init_fn=default_worker_init_fn,
+        )
+    
+if __name__ == "__main__":
+    main()

nl_tasks/src/merge.py

@@ -0,0 +1,82 @@
+import torch
+# import wandb
+import os
+import yaml
+from peft import LoraConfig, get_peft_model_state_dict
+from torch.utils.data import DataLoader
+import time
+
+from typing import List, Tuple
+
+# import prodigyopt
+
+
+###
+import copy
+from dataclasses import field, dataclass, asdict
+from typing import Sequence, Literal, Dict
+
+import transformers
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+from transformers import Trainer
+from transformers.modeling_utils import *
+from transformers.trainer import _is_peft_model
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.data.data_collator import DataCollator
+
+from transformers.training_args import TrainingArguments
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.trainer_callback import TrainerCallback
+from transformers.trainer_utils import EvalPrediction
+from torch.utils.data import Dataset, IterableDataset
+from datasets import load_dataset
+##
+#from ..pipeline.flux_omini import transformer_forward, encode_images
+# from ...omini.rotation import RotationTuner, RotationConfig
+from rpeft.rotation import RotationTuner, RotationConfig
+from rpeft import get_peft_model, PeftModel
+from .config import MainConfig, convert_to_trainer_args
+import pyrallis
+from omegaconf import OmegaConf
+
+import argparse
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "</s>"
+DEFAULT_UNK_TOKEN = "</s>"
+PROMPT = (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+)
+
+# parser = argparse.ArgumentParser(description='Merge Adapter to Base Model')
+# parser.add_argument('--base_mode', type=str)
+# parser.add_argument('--adapter_path', type=str)
+# parser.add_argument('--output_path', type=str)
+# args = parser.parse_args()
+
+@pyrallis.wrap()
+def main(mainCfg: MainConfig):
+    print('='*120)
+    model_name = mainCfg.model.model_name
+    # adapter = mainCfg.trainer_args.output_dir + '/ft2'
+    # output_path = mainCfg.trainer_args.output_dir + '/merge/'
+    adapter = mainCfg.model.merge_adapter_path
+    output_path = mainCfg.model.merge_output_path
+
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto",)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, device_map='auto')
+
+    # config = PeftConfig.from_pretrained(args.adapter)
+    model = PeftModel.from_pretrained(model, adapter)
+    model = model.merge_and_unload()
+    model.save_pretrained(output_path, safe_serialization=False)
+    tokenizer.save_pretrained(output_path)
+    # print(model)
+    print('merge.py ends', adapter, output_path)
+    return
+
+if __name__ == "__main__":
+    main()

nl_tasks/src/peft_merge.py

@@ -0,0 +1,82 @@
+import torch
+# import wandb
+import os
+import yaml
+from peft import LoraConfig, get_peft_model_state_dict
+from torch.utils.data import DataLoader
+import time
+
+from typing import List, Tuple
+
+# import prodigyopt
+
+
+###
+import copy
+from dataclasses import field, dataclass, asdict
+from typing import Sequence, Literal, Dict
+
+import transformers
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+from transformers import Trainer
+from transformers.modeling_utils import *
+from transformers.trainer import _is_peft_model
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.data.data_collator import DataCollator
+
+from transformers.training_args import TrainingArguments
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.trainer_callback import TrainerCallback
+from transformers.trainer_utils import EvalPrediction
+from torch.utils.data import Dataset, IterableDataset
+from datasets import load_dataset
+##
+#from ..pipeline.flux_omini import transformer_forward, encode_images
+# from ...omini.rotation import RotationTuner, RotationConfig
+from rpeft.rotation import RotationTuner, RotationConfig
+from peft import get_peft_model, PeftModel
+from .config import MainConfig, convert_to_trainer_args
+import pyrallis
+from omegaconf import OmegaConf
+
+import argparse
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "</s>"
+DEFAULT_UNK_TOKEN = "</s>"
+PROMPT = (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+)
+
+# parser = argparse.ArgumentParser(description='Merge Adapter to Base Model')
+# parser.add_argument('--base_mode', type=str)
+# parser.add_argument('--adapter_path', type=str)
+# parser.add_argument('--output_path', type=str)
+# args = parser.parse_args()
+
+@pyrallis.wrap()
+def main(mainCfg: MainConfig):
+    print('='*120)
+    model_name = mainCfg.model.model_name
+    # adapter = mainCfg.trainer_args.output_dir + '/ft2'
+    # output_path = mainCfg.trainer_args.output_dir + '/merge/'
+    adapter = mainCfg.model.merge_adapter_path
+    output_path = mainCfg.model.merge_output_path
+
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto",)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, device_map='auto')
+
+    # config = PeftConfig.from_pretrained(args.adapter)
+    model = PeftModel.from_pretrained(model, adapter)
+    model = model.merge_and_unload()
+    model.save_pretrained(output_path, safe_serialization=False)
+    tokenizer.save_pretrained(output_path)
+    # print(model)
+    print('peft_merge.py ends', adapter, output_path)
+    return
+
+if __name__ == "__main__":
+    main()

nl_tasks/src/testLlama.py

@@ -0,0 +1,702 @@
+#
+import sys
+#print('sys.path: ___ ', sys.path)
+#print(f"Current Python Executable: {sys.executable}")
+
+### dynamo warning
+import warnings
+
+# Ignore FutureWarning: prims_common.check, Online Softmax
+warnings.filterwarnings("ignore", category=FutureWarning, module='torch._inductor.lowering')
+warnings.filterwarnings("ignore", message=".*Online softmax is disabled on the fly.*", category=UserWarning)
+
+warnings.filterwarnings("ignore", message=".*Our suggested max number of worker in current system is 1.*", category=UserWarning)
+warnings.filterwarnings("ignore", message=".*will be initialized from a multivariate normal distribution.*")
+warnings.filterwarnings("ignore", message=".*that differ from the model config and generation config.*", category=UserWarning)
+warnings.filterwarnings("ignore", message=".*torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch..*", category=UserWarning)
+
+import torch
+torch.backends.cuda.matmul.fp32_precision = 'tf32'
+# import wandb
+import os
+torch.set_num_threads(1)
+os.environ["OMP_NUM_THREADS"]="1"
+os.environ["MKL_NUM_THREADS"]="1"
+import torch
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"PyTorch built with CUDA version: {torch.version.cuda}")
+
+import yaml
+#from peft import LoraConfig, get_peft_model_state_dict
+from torch.utils.data import DataLoader
+import time
+from datetime import datetime
+import math
+
+from typing import List, Tuple
+
+# import prodigyopt
+
+
+###
+import copy
+from dataclasses import field, dataclass, asdict
+from typing import Sequence, Literal, Dict
+
+import transformers
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+from transformers import Trainer
+from transformers.modeling_utils import *
+from transformers.trainer import _is_peft_model
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.data.data_collator import DataCollator
+
+from transformers.training_args import TrainingArguments
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.trainer_callback import TrainerCallback
+from transformers.trainer_utils import EvalPrediction
+from torch.utils.data import Dataset, IterableDataset
+from datasets import load_dataset
+##
+#from ..pipeline.flux_omini import transformer_forward, encode_images
+# from ...omini.rotation import RotationTuner, RotationConfig
+
+
+from rpeft.rotation import RotationTuner, RotationConfig
+from rpeft import get_peft_model, PeftModel
+from .config import MainConfig, convert_to_trainer_args
+import pyrallis
+from omegaconf import OmegaConf
+import torch.optim as optim
+import wandb
+from torch.nn.utils.rnn import pad_sequence
+
+IGNORE_INDEX = -100
+PROMPT = (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+)
+
+
+import platform
+from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
+class ExperimentMonitorCallback(TrainerCallback):
+    """
+    Callback to monitor training performance and log system stats to a JSON file.
+    It captures:
+    1. Experiment Metadata (GPU info, Batch size, Learning rate, etc.)
+    2. Runtime Metrics (Avg time/step, Throughput)
+    3. Memory Metrics (Allocated, Reserved, and Peak usage)
+    """
+
+    def __init__(self, log_file_path: str, run_name: str = "experiment", log_interval: int = 100):
+        # English comments as requested
+        self.log_file_path = log_file_path
+        self.run_name = run_name
+        self.log_interval = log_interval
+        
+        # Timing variables
+        self.start_time = None
+        self.last_log_time = None
+        
+        # Data container to be saved
+        self.log_data = {
+            "metadata": {},
+            "metrics": []
+        }
+
+    def _get_gpu_info(self):
+        # Helper to get GPU details if available
+        if torch.cuda.is_available():
+            return {
+                "name": torch.cuda.get_device_name(0),
+                "count": torch.cuda.device_count(),
+                "capability": torch.cuda.get_device_capability(0)
+            }
+        return "CPU_ONLY"
+
+    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        # Initialize timing
+        self.start_time = time.perf_counter()
+        self.last_log_time = self.start_time
+        
+        # Reset peak memory stats to ensure we capture peaks specific to this run
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+
+        # Capture experiment metadata
+        self.log_data["metadata"] = {
+            "run_name": self.run_name,
+            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            "python_version": platform.python_version(),
+            "pytorch_version": torch.__version__,
+            "gpu_info": self._get_gpu_info(),
+            "configuration": {
+                "batch_size_per_device": args.per_device_train_batch_size,
+                "learning_rate": args.learning_rate,
+                "max_steps": args.max_steps,
+                "num_train_epochs": args.num_train_epochs,
+                "fp16": args.fp16,
+                "bf16": args.bf16,
+                "optim": args.optim,
+            }
+        }
+        
+        # Create/Overwrite the file with initial metadata
+        self._save_log()
+        # print(f"[{self.run_name}] Experiment started. Logging to {self.log_file_path}")
+
+    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        current_step = state.global_step
+        
+        # Perform logging only at specified intervals
+        if current_step > 0 and current_step % self.log_interval == 0:
+            current_time = time.perf_counter()
+            
+            # Calculate time elapsed since the last log
+            elapsed_since_last = current_time - self.last_log_time
+            avg_time_per_step = elapsed_since_last / self.log_interval
+            
+            # Memory Statistics (in GB)
+            mem_stats = {}
+            if torch.cuda.is_available():
+                # Current usage
+                mem_stats["allocated_gb"] = torch.cuda.memory_allocated() / 1024**3
+                mem_stats["reserved_gb"] = torch.cuda.memory_reserved() / 1024**3
+                # Peak usage since start (Long-term peak)
+                mem_stats["peak_allocated_gb"] = torch.cuda.max_memory_allocated() / 1024**3
+            
+            # Construct metric entry
+            metric_entry = {
+                "step": current_step,
+                "epoch": state.epoch,
+                "timestamp": datetime.now().isoformat(),
+                "performance": {
+                    "avg_time_per_step_s": round(avg_time_per_step, 4),
+                    "steps_per_second": round(1.0 / avg_time_per_step, 2)
+                },
+                "memory": mem_stats
+            }
+
+            # Append to internal list and save to file
+            self.log_data["metrics"].append(metric_entry)
+            self._save_log()
+            
+            # Update last log time
+            self.last_log_time = current_time
+            
+            # Optional: Print a brief summary to console
+            print(f" -> Step {current_step}: {avg_time_per_step*1000:.1f}s/step |"\
+                f"Peak Mem: {mem_stats.get('peak_allocated_gb', 0):.2f} GB |"\
+                f"Reserved: {mem_stats.get('reserved_gb', 0):.2f} GB")
+
+    def _save_log(self):
+        # Dump the entire data structure to JSON
+        # For very long training runs, appending to a JSONL (lines) file might be more efficient,
+        # but standard JSON is easier to read for analysis.
+        try:
+            with open(self.log_file_path, 'w', encoding='utf-8') as f:
+                json.dump(self.log_data, f, indent=4)
+        except Exception as e:
+            print(f"Error saving experiment log: {e}")
+            
+def get_rank():
+    try:
+        rank = int(os.environ.get("LOCAL_RANK"))
+    except:
+        rank = 0
+    return rank
+
+
+def get_config():
+    config_path = os.environ.get("OMINI_CONFIG")
+    assert config_path is not None, "Please set the OMINI_CONFIG environment variable"
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def init_wandb(wandb_config, run_name):
+    import wandb
+
+    try:
+        assert os.environ.get("WANDB_API_KEY") is not None
+        wandb.init(
+            project=wandb_config["project"],
+            name=run_name,
+            config={},
+        )
+    except Exception as e:
+        print("Failed to initialize WanDB:", e)
+
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+        
+
+def smart_tokenizer_and_embedding_resize(
+        special_tokens_dict: Dict,
+        tokenizer: transformers.PreTrainedTokenizer,
+        model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+def preprocess(
+        sources: Sequence[str],
+        targets: Sequence[str],
+        tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+
+    
+@dataclass
+class DataCollatorForSupervisedDataset():
+    tokenizer: transformers.PreTrainedTokenizer
+    max_length: int = field(default=512)
+    mode: str = field(default="fixed")  # "dynamic" or "fixed"
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        # Extract inputs and labels
+        # Assuming instances is a list of dicts like {'input_ids': [...], 'labels': [...]}
+        input_ids_list = [torch.tensor(x["input_ids"], dtype=torch.long) for x in instances]
+        labels_list = [torch.tensor(x["labels"], dtype=torch.long) for x in instances]
+
+        # 1. Determine padding logic
+        if self.mode == "dynamic":
+            # Dynamic padding: pad to the longest sequence in the batch
+            # But cap it at self.max_length to prevent OOM
+            batch_max_len = max([len(x) for x in input_ids_list])
+            target_len = min(batch_max_len, self.max_length)
+        else:
+            # Fixed padding: always pad to max_length
+            target_len = self.max_length
+
+        # 2. Helper to pad and truncate
+        def pad_and_truncate(tensors, padding_value):
+            # First, pad everything using PyTorch's optimized utility (batch_first=True)
+            padded = pad_sequence(tensors, batch_first=True, padding_value=padding_value)
+            
+            # Handle truncation/extending to exact target_len
+            curr_len = padded.shape[1]
+            if curr_len > target_len:
+                # Truncate if too long (rare if filtered beforehand)
+                return padded[:, :target_len]
+            elif curr_len < target_len:
+                # Pad more if shorter than target_len (happens in fixed mode)
+                diff = target_len - curr_len
+                padding = torch.full((padded.shape[0], diff), padding_value, dtype=padded.dtype)
+                return torch.cat([padded, padding], dim=1)
+            else:
+                return padded
+
+        # 3. Apply padding
+        # Critical: tokenizer.pad_token_id must NOT be None here
+        if self.tokenizer.pad_token_id is None:
+            raise ValueError("Tokenizer.pad_token_id is None. Please set it to eos_token_id or unk_token_id.")
+            
+        input_ids = pad_and_truncate(input_ids_list, self.tokenizer.pad_token_id)
+        labels = pad_and_truncate(labels_list, IGNORE_INDEX)
+
+        # 4. Create Attention Mask explicitly
+        # .ne() creates Bools, .long() casts to 0s and 1s for compatibility
+        attention_mask = input_ids.ne(self.tokenizer.pad_token_id).long()
+
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": attention_mask
+        }
+        
+def train_tokenize_function(examples, tokenizer, query, response):
+    sources = [PROMPT.format_map(dict(instruction=instruction)) for instruction in examples[query]]
+    targets = [f"{output}{tokenizer.eos_token}" for output in examples[response]]
+    data_dict = preprocess(sources, targets, tokenizer)
+    return data_dict
+
+
+
+### Trainer
+def default_worker_init_fn(worker_id):
+    # mỗi worker chỉ 1 thread cho BLAS
+    try:
+        import numpy as _np
+    except Exception:
+        _np = None
+    torch.set_num_threads(1)
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
+    os.environ.setdefault("MKL_NUM_THREADS", "1")
+    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
+    # Optional: bind CPU affinity per worker to avoid contention (NUMA-aware)
+    try:
+        cpu_count = os.cpu_count() or 1
+        # chia đều CPU cho workers
+        num_workers = getattr(torch.utils.data, "_num_workers", None)
+        # fallback: if not available, compute from environment variable or pass externally
+        # We'll do a simple round-robin assignment using worker_id
+        # assign a small mask of cores to this worker (e.g., chunk size 4)
+        chunk = max(1, cpu_count // max(1, min(64, cpu_count)))
+        start = (worker_id * chunk) % cpu_count
+        end = start + chunk
+        mask = set(range(start, min(end, cpu_count)))
+        try:
+            os.sched_setaffinity(0, mask)
+        except Exception:
+            pass
+    except Exception:
+        pass
+
+def set_seed(seed: int):
+    # random.seed(seed)
+    # np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    transformers.set_seed(seed)
+
+
+@pyrallis.wrap()
+def main(mainCfg: MainConfig):
+    #mainCfg = get_config()
+    #print(mainCfg)
+    print('='*120)
+    # print(OmegaConf.to_yaml(mainCfg))
+    # print('-'*40)
+    # 
+    # print((training_args))
+    set_seed(mainCfg.seed)
+    training_args = convert_to_trainer_args(mainCfg)
+
+    # wandb 
+    ENTITY = "nvan-13-korea-university" 
+    PROJECT = os.environ.get("WANDB_PROJECT")
+    api = wandb.Api()
+    try:
+        runs_list = api.runs(f"{ENTITY}/{PROJECT}")
+        next_run_num = len(runs_list) + 1
+    except Exception as e:
+        next_run_num = 1
+
+    training_args.run_name = f'[{next_run_num}]lr={mainCfg.trainer_args.learning_rate:.1e},b={mainCfg.trainer_args.per_device_train_batch_size},'\
+                            f'n={mainCfg.rotation_adapter_config.num_rotations},r={mainCfg.rotation_adapter_config.r},'\
+                            f'init={mainCfg.run_text}'
+    # training_args.project = f'Rotation-Llama2-{mainCfg.data.dataset_name}'
+
+    # print('-'*40)
+    # print(training_args.to_json_string())
+    # exit()
+    
+    model = AutoModelForCausalLM.from_pretrained(mainCfg.model.model_name,
+                                                 device_map="auto", low_cpu_mem_usage=True,
+                                                 dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+                                                 # attn_implementation="sdpa",
+                                                 )
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    print("DEVICE", model.device)
+
+    # for name, param in model.named_parameters():
+    #     if 'q_proj' in name and 'layers.5' in name:
+    #         print(f"Name: {name} | {param.shape} ")
+            # print(f"Name (pretrained): {name} | {param.shape} | {param.data[0:5,0:5]}")
+    # print('model', model)
+    # exit()
+
+    total_params_now = sum(p.numel() for p in model.parameters())
+    print(f'#params of the pretrained model, {total_params_now:,}')
+    # print(model)
+    if mainCfg.model.adapter_path is not None:
+        print('___ Loading from:  ', mainCfg.model.adapter_path)
+        model = PeftModel.from_pretrained(model, mainCfg.model.adapter_path, is_trainable = True)
+    elif mainCfg.rotation_adapter_config.r is not None:
+        import peft
+        if mainCfg.run_text == 'loco':
+            rotation_adapter_config = asdict(mainCfg.rotation_adapter_config)
+
+            for adapter_name in mainCfg.data.adapter_names:
+                rotation_config = RotationConfig(**rotation_adapter_config)
+                model = get_peft_model(model, rotation_config, adapter_name=adapter_name)
+            print('loaded a LoCo model, batch = ', training_args.per_device_train_batch_size)
+        elif mainCfg.run_text == 'boft':
+            from peft import BOFTConfig
+            boft_config = BOFTConfig(
+                boft_block_size=mainCfg.rotation_adapter_config.r,
+                boft_n_butterfly_factor=2*mainCfg.rotation_adapter_config.num_rotations,
+                target_modules=["q_proj", "v_proj",],
+                boft_dropout=0.05, #mainCfg.rotation_adapter_config.drop_out,
+                bias="none",
+                # task_type="CAUSAL_LM",
+            )
+
+            for adapter_name in mainCfg.data.adapter_names:
+                model = peft.get_peft_model(model, boft_config, adapter_name=adapter_name)
+            print('loaded a BOFT model, batch = ', training_args.per_device_train_batch_size)
+        elif mainCfg.run_text == 'hra':
+            from peft import HRAConfig
+            hra_config = HRAConfig(
+                r=2*mainCfg.rotation_adapter_config.r,
+                target_modules=["q_proj", "v_proj",],
+                init_weights=True,
+                # task_type="CAUSAL_LM",
+            )
+
+            for adapter_name in mainCfg.data.adapter_names:
+                model = peft.get_peft_model(model, hra_config, adapter_name=adapter_name)
+            print('loaded a HRA model, batch = ', training_args.per_device_train_batch_size)
+        elif mainCfg.run_text == 'oft':
+            from peft import HRAConfig, OFTConfig
+
+            oft_config = OFTConfig(
+                # r=16,
+                oft_block_size=4*mainCfg.rotation_adapter_config.r,
+                use_cayley_neumann=True,
+                target_modules=["q_proj", "v_proj",],
+                module_dropout=0.05, # mainCfg.rotation_adapter_config.drop_out,
+                # task_type="CAUSAL_LM",
+                bias="none",
+            )
+
+            for adapter_name in mainCfg.data.adapter_names:
+                model = peft.get_peft_model(model, oft_config, adapter_name=adapter_name)
+            print('loaded a OFT model, batch = ', training_args.per_device_train_batch_size)
+        else:
+            raise KeyError('wrong model names')
+
+    else:
+        print("Full Parameter Fine-Tuning")
+    model = model.to(DEVICE)
+        
+    # print('model', model)
+    model.print_trainable_parameters()
+
+    # print("Program starts")
+    # time.sleep(300)
+    # exit()
+
+    # for name, param in model.named_parameters():
+    #     if 'q_proj' in name and 'rotation' in name and 'layers.5' in name:
+    #         print(f"Name: {name} | {param.shape} ")
+    #         print(f"Name (pretrained): {name} | {param.shape} ")
+    #         X = param.data
+    # print('model', type(model), X.shape)
+    # visualize_value_distribution(X)
+    # exit()
+
+    rotation_layers = filter(
+                lambda p: p.requires_grad, model.parameters()
+            )    
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        mainCfg.model.model_name,
+        model_max_length=mainCfg.model.model_max_seq_length,
+        padding_side="right",
+        use_fast=True,
+    )
+
+    if tokenizer.pad_token is None:
+        if tokenizer.unk_token_id is not None:
+            tokenizer.pad_token_id = tokenizer.unk_token_id
+            tokenizer.pad_token = tokenizer.unk_token
+            print("Set PAD token to UNK token.")
+        elif tokenizer.eos_token_id is not None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+            tokenizer.pad_token = tokenizer.eos_token
+            print("Set PAD token to EOS token.")
+
+        if model is not None:
+            model.config.pad_token_id = tokenizer.pad_token_id
+            if model.config.pad_token_id != tokenizer.pad_token_id:
+                raise ValueError("Failed to sync pad_token_id between tokenizer and model config")
+
+    # local MetaMathQA-40K
+    raw_datasets = load_dataset("json", data_files=mainCfg.data.path, split=mainCfg.data.dataset_split)
+
+    train_dataset = raw_datasets.map(
+        train_tokenize_function,
+        batched=True,
+        batch_size=30000,
+        num_proc=32,
+        remove_columns=raw_datasets.column_names,
+        load_from_cache_file=True,
+        desc="Running tokenizer on train dataset",
+        fn_kwargs={"tokenizer": tokenizer, "query": mainCfg.data.dataset_field[0],
+                   "response": mainCfg.data.dataset_field[1]}
+    )
+
+    # valid_dataset = raw_valid_datasets.map(
+    #     train_tokenize_function,
+    #     batched=True,
+    #     batch_size=30000,
+    #     num_proc=32,
+    #     remove_columns=raw_train_datasets.column_names,
+    #     load_from_cache_file=True,
+    #     desc="Running tokenizer on train dataset",
+    #     fn_kwargs={"tokenizer": tokenizer, "query": mainCfg.data.dataset_field[0],
+    #                "response": mainCfg.data.dataset_field[1]}
+    # )
+    print('- dataset size: ', len(train_dataset))
+
+
+    # print('dataset', type(train_dataset))
+    # print('process', len(train_dataset))
+    # print(f"Sample features: {train_dataset.column_names}, {train_dataset.num_rows}")
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer, max_length=mainCfg.model.model_max_seq_length, 
+                                                     #mode=mainCfg.model.data_collator_mode,
+                                                     )
+    data_module = dict(train_dataset=train_dataset, data_collator=data_collator)
+
+    optimizer = optim.AdamW(
+        rotation_layers, 
+        lr=mainCfg.trainer_args.learning_rate, #
+        eps=1e-8
+    )
+    # print('model x', model)
+    start_time = datetime.now()
+    print('start time: ', start_time.strftime("%Y-%m-%d %H:%M:%S"))
+
+    monitor = ExperimentMonitorCallback(
+        log_file_path="./training_metrics_bs8.json",
+        run_name="Experiment_BatchSize_8",
+        log_interval=10  # Will calculate average over every 100 steps
+    )
+    training_args.remove_unused_columns = False
+    training_args.torch_compile=False
+    trainer = MyTrainer(model=model, processing_class=tokenizer,
+                        lamda=mainCfg.model.lambda_reg,
+                        optimizers=(optimizer, None),
+                        args=training_args, **data_module,
+                        callbacks=[monitor],
+                        )
+    model.config.use_cache = False
+
+    trainer.train()
+    
+    end_time = datetime.now()
+    print('end time: ', end_time.strftime("%Y-%m-%d %H:%M:%S"), '| duration: ', end_time - start_time)
+
+    # Save Model (Includes Adapter weights & Config)
+    # trainer.save_model(os.path.join(training_args.output_dir, 'ft'))
+    # Save Tokenizer
+    tokenizer.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
+    # Save Training State (Metrics & Logs)
+    trainer.save_state()
+
+    # save peft_config. Or model.base_model.peft_config['default']
+    # model.peft_config.save_pretrained(os.path.join(training_args.output_dir, 'ft'))
+
+    # the easiest way
+    model.save_pretrained(os.path.join(training_args.output_dir, 'ft2'))
+    return
+
+
+
+class MyTrainer(Trainer):
+
+    def __init__(
+            self,
+            model: Union[PreTrainedModel, nn.Module] = None,
+            args: TrainingArguments = None,
+            data_collator: Optional[DataCollator] = None,
+            train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None,
+            eval_dataset: Optional[Union[Dataset, Dict[str, Dataset], "datasets.Dataset"]] = None,
+            processing_class: Optional[PreTrainedTokenizerBase] = None,
+            model_init: Optional[Callable[[], PreTrainedModel]] = None,
+            compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+            callbacks: Optional[List[TrainerCallback]] = None,
+            optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+            preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
+            #run_name: Optional[str] = None,
+            #report_to: Optional[Union[str, list[str]]] = None,
+            # project
+            lamda: float = 1e-4
+    ):
+        super().__init__(model=model, args=args, data_collator=data_collator,
+                         train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=processing_class,
+                         model_init=model_init, compute_metrics=compute_metrics, callbacks=callbacks,
+                         optimizers=optimizers, preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+                         #run_name=run_name, report_to=report_to
+                         )
+        self.lamda = lamda
+
+
+    def get_train_dataloader(self):
+        # get dataset & sampler from super
+        train_dataset = self.train_dataset
+        sampler = self._get_train_sampler()
+
+        # compute effective batch size per step (HF has some routines; we use per_device_train_batch_size)
+        batch_size = self.args.train_batch_size if hasattr(self.args, "train_batch_size") else self.args.per_device_train_batch_size
+
+        # recommended num_workers: start moderate (16), you can tune upward
+        num_workers = getattr(self.args, "dataloader_num_workers", 16)
+        pin_memory = getattr(self.args, "dataloader_pin_memory", True)
+        prefetch_factor = getattr(self.args, "dataloader_prefetch_factor", 2)
+        persistent_workers = getattr(self.args, "dataloader_persistent_workers", True)
+
+        return DataLoader(
+            train_dataset,
+            batch_size=batch_size,
+            sampler=sampler,
+            collate_fn=self.data_collator,
+            drop_last=self.args.dataloader_drop_last if hasattr(self.args, "dataloader_drop_last") else False,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            persistent_workers=persistent_workers,
+            prefetch_factor=prefetch_factor,
+            worker_init_fn=default_worker_init_fn,
+        )
+    
+if __name__ == "__main__":
+    main()