LMFuser 0.1.2__tar.gz → 0.2.0__tar.gz

{lmfuser-0.1.2 → lmfuser-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: LMFuser
-Version: 0.1.2
+Version: 0.2.0
 Summary: The LMFuser training framework.
 Project-URL: Homepage, https://github.com/TYTTYTTYT/LMFuser
 Project-URL: Documentation, https://github.com/TYTTYTTYT/LMFuser
@@ -17,7 +17,7 @@ Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Topic :: Utilities
 Requires-Python: >=3.11
 Requires-Dist: hyperargs>=0.1.3
-Requires-Dist: lmfuser-data>=0.1.1
+Requires-Dist: lmfuser-data>=0.2.0
 Requires-Dist: numpy
 Requires-Dist: pandas
 Requires-Dist: pyarrow

{lmfuser-0.1.2 → lmfuser-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "LMFuser"
-version = "0.1.2"
+version = "0.2.0"
 requires-python = ">= 3.11"
 description = "The LMFuser training framework."
 readme = "README.md"
@@ -20,7 +20,7 @@ dependencies = [
   "pandas",
   "pyarrow",
   "torch>=2.4.0",
-  "LMFuser-data>=0.1.1",
+  "LMFuser-data>=0.2.0",
   "tqdm",
   "wandb",
   "HyperArgs>=0.1.3",

{lmfuser-0.1.2 → lmfuser-0.2.0}/src/lmfuser/runners/ddp_runner.py

@@ -66,11 +66,19 @@ class Wrapper(nn.Module):
         self.module = model.to(get_default_device())
         self.forward = model.forward
 
+    def no_sync(self):
+        # duck-type DDP for the accumulation path: a single process has no
+        # gradient sync to skip (latent crash: single GPU + grad accumulation)
+        from contextlib import nullcontext
+        return nullcontext()
+
     def __getattr__(self, name: str) -> Any:
         try:
             return super().__getattr__(name)
-        except:
-            return self.module.__getattr__(name)
+        except AttributeError:
+            # full lookup (instance dict included) — nn.Module.__getattr__ alone
+            # misses plain attributes like `.config`
+            return getattr(self.module, name)
 
 
 class DDPWraper(nn.Module):
@@ -85,8 +93,8 @@ class DDPWraper(nn.Module):
     def __getattr__(self, name: str) -> Any:
         try:
             return super().__getattr__(name)
-        except:
-            return self.model.module.__getattr__(name)
+        except AttributeError:
+            return getattr(self.model.module, name)
 
 
 class FSDP2Wrapper(nn.Module):
@@ -141,6 +149,11 @@ class DDPRunnerConfig(RunerConf):
     shuffle_dataset = BoolArg(default=True)
     row_prefetch = IntArg(0, min_value=0)
     num_row_workers = IntArg(1, min_value=1)
+    # batch-mode loader (train_dataloader_type: batch): TOTAL workers per rank,
+    # shared-memory ring depth, and per-slot capacity
+    num_batch_workers = IntArg(4, min_value=1)
+    batch_queue_depth = IntArg(4, min_value=2)
+    batch_slot_mb = IntArg(128, min_value=1)
 
     resume_training = BoolArg(default=False)
     resume_path = StrArg(default=None, allow_none=True)
@@ -202,6 +215,9 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             worker_timeout=config.worker_timeout.value(), # type: ignore
             world_size=get_world_size(),
             rank=get_global_rank(),
+            num_batch_workers=config.num_batch_workers.value(), # type: ignore
+            batch_queue_depth=config.batch_queue_depth.value(), # type: ignore
+            batch_slot_mb=config.batch_slot_mb.value(), # type: ignore
         )
 
         self.eval_data_loaders = config.task_conf.get_eval_dataloaders(
@@ -772,6 +788,7 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             self.model_loader.model_path = test_model_path
             self._model = None
 
+        self.model.eval()  # ensure eval mode after potential model reload
         batch_list: list[dict[str, list[Any]]] = []
         with ExitStack() as stack:
             stack.enter_context(torch.no_grad())

{lmfuser-0.1.2 → lmfuser-0.2.0}/src/lmfuser/task.py

@@ -5,7 +5,7 @@ import torch
 from torch import nn
 from lmfuser_data.interfaces import Batch, Row
 from lmfuser_data.scanners import Scanner
-from lmfuser_data import DataLoader, PyTorchDataLoader
+from lmfuser_data import DataLoader, PyTorchDataLoader, BatchDataLoader
 from lmfuser_data.interfaces import SubclassTracer
 from hyperargs import Conf, StrArg, FloatArg, IntArg, OptionArg, add_dependency, monitor_on
 
@@ -53,7 +53,7 @@ class TaskBase(Conf, SubclassTracer):
 
     scanner_type = OptionArg(default='C4Scanner', option_fn=scanner_type_list)
 
-    train_dataloader_type = OptionArg(default='single file', options=['single file', 'sharded', 'empty'])
+    train_dataloader_type = OptionArg(default='single file', options=['single file', 'sharded', 'batch', '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'])
 
@@ -106,8 +106,11 @@ class TaskBase(Conf, SubclassTracer):
         worker_timeout: float,
         num_workers: int,
         rank: int,
-        world_size: int
-    ) -> None | DataLoader | PyTorchDataLoader | EmptyDataLoader:
+        world_size: int,
+        num_batch_workers: int = 4,
+        batch_queue_depth: int = 4,
+        batch_slot_mb: int = 128,
+    ) -> None | DataLoader | PyTorchDataLoader | BatchDataLoader | EmptyDataLoader:
         if self.num_train_data_path.value() == 0:
             return None
         if self._train_dataloader is not None:
@@ -118,9 +121,30 @@ class TaskBase(Conf, SubclassTracer):
         assert scanner_type is not None, 'scanner_type is None'
 
         dataloader_type = self.train_dataloader_type.value()
-        assert dataloader_type in ('sharded', 'single file'), f'Unknown dataloader type: {dataloader_type}'
+        assert dataloader_type in ('sharded', 'single file', 'batch', 'empty'), \
+            f'Unknown dataloader type: {dataloader_type}'
 
-        if dataloader_type == 'sharded':
+        if dataloader_type == 'batch':
+            self._train_dataloader = BatchDataLoader(
+                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,
+                map_fn=self.get_row_processor(),
+                flow_fn=self.get_flow_processor(),
+                collate_fn=self.get_collate_fn(),
+                batch_map_fn=self.get_batch_processor(),
+                ignore_error=ignore_error,
+                num_workers=num_batch_workers,
+                queue_depth=batch_queue_depth,
+                slot_mb=batch_slot_mb,
+                num_ranks=world_size,
+                rank_idx=rank,
+                worker_timeout=worker_timeout,
+            )
+        elif dataloader_type == 'sharded':
             self._train_dataloader = DataLoader(
                 batch_size=batch_size,
                 path_list=path_list, # type: ignore
@@ -402,8 +426,11 @@ class Tasks(Conf):
         worker_timeout: float,
         num_workers: int,
         rank: int,
-        world_size: int
-    ) -> list[DataLoader | None | PyTorchDataLoader | EmptyDataLoader]:
+        world_size: int,
+        num_batch_workers: int = 4,
+        batch_queue_depth: int = 4,
+        batch_slot_mb: int = 128,
+    ) -> list[DataLoader | None | PyTorchDataLoader | BatchDataLoader | EmptyDataLoader]:
         return [
             task.conf._get_train_dataloader(
                 batch_size=batch_size,
@@ -417,6 +444,9 @@ class Tasks(Conf):
                 num_workers=num_workers,
                 rank=rank,
                 world_size=world_size,
+                num_batch_workers=num_batch_workers,
+                batch_queue_depth=batch_queue_depth,
+                batch_slot_mb=batch_slot_mb,
             )
             for task in self.tasks
         ]

{lmfuser-0.1.2 → lmfuser-0.2.0}/src/lmfuser/utils.py

@@ -42,7 +42,11 @@ def dist_init() -> None:
     device_type = get_default_device_type()
 
     if device_type == 'cuda':
-        dist.init_process_group(backend='nccl')
+        import datetime as _dt
+        # 1h timeout: a transient data stall on one rank (e.g. a shard
+        # download hang) must not nuke the whole DDP group (default 10min
+        # killed a pretrain run mid-flight)
+        dist.init_process_group(backend='nccl', timeout=_dt.timedelta(seconds=3600))
         torch.cuda.set_device(get_local_rank())
     elif device_type == 'npu':
         dist.init_process_group(backend='hccl')