diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index 3c5f4c4d0..c7e30a0c5 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -680,6 +680,8 @@ def test_litgpt_variants(name, device): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") + if device == "cuda" and name == "falcon-40b-like": + pytest.skip("NVFuser reenable when https://github.com/NVIDIA/Fuser/issues/3505 is fixed, Thunder issue #1504") if device == "cuda" and name == "falcon-7b-like": pytest.skip("NVFuser reenable when https://github.com/NVIDIA/Fuser/issues/3292 is fixed") @@ -783,7 +785,7 @@ def sample(logits): ("cpu", "cuda"), ) def test_tom_overrides_proxy(device): - from litgpt.config import Config + from thunder.tests.litgpt_model import Config from litgpt.model import GPT if device == "cuda" and not torch.cuda.is_available(): diff --git a/thunder/tests/test_torch_compile_executor.py b/thunder/tests/test_torch_compile_executor.py index c0bd1b351..ea02dc8b5 100644 --- a/thunder/tests/test_torch_compile_executor.py +++ b/thunder/tests/test_torch_compile_executor.py @@ -20,6 +20,7 @@ def test_supported_ops_are_in_pytorch_executor(): # appropriate visual studio config. @pytest.mark.skipif(not is_inductor_supported() or platform.system() == "Windows", reason="inductor unsupported") def test_torch_compile_litgpt(): + from thunder.tests.litgpt_model import Config from litgpt.model import GPT model = GPT.from_name("llama1-like", n_layer=1) @@ -40,7 +41,7 @@ def test_torch_compile_litgpt(): @requiresCUDA @pytest.mark.skipif(not device_supports_bf16(torch.device("cuda")), reason="bf16 is not supported") def test_torch_compile_cat_nvfuser_phi2_tanh(): - from litgpt.config import Config + from thunder.tests.litgpt_model import Config from litgpt.model import GPT device = torch.device("cuda") diff --git a/thunder/tests/test_transforms.py b/thunder/tests/test_transforms.py index f6b56bc90..42a0ba0e6 100644 --- a/thunder/tests/test_transforms.py +++ b/thunder/tests/test_transforms.py @@ -74,7 +74,8 @@ def _test_equal_nvtx_push_and_pop(trc): @requiresCUDA def test_materialization(): from thunder.transforms import MaterializationTransform - from litgpt.config import Config + from thunder.tests.litgpt_model import Config + from litgpt.model import GPT config = Config.from_name("llama2-like") @@ -121,7 +122,7 @@ def test_materialization(): def test_quantization_on_meta(): from thunder.transforms import MaterializationTransform from thunder.transforms.quantization import BitsAndBytesLinearQuant4bit, get_bitsandbytes_executor - from litgpt.config import Config + from thunder.tests.litgpt_model import Config from litgpt.model import GPT bitsandbytes_executor = get_bitsandbytes_executor()