fix: add explicit imports for DistributedSampler, DataLoader, IterableDatasetShard, ipdb

root_gainlora/runtest_colab.ipynb

@@ -226,10 +226,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "id": "6257b4db",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/content\n"
+     ]
+    }
+   ],
    "source": [
     "# Cell for remove run\n",
     "%cd ../..\n",
@@ -288,7 +296,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "id": "ad592075",
    "metadata": {},
    "outputs": [
@@ -296,67 +304,49 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      " the above exception, another exception occurred:\n",
-      "\n",
+      "ing: False -----\n",
+      "tmp\n",
       "Traceback (most recent call last):\n",
-      "  File \"/content/Continual/root_gainlora/src/run_t5.py\", line 917, in <module>\n",
+      "  File \"/content/Continual/root_gainlora/src/run_t5.py\", line 923, in <module>\n",
       "    main()\n",
-      "  File \"/content/Continual/root_gainlora/src/run_t5.py\", line 426, in main\n",
-      "    raw_datasets = load_dataset(\n",
-      "                   ^^^^^^^^^^^^^\n",
-      "  File \"/usr/local/lib/python3.12/dist-packages/datasets/load.py\", line 2606, in load_dataset\n",
-      "    builder_instance = load_dataset_builder(\n",
-      "                       ^^^^^^^^^^^^^^^^^^^^^\n",
-      "  File \"/usr/local/lib/python3.12/dist-packages/datasets/load.py\", line 2277, in load_dataset_builder\n",
-      "    dataset_module = dataset_module_factory(\n",
-      "                     ^^^^^^^^^^^^^^^^^^^^^^^\n",
-      "  File \"/usr/local/lib/python3.12/dist-packages/datasets/load.py\", line 1823, in dataset_module_factory\n",
-      "    ).get_module()\n",
-      "      ^^^^^^^^^^^^\n",
-      "  File \"/usr/local/lib/python3.12/dist-packages/datasets/load.py\", line 968, in get_module\n",
-      "    trust_remote_code = resolve_trust_remote_code(self.trust_remote_code, self.name)\n",
-      "                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-      "  File \"/usr/local/lib/python3.12/dist-packages/datasets/load.py\", line 134, in resolve_trust_remote_code\n",
-      "    raise ValueError(\n",
-      "ValueError: The repository for cl_dataset contains custom code which must be executed to correctly load the dataset. You can inspect the repository content at https://hf.co/datasets/cl_dataset.\n",
-      "Please pass the argument `trust_remote_code=True` to allow custom code to be run.\n",
+      "  File \"/content/Continual/root_gainlora/src/run_t5.py\", line 743, in main\n",
+      "    trainer.get_reg_matrix()\n",
+      "  File \"/content/Continual/root_gainlora/src/cl_trainer_gainlora_inflora.py\", line 143, in get_reg_matrix\n",
+      "    self.feature_list, self.feature_trans_list, self._cur_task = self.load_previous_reg_matrix()\n",
+      "                                                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/content/Continual/root_gainlora/src/cl_trainer_gainlora_inflora.py\", line 118, in load_previous_reg_matrix\n",
+      "    all_dirs = os.listdir(log_path)\n",
+      "               ^^^^^^^^^^^^^^^^^^^^\n",
+      "FileNotFoundError: [Errno 2] No such file or directory: 'tmp'\n",
       "\n",
       "RC: 0\n"
      ]
     }
    ],
    "source": [
-    "# Run first task with --max_steps 5, only show last 50 lines  \n",
+    "# Test with real output path and --max_steps 3\n",
     "import subprocess\n",
-    "result = subprocess.run(\n",
-    "    ['bash', '-c', '''\n",
+    "r = subprocess.run(['bash', '-c', '''\n",
     "cd /content/Continual/root_gainlora\n",
+    "rm -rf /root/.cache/huggingface/modules/datasets_modules/ 2>/dev/null || true\n",
     "CUDA_VISIBLE_DEVICES=0 python src/run_t5.py \\\n",
-    "   --do_train --predict_with_generate \\\n",
-    "   --model_name_or_path google/flan-t5-large \\\n",
+    "   --do_train --predict_with_generate --model_name_or_path google/flan-t5-large \\\n",
     "   --data_dir CL_Benchmark \\\n",
     "   --task_order yelp,amazon,mnli,cb,copa,qqp,rte,imdb,sst2,dbpedia,agnews,yahoo,multirc,boolq,wic \\\n",
     "   --task_config_dir configs/gen_script_long_order3_t5_configs/yelp \\\n",
-    "   --output_dir /tmp/test_run \\\n",
-    "   --per_device_train_batch_size 8 --per_device_eval_batch_size 128 \\\n",
-    "   --gradient_accumulation_steps 4 --learning_rate 0.0003 --num_train_epochs 1 \\\n",
-    "   --run_name test --max_source_length 512 --max_target_length 50 --generation_max_length 50 \\\n",
-    "   --add_task_name False --add_dataset_name False --overwrite_output_dir --overwrite_cache \\\n",
-    "   --lr_scheduler_type constant --warmup_steps 0 --logging_strategy steps --logging_steps 10 \\\n",
-    "   --metric_for_best_model eval_exact_match --evaluation_strategy steps --save_strategy steps \\\n",
-    "   --save_total_limit 1 --load_best_model_at_end --lora_r 8 --lora_alpha 32 --lora_dropout 0.0 \\\n",
-    "   --add_instruction_replay --data_replay_freq -1 --replay_after_n_epoch 0 --mlp_hidden_dim 100 \\\n",
-    "   --model_name gainlora_inflora --threshold 0.995 --transthreshold 0.995 \\\n",
-    "   --max_steps 3 2>&1 | tail -50\n",
-    "'''],\n",
-    "    capture_output=True, text=True, timeout=300\n",
-    ")\n",
-    "# Only print last 1500 chars\n",
-    "out = result.stdout[-1500:] if result.stdout else \"\"\n",
-    "err = result.stderr[-500:] if result.stderr else \"\"\n",
-    "print(out)\n",
-    "if err: print(\"ERR:\", err)\n",
-    "print(\"RC:\", result.returncode)"
+    "   --output_dir logs_and_outputs/gen_script_long_order3_t5_gainlora_inflora/outputs/1-yelp \\\n",
+    "   --per_device_train_batch_size 8 --per_device_eval_batch_size 128 --gradient_accumulation_steps 4 \\\n",
+    "   --learning_rate 0.0003 --num_train_epochs 1 --run_name test --max_source_length 512 \\\n",
+    "   --max_target_length 50 --generation_max_length 50 --add_task_name False --add_dataset_name False \\\n",
+    "   --overwrite_output_dir --overwrite_cache --lr_scheduler_type constant --warmup_steps 0 \\\n",
+    "   --logging_strategy steps --logging_steps 10 --metric_for_best_model eval_exact_match \\\n",
+    "   --evaluation_strategy steps --save_strategy steps --save_total_limit 1 --load_best_model_at_end \\\n",
+    "   --lora_r 8 --lora_alpha 32 --lora_dropout 0.0 --add_instruction_replay --data_replay_freq -1 \\\n",
+    "   --replay_after_n_epoch 0 --mlp_hidden_dim 100 --model_name gainlora_inflora \\\n",
+    "   --threshold 0.995 --transthreshold 0.995 --max_steps 3 2>&1 | tail -15\n",
+    "'''], capture_output=True, text=True, timeout=300, cwd='/content')\n",
+    "print(r.stdout[-800:])\n",
+    "print(\"RC:\", r.returncode)"
    ]
   }
  ],

root_gainlora/src/cl_trainer_gainlora_inflora.py

@@ -4,13 +4,22 @@ from transformers.trainer_seq2seq import Seq2SeqTrainer
 from transformers.trainer import *
 from transformers.trainer_callback import TrainerCallback
 import numpy as np
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+try:
+    from transformers.trainer_pt_utils import IterableDatasetShard
+except ImportError:
+    from torch.utils.data import IterableDataset as IterableDatasetShard
 
 from cl_collator import SUPPORTED_DECODER_MODELS, check_model
 from cl_dataset import ANSWER_PREFIX
 import cupy as cp
 from torch.utils.dlpack import to_dlpack, from_dlpack
 from cupy import fromDlpack
-import ipdb
+try:
+    import ipdb
+except ImportError:
+    ipdb = None
 
 def skip_instructions(model, predictions_ids, tokenizer, ignore_idx=-100):
     predictions_ids = np.where(predictions_ids == ignore_idx, tokenizer.pad_token_id, predictions_ids)