Fix code quality

mosaicml · Oct 24, 2023 · e85fc95 · e85fc95
1 parent 8848934
commit e85fc95
Show file tree

Hide file tree

Showing 9 changed files with 8 additions and 11 deletions.
diff --git a/tests/fixtures/autouse.py b/tests/fixtures/autouse.py
@@ -27,10 +27,12 @@ def clear_cuda_cache(request: pytest.FixtureRequest):
         torch.cuda.empty_cache()
         gc.collect()  # Only gc on GPU tests as it 2x slows down CPU tests
 
+
 @pytest.fixture
 def random_seed() -> int:
     return 17
 
+
 @pytest.fixture(autouse=True)
 def seed_all(random_seed: int):
     """Sets the seed for reproducibility."""

diff --git a/tests/fixtures/data.py b/tests/fixtures/data.py
@@ -47,8 +47,11 @@ def tiny_ft_dataloader(tiny_ft_dataset_path: str,
         'timeout': 0
     })
 
-    return build_finetuning_dataloader(
+    dataloader = build_finetuning_dataloader(
         dataloader_cfg,
         mpt_tokenizer,
         device_batch_size,
-    )
+    ).dataloader
+
+    assert isinstance(dataloader, DataLoader)
+    return dataloader
diff --git a/tests/test_flash_triton_torch.py b/tests/test_flash_triton_torch.py
@@ -3,7 +3,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import OmegaConf as om
 
 

diff --git a/tests/test_hf_config.py b/tests/test_hf_config.py
@@ -9,7 +9,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import DictConfig
 from omegaconf import OmegaConf as om
 from transformers import AutoModelForCausalLM

diff --git a/tests/test_hf_v_mpt.py b/tests/test_hf_v_mpt.py
@@ -5,7 +5,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import OmegaConf as om
 
 from llmfoundry import COMPOSER_MODEL_REGISTRY
@@ -52,8 +51,6 @@ def test_compare_hf_v_mpt(attn_impl: str, dropout: float, alibi: bool,
     batch_size = 2  # set batch size
     device = 'cuda'  # set decive
 
-
-
     # get hf gpt2 cfg
     hf_cfg = om.create({
         'model': {

diff --git a/tests/test_init_fn.py b/tests/test_init_fn.py
@@ -8,7 +8,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
 from torch import nn

diff --git a/tests/test_llama_patch.py b/tests/test_llama_patch.py
@@ -6,7 +6,6 @@
 import pytest
 import torch
 import transformers
-from composer.utils import reproducibility
 from transformers.models.llama.modeling_llama import LlamaAttention
 
 from llmfoundry.models.layers.llama_attention_monkeypatch import (

diff --git a/tests/test_model.py b/tests/test_model.py
@@ -16,7 +16,7 @@
 from composer.core.precision import Precision, get_precision_context
 from composer.optim import DecoupledAdamW
 from composer.trainer.dist_strategy import prepare_fsdp_module
-from composer.utils import dist, get_device, reproducibility
+from composer.utils import dist, get_device
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
 from transformers import (AutoModelForCausalLM, AutoTokenizer, PreTrainedModel,

diff --git a/tests/test_onnx.py b/tests/test_onnx.py
@@ -4,7 +4,6 @@
 import pathlib
 
 import torch
-from composer.utils import reproducibility
 from transformers import AutoModelForCausalLM
 
 from llmfoundry import MPTConfig, MPTForCausalLM