LMFuser 0.2.0__tar.gz → 0.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: LMFuser
-Version: 0.2.0
+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

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

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

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

@@ -59,12 +59,20 @@ 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
@@ -83,12 +91,23 @@ class Wrapper(nn.Module):
 
 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:
@@ -98,11 +117,28 @@ class DDPWraper(nn.Module):
 
 
 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:
@@ -137,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')
@@ -378,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.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')