LMFuser 0.0.7__tar.gz → 0.0.9__tar.gz

{lmfuser-0.0.7 → lmfuser-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: LMFuser
-Version: 0.0.7
+Version: 0.0.9
 Summary: The LMFuser training framework.
 Project-URL: Homepage, https://github.com/TYTTYTTYT/LMFuser
 Project-URL: Documentation, https://github.com/TYTTYTTYT/LMFuser

{lmfuser-0.0.7 → lmfuser-0.0.9}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "LMFuser"
-version = "0.0.7"
+version = "0.0.9"
 requires-python = ">= 3.11"
 description = "The LMFuser training framework."
 readme = "README.md"

{lmfuser-0.0.7 → lmfuser-0.0.9}/src/lmfuser/runners/ddp_runner.py

@@ -190,6 +190,20 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             rank=get_global_rank(),
         )
 
+        self.test_data_loaders = config.task_conf.get_test_dataloaders(
+            batch_size=config.sub_batch_size.value(), # type: ignore
+            seed=config.seed.value(), # type: ignore
+            shuffle=config.shuffle_dataset.value(), # type: ignore
+            prefetch_factor=config.row_prefetch.value(), # type: ignore
+            num_workers=config.num_row_workers.value(), # type: ignore
+            ignore_error=config.ignore_data_error.value(), # type: ignore
+            qps=config.data_row_qps.value(),
+            instruct_timeout=config.instruct_timeout.value(), # type: ignore
+            worker_timeout=config.worker_timeout.value(), # type: ignore
+            world_size=get_world_size(),
+            rank=get_global_rank(),
+        )
+
         self.train_task_idxs: list[int] = []
         for idx, loader in enumerate(self.train_data_loaders):
             if loader is not None:
@@ -200,6 +214,11 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             if loader is not None:
                 self.eval_task_idxs.append(idx)
 
+        self.test_task_idxs: list[int] = []
+        for idx, loader in enumerate(self.test_data_loaders):
+            if loader is not None:
+                self.test_task_idxs.append(idx)
+
         assert len(self.train_task_idxs) + len(self.eval_task_idxs) > 0, 'At least one train or eval task must be provided.'
 
         self.task_weights: list[float] = [w.value() for w in config.task_conf.task_weights] # type: ignore
@@ -210,6 +229,9 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         self.train_iters = [iter(loader) if loader is not None else None for loader in self.train_data_loaders]
         self.eval_iters = [iter(loader) if loader is not None else None for loader in self.eval_data_loaders]
 
