LMFuser 0.1.3__tar.gz → 0.3.0__tar.gz

{lmfuser-0.1.3 → lmfuser-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: LMFuser
-Version: 0.1.3
+Version: 0.3.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.3 → lmfuser-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "LMFuser"
-version = "0.1.3"
+version = "0.3.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.3 → lmfuser-0.3.0}/src/lmfuser/runners/ddp_runner.py

@@ -59,42 +59,86 @@ logger = logging.getLogger(__name__)
 T = TypeVar('T', bound=Conf)
 
 
+def _compile_kwargs(compile_mode: str) -> dict:
+    return {} if compile_mode == 'default' else {'mode': compile_mode}
+
+
 class Wrapper(nn.Module):
 
-    def __init__(self, model: nn.Module) -> None:
+    def __init__(self, model: nn.Module, compile_mode: str = 'disable') -> None:
         super().__init__()
         self.module = model.to(get_default_device())
-        self.forward = model.forward
+        if compile_mode != 'disable':
+            logger.critical(f'torch.compile enabled (mode={compile_mode})')
+            self.forward = torch.compile(model.forward, **_compile_kwargs(compile_mode))
+        else:
+            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):
 
-    def __init__(self, model: nn.Module) -> None:
+    def __init__(self, model: nn.Module, compile_mode: str = 'disable') -> None:
         super().__init__()
-        self.model = DDP(
-            model.to(get_default_device()), find_unused_parameters=True
-        ) if get_world_size() > 1 else Wrapper(model)
-        self.forward = self.model.forward
+        if get_world_size() > 1:
+            self.model = DDP(
+                model.to(get_default_device()), find_unused_parameters=True
+            )
+            if compile_mode != 'disable':
+                # compile the DDP module so Dynamo's DDPOptimizer splits the
+                # graph at gradient-bucket boundaries (keeps comm overlap).
+                # OptimizedModule forwards attribute access to the DDP module,
+                # so .module / .no_sync() paths keep working.
+                logger.critical(f'torch.compile enabled (mode={compile_mode})')
+                self.model = torch.compile(self.model, **_compile_kwargs(compile_mode))
+            self.forward = self.model.forward
+        else:
+            self.model = Wrapper(model, compile_mode=compile_mode)
+            self.forward = self.model.forward
 
     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):
-    def __init__(self, model: FSDPModule) -> None:
+    def __init__(self, model: FSDPModule, compile_mode: str = 'disable') -> None:
         super().__init__()
         self.model = model
+        if compile_mode == 'reduce-overhead':
+            logger.critical(
+                'compile_mode=reduce-overhead is unsupported with fsdp2 '
+                '(cudagraphs need static addresses; FSDP all-gathers allocate '
+                'dynamically) — downgrading to "default".')
+            compile_mode = 'default'
+        if compile_mode != 'disable':
+            logger.critical(f'torch.compile enabled (mode={compile_mode})')
+            # compile the callable, NOT stored as a submodule: registering the
+            # OptimizedModule would duplicate every parameter under a second
+            # name and break optimizer construction / state_dict.
+            object.__setattr__(self, '_compiled_fwd',
+                               torch.compile(model, **_compile_kwargs(compile_mode)))
+        else:
+            object.__setattr__(self, '_compiled_fwd', None)
 
     def forward(self, *args: Any, **kwargs: Any) -> Any:
+        if self._compiled_fwd is not None:
+            return self._compiled_fwd(*args, **kwargs)
         return self.model(*args, **kwargs) # type: ignore
 
     def __getattr__(self, name: str) -> Any:
@@ -129,6 +173,14 @@ class DDPRunnerConfig(RunerConf):
     lr_scheduler: LRSchedulerConfig = LRSchedulerConfig()
 
     dp_type = OptionArg(default='ddp', options=['ddp', 'fsdp2'])
