LMFuser 0.3.0__tar.gz → 0.4.1__tar.gz

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

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

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

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

{lmfuser-0.3.0 → lmfuser-0.4.1}/src/lmfuser/optimizers.py

@@ -10,7 +10,7 @@ from torch.optim import (
     Adadelta,
     Adagrad
 )
-from hyperargs import Conf, OptionArg, FloatArg, BoolArg, add_dependency, monitor_on
+from hyperargs import Conf, OptionArg, FloatArg, BoolArg, add_dependency, monitor_on, IntArg
 
 
 class OptimizerConfigBase(Conf):
@@ -32,6 +32,8 @@ class AdamWConfig(OptimizerConfigBase):
     eps = FloatArg(1e-8, min_value=0.0, max_value=1.0)
     weight_decay = FloatArg(0.01, min_value=0.0, max_value=1.0)
     amsgrad = BoolArg(False)
+    # fused CUDA kernel (single multi-tensor launch); ignored on CPU
+    fused = BoolArg(default=False)
 
     def init_optimzier(
         self, 
@@ -43,7 +45,8 @@ class AdamWConfig(OptimizerConfigBase):
             betas=(self.beta1.value(), self.beta2.value()), # type: ignore
             eps=self.eps.value(), # type: ignore
             weight_decay=self.weight_decay.value(), # type: ignore
-            amsgrad=self.amsgrad.value() # type: ignore
+            amsgrad=self.amsgrad.value(), # type: ignore
+            fused=self.fused.value() or None, # type: ignore
         )
 
 

{lmfuser-0.3.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
@@ -63,6 +65,88 @@ def _compile_kwargs(compile_mode: str) -> dict:
     return {} if compile_mode == 'default' else {'mode': compile_mode}
 
 
+class _TqdmLogHandler(logging.Handler):
+    """Route logging through tqdm.write so log lines print above the progress
+    bar instead of tearing through it mid-refresh."""
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            tqdm.write(self.format(record))
+        except Exception:
+            self.handleError(record)
+
+
+class _DevicePrefetcher:
+    """Background fetch + side-stream H2D double buffering over a dataloader.
+
+    next() returns a device-resident batch whose copy has completed (event
+    sync on the compute stream, not the host)."""
+
+    def __init__(self, loader, device, precision: str) -> None:
+        import threading, queue as _q
+        self._loader = loader
+        self._device = device
+        self._precision = precision
+        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()
+
+    def _cast(self, v: Tensor) -> Tensor:
+        if torch.is_floating_point(v):
+            dt = {'fp32': torch.float32, 'fp16': torch.float16, 'bf16': torch.bfloat16}[self._precision]
+            if v.dtype != dt:
+                v = v.to(dt)
+        return v
+
+    def _run(self) -> None:
+        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):
+        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:
+        self._stop = True
+
+
 class Wrapper(nn.Module):
 
     def __init__(self, model: nn.Module, compile_mode: str = 'disable') -> None:
@@ -91,12 +175,17 @@ class Wrapper(nn.Module):
 
 class DDPWraper(nn.Module):
 
-    def __init__(self, model: nn.Module, compile_mode: str = 'disable') -> None:
+    def __init__(self, model: nn.Module, compile_mode: str = 'disable', find_unused: bool = True, bf16_grads: bool = False) -> None:
         super().__init__()
         if get_world_size() > 1:
+            logger.critical(f'DDP options: find_unused={find_unused} bf16_grads={bf16_grads}')
             self.model = DDP(
-                model.to(get_default_device()), find_unused_parameters=True
+                model.to(get_default_device()),
+                find_unused_parameters=find_unused,
             )
+            if bf16_grads:
+                from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
+                self.model.register_comm_hook(None, default_hooks.bf16_compress_hook)
             if compile_mode != 'disable':
                 # compile the DDP module so Dynamo's DDPOptimizer splits the
                 # graph at gradient-bucket boundaries (keeps comm overlap).
@@ -181,6 +270,29 @@ class DDPRunnerConfig(RunerConf):
     # 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'])
+    # gather/log training metrics every N steps. At freq=1 the per-step cost is
+    # substantial for fast steps: a metrics all_gather + a dist_avg collective +
+    # ~6 wandb/logger calls + .item() syncs + a barrier — about half the step
+    # time for a 116M model. Logged values are point samples at the logged step.
+    metric_sync_freq = IntArg(1, min_value=1)
+    # DDP tuning. find_unused_parameters=1 (default, safest) supports tasks
+    # where a subset of parameters gets no grad on some steps (GAN/GRPO mode
+    # switching) at the cost of a per-step graph traversal and weaker comm
+    # overlap. Plain single-player training should set 0.
+    ddp_find_unused = BoolArg(default=True)
+    # compress gradient allreduce to bf16 (halves comm volume; fp32 master
+    # weights are untouched — only the wire format changes)
+    ddp_bf16_grads = BoolArg(default=False)
+    # double-buffered device prefetch: a background thread fetches batch N+1
+    # and copies it to the GPU on a side stream (pinned staging) while the GPU
+    # computes step N. Removes the serialized fetch+H2D segments from the step
+    # wall (measured ~40-120ms/step). Tensor-only batches.
+    device_prefetch = BoolArg(default=False)
+    # diagnostic: log per-step wall time split into fetch / h2d / traincall
+    # segments every metric_sync_freq steps. fetch = waiting on the loader,
+    # h2d = host-to-device copy (0 with device_prefetch), traincall = the
+    # task.train_step call including any GPU-queue wait it absorbs.
+    step_timing = BoolArg(default=False)
     model_precision = OptionArg(options=['fp32', 'fp16', 'bf16'], default='fp32')
     use_amp = BoolArg(default=False)
     amp_precision = OptionArg(options=['fp16', 'bf16'], default='fp16')
@@ -245,9 +357,22 @@ 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),
             batch_size=config.sub_batch_size.value(), # type: ignore
             seed=config.seed.value(), # type: ignore
             shuffle=config.shuffle_dataset.value(), # type: ignore
@@ -365,6 +490,41 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         else:
             raise ValueError(f'dp_type must be either "ddp" or "fsdp2", got "{self.dp_type}" instead.')
 
+    def _collect_data_states(self) -> list | None:
+        """Per-task shard-cursor tables, merged across ranks (rank0 only;
+        other ranks and loaders without cursor support return/contribute
+        None). A COLLECTIVE call when world_size > 1 — every rank must
+        reach it."""
+        per_task = [
+            loader.state_dict()
+            if loader is not None and hasattr(loader, 'state_dict') else None
+            for loader in self.train_data_loaders
+        ]
+        # 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 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)):
+                m: dict = {}
+                for rank_states in gathered:
+                    state = rank_states[task_idx] if rank_states else None
+                    if state:
+                        for src, table in state.items():
+                            m.setdefault(src, {}).update(table)
+                merged.append(m or None)
+            return merged
+        return per_task
+
     def save(self, directory: str | os.PathLike, *args, **kwargs) -> None:
         if isinstance(self.model, (DDPWraper, Wrapper)):
             model = self.model.model.module
@@ -390,25 +550,38 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         else:
             raise ValueError(f'Unknown model type {type(self.model)}')
         
-        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
-
-            optimizer_path = path / 'optimizer.pt'
-            torch.save(optim_states, optimizer_path)
-
-            scheduler_path = path / 'scheduler.pt'
-            torch.save(self.scheduler.state_dict(), scheduler_path)
+        data_states = self._collect_data_states()
 
-            runner_path = path / 'runner.json'
-            with open(runner_path, 'w') as f:
+        if get_global_rank() == 0:
+            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:
+                with open(staging / 'data_state.json', 'w') as f:
+                    json.dump(data_states, f)
+
+            torch.save(optim_states, staging / 'optimizer.pt')
+            torch.save(self.scheduler.state_dict(), staging / 'scheduler.pt')
+
+            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)
@@ -429,7 +602,11 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                 self._model = FSDP2Wrapper(model, compile_mode=compile_mode)
             else:
                 logger.critical(f'wrapping model with DDPWraper')
-                self._model = DDPWraper(model, compile_mode=compile_mode)
+                self._model = DDPWraper(
+                    model, compile_mode=compile_mode,
+                    find_unused=self.config.ddp_find_unused.value(),
+                    bf16_grads=self.config.ddp_bf16_grads.value(),
+                )
         assert self._model is not None
         return self._model
 
@@ -559,11 +736,12 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         epochs = [loader.epoch for loader in self.train_data_loaders if loader is not None]
         return max(epochs) + self.pre_epoch
 
-    def step_log(self, data: dict[str, Any]) -> None:
+    def step_log(self, data: dict[str, Any], console: bool = True) -> None:
         self._wandb
         if get_global_rank() != 0:
             return
-        self.logger.critical(f'step:{self.step}\t{data}')
+        if console:
+            self.logger.critical(f'step:{self.step}\t{data}')
         wandb.log(data, step=self.step)
 
     def _one_train_step(self, **kwargs: Any) -> None:
@@ -577,6 +755,9 @@ class DDPRunner(Runner[DDPRunnerConfig]):
 
         # calculate loss
         running_loss: float = 0.0
+        running_loss_t = None
+        _msf = self.config.metric_sync_freq.value() or 1
+        log_this = (self.step % _msf == 0)
         batch_datas: defaultdict[Hashable, List[float]] = defaultdict(list)
         subbatch_result_lst: list[Batch] = []
         for acc_idx in range(self.config._num_acc_steps):
@@ -604,13 +785,41 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                         self.model.set_requires_gradient_sync(True, recurse=True)
 
                 # compute loss for each sub_batch
+                import time as _time
+                _tf0 = _time.perf_counter()
+                if self.config.device_prefetch.value():
+                    if not hasattr(self, '_prefetchers'):
+                        self._prefetchers = {}
+                    pf = self._prefetchers.get(task_id)
+                    if pf is None:
+                        pf = _DevicePrefetcher(
+                            self.train_data_loaders[task_id],
+                            get_default_device(),
+                            self.config.model_precision.value(),
+                        )
+                        self._prefetchers[task_id] = pf
+                    _dev_batch = pf.next()
+                    _tf1 = _tf2 = _time.perf_counter()
+                else:
+                    _raw_batch = self._next_train_batch(task_id)
+                    _tf1 = _time.perf_counter()
+                    _dev_batch = self._batch_to_device(_raw_batch)
+                    _tf2 = _time.perf_counter()
+                if self.config.step_timing.value():
+                    if not hasattr(self, '_phase_acc'):
+                        self._phase_acc = [0.0, 0.0, 0.0, 0]
+                    self._phase_acc[0] += _tf1 - _tf0
+                    self._phase_acc[1] += _tf2 - _tf1
+                _tf2b = _time.perf_counter()
                 subbatch_result = task.train_step(
                     model=self.model, # type: ignore
-                    batch=self._batch_to_device(self._next_train_batch(task_id)),
-                    step=self.step, 
+                    batch=_dev_batch,
+                    step=self.step,
                     device=get_local_rank(),
                     acc_step=acc_idx,
                 )
+                if self.config.step_timing.value():
+                    self._phase_acc[2] += _time.perf_counter() - _tf2b
                 if isinstance(subbatch_result, torch.Tensor):
                     subbatch_result = {'loss': subbatch_result}
                 if 'loss' not in subbatch_result:
@@ -621,7 +830,7 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                 subbatch_result_lst.append(subbatch_result) # type: ignore
                 assert isinstance(loss, Tensor)
                 loss = loss / self.config._num_acc_steps
-                running_loss += loss.item()
+                running_loss_t = loss.detach() if running_loss_t is None else running_loss_t + loss.detach()
                 for k, v in subbatch_result.items():
                     if k == 'loss':
                         continue
@@ -635,30 +844,66 @@ class DDPRunner(Runner[DDPRunnerConfig]):
                 else:
                     loss.backward()
 
-        running_loss = dist_avg(running_loss)
-        self.step_log({f'{task.__class__.__name__}/train/loss': running_loss})
-        self.step_log({'train/epoch': self.epoch})
-        for k, v in batch_datas.items():
-            try:
-                avg = sum(v) / len(v)
-            except:
-                continue
-            self.step_log({f'{task.__class__.__name__}/train/{k}': avg})
-        self._pbar_train.set_description(
-            f'train loss: {running_loss:.3g}', refresh=True
-        )
+        if log_this and getattr(self, '_phase_acc', None):
+            n = max(self._phase_acc[3], 1)
+            fetch_ms = self._phase_acc[0] / n * 1000
+            h2d_ms = self._phase_acc[1] / n * 1000
+            traincall_ms = self._phase_acc[2] / n * 1000
+            logger.critical(
+                f'[timing] fetch={fetch_ms:.1f}ms h2d={h2d_ms:.1f}ms '
+                f'traincall={traincall_ms:.1f}ms per-step (over {n} steps)')
+            # wandb: rank0's own segments plus the all-rank max of each — the
+            # max is the diagnostic one (in DDP the slowest rank paces the
+            # whole step, e.g. a data-starved rank shows up as fetch_ms_max)
+            t = torch.tensor([fetch_ms, h2d_ms, traincall_ms], device=get_default_device())
+            if torch.distributed.is_initialized():
+                torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX)
+            mx = t.tolist()
+            self.step_log({
+                'timing/fetch_ms': fetch_ms,
+                'timing/h2d_ms': h2d_ms,
+                'timing/traincall_ms': traincall_ms,
+                'timing/fetch_ms_max': mx[0],
+                'timing/h2d_ms_max': mx[1],
+                'timing/traincall_ms_max': mx[2],
+            })
+            self._phase_acc = [0.0, 0.0, 0.0, 0]
+        if getattr(self, '_phase_acc', None) is not None:
+            self._phase_acc[3] += 1
+        if log_this:
+            running_loss = running_loss_t.item() if running_loss_t is not None else 0.0
+            running_loss = dist_avg(running_loss)
+            self.step_log({f'{task.__class__.__name__}/train/loss': running_loss})
+            self.step_log({'train/epoch': self.epoch})
+        if log_this:
+            for k, v in batch_datas.items():
+                try:
+                    avg = sum(v) / len(v)
+                except:
+                    continue
+                self.step_log({f'{task.__class__.__name__}/train/{k}': avg})
+            # the sampled step is part of the label: loss refreshes every
+            # metric_sync_freq steps while the bar ticks every step, and an
+            # unlabeled value reads as a frozen metric
+            self._pbar_train.set_postfix_str(
+                f'loss={running_loss:.3g}@{self.step}', refresh=True
+            )
 
         grad_norm_clip_val = self.config.grad_norm_clip.value()
         if grad_norm_clip_val is not None:
             if self.scaler is not None:
                 self.scaler.unscale_(self.optimizer)
-            norm = clip_grad_norm_(
+            norm_t = clip_grad_norm_(
                 parameters=self.model.parameters(),  # type: ignore
                 max_norm=grad_norm_clip_val
-            ).item()
-            norm = dist_avg(norm)
-            self.step_log({f'{task.__class__.__name__}/train/grad_norm': norm})
+            )
+            if log_this:
+                norm = dist_avg(norm_t.item())
+                self.step_log({f'{task.__class__.__name__}/train/grad_norm': norm})
 
+        # pure host-side work — no .item() sync, no collective — so it runs and
+        # reports every step (wandb enqueue is async; console line only on
+        # log steps to keep the terminal readable)
         num_hot_params = 0
         num_freeze_params = 0
         for param in self.model.parameters(): # type: ignore
@@ -670,13 +915,12 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         num_total_params = num_hot_params + num_freeze_params
         if num_total_params == 0:
             raise RuntimeError('The model contains no parameters.')
-
         self.step_log({
             f'{task.__class__.__name__}/train/num_hot_params': num_hot_params,
             f'{task.__class__.__name__}/train/num_freeze_params': num_freeze_params,
             f'{task.__class__.__name__}/train/num_total_params': num_total_params,
             f'{task.__class__.__name__}/train/hot_ratio': num_hot_params / num_total_params,
-        })
+        }, console=log_this)
 
         if self.scaler is not None:
             self.scaler.step(self.optimizer)
@@ -688,21 +932,22 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             current_lr = current_lr[0]
         self.step_log(
             {f'{task.__class__.__name__}/train/learning_rate': current_lr}
-        )
+        ) if log_this else None
         self.scheduler.step()
 
-        other_step_results = defaultdict(list)
-        for r in subbatch_result_lst:
-            for k, v in r.items():
-                if isinstance(v, (float, int)):
-                    other_step_results[k].append(float(v))
-        all_other_step_results = batch_all_gather(other_step_results)
-        for k, v in all_other_step_results.items():
-            try:
-                avg = sum(v) / len(v)
-            except:
-                continue
-            self.step_log({f'{task.__class__.__name__}/train/{k}': avg})
+        if log_this:
+            other_step_results = defaultdict(list)
+            for r in subbatch_result_lst:
+                for k, v in r.items():
+                    if isinstance(v, (float, int)):
+                        other_step_results[k].append(float(v))
+            all_other_step_results = batch_all_gather(other_step_results)
+            for k, v in all_other_step_results.items():
+                try:
+                    avg = sum(v) / len(v)
+                except:
+                    continue
+                self.step_log({f'{task.__class__.__name__}/train/{k}': avg})
 
         save_step_freq = self.config.save_step_freq.value()
         assert save_step_freq is not None
@@ -712,7 +957,7 @@ class DDPRunner(Runner[DDPRunnerConfig]):
             assert ckpt_path is not None
             self.save(ckpt_path)
             self._pbar_train.set_description('model saved!', True)
-        torch.distributed.barrier() if get_world_size() > 1 else ...
+            torch.distributed.barrier() if get_world_size() > 1 else ...
 
         eval_step_freq = self.config.eval_step_freq.value()
         assert eval_step_freq is not None
@@ -726,6 +971,16 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         self.eval() # evaluate first
         self._prepare_train()
 
+        # route log lines through tqdm.write for the training run so they
+        # print cleanly above the bar (plain StreamHandlers tear through it)
+        root_logger = logging.getLogger()
+        if not any(isinstance(h, _TqdmLogHandler) for h in root_logger.handlers):
+            tqdm_handler = _TqdmLogHandler()
+            tqdm_handler.setFormatter(logging.Formatter(logging.BASIC_FORMAT))
+            for h in [x for x in root_logger.handlers if type(x) is logging.StreamHandler]:
+                root_logger.removeHandler(h)
+            root_logger.addHandler(tqdm_handler)
+
         self._last_epoch = self.epoch
         stop_metric = self.config.stop_by.value()
         assert stop_metric in ('step', 'epoch')
@@ -753,6 +1008,8 @@ class DDPRunner(Runner[DDPRunnerConfig]):
         while not self._should_stop():
             self._one_train_step()
             if stop_metric == 'step':
+                # per-step update is fine: tqdm only renders every mininterval
+                # (0.1s); the call itself is just a counter bump
                 self._pbar_train.update(1)
             elif stop_metric == 'epoch':
                 current_epoch = self.epoch
@@ -941,16 +1198,29 @@ 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')
 
+        data_state_path = Path(directory) / 'data_state.json'
+        if data_state_path.exists():
+            with open(data_state_path, 'r') as f:
+                self._resume_data_states = json.load(f)
+            n = sum(
+                len(tab)
+                for task_state in self._resume_data_states if task_state
+                for tab in task_state.values()
+            )
+            logger.info(f'loaded data stream state ({n} shard cursors)')
+        else:
+            logger.info('no data_state.json in checkpoint — data stream restarts from scratch')
+
         runner_path = Path(directory) / 'runner.json'
         if runner_path.exists():
             with open(runner_path, 'r') as f:

{lmfuser-0.3.0 → lmfuser-0.4.1}/src/lmfuser/task.py

@@ -110,6 +110,7 @@ class TaskBase(Conf, SubclassTracer):
         num_batch_workers: int = 4,
         batch_queue_depth: int = 4,
         batch_slot_mb: int = 128,
+        resume_state: dict | None = None,
     ) -> None | DataLoader | PyTorchDataLoader | BatchDataLoader | EmptyDataLoader:
         if self.num_train_data_path.value() == 0:
             return None
@@ -143,6 +144,9 @@ class TaskBase(Conf, SubclassTracer):
                 num_ranks=world_size,
                 rank_idx=rank,
                 worker_timeout=worker_timeout,
+                # only forward when present: lmfuser-data < 0.3.0 has no
+                # resume_state parameter
+                **({'resume_state': resume_state} if resume_state else {}),
             )
         elif dataloader_type == 'sharded':
             self._train_dataloader = DataLoader(
@@ -430,6 +434,7 @@ class Tasks(Conf):
         num_batch_workers: int = 4,
         batch_queue_depth: int = 4,
         batch_slot_mb: int = 128,
+        resume_states: list[dict | None] | None = None,
     ) -> list[DataLoader | None | PyTorchDataLoader | BatchDataLoader | EmptyDataLoader]:
         return [
             task.conf._get_train_dataloader(
@@ -447,8 +452,9 @@ class Tasks(Conf):
                 num_batch_workers=num_batch_workers,
                 batch_queue_depth=batch_queue_depth,
                 batch_slot_mb=batch_slot_mb,
+                resume_state=(resume_states[i] if resume_states else None),
             )
-            for task in self.tasks
+            for i, task in enumerate(self.tasks)
         ]
 
     def get_eval_dataloaders(

{lmfuser-0.3.0 → lmfuser-0.4.1}/src/lmfuser/utils.py

@@ -207,16 +207,11 @@ def batch_all_gather(batch: dict[str, Any]) -> dict[str, Any]:
     Returns:
         Dict[Dict[str, Union[Tensor, List[Any]]]]: 汇总后的batch
     """
-    logger.info(f'Begin all gather on rank {get_global_rank()}')
     if not dist.is_initialized():
         return batch
 
     gathered = {}
     for k, v_list in batch.items():
-        logger.info(f'Begin gather key {k} on rank {get_global_rank()}')
-        if isinstance(v_list, list):
-            logger.info(f'Key {k} with length {len(v_list)} on rank {get_global_rank()}')
-        logger.info(f'after barrier on rank {get_global_rank()}')
         if isinstance(v_list, Tensor):
             gathered[k] = tensor_all_gather(v_list)
         else:

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')