[test] merge old components to test to model zoo

colossalai/testing/__init__.py

@@ -9,6 +9,7 @@
 )
 from .pytest_wrapper import run_on_environment_flag
 from .utils import (
+    DummyDataloader,
     clear_cache_before_run,
     free_port,
     parameterize,
@@ -34,4 +35,5 @@
     "run_on_environment_flag",
     "check_state_dict_equal",
     "assert_hf_output_close",
+    "DummyDataloader",
 ]

colossalai/testing/utils.py

@@ -273,3 +273,24 @@ def _clear_cache(*args, **kwargs):
         return _clear_cache
 
     return _wrap_func
+
+
+class DummyDataloader:
+    def __init__(self, data_gen_fn: Callable, length: int = 10):
+        self.data_gen_fn = data_gen_fn
+        self.length = length
+        self.step = 0
+
+    def __iter__(self):
+        self.step = 0
+        return self
+
+    def __next__(self):
+        if self.step < self.length:
+            self.step += 1
+            return self.data_gen_fn()
+        else:
+            raise StopIteration
+
+    def __len__(self):
+        return self.length