From c9e37496cb623c79f835817bcd6cd271d37ff370 Mon Sep 17 00:00:00 2001
From: Dan Saunders <danjsaund@gmail.com>
Date: Fri, 10 Jan 2025 17:32:33 +0000
Subject: [PATCH] Fix

---
 tests/e2e/integrations/test_cut_cross_entropy.py | 2 +-
 tests/e2e/patched/test_4d_multipack_llama.py     | 2 +-
 tests/e2e/patched/test_fa_xentropy.py            | 2 +-
 tests/e2e/patched/test_falcon_samplepack.py      | 2 +-
 tests/e2e/patched/test_fused_llama.py            | 2 +-
 tests/e2e/patched/test_llama_s2_attention.py     | 2 +-
 tests/e2e/patched/test_lora_llama_multipack.py   | 2 +-
 tests/e2e/patched/test_mistral_samplepack.py     | 2 +-
 tests/e2e/patched/test_mixtral_samplepack.py     | 2 +-
 tests/e2e/patched/test_phi_multipack.py          | 2 +-
 tests/e2e/patched/test_resume.py                 | 2 +-
 tests/e2e/patched/test_unsloth_qlora.py          | 2 +-
 tests/e2e/test_embeddings_lr.py                  | 2 +-
 tests/e2e/test_falcon.py                         | 2 +-
 tests/e2e/test_llama.py                          | 2 +-
 tests/e2e/test_llama_pretrain.py                 | 2 +-
 tests/e2e/test_llama_vision.py                   | 2 +-
 tests/e2e/test_lora_llama.py                     | 2 +-
 tests/e2e/test_mamba.py                          | 2 +-
 tests/e2e/test_mistral.py                        | 2 +-
 tests/e2e/test_mixtral.py                        | 2 +-
 tests/e2e/test_optimizers.py                     | 2 +-
 tests/e2e/test_packing_loss.py                   | 2 +-
 tests/e2e/test_phi.py                            | 2 +-
 tests/e2e/test_relora_llama.py                   | 2 +-
 tests/e2e/test_reward_model_llama.py             | 2 +-
 26 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/tests/e2e/integrations/test_cut_cross_entropy.py b/tests/e2e/integrations/test_cut_cross_entropy.py
index 3205a479c..291a4a4ec 100644
--- a/tests/e2e/integrations/test_cut_cross_entropy.py
+++ b/tests/e2e/integrations/test_cut_cross_entropy.py
@@ -4,7 +4,7 @@ Simple end-to-end test for Cut Cross Entropy integration
 
 import pytest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils import get_pytorch_version
diff --git a/tests/e2e/patched/test_4d_multipack_llama.py b/tests/e2e/patched/test_4d_multipack_llama.py
index e0b9a453e..da27069ac 100644
--- a/tests/e2e/patched/test_4d_multipack_llama.py
+++ b/tests/e2e/patched/test_4d_multipack_llama.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_fa_xentropy.py b/tests/e2e/patched/test_fa_xentropy.py
index 80236d3e8..2bfd36d15 100644
--- a/tests/e2e/patched/test_fa_xentropy.py
+++ b/tests/e2e/patched/test_fa_xentropy.py
@@ -8,7 +8,7 @@ import os
 import pytest
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_falcon_samplepack.py b/tests/e2e/patched/test_falcon_samplepack.py
index 79ef12996..62ee4f717 100644
--- a/tests/e2e/patched/test_falcon_samplepack.py
+++ b/tests/e2e/patched/test_falcon_samplepack.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_fused_llama.py b/tests/e2e/patched/test_fused_llama.py
index 817f7ec61..e7ab510c9 100644
--- a/tests/e2e/patched/test_fused_llama.py
+++ b/tests/e2e/patched/test_fused_llama.py
@@ -9,7 +9,7 @@ import unittest
 import pytest
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_llama_s2_attention.py b/tests/e2e/patched/test_llama_s2_attention.py
index 147bca788..8d0ba6c2a 100644
--- a/tests/e2e/patched/test_llama_s2_attention.py
+++ b/tests/e2e/patched/test_llama_s2_attention.py
@@ -8,7 +8,7 @@ import unittest
 
 import pytest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_lora_llama_multipack.py b/tests/e2e/patched/test_lora_llama_multipack.py
index f321dbddd..bc18e3d81 100644
--- a/tests/e2e/patched/test_lora_llama_multipack.py
+++ b/tests/e2e/patched/test_lora_llama_multipack.py
@@ -9,7 +9,7 @@ import unittest
 import pytest
 from transformers.utils import is_auto_gptq_available, is_torch_bf16_gpu_available
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_mistral_samplepack.py b/tests/e2e/patched/test_mistral_samplepack.py
index 7d1d4b94e..c7fd0ecbc 100644
--- a/tests/e2e/patched/test_mistral_samplepack.py
+++ b/tests/e2e/patched/test_mistral_samplepack.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_mixtral_samplepack.py b/tests/e2e/patched/test_mixtral_samplepack.py
index 4f960a452..156dac7e8 100644
--- a/tests/e2e/patched/test_mixtral_samplepack.py
+++ b/tests/e2e/patched/test_mixtral_samplepack.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_phi_multipack.py b/tests/e2e/patched/test_phi_multipack.py
index 1c4104818..ce466460e 100644
--- a/tests/e2e/patched/test_phi_multipack.py
+++ b/tests/e2e/patched/test_phi_multipack.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_resume.py b/tests/e2e/patched/test_resume.py
index 2523cdafd..f6a3e0109 100644
--- a/tests/e2e/patched/test_resume.py
+++ b/tests/e2e/patched/test_resume.py
@@ -9,7 +9,7 @@ import subprocess
 
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/patched/test_unsloth_qlora.py b/tests/e2e/patched/test_unsloth_qlora.py
index 857bf930e..da5eaffb6 100644
--- a/tests/e2e/patched/test_unsloth_qlora.py
+++ b/tests/e2e/patched/test_unsloth_qlora.py
@@ -6,7 +6,7 @@ import os
 
 import pytest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_embeddings_lr.py b/tests/e2e/test_embeddings_lr.py
index 60470b898..4261ccc26 100644
--- a/tests/e2e/test_embeddings_lr.py
+++ b/tests/e2e/test_embeddings_lr.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_falcon.py b/tests/e2e/test_falcon.py
index 6d31d4caf..ddcb66275 100644
--- a/tests/e2e/test_falcon.py
+++ b/tests/e2e/test_falcon.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_llama.py b/tests/e2e/test_llama.py
index 5a2c1e552..a94828490 100644
--- a/tests/e2e/test_llama.py
+++ b/tests/e2e/test_llama.py
@@ -7,7 +7,7 @@ import os
 
 from e2e.utils import check_model_output_exists
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_llama_pretrain.py b/tests/e2e/test_llama_pretrain.py
index be2c7b50f..68cd490be 100644
--- a/tests/e2e/test_llama_pretrain.py
+++ b/tests/e2e/test_llama_pretrain.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_llama_vision.py b/tests/e2e/test_llama_vision.py
index fed4b24ac..91f101e44 100644
--- a/tests/e2e/test_llama_vision.py
+++ b/tests/e2e/test_llama_vision.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_lora_llama.py b/tests/e2e/test_lora_llama.py
index ba0696c44..696c47aed 100644
--- a/tests/e2e/test_lora_llama.py
+++ b/tests/e2e/test_lora_llama.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_mamba.py b/tests/e2e/test_mamba.py
index 33732602e..0591fb23a 100644
--- a/tests/e2e/test_mamba.py
+++ b/tests/e2e/test_mamba.py
@@ -8,7 +8,7 @@ import unittest
 
 import pytest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_mistral.py b/tests/e2e/test_mistral.py
index 66fedd6a9..a304e9b4a 100644
--- a/tests/e2e/test_mistral.py
+++ b/tests/e2e/test_mistral.py
@@ -8,7 +8,7 @@ import unittest
 
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_mixtral.py b/tests/e2e/test_mixtral.py
index c984c3fa0..6e06626f6 100644
--- a/tests/e2e/test_mixtral.py
+++ b/tests/e2e/test_mixtral.py
@@ -9,7 +9,7 @@ import unittest
 import torch
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_optimizers.py b/tests/e2e/test_optimizers.py
index 76da6bb6a..453872538 100644
--- a/tests/e2e/test_optimizers.py
+++ b/tests/e2e/test_optimizers.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_packing_loss.py b/tests/e2e/test_packing_loss.py
index 71e57509a..13244a215 100644
--- a/tests/e2e/test_packing_loss.py
+++ b/tests/e2e/test_packing_loss.py
@@ -8,7 +8,7 @@ import unittest
 
 from transformers.utils import is_torch_bf16_gpu_available
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_phi.py b/tests/e2e/test_phi.py
index a94af7d30..54f564d0e 100644
--- a/tests/e2e/test_phi.py
+++ b/tests/e2e/test_phi.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_relora_llama.py b/tests/e2e/test_relora_llama.py
index 8e02d8a06..6c785dc86 100644
--- a/tests/e2e/test_relora_llama.py
+++ b/tests/e2e/test_relora_llama.py
@@ -7,7 +7,7 @@ import os
 import unittest
 from pathlib import Path
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
diff --git a/tests/e2e/test_reward_model_llama.py b/tests/e2e/test_reward_model_llama.py
index 469b9cf0d..4cd8602f3 100644
--- a/tests/e2e/test_reward_model_llama.py
+++ b/tests/e2e/test_reward_model_llama.py
@@ -6,7 +6,7 @@ import logging
 import os
 import unittest
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.cli.args import TrainerCliArgs
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config