Fixed FT test case on qaic

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

Author: quic-meetkuma

QEfficient/cloud/finetune.py

@@ -288,7 +288,7 @@ def main(
     scheduler = StepLR(optimizer, step_size=1, gamma=train_config.gamma)
     if train_config.enable_ddp:
         model = nn.parallel.DistributedDataParallel(model, device_ids=[dist.get_rank()])
-    train(
+    results = train(
         model,
         train_dataloader,
         eval_dataloader,
@@ -303,6 +303,7 @@ def main(
     )
     if train_config.enable_ddp:
         dist.destroy_process_group()
+    return results
 
 
 if __name__ == "__main__":

tests/finetune/test_finetune.py

@@ -8,6 +8,7 @@
 import os
 import shutil
 
+import numpy as np
 import pytest
 import torch.optim as optim
 from torch.utils.data import DataLoader
@@ -22,12 +23,10 @@ def clean_up(path):
         shutil.rmtree(path)
 
 
-configs = [pytest.param("meta-llama/Llama-3.2-1B", 1, 1, 1, None, True, True, "cpu", id="llama_config")]
+configs = [pytest.param("meta-llama/Llama-3.2-1B", 10, 20, 1, None, True, True, "qaic", id="llama_config")]
 
 
-# TODO:enable this once docker is available
 @pytest.mark.on_qaic
-@pytest.mark.skip(reason="eager docker not available in sdk")
 @pytest.mark.parametrize(
     "model_name,max_eval_step,max_train_step,intermediate_step_save,context_length,run_validation,use_peft,device",
     configs,
@@ -65,7 +64,13 @@ def test_finetune(
         "device": device,
     }
 
-    finetune(**kwargs)
+    results = finetune(**kwargs)
+
+    assert np.allclose(results["avg_train_prep"], 1.002326), "Train perplexity is not matching."
+    assert np.allclose(results["avg_train_loss"], 0.00232327), "Train loss is not matching."
+    assert np.allclose(results["avg_eval_prep"], 1.0193923), "Eval perplexity is not matching."
+    assert np.allclose(results["avg_eval_loss"], 0.0192067), "Eval loss is not matching."
+    assert results["avg_epoch_time"] < 30, "Training should complete within 30 seconds."
 
     train_config_spy.assert_called_once()
     generate_dataset_config_spy.assert_called_once()
@@ -99,8 +104,11 @@ def test_finetune(
 
     args, kwargs = update_config_spy.call_args
     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."
+    )
 
-    saved_file = os.path.join(train_config.output_dir, "adapter_model.safetensors")
+    saved_file = os.path.join(train_config.output_dir, "complete_epoch_1/adapter_model.safetensors")
     assert os.path.isfile(saved_file)
 
     clean_up(train_config.output_dir)