try tests for kd on l40s

2025-01-14 23:56:00 -05:00
parent 35a84f2cb8
commit 483defb9ae
2 changed files with 2 additions and 1 deletions
--- a/cicd/tests.py
+++ b/cicd/tests.py
@@ -59,7 +59,7 @@ VOLUME_CONFIG = {
 }

 N_GPUS = int(os.environ.get("N_GPUS", 1))
-GPU_CONFIG = modal.gpu.A10G(count=N_GPUS)
+GPU_CONFIG = modal.gpu.L40S(count=N_GPUS)


 def run_cmd(cmd: str, run_folder: str):
--- a/tests/e2e/integrations/test_kd.py
+++ b/tests/e2e/integrations/test_kd.py
@@ -33,6 +33,7 @@ def min_cfg(temp_dir):
        "dataloader_prefetch_factor": 8,
        "dataloader_num_workers": 4,
        "dataloader_pin_memory": True,
+        # "dataset_prepared_path": str(Path(temp_dir) / "last_run_prepared"),
        "datasets": [
            {
                "path": "axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample",