LMFuser 0.4.0__tar.gz → 0.4.1__tar.gz

{lmfuser-0.4.0 → lmfuser-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: LMFuser
-Version: 0.4.0
+Version: 0.4.1
 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.3.0
+Requires-Dist: lmfuser-data>=0.3.1
 Requires-Dist: numpy
 Requires-Dist: pandas
 Requires-Dist: pyarrow

{lmfuser-0.4.0 → lmfuser-0.4.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "LMFuser"
-version = "0.4.0"
+version = "0.4.1"
 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.3.0",
+  "LMFuser-data>=0.3.1",
   "tqdm",
   "wandb",
   "HyperArgs>=0.1.3",

{lmfuser-0.4.0 → lmfuser-0.4.1}/src/lmfuser/runners/ddp_runner.py

@@ -5,6 +5,8 @@ from collections import defaultdict
 import logging
 import os
 import random
+import hashlib
+import shutil
 from pathlib import Path
 import json
 from logging import Logger, getLogger
@@ -88,6 +90,7 @@ class _DevicePrefetcher:
         self._stream = torch.cuda.Stream(device=device)
         self._queue: _q.Queue = _q.Queue(maxsize=2)
         self._stop = False
+        self._error: BaseException | None = None
         self._thread = threading.Thread(target=self._run, daemon=True)
         self._thread.start()
 
@@ -99,30 +102,45 @@ class _DevicePrefetcher:
         return v
 
     def _run(self) -> None:
-        it = iter(self._loader)
-        while not self._stop:
-            try:
-                batch = next(it)
-            except StopIteration:
-                it = iter(self._loader)
-                batch = next(it)
-            with torch.cuda.stream(self._stream):
-                dev = {}
-                ok = True
-                for k, v in batch.items():
-                    if isinstance(v, Tensor):
-                        v = self._cast(v)
-                        dev[k] = v.pin_memory().to(self._device, non_blocking=True)
-                    else:
-                        dev[k] = v
-                ev = torch.cuda.Event()
-                ev.record(self._stream)
-            self._queue.put((dev, ev))
+        try:
+            it = iter(self._loader)
+            while not self._stop:
+                try:
+                    batch = next(it)
+                except StopIteration:
+                    it = iter(self._loader)
+                    batch = next(it)
+                with torch.cuda.stream(self._stream):
+                    dev = {}
+                    for k, v in batch.items():
+                        if isinstance(v, Tensor):
+                            v = self._cast(v)
+                            dev[k] = v.pin_memory().to(self._device, non_blocking=True)
+                        else:
+                            dev[k] = v
+                    ev = torch.cuda.Event()
+                    ev.record(self._stream)
+                self._queue.put((dev, ev))
+        except BaseException as e:     # never die silently: next() would block forever
+            self._error = e
+            self._queue.put(None)
 
     def next(self):
-        dev, ev = self._queue.get()
-        # make the COMPUTE stream wait for the copy — no host sync
-        torch.cuda.current_stream().wait_event(ev)
+        item = self._queue.get()
+        if item is None:
+            raise RuntimeError('device prefetch thread died') from self._error
+        dev, ev = item
+        compute = torch.cuda.current_stream()
+        # order the compute stream behind the copy...
+        compute.wait_event(ev)
+        for v in dev.values():
+            if isinstance(v, Tensor):
+                # ...and keep the allocator from recycling these blocks (they
+                # belong to the side stream's pool) until the compute stream is
+                # actually done with them. Without this the next prefetched
+                # batch can be DMA'd into the buffer the current step is still
+                # reading — silent, intermittent input corruption.
+                v.record_stream(compute)
         return dev
 
     def close(self) -> None:
@@ -339,7 +357,19 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             resume_path = config.resume_path.value()
             assert resume_path is not None, 'resume_path is None'
             self.load(resume_path)
-        self.config.seed = self.config.seed.parse(hash(f'original_seed_{self.config.seed.value()}|step_{self.step}'))
+        # Derive the data seed deterministically. `hash()` on a str is
+        # SipHash-randomized per interpreter (PYTHONHASHSEED), so it produced a
+        # DIFFERENT seed on every rank and every restart: ranks disagreed on
+        # task sampling, and the run was irreproducible.
+        # The step is folded in only when the data stream cannot place itself:
+        # with restored shard cursors the row permutations must stay identical
+        # to the ones the cursors were recorded against.
+        _base = self.config.seed.value()
+        _stable = getattr(self, '_resume_data_states', None) is not None
+        _key = f'original_seed_{_base}' if _stable else f'original_seed_{_base}|step_{self.step}'
+        self.config.seed = self.config.seed.parse(
+            int.from_bytes(hashlib.sha256(_key.encode()).digest()[:4], 'big') & 0x7FFFFFFF
+        )
 
         self.train_data_loaders = config.task_conf.get_train_dataloaders(
             resume_states=getattr(self, '_resume_data_states', None),
@@ -470,12 +500,18 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             if loader is not None and hasattr(loader, 'state_dict') else None
             for loader in self.train_data_loaders
         ]
-        if not any(s is not None for s in per_task):
+        # NOTE: the emptiness check must NOT short-circuit ahead of the
+        # all_gather_object below — that is a collective, and a rank returning
+        # early while the others enter it hangs the job until the NCCL timeout.
+        empty = not any(s is not None for s in per_task)
+        if empty and get_world_size() <= 1:
             return None
         if get_world_size() > 1:
             gathered: list[Any] = [None] * get_world_size()
             torch.distributed.all_gather_object(gathered, per_task)
-            if get_global_rank() != 0:
+            if get_global_rank() != 0 or all(
+                not any(s is not None for s in (r or [])) for r in gathered
+            ):
                 return None
             merged: list = []
             for task_idx in range(len(per_task)):
@@ -517,30 +553,35 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         data_states = self._collect_data_states()
 
         if get_global_rank() == 0:
-            path = Path(directory) / str(self.step)
-            os.makedirs(path, exist_ok=True)
-            self.model_loader.save_model(model, path) # type: ignore
+            final = Path(directory) / str(self.step)
+            # Build the checkpoint in a staging directory and publish it with a
+            # single rename. A crash midway through the multi-GB optimizer save
+            # used to leave weights without runner.json, and `load()` treats a
+            # missing runner.json as "start from step 1" — silently restarting
+            # a long run (LR warmup and all) instead of resuming it.
+            staging = Path(directory) / f'.{self.step}.incomplete'
+            shutil.rmtree(staging, ignore_errors=True)
+            os.makedirs(staging, exist_ok=True)
+
+            self.model_loader.save_model(model, staging) # type: ignore
 
             if data_states is not None:
-                tmp_path = path / 'data_state.json.tmp'
-                with open(tmp_path, 'w') as f:
+                with open(staging / 'data_state.json', 'w') as f:
                     json.dump(data_states, f)
-                os.replace(tmp_path, path / 'data_state.json')
 
-            optimizer_path = path / 'optimizer.pt'
-            torch.save(optim_states, optimizer_path)
+            torch.save(optim_states, staging / 'optimizer.pt')
+            torch.save(self.scheduler.state_dict(), staging / 'scheduler.pt')
 
-            scheduler_path = path / 'scheduler.pt'
-            torch.save(self.scheduler.state_dict(), scheduler_path)
-
-            runner_path = path / 'runner.json'
-            with open(runner_path, 'w') as f:
+            with open(staging / 'runner.json', 'w') as f:
                 f.write(json.dumps({
                     'step': self.step + 1,
                     'epoch': self.epoch,
                     'config': self.config.to_dict(),
                 }, indent=4))
 
+            shutil.rmtree(final, ignore_errors=True)   # re-save of the same step
+            os.replace(staging, final)
+
     @property
     def model(self) -> DDPWraper | FSDP2Wrapper:
         model = getattr(self, '_model', None)
@@ -1157,13 +1198,13 @@ class DDPRunner(Runner[DDPRunnerConfig]):
 
         optimizer_path = Path(directory) / 'optimizer.pt'
         if optimizer_path.exists():
-            self._optimizer_states = torch.load(optimizer_path)
+            self._optimizer_states = torch.load(optimizer_path, map_location='cpu')
         else:
             logger.warning(f'optimizer.pt not found in {directory}, skip loading optimizer')
 
         scheduler_path = Path(directory) / 'scheduler.pt'
         if scheduler_path.exists():
-            self._scheduler_states = torch.load(scheduler_path)
+            self._scheduler_states = torch.load(scheduler_path, map_location='cpu')
         else:
             logger.warning(f'scheduler.pt not found in {directory}, skip loading scheduler')
 

lmfuser-0.4.1/tests/test_runner_robustness.py

@@ -0,0 +1,114 @@
+"""Failure-mode tests for DDPRunner bookkeeping (0.4.1).
+
+Run:  python tests/test_runner_robustness.py
+
+Covers issues found in review, all of which failed silently:
+  1. the data seed is deterministic — equal across ranks and across processes
+  2. the data seed is stable across a resume that restored data cursors
+     (the shard-cursor row permutations are derived from it)
+  3. checkpoints are published atomically — a crash mid-save never leaves a
+     directory that load() would treat as "start from step 1"
+  4. optimizer/scheduler state loads onto CPU, not every rank onto cuda:0
+"""
+import os
+import sys
+import json
+import hashlib
+import shutil
+import subprocess
+import tempfile
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+
+
+def _derive(base: int, step, stable: bool) -> int:
+    """Mirror of the runner's seed derivation."""
+    key = f'original_seed_{base}' if stable else f'original_seed_{base}|step_{step}'
+    return int.from_bytes(hashlib.sha256(key.encode()).digest()[:4], 'big') & 0x7FFFFFFF
+
+
+def test_seed_is_deterministic_across_processes() -> None:
+    """`hash()` on a str is PYTHONHASHSEED-randomized: every rank used to get a
+    different seed, so ranks disagreed on task sampling and no run was
+    reproducible."""
+    snippet = (
+        'import hashlib;'
+        "k='original_seed_42|step_1';"
+        "print(int.from_bytes(hashlib.sha256(k.encode()).digest()[:4],'big') & 0x7FFFFFFF)"
+    )
+    seeds = {
+        subprocess.run([sys.executable, '-c', snippet], capture_output=True,
+                       text=True).stdout.strip()
+        for _ in range(4)
+    }
+    assert len(seeds) == 1, f'seed differs across processes: {seeds}'
+
+    old = {
+        subprocess.run([sys.executable, '-c',
+                        "print(hash('original_seed_42|step_1'))"],
+                       capture_output=True, text=True).stdout.strip()
+        for _ in range(4)
+    }
+    assert len(old) > 1, 'expected the old hash() to be randomized — test is moot'
+    print(f'PASS 1: seed identical across processes ({seeds.pop()}); '
+          f'old hash() gave {len(old)} different values')
+
+
+def test_seed_stable_when_cursors_restored() -> None:
+    """Shard cursors index a permutation derived from the seed; folding the
+    step into it on resume would silently invalidate every cursor."""
+    fresh_step1 = _derive(42, 1, stable=False)
+    fresh_step9 = _derive(42, 900_000, stable=False)
+    assert fresh_step1 != fresh_step9, 'without cursors the seed should follow the step'
+
+    resumed_a = _derive(42, 1, stable=True)
+    resumed_b = _derive(42, 900_000, stable=True)
+    assert resumed_a == resumed_b, 'with cursors restored the seed must not move'
+    print('PASS 2: seed stable across a cursor-restoring resume, step-varying otherwise')
+
+
+def test_checkpoint_is_atomic() -> None:
+    """A checkpoint directory must never be visible while incomplete: load()
+    treats a missing runner.json as step 1 and silently restarts the run."""
+    src = open(os.path.join(os.path.dirname(__file__), '..', 'src', 'lmfuser',
+                            'runners', 'ddp_runner.py')).read()
+
+    assert '.incomplete' in src, 'save() does not stage the checkpoint'
+    assert 'os.replace(staging, final)' in src, 'save() does not publish by rename'
+    # the publish must come after every writer
+    pub = src.index('os.replace(staging, final)')
+    for writer in ("save_model(model, staging)", "staging / 'optimizer.pt'",
+                   "staging / 'scheduler.pt'", "staging / 'runner.json'"):
+        assert src.index(writer) < pub, f'{writer} runs after the rename'
+
+    # a staged directory is invisible to a "latest numeric checkpoint" scan,
+    # which is how supervisors pick what to resume from
+    tmp = tempfile.mkdtemp(prefix='ckpt_atomic_')
+    try:
+        os.makedirs(os.path.join(tmp, '.25000.incomplete'))
+        os.makedirs(os.path.join(tmp, '20000'))
+        numeric = [d for d in os.listdir(tmp) if d.isdigit()]
+        assert numeric == ['20000'], f'staging dir leaked into the scan: {numeric}'
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+    print('PASS 3: checkpoint staged then published by rename; staging dir not resumable')
+
+
+def test_state_loads_on_cpu() -> None:
+    """Without map_location every rank deserializes rank0's CUDA optimizer
+    state onto cuda:0 — ~4GB x N ranks transiently, OOM on the larger models."""
+    src = open(os.path.join(os.path.dirname(__file__), '..', 'src', 'lmfuser',
+                            'runners', 'ddp_runner.py')).read()
+    for name in ('optimizer_path', 'scheduler_path'):
+        idx = src.index(f'torch.load({name}')
+        call = src[idx:idx + 120]
+        assert "map_location='cpu'" in call, f'torch.load({name}) lacks map_location'
+    print('PASS 4: optimizer/scheduler state deserializes on CPU')
+
+
+if __name__ == '__main__':
+    test_seed_is_deterministic_across_processes()
+    test_seed_stable_when_cursors_restored()
+    test_checkpoint_is_atomic()
+    test_state_loads_on_cpu()
+    print('ALL PASS')