Updated jenkins file to run finetuning tests and dump in separate file. Addressed comments.

Signed-off-by: Meet Patel <[email protected]>

Author: quic-meetkuma

QEfficient/cloud/finetune.py

@@ -47,11 +47,11 @@
 warnings.filterwarnings("ignore")
 
 
-def setup_distributed_training(config: TrainConfig) -> None:
+def setup_distributed_training(train_config: TrainConfig) -> None:
     """Initialize distributed training environment if enabled.
 
     Args:
-        config (TrainConfig): Training configuration object.
+        train_config (TrainConfig): Training configuration object.
 
     Notes:
         - If distributed data parallel (DDP) is disabled, this function does nothing.
@@ -61,14 +61,14 @@ def setup_distributed_training(config: TrainConfig) -> None:
     Raises:
         AssertionError: If device is CPU or includes an index with DDP enabled.
     """
-    if not config.enable_ddp:
+    if not train_config.enable_ddp:
         return
 
-    torch_device = torch.device(config.device)
+    torch_device = torch.device(train_config.device)
     assert torch_device.type != "cpu", "Host doesn't support single-node DDP"
     assert torch_device.index is None, f"DDP requires only device type, got: {torch_device}"
 
-    dist.init_process_group(backend=config.dist_backend)
+    dist.init_process_group(backend=train_config.dist_backend)
     # from here onward "qaic/cuda" will automatically map to "qaic:i/cuda:i", where i = process rank
     getattr(torch, torch_device.type).set_device(dist.get_rank())
 

QEfficient/finetune/configs/training.py

@@ -19,6 +19,7 @@ class TrainConfig:
         batch_size_training (int): Batch size for training (default: 1).
         context_length (Optional[int]): Maximum sequence length for inputs (default: None).
         gradient_accumulation_steps (int): Steps for gradient accumulation (default: 4).
+        gradient checkpointing (bool): Enable gradient checkpointing to save the memory by compromising the speed. (default: False).
         num_epochs (int): Number of training epochs (default: 1).
         max_train_step (int): Maximum training steps (default: 0, unlimited if 0).
         max_eval_step (int): Maximum evaluation steps (default: 0, unlimited if 0).
@@ -32,6 +33,7 @@ class TrainConfig:
         use_autocast (bool): Use autocast for mixed precision (default: True).
         val_batch_size (int): Batch size for validation (default: 1).
         dataset (str): Dataset name for training (default: "samsum_dataset").
+        task_type (str): Type of task for which the finetuning is to be done. Options: "generation" and "seq_classification". (default: "generation")
         peft_method (str): Parameter-efficient fine-tuning method (default: "lora").
         use_peft (bool): Whether to use PEFT (default: True).
         from_peft_checkpoint (str): Path to PEFT checkpoint (default: "").

scripts/Jenkinsfile

@@ -65,6 +65,23 @@ pipeline {
                        }
                    }
                }
+               stage('Run Non-CLI QAIC Finetuning Tests') {
+                   steps {
+                       timeout(time: 200, unit: 'MINUTES') {
+                           sh '''
+                           sudo docker exec ${BUILD_TAG} bash -c "
+                           cd /efficient-transformers &&
+                           . preflight_qeff/bin/activate &&
+                           mkdir -p $PWD/Non_cli_qaic_finetuning &&
+                           export TOKENIZERS_PARALLELISM=false &&
+                           export QEFF_HOME=$PWD/Non_cli_qaic_finetuning &&
+                           pytest tests -m '(not cli) and (on_qaic) and (not qnn) and (finetune)' --ignore tests/vllm -n 4 --junitxml=tests/tests_log3.xml &&
+                           junitparser merge tests/tests_log3.xml tests/tests_log.xml &&
+                           deactivate"
+                           '''
+                       }
+                   }
+               }
            }
        }
        stage('CLI Tests') {

tests/finetune/test_finetune.py

@@ -23,7 +23,19 @@ def clean_up(path):
         shutil.rmtree(path)
 
 
-configs = [pytest.param("meta-llama/Llama-3.2-1B", 10, 20, 1, None, True, True, "qaic", id="llama_config")]
+configs = [
+    pytest.param(
+        "meta-llama/Llama-3.2-1B",  # model_name
+        10,  # max_eval_step
+        20,  # max_train_step
+        1,  # intermediate_step_save
+        None,  # context_length
+        True,  # run_validation
+        True,  # use_peft
+        "qaic",  # device
+        id="llama_config",  # config name
+    )
+]
 
 
 @pytest.mark.on_qaic
@@ -105,7 +117,8 @@ def test_finetune(
     args, kwargs = update_config_spy.call_args_list[0]
     train_config = args[0]
     assert max_train_step >= train_config.gradient_accumulation_steps, (
-        "Total training step should be more than 4 which is gradient accumulation steps."
+        "Total training step should be more than "
+        f"{train_config.gradient_accumulation_steps} which is gradient accumulation steps."
     )
 
     saved_file = os.path.join(train_config.output_dir, "complete_epoch_1/adapter_model.safetensors")