+        self._all_eval_results: dict[str, list[dict[str, Any]]] = {}
+        self._test_results: dict[str, dict[str, Any]] = {}
+
     def sample_train_task_id(self) -> int:
         '''
         by default, randomly select a task from the task list.
@@ -261,6 +283,7 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                 logger.critical(f'casting model to fp16')
                 model = model.half()
             self._model = DDPWraper(model)
+        assert self._model is not None
         return self._model
 
     def _should_stop(self) -> bool:
@@ -550,7 +573,7 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                 dataloader,
                 dynamic_ncols=True,
                 desc=f'evaluating task {task.__class__.__name__}',
-                position=2,
+                position=3,
                 disable=True if get_global_rank() != 0 else False,
                 leave=False,
                 unit='batch'
@@ -577,6 +600,64 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         for k, v in metrics.items():
             self.step_log({f'{task.__class__.__name__}/dev/{k}': v})
 
+        if task.__class__.__name__ not in self._all_eval_results:
+            self._all_eval_results[task.__class__.__name__] = []
+        self._all_eval_results[task.__class__.__name__].append({'step': self.step, 'metrics': metrics})
+
+    def _test_one_task(self, task_id: int, **kwargs: Any) -> None:
+        os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+        dataloader = self.test_data_loaders[task_id]
+        task = self.tasks[task_id]
+
+        ckpt_path = self.config.checkpoint_directory.value()
+        assert ckpt_path is not None
+        test_model_path = task.set_test_model_path(ckpt_path, self._all_eval_results)
+        if test_model_path is not None:
+            if isinstance(test_model_path, int):
+                test_model_path = Path(ckpt_path) / str(test_model_path)
+            
+            self.model_loader.model_path = test_model_path
+            self._model = None
+
+        batch_list: list[dict[str, list[Any]]] = []
+        with ExitStack() as stack:
+            stack.enter_context(torch.no_grad())
+            if get_world_size() > 1:
+                stack.enter_context(self.model.model.no_sync())
+            for batch in tqdm(
+                dataloader,
+                dynamic_ncols=True,
+                desc=f'testing task {task.__class__.__name__}',
+                position=4,
+                disable=True if get_global_rank() != 0 else False,
+                leave=False,
+                unit='batch'
+            ):
+                result = task.eval_step(
+                    self.model, # type: ignore
+                    self._batch_to_device(batch),
+                    self.step,
+                    get_default_device()
+                )
+                batch_list.append(result)
+
+            if len(batch_list) == 0:
+                raise RuntimeError(f'No test data found in rank {get_global_rank()} '
+                                f'for task {task.__class__.__name__}')
+            cat_result = batch_list[0]
+            for batch in batch_list[1:]:
+                for k, v in batch.items():
+                    cat_result[k] += v
+        all_result = batch_all_gather(cat_result)
+
+        with torch.no_grad():
+            metrics = task.cal_dev_metric(all_result)
+        for k, v in metrics.items():
+            self.step_log({f'{task.__class__.__name__}/test/{k}': v})
+
+        self._test_results[task.__class__.__name__] = {'step': self.step, 'metrics': metrics}
+
     def eval(self, *args: Any, **kwargs: Any) -> None:
         for task_id in tqdm(
             self.eval_task_idxs, 
@@ -589,6 +670,48 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         ):
             self._eval_one_task(task_id)
 
+        if get_global_rank() != 0:
+            return
+        ckpt_path = self.config.checkpoint_directory.value()
+        assert ckpt_path is not None
+        try:
+            eval_results = json.dumps(self._all_eval_results, indent=2)
+            ckpt_dir = Path(ckpt_path)
+            if not ckpt_dir.exists():
+                ckpt_dir.mkdir(parents=True)
+            eval_path = ckpt_dir / 'eval_results.json'
+            with open(eval_path, 'w') as f:
+                f.write(eval_results)
+        except:
+            logger.warning(f'Failed to save eval results to {ckpt_path}')
+
+    def test(self, *args: Any, **kwargs: Any) -> None:
+        for task_id in tqdm(
+            self.test_task_idxs,
+            dynamic_ncols=True, 
+            desc=f'testing {len(self.test_task_idxs)} tasks...',
+            position=3,
+            disable=True if get_global_rank() != 0 else False,
+            leave=False,
+            unit='task'
+        ):
+            self._test_one_task(task_id)
+
+        if get_global_rank() != 0:
+            return
+        ckpt_path = self.config.checkpoint_directory.value()
+        assert ckpt_path is not None
+        try:
+            test_results = json.dumps(self._test_results, indent=2)
+            ckpt_dir = Path(ckpt_path)
+            if not ckpt_dir.exists():
+                ckpt_dir.mkdir(parents=True)
+            test_path = ckpt_dir / 'test_results.json'
+            with open(test_path, 'w') as f:
+                f.write(test_results)
+        except:
+            logger.warning(f'Failed to save test results to {ckpt_path}')
+
     def produce(self, *args: Any, **kwargs: Any) -> None:
         raise NotImplementedError('produce method not implemented')
 

{lmfuser-0.0.7 → lmfuser-0.0.9}/src/lmfuser/runners/runner.py

@@ -25,7 +25,11 @@ class Runner(ABC, Generic[ConfType]):
 
     @abstractmethod
     def eval(self, *args: Any, **kwargs: Any) -> None:
-        raise NotImplementedError('train method not implemented')
+        raise NotImplementedError('eval method not implemented')
+
+    @abstractmethod
+    def test(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError('test method not implemented')
 
     @abstractmethod
     def produce(self, *args: Any, **kwargs: Any) -> None:

{lmfuser-0.0.7 → lmfuser-0.0.9}/src/lmfuser/task.py

@@ -1,5 +1,5 @@
 from typing import Any, Callable
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
 
 import torch
 from torch import nn
@@ -12,8 +12,32 @@ from hyperargs import Conf, StrArg, FloatArg, IntArg, OptionArg, add_dependency,
 def scanner_type_list() -> list[str]:
     return list(Scanner.all_subclass_names())
 
+
+class EmptyDataLoader:
+    '''
+        EmptyDataLoader is a class that provides empty data for running with no data requierment.
+    '''
+    def __init__(self, init_step: int = 0) -> None:
+        self.init_step = init_step
+
+    @property
+    def epoch(self) -> int:
+        return 0
+
+    def __iter__(self) -> Iterator[Batch]:
+        def it_wrap() -> Iterator[Batch]:
+            while True:
+                yield {'step': torch.tensor(self.init_step)}
+                self.init_step += 1
+        return it_wrap()
+
+
 @add_dependency('num_train_data_path', 'train_data_path_list')
 @add_dependency('num_train_data_path', 'train_data_weights')
+@add_dependency('num_eval_data_path', 'eval_data_path_list')
+@add_dependency('num_eval_data_path', 'eval_data_weights')
+@add_dependency('num_test_data_path', 'test_data_path_list')
+@add_dependency('num_test_data_path', 'test_data_weights')
 class TaskBase(Conf, SubclassTracer):
     num_train_data_path = IntArg(1, min_value=0)
     train_data_path_list = [StrArg('Enther the path to the data file.')]
@@ -23,13 +47,19 @@ class TaskBase(Conf, SubclassTracer):
     eval_data_path_list = [StrArg('Enther the path to the data file.')]
     eval_data_weights = [FloatArg(1.0, min_value=0.0, max_value=1.0)]
 
+    num_test_data_path = IntArg(0, min_value=0)
+    test_data_path_list: list[StrArg] = []
+    test_data_weights: list[FloatArg] = []
+
     scanner_type = OptionArg(default='C4Scanner', option_fn=scanner_type_list)
 
-    train_dataloader_type = OptionArg(default='single file', options=['single file', 'sharded'])
-    eval_dataloader_type = OptionArg(default='single file', options=['single file', 'sharded'])
+    train_dataloader_type = OptionArg(default='single file', options=['single file', 'sharded', 'empty'])
+    eval_dataloader_type = OptionArg(default='single file', options=['single file', 'sharded', 'empty'])
+    test_dataloader_type = OptionArg(default='single file', options=['single file', 'sharded', 'empty'])
 
-    _train_dataloader: DataLoader | None | PyTorchDataLoader = None
-    _eval_dataloader: DataLoader | None | PyTorchDataLoader = None
+    _train_dataloader: DataLoader | None | PyTorchDataLoader | EmptyDataLoader = None
+    _eval_dataloader: DataLoader | None | PyTorchDataLoader | EmptyDataLoader = None
+    _test_dataloader: DataLoader | None | PyTorchDataLoader | EmptyDataLoader = None
 
     @monitor_on('num_train_data_path')
     def set_train_path_list(self) -> None:
@@ -50,8 +80,19 @@ class TaskBase(Conf, SubclassTracer):
             self.eval_data_path_list = self.eval_data_path_list[:num]
             self.eval_data_weights = self.eval_data_weights[:num]
         elif len(self.eval_data_path_list) < num:
-            self.eval_data_path_list += [StrArg('Enther the path to the data file.')] * (num - len(self.train_data_path_list))
-            self.eval_data_weights += [FloatArg(1.0, min_value=0.0, max_value=1.0)] * (num - len(self.train_data_weights))
+            self.eval_data_path_list += [StrArg('Enther the path to the data file.')] * (num - len(self.eval_data_path_list))
+            self.eval_data_weights += [FloatArg(1.0, min_value=0.0, max_value=1.0)] * (num - len(self.eval_data_weights))
+
+    @monitor_on('num_test_data_path')
+    def set_test_path_list(self) -> None:
+        num = self.num_test_data_path.value()
+        assert isinstance(num, int)
+        if len(self.test_data_path_list) > num:
+            self.test_data_path_list = self.test_data_path_list[:num]
+            self.test_data_path_list = self.test_data_path_list[:num]
+        elif len(self.test_data_path_list) < num:
+            self.test_data_path_list += [StrArg('Enther the path to the data file.')] * (num - len(self.test_data_path_list))
+            self.test_data_weights += [FloatArg(1.0, min_value=0.0, max_value=1.0)] * (num - len(self.test_data_weights))
 
     def _get_train_dataloader(
         self,
@@ -66,7 +107,7 @@ class TaskBase(Conf, SubclassTracer):
         num_workers: int,
         rank: int,
         world_size: int
-    ) -> None | DataLoader | PyTorchDataLoader:
+    ) -> None | DataLoader | PyTorchDataLoader | EmptyDataLoader:
         if self.num_train_data_path.value() == 0:
             return None
         if self._train_dataloader is not None:
@@ -113,6 +154,8 @@ class TaskBase(Conf, SubclassTracer):
                 collate_fn=self.get_collate_fn(),
                 drop_last=False
             )
+        elif dataloader_type == 'empty':
+            self._train_dataloader = EmptyDataLoader(init_step=0)
         else:
             raise ValueError(f'Unknown dataloader type: {dataloader_type}')
 
@@ -131,7 +174,7 @@ class TaskBase(Conf, SubclassTracer):
         num_workers: int,
         rank: int,
         world_size: int
-    ) -> None | DataLoader | PyTorchDataLoader:
+    ) -> None | DataLoader | PyTorchDataLoader | EmptyDataLoader:
         if self.num_eval_data_path.value() == 0:
             return None
         if self._eval_dataloader is not None:
@@ -181,11 +224,83 @@ class TaskBase(Conf, SubclassTracer):
                 collate_fn=self.get_collate_fn(),
                 drop_last=False
             )
+        elif dataloader_type == 'empty':
+            self._eval_dataloader = EmptyDataLoader(init_step=0)
         else:
             raise ValueError(f'Unknown dataloader type: {dataloader_type}')
 
         return self._eval_dataloader
 
+    def _get_test_dataloader(
+        self,
+        batch_size: int,
+        seed: int,
+        shuffle: bool,
+        prefetch_factor: int,
+        ignore_error: bool,
+        qps: float | None,
+        instruct_timeout: float,
+        worker_timeout: float,
+        num_workers: int,
+        rank: int,
+        world_size: int
+    ) -> None | DataLoader | PyTorchDataLoader | EmptyDataLoader:
+        if self.num_test_data_path.value() == 0:
+            return None
+        if self._test_dataloader is not None:
+            return self._test_dataloader
+        path_list = [p.value() for p in self.test_data_path_list]
+        weight_list = [w.value() for w in self.test_data_weights]
+        scanner_type = self.scanner_type.value()
+        assert scanner_type is not None, 'scanner_type is None'
+
+        dataloader_type = self.test_dataloader_type.value()
+        assert dataloader_type is not None, 'dataloader_type is None'
+
+        dataloader_type = self.test_dataloader_type.value()
+        assert dataloader_type in ('sharded', 'single file'), f'Unknown dataloader type: {dataloader_type}'
+
+        if dataloader_type == 'sharded':
+            self._test_dataloader = DataLoader(
+                batch_size=batch_size,
+                path_list=path_list, # type: ignore
+                distributor_weights=weight_list, # type: ignore
+                scanner_type=Scanner.get_subclass(scanner_type),
+                seed=seed,
+                shuffle=shuffle,
+                pre_fetch_factor=prefetch_factor,
+                ignore_error=ignore_error,
+                qps=qps,
+                instruct_timeout=instruct_timeout,
+                worker_timeout=worker_timeout,
+                num_workers=num_workers,
+                map_fn=self.get_row_processor(),
+                flow_fn=self.get_flow_processor(),
+                batch_map_fn=self.get_batch_processor(),
+                rank_idx=rank,
+                num_ranks=world_size,
+            )
+        elif dataloader_type == 'single file':
+            self._test_dataloader = PyTorchDataLoader(
+                batch_size=batch_size,
+                path_list=path_list, # type: ignore
+                scanner_type=Scanner.get_subclass(scanner_type),
+                seed=seed,
+                shuffle=shuffle,
+                pre_fetch_factor=prefetch_factor,
+                num_workers=num_workers,
+                num_ranks=world_size,
+                rank_idx=rank,
+                collate_fn=self.get_collate_fn(),
+                drop_last=False
+            )
+        elif dataloader_type == 'empty':
+            self._test_dataloader = EmptyDataLoader(init_step=0)
+        else:
+            raise ValueError(f'Unknown dataloader type: {dataloader_type}')
+
+        return self._test_dataloader
+
     def train_step(
         self, model: nn.Module,
         batch: Batch,
@@ -209,6 +324,17 @@ class TaskBase(Conf, SubclassTracer):
     def cal_dev_metric(self, eval_outputs: dict[str, list[Any]]) -> dict[str, Any]:
         raise NotImplementedError('Please implement this method in child class')
 
+    def set_test_model_path(self, checkpoints_path: str, dev_metrics: dict[str, list[dict[str, Any]]]) -> str | int | None:
+        '''Set the path of the test model based on the dev metrics.
+        Args:
+            checkpoints_path (str): The path to the checkpoints directory.
+            dev_metrics (dict[str, list[dict[str, Any]]]): The development metrics for each task.
+
+        Returns:
+            str | int | None: The path or the checkpoint step to the test model, or None to use the last step model.
+        '''
+        return None
+
     def get_row_processor(self) -> Callable[[Row], Row] | None:
         return None
 
@@ -277,7 +403,7 @@ class Tasks(Conf):
         num_workers: int,
         rank: int,
         world_size: int
-    ) -> list[DataLoader | None | PyTorchDataLoader]:
+    ) -> list[DataLoader | None | PyTorchDataLoader | EmptyDataLoader]:
         return [
             task.conf._get_train_dataloader(
                 batch_size=batch_size,
@@ -308,7 +434,7 @@ class Tasks(Conf):
         num_workers: int,
         rank: int,
         world_size: int
-    ) -> list[DataLoader | None | PyTorchDataLoader]:
+    ) -> list[DataLoader | None | PyTorchDataLoader | EmptyDataLoader]:
         return [
             task.conf._get_eval_dataloader(
                 batch_size=batch_size,
@@ -326,9 +452,33 @@ class Tasks(Conf):
             for task in self.tasks
         ]
 
-if __name__ == '__main__':
-    class TaskTemp(Task):
-        pass
-
-    conf = Tasks.parse_command_line()
-    print(conf)
+    def get_test_dataloaders(
+        self,
+        batch_size: int,
+        seed: int,
+        shuffle: bool,
+        prefetch_factor: int,
+        ignore_error: bool,
+        qps: float | None,
+        instruct_timeout: float,
+        worker_timeout: float,
+        num_workers: int,
+        rank: int,
+        world_size: int
+    ) -> list[DataLoader | None | PyTorchDataLoader | EmptyDataLoader]:
+        return [
+            task.conf._get_test_dataloader(
+                batch_size=batch_size,
+                seed=seed,
+                shuffle=shuffle,
+                prefetch_factor=prefetch_factor,
+                ignore_error=ignore_error,
+                qps=qps,
+                instruct_timeout=instruct_timeout,
+                worker_timeout=worker_timeout,
+                num_workers=num_workers,
+                rank=rank,
+                world_size=world_size,
+            )
+            for task in self.tasks
+        ]