+    # torch.compile for the training model. 'default' = Inductor fusion;
+    # 'reduce-overhead' additionally wraps CUDA graphs (ddp only — with fsdp2 it
+    # is downgraded to 'default': cudagraphs' static-address requirement is
+    # incompatible with FSDP's dynamically allocated all-gather buffers).
+    # Explicit-by-config on purpose: anything that changes numerics or the
+    # execution path must be declared in the YAML, never auto-detected.
+    # NOTE 'disable' not 'off': YAML 1.1 parses bare off/on/yes/no as booleans
+    compile_mode = OptionArg(default='disable', options=['disable', 'default', 'reduce-overhead'])
     model_precision = OptionArg(options=['fp32', 'fp16', 'bf16'], default='fp32')
     use_amp = BoolArg(default=False)
     amp_precision = OptionArg(options=['fp16', 'bf16'], default='fp16')
@@ -141,6 +193,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 +259,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(
@@ -362,12 +422,14 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                 model = model.bfloat16() # type: ignore
             elif self.config.model_precision.value() == 'fp32':
                 model = model.float() # type: ignore
+            compile_mode = self.config.compile_mode.value()
+            assert compile_mode is not None
             if isinstance(model, FSDPModule):
                 logger.critical(f'wrapping model with FSDPModule')
-                self._model = FSDP2Wrapper(model)
+                self._model = FSDP2Wrapper(model, compile_mode=compile_mode)
             else:
                 logger.critical(f'wrapping model with DDPWraper')
-                self._model = DDPWraper(model)
+                self._model = DDPWraper(model, compile_mode=compile_mode)
         assert self._model is not None
         return self._model
 

{lmfuser-0.1.3 → lmfuser-0.3.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.3 → lmfuser-0.3.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')

lmfuser-0.3.0/tests/test_fsdp2_compile.py

@@ -0,0 +1,72 @@
+"""FSDP2Wrapper + compile_mode coverage (synthetic, single GPU).
+
+Run:  torchrun --nproc_per_node=1 tests/test_fsdp2_compile.py
+
+Covers, per compile_mode in {disable, default, reduce-overhead}:
+  * forward/backward/optimizer step run under fully_shard;
+  * parameters are NOT duplicated by compilation (the OptimizedModule must not
+    be registered as a submodule);
+  * 'reduce-overhead' downgrades to 'default' with a critical log.
+"""
+import logging
+import os
+import sys
+
+import torch
+from torch import nn
+import torch.distributed as dist
+from torch.distributed.fsdp import fully_shard
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+from lmfuser.runners.ddp_runner import FSDP2Wrapper   # noqa: E402
+
+
+def build_model() -> nn.Module:
+    torch.manual_seed(0)
+    return nn.Sequential(
+        nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 64),
+    )
+
+
+def run_mode(mode: str) -> None:
+    model = build_model().cuda()
+    n_params_before = len(list(model.parameters()))
+    fully_shard(model)
+
+    records: list[str] = []
+    handler = logging.Handler()
+    handler.emit = lambda r: records.append(r.getMessage())   # type: ignore
+    logging.getLogger('lmfuser.runners.ddp_runner').addHandler(handler)
+
+    wrapper = FSDP2Wrapper(model, compile_mode=mode)
+
+    # ---- parameter registration must not be duplicated by compile ----
+    n_params_after = len(list(wrapper.parameters()))
+    assert n_params_after == n_params_before, \
+        f'[{mode}] parameter duplication: {n_params_before} -> {n_params_after}'
+    opt = torch.optim.AdamW(wrapper.parameters(), lr=1e-3)   # raises on duplicates
+
+    # ---- train a few steps ----
+    for _ in range(3):
+        x = torch.randn(8, 64, device='cuda')
+        loss = wrapper(x).square().mean()
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+    assert torch.isfinite(loss).item(), f'[{mode}] non-finite loss'
+
+    # ---- downgrade warning ----
+    if mode == 'reduce-overhead':
+        assert any('downgrading' in m for m in records), \
+            '[reduce-overhead] downgrade warning not logged'
+
+    print(f'PASS fsdp2 compile_mode={mode} (loss={loss.item():.4f})')
+
+
+if __name__ == '__main__':
+    dist.init_process_group(backend='nccl')
+    torch.cuda.set_device(0)
+    for mode in ('disable', 'default', 'reduce-overhead'):
+        run_mode(mode)
+    dist.destroy_process_group()
+    print('FSDP2_COMPILE_TESTS_PASS')