TorchRunner

danling.runners.TorchRunner

Bases: Fp8Mixin, BaseRunner

PyTorch-native runner for training, evaluation, and inference.

Use this runner for single-model PyTorch training with optional DDP, autocast/FP8, torch.compile, stateful dataloaders, metric logging, and file or torch.distributed.checkpoint persistence.

Users must provide self.model before construction completes. Most training tasks also provide self.criterion, and either self.optimizer or config.optim. Datasets may be supplied through self.datasets and will be materialized into StatefulDataLoader instances during __post_init__.

The default batch contract is intentionally simple: mappings use input/target, sequences use index 0/1, and any other value is treated as model input with no target. Override train_step, evaluate_step, or infer_step when a task needs a different contract.

Attributes:

Name	Type	Description
model	Module	Local model module after materialization (possibly DDP-wrapped).
ema	Module | None	Optional EMA model, or an eval-only shadow model without update().
criterion	Callable | None	Loss callable used by default train/evaluate steps.
optimizer	Optimizer | None	Optimizer used by the runner or backend engine.
scheduler	Any | None	Optional LR scheduler.
optimizer_container	OptimizerContainer | None	Helper that owns optimizer step, clipping, non-finite checks, and step-scheduler dispatch.
compiler	Compiler	torch.compile policy object.
scheduler_interval	str	Effective scheduler interval ("step" or epoch/metric-style interval).
scheduler_monitor	str | None	Optional metric path used for metric schedulers.

Source code in danling/runners/torch_runner.py

class TorchRunner(Fp8Mixin, BaseRunner):
    r"""
    PyTorch-native runner for training, evaluation, and inference.

    Use this runner for single-model PyTorch training with optional DDP,
    autocast/FP8, `torch.compile`, stateful dataloaders, metric logging, and
    file or torch.distributed.checkpoint persistence.

    Users must provide `self.model` before construction completes. Most
    training tasks also provide `self.criterion`, and either `self.optimizer`
    or `config.optim`. Datasets may be supplied through `self.datasets` and
    will be materialized into `StatefulDataLoader` instances during
    `__post_init__`.

    The default batch contract is intentionally simple:
    mappings use `input`/`target`, sequences use index 0/1, and any other value
    is treated as model input with no target. Override `train_step`,
    `evaluate_step`, or `infer_step` when a task needs a different contract.

    Attributes:
        model: Local model module after materialization (possibly DDP-wrapped).
        ema: Optional EMA model, or an eval-only shadow model without `update()`.
        criterion: Loss callable used by default train/evaluate steps.
        optimizer: Optimizer used by the runner or backend engine.
        scheduler: Optional LR scheduler.
        optimizer_container: Helper that owns optimizer step, clipping,
            non-finite checks, and step-scheduler dispatch.
        compiler: `torch.compile` policy object.
        scheduler_interval: Effective scheduler interval (`"step"` or
            epoch/metric-style interval).
        scheduler_monitor: Optional metric path used for metric schedulers.
    """

    model: nn.Module
    ema: nn.Module | None = None
    criterion: Callable | None = None
    optimizer: optim.Optimizer | None = None
    scheduler: Any | None = None
    optimizer_container: OptimizerContainer | None = None
    grad_scaler: torch.amp.GradScaler | None = None
    compiler: Compiler
    scheduler_interval: str = "step"
    scheduler_monitor: str | None = None
    _train_pg_timeout_reduced: bool = False
    _profiler_context: Any | None = None
    diagnostics: RuntimeDiagnostics | None = None
    _profiler: Any | None = None
    _profiler_trace_dir: str | None = None
    _profiler_operator_table_rows: int = 100
    _pending_loss_normalizer: float | None = None
    _pending_loss_reduction: str | None = None
    _accumulation_divisor_local: float = 0.0
    _accumulation_divisor_total: float | None = None
    _accumulation_mode: str | None = None
    _train_window_will_flush: bool = False
    _optimizer_parameter_cache: OptimizerParameterCache | None = None
    _supports_torchft_runtime: bool = True

    _VALID_CHECKPOINT_BACKENDS = frozenset({"file", "dcp"})

    @classmethod
    def _validate_checkpoint_backend(cls, backend: str) -> str:
        """Normalize and validate a resolved checkpoint backend value."""
        backend = str(backend).strip().lower()
        if backend not in cls._VALID_CHECKPOINT_BACKENDS:
            raise ValueError(f"invalid checkpoint backend: {backend!r}. Expected one of: 'auto', 'file', 'dcp'.")
        return backend

    def __init__(self, config) -> None:
        if not isinstance(config, RunnerConfig):
            config = RunnerConfig(config)
        if isinstance(config.get("device"), torch.device):
            config["device"] = str(config["device"])
        config.stack = normalize_stack_name(config.get("stack", "ddp"))
        checkpoint_backend = str(config.get("ckpt.backend")).strip().lower()
        if checkpoint_backend == "auto":
            checkpoint_backend = "dcp" if self.world_size > 1 or config.get("ft.enabled", False) else "file"
        config["ckpt"]["backend"] = self._validate_checkpoint_backend(checkpoint_backend)
        super().__init__(config)

    def __post_init__(self):
        self._construct_model()
        self._pending_loss_normalizer = None
        self._accumulation_divisor_local = 0.0
        self._accumulation_mode = None
        self._train_window_will_flush = False
        self._optimizer_parameter_cache = None
        if self.model is None:
            raise ValueError("cannot initialize TorchRunner: model is not initialized")
        if self.datasets:
            self.build_dataloaders()
        if self.fault_tolerance is not None and self.fault_tolerance.enabled and not self._supports_torchft_runtime:
            raise NotImplementedError(
                "TorchFT integration is currently supported by TorchRunner/DDP and ParallelRunner FSDP only"
            )
        self.compiler = Compiler(self.config.compile)
        self.setup_fp8()
        self.setup_grad_scaler()
        self.materialize_model()
        self.build_optimizer()
        self.build_scheduler()
        self._finalize_runtime_components()
        sched_cfg = self._get_scheduler_config()
        interval = sched_cfg.get("interval") if sched_cfg is not None else None
        monitor = sched_cfg.get("monitor") if sched_cfg is not None else None
        self.scheduler_interval = normalize_scheduler_interval(interval, self.scheduler)
        self.scheduler_monitor = None if monitor is None else str(monitor)
        self._bind_optimizer_container()
        self.auto_restore()
        self.diagnostics = RuntimeDiagnostics(self)
        self._init_profiler()
        super().__post_init__()

    def _finalize_runtime_components(self) -> None:
        """Hook for backend-specific engine/materialization after optimizer and scheduler build."""

    def init_distributed(self) -> None:
        r"""
        Initialize the distributed environment.

        The default implementation initializes the default torch.distributed
        process group from `WORLD_SIZE`/`RANK`/`LOCAL_RANK` environment
        variables when `WORLD_SIZE > 1` or TorchFT is enabled, sets the selected CUDA device,
        broadcasts `self.timestamp` from rank 0, and seeds
        `elastic_state.restart_count` from `TORCHELASTIC_RESTART_COUNT`.

        **Called when:** once during `BaseRunner.__init__`, before
        `init_checkpoint_manager`, `init_fault_tolerance`, and
        `init_garbage_collection`. The runner is partially constructed at
        this point — `self.config`, `self.workspace`, `self.timestamp`, the
        dataloader container, and the default `FileCheckpointManager` are
        bound, but the model is not materialized and optimizers/dataloaders
        are not built.

        **Precondition:** environment variables `WORLD_SIZE`, `RANK`,
        `LOCAL_RANK` are set when running distributed. The default
        torch.distributed process group is **not** already initialized when
        `WORLD_SIZE > 1` or TorchFT is enabled — the runner owns process-group lifecycle.

        Raises:
            RuntimeError: the default process group is already initialized
                when `WORLD_SIZE > 1` or TorchFT is enabled.
            ValueError: `dist.init_timeout_seconds` is non-positive.

        **Side effects:** when `WORLD_SIZE > 1` or TorchFT is enabled, calls
        `dist.init_process_group(...)`, sets the active CUDA device when
        the runtime device is CUDA, and broadcasts `self.timestamp` from rank 0.
        Reads `TORCHELASTIC_RESTART_COUNT` into `elastic_state.restart_count`.

        !!! danger "Do not"
            - Initialize a process group via `dist.init_process_group(...)`
              outside the runner; the runner owns its lifecycle.
            - Build the model or dataloaders here; those happen in
              `__post_init__`.
            - Bind the checkpoint manager here; `init_checkpoint_manager`
              runs next.

        **Backend notes:**

        - `ParallelRunner` extends this hook: after calling `super()`, it
          builds the parallel topology (`build_topology`) and initializes
          per-axis process groups via `init_device_mesh`.
        - `DeepSpeedRunner` inherits the default; DeepSpeed reuses the
          default process group initialized here.
        """

        backend = self.config.get("dist.backend") or os.getenv("BACKEND")
        init_method = self.config.get("dist.init_method") or os.getenv("INIT_METHOD")
        init_timeout = self._dist_timeout("dist.init_timeout_seconds")
        world_size = int(os.getenv("WORLD_SIZE", "1"))
        rank = int(os.getenv("RANK", "0"))
        runtime_device = self.device
        if backend is None:
            backend = dist.Backend.default_device_backend_map.get(runtime_device.type)
        use_cuda_runtime = torch.cuda.is_available() and runtime_device.type == "cuda"
        runtime_device_index = runtime_device.index if runtime_device.index is not None else self.local_rank
        dist_ready = dist.is_available() and dist.is_initialized()
        requires_process_group = world_size > 1 or bool(self.config.get("ft.enabled", False))
        if requires_process_group and dist_ready:
            raise RuntimeError(
                "default process group is already initialized; Runner requires owning process-group lifecycle"
            )
        if requires_process_group:
            if use_cuda_runtime:
                torch.cuda.set_device(runtime_device_index)
            init_kwargs: dict[str, Any] = {
                "backend": backend,
                "init_method": init_method,
                "world_size": world_size,
                "rank": rank,
            }
            if init_timeout is not None:
                init_kwargs["timeout"] = init_timeout
            dist.init_process_group(**init_kwargs)
            dist_ready = bool(dist.is_available() and dist.is_initialized())

        if dist_ready and use_cuda_runtime:
            torch.cuda.set_device(runtime_device_index)

        if dist_ready and self.world_size > 1:
            object_list = [self.timestamp]
            dist.broadcast_object_list(object_list)
            self.timestamp = str(object_list[0])

        restart_count = os.getenv("TORCHELASTIC_RESTART_COUNT")
        if restart_count is not None:
            self.elastic_state.restart_count = int(restart_count)

        self._train_pg_timeout_reduced = False

    def init_checkpoint_manager(self) -> None:
        """
        Bind the checkpoint manager corresponding to `config.ckpt.backend`.

        The default dispatches by backend: when the backend is `"dcp"`, it
        binds a `TorchDistributedCheckpointManager` (or
        `TorchFTCheckpointManager` when dataloader checkpoints are
        enabled). For `"file"` it leaves the `FileCheckpointManager` already
        bound by `BaseRunner.__init__` in place.

        **Called when:** once during `BaseRunner.__init__`, after
        `init_distributed` and before `init_fault_tolerance`. The default
        `FileCheckpointManager` is already bound at this point — overrides
        should swap it via `set_checkpoint_manager(...)`, not by direct
        attribute assignment.

        **Precondition:** `config.ckpt.backend` is normalized to one
        of `{"file", "dcp"}` (TorchRunner does this in `__init__`). When
        the backend is `"dcp"`, the default process group is initialized
        for distributed runs.

        **Side effects:** swaps `self.checkpoint_manager` via
        `set_checkpoint_manager(...)` when the backend differs from
        `"file"`. The prior manager is closed with a zero timeout.

        !!! danger "Do not"
            - Set `self.checkpoint_manager` directly; use
              `set_checkpoint_manager` so the prior manager is closed
              cleanly.
            - Initialize fault tolerance here; `init_fault_tolerance` runs
              next.
            - Bind the model or dataloaders here.

        **Backend notes:**

        - `DeepSpeedRunner` coerces `config.ckpt.backend` to `"file"`
          in `__init__`, so this hook is a no-op for that backend.
        - `ParallelRunner` coerces the backend to `"dcp"`, so this hook
          always binds `TorchDistributedCheckpointManager` or
          `TorchFTCheckpointManager`.
        """
        checkpoint_backend = str(self.config.get("ckpt.backend")).lower()
        if checkpoint_backend == "dcp":
            ft_checkpoint_enabled = bool(
                self.config.get("ft.enabled", False) or self.config.get("ckpt.dataloader_checkpoint.enabled", False)
            )
            manager_cls = TorchFTCheckpointManager if ft_checkpoint_enabled else TorchDistributedCheckpointManager
            self.set_checkpoint_manager(manager_cls(self))
            return
        # Backend is normalized to {"file", "dcp"} in `__init__`; "file" is the
        # remaining case and reuses the default `FileCheckpointManager` that
        # `BaseRunner.__init__` already bound.

    def _dist_timeout(self, key: str) -> timedelta | None:
        value = self.config.get(key)
        if value is None:
            return None
        seconds = int(value)
        if seconds <= 0:
            raise ValueError(f"{key} must be a positive integer, got {seconds}")
        return timedelta(seconds=seconds)

    def _timeout_process_groups(self) -> tuple[Any | None, ...]:
        groups: list[Any | None] = [None]
        if self.fault_tolerance is not None and self.fault_tolerance.replicate_process_group is not None:
            groups.append(self.fault_tolerance.replicate_process_group)
        return tuple(groups)

    def _set_process_group_timeout(self, timeout: timedelta) -> None:
        if not (dist.is_available() and dist.is_initialized()):
            return
        set_pg_timeout = getattr(dist_c10d, "_set_pg_timeout", None)
        if not callable(set_pg_timeout):
            warn(
                "torch.distributed does not expose process-group timeout mutation; "
                "skipping dist.train_timeout_seconds update",
                RuntimeWarning,
                stacklevel=2,
            )
            return

        for group in self._timeout_process_groups():
            backend = str(dist.get_backend() if group is None else dist.get_backend(group)).lower()
            if backend != "nccl":
                continue

            barrier_kwargs = {} if group is None else {"group": group}
            if torch.cuda.is_available():
                dist.barrier(device_ids=[torch.cuda.current_device()], **barrier_kwargs)
                torch.cuda.synchronize()
            else:
                dist.barrier(**barrier_kwargs)

            try:
                set_pg_timeout(timeout, group)
            except TypeError:
                if group is not None:
                    warn(
                        "torch.distributed does not support subgroup timeout mutation; "
                        "skipping dist.train_timeout_seconds update for a non-default process group",
                        RuntimeWarning,
                        stacklevel=2,
                    )
                    continue
                set_pg_timeout(timeout)
            except Exception as exc:
                group_name = "default" if group is None else "subgroup"
                warn(f"failed to update {group_name} process-group timeout: {exc}", RuntimeWarning, stacklevel=2)

    def _maybe_reduce_train_process_group_timeout(self) -> None:
        if self._train_pg_timeout_reduced:
            return
        if self.train_state.global_step != 1:
            return
        timeout = self._dist_timeout("dist.train_timeout_seconds")
        if timeout is None:
            return
        self._set_process_group_timeout(timeout)
        self._train_pg_timeout_reduced = True

    def destroy_process_group(self) -> None:
        if not (dist.is_available() and dist.is_initialized()):
            return
        try:
            dist.destroy_process_group()
        except Exception as exc:
            warn(f"failed to destroy default process group: {exc}", RuntimeWarning, stacklevel=2)

    def _profiler_activities(self, configured: object | None) -> list[torch.profiler.ProfilerActivity]:
        if configured is None:
            activities = [torch.profiler.ProfilerActivity.CPU]
            if torch.cuda.is_available() and self.device.type == "cuda":
                activities.append(torch.profiler.ProfilerActivity.CUDA)
            return activities

        if isinstance(configured, str):
            configured = (configured,)
        if not isinstance(configured, Sequence):
            raise ValueError("profiling.activities must be a string or sequence of strings")

        aliases = {"cpu": "CPU", "cuda": "CUDA", "gpu": "CUDA", "xpu": "XPU"}
        activities = []
        for activity in configured:
            if isinstance(activity, torch.profiler.ProfilerActivity):
                activities.append(activity)
                continue
            name = aliases.get(str(activity).strip().lower(), str(activity).strip().upper())
            if not hasattr(torch.profiler.ProfilerActivity, name):
                raise ValueError(f"unsupported profiling activity: {activity!r}")
            activities.append(getattr(torch.profiler.ProfilerActivity, name))
        return activities

    def _init_profiler(self) -> None:
        profiling = self.config.get("profiling")
        if not isinstance(profiling, Mapping) or not bool(profiling.get("enabled", False)):
            return

        wait = int(profiling.get("wait", 1))
        warmup = int(profiling.get("warmup", 1))
        active = int(profiling.get("active", 3))
        repeat = profiling.get("repeat")
        if wait < 0:
            raise ValueError(f"profiling.wait must be a non-negative integer, got {wait}")
        if warmup < 0:
            raise ValueError(f"profiling.warmup must be a non-negative integer, got {warmup}")
        if active <= 0:
            raise ValueError(f"profiling.active must be a positive integer, got {active}")
        if repeat is not None:
            repeat = int(repeat)
            if repeat <= 0:
                raise ValueError(f"profiling.repeat must be a positive integer, got {repeat}")

        activities = self._profiler_activities(profiling.get("activities"))

        schedule_kwargs: dict[str, Any] = {"wait": wait, "warmup": warmup, "active": active}
        if repeat is not None:
            schedule_kwargs["repeat"] = repeat

        trace_dir = os.fsdecode(str(profiling.get("trace_dir", "profiles")))
        if not os.path.isabs(trace_dir):
            trace_dir = os.path.join(self.workspace.dir, trace_dir)
        trace_dir = os.path.join(trace_dir, self.timestamp, f"rank-{self.rank:05d}")
        os.makedirs(trace_dir, exist_ok=True)
        self._profiler_trace_dir = trace_dir
        self._profiler_operator_table_rows = int(profiling.get("operator_table_rows", 100))
        profile_kwargs: dict[str, Any] = {
            "activities": activities,
            "schedule": torch.profiler.schedule(**schedule_kwargs),
            "on_trace_ready": torch.profiler.tensorboard_trace_handler(trace_dir),
            "record_shapes": bool(profiling.get("record_shapes", False)),
            "profile_memory": bool(profiling.get("profile_memory", False)),
            "with_stack": bool(profiling.get("with_stack", False)),
            "with_flops": bool(profiling.get("with_flops", False)),
            "with_modules": bool(profiling.get("with_modules", False)),
            "acc_events": bool(profiling.get("acc_events", False)),
        }
        if profiling.get("use_cuda") is not None:
            profile_kwargs["use_cuda"] = bool(profiling.get("use_cuda"))
        post_processing_timeout_seconds = profiling.get("post_processing_timeout_seconds")
        if post_processing_timeout_seconds is not None:
            profile_kwargs["post_processing_timeout_s"] = float(post_processing_timeout_seconds)

        profiler_context = torch.profiler.profile(
            **profile_kwargs,
        )
        profiler = profiler_context.__enter__()
        if hasattr(profiler, "step_num"):
            profiler.step_num = self.train_state.global_step
        self._profiler_context = profiler_context
        self._profiler = profiler

    def _step_profiler(self) -> None:
        if self.diagnostics is not None:
            self.diagnostics.step()
        if self._profiler is None:
            return
        self._profiler.step()

    def _close_profiler(self) -> None:
        profiler_context = self._profiler_context
        profiler = self._profiler
        trace_dir = self._profiler_trace_dir
        operator_table_rows = self._profiler_operator_table_rows
        self._profiler_context = None
        self._profiler = None
        self._profiler_trace_dir = None
        if profiler_context is None:
            return
        profiler_context.__exit__(None, None, None)
        operator_table_path = None
        if profiler is not None and trace_dir is not None and operator_table_rows > 0:
            operator_table_path = os.path.join(trace_dir, "operator_table.txt")
            try:
                try:
                    table = profiler.key_averages().table(
                        sort_by="self_cuda_time_total",
                        row_limit=operator_table_rows,
                    )
                except Exception:
                    table = profiler.key_averages().table(
                        sort_by="self_cpu_time_total",
                        row_limit=operator_table_rows,
                    )
                with open(operator_table_path, "w", encoding="utf-8") as fp:
                    fp.write(table)
                    fp.write("\n")
            except Exception as exc:
                operator_table_path = None
                if "Profiler must be initialized" not in str(exc):
                    warn(f"failed to write profiler operator table: {exc}", RuntimeWarning, stacklevel=2)
        if trace_dir is not None and (not self.distributed or self.is_main_process):
            fields = {"trace_dir": trace_dir, "operator_table": operator_table_path}
            print(self._format_summary("profiler artifacts", fields))

    @on_main_process
    def init_tensorboard(self, *args, **kwargs) -> None:
        r"""
        Set up TensorBoard SummaryWriter.
        """

        try:
            from torch.utils.tensorboard.writer import SummaryWriter  # pylint: disable=C0415
        except ImportError as error:
            raise ImportError(
                "TensorBoard logging requires the optional TensorBoard dependencies; "
                "install them with `pip install 'danling[tensorboard]'`."
            ) from error

        tensorboard_config = self.config.tensorboard
        for key in ("log_dir", "comment", "purge_step", "max_queue", "flush_secs", "filename_suffix"):
            if key not in kwargs and tensorboard_config.get(key) is not None:
                kwargs[key] = tensorboard_config[key]
        if "log_dir" not in kwargs:
            kwargs["log_dir"] = os.path.join(self.workspace.dir, "tensorboard", self.timestamp)

        self.writer = SummaryWriter(*args, **kwargs)
        self.writer.add_scalar = catch(OSError, verbose=False)(self.writer.add_scalar)  # type: ignore[method-assign]

    def set_seed(self, seed: int | None = None, bias: int | bool | None = None) -> int:
        r"""
        Set up random seed.

        Args:
            seed: Random seed to set.
                Defaults to `self.config.seed` (`config.seed`).

            bias: Make the seed different for each processes.
                This is used to ensure the data augmentation are applied differently on every processes.
                Defaults to `self.rank`.
                Set to `False` to disable this feature.
        Returns:
            Random seed set.
        """

        base_seed = seed if seed is not None else self.config.seed  # type: ignore[assignment]
        if base_seed is None:
            base_seed = random.randint(0, 2**32 - 1)
            if self.distributed and dist.is_initialized():
                object_list = [base_seed]
                dist.broadcast_object_list(object_list)
                base_seed = object_list[0]
        base_seed = int(base_seed)
        # Keep `config.seed` as the global/base seed (before per-rank bias).
        self.config.seed = base_seed

        process_seed = base_seed
        if bias is None:
            if self.fault_tolerance is not None:
                _, bias = self.fault_tolerance.data_parallel_info(self.world_size, self.rank)
            else:
                bias = self.rank
        if bias:
            process_seed += int(bias)

        torch.manual_seed(process_seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(process_seed)
        if np_random is not None:
            np_random.seed(process_seed)
        random.seed(process_seed)
        self.rng_state.python = random.getstate()
        self.rng_state.numpy = np_random.get_state() if np_random is not None else None
        self.rng_state.torch_cpu = torch.get_rng_state()
        if torch.cuda.is_available():
            self.rng_state.torch_cuda = torch.cuda.get_rng_state_all()
        else:
            self.rng_state.torch_cuda = None
        return process_seed

    def set_deterministic(self) -> None:
        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
        cudnn.benchmark = False
        cudnn.deterministic = True
        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = False
            torch.backends.cudnn.allow_tf32 = False
        torch.use_deterministic_algorithms(True, warn_only=True)

    @staticmethod
    def _normalized_precision_name(precision: str | None) -> str:
        if precision is None:
            return ""
        return str(precision).strip().lower().replace("-", "_")

    def runner_owns_grad_scaling(self) -> bool:
        """Whether this runner should own AMP gradient scaling."""

        return True

    def setup_grad_scaler(self) -> None:
        """Bind the AMP grad scaler for runner-owned fp16 training."""

        self.grad_scaler = None
        precision = self._normalized_precision_name(self.precision)
        if precision not in {"fp16", "float16", "half"}:
            return
        if not self.runner_owns_grad_scaling():
            return
        if self.fp8_enabled:
            raise ValueError("precision='fp16' cannot be combined with FP8 autocast")
        if self.device.type != "cuda":
            raise ValueError("runner-owned fp16 precision requires a CUDA device; use bf16 or a backend-owned scaler")
        self.grad_scaler = torch.amp.GradScaler(device=self.device.type)

    def update_ema(self) -> None:
        """Update EMA weights after a successful optimizer step."""

        if self.ema is None:
            return
        if self.model is None:
            raise ValueError("cannot update EMA: model is not initialized")
        update = getattr(self.ema, "update", None)
        if update is None:
            return
        if not callable(update):
            raise TypeError(
                "EMA module has an `update` attribute but it is not callable; "
                "implement `ema.update(model)`, remove the attribute, or override `update_ema()`."
            )
        update(self.unwrap(self.model))

    def validate_ema_update_contract(self) -> None:
        update = None if self.ema is None else getattr(self.ema, "update", None)
        if update is not None and not callable(update):
            raise TypeError(
                "EMA module has an `update` attribute but it is not callable; "
                "implement `ema.update(model)`, remove the attribute, or override `update_ema()`."
            )

    def _record_grad_norm(self, grad_norm: float | None) -> None:
        if grad_norm is not None:
            self.meters.grad_norm.update(grad_norm)

    def build_model(self) -> nn.Module:
        r"""Build a model under the runner's meta-device context when opted in.

        Override this instead of binding ``self.model`` in ``__init__`` when
        ``model_init.meta=True``. Model-specific configuration stays on the
        runner's existing Chanfig configuration.
        """
        raise NotImplementedError("model_init.meta=True requires a model or an overridden build_model()")

    def _construct_model(self) -> None:
        if self.model is not None or not self.config.model_init.meta:
            return
        with torch.device("meta"):
            model = self.build_model()
        if not isinstance(model, nn.Module):
            raise TypeError("build_model() must return an nn.Module")
        self.model = model

    def initialize_model(self, model: nn.Module) -> None:
        r"""Initialize or load newly allocated model storage before optimizer creation.

        This runs only for a meta model. Override it to initialize parameters
        and persistent/non-persistent buffers, or to load model-specific state.
        The default calls the model's explicit ``init_weights()`` protocol;
        it does not guess initialization order or reset pretrained models.
        Ordinary checkpoint/pretrained restoration still runs in ``auto_restore``.
        """
        initialize = getattr(model, "init_weights", None)
        if not callable(initialize):
            raise NotImplementedError("meta models require init_weights() or an overridden initialize_model()")
        initialize()

    def model_init_device(self, model: nn.Module) -> torch.device:
        r"""Select the parameter-allocation device for a meta model."""
        return self.device

    def materialize_model_values(self, model: nn.Module) -> nn.Module:
        r"""Allocate a meta model after transforms, retaining tied tensor identities.

        Already initialized models keep the ordinary ``to(device)`` path.
        Mixed initialized/meta tensors require a model-specific override since
        ``to_empty`` would otherwise discard initialized values.
        """
        tensors = tuple(model.parameters()) + tuple(model.buffers())
        if not any(tensor.is_meta for tensor in tensors):
            return model.to(self.device)
        if not self.config.model_init.meta:
            raise ValueError("meta model materialization requires model_init.meta=True")
        if any(not tensor.is_meta for tensor in tensors):
            raise ValueError("mixed initialized/meta model tensors require overriding materialize_model_values()")
        if self.optimizer is not None:
            raise ValueError("build the optimizer after meta model materialization")

        aliases: dict[int, list[tuple[nn.Module, str, str]]] = {}
        for module in model.modules():
            for collection in ("_parameters", "_buffers"):
                for name, tensor in getattr(module, collection).items():
                    if tensor is not None:
                        aliases.setdefault(id(tensor), []).append((module, collection, name))
        allocation_device = torch.device(self.model_init_device(model))
        model.to_empty(device=allocation_device)
        for references in aliases.values():
            owner, collection, name = references[0]
            tensor = getattr(owner, collection)[name]
            if collection == "_buffers" and allocation_device != self.device:
                tensor = torch.empty_like(tensor, device=self.device)
            for owner, collection, name in references:
                getattr(owner, collection)[name] = tensor
        with torch.no_grad():
            self.initialize_model(self.unwrap(model))
        if any(tensor.is_meta for tensor in (*model.parameters(), *model.buffers())):
            raise RuntimeError("initialize_model() left meta tensors in the model")
        buffers: dict[int, torch.Tensor] = {}
        for module in model.modules():
            for name, buffer in module._buffers.items():
                if buffer is not None:
                    if id(buffer) not in buffers:
                        buffers[id(buffer)] = buffer.to(self.device)
                    module._buffers[name] = buffers[id(buffer)]
        return model

    def compile_regions(self, model: nn.Module) -> Mapping[str, Sequence[nn.Module]]:
        r"""Return model-owned named module groups, resolved after transforms.

        Override this or implement ``model.compile_regions()``. Members of a
        group share one forward implementation and compiler cache. Modules
        retain their identities, parameter names, and ``Module.__call__`` hooks.
        """
        regions = getattr(model, "compile_regions", None)
        return regions() if callable(regions) else {}

    def compile_model(self, model: nn.Module) -> nn.Module:
        r"""Compile selected named groups, or the whole model when none are selected."""
        requested = self.config.compile.regions
        if not self.compiler.enabled or not requested:
            return self.compiler.compile(model)
        if isinstance(requested, str) or len(set(requested)) != len(requested):
            raise ValueError("compile.regions must contain distinct region names")
        available = self.compile_regions(model)
        groups = []
        seen: set[int] = set()
        for name in requested:
            if name not in available:
                raise ValueError(f"unknown compile region {name!r}; available regions: {tuple(available)}")
            modules = tuple(available[name])
            if not modules or any(not isinstance(module, nn.Module) for module in modules):
                raise ValueError(f"compile region {name!r} must contain nn.Module instances")
            if seen.intersection(id(module) for module in modules):
                raise ValueError("selected compile regions must not share modules")
            seen.update(id(module) for module in modules)
            groups.append(modules)
        for modules in groups:
            self.compiler.compile_modules(modules)
        return model

    def materialize_model(self) -> None:
        """
        Move the model to the runtime device, optionally compile, and wrap
        with DDP when distributed.

        The default is a single-module DDP-style materialization: it moves
        `self.model` to `self.device`, applies any FP8 module policy when
        FP8 is enabled, runs `torch.compile` via `self.compiler`, and wraps the
        result with `nn.parallel.DistributedDataParallel` when world size > 1.
        The configured DDP-optimizer context covers both DDP construction and
        the compiled callable's lazy compilation and recompilation.

        **Called when:** once during `__post_init__`, after `setup_fp8()`
        and before `build_optimizer()`. The order matters — the optimizer
        must see post-wrap parameters.

        **Precondition:** `self.model` is set (typically by the user before
        constructing the runner). `self.device` resolves to the runtime
        device.

        Raises:
            ValueError: `self.model` is not initialized.

        **Side effects:** moves `self.model` to `self.device`; applies FP8
        module policy when `self.fp8_enabled`; compiles via
        `self.compiler.compile(...)`; wraps with `DistributedDataParallel` for
        world size > 1 under the same configured DDP-optimizer context. Moves
        `self.ema` to device when EMA is bound.

        !!! danger "Do not"
            - Build the optimizer or scheduler here; they run after this
              hook.
            - Skip the device move when overriding (tensors must live on
              `self.device` before the forward pass).
            - Re-wrap an already-wrapped model (e.g., DDP-wrap a DDP module).

        **Backend notes:**

        - `DeepSpeedRunner` overrides this hook to move the model to device
          and compile only; the DeepSpeed engine wraps the model later in
          `_finalize_runtime_components`.
        - `ParallelRunner` overrides this hook for FSDP2, pipeline-parallel
          schedules, and tensor/expert/context parallelism (via the
          `parallelize_model` and `apply_activation_checkpointing` hooks).
        """
        if self.model is None:
            raise ValueError("cannot materialize model: model is not initialized")

        model = self.materialize_model_values(self.model)
        self.model = model
        if self.fp8_enabled:
            self.apply_fp8_module_policy_to_model_parts()
            model = self.model
        should_wrap_ddp = self.distributed and not isinstance(
            model, (nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)
        )
        model = self.compile_model(model)
        if should_wrap_ddp:
            with self.compiler.ddp_optimizer():
                model = nn.parallel.DistributedDataParallel(
                    model,
                    find_unused_parameters=bool(self.config.ddp.find_unused_parameters),
                    static_graph=bool(self.config.ddp.static_graph),
                )
        self.model = model

        if self.ema is not None:
            self.ema = self.ema.to(self.device)

    def unwrap(self, model: nn.Module) -> nn.Module:
        """Return the native module beneath compile and data-parallel wrappers."""
        from torch._dynamo.eval_frame import OptimizedModule

        while isinstance(model, (OptimizedModule, nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)):
            model = model._orig_mod if isinstance(model, OptimizedModule) else model.module
        return model

    def _iter_unique_parameters(self, modules: Sequence[nn.Module]) -> Iterator[nn.Parameter]:
        seen: set[int] = set()
        for module in modules:
            for parameter in module.parameters():
                parameter_id = id(parameter)
                if parameter_id in seen:
                    continue
                seen.add(parameter_id)
                yield parameter

    def _iter_unique_named_parameters(
        self, modules: Sequence[nn.Module], prefixes: Sequence[str] | None = None
    ) -> Iterator[tuple[str, nn.Parameter]]:
        seen: set[int] = set()
        if prefixes is None:
            prefixes = ("",) * len(modules)
        if len(prefixes) != len(modules):
            raise ValueError("prefix count must match module count")
        for module, prefix in zip(modules, prefixes):
            for name, parameter in module.named_parameters():
                parameter_id = id(parameter)
                if parameter_id in seen:
                    continue
                seen.add(parameter_id)
                yield f"{prefix}{name}", parameter

    def iter_optimizer_parameters(self) -> Iterator[nn.Parameter]:
        if self.model is None:
            return
        yield from self._iter_unique_parameters((self.unwrap(self.model),))

    def iter_optimizer_named_parameters(self) -> Iterator[tuple[str, nn.Parameter]]:
        if self.model is None:
            return
        yield from self._iter_unique_named_parameters((self.unwrap(self.model),))

    def _optimizer_param_group_options(
        self,
        group_cfg: Mapping[str, Any],
        optimizer_cfg: Mapping[str, Any],
        *,
        index: int,
    ) -> dict[str, Any]:
        options = {
            str(key): value
            for key, value in group_cfg.items()
            if key not in {"pattern", "params", "lr_multiplier", "weight_decay_multiplier", "beta1", "beta2"}
        }
        if "lr_multiplier" in group_cfg:
            if optimizer_cfg.get("lr") is None:
                raise ValueError(f"optim.param_groups[{index}].lr_multiplier requires optim.lr")
            options["lr"] = float(optimizer_cfg["lr"]) * float(group_cfg["lr_multiplier"])
        if "weight_decay_multiplier" in group_cfg:
            if optimizer_cfg.get("weight_decay") is None:
                raise ValueError(f"optim.param_groups[{index}].weight_decay_multiplier requires optim.weight_decay")
            options["weight_decay"] = float(optimizer_cfg["weight_decay"]) * float(group_cfg["weight_decay_multiplier"])

        beta1 = group_cfg.get("beta1")
        beta2 = group_cfg.get("beta2")
        if beta1 is not None or beta2 is not None:
            if optimizer_cfg.get("betas") is None:
                raise ValueError(f"optim.param_groups[{index}].beta1/beta2 requires optim.betas")
            beta1_default, beta2_default = optimizer_cfg["betas"]
            options["betas"] = (
                float(beta1_default if beta1 is None else beta1),
                float(beta2_default if beta2 is None else beta2),
            )
        return options

    def _build_optimizer_param_groups(
        self, optimizer_cfg: Mapping[str, Any]
    ) -> list[nn.Parameter] | list[dict[str, Any]]:
        group_configs = optimizer_cfg.get("param_groups")
        if group_configs is None:
            return list(self.iter_optimizer_parameters())
        if isinstance(group_configs, (str, bytes, Mapping)) or not isinstance(group_configs, Sequence):
            raise ValueError("optim.param_groups must be a sequence of mappings")

        named_parameters = list(self.iter_optimizer_named_parameters())
        if not named_parameters:
            return []

        assigned: set[int] = set()
        param_groups: list[dict[str, Any]] = []
        for index, group_cfg in enumerate(group_configs):
            if not isinstance(group_cfg, Mapping):
                raise ValueError(f"optim.param_groups[{index}] must be a mapping")
            pattern = group_cfg.get("pattern")
            if pattern is None:
                raise ValueError(f"optim.param_groups[{index}] requires `pattern`")
            regex = re.compile(str(pattern))
            parameters = [
                parameter
                for name, parameter in named_parameters
                if id(parameter) not in assigned and regex.search(name) is not None
            ]
            if not parameters:
                warn(
                    f"optim.param_groups[{index}] pattern {pattern!r} matched no parameters",
                    RuntimeWarning,
                    stacklevel=2,
                )
                continue
            assigned.update(id(parameter) for parameter in parameters)
            param_groups.append(
                {
                    "params": parameters,
                    **self._optimizer_param_group_options(group_cfg, optimizer_cfg, index=index),
                }
            )

        unmatched = [parameter for _name, parameter in named_parameters if id(parameter) not in assigned]
        if unmatched:
            param_groups.append({"params": unmatched})
        return param_groups

    def build_optimizer(self) -> None:
        """
        Auto-build the optimizer from `config.optim` when `self.optimizer`
        is absent.

        The default iterates parameters via `iter_optimizer_parameters` and
        dispatches to the `OPTIMIZERS` registry with the merged config. If
        `optim.param_groups` is configured, entries are matched by regex
        `search` against `iter_optimizer_named_parameters`; unmatched
        parameters keep the optimizer-level defaults.

        **Called when:** once during `TorchRunner.__post_init__`, after
        `materialize_model` (so parameters reflect DDP/FSDP wrapping) and
        before `build_scheduler`.

        **Precondition:** `self.model` is materialized and on `self.device`.
        `self.optimizer` is `None` (the auto-build is skipped when the user
        has already bound an optimizer).

        **Side effects:** sets `self.optimizer` to the registry-built
        instance.

        !!! danger "Do not"
            - Run before `materialize_model`; parameters won't reflect
              DDP/FSDP wrapping.
            - Build a scheduler here.
            - Override parameter enumeration here; override
              `iter_optimizer_parameters` / `iter_optimizer_named_parameters`
              instead so subclass topology (e.g., `ParallelRunner.model_parts`)
              is preserved.

        **Backend notes:**

        - `DeepSpeedRunner` inherits this hook; DeepSpeed may replace the
          optimizer with a DeepSpeed-managed instance during
          `_finalize_runtime_components`.
        - `ParallelRunner` inherits this hook but overrides
          `iter_optimizer_parameters` to enumerate `self.model_parts`.
        """
        if self.optimizer is not None or self.model is None:
            return
        optimizer_cfg = self.config.get("optim")
        if not isinstance(optimizer_cfg, Mapping) or not optimizer_cfg:
            return
        optimizer_kwargs = {str(key): value for key, value in optimizer_cfg.items() if value is not None}
        optimizer_kwargs.pop("param_groups", None)
        if "type" not in optimizer_kwargs:
            return
        parameters = self._build_optimizer_param_groups(optimizer_cfg)
        if not parameters:
            return
        self.optimizer = OPTIMIZERS.build(params=parameters, **optimizer_kwargs)

    def _get_scheduler_config(self) -> Mapping[str, Any] | None:
        scheduler_cfg = self.config.get("sched")
        if not isinstance(scheduler_cfg, Mapping):
            return None
        return scheduler_cfg

    def build_scheduler(self) -> None:
        """
        Auto-build the LR scheduler from `config.sched` when
        `self.scheduler` is absent.

        The default pops `interval` and `monitor` from the config (those
        drive runner-level dispatch, not scheduler construction), defaults
        `total_steps` to `self.steps` when computable, and dispatches to
        the `SCHEDULERS` registry with `self.optimizer` and the merged
        config.

        **Called when:** once during `TorchRunner.__post_init__`, after
        `build_optimizer`.

        **Precondition:** `self.optimizer` is bound. `self.scheduler` is
        `None` (the auto-build is skipped when the user has already bound a
        scheduler).

        **Side effects:** sets `self.scheduler` to the registry-built
        instance.

        !!! danger "Do not"
            - Run before `build_optimizer`; the scheduler must wrap an
              optimizer.
            - Set scheduler interval or monitor here; configure them via
              `config.sched.interval` / `config.sched.monitor`.

        **Backend notes:**

        - `DeepSpeedRunner` inherits this hook; the scheduler may be handed
          to the DeepSpeed engine in `_finalize_runtime_components` when
          its effective interval is `"step"`. Otherwise the runner retains
          it.
        """
        if self.scheduler is not None or self.optimizer is None:
            return
        sched_cfg = self._get_scheduler_config()
        if not isinstance(sched_cfg, Mapping) or not sched_cfg:
            return
        scheduler_kwargs = {str(key): value for key, value in sched_cfg.items() if value is not None}
        scheduler_kwargs.pop("interval", None)
        scheduler_kwargs.pop("monitor", None)
        if "type" not in scheduler_kwargs:
            return
        if "total_steps" not in scheduler_kwargs:
            steps = self.steps
            if steps is not None:
                scheduler_kwargs["total_steps"] = steps
        self.scheduler = SCHEDULERS.build(self.optimizer, **scheduler_kwargs)

    def _bind_optimizer_container(self) -> None:
        if self.optimizer is None:
            self.optimizer_container = None
            return
        self.optimizer_container = OptimizerContainer(
            self.optimizer,
            scheduler=self.scheduler,
            scheduler_interval=self.scheduler_interval,
        )

    def _resolve_scheduler_metric(self, result: Mapping[str, Any]) -> Any:
        def scalarize(value: Any) -> Any:
            if isinstance(value, torch.Tensor):
                if value.numel() != 1:
                    raise ValueError(
                        "scheduler monitor must resolve to a scalar metric, "
                        f"but got tensor with shape {tuple(value.shape)}"
                    )
                return value.item()
            return value

        monitor = self.scheduler_monitor or self.config.score.metric

        if "." in monitor:
            value: Any = result
            for key in monitor.split("."):
                if not isinstance(value, Mapping) or key not in value:
                    raise ValueError(
                        f"could not resolve sched.monitor={monitor!r} from aggregated result {dict(result)!r}"
                    )
                value = value[key]
            return scalarize(value)

        score_split = self.score_split
        if score_split is not None:
            split_result = result.get(score_split)
            if isinstance(split_result, Mapping) and monitor in split_result:
                return scalarize(split_result[monitor])

        if monitor in result and not isinstance(result[monitor], Mapping):
            return scalarize(result[monitor])

        matches: list[tuple[str, Any]] = []
        for split_name, split_result in result.items():
            if isinstance(split_result, Mapping) and monitor in split_result:
                matches.append((split_name, split_result[monitor]))

        if len(matches) == 1:
            return scalarize(matches[0][1])
        if len(matches) > 1:
            splits = ", ".join(split_name for split_name, _ in matches)
            raise ValueError(
                f"ambiguous sched.monitor={monitor!r}: matched multiple splits ({splits}). "
                "Use '<split>.<metric>' to disambiguate."
            )

        raise ValueError(f"could not resolve sched.monitor={monitor!r} from aggregated result {dict(result)!r}")

    def _step_epoch_scheduler(self, result: Mapping[str, Any]) -> bool:
        if self.scheduler is None or self.scheduler_interval != "epoch":
            return False

        scheduler_metric = SCHEDULER_METRIC_UNSET
        if scheduler_requires_metric(self.scheduler):
            scheduler_metric = self._resolve_scheduler_metric(result)

        if self.optimizer_container is not None:
            return self.optimizer_container.step_scheduler(scheduler_metric=scheduler_metric)
        return step_scheduler(self.scheduler, scheduler_metric=scheduler_metric)

    def build_dataloaders(self):
        """
        Build dataloaders for dataset splits not already materialized.

        The default iterates `self.datasets`, merges `config.dataloader`
        defaults with split-specific overrides (`config.dataloader.<split>`),
        constructs a sampler via `build_datasampler`, and wraps each dataset
        in a `StatefulDataLoader` using `self.collate_fn`. Train splits
        default to `shuffle=True` and `drop_last=True`; non-train splits
        default to the opposite.

        **Called when:** once during `TorchRunner.__post_init__` when
        `self.datasets` is non-empty.

        **Precondition:** `self.datasets` is populated (typically by the
        user before constructing the runner). `self.dataloaders` is bound
        to a default-constructed `DataLoaderDict`.

        **Side effects:** populates `self.dataloaders[split]` for each
        split in `self.datasets` not already materialized. Existing entries
        in `self.dataloaders` are left untouched.

        !!! danger "Do not"
            - Override sampler logic here; override `build_datasampler`
              instead.
            - Override collation; set `self.collate_fn` or override
              `collate_fn` (classmethod) instead.
            - Bind the optimizer or scheduler here.

        **Backend notes:**

        - `ParallelRunner` substitutes `self.dataloaders` with a proxying
          dict in `__init__` so non-first/last pipeline stages receive a
          `StepProxyLoader` view. The build logic itself is inherited.
        """
        datasets = {k: d for k, d in self.datasets.items() if k not in self.dataloaders}
        dataloader_config = self.config.get("dataloader", NestedDict())
        default_kwargs = NestedDict({k: v for k, v in dataloader_config.items() if k not in self.datasets})
        split_kwargs = NestedDict({k: v for k, v in dataloader_config.items() if k in self.datasets})
        if self.config.deterministic and default_kwargs.get("worker_init_fn") is None:
            default_kwargs["worker_init_fn"] = _seed_dataloader_worker
        # Worker/iterator seeds must not consume the model RNG when a loader is
        # rebuilt on checkpoint resume, even with nondeterministic kernels.
        for k, dataset in datasets.items():
            kwargs = NestedDict(default_kwargs)
            if k in split_kwargs:
                kwargs.merge(split_kwargs[k], overwrite=True)
            if kwargs.get("generator") is None:
                generator = torch.Generator()
                if self.config.seed is None:
                    generator.seed()
                else:
                    generator.manual_seed(int(self.config.seed) + int(self.rank))
                kwargs["generator"] = generator
            is_train_split = k in self.train_splits
            shuffle = kwargs.pop("shuffle", is_train_split)
            collate_fn = kwargs.pop("collate_fn", self.collate_fn)
            batch_sampler = kwargs.pop("batch_sampler", None)
            if batch_sampler is not None:
                kwargs.pop("batch_size", None)
                kwargs.pop("drop_last", None)
                kwargs.pop("sampler", None)
                self.dataloaders[k] = StatefulDataLoader(
                    dataset,
                    batch_sampler=batch_sampler,
                    collate_fn=collate_fn,
                    **kwargs,
                )
                continue
            kwargs.setdefault("drop_last", is_train_split)
            sampler = kwargs.pop("sampler", None)
            if sampler is None:
                sampler = self.build_datasampler(dataset, split=k, shuffle=shuffle)
                if type(sampler) is StatefulRandomSampler:
                    # Loader worker seeds and sampler replay must not advance each other's RNG.
                    sampler.generator = torch.Generator().set_state(kwargs["generator"].get_state())
            self.dataloaders[k] = StatefulDataLoader(dataset, sampler=sampler, collate_fn=collate_fn, **kwargs)

    def build_datasampler(self, dataset: Any, *, split: str, shuffle: bool) -> Any:
        """
        Build the sampler for one dataset split.

        **Called when:** `build_dataloaders` materializes a split from
        `self.datasets`.

        Args:
            dataset: Dataset object for the split.
            split: Split name being materialized.
            shuffle: Whether this split should be sampled in shuffled order.

        Returns:
            A local random/sequential sampler in single-process mode, or a
            `DistributedSampler` in distributed mode.

        **Backend notes:**

        - `ParallelRunner` overrides replica/rank selection so data-parallel
          sampling follows its topology instead of raw global rank.
        """
        if self.distributed:
            num_replicas = self.world_size
            rank = self.rank
            if self.fault_tolerance is not None:
                num_replicas, rank = self.fault_tolerance.data_parallel_info(num_replicas, rank)
            return utils.data.distributed.DistributedSampler(
                dataset,
                num_replicas=num_replicas,
                rank=rank,
                shuffle=shuffle,
            )
        return StatefulRandomSampler(dataset) if shuffle else utils.data.SequentialSampler(dataset)

    @staticmethod
    def collate_fn(batch):
        return utils.data.dataloader.default_collate(batch)

    def to_device(self, data: Any):
        """Move one batch to runtime device; override in subclasses for custom fast paths."""
        return to_device(data, self.device)

    def _step_mode_split_budget(
        self,
        *,
        remaining_steps: int,
        remaining_splits: int,
        loader: Any,
    ) -> int:
        if remaining_steps <= 0:
            return 0
        if remaining_splits <= 0:
            return remaining_steps

        fair_share = max((remaining_steps + remaining_splits - 1) // remaining_splits, 1)
        loader_length = self._loader_length(loader)
        if loader_length is None:
            return fair_share

        loader_step_budget = max((loader_length + self.accum_steps - 1) // self.accum_steps, 1)
        return min(fair_share, loader_step_budget, remaining_steps)

    @staticmethod
    def _set_loader_epoch(loader: Any, epoch: int) -> None:
        batch_sampler = getattr(loader, "batch_sampler", None)
        if hasattr(batch_sampler, "set_epoch"):
            batch_sampler.set_epoch(epoch)  # type: ignore[union-attr]
        sampler = getattr(loader, "sampler", None)
        if hasattr(sampler, "set_epoch"):
            sampler.set_epoch(epoch)  # type: ignore[union-attr]

    def loop_time(self, *, sync: bool = False) -> float:
        if sync and torch.cuda.is_available() and self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        return perf_counter()

    @property
    def reports_batch_telemetry(self) -> bool:
        return True

    @staticmethod
    def _as_float_or_none(value: Any) -> float | None:
        if isinstance(value, (bool, int, float)):
            return float(value)
        if torch.is_tensor(value) and value.numel() == 1:
            return float(value.detach().item())
        return None

    def _mapping_loss_normalizer(self, mapping: Mapping[str, Any] | None) -> float | None:
        if mapping is None:
            return None
        for key in ("loss_normalizer", "num_valid_tokens", "valid_tokens", "num_tokens", "token_count"):
            if key in mapping:
                normalizer = self._as_float_or_none(mapping[key])
                if normalizer is not None:
                    if not math.isfinite(normalizer) or normalizer < 0:
                        raise ValueError("loss normalizer must be finite and nonnegative")
                    return normalizer
        return None

    def _tensor_loss_normalizer(self, target: Any) -> float | None:
        if not torch.is_tensor(target):
            return None
        if isinstance(self.criterion, nn.CrossEntropyLoss) and target.is_floating_point():
            class_dim = 1 if target.ndim > 1 else 0
            return float(target.numel() // target.shape[class_dim])
        ignore_index = getattr(self.criterion, "ignore_index", None)
        if ignore_index is not None:
            valid = target != int(ignore_index)
            weight = getattr(self.criterion, "weight", None)
            if isinstance(self.criterion, (nn.CrossEntropyLoss, nn.NLLLoss)) and weight is not None:
                return float(weight.to(device=target.device)[target[valid]].sum().item())
            return float(valid.sum().item())
        if getattr(self.criterion, "reduction", None) == "mean":
            return float(target.numel())
        return None

    def _get_loss_normalizer(self, data: Any) -> float | None:
        if isinstance(data, Mapping):
            explicit = self._mapping_loss_normalizer(data)
            if explicit is not None:
                return explicit

            target = data.get("target")
            if isinstance(target, Mapping):
                explicit = self._mapping_loss_normalizer(target)
                if explicit is not None:
                    return explicit
            if target is not None:
                normalizer = self._tensor_loss_normalizer(target)
                if normalizer is not None:
                    return normalizer

            inputs = data.get("input")
            if isinstance(inputs, Mapping):
                attention_mask = inputs.get("attention_mask")
                if isinstance(attention_mask, torch.Tensor):
                    return float(attention_mask.detach().sum().item())
            return None

        if isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
            target = data[1] if len(data) > 1 else None
            if isinstance(target, Mapping):
                explicit = self._mapping_loss_normalizer(target)
                if explicit is not None:
                    return explicit
            if target is not None:
                normalizer = self._tensor_loss_normalizer(target)
                if normalizer is not None:
                    return normalizer
        return None

    def _explicit_loss_normalizer(self, data: Any) -> float | None:
        mappings: list[Any] = []
        if isinstance(data, Mapping):
            mappings.extend((data, data.get("target")))
        elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)) and len(data) > 1:
            mappings.append(data[1])
        for mapping in mappings:
            if isinstance(mapping, Mapping):
                normalizer = self._mapping_loss_normalizer(mapping)
                if normalizer is not None:
                    return normalizer
        return None

    def _get_loss_reduction(self, data: Any = None) -> str:
        """Describe the scalar loss before the accumulation window is normalized.

        A mean is weighted by its actual denominator. A sum stays additive;
        providing an explicit normalizer instead declares a summed numerator
        whose desired denominator is that metadata value. Unknown criteria keep
        the historical microbatch-mean behavior.
        """
        if getattr(self.criterion, "reduction", None) != "sum":
            return "mean"
        return "normalized_sum" if self._explicit_loss_normalizer(data) is not None else "sum"

    def _compute_loss(self, pred: Any, target: Any) -> torch.Tensor | None:
        """Evaluate a criterion while giving an empty CE/NLL mean a zero VJP."""
        criterion = self.criterion
        if criterion is None:
            return None
        if isinstance(criterion, (nn.CrossEntropyLoss, nn.NLLLoss)):
            normalizer = self._tensor_loss_normalizer(target)
            if normalizer == 0 and criterion.reduction == "mean":
                if isinstance(criterion, nn.CrossEntropyLoss):
                    return F.cross_entropy(
                        pred,
                        target,
                        weight=criterion.weight,
                        ignore_index=criterion.ignore_index,
                        reduction="sum",
                        label_smoothing=criterion.label_smoothing,
                    )
                return F.nll_loss(
                    pred, target, weight=criterion.weight, ignore_index=criterion.ignore_index, reduction="sum"
                )
        return criterion(pred, target)

    def _loss_for_reporting(self, loss: torch.Tensor | None, normalizer: float | None, data: Any):
        self._loss_reporting_reduction = self._get_loss_reduction(data)
        if loss is not None and self._loss_reporting_reduction == "normalized_sum":
            return loss / normalizer if normalizer else loss * 0
        return loss

    def _update_loss_meter(self, loss: torch.Tensor, normalizer: float | None) -> None:
        weight = 1.0 if self._loss_reporting_reduction == "sum" or normalizer is None else normalizer
        self.meters.loss.update(loss.detach(), n=weight)

    def _sum_loss_for_reporting(self, loss: torch.Tensor) -> torch.Tensor:
        return self.all_reduce(loss, op=dist.ReduceOp.SUM)

    def get_epoch_result(self) -> RoundDict:
        result = super().get_epoch_result()
        if getattr(self, "_loss_reporting_reduction", None) == "sum":
            meter = self.meters.get("loss")
            loss = torch.as_tensor(
                0.0 if meter is None else meter.sum, dtype=torch.float64, device=self.all_reduce_device()
            )
            result["loss"] = self._sum_loss_for_reporting(loss).item()
        return result

    def _loss_normalizer_sync_divisor(self) -> int:
        if self.fault_tolerance is not None and self.fault_tolerance.replicate_process_group is not None:
            return max(self.world_size, 1) * max(
                int(dist.get_world_size(group=self.fault_tolerance.replicate_process_group)), 1
            )
        if dist.is_available() and dist.is_initialized():
            return max(self.world_size, 1)
        return 1

    def _reduce_loss_normalizer_total(self, local_total: float) -> float:
        if self._loss_normalizer_sync_divisor() <= 1:
            return local_total
        if not (dist.is_available() and dist.is_initialized()):
            return local_total

        device = self.all_reduce_device()
        total_tensor = torch.tensor(local_total, dtype=torch.float64, device=device)
        self.all_reduce(total_tensor, op=dist.ReduceOp.SUM)
        return float(total_tensor.item())

    def all_reduce_device(self) -> torch.device:
        if dist.is_available() and dist.is_initialized():
            group = self.all_reduce_group()
            if self.fault_tolerance is not None and group is self.fault_tolerance.replicate_process_group:
                # FT reductions first traverse the native local group. A Gloo
                # replica group also accepts CUDA; NCCL cannot accept CPU.
                group = None
            if group is not None:
                try:
                    backend = str(dist.get_backend(group=group)).lower()
                except TypeError:
                    backend = str(dist.get_backend(group)).lower()
                except (RuntimeError, ValueError):
                    return torch.device("cpu")
            else:
                backend = str(dist.get_backend()).lower()
            if "gloo" in backend or "mpi" in backend:
                return torch.device("cpu")
            if torch.cuda.is_available() and self.device.type == "cuda":
                return self.device
        return torch.device("cpu")

    def all_reduce_group(self):
        if self.fault_tolerance is not None and self.fault_tolerance.replicate_process_group is not None:
            return self.fault_tolerance.replicate_process_group
        return None

    def all_reduce(self, tensor: torch.Tensor, *, op=dist.ReduceOp.SUM) -> torch.Tensor:
        """Reduce tensor over the runner's replica/data-parallel collective domain."""
        if dist.is_available() and dist.is_initialized():
            group = self.all_reduce_group()
            ft_group = None if self.fault_tolerance is None else self.fault_tolerance.replicate_process_group
            if group is not None and group is ft_group:
                dist.all_reduce(tensor, op=op)
            dist.all_reduce(tensor, op=op, group=group)
        return tensor

    def _sync_optimizer_skip_decision(self, should_skip: bool) -> bool:
        if not (dist.is_available() and dist.is_initialized()):
            return should_skip
        payload = torch.tensor(float(should_skip), device=self.all_reduce_device())
        self.all_reduce(payload, op=dist.ReduceOp.MAX)
        return payload.item() > 0

    def reduce(self, tensor: torch.Tensor) -> torch.Tensor:
        """Average-reduce tensor over the runner's collective domain."""
        tensor = _local_reduction_tensor(tensor)
        if not (dist.is_available() and dist.is_initialized()):
            return tensor
        group = self.all_reduce_group()
        group_size = max(self.world_size if group is None else dist.get_world_size(group=group), 1)
        if (
            group is not None
            and self.fault_tolerance is not None
            and group is self.fault_tolerance.replicate_process_group
        ):
            group_size *= max(self.world_size, 1)
        if group_size <= 1:
            return tensor

        original_device = tensor.device
        payload_device = self.all_reduce_device()
        payload = tensor if original_device == payload_device else tensor.to(payload_device)
        self.all_reduce(payload, op=dist.ReduceOp.SUM)
        payload = payload / group_size
        if payload.device != original_device:
            payload = payload.to(original_device)
        return payload

    def reduce_loss_for_logging(self, loss: torch.Tensor | None, loss_n: float | None) -> torch.Tensor | None:
        """Detach and all-reduce weighted loss tensor for logging."""
        if loss is None:
            return None
        loss_value = _local_reduction_tensor(loss.detach()).to(dtype=torch.float64)
        if loss_value.ndim > 0:
            loss_value = loss_value.mean()
        normalizer = float(loss_n) if loss_n is not None else 1.0
        summed = getattr(self, "_loss_reporting_reduction", None) == "sum"
        if summed:
            normalizer = 1.0
        payload_device = self.all_reduce_device()
        payload = torch.stack(
            (
                loss_value.to(device=payload_device) * normalizer,
                torch.tensor(normalizer, dtype=torch.float64, device=payload_device),
            )
        )
        self.all_reduce(payload, op=dist.ReduceOp.SUM)
        if payload[1].item() <= 0:
            return None
        return payload[0] if summed else payload[0] / payload[1]

    def _reset_accumulation_normalization(self) -> None:
        """Clear the per-window accumulation state.

        Called at loop start and on every optimizer flush. After reset, the
        next call to `_scaled_loss_for_backward` re-classifies the window mode
        from its first batch's normalizer.
        """
        self._pending_loss_normalizer = None
        self._pending_loss_reduction = None
        self._accumulation_divisor_local = 0.0
        self._accumulation_divisor_total = None
        self._accumulation_mode = None

    def _loss_scale_for_backward(self) -> float:
        """Consume the pending loss-normalizer signal and return this micro-step's loss scale."""
        loss_normalizer = self._pending_loss_normalizer
        reduction = self._pending_loss_reduction or self._get_loss_reduction()
        mode = reduction if reduction != "mean" else ("weighted" if loss_normalizer is not None else "uniform")
        if self._accumulation_mode is None:
            self._accumulation_mode = mode

        if self._accumulation_mode != mode:
            raise ValueError("loss reduction or normalization changed within the current accumulation window")
        if mode in {"weighted", "normalized_sum"} and loss_normalizer is None:
            raise ValueError(
                "loss normalizer became unavailable within the current accumulation window. "
                "Override `train_step()` or provide consistent batch metadata for weighted normalization."
            )
        normalizer = float(loss_normalizer) if loss_normalizer is not None else 1.0
        if not math.isfinite(normalizer) or normalizer < 0:
            raise ValueError("loss normalizer must be finite and nonnegative")
        self._accumulation_divisor_local += normalizer
        self._pending_loss_normalizer = None
        self._pending_loss_reduction = None
        return normalizer if mode == "weighted" else 1.0

    def _scaled_loss_for_backward(self, loss: torch.Tensor) -> torch.Tensor:
        """Scale and accumulate loss for one micro-step inside an accumulation window.

        Means are multiplied by their real denominator and summed before one
        global division. A sum remains additive unless an explicit batch
        normalizer declares the denominator of that summed numerator. Empty
        microbatches contribute zero numerator and zero denominator. Unknown
        criteria without metadata retain uniform microbatch averaging.

        Loop overrides may publish ``_pending_loss_normalizer`` before calling
        ``backward``; a window must use a consistent reduction convention.
        """
        scale = self._loss_scale_for_backward()
        return loss * scale

    def _gradient_scale_for_step(self) -> float | None:
        if self._accumulation_mode is None:
            return None
        total = self._reduce_loss_normalizer_total(self._accumulation_divisor_local)
        self._accumulation_divisor_total = total
        if total <= 0:
            return 0.0
        if self._accumulation_mode == "sum":
            return float(max(self._loss_normalizer_sync_divisor(), 1))
        return float(max(self._loss_normalizer_sync_divisor(), 1)) / total

    def _accumulation_is_empty(self) -> bool:
        return self._accumulation_mode is not None and self._accumulation_divisor_total == 0

    def _optimizer_parameters_for_scaling(self) -> list[nn.Parameter]:
        if self.optimizer is None:
            return []
        if self.optimizer_container is not None:
            return self.optimizer_container.parameter_cache.get_parameters_for_clipping(self.optimizer)

        parameter_cache = self._optimizer_parameter_cache
        if parameter_cache is None:
            parameter_cache = OptimizerParameterCache(self.optimizer)
            self._optimizer_parameter_cache = parameter_cache
        else:
            parameter_cache.bind(self.optimizer)
        return parameter_cache.get_parameters_for_clipping()

    def _scale_optimizer_gradients(self, scale: float) -> None:
        if scale == 1.0 or self.optimizer is None:
            return
        parameters = self._optimizer_parameters_for_scaling()
        for parameter in parameters:
            grad = parameter.grad
            if grad is None:
                continue
            grad.mul_(float(scale))

    def _optimizer_has_nonfinite_grad(self) -> bool:
        if self.optimizer_container is not None:
            return self.optimizer_container.has_nan_inf_grad()

        parameters = self._optimizer_parameters_for_scaling()
        for parameter in parameters:
            grad = parameter.grad
            if grad is None:
                continue
            if grad.is_sparse:
                grad = grad.coalesce().values()
            elif grad.layout != torch.strided:
                grad = grad.values()
            if not torch.isfinite(grad).all():
                return True
        return False

    def _zero_optimizer_grad(self) -> None:
        if self.optimizer_container is not None:
            self.optimizer_container.zero_grad()
        elif self.optimizer is not None:
            self.optimizer.zero_grad()

    def _sync_fault_tolerant_gradients(self) -> None:
        """Average already-local-DDP-reduced gradients across FT replicas."""
        if self.fault_tolerance is None:
            return
        group = self.fault_tolerance.replicate_process_group
        if group is None:
            return
        degree = max(int(dist.get_world_size(group=group)), 1)
        for parameter in self._optimizer_parameters_for_scaling():
            present = torch.tensor(int(parameter.grad is not None), device=parameter.device)
            dist.all_reduce(present, group=group, op=dist.ReduceOp.MAX)
            if not bool(present.item()):
                continue
            gradient = torch.zeros_like(parameter) if parameter.grad is None else parameter.grad
            if gradient.is_sparse:
                gradient = gradient.to_dense()
            dist.all_reduce(gradient, group=group, op=dist.ReduceOp.SUM)
            gradient.div_(degree)
            parameter.grad = gradient

    def _clip_optimizer_gradients(
        self,
        *,
        max_grad_value: float | None,
        max_grad_norm: float | None,
    ) -> float | None:
        if self.optimizer_container is not None:
            return self.optimizer_container.clip_gradients(
                max_grad_value=max_grad_value,
                max_grad_norm=max_grad_norm,
                pipeline_group=self._gradient_clipping_group(),
                device=self.device,
            )

        if self.optimizer is None or (max_grad_value is None and max_grad_norm is None):
            return None

        parameters = self._optimizer_parameters_for_scaling()
        if max_grad_value is not None:
            clip_grad_value_(parameters, max_grad_value)
        if max_grad_norm is not None:
            grad_norm = clip_grad_norm(
                parameters, max_grad_norm, pipeline_group=self._gradient_clipping_group(), device=self.device
            )
            return float(grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm)
        return None

    def _gradient_clipping_group(self):
        return None

    def _optimizer_step_and_zero_grad(self) -> None:
        if self.optimizer is None:
            raise ValueError("cannot step optimizer: optimizer is not configured")
        if self.grad_scaler is not None:
            has_gradients = any(parameter.grad is not None for parameter in self._optimizer_parameters_for_scaling())
            if has_gradients:
                self.grad_scaler.step(self.optimizer)
            self._update_grad_scaler(committed=has_gradients)
        else:
            self.optimizer.step()
        if self.optimizer_container is not None and self.optimizer_container.scheduler_interval == "step":
            self.optimizer_container.step_scheduler()
        self._zero_optimizer_grad()

    def _update_grad_scaler(self, *, overflow: bool = False, committed: bool = True) -> None:
        """Advance AMP only for an update, using a global overflow decision.

        A finite rejected step must clear the scaler's per-optimizer stage
        without advancing its growth tracker. For overflow, load a backed-off
        public state on every optimizer rank; local found-inf buffers alone do
        not cover independent pipeline stages.
        """
        scaler = self.grad_scaler
        if scaler is None:
            return
        if overflow:
            state = scaler.state_dict()
            if state:
                state["scale"] *= state["backoff_factor"]
                state["_growth_tracker"] = 0
                scaler.update(new_scale=state["scale"])
                scaler.load_state_dict(state)
                return
        if not committed:
            scaler.update(new_scale=scaler.get_scale())
        else:
            scaler.update()

    @contextmanager
    def train_context(self):
        """Context for one training micro-step (autocast + optional DDP no_sync)."""
        self._begin_optimizer_window()
        with self._train_step_context(no_sync_targets=self._train_no_sync_targets()):
            yield

    def _begin_optimizer_window(self) -> None:
        """Allow fault-tolerant healing before reading a new window's data."""
        if self.fault_tolerance is not None:
            self.fault_tolerance.begin_step()

    def _should_train_no_sync(self) -> bool:
        if self._train_window_will_flush:
            return False
        micro_steps = self.train_state.micro_step + 1
        return self.accum_steps > 1 and micro_steps % self.accum_steps != 0

    def infer_context(self):
        """Precision context used by train/eval/infer forward passes."""

        if self.fp8_enabled:
            return self.fp8_autocast()

        precision = self.precision
        if precision is None:
            return nullcontext()
        return torch.autocast(self.device.type, dtype=get_precision(precision))

    @staticmethod
    def _metric_state_ids(metrics) -> set[int]:
        state = {id(metrics)} if metrics is not None else set()
        if type(metrics) in (MultiTaskMetrics, StreamMetrics):
            for child in metrics.values():
                state.update(TorchRunner._metric_state_ids(child))
        elif type(metrics) is MetricMeter and isinstance(metrics.metric, TorchMetric):
            state.update(id(module) for module in metrics.metric.modules() if isinstance(module, TorchMetric))
        return state

    def build_evaluation_metrics(self):
        """Provide a collector for an isolated training-time validation pass.

        A separate evaluation collector retains its normal reset/update
        semantics. Native collectors shared with training are rebuilt from
        their metric definitions. Override this factory for a shared custom
        collector; it must return an independent instance, not a training one.
        TorchMetrics Metric callables use their public clone/reset lifecycle.
        Other metric functions and preprocessing callbacks must be stateless;
        override this factory when their callable objects carry mutable state.
        """
        metrics = self.evaluate_metrics
        training = self._metric_state_ids(self.train_metrics) | self._metric_state_ids(self.metrics)
        if not (self._metric_state_ids(metrics) & training):
            return metrics

        def fresh_metric(metric):
            if isinstance(metric, TorchMetric):
                metric = metric.clone()
                metric.reset()
            return metric

        def fresh(source):
            if type(source) is GlobalMetrics:
                return GlobalMetrics(
                    preprocess=source.preprocess,
                    distributed=source.distributed,
                    device=source.device,
                    **source.metrics,
                )
            if type(source) is StreamMetrics:
                return StreamMetrics(
                    preprocess=source.preprocess,
                    distributed=source.distributed,
                    device=source.device,
                    **{name: fresh_metric(meter.metric) for name, meter in source.all_items()},
                )
            if type(source) is MultiTaskMetrics:
                return MultiTaskMetrics(
                    aggregate=source.getattr("_aggregate", None),
                    aggregate_weights=source.getattr("_aggregate_weights", None),
                    **{name: fresh(child) for name, child in source.items()},
                )
            raise TypeError("shared custom metrics require build_evaluation_metrics() to return a fresh collector")

        return fresh(metrics)

    def _evaluation_data_parallel_info(self) -> tuple[int, int]:
        """Validation covers the dataset within each native replica group."""
        return self.world_size, self.rank

    def build_evaluation_dataloader(self, split: str):
        """Provide a validation loader starting at its sample-index origin.

        Only training-time validation calls this factory. Native
        map-style loaders with standard samplers are reconstructed from public
        configuration at epoch zero. Custom loaders, stateful datasets and
        iterable streams must override this hook to define their fresh origin.
        Standalone evaluation continues to consume the caller's original loader.
        """
        # ParallelRunner exposes StepProxyLoader views via __getitem__; its
        # items remain the native loaders. Replacing raw entries lets that
        # mapping create fresh stage-specific proxies without unwrapping them.
        source = dict(self.dataloaders.items())[split]
        if type(source) in (list, tuple, range):
            return source
        if type(source) not in (utils.data.DataLoader, StatefulDataLoader):
            raise TypeError("custom validation loaders must provide build_evaluation_dataloader(split)")
        dataset = source.dataset
        if isinstance(dataset, utils.data.IterableDataset) or callable(getattr(dataset, "load_state_dict", None)):
            raise TypeError("stateful or iterable validation datasets must provide build_evaluation_dataloader(split)")

        def generator(original):
            seed = torch.initial_seed() if self.config.seed is None else int(self.config.seed) + self.rank
            return torch.Generator().manual_seed(original.initial_seed() if original is not None else seed)

        sampler = source.sampler
        if type(sampler) is utils.data.SequentialSampler:
            sampler = utils.data.SequentialSampler(dataset)
        elif type(sampler) is utils.data.distributed.DistributedSampler:
            degree, rank = self._evaluation_data_parallel_info()
            sampler = utils.data.distributed.DistributedSampler(
                dataset,
                num_replicas=degree,
                rank=rank,
                shuffle=sampler.shuffle,
                seed=sampler.seed,
                drop_last=sampler.drop_last,
            )
        elif type(sampler) in (utils.data.RandomSampler, StatefulRandomSampler):
            sampler = type(sampler)(
                dataset,
                replacement=sampler.replacement,
                num_samples=sampler.num_samples,
                generator=generator(sampler.generator),
            )
        else:
            raise TypeError("custom validation samplers must provide build_evaluation_dataloader(split)")
        kwargs = {
            name: getattr(source, name)
            for name in (
                "num_workers",
                "collate_fn",
                "pin_memory",
                "timeout",
                "worker_init_fn",
                "multiprocessing_context",
                "prefetch_factor",
                "persistent_workers",
                "pin_memory_device",
                "in_order",
                "snapshot_every_n_steps",
            )
            if hasattr(source, name)
        }
        kwargs["generator"] = generator(source.generator)
        if source.batch_size is None and source.batch_sampler is not None:
            batch_sampler = source.batch_sampler
            if type(batch_sampler) not in (utils.data.BatchSampler, StatefulBatchSampler):
                raise TypeError("custom validation batching must provide build_evaluation_dataloader(split)")
            kwargs["batch_sampler"] = type(batch_sampler)(sampler, batch_sampler.batch_size, batch_sampler.drop_last)
        else:
            kwargs.update(sampler=sampler, batch_size=source.batch_size, drop_last=source.drop_last)
        return type(source)(dataset, **kwargs)

    @contextmanager
    def _evaluation_context(self, evaluate_splits: Sequence[str]):
        """Isolate validation inserted at a safe training boundary.

        Callers must complete any distributed healing before entering. Factories
        execute after the RNG snapshot and are covered by exception restoration.
        This scope does not begin an optimizer/FT transaction or change progress.
        Checkpoints within it export the entering training RNG and loader state.
        """
        names = (
            "meters",
            "metrics",
            "train_metrics",
            "evaluate_metrics",
            "split",
            "_evaluation_result_index",
            "_loss_reporting_reduction",
            "_pending_loss_normalizer",
            "_pending_loss_reduction",
            "_accumulation_mode",
            "_accumulation_divisor_local",
            "_accumulation_divisor_total",
            "_pipeline_loss_reduction",
            "_pipeline_loss_divisor_local",
            "_pipeline_loss_weighting",
            "_pipeline_backward_scaled",
            "_pipeline_batch_normalizer",
            "_pipeline_target_normalizer",
            "_pipeline_normalizer_ratio",
        )
        absent = object()
        saved = {name: getattr(self, name, absent) for name in names}
        mode = self.mode
        modules: dict[int, tuple[nn.Module, bool]] = {}
        for model in [self.model, self.ema, *(getattr(self, "model_parts", None) or ())]:
            if isinstance(model, nn.Module):
                for module in model.modules():
                    modules.setdefault(id(module), (module, module.training))
        python_rng = random.getstate()
        numpy_rng = np_random.get_state() if np_random is not None else None
        cpu_rng = torch.get_rng_state()
        devices = [torch.cuda.current_device()] if torch.cuda.is_initialized() else []
        if self.device.type == "cuda":
            device = self.device.index if self.device.index is not None else torch.cuda.current_device()
            if device not in devices:
                devices.append(device)
        cuda_rng = {device: torch.cuda.get_rng_state(device) for device in devices}
        stored_rng = self.rng_state.state_dict()
        previous_rng = getattr(self, "_evaluation_training_rng", None)
        previous_loaders = getattr(self, "_evaluation_training_dataloaders", None)
        original_loaders = self.dataloaders
        try:
            with torch.random.fork_rng(devices=devices):
                self._evaluation_training_rng = previous_rng or {
                    "python": python_rng,
                    "numpy": numpy_rng,
                    "torch_cpu": cpu_rng,
                    "cuda": cuda_rng,
                }
                self._evaluation_training_dataloaders = (
                    previous_loaders if previous_loaders is not None else copy.deepcopy(self.dataloaders.state_dict())
                )
                metrics = self.build_evaluation_metrics()
                training_metrics = self._metric_state_ids(self.train_metrics) | self._metric_state_ids(self.metrics)
                if self._metric_state_ids(metrics) & training_metrics:
                    raise TypeError("build_evaluation_metrics() must not return a training collector")
                loaders = {split: self.build_evaluation_dataloader(split) for split in evaluate_splits}
                self.evaluate_metrics = metrics
                self._evaluation_result_index = self.train_state.global_step
                self.dataloaders = copy.copy(original_loaders)
                for split, loader in loaders.items():
                    self.dataloaders[split] = loader
                yield
        finally:
            self._mode = mode
            for module, training in modules.values():
                module.training = training
            self.dataloaders = original_loaders
            for name, value in saved.items():
                if value is absent:
                    self.__dict__.pop(name, None)
                else:
                    setattr(self, name, value)
            random.setstate(python_rng)
            if numpy_rng is not None:
                np_random.set_state(numpy_rng)
            self.rng_state.load_state_dict(stored_rng)
            self._evaluation_training_rng = previous_rng
            self._evaluation_training_dataloaders = previous_loaders

    def _train_no_sync_targets(self) -> tuple[nn.Module, ...]:
        if isinstance(self.model, nn.parallel.DistributedDataParallel):
            return (self.model,)
        return ()

    @contextmanager
    def _train_step_context(self, *, no_sync_targets: tuple[nn.Module, ...] | list[nn.Module] = ()):
        with ExitStack() as stack:
            if self.diagnostics is not None:
                stack.enter_context(self.diagnostics.span("train.microbatch"))
            stack.enter_context(self.infer_context())
            if self._should_train_no_sync() and no_sync_targets:
                for module in no_sync_targets:
                    no_sync = getattr(module, "no_sync", None)
                    if callable(no_sync):
                        stack.enter_context(no_sync())
                        continue

                    set_requires_gradient_sync = getattr(module, "set_requires_gradient_sync", None)
                    if callable(set_requires_gradient_sync):
                        set_requires_gradient_sync(False)
                        stack.callback(set_requires_gradient_sync, True)
                        continue

                    raise TypeError(
                        "cannot disable gradient synchronization for "
                        f"{type(module).__name__}: expected `no_sync()` or `set_requires_gradient_sync(...)`"
                    )
            yield

    def train_step(self, data: Any) -> tuple[Any, torch.Tensor | None]:
        """
        Run one training micro-step.

        The default implementation runs forward → loss → metric update → backward
        → step for one micro-batch.

        **Called when:** once per micro-batch by `train_epoch`/`train_steps`. The
        caller seeds the loop's accumulation state before each invocation; this
        method consumes that state through `backward()` and `step()`.

        **Precondition:** `self.model`, `self.optimizer`, and `self.criterion`
        are bound; `self.mode == RunnerMode.train`.

        Args:
            data: One micro-batch. The default unpacks `data["input"]` /
                `data.get("target")` for mappings, `(data[0], data[1])` for
                non-string sequences, and `(data, None)` otherwise. Override
                `train_step` if your batch shape differs.

        Returns:
            `(pred, loss)`. `pred` is the model output (used by `metrics.update`).
            `loss` is the scalar loss returned to the caller for reduced logging.
            The default raises when `criterion` is missing or returns `None`;
            overrides may return `(pred, None)` to signal no loss available, in
            which case the caller skips loss bookkeeping.

        Raises:
            ValueError: `self.model` is not initialized, or `criterion` is missing
                or returned `None`.

        **Side effects:** moves `data` to `self.device`, runs forward under
        `train_context()` (autocast + optional DDP no-sync), updates
        `self.metrics` when bound, then calls `self.backward(loss)` and
        `self.step()` to scale gradients, advance accumulation state, and flush
        the optimizer on accumulation boundaries.

        !!! danger "Do not"
            - Zero gradients (`optimizer_step` does this on flush).
            - Call `self.optimizer.step()` directly (use `self.step()`).
            - Mutate `train_state.global_step` or `train_state.micro_step`.
            - Implement gradient scaling here (override `_backward()` instead).
            - Call `save_checkpoint()` (cadence is owned by the loop method).

        **Backend notes:**

        - `DeepSpeedRunner` inherits the default; `_backward`/`step` route
          through the DeepSpeed engine.
        - `ParallelRunner` overrides this method when a pipeline schedule is
          set; the schedule owns micro-batching and loss reduction.
        """
        data = self.to_device(data)
        with self.train_context():
            if isinstance(data, Mapping):
                inputs = data["input"]
                target = data.get("target")
            elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
                inputs = data[0]
                target = data[1] if len(data) > 1 else None
            else:
                inputs = data
                target = None

            if self.model is None:
                raise ValueError("cannot run train_step: model is not initialized")
            pred = self.model(**inputs) if isinstance(inputs, Mapping) else self.model(inputs)
            loss = self._compute_loss(pred, target)
            if loss is None:
                raise ValueError("cannot run train_step: criterion did not produce a loss")
            skipped = self.backward(loss)
            if not skipped and self.metrics is not None and pred is not None and target is not None:
                self.metrics.update(pred, target)
            self.step()
        return pred, loss.detach().new_zeros(()) if skipped else loss

    def backward(self, loss: torch.Tensor) -> bool:
        """Backpropagate one loss, or a synchronized zero when any replica is non-finite.

        Returns:
            Whether this micro-batch was replaced with zero.
        """

        zero_weight = self._pending_loss_normalizer == 0
        globally_nonfinite = False
        fault_tolerant = self.fault_tolerance is not None and self.fault_tolerance.enabled
        nonfinite = False
        if self.skip_nonfinite_loss or fault_tolerant:
            # Every rank votes, including an empty rank whose criterion may
            # legitimately be undefined. Zero weight makes that vote finite.
            nonfinite = not zero_weight and not bool(torch.isfinite(loss.detach()).all().item())
        if self.skip_nonfinite_loss:
            globally_nonfinite = self._sync_optimizer_skip_decision(nonfinite)
        if fault_tolerant and self.fault_tolerance is not None and (nonfinite or globally_nonfinite):
            self.fault_tolerance.record_nonfinite_loss()
        if zero_weight or globally_nonfinite:
            self._backward_zero_vjp(loss, preserve_missing=globally_nonfinite)
            return True
        self._backward(loss)
        return False

    def _backward_zero_vjp(self, loss: torch.Tensor, *, preserve_missing: bool = False) -> None:
        # Replace only this contribution before AccumulateGrad and distributed
        # reducers see it. Prior window gradients then synchronize normally.
        parameters = [parameter for parameter in self.iter_optimizer_parameters() if parameter.requires_grad]
        presence = None
        if preserve_missing:
            presence = torch.tensor(
                [float(parameter.grad is not None) for parameter in parameters], device=self.all_reduce_device()
            )
            if not self._should_train_no_sync():
                self.all_reduce(presence, op=dist.ReduceOp.MAX)
        handles = [parameter.register_hook(torch.zeros_like) for parameter in parameters]
        zero = loss.detach().new_zeros((), requires_grad=True)
        if loss.requires_grad:
            zero = zero + torch.where(torch.isfinite(loss), loss, torch.zeros_like(loss)).sum() * 0
        else:
            for parameter in parameters:
                zero = zero + parameter.reshape(-1)[:1].real.nan_to_num().sum() * 0
        try:
            self._backward(zero)
        finally:
            for handle in handles:
                handle.remove()
        if presence is not None:
            for parameter, existed in zip(parameters, presence.tolist()):
                if not existed:
                    parameter.grad = None

    def _backward(self, loss: torch.Tensor) -> None:
        """
        Execute the backend backward pass on one micro-step loss.

        **Called when:** the default `train_step` has produced a loss tensor.
        The method receives the raw micro-step loss; accumulation scaling and
        loss-normalizer weighting are applied before `Tensor.backward()`.

        Args:
            loss: The loss tensor for this micro-step.

        **Side effects:** accumulates gradients on model parameters.

        !!! danger "Do not"
            - Advance the optimizer here; optimizer stepping belongs to
              `step()`/`optimizer_step()`.
            - Mutate `train_state` counters.

        **Backend notes:**

        - `DeepSpeedRunner` overrides this hook to call the DeepSpeed engine's
          backward method.
        """
        scaled_loss = self._scaled_loss_for_backward(loss)
        if self.grad_scaler is not None:
            self.grad_scaler.scale(scaled_loss).backward()
            return
        scaled_loss.backward()

    def step(self) -> None:
        """
        Advance the accumulation state machine after one training micro-step.

        **Called when:** `train_step` finishes backward for a micro-batch.

        **Side effects:** increments `train_state.micro_step` and calls
        `optimizer_step()` only when the accumulation boundary is reached or
        the surrounding loop marks the current batch as the final flush in a
        partial window.

        !!! danger "Do not"
            - Call this from evaluation/inference paths.
            - Call `optimizer_step()` in addition to this method from the same
              micro-step.
            - Adjust `train_state.micro_step` in `train_step` overrides.
        """
        micro_steps = self.train_state.micro_step + 1
        self.train_state.micro_step = micro_steps
        if self._train_window_will_flush:
            self.optimizer_step()
            remainder = micro_steps % self.accum_steps
            if self.accum_steps > 1 and remainder != 0:
                self.train_state.micro_step += self.accum_steps - remainder
            return
        if self.accum_steps <= 1 or micro_steps % self.accum_steps == 0:
            self.optimizer_step()

    def optimizer_step(self) -> bool:
        """
        Perform one backend optimizer update.

        The default Torch implementation waits for checkpoint staging, applies
        accumulated-loss gradient scaling, optional grad clipping, non-finite
        grad skip logic, optimizer/scheduler stepping through
        `OptimizerContainer`, gradient zeroing, profiler advancement, and
        garbage-collection cadence.

        **Called when:** `step()` reaches an accumulation boundary, or
        `_flush_pending_optimizer_step()` flushes a partial boundary before
        shutdown.

        Returns:
            `True` when an optimizer update is applied, otherwise `False`.

        **Side effects:** may update optimizer/scheduler state; increments
        `train_state.global_step` only when an update is actually applied.

        !!! danger "Do not"
            - Increment `global_step` on skipped updates.
            - Forget to zero gradients after a successful update or skipped
              non-finite update.
            - Bypass `checkpoint_manager.maybe_wait_for_staging()`.

        **Backend notes:**

        - `DeepSpeedRunner` overrides this hook because the DeepSpeed engine
          owns the concrete optimizer update.
        """
        if self.optimizer_container is None and self.optimizer is None:
            raise ValueError(
                "cannot perform optimizer step: no optimizer is configured; "
                "set `self.optimizer`, implement `build_optimizer()`, or override `optimizer_step()`"
            )

        self.validate_ema_update_contract()
        self._begin_optimizer_window()
        context = self.diagnostics.span("checkpoint.staging_wait") if self.diagnostics is not None else nullcontext()
        with context:
            self.checkpoint_manager.maybe_wait_for_staging()
        self._sync_fault_tolerant_gradients()
        if self.grad_scaler is not None and self.optimizer is not None:
            self.grad_scaler.unscale_(self.optimizer)
        grad_scale = self._gradient_scale_for_step()
        if grad_scale is not None:
            self._scale_optimizer_gradients(grad_scale)
        empty_window = self._accumulation_is_empty()
        max_grad_value = self.max_grad_value
        max_grad_norm = self.max_grad_norm
        skip_nonfinite_grad = self.skip_nonfinite_grad
        fault_tolerant = self.fault_tolerance is not None and self.fault_tolerance.enabled
        has_nonfinite_grad = False
        if skip_nonfinite_grad or self.grad_scaler is not None or fault_tolerant:
            has_nonfinite_grad = self._optimizer_has_nonfinite_grad()
            has_nonfinite_grad = self._sync_optimizer_skip_decision(has_nonfinite_grad)
        commit = not has_nonfinite_grad and not empty_window
        if fault_tolerant and self.fault_tolerance is not None:
            commit = self.fault_tolerance.should_commit(nonfinite=has_nonfinite_grad) and not empty_window
        if not commit:
            self._zero_optimizer_grad()
            self._update_grad_scaler(overflow=has_nonfinite_grad, committed=False)
            self._reset_accumulation_normalization()
            return False

        grad_norm = self._clip_optimizer_gradients(max_grad_value=max_grad_value, max_grad_norm=max_grad_norm)
        self._record_grad_norm(grad_norm)
        self._optimizer_step_and_zero_grad()
        self.update_ema()

        self._reset_accumulation_normalization()
        self.train_state.global_step += 1
        self._step_profiler()
        self._maybe_reduce_train_process_group_timeout()
        self.supervisor.maybe_collect_garbage(self.train_state.global_step, scope="train")
        return True

    def _flush_pending_optimizer_step(self) -> bool:
        """
        Flush a partial accumulation window at loop boundaries.

        Returns:
            `True` when a boundary flush produced an optimizer update.
        """
        if self.accum_steps <= 1:
            return False
        remainder = self.train_state.micro_step % self.accum_steps
        if remainder == 0:
            return False
        if self.distributed and self._train_no_sync_targets():
            self._discard_pending_optimizer_step(remainder)
            return False
        stepped = self.optimizer_step()
        # Boundary flush clears current accumulation window; realign to the next
        # accumulation boundary so the next loop starts with a fresh full window.
        self.train_state.micro_step += self.accum_steps - remainder
        return stepped

    def _discard_pending_optimizer_step(self, remainder: int | None = None) -> None:
        if self.accum_steps <= 1:
            return
        if remainder is None:
            remainder = self.train_state.micro_step % self.accum_steps
        if remainder == 0:
            return
        if self.optimizer_container is not None:
            self.optimizer_container.zero_grad()
        elif self.optimizer is not None:
            self.optimizer.zero_grad()
        self._reset_accumulation_normalization()
        self.train_state.micro_step -= remainder

    def prepare_for_shutdown_checkpoint(self) -> None:
        self._flush_pending_optimizer_step()

    def can_checkpoint_for_shutdown(self) -> bool:
        """A fetched batch and its accumulation window must be consumed before saving."""
        return getattr(self, "_train_checkpoint_ready", True)

    def synchronize_termination_signal(self, signum: int | None) -> int | None:
        """Let one signaled rank stop its torchrun replica at a shared safe boundary."""
        if not (dist.is_available() and dist.is_initialized()):
            return signum
        # TorchFT replicas have independent default process groups. A shutdown
        # vote must not wait for its managed group or stop the other replicas.
        backend = str(dist.get_backend()).lower()
        device = torch.device("cuda", torch.cuda.current_device()) if "nccl" in backend else torch.device("cpu")
        pending = torch.tensor(0 if signum is None else signum, dtype=torch.int32, device=device)
        dist.all_reduce(pending, op=dist.ReduceOp.MAX)
        return int(pending.item()) or None

    def _next_train_batch(self, iterator):
        diagnostics = getattr(self, "diagnostics", None)
        context = diagnostics.span("data.fetch") if diagnostics is not None else nullcontext()
        with context:
            return next(iterator, None)

    def _iter_train_batches(self, loader: Any) -> Iterator[tuple[int, Any, bool]]:
        iterator = None
        generation = -1

        def begin_window():
            nonlocal iterator, generation
            self._train_checkpoint_ready = True
            self.supervisor.maybe_handle_termination_signal()
            # Recovery may restore loader/RNG state. Start it before fetching
            # the window's first batch, rather than after selecting its data.
            self._begin_optimizer_window()
            current_generation = getattr(self, "_dataloader_state_generation", 0)
            if iterator is None or generation != current_generation:
                iterator = iter(enumerate(loader))
                generation = current_generation
            self._train_checkpoint_ready = False
            return self._next_train_batch(iterator)

        current = begin_window()
        while current is not None:
            iteration, data = current
            next_micro_step = self.train_state.micro_step + 1
            reaches_accum_boundary = self.accum_steps <= 1 or next_micro_step % self.accum_steps == 0
            # Only peek when a partial accumulation window needs an end-of-loader
            # check. At optimizer/checkpoint boundaries the loader must describe
            # the last consumed batch, without an uncheckpointed lookahead batch.
            next_item = None if reaches_accum_boundary else self._next_train_batch(iterator)
            yield iteration, data, reaches_accum_boundary or next_item is None
            current = begin_window() if reaches_accum_boundary else next_item
        self._train_checkpoint_ready = True

    def _resolve_requested_splits(
        self,
        requested_splits: list[str] | None,
        available_splits: list[str],
        *,
        kind: str,
    ) -> list[str]:
        if requested_splits is None:
            return available_splits

        splits = self._sorted_unique(requested_splits)
        unknown_splits = sorted(set(splits).difference(available_splits))
        if unknown_splits:
            raise ValueError(
                f"unknown {kind} split(s): {unknown_splits}; " f"available {kind} split(s): {available_splits}"
            )
        return splits

    def train(
        self,
        train_splits: list[str] | None = None,
        evaluate_splits: list[str] | None = None,
    ) -> RoundDict:
        """
        Run the full training workflow.

        Selects epoch mode or step mode from `self.is_step_mode`, validates
        explicit split lists against the runner's configured/inferred splits,
        and delegates to `train_epochs` or `train_steps`.

        **Called when:** user code starts training.

        Args:
            train_splits: Optional training splits. When `None`, use `self.train_splits`.
            evaluate_splits: Optional evaluation splits. When `None`, use `self.evaluate_splits`.

        Returns:
            Aggregated runner results (`self.results`).

        Raises:
            ValueError: no valid training split can be resolved.

        **Side effects:** prints selected splits and runs the selected training
        loop. Checkpointing, result writing, scheduler stepping, and early stop
        are owned by the delegated loop method.
        """

        train_splits = self._resolve_requested_splits(train_splits, self.train_splits, kind="training")
        if not train_splits:
            raise ValueError("cannot start training: no valid training split was resolved")

        evaluate_splits = self._resolve_requested_splits(evaluate_splits, self.evaluate_splits, kind="evaluation")

        self.log_startup_summary(train_splits=train_splits, evaluate_splits=evaluate_splits)
        try:
            if self.is_step_mode:
                result = self.train_steps(train_splits=train_splits, evaluate_splits=evaluate_splits)
            else:
                result = self.train_epochs(train_splits=train_splits, evaluate_splits=evaluate_splits)
        except Exception as exc:
            self.log_crash_summary(exc)
            raise
        self.log_final_summary()
        return result

    @contextmanager
    def _training_evaluation_splits(self, splits: list[str]):
        previous = getattr(self, "_active_evaluation_splits", None)
        self._active_evaluation_splits = splits
        try:
            yield
        finally:
            if previous is None:
                del self._active_evaluation_splits
            else:
                self._active_evaluation_splits = previous

    def _maybe_evaluate_during_training(
        self,
        splits: list[str],
        *,
        force: bool = False,
        checkpoint: bool = True,
    ) -> bool:
        """Publish one complete validation row at a safe optimizer boundary.

        This is training orchestration; public evaluation methods continue to
        return/log a split without publishing model-selection history.
        """
        interval = self.config.evaluation.interval
        step = self.train_state.global_step
        if not self._training_evaluation_enabled or not splits or self.train_state.last_evaluation_step == step:
            return False
        if not force and (interval <= 0 or step == 0 or step % interval != 0):
            return False

        result = RoundDict()
        with self._evaluation_context(splits):
            for split in splits:
                if self.config.evaluation.steps is None:
                    result[split] = self.evaluate_epoch(split)
                else:
                    result[split] = self.evaluate_steps(split, steps=self.config.evaluation.steps)

        self.append_result(result, index=step)
        self.train_state.last_evaluation_step = step
        if checkpoint:
            self.save_result()
            self._save_evaluation_checkpoint()
        self.supervisor.maybe_handle_termination_signal()
        return True

    def _save_evaluation_checkpoint(self, *, epochs: int | None = None, at_epoch_end: bool = False) -> None:
        # Every collective participant must enter a best-triggered save. The
        # writer still decides whether to publish its own native best alias.
        improved = self.is_best
        if dist.is_available() and dist.is_initialized():
            vote = torch.tensor(float(improved), device=self.all_reduce_device())
            dist.all_reduce(vote, op=dist.ReduceOp.MAX)
            group = None if self.fault_tolerance is None else self.fault_tolerance.replicate_process_group
            if group is not None:
                dist.all_reduce(vote, op=dist.ReduceOp.MAX, group=group)
            improved = bool(vote.item())
        previous = getattr(self, "_checkpoint_from_evaluation", False)
        self._checkpoint_from_evaluation = not at_epoch_end
        try:
            self.save_checkpoint(epochs=epochs, force=improved)
        finally:
            self._checkpoint_from_evaluation = previous

    def train_epochs(
        self,
        train_splits: list[str] | None = None,
        evaluate_splits: list[str] | None = None,
    ) -> RoundDict:
        """
        Run epoch-mode training until `self.epochs` is reached.

        Each epoch runs all train splits, then all evaluation splits, advances
        epoch/metric schedulers, appends and writes results, saves periodic
        checkpoints on `ckpt.interval`, and refreshes the best checkpoint
        whenever the score improves.

        **Called when:** `train` dispatches while `config.epochs` is set, or
        user code explicitly wants epoch-mode semantics.

        Args:
            train_splits: Training splits for each epoch.
            evaluate_splits: Evaluation splits after each epoch.

        Returns:
            Aggregated runner results (`self.results`).

        Raises:
            ValueError: `config.epochs` is not set.
        """
        if train_splits is None:
            train_splits = self.train_splits
        if evaluate_splits is None:
            evaluate_splits = self.evaluate_splits

        total_epochs = self.epochs
        if total_epochs is None:
            raise ValueError("cannot run epoch-mode training: config.epochs is not set")
        # Even a completed resume needs healing and membership before final
        # collective checkpoint/export; no batch is fetched by this handshake.
        self._begin_optimizer_window()
        print(f"train: epoch mode start epoch={self.train_state.epoch} total_epochs={total_epochs}")
        periodic = self._training_evaluation_enabled
        early_stop_counter = 0
        patience = self.patience
        for epoch in range(self.train_state.epoch, total_epochs):
            self.supervisor.maybe_handle_termination_signal()
            self.train_state.epoch = epoch
            result = RoundDict()
            with self._training_evaluation_splits(evaluate_splits):
                for split in train_splits:
                    result[split] = self.train_epoch(split)
                    self.supervisor.maybe_handle_termination_signal()
            if periodic:
                self._maybe_evaluate_during_training(evaluate_splits, force=True, checkpoint=False)
                row = self.results.get(self.train_state.global_step, {})
                result.update({split: row[split] for split in evaluate_splits if split in row})
            else:
                for split in evaluate_splits:
                    result[split] = self.evaluate_epoch(split)
                    self.supervisor.maybe_handle_termination_signal()
            self._step_epoch_scheduler(result)
            index = self.train_state.global_step if periodic else epoch
            self.append_result(result, index=index)
            print(self.format_epoch_result(result, epochs=epoch, total_epochs=total_epochs))
            self.save_result()
            self.train_state.epoch = epoch + 1
            # Call every epoch: the manager persists periodic history on cadence and
            # publishes `best.pth` whenever the score improved this epoch.
            if periodic:
                self._save_evaluation_checkpoint(epochs=epoch, at_epoch_end=True)
            else:
                self.save_checkpoint(epochs=epoch)
            early_stop_counter = 0 if self.is_best else early_stop_counter + 1
            if early_stop_counter > patience:
                print("train: early-stop triggered")
                break
        self.save_checkpoint(last_step=True)
        self.save_model_checkpoint()
        return self.results

    def train_epoch(self, split: str = "train") -> RoundDict:
        """
        Run one full dataloader pass for a training split.

        This is the per-split epoch loop. It sets train mode, resets meters and
        train metrics, manages accumulation-window normalization, invokes
        `train_step` for each micro-batch, emits step logs, and records
        interval/epoch telemetry.

        **Called when:** `train_epochs` processes one train split.

        Args:
            split: Training split name.

        Returns:
            Epoch-level metric mapping for this split.

        **Side effects:** updates optimizer state through `train_step`,
        advances `train_state.global_step` on optimizer flushes, and writes
        step logs.

        !!! danger "Do not"
            - Call this for evaluation data; use `evaluate_epoch`.
            - Override this just to change one batch's forward/loss logic;
              override `train_step`.
            - Manually manage gradient zeroing inside `train_step`; this loop
              and `optimizer_step` own accumulation boundaries.
            - Increment `train_state.epoch`; the surrounding `train_epochs`
              loop owns epoch progress.
            - Save result or checkpoint aliases here; `train_epochs` owns
              epoch-level persistence.

        See Also:
            [`train_steps`][danling.runners.TorchRunner.train_steps]:
                Step-mode counterpart that consumes splits against a global
                step budget instead of one epoch per split.
        """
        loader = self.dataloaders[split]
        loader_length = self._loader_length(loader)
        length = loader_length - 1 if loader_length is not None else None
        last_loss: torch.Tensor | None = None
        last_loss_n: float | None = None
        self._set_loader_epoch(loader, self.train_state.epoch)
        self.mode = RunnerMode.train
        self.split = split
        self.meters.reset()
        self.metrics = self.train_metrics
        if self.metrics is not None:
            self.metrics.reset()
        telemetry = LoopTelemetry(self, start_time=self.loop_time())
        self._reset_accumulation_normalization()
        if self.optimizer_container is not None:
            self.optimizer_container.zero_grad()
        elif self.optimizer is not None:
            self.optimizer.zero_grad()

        for iteration, data, will_flush in self._iter_train_batches(loader):
            self.supervisor.maybe_handle_termination_signal()
            step_before = self.train_state.global_step
            loss_n = self._get_loss_normalizer(data)
            self._pending_loss_normalizer = loss_n
            self._pending_loss_reduction = self._get_loss_reduction(data)
            self._train_window_will_flush = will_flush
            try:
                _, loss = self.train_step(data)
            finally:
                self._train_window_will_flush = False
                self._pending_loss_normalizer = None
                self._pending_loss_reduction = None

            loss = self._loss_for_reporting(loss, loss_n, data)
            self._train_checkpoint_ready = will_flush
            self.supervisor.mark_heartbeat_progress()
            self.supervisor.maybe_handle_termination_signal()
            current_time = self.loop_time()
            if self.scheduler is not None and hasattr(self.scheduler, "get_last_lr"):
                self.meters.lr.update(self.scheduler.get_last_lr()[0])
            if loss is not None:
                self._update_loss_meter(loss, loss_n)
            telemetry.observe(iteration=iteration, data=data, current_time=current_time)

            if self.train_state.global_step > step_before:
                splits = getattr(self, "_active_evaluation_splits", self.evaluate_splits)
                self._maybe_evaluate_during_training(splits)

            if self.log_interval > 0 and (
                (iteration > 0 and iteration % self.log_interval == 0) or iteration == length
            ):
                telemetry.emit_log(split=split, iteration=iteration, length=length, loss=loss, loss_n=loss_n)
            last_loss = loss
            last_loss_n = loss_n

        if (
            length is None
            and self.log_interval > 0
            and telemetry.last_iteration is not None
            and telemetry.last_iteration != telemetry.last_print_iteration
        ):
            assert telemetry.last_iteration is not None
            telemetry.emit_log(
                split=split,
                iteration=telemetry.last_iteration,
                length=length,
                loss=last_loss,
                loss_n=last_loss_n,
                reset_peak_stats=False,
            )
        result = self.get_epoch_result()
        telemetry.finalize_result(result, elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time)
        return result

    def train_steps(
        self,
        train_splits: list[str] | None = None,
        evaluate_splits: list[str] | None = None,
    ) -> RoundDict:
        """
        Run step-mode training for the configured global step budget.

        Step mode consumes train splits in sorted split order until
        `train_state.global_step >= self.steps`, then optionally evaluates
        configured evaluation splits with `evaluate_steps`.

        **Called when:** `train` dispatches while `config.epochs` is unset, or
        user code explicitly wants a global-step budget.

        Args:
            train_splits: Training splits to consume in order.
            evaluate_splits: Evaluation splits to run after training steps finish.

        Returns:
            Aggregated runner results (`self.results`).

        Raises:
            ValueError: total step budget cannot be resolved.

        **Side effects:** updates epoch as an outer split-round counter,
        appends one result row indexed by `global_step`, writes result files,
        and saves the final checkpoint.

        !!! danger "Do not"
            - Assume a split is consumed exactly once; step mode can resume a
              split iterator across outer rounds.
            - Mutate `train_state.global_step` outside optimizer stepping.

        See Also:
            [`train_epoch`][danling.runners.TorchRunner.train_epoch]:
                Per-split epoch loop used by epoch-mode training.
        """
        if train_splits is None:
            train_splits = self.train_splits
        if evaluate_splits is None:
            evaluate_splits = self.evaluate_splits

        total_steps = self.steps
        if total_steps is None:
            raise ValueError("cannot run step-mode training: config.steps could not be resolved")
        self._begin_optimizer_window()
        print(f"train: step mode start global_step={self.train_state.global_step} steps={total_steps}")
        periodic = self._training_evaluation_enabled
        result = RoundDict()
        step_mode_iterators: dict[str, Iterator[tuple[int, Any, bool]] | None] = dict.fromkeys(train_splits)
        step_mode_sampler_epochs = self.train_state.sampler_epochs
        for split in train_splits:
            step_mode_sampler_epochs.setdefault(split, self.train_state.epoch)
        while self.train_state.global_step < total_steps:
            self.supervisor.maybe_handle_termination_signal()
            round_start_step = self.train_state.global_step
            plan_epoch = self.train_state.epoch
            restart_round = False
            round_finished_restored_split = False
            round_result = RoundDict()
            total_train_splits = len(train_splits)
            for split_index, split in enumerate(train_splits):
                if split_index < self.train_state.step_mode_split_index:
                    continue
                self.supervisor.maybe_handle_termination_signal()
                self.mode = RunnerMode.train
                self.split = split
                remaining = total_steps - self.train_state.global_step
                if remaining <= 0:
                    break
                loader = self.dataloaders[split]
                remaining_splits = total_train_splits - split_index
                split_steps = self._step_mode_split_budget(
                    remaining_steps=remaining,
                    remaining_splits=remaining_splits,
                    loader=loader,
                )
                if split_steps <= 0:
                    break
                start_global_step = self.train_state.global_step
                target_global_step = self.train_state.step_mode_target or start_global_step + split_steps
                # A checkpoint can be saved by the last update of a split,
                # before the loop clears its target and advances the round.
                round_finished_restored_split |= target_global_step <= start_global_step
                self.train_state.step_mode_split_index = split_index
                self.train_state.step_mode_target = target_global_step
                self.meters.reset()
                self.metrics = self.train_metrics
                if self.metrics is not None:
                    self.metrics.reset()
                telemetry = LoopTelemetry(self, start_time=self.loop_time())
                self._reset_accumulation_normalization()
                if self.optimizer_container is not None:
                    self.optimizer_container.zero_grad()
                elif self.optimizer is not None:
                    self.optimizer.zero_grad()
                checkpoint_cadence = self.checkpoint_interval
                batch_iteration = -1

                while self.train_state.global_step < target_global_step:
                    self.supervisor.maybe_handle_termination_signal()
                    generation = getattr(self, "_dataloader_state_generation", 0)
                    # Quorum can restore newer progress as well as model and
                    # loader state. Revalidate this plan before any iterator
                    # creation or fetch; the generator's begin is idempotent.
                    self._begin_optimizer_window()
                    if getattr(self, "_dataloader_state_generation", 0) != generation:
                        # Healing can reload only some local ranks, even when
                        # their shared training plan is unchanged. Invalidate
                        # local loader caches without changing collective loop
                        # ordering on the ranks that did not need recovery.
                        step_mode_iterators = dict.fromkeys(train_splits)
                        step_mode_sampler_epochs = self.train_state.sampler_epochs
                        for train_split in train_splits:
                            step_mode_sampler_epochs.setdefault(train_split, self.train_state.epoch)
                    if (
                        self.train_state.global_step >= total_steps
                        or self.train_state.epoch != plan_epoch
                        or self.train_state.step_mode_split_index != split_index
                        or self.train_state.step_mode_target != target_global_step
                    ):
                        restart_round = True
                        break
                    if self.train_state.global_step >= target_global_step:
                        break
                    batch: tuple[int, Any, bool] | None = None
                    iterator = step_mode_iterators[split]
                    recreated = 0
                    while True:
                        if iterator is None:
                            # A restored cursor can already be at the end of
                            # its pass before this process creates an iterator.
                            # Allow that pass to finish, then try one fresh pass.
                            if recreated >= 2:
                                break
                            self._set_loader_epoch(loader, step_mode_sampler_epochs[split])
                            iterator = self._iter_train_batches(loader)
                            step_mode_iterators[split] = iterator
                            recreated += 1
                        try:
                            batch = next(iterator)
                            break
                        except StopIteration:
                            iterator = None
                            step_mode_iterators[split] = None
                            step_mode_sampler_epochs[split] += 1
                    if batch is None:
                        break
                    _, data, will_flush = batch
                    batch_iteration += 1
                    self.supervisor.maybe_handle_termination_signal()
                    step_before = self.train_state.global_step
                    loss_n = self._get_loss_normalizer(data)
                    self._pending_loss_normalizer = loss_n
                    self._pending_loss_reduction = self._get_loss_reduction(data)
                    self._train_window_will_flush = will_flush
                    try:
                        _, loss = self.train_step(data)
                    finally:
                        self._train_window_will_flush = False
                        self._pending_loss_normalizer = None
                        self._pending_loss_reduction = None

                    loss = self._loss_for_reporting(loss, loss_n, data)
                    self._train_checkpoint_ready = will_flush
                    self.supervisor.mark_heartbeat_progress()
                    self.supervisor.maybe_handle_termination_signal()
                    current_time = self.loop_time()
                    if self.scheduler is not None and hasattr(self.scheduler, "get_last_lr"):
                        self.meters.lr.update(self.scheduler.get_last_lr()[0])
                    if loss is not None:
                        self._update_loss_meter(loss, loss_n)
                    telemetry.observe(iteration=batch_iteration, data=data, current_time=current_time)

                    step_after = self.train_state.global_step
                    evaluated = step_after > step_before and self._maybe_evaluate_during_training(evaluate_splits)
                    if (
                        not evaluated
                        and checkpoint_cadence > 0
                        and step_after != step_before
                        and step_after % checkpoint_cadence == 0
                    ):
                        self.save_checkpoint(save_best=False) if periodic else self.save_checkpoint()

                    global_step = step_after if step_after != step_before else None
                    is_log_step = (
                        self.log_interval > 0
                        and global_step is not None
                        and global_step > 0
                        and global_step % self.log_interval == 0
                    )
                    is_boundary_step = global_step in (target_global_step, total_steps)
                    if self.log_interval > 0 and global_step is not None and (is_log_step or is_boundary_step):
                        telemetry.emit_log(
                            split=split,
                            iteration=batch_iteration,
                            length=total_steps,
                            loss=loss,
                            loss_n=loss_n,
                            display_iteration=global_step,
                        )

                if restart_round:
                    break
                round_result[split] = self.get_epoch_result()
                telemetry.finalize_result(
                    round_result[split], elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time
                )
                self.train_state.step_mode_split_index = split_index + 1
                self.train_state.step_mode_target = 0
                self.supervisor.maybe_handle_termination_signal()

            if restart_round:
                step_mode_iterators = dict.fromkeys(train_splits)
                step_mode_sampler_epochs = self.train_state.sampler_epochs
                for split in train_splits:
                    step_mode_sampler_epochs.setdefault(split, self.train_state.epoch)
                continue
            if self.train_state.global_step == round_start_step and not round_finished_restored_split:
                remaining_steps = total_steps - self.train_state.global_step
                warn(
                    f"step-mode training made no progress after one full split pass "
                    f"(target={total_steps}, reached={self.train_state.global_step}, remaining={remaining_steps})",
                    RuntimeWarning,
                    stacklevel=2,
                )
                break
            self._step_epoch_scheduler(round_result)
            result = round_result
            self.train_state.epoch += 1
            self.train_state.step_mode_split_index = 0
            self.train_state.step_mode_target = 0
        remaining_steps = total_steps - self.train_state.global_step
        if remaining_steps > 0:
            warn(
                f"step-mode training finished with {remaining_steps} step(s) remaining "
                f"(target={total_steps}, reached={self.train_state.global_step})",
                RuntimeWarning,
                stacklevel=2,
            )
        if periodic:
            self._maybe_evaluate_during_training(evaluate_splits, force=True, checkpoint=False)
            row = self.results.get(self.train_state.global_step, {})
            result.update({split: row[split] for split in evaluate_splits if split in row})
        else:
            for split in evaluate_splits:
                result[split] = self.evaluate_steps(split=split)
        self.append_result(result, index=self.train_state.global_step)
        self.save_result()
        self.save_checkpoint(last_step=True)
        self.save_model_checkpoint()
        return self.results

    def evaluate_step(self, data: Any) -> tuple[Any, torch.Tensor | None]:
        """
        Run one evaluation micro-step.

        The default implementation runs forward → optional loss → optional
        metric update under `infer_context()`. No backward pass and no optimizer step.

        **Called when:** once per micro-batch by `evaluate_epoch`/`evaluate_steps`,
        which run under `torch.inference_mode()`.

        **Precondition:** at least one of `self.model` or `self.ema` is bound.
        `self.mode == RunnerMode.evaluate`. The default prefers `self.ema` over
        `self.model` when both are available.

        Args:
            data: One micro-batch. The default unpacks `data["input"]` /
                `data.get("target")` for mappings, `(data[0], data[1])` for
                non-string sequences, and `(data, None)` otherwise. Override
                `evaluate_step` if your batch shape differs.

        Returns:
            `(pred, loss)`. `pred` is the model output (used by `metrics.update`).
            `loss` is the scalar loss returned to the caller for reduced
            logging, or `None` when no `criterion` is set.

        Raises:
            ValueError: neither `self.model` nor `self.ema` is initialized.

        **Side effects:** moves `data` to `self.device`, runs forward through
        `self.ema or self.model` under `infer_context()`, computes loss when
        `criterion` is set, and updates `self.metrics` when bound.

        !!! danger "Do not"
            - Call `self.backward(...)` or `self.step()` (no optimizer here).
            - Mutate `train_state.global_step` or `train_state.micro_step`.
            - Switch the runner mode (the loop owns `self.mode`).
            - Call `save_checkpoint()` (cadence is owned by training loops only).

        **Backend notes:**

        - `ParallelRunner` overrides this method when a pipeline schedule is
          set; the schedule owns micro-batching and pipeline-stage loss
          reduction.
        """
        data = self.to_device(data)
        if isinstance(data, Mapping):
            inputs = data["input"]
            target = data.get("target")
        elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
            inputs = data[0]
            target = data[1] if len(data) > 1 else None
        else:
            inputs = data
            target = None

        if self.model is None and self.ema is None:
            raise ValueError("cannot run evaluate_step: model is not initialized")
        model = self.ema or self.model
        with self.infer_context():
            pred = model(**inputs) if isinstance(inputs, Mapping) else model(inputs)
            loss = self._compute_loss(pred, target)

        if self.metrics is not None and pred is not None and target is not None:
            self.metrics.update(pred, target)

        return pred, loss

    def evaluate(self, evaluate_splits: list[str] | None = None) -> RoundDict:
        """
        Run evaluation across splits with epoch-mode semantics.

        **Called when:** user code explicitly evaluates a runner, or training
        code delegates to evaluation helpers.

        Args:
            evaluate_splits: Optional evaluation splits. When `None`, use `self.evaluate_splits`.

        Returns:
            Mapping of split -> evaluation result for this call.

        Raises:
            ValueError: no valid evaluation split can be resolved.

        **Side effects:** sets evaluation mode per split, prints a formatted
        aggregate result, and writes scalar outputs through `evaluate_epoch`.
        """

        evaluate_splits = self._resolve_requested_splits(evaluate_splits, self.evaluate_splits, kind="evaluation")
        if not evaluate_splits:
            raise ValueError("cannot start evaluation: no valid evaluation split was resolved")
        print("evaluate: start")
        print(f"evaluate: splits={evaluate_splits}")
        result = RoundDict()
        for split in evaluate_splits:
            result[split] = self.evaluate_epoch(split=split)
        display_epoch = self.train_state.epoch
        if self.epochs is not None and display_epoch > 0:
            display_epoch -= 1
        print(self.format_epoch_result(result, epochs=display_epoch))
        return result

    @torch.inference_mode()
    def evaluate_epoch(self, split: str = "val") -> RoundDict:
        """
        Run one full dataloader pass for an evaluation split.

        Sets evaluation mode, resets meters/evaluation metrics, runs
        `evaluate_step` for every batch under inference mode, emits step logs,
        and writes the split result at the current epoch index.

        **Called when:** `evaluate` or `train_epochs` evaluates a split.

        Args:
            split: Evaluation split name.

        Returns:
            Epoch-level metric mapping for this split.

        **Side effects:** updates evaluation meters/metrics, emits logs, writes
        scalar results, and records telemetry. It does not update optimizer or
        training progress counters.
        """
        loader = self.dataloaders[split]
        loader_length = self._loader_length(loader)
        length = loader_length - 1 if loader_length is not None else None

        last_loss: torch.Tensor | None = None
        last_loss_n: float | None = None
        self.mode = RunnerMode.evaluate
        self.split = split
        self.meters = AverageMeters()
        self.metrics = self.evaluate_metrics
        if self.metrics is not None:
            self.metrics.reset()
        telemetry = LoopTelemetry(self, start_time=self.loop_time())
        consumed = 0
        for iteration, data in enumerate(loader):
            consumed = iteration + 1
            self.supervisor.maybe_handle_termination_signal()
            loss_n = self._get_loss_normalizer(data)
            if loss_n is not None and loss_n < 0:
                loss_n = None
            _, loss = self.evaluate_step(data)
            loss = self._loss_for_reporting(loss, loss_n, data)
            self.supervisor.mark_heartbeat_progress()
            self.supervisor.maybe_handle_termination_signal()
            current_time = self.loop_time()
            if loss is not None:
                self._update_loss_meter(loss, loss_n)
            telemetry.observe(
                iteration=iteration,
                data=data,
                current_time=current_time,
            )
            self.supervisor.maybe_collect_garbage(iteration + 1, scope=f"evaluate:{split}")

            if self.log_interval > 0 and (
                (iteration > 0 and iteration % self.log_interval == 0) or iteration == length
            ):
                telemetry.emit_log(split=split, iteration=iteration, length=length, loss=loss, loss_n=loss_n)
            last_loss = loss
            last_loss_n = loss_n

        if (
            length is None
            and self.log_interval > 0
            and telemetry.last_iteration is not None
            and telemetry.last_iteration != telemetry.last_print_iteration
        ):
            assert telemetry.last_iteration is not None
            telemetry.emit_log(
                split=split,
                iteration=telemetry.last_iteration,
                length=length,
                loss=last_loss,
                loss_n=last_loss_n,
                reset_peak_stats=False,
            )
        result = self.get_epoch_result()
        telemetry.finalize_result(
            result, elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time, steps=consumed
        )
        self.write_result(result, split, getattr(self, "_evaluation_result_index", self.train_state.epoch))
        return result

    @torch.inference_mode()
    def evaluate_steps(self, split: str = "val", steps: int | None = None) -> RoundDict:
        """
        Run bounded evaluation steps on one split.

        Used by step-mode training to evaluate a small fixed number of batches
        without requiring a full evaluation pass.

        **Called when:** `train_steps` evaluates configured splits after the
        step budget finishes, or user code requests bounded evaluation.

        Args:
            split: Evaluation split name.
            steps: Number of batches to evaluate. When `None`, defaults to `max(self.steps // 20, 1)`.

        Returns:
            Step-bounded evaluation metrics.

        Raises:
            ValueError: step budget cannot be inferred, `steps` is negative, or
                the dataloader exhausts before the requested number of steps.

        **Side effects:** writes scalar results at `train_state.global_step`.
        """
        if steps is None:
            total_steps = self.steps
            if total_steps is None:
                raise ValueError("cannot infer evaluation steps: step budget is unavailable; pass `steps`")
            steps = max(total_steps // 20, 1)
        if steps < 0:
            raise ValueError(f"invalid steps: expected a non-negative value, got {steps}")
        loader = self.dataloaders[split]
        length = steps - 1

        self.mode = RunnerMode.evaluate
        self.split = split
        if steps == 0:
            self.meters = AverageMeters()
            self.metrics = self.evaluate_metrics
            if self.metrics is not None:
                self.metrics.reset()
            result = self.get_epoch_result()
            self.write_result(result, split, self.train_state.global_step)
            return result

        self.meters = AverageMeters()
        self.metrics = self.evaluate_metrics
        if self.metrics is not None:
            self.metrics.reset()
        telemetry = LoopTelemetry(self, start_time=self.loop_time())
        consumed = 0
        batches = loader if steps is None else islice(loader, steps)
        for iteration, data in enumerate(batches):
            consumed = iteration + 1
            self.supervisor.maybe_handle_termination_signal()
            loss_n = self._get_loss_normalizer(data)
            if loss_n is not None and loss_n < 0:
                loss_n = None
            _, loss = self.evaluate_step(data)
            loss = self._loss_for_reporting(loss, loss_n, data)
            self.supervisor.mark_heartbeat_progress()
            self.supervisor.maybe_handle_termination_signal()
            current_time = self.loop_time()
            if loss is not None:
                self._update_loss_meter(loss, loss_n)
            telemetry.observe(iteration=iteration, data=data, current_time=current_time)
            self.supervisor.maybe_collect_garbage(iteration + 1, scope=f"evaluate:{split}")

            if self.log_interval > 0 and (
                (iteration > 0 and iteration % self.log_interval == 0) or iteration == length
            ):
                telemetry.emit_log(split=split, iteration=iteration, length=length, loss=loss, loss_n=loss_n)

        if steps is not None and consumed < steps:
            raise ValueError(
                f"evaluate steps exhausted early on split '{split}': requested {steps} step(s), got {consumed}"
            )
        result = self.get_epoch_result()
        telemetry.finalize_result(
            result, elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time, steps=consumed
        )
        self.write_result(result, split, self.train_state.global_step)
        return result

    @torch.inference_mode()
    def infer_step(self, data: Any) -> list[float]:
        """
        Run one inference micro-step.

        The default implementation runs forward through `self.ema or self.model`,
        detaches scalar-per-example predictions, squeezes the trailing
        dimension, moves them to CPU, and returns them as a Python list.

        **Called when:** once per micro-batch by `infer`/`_iter_infer_batches`.
        The method is decorated with `torch.inference_mode()`.

        **Precondition:** at least one of `self.model` or `self.ema` is bound.
        `self.mode == RunnerMode.infer`.

        Args:
            data: One micro-batch. The default unpacks `data["input"]` for
                mappings, `data[0]` for non-string sequences, and `data`
                itself otherwise. Override `infer_step` if your batch shape
                differs or you need to pass auxiliary tensors to the model.

        Returns:
            List of CPU floats for scalar-per-example predictions. The
            default converts with `pred.squeeze(-1).detach().cpu().tolist()`.
            Override if your model emits multi-dim tensors, mappings, or
            non-numeric outputs.

        Raises:
            ValueError: neither `self.model` nor `self.ema` is initialized.

        **Side effects:** moves `data` to `self.device`, runs forward through
        `self.ema or self.model` under `infer_context()`, then converts the
        output to a CPU list.

        !!! danger "Do not"
            - Compute or accumulate metrics (inference is metric-free).
            - Mutate runner state counters.
            - Return a `torch.Tensor` (callers expect `list[float]` for
              batched aggregation and streaming).
            - Call `self.backward(...)` or `self.step()`.

        **Backend notes:**

        - `ParallelRunner` overrides this method when a pipeline schedule is
          set; non-first-stage ranks pass `data=None` and the schedule routes
          activations through pipeline communication.
        """
        data = self.to_device(data)
        if isinstance(data, Mapping):
            inputs = data["input"]
        elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
            inputs = data[0]
        else:
            inputs = data

        if self.model is None and self.ema is None:
            raise ValueError("cannot run infer_step: model is not initialized")
        model = self.ema or self.model
        with self.infer_context():
            pred = model(**inputs) if isinstance(inputs, Mapping) else model(inputs)
        values = pred.squeeze(-1).detach().cpu().tolist()
        if isinstance(values, list):
            return values
        return [float(values)]

    def infer(
        self,
        split: str = "infer",
        *,
        steps: int | None = None,
        stream: bool | None = None,
    ) -> list[float] | Iterator[list[float]]:
        """
        Run inference on one split.

        In non-stream mode this consumes all requested batches and returns a
        flattened Python list. In stream mode it returns an iterator of
        per-batch outputs and leaves consumption to the caller.

        **Called when:** user code requests prediction-only execution.

        Args:
            split: Inference split name.
            steps: Optional max number of batches to consume.
            stream: `True` returns a generator of per-batch outputs, `False` returns a flattened list.
                When `None`, stream only for unsized loaders without explicit `steps`.

        Returns:
            Flattened predictions or a streaming iterator of batch predictions.

        Raises:
            ValueError: `steps` is negative, or non-stream inference is
                requested for an unsized loader without an explicit step count.

        **Side effects:** sets inference mode/split. It does not update metrics
        or optimizer state.
        """

        self.mode = RunnerMode.infer
        self.split = split
        loader = self.dataloaders[split]
        if steps is not None and steps < 0:
            raise ValueError(f"invalid steps: expected a non-negative value, got {steps}")

        loader_length = self._loader_length(loader)
        if stream is None:
            stream = steps is None and loader_length is None

        if not stream and loader_length is None and steps is None:
            raise ValueError("infer with stream=False requires `steps` for unsized loaders")

        iterator = self._iter_infer_batches(loader, steps=steps, split=split)
        if stream:
            return iterator

        total = steps if steps is not None else loader_length
        output: list = []
        for values in tqdm(iterator, total=total, disable=self.distributed and not self.is_main_process):
            output.extend(values)
        if self.distributed:
            output = self._gather_infer_predictions(loader, output)
        return output

    def _gather_infer_predictions(self, loader: Any, output: list) -> list:
        """
        Gather sharded inference outputs to every rank and restore dataset order.

        Under DDP the loader's `DistributedSampler` shards the split, so each rank's default `infer_step`
        produces one prediction per sampled item. Pair each local output with its global sample index,
        all-gather the `(index, output)` pairs across ranks, drop the sampler's duplicate padding (same
        index seen twice -> identical output), and return the outputs in dataset order so every rank sees
        the whole split.
        """
        import torch.distributed as dist

        sampler = getattr(loader, "sampler", None)
        if sampler is None or not dist.is_available() or not dist.is_initialized():
            return output
        indices = list(sampler)
        if len(output) > len(indices):
            raise ValueError(
                "cannot gather distributed inference output: infer_step produced more predictions "
                f"({len(output)}) than the local sampler has indices ({len(indices)})"
            )
        paired = list(zip(indices[: len(output)], output))
        gathered: list = [None] * self.world_size
        dist.all_gather_object(gathered, paired)
        merged: dict = {}
        for rank_pairs in gathered:
            for index, value in rank_pairs:
                merged[index] = value
        return [merged[index] for index in sorted(merged)]

    def _iter_infer_batches(self, loader: Any, *, steps: int | None, split: str) -> Iterator[list[float]]:
        batches = loader if steps is None else islice(loader, steps)
        for iteration, data in enumerate(batches):
            values = self.infer_step(data)
            self.supervisor.mark_heartbeat_progress()
            yield values
            self.supervisor.maybe_collect_garbage(iteration + 1, scope=f"infer:{split}")

    def _export_checkpoint_metadata(self, cls: type = dict) -> Mapping[str, Any]:
        return cls()

    def _export_checkpoint_components(self, cls: type = dict) -> Mapping[str, Any]:
        if self.model is None:
            raise ValueError("cannot build checkpoint state: model is not initialized")
        state = cls()
        state["ema"] = self.ema.state_dict() if self.ema else None
        state["grad_scaler"] = self.grad_scaler.state_dict() if self.grad_scaler is not None else None
        state["optimizer"] = self.optimizer.state_dict() if self.optimizer else None
        state["scheduler"] = self.scheduler.state_dict() if self.scheduler else None
        state["model"] = self.unwrap(self.model).state_dict()
        return state

    def _checkpoint_rng_state(self) -> Mapping[str, Any]:
        rng = dict(super()._checkpoint_rng_state())
        rng["torch_cpu"] = torch.get_rng_state()
        rng["torch_cuda"] = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
        training = getattr(self, "_evaluation_training_rng", None)
        if training is not None:
            for name in ("python", "numpy", "torch_cpu"):
                rng[name] = training[name]
            if rng["torch_cuda"] is not None:
                for device, state in training["cuda"].items():
                    rng["torch_cuda"][device] = state
        return rng

    def _checkpoint_dataloader_state(self) -> Mapping[str, Any]:
        training = getattr(self, "_evaluation_training_dataloaders", None)
        return super()._checkpoint_dataloader_state() if training is None else training

    def _export_checkpoint_local_state(self) -> Mapping[str, Any]:
        """Snapshot replica-local resume state without gathering model/optimizer weights."""
        state = BaseRunner.state_dict(self)
        local = {"state": {"rng": state["state"]["rng"]}, "dataloaders": state["dataloaders"]}
        fault_tolerance = self.fault_tolerance
        if fault_tolerance is not None and fault_tolerance.enabled:
            state_dict = getattr(fault_tolerance, "state_dict", None)
            if state_dict is not None:
                local["fault_tolerance"] = state_dict()
        return local

    def state_dict(self, cls: type = dict) -> Mapping:
        """
        Return the TorchRunner checkpoint payload.

        Extends `BaseRunner.state_dict` with backend metadata plus EMA,
        optimizer, scheduler, and unwrapped model state.

        **Called when:** checkpoint managers persist a TorchRunner checkpoint.

        Args:
            cls: Mapping factory used for nested payloads.

        Returns:
            Mapping containing base runner state and torch component state.

        **Side effects:** snapshots Python/NumPy/Torch RNG state before
        exporting.
        """
        state = cls(super().state_dict(cls))
        state.update(self._export_checkpoint_metadata(cls))
        state.update(self._export_checkpoint_components(cls))
        if self.fault_tolerance is not None and self.fault_tolerance.enabled:
            state["fault_tolerance"] = self.fault_tolerance.state_dict()
        return state

    def _restore_model_checkpoint(self, state_dict: Mapping[str, Any], *args, **kwargs) -> None:
        if self.model is None:
            raise ValueError("cannot load model weights: model is not initialized")
        self.unwrap(self.model).load_state_dict(state_dict, *args, **kwargs)

    def load_model(self, state_dict: Mapping[str, Any], *args, **kwargs) -> None:
        self._restore_model_checkpoint(state_dict, *args, **kwargs)

    def _restore_optimizer_checkpoint(self, state_dict: Mapping[str, Any], *args, **kwargs) -> None:
        if self.optimizer is None:
            return
        self.optimizer.load_state_dict(dict(state_dict), *args, **kwargs)

    def _checkpoint_optimizer_load_template(self) -> Mapping[str, Any] | None:
        """Build native optimizer storage only when reading a trained DCP state."""
        if self.optimizer is None:
            return None
        TorchDistributedCheckpointManager.export_model_optimizer_state(
            model=self.unwrap(self.model), optimizer=self.optimizer, initialize_optimizer_state=True
        )
        return self.optimizer.state_dict()

    def load_optimizer(self, state_dict: Mapping[str, Any] | None, *args, **kwargs) -> None:
        if self.optimizer is None:
            return
        optimizer_state = self._require_checkpoint_component_state("optimizer", state_dict)
        self._restore_optimizer_checkpoint(optimizer_state, *args, **kwargs)

    def load_state_dict(self, checkpoint: Mapping[str, Any]) -> None:
        """
        Restore base runner state plus Torch RNG state.

        Model, optimizer, scheduler, and dataloader components are restored by
        `load_checkpoint`; this method owns only runner/RNG state.
        """
        super().load_state_dict(checkpoint)
        if self.fault_tolerance is not None and "fault_tolerance" in checkpoint:
            self.fault_tolerance.load_state_dict(checkpoint["fault_tolerance"])
        if self.grad_scaler is not None:
            scaler_state = checkpoint.get("grad_scaler")
            if scaler_state is not None:
                self.grad_scaler.load_state_dict(self._require_checkpoint_component_state("grad_scaler", scaler_state))
        state_dict = checkpoint.get("state") or {}
        rng_state = state_dict.get("rng")
        if isinstance(rng_state, Mapping) and "torch_cpu" in rng_state and self.rng_state.torch_cpu is not None:
            torch.set_rng_state(self.rng_state.torch_cpu)
        if (
            torch.cuda.is_available()
            and isinstance(rng_state, Mapping)
            and "torch_cuda" in rng_state
            and self.rng_state.torch_cuda is not None
        ):
            torch.cuda.set_rng_state_all(self.rng_state.torch_cuda)

    def load_checkpoint(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> None:
        """
        Load a full checkpoint and rebind optimizer/scheduler helpers.

        This delegates component restore to `BaseRunner.load_checkpoint`, then
        rebuilds the `OptimizerContainer` so scheduler and optimizer state stay
        bound after restore.
        """
        super().load_checkpoint(checkpoint, *args, **kwargs)
        self._bind_optimizer_container()

    # `save_checkpoint` is inherited from `BaseRunner`; collective vs main-only
    # dispatch is owned by `checkpoint_manager.is_collective`.

    def load_pretrained(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> None:
        if not isinstance(checkpoint, Mapping) and str(self.config.get("ckpt.backend")).lower() == "dcp":
            checkpoint_path = os.fsdecode(checkpoint)
            if checkpoint_path.endswith(".pointer") or not os.path.isfile(checkpoint_path):
                ckpt = self.checkpoint_manager.load_model_checkpoint(checkpoint)
                if ckpt.get("ema") is not None:
                    self.load_model(ckpt["ema"], *args, **kwargs)
                elif "model" in ckpt:
                    self.load_model(ckpt["model"], *args, **kwargs)
                elif "model_parts" in ckpt:
                    self.load_model(ckpt["model_parts"], *args, **kwargs)
                else:
                    raise ValueError(
                        "cannot load pretrained weights: checkpoint has no EMA or model state\n"
                        "Use `load_checkpoint` for full checkpoint restore instead of `load_pretrained`"
                    )
                self.config.pretrained = os.fsdecode(checkpoint)
                self.log_restore_summary(
                    kind="pretrained",
                    source=checkpoint,
                    optimizer="skipped",
                    scheduler="skipped",
                )
                return
        super().load_pretrained(checkpoint, *args, **kwargs)

    def read_checkpoint(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> Mapping[str, Any]:
        """Read checkpoint payload from mapping/file/DCP directory input."""
        if isinstance(checkpoint, Mapping):
            return checkpoint

        checkpoint_path = os.fsdecode(checkpoint)
        if str(self.config.get("ckpt.backend")).lower() == "dcp" and (
            checkpoint_path.endswith(".pointer") or not os.path.isfile(checkpoint_path)
        ):
            return self.checkpoint_manager.load_checkpoint(checkpoint)
        return super().read_checkpoint(checkpoint, *args, **kwargs)

    @classmethod
    def read_config(
        cls,
        checkpoint: Mapping | bytes | str | os.PathLike,
        *args,
        **kwargs,
    ) -> RunnerConfig:
        """Read runner config from checkpoint payload, including DCP directory inputs."""
        if isinstance(checkpoint, Mapping):
            return super().read_config(checkpoint, *args, **kwargs)

        if TorchDistributedCheckpointManager.is_checkpoint_path(checkpoint):
            return TorchDistributedCheckpointManager.read_config(checkpoint)

        return super().read_config(checkpoint, *args, **kwargs)

    @property
    def device(self) -> torch.device:
        """Use the configured device, falling back to local CUDA or CPU selection."""
        configured_device = self.config.get("device")
        if configured_device is not None:
            device = torch.device(configured_device)
            if device.type == "cuda" and device.index is None:
                return torch.device("cuda", self.local_rank)
            return device
        if torch.cuda.is_available():
            return torch.device("cuda", self.local_rank)
        return torch.device("cpu")

    @property
    def mode(self) -> RunnerMode:
        return self._mode

    @mode.setter
    def mode(self, mode: str | RunnerMode) -> None:
        if isinstance(mode, str):
            mode = RunnerMode(mode)
        if getattr(self, "_mode", None) == mode:
            return
        self._mode = mode

        is_train = mode == RunnerMode.train
        model_parts = getattr(self, "model_parts", None)
        if isinstance(model_parts, Sequence) and model_parts:
            for model_part in model_parts:
                if not isinstance(model_part, nn.Module):
                    continue
                model_part.train(is_train)
        elif self.model is not None:
            self.model.train(is_train)
        if self.ema is not None:
            self.ema.train(is_train)

    @property
    def rank(self) -> int:
        if dist.is_available() and dist.is_initialized():
            return dist.get_rank()
        return int(os.getenv("RANK", "0"))

    @property
    def world_size(self) -> int:
        r"""
        Number of Processes.
        """
        if dist.is_available() and dist.is_initialized():
            return dist.get_world_size()
        return int(os.getenv("WORLD_SIZE", "1"))

    @property
    def distributed(self) -> bool:
        if self.world_size > 1:
            return True
        runtime = self.fault_tolerance
        return runtime is not None and runtime.enabled and runtime.group_size > 1

    def close(self, timeout: float | None = None) -> bool:
        """Close runner resources."""
        try:
            self._close_profiler()
            drained = super().close(timeout=timeout)
        except Exception:
            self._close_profiler()
            if self.diagnostics is not None:
                self.diagnostics.close()
            self.destroy_process_group()
            raise
        if not drained:
            return False
        if self.diagnostics is not None:
            self.diagnostics.close()
        self.destroy_process_group()
        return drained

device property

Python

device: device

Use the configured device, falling back to local CUDA or CPU selection.

world_size property

Python

world_size: int

Number of Processes.

init_distributed

Python

init_distributed() -> None

Initialize the distributed environment.

The default implementation initializes the default torch.distributed process group from WORLD_SIZE/RANK/LOCAL_RANK environment variables when WORLD_SIZE > 1 or TorchFT is enabled, sets the selected CUDA device, broadcasts self.timestamp from rank 0, and seeds elastic_state.restart_count from TORCHELASTIC_RESTART_COUNT.

Called when: once during BaseRunner.__init__, before init_checkpoint_manager, init_fault_tolerance, and init_garbage_collection. The runner is partially constructed at this point — self.config, self.workspace, self.timestamp, the dataloader container, and the default FileCheckpointManager are bound, but the model is not materialized and optimizers/dataloaders are not built.

Precondition: environment variables WORLD_SIZE, RANK, LOCAL_RANK are set when running distributed. The default torch.distributed process group is not already initialized when WORLD_SIZE > 1 or TorchFT is enabled — the runner owns process-group lifecycle.

Raises:

Type	Description
RuntimeError	the default process group is already initialized when WORLD_SIZE > 1 or TorchFT is enabled.
ValueError	dist.init_timeout_seconds is non-positive.

Side effects: when WORLD_SIZE > 1 or TorchFT is enabled, calls dist.init_process_group(...), sets the active CUDA device when the runtime device is CUDA, and broadcasts self.timestamp from rank 0. Reads TORCHELASTIC_RESTART_COUNT into elastic_state.restart_count.

Do not

• Initialize a process group via dist.init_process_group(...) outside the runner; the runner owns its lifecycle.
• Build the model or dataloaders here; those happen in __post_init__.
• Bind the checkpoint manager here; init_checkpoint_manager runs next.

Backend notes:

• ParallelRunner extends this hook: after calling super(), it builds the parallel topology (build_topology) and initializes per-axis process groups via init_device_mesh.
• DeepSpeedRunner inherits the default; DeepSpeed reuses the default process group initialized here.

Source code in danling/runners/torch_runner.py

def init_distributed(self) -> None:
    r"""
    Initialize the distributed environment.

    The default implementation initializes the default torch.distributed
    process group from `WORLD_SIZE`/`RANK`/`LOCAL_RANK` environment
    variables when `WORLD_SIZE > 1` or TorchFT is enabled, sets the selected CUDA device,
    broadcasts `self.timestamp` from rank 0, and seeds
    `elastic_state.restart_count` from `TORCHELASTIC_RESTART_COUNT`.

    **Called when:** once during `BaseRunner.__init__`, before
    `init_checkpoint_manager`, `init_fault_tolerance`, and
    `init_garbage_collection`. The runner is partially constructed at
    this point — `self.config`, `self.workspace`, `self.timestamp`, the
    dataloader container, and the default `FileCheckpointManager` are
    bound, but the model is not materialized and optimizers/dataloaders
    are not built.

    **Precondition:** environment variables `WORLD_SIZE`, `RANK`,
    `LOCAL_RANK` are set when running distributed. The default
    torch.distributed process group is **not** already initialized when
    `WORLD_SIZE > 1` or TorchFT is enabled — the runner owns process-group lifecycle.

    Raises:
        RuntimeError: the default process group is already initialized
            when `WORLD_SIZE > 1` or TorchFT is enabled.
        ValueError: `dist.init_timeout_seconds` is non-positive.

    **Side effects:** when `WORLD_SIZE > 1` or TorchFT is enabled, calls
    `dist.init_process_group(...)`, sets the active CUDA device when
    the runtime device is CUDA, and broadcasts `self.timestamp` from rank 0.
    Reads `TORCHELASTIC_RESTART_COUNT` into `elastic_state.restart_count`.

    !!! danger "Do not"
        - Initialize a process group via `dist.init_process_group(...)`
          outside the runner; the runner owns its lifecycle.
        - Build the model or dataloaders here; those happen in
          `__post_init__`.
        - Bind the checkpoint manager here; `init_checkpoint_manager`
          runs next.

    **Backend notes:**

    - `ParallelRunner` extends this hook: after calling `super()`, it
      builds the parallel topology (`build_topology`) and initializes
      per-axis process groups via `init_device_mesh`.
    - `DeepSpeedRunner` inherits the default; DeepSpeed reuses the
      default process group initialized here.
    """

    backend = self.config.get("dist.backend") or os.getenv("BACKEND")
    init_method = self.config.get("dist.init_method") or os.getenv("INIT_METHOD")
    init_timeout = self._dist_timeout("dist.init_timeout_seconds")
    world_size = int(os.getenv("WORLD_SIZE", "1"))
    rank = int(os.getenv("RANK", "0"))
    runtime_device = self.device
    if backend is None:
        backend = dist.Backend.default_device_backend_map.get(runtime_device.type)
    use_cuda_runtime = torch.cuda.is_available() and runtime_device.type == "cuda"
    runtime_device_index = runtime_device.index if runtime_device.index is not None else self.local_rank
    dist_ready = dist.is_available() and dist.is_initialized()
    requires_process_group = world_size > 1 or bool(self.config.get("ft.enabled", False))
    if requires_process_group and dist_ready:
        raise RuntimeError(
            "default process group is already initialized; Runner requires owning process-group lifecycle"
        )
    if requires_process_group:
        if use_cuda_runtime:
            torch.cuda.set_device(runtime_device_index)
        init_kwargs: dict[str, Any] = {
            "backend": backend,
            "init_method": init_method,
            "world_size": world_size,
            "rank": rank,
        }
        if init_timeout is not None:
            init_kwargs["timeout"] = init_timeout
        dist.init_process_group(**init_kwargs)
        dist_ready = bool(dist.is_available() and dist.is_initialized())

    if dist_ready and use_cuda_runtime:
        torch.cuda.set_device(runtime_device_index)

    if dist_ready and self.world_size > 1:
        object_list = [self.timestamp]
        dist.broadcast_object_list(object_list)
        self.timestamp = str(object_list[0])

    restart_count = os.getenv("TORCHELASTIC_RESTART_COUNT")
    if restart_count is not None:
        self.elastic_state.restart_count = int(restart_count)

    self._train_pg_timeout_reduced = False

init_checkpoint_manager

Python

init_checkpoint_manager() -> None

Bind the checkpoint manager corresponding to config.ckpt.backend.

The default dispatches by backend: when the backend is "dcp", it binds a TorchDistributedCheckpointManager (or TorchFTCheckpointManager when dataloader checkpoints are enabled). For "file" it leaves the FileCheckpointManager already bound by BaseRunner.__init__ in place.

Called when: once during BaseRunner.__init__, after init_distributed and before init_fault_tolerance. The default FileCheckpointManager is already bound at this point — overrides should swap it via set_checkpoint_manager(...), not by direct attribute assignment.

Precondition: config.ckpt.backend is normalized to one of {"file", "dcp"} (TorchRunner does this in __init__). When the backend is "dcp", the default process group is initialized for distributed runs.

Side effects: swaps self.checkpoint_manager via set_checkpoint_manager(...) when the backend differs from "file". The prior manager is closed with a zero timeout.

Do not

• Set self.checkpoint_manager directly; use set_checkpoint_manager so the prior manager is closed cleanly.
• Initialize fault tolerance here; init_fault_tolerance runs next.
• Bind the model or dataloaders here.

Backend notes:

• DeepSpeedRunner coerces config.ckpt.backend to "file" in __init__, so this hook is a no-op for that backend.
• ParallelRunner coerces the backend to "dcp", so this hook always binds TorchDistributedCheckpointManager or TorchFTCheckpointManager.

Source code in danling/runners/torch_runner.py

def init_checkpoint_manager(self) -> None:
    """
    Bind the checkpoint manager corresponding to `config.ckpt.backend`.

    The default dispatches by backend: when the backend is `"dcp"`, it
    binds a `TorchDistributedCheckpointManager` (or
    `TorchFTCheckpointManager` when dataloader checkpoints are
    enabled). For `"file"` it leaves the `FileCheckpointManager` already
    bound by `BaseRunner.__init__` in place.

    **Called when:** once during `BaseRunner.__init__`, after
    `init_distributed` and before `init_fault_tolerance`. The default
    `FileCheckpointManager` is already bound at this point — overrides
    should swap it via `set_checkpoint_manager(...)`, not by direct
    attribute assignment.

    **Precondition:** `config.ckpt.backend` is normalized to one
    of `{"file", "dcp"}` (TorchRunner does this in `__init__`). When
    the backend is `"dcp"`, the default process group is initialized
    for distributed runs.

    **Side effects:** swaps `self.checkpoint_manager` via
    `set_checkpoint_manager(...)` when the backend differs from
    `"file"`. The prior manager is closed with a zero timeout.

    !!! danger "Do not"
        - Set `self.checkpoint_manager` directly; use
          `set_checkpoint_manager` so the prior manager is closed
          cleanly.
        - Initialize fault tolerance here; `init_fault_tolerance` runs
          next.
        - Bind the model or dataloaders here.

    **Backend notes:**

    - `DeepSpeedRunner` coerces `config.ckpt.backend` to `"file"`
      in `__init__`, so this hook is a no-op for that backend.
    - `ParallelRunner` coerces the backend to `"dcp"`, so this hook
      always binds `TorchDistributedCheckpointManager` or
      `TorchFTCheckpointManager`.
    """
    checkpoint_backend = str(self.config.get("ckpt.backend")).lower()
    if checkpoint_backend == "dcp":
        ft_checkpoint_enabled = bool(
            self.config.get("ft.enabled", False) or self.config.get("ckpt.dataloader_checkpoint.enabled", False)
        )
        manager_cls = TorchFTCheckpointManager if ft_checkpoint_enabled else TorchDistributedCheckpointManager
        self.set_checkpoint_manager(manager_cls(self))
        return

init_tensorboard

Python

init_tensorboard(*args, **kwargs) -> None

Set up TensorBoard SummaryWriter.

Source code in danling/runners/torch_runner.py

@on_main_process
def init_tensorboard(self, *args, **kwargs) -> None:
    r"""
    Set up TensorBoard SummaryWriter.
    """

    try:
        from torch.utils.tensorboard.writer import SummaryWriter  # pylint: disable=C0415
    except ImportError as error:
        raise ImportError(
            "TensorBoard logging requires the optional TensorBoard dependencies; "
            "install them with `pip install 'danling[tensorboard]'`."
        ) from error

    tensorboard_config = self.config.tensorboard
    for key in ("log_dir", "comment", "purge_step", "max_queue", "flush_secs", "filename_suffix"):
        if key not in kwargs and tensorboard_config.get(key) is not None:
            kwargs[key] = tensorboard_config[key]
    if "log_dir" not in kwargs:
        kwargs["log_dir"] = os.path.join(self.workspace.dir, "tensorboard", self.timestamp)

    self.writer = SummaryWriter(*args, **kwargs)
    self.writer.add_scalar = catch(OSError, verbose=False)(self.writer.add_scalar)  # type: ignore[method-assign]

set_seed

Python

set_seed(seed: int | None = None, bias: int | bool | None = None) -> int

Set up random seed.

Parameters:

Name	Type	Description	Default
seed	int | None	Random seed to set. Defaults to self.config.seed (config.seed).	None
bias	int | bool | None	Make the seed different for each processes. This is used to ensure the data augmentation are applied differently on every processes. Defaults to self.rank. Set to False to disable this feature.	None

Source code in danling/runners/torch_runner.py

def set_seed(self, seed: int | None = None, bias: int | bool | None = None) -> int:
    r"""
    Set up random seed.

    Args:
        seed: Random seed to set.
            Defaults to `self.config.seed` (`config.seed`).

        bias: Make the seed different for each processes.
            This is used to ensure the data augmentation are applied differently on every processes.
            Defaults to `self.rank`.
            Set to `False` to disable this feature.
    Returns:
        Random seed set.
    """

    base_seed = seed if seed is not None else self.config.seed  # type: ignore[assignment]
    if base_seed is None:
        base_seed = random.randint(0, 2**32 - 1)
        if self.distributed and dist.is_initialized():
            object_list = [base_seed]
            dist.broadcast_object_list(object_list)
            base_seed = object_list[0]
    base_seed = int(base_seed)
    # Keep `config.seed` as the global/base seed (before per-rank bias).
    self.config.seed = base_seed

    process_seed = base_seed
    if bias is None:
        if self.fault_tolerance is not None:
            _, bias = self.fault_tolerance.data_parallel_info(self.world_size, self.rank)
        else:
            bias = self.rank
    if bias:
        process_seed += int(bias)

    torch.manual_seed(process_seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(process_seed)
    if np_random is not None:
        np_random.seed(process_seed)
    random.seed(process_seed)
    self.rng_state.python = random.getstate()
    self.rng_state.numpy = np_random.get_state() if np_random is not None else None
    self.rng_state.torch_cpu = torch.get_rng_state()
    if torch.cuda.is_available():
        self.rng_state.torch_cuda = torch.cuda.get_rng_state_all()
    else:
        self.rng_state.torch_cuda = None
    return process_seed

runner_owns_grad_scaling

Python

runner_owns_grad_scaling() -> bool

Whether this runner should own AMP gradient scaling.

Source code in danling/runners/torch_runner.py

def runner_owns_grad_scaling(self) -> bool:
    """Whether this runner should own AMP gradient scaling."""

    return True

setup_grad_scaler

Python

setup_grad_scaler() -> None

Bind the AMP grad scaler for runner-owned fp16 training.

Source code in danling/runners/torch_runner.py

def setup_grad_scaler(self) -> None:
    """Bind the AMP grad scaler for runner-owned fp16 training."""

    self.grad_scaler = None
    precision = self._normalized_precision_name(self.precision)
    if precision not in {"fp16", "float16", "half"}:
        return
    if not self.runner_owns_grad_scaling():
        return
    if self.fp8_enabled:
        raise ValueError("precision='fp16' cannot be combined with FP8 autocast")
    if self.device.type != "cuda":
        raise ValueError("runner-owned fp16 precision requires a CUDA device; use bf16 or a backend-owned scaler")
    self.grad_scaler = torch.amp.GradScaler(device=self.device.type)

update_ema

Python

update_ema() -> None

Update EMA weights after a successful optimizer step.

Source code in danling/runners/torch_runner.py

def update_ema(self) -> None:
    """Update EMA weights after a successful optimizer step."""

    if self.ema is None:
        return
    if self.model is None:
        raise ValueError("cannot update EMA: model is not initialized")
    update = getattr(self.ema, "update", None)
    if update is None:
        return
    if not callable(update):
        raise TypeError(
            "EMA module has an `update` attribute but it is not callable; "
            "implement `ema.update(model)`, remove the attribute, or override `update_ema()`."
        )
    update(self.unwrap(self.model))

build_model

Python

build_model() -> Module

Build a model under the runner’s meta-device context when opted in.

Override this instead of binding self.model in __init__ when model_init.meta=True. Model-specific configuration stays on the runner’s existing Chanfig configuration.

Source code in danling/runners/torch_runner.py

def build_model(self) -> nn.Module:
    r"""Build a model under the runner's meta-device context when opted in.

    Override this instead of binding ``self.model`` in ``__init__`` when
    ``model_init.meta=True``. Model-specific configuration stays on the
    runner's existing Chanfig configuration.
    """
    raise NotImplementedError("model_init.meta=True requires a model or an overridden build_model()")

initialize_model

Python

initialize_model(model: Module) -> None

Initialize or load newly allocated model storage before optimizer creation.

This runs only for a meta model. Override it to initialize parameters and persistent/non-persistent buffers, or to load model-specific state. The default calls the model’s explicit init_weights() protocol; it does not guess initialization order or reset pretrained models. Ordinary checkpoint/pretrained restoration still runs in auto_restore.

Source code in danling/runners/torch_runner.py

def initialize_model(self, model: nn.Module) -> None:
    r"""Initialize or load newly allocated model storage before optimizer creation.

    This runs only for a meta model. Override it to initialize parameters
    and persistent/non-persistent buffers, or to load model-specific state.
    The default calls the model's explicit ``init_weights()`` protocol;
    it does not guess initialization order or reset pretrained models.
    Ordinary checkpoint/pretrained restoration still runs in ``auto_restore``.
    """
    initialize = getattr(model, "init_weights", None)
    if not callable(initialize):
        raise NotImplementedError("meta models require init_weights() or an overridden initialize_model()")
    initialize()

model_init_device

Python

model_init_device(model: Module) -> device

Select the parameter-allocation device for a meta model.

Source code in danling/runners/torch_runner.py

def model_init_device(self, model: nn.Module) -> torch.device:
    r"""Select the parameter-allocation device for a meta model."""
    return self.device

materialize_model_values

Python

materialize_model_values(model: Module) -> Module

Allocate a meta model after transforms, retaining tied tensor identities.

Already initialized models keep the ordinary to(device) path. Mixed initialized/meta tensors require a model-specific override since to_empty would otherwise discard initialized values.

Source code in danling/runners/torch_runner.py

def materialize_model_values(self, model: nn.Module) -> nn.Module:
    r"""Allocate a meta model after transforms, retaining tied tensor identities.

    Already initialized models keep the ordinary ``to(device)`` path.
    Mixed initialized/meta tensors require a model-specific override since
    ``to_empty`` would otherwise discard initialized values.
    """
    tensors = tuple(model.parameters()) + tuple(model.buffers())
    if not any(tensor.is_meta for tensor in tensors):
        return model.to(self.device)
    if not self.config.model_init.meta:
        raise ValueError("meta model materialization requires model_init.meta=True")
    if any(not tensor.is_meta for tensor in tensors):
        raise ValueError("mixed initialized/meta model tensors require overriding materialize_model_values()")
    if self.optimizer is not None:
        raise ValueError("build the optimizer after meta model materialization")

    aliases: dict[int, list[tuple[nn.Module, str, str]]] = {}
    for module in model.modules():
        for collection in ("_parameters", "_buffers"):
            for name, tensor in getattr(module, collection).items():
                if tensor is not None:
                    aliases.setdefault(id(tensor), []).append((module, collection, name))
    allocation_device = torch.device(self.model_init_device(model))
    model.to_empty(device=allocation_device)
    for references in aliases.values():
        owner, collection, name = references[0]
        tensor = getattr(owner, collection)[name]
        if collection == "_buffers" and allocation_device != self.device:
            tensor = torch.empty_like(tensor, device=self.device)
        for owner, collection, name in references:
            getattr(owner, collection)[name] = tensor
    with torch.no_grad():
        self.initialize_model(self.unwrap(model))
    if any(tensor.is_meta for tensor in (*model.parameters(), *model.buffers())):
        raise RuntimeError("initialize_model() left meta tensors in the model")
    buffers: dict[int, torch.Tensor] = {}
    for module in model.modules():
        for name, buffer in module._buffers.items():
            if buffer is not None:
                if id(buffer) not in buffers:
                    buffers[id(buffer)] = buffer.to(self.device)
                module._buffers[name] = buffers[id(buffer)]
    return model

compile_regions

Python

compile_regions(model: Module) -> Mapping[str, Sequence[Module]]

Return model-owned named module groups, resolved after transforms.

Override this or implement model.compile_regions(). Members of a group share one forward implementation and compiler cache. Modules retain their identities, parameter names, and Module.__call__ hooks.

Source code in danling/runners/torch_runner.py

def compile_regions(self, model: nn.Module) -> Mapping[str, Sequence[nn.Module]]:
    r"""Return model-owned named module groups, resolved after transforms.

    Override this or implement ``model.compile_regions()``. Members of a
    group share one forward implementation and compiler cache. Modules
    retain their identities, parameter names, and ``Module.__call__`` hooks.
    """
    regions = getattr(model, "compile_regions", None)
    return regions() if callable(regions) else {}

compile_model

Python

compile_model(model: Module) -> Module

Compile selected named groups, or the whole model when none are selected.

Source code in danling/runners/torch_runner.py

def compile_model(self, model: nn.Module) -> nn.Module:
    r"""Compile selected named groups, or the whole model when none are selected."""
    requested = self.config.compile.regions
    if not self.compiler.enabled or not requested:
        return self.compiler.compile(model)
    if isinstance(requested, str) or len(set(requested)) != len(requested):
        raise ValueError("compile.regions must contain distinct region names")
    available = self.compile_regions(model)
    groups = []
    seen: set[int] = set()
    for name in requested:
        if name not in available:
            raise ValueError(f"unknown compile region {name!r}; available regions: {tuple(available)}")
        modules = tuple(available[name])
        if not modules or any(not isinstance(module, nn.Module) for module in modules):
            raise ValueError(f"compile region {name!r} must contain nn.Module instances")
        if seen.intersection(id(module) for module in modules):
            raise ValueError("selected compile regions must not share modules")
        seen.update(id(module) for module in modules)
        groups.append(modules)
    for modules in groups:
        self.compiler.compile_modules(modules)
    return model

materialize_model

Python

materialize_model() -> None

Move the model to the runtime device, optionally compile, and wrap with DDP when distributed.

The default is a single-module DDP-style materialization: it moves self.model to self.device, applies any FP8 module policy when FP8 is enabled, runs torch.compile via self.compiler, and wraps the result with nn.parallel.DistributedDataParallel when world size > 1. The configured DDP-optimizer context covers both DDP construction and the compiled callable’s lazy compilation and recompilation.

Called when: once during __post_init__, after setup_fp8() and before build_optimizer(). The order matters — the optimizer must see post-wrap parameters.

Precondition: self.model is set (typically by the user before constructing the runner). self.device resolves to the runtime device.

Raises:

Type	Description
ValueError	self.model is not initialized.

Side effects: moves self.model to self.device; applies FP8 module policy when self.fp8_enabled; compiles via self.compiler.compile(...); wraps with DistributedDataParallel for world size > 1 under the same configured DDP-optimizer context. Moves self.ema to device when EMA is bound.

Do not

• Build the optimizer or scheduler here; they run after this hook.
• Skip the device move when overriding (tensors must live on self.device before the forward pass).
• Re-wrap an already-wrapped model (e.g., DDP-wrap a DDP module).

Backend notes:

• DeepSpeedRunner overrides this hook to move the model to device and compile only; the DeepSpeed engine wraps the model later in _finalize_runtime_components.
• ParallelRunner overrides this hook for FSDP2, pipeline-parallel schedules, and tensor/expert/context parallelism (via the parallelize_model and apply_activation_checkpointing hooks).

Source code in danling/runners/torch_runner.py

def materialize_model(self) -> None:
    """
    Move the model to the runtime device, optionally compile, and wrap
    with DDP when distributed.

    The default is a single-module DDP-style materialization: it moves
    `self.model` to `self.device`, applies any FP8 module policy when
    FP8 is enabled, runs `torch.compile` via `self.compiler`, and wraps the
    result with `nn.parallel.DistributedDataParallel` when world size > 1.
    The configured DDP-optimizer context covers both DDP construction and
    the compiled callable's lazy compilation and recompilation.

    **Called when:** once during `__post_init__`, after `setup_fp8()`
    and before `build_optimizer()`. The order matters — the optimizer
    must see post-wrap parameters.

    **Precondition:** `self.model` is set (typically by the user before
    constructing the runner). `self.device` resolves to the runtime
    device.

    Raises:
        ValueError: `self.model` is not initialized.

    **Side effects:** moves `self.model` to `self.device`; applies FP8
    module policy when `self.fp8_enabled`; compiles via
    `self.compiler.compile(...)`; wraps with `DistributedDataParallel` for
    world size > 1 under the same configured DDP-optimizer context. Moves
    `self.ema` to device when EMA is bound.

    !!! danger "Do not"
        - Build the optimizer or scheduler here; they run after this
          hook.
        - Skip the device move when overriding (tensors must live on
          `self.device` before the forward pass).
        - Re-wrap an already-wrapped model (e.g., DDP-wrap a DDP module).

    **Backend notes:**

    - `DeepSpeedRunner` overrides this hook to move the model to device
      and compile only; the DeepSpeed engine wraps the model later in
      `_finalize_runtime_components`.
    - `ParallelRunner` overrides this hook for FSDP2, pipeline-parallel
      schedules, and tensor/expert/context parallelism (via the
      `parallelize_model` and `apply_activation_checkpointing` hooks).
    """
    if self.model is None:
        raise ValueError("cannot materialize model: model is not initialized")

    model = self.materialize_model_values(self.model)
    self.model = model
    if self.fp8_enabled:
        self.apply_fp8_module_policy_to_model_parts()
        model = self.model
    should_wrap_ddp = self.distributed and not isinstance(
        model, (nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)
    )
    model = self.compile_model(model)
    if should_wrap_ddp:
        with self.compiler.ddp_optimizer():
            model = nn.parallel.DistributedDataParallel(
                model,
                find_unused_parameters=bool(self.config.ddp.find_unused_parameters),
                static_graph=bool(self.config.ddp.static_graph),
            )
    self.model = model

    if self.ema is not None:
        self.ema = self.ema.to(self.device)

unwrap

Python

unwrap(model: Module) -> Module

Return the native module beneath compile and data-parallel wrappers.

Source code in danling/runners/torch_runner.py

def unwrap(self, model: nn.Module) -> nn.Module:
    """Return the native module beneath compile and data-parallel wrappers."""
    from torch._dynamo.eval_frame import OptimizedModule

    while isinstance(model, (OptimizedModule, nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)):
        model = model._orig_mod if isinstance(model, OptimizedModule) else model.module
    return model

build_optimizer

Python

build_optimizer() -> None

Auto-build the optimizer from config.optim when self.optimizer is absent.

The default iterates parameters via iter_optimizer_parameters and dispatches to the OPTIMIZERS registry with the merged config. If optim.param_groups is configured, entries are matched by regex search against iter_optimizer_named_parameters; unmatched parameters keep the optimizer-level defaults.

Called when: once during TorchRunner.__post_init__, after materialize_model (so parameters reflect DDP/FSDP wrapping) and before build_scheduler.

Precondition: self.model is materialized and on self.device. self.optimizer is None (the auto-build is skipped when the user has already bound an optimizer).

Side effects: sets self.optimizer to the registry-built instance.

Do not

• Run before materialize_model; parameters won’t reflect DDP/FSDP wrapping.
• Build a scheduler here.
• Override parameter enumeration here; override iter_optimizer_parameters / iter_optimizer_named_parameters instead so subclass topology (e.g., ParallelRunner.model_parts) is preserved.

Backend notes:

• DeepSpeedRunner inherits this hook; DeepSpeed may replace the optimizer with a DeepSpeed-managed instance during _finalize_runtime_components.
• ParallelRunner inherits this hook but overrides iter_optimizer_parameters to enumerate self.model_parts.

Source code in danling/runners/torch_runner.py

def build_optimizer(self) -> None:
    """
    Auto-build the optimizer from `config.optim` when `self.optimizer`
    is absent.

    The default iterates parameters via `iter_optimizer_parameters` and
    dispatches to the `OPTIMIZERS` registry with the merged config. If
    `optim.param_groups` is configured, entries are matched by regex
    `search` against `iter_optimizer_named_parameters`; unmatched
    parameters keep the optimizer-level defaults.

    **Called when:** once during `TorchRunner.__post_init__`, after
    `materialize_model` (so parameters reflect DDP/FSDP wrapping) and
    before `build_scheduler`.

    **Precondition:** `self.model` is materialized and on `self.device`.
    `self.optimizer` is `None` (the auto-build is skipped when the user
    has already bound an optimizer).

    **Side effects:** sets `self.optimizer` to the registry-built
    instance.

    !!! danger "Do not"
        - Run before `materialize_model`; parameters won't reflect
          DDP/FSDP wrapping.
        - Build a scheduler here.
        - Override parameter enumeration here; override
          `iter_optimizer_parameters` / `iter_optimizer_named_parameters`
          instead so subclass topology (e.g., `ParallelRunner.model_parts`)
          is preserved.

    **Backend notes:**

    - `DeepSpeedRunner` inherits this hook; DeepSpeed may replace the
      optimizer with a DeepSpeed-managed instance during
      `_finalize_runtime_components`.
    - `ParallelRunner` inherits this hook but overrides
      `iter_optimizer_parameters` to enumerate `self.model_parts`.
    """
    if self.optimizer is not None or self.model is None:
        return
    optimizer_cfg = self.config.get("optim")
    if not isinstance(optimizer_cfg, Mapping) or not optimizer_cfg:
        return
    optimizer_kwargs = {str(key): value for key, value in optimizer_cfg.items() if value is not None}
    optimizer_kwargs.pop("param_groups", None)
    if "type" not in optimizer_kwargs:
        return
    parameters = self._build_optimizer_param_groups(optimizer_cfg)
    if not parameters:
        return
    self.optimizer = OPTIMIZERS.build(params=parameters, **optimizer_kwargs)

build_scheduler

Python

build_scheduler() -> None

Auto-build the LR scheduler from config.sched when self.scheduler is absent.

The default pops interval and monitor from the config (those drive runner-level dispatch, not scheduler construction), defaults total_steps to self.steps when computable, and dispatches to the SCHEDULERS registry with self.optimizer and the merged config.

Called when: once during TorchRunner.__post_init__, after build_optimizer.

Precondition: self.optimizer is bound. self.scheduler is None (the auto-build is skipped when the user has already bound a scheduler).

Side effects: sets self.scheduler to the registry-built instance.

Do not

• Run before build_optimizer; the scheduler must wrap an optimizer.
• Set scheduler interval or monitor here; configure them via config.sched.interval / config.sched.monitor.

Backend notes:

• DeepSpeedRunner inherits this hook; the scheduler may be handed to the DeepSpeed engine in _finalize_runtime_components when its effective interval is "step". Otherwise the runner retains it.

Source code in danling/runners/torch_runner.py

def build_scheduler(self) -> None:
    """
    Auto-build the LR scheduler from `config.sched` when
    `self.scheduler` is absent.

    The default pops `interval` and `monitor` from the config (those
    drive runner-level dispatch, not scheduler construction), defaults
    `total_steps` to `self.steps` when computable, and dispatches to
    the `SCHEDULERS` registry with `self.optimizer` and the merged
    config.

    **Called when:** once during `TorchRunner.__post_init__`, after
    `build_optimizer`.

    **Precondition:** `self.optimizer` is bound. `self.scheduler` is
    `None` (the auto-build is skipped when the user has already bound a
    scheduler).

    **Side effects:** sets `self.scheduler` to the registry-built
    instance.

    !!! danger "Do not"
        - Run before `build_optimizer`; the scheduler must wrap an
          optimizer.
        - Set scheduler interval or monitor here; configure them via
          `config.sched.interval` / `config.sched.monitor`.

    **Backend notes:**

    - `DeepSpeedRunner` inherits this hook; the scheduler may be handed
      to the DeepSpeed engine in `_finalize_runtime_components` when
      its effective interval is `"step"`. Otherwise the runner retains
      it.
    """
    if self.scheduler is not None or self.optimizer is None:
        return
    sched_cfg = self._get_scheduler_config()
    if not isinstance(sched_cfg, Mapping) or not sched_cfg:
        return
    scheduler_kwargs = {str(key): value for key, value in sched_cfg.items() if value is not None}
    scheduler_kwargs.pop("interval", None)
    scheduler_kwargs.pop("monitor", None)
    if "type" not in scheduler_kwargs:
        return
    if "total_steps" not in scheduler_kwargs:
        steps = self.steps
        if steps is not None:
            scheduler_kwargs["total_steps"] = steps
    self.scheduler = SCHEDULERS.build(self.optimizer, **scheduler_kwargs)

build_dataloaders

Python

build_dataloaders()

Build dataloaders for dataset splits not already materialized.

The default iterates self.datasets, merges config.dataloader defaults with split-specific overrides (config.dataloader.<split>), constructs a sampler via build_datasampler, and wraps each dataset in a StatefulDataLoader using self.collate_fn. Train splits default to shuffle=True and drop_last=True; non-train splits default to the opposite.

Called when: once during TorchRunner.__post_init__ when self.datasets is non-empty.

Precondition: self.datasets is populated (typically by the user before constructing the runner). self.dataloaders is bound to a default-constructed DataLoaderDict.

Side effects: populates self.dataloaders[split] for each split in self.datasets not already materialized. Existing entries in self.dataloaders are left untouched.

Do not

• Override sampler logic here; override build_datasampler instead.
• Override collation; set self.collate_fn or override collate_fn (classmethod) instead.
• Bind the optimizer or scheduler here.

Backend notes:

• ParallelRunner substitutes self.dataloaders with a proxying dict in __init__ so non-first/last pipeline stages receive a StepProxyLoader view. The build logic itself is inherited.

Source code in danling/runners/torch_runner.py

def build_dataloaders(self):
    """
    Build dataloaders for dataset splits not already materialized.

    The default iterates `self.datasets`, merges `config.dataloader`
    defaults with split-specific overrides (`config.dataloader.<split>`),
    constructs a sampler via `build_datasampler`, and wraps each dataset
    in a `StatefulDataLoader` using `self.collate_fn`. Train splits
    default to `shuffle=True` and `drop_last=True`; non-train splits
    default to the opposite.

    **Called when:** once during `TorchRunner.__post_init__` when
    `self.datasets` is non-empty.

    **Precondition:** `self.datasets` is populated (typically by the
    user before constructing the runner). `self.dataloaders` is bound
    to a default-constructed `DataLoaderDict`.

    **Side effects:** populates `self.dataloaders[split]` for each
    split in `self.datasets` not already materialized. Existing entries
    in `self.dataloaders` are left untouched.

    !!! danger "Do not"
        - Override sampler logic here; override `build_datasampler`
          instead.
        - Override collation; set `self.collate_fn` or override
          `collate_fn` (classmethod) instead.
        - Bind the optimizer or scheduler here.

    **Backend notes:**

    - `ParallelRunner` substitutes `self.dataloaders` with a proxying
      dict in `__init__` so non-first/last pipeline stages receive a
      `StepProxyLoader` view. The build logic itself is inherited.
    """
    datasets = {k: d for k, d in self.datasets.items() if k not in self.dataloaders}
    dataloader_config = self.config.get("dataloader", NestedDict())
    default_kwargs = NestedDict({k: v for k, v in dataloader_config.items() if k not in self.datasets})
    split_kwargs = NestedDict({k: v for k, v in dataloader_config.items() if k in self.datasets})
    if self.config.deterministic and default_kwargs.get("worker_init_fn") is None:
        default_kwargs["worker_init_fn"] = _seed_dataloader_worker
    # Worker/iterator seeds must not consume the model RNG when a loader is
    # rebuilt on checkpoint resume, even with nondeterministic kernels.
    for k, dataset in datasets.items():
        kwargs = NestedDict(default_kwargs)
        if k in split_kwargs:
            kwargs.merge(split_kwargs[k], overwrite=True)
        if kwargs.get("generator") is None:
            generator = torch.Generator()
            if self.config.seed is None:
                generator.seed()
            else:
                generator.manual_seed(int(self.config.seed) + int(self.rank))
            kwargs["generator"] = generator
        is_train_split = k in self.train_splits
        shuffle = kwargs.pop("shuffle", is_train_split)
        collate_fn = kwargs.pop("collate_fn", self.collate_fn)
        batch_sampler = kwargs.pop("batch_sampler", None)
        if batch_sampler is not None:
            kwargs.pop("batch_size", None)
            kwargs.pop("drop_last", None)
            kwargs.pop("sampler", None)
            self.dataloaders[k] = StatefulDataLoader(
                dataset,
                batch_sampler=batch_sampler,
                collate_fn=collate_fn,
                **kwargs,
            )
            continue
        kwargs.setdefault("drop_last", is_train_split)
        sampler = kwargs.pop("sampler", None)
        if sampler is None:
            sampler = self.build_datasampler(dataset, split=k, shuffle=shuffle)
            if type(sampler) is StatefulRandomSampler:
                # Loader worker seeds and sampler replay must not advance each other's RNG.
                sampler.generator = torch.Generator().set_state(kwargs["generator"].get_state())
        self.dataloaders[k] = StatefulDataLoader(dataset, sampler=sampler, collate_fn=collate_fn, **kwargs)

build_datasampler

Python

build_datasampler(dataset: Any, *, split: str, shuffle: bool) -> Any

Build the sampler for one dataset split.

Called when: build_dataloaders materializes a split from self.datasets.

Parameters:

Name	Type	Description	Default
dataset	Any	Dataset object for the split.	required
split	str	Split name being materialized.	required
shuffle	bool	Whether this split should be sampled in shuffled order.	required

Returns:

Type	Description
Any	A local random/sequential sampler in single-process mode, or a
Any	DistributedSampler in distributed mode.

Backend notes:

• ParallelRunner overrides replica/rank selection so data-parallel sampling follows its topology instead of raw global rank.

Source code in danling/runners/torch_runner.py

def build_datasampler(self, dataset: Any, *, split: str, shuffle: bool) -> Any:
    """
    Build the sampler for one dataset split.

    **Called when:** `build_dataloaders` materializes a split from
    `self.datasets`.

    Args:
        dataset: Dataset object for the split.
        split: Split name being materialized.
        shuffle: Whether this split should be sampled in shuffled order.

    Returns:
        A local random/sequential sampler in single-process mode, or a
        `DistributedSampler` in distributed mode.

    **Backend notes:**

    - `ParallelRunner` overrides replica/rank selection so data-parallel
      sampling follows its topology instead of raw global rank.
    """
    if self.distributed:
        num_replicas = self.world_size
        rank = self.rank
        if self.fault_tolerance is not None:
            num_replicas, rank = self.fault_tolerance.data_parallel_info(num_replicas, rank)
        return utils.data.distributed.DistributedSampler(
            dataset,
            num_replicas=num_replicas,
            rank=rank,
            shuffle=shuffle,
        )
    return StatefulRandomSampler(dataset) if shuffle else utils.data.SequentialSampler(dataset)

to_device

Python

to_device(data: Any)

Move one batch to runtime device; override in subclasses for custom fast paths.

Source code in danling/runners/torch_runner.py

def to_device(self, data: Any):
    """Move one batch to runtime device; override in subclasses for custom fast paths."""
    return to_device(data, self.device)

all_reduce

Python

all_reduce(tensor: Tensor, *, op=SUM) -> Tensor

Reduce tensor over the runner’s replica/data-parallel collective domain.

Source code in danling/runners/torch_runner.py

def all_reduce(self, tensor: torch.Tensor, *, op=dist.ReduceOp.SUM) -> torch.Tensor:
    """Reduce tensor over the runner's replica/data-parallel collective domain."""
    if dist.is_available() and dist.is_initialized():
        group = self.all_reduce_group()
        ft_group = None if self.fault_tolerance is None else self.fault_tolerance.replicate_process_group
        if group is not None and group is ft_group:
            dist.all_reduce(tensor, op=op)
        dist.all_reduce(tensor, op=op, group=group)
    return tensor

reduce

Python

reduce(tensor: Tensor) -> Tensor

Average-reduce tensor over the runner’s collective domain.

Source code in danling/runners/torch_runner.py

def reduce(self, tensor: torch.Tensor) -> torch.Tensor:
    """Average-reduce tensor over the runner's collective domain."""
    tensor = _local_reduction_tensor(tensor)
    if not (dist.is_available() and dist.is_initialized()):
        return tensor
    group = self.all_reduce_group()
    group_size = max(self.world_size if group is None else dist.get_world_size(group=group), 1)
    if (
        group is not None
        and self.fault_tolerance is not None
        and group is self.fault_tolerance.replicate_process_group
    ):
        group_size *= max(self.world_size, 1)
    if group_size <= 1:
        return tensor

    original_device = tensor.device
    payload_device = self.all_reduce_device()
    payload = tensor if original_device == payload_device else tensor.to(payload_device)
    self.all_reduce(payload, op=dist.ReduceOp.SUM)
    payload = payload / group_size
    if payload.device != original_device:
        payload = payload.to(original_device)
    return payload

reduce_loss_for_logging

Python

reduce_loss_for_logging(loss: Tensor | None, loss_n: float | None) -> Tensor | None

Detach and all-reduce weighted loss tensor for logging.

Source code in danling/runners/torch_runner.py

def reduce_loss_for_logging(self, loss: torch.Tensor | None, loss_n: float | None) -> torch.Tensor | None:
    """Detach and all-reduce weighted loss tensor for logging."""
    if loss is None:
        return None
    loss_value = _local_reduction_tensor(loss.detach()).to(dtype=torch.float64)
    if loss_value.ndim > 0:
        loss_value = loss_value.mean()
    normalizer = float(loss_n) if loss_n is not None else 1.0
    summed = getattr(self, "_loss_reporting_reduction", None) == "sum"
    if summed:
        normalizer = 1.0
    payload_device = self.all_reduce_device()
    payload = torch.stack(
        (
            loss_value.to(device=payload_device) * normalizer,
            torch.tensor(normalizer, dtype=torch.float64, device=payload_device),
        )
    )
    self.all_reduce(payload, op=dist.ReduceOp.SUM)
    if payload[1].item() <= 0:
        return None
    return payload[0] if summed else payload[0] / payload[1]

train_context

Python

train_context()

Context for one training micro-step (autocast + optional DDP no_sync).

Source code in danling/runners/torch_runner.py

@contextmanager
def train_context(self):
    """Context for one training micro-step (autocast + optional DDP no_sync)."""
    self._begin_optimizer_window()
    with self._train_step_context(no_sync_targets=self._train_no_sync_targets()):
        yield

infer_context

Python

infer_context()

Precision context used by train/eval/infer forward passes.

Source code in danling/runners/torch_runner.py

def infer_context(self):
    """Precision context used by train/eval/infer forward passes."""

    if self.fp8_enabled:
        return self.fp8_autocast()

    precision = self.precision
    if precision is None:
        return nullcontext()
    return torch.autocast(self.device.type, dtype=get_precision(precision))

build_evaluation_metrics

Python

build_evaluation_metrics()

Provide a collector for an isolated training-time validation pass.

A separate evaluation collector retains its normal reset/update semantics. Native collectors shared with training are rebuilt from their metric definitions. Override this factory for a shared custom collector; it must return an independent instance, not a training one. TorchMetrics Metric callables use their public clone/reset lifecycle. Other metric functions and preprocessing callbacks must be stateless; override this factory when their callable objects carry mutable state.

Source code in danling/runners/torch_runner.py

def build_evaluation_metrics(self):
    """Provide a collector for an isolated training-time validation pass.

    A separate evaluation collector retains its normal reset/update
    semantics. Native collectors shared with training are rebuilt from
    their metric definitions. Override this factory for a shared custom
    collector; it must return an independent instance, not a training one.
    TorchMetrics Metric callables use their public clone/reset lifecycle.
    Other metric functions and preprocessing callbacks must be stateless;
    override this factory when their callable objects carry mutable state.
    """
    metrics = self.evaluate_metrics
    training = self._metric_state_ids(self.train_metrics) | self._metric_state_ids(self.metrics)
    if not (self._metric_state_ids(metrics) & training):
        return metrics

    def fresh_metric(metric):
        if isinstance(metric, TorchMetric):
            metric = metric.clone()
            metric.reset()
        return metric

    def fresh(source):
        if type(source) is GlobalMetrics:
            return GlobalMetrics(
                preprocess=source.preprocess,
                distributed=source.distributed,
                device=source.device,
                **source.metrics,
            )
        if type(source) is StreamMetrics:
            return StreamMetrics(
                preprocess=source.preprocess,
                distributed=source.distributed,
                device=source.device,
                **{name: fresh_metric(meter.metric) for name, meter in source.all_items()},
            )
        if type(source) is MultiTaskMetrics:
            return MultiTaskMetrics(
                aggregate=source.getattr("_aggregate", None),
                aggregate_weights=source.getattr("_aggregate_weights", None),
                **{name: fresh(child) for name, child in source.items()},
            )
        raise TypeError("shared custom metrics require build_evaluation_metrics() to return a fresh collector")

    return fresh(metrics)

build_evaluation_dataloader

Python

build_evaluation_dataloader(split: str)

Provide a validation loader starting at its sample-index origin.

Only training-time validation calls this factory. Native map-style loaders with standard samplers are reconstructed from public configuration at epoch zero. Custom loaders, stateful datasets and iterable streams must override this hook to define their fresh origin. Standalone evaluation continues to consume the caller’s original loader.

Source code in danling/runners/torch_runner.py

def build_evaluation_dataloader(self, split: str):
    """Provide a validation loader starting at its sample-index origin.

    Only training-time validation calls this factory. Native
    map-style loaders with standard samplers are reconstructed from public
    configuration at epoch zero. Custom loaders, stateful datasets and
    iterable streams must override this hook to define their fresh origin.
    Standalone evaluation continues to consume the caller's original loader.
    """
    # ParallelRunner exposes StepProxyLoader views via __getitem__; its
    # items remain the native loaders. Replacing raw entries lets that
    # mapping create fresh stage-specific proxies without unwrapping them.
    source = dict(self.dataloaders.items())[split]
    if type(source) in (list, tuple, range):
        return source
    if type(source) not in (utils.data.DataLoader, StatefulDataLoader):
        raise TypeError("custom validation loaders must provide build_evaluation_dataloader(split)")
    dataset = source.dataset
    if isinstance(dataset, utils.data.IterableDataset) or callable(getattr(dataset, "load_state_dict", None)):
        raise TypeError("stateful or iterable validation datasets must provide build_evaluation_dataloader(split)")

    def generator(original):
        seed = torch.initial_seed() if self.config.seed is None else int(self.config.seed) + self.rank
        return torch.Generator().manual_seed(original.initial_seed() if original is not None else seed)

    sampler = source.sampler
    if type(sampler) is utils.data.SequentialSampler:
        sampler = utils.data.SequentialSampler(dataset)
    elif type(sampler) is utils.data.distributed.DistributedSampler:
        degree, rank = self._evaluation_data_parallel_info()
        sampler = utils.data.distributed.DistributedSampler(
            dataset,
            num_replicas=degree,
            rank=rank,
            shuffle=sampler.shuffle,
            seed=sampler.seed,
            drop_last=sampler.drop_last,
        )
    elif type(sampler) in (utils.data.RandomSampler, StatefulRandomSampler):
        sampler = type(sampler)(
            dataset,
            replacement=sampler.replacement,
            num_samples=sampler.num_samples,
            generator=generator(sampler.generator),
        )
    else:
        raise TypeError("custom validation samplers must provide build_evaluation_dataloader(split)")
    kwargs = {
        name: getattr(source, name)
        for name in (
            "num_workers",
            "collate_fn",
            "pin_memory",
            "timeout",
            "worker_init_fn",
            "multiprocessing_context",
            "prefetch_factor",
            "persistent_workers",
            "pin_memory_device",
            "in_order",
            "snapshot_every_n_steps",
        )
        if hasattr(source, name)
    }
    kwargs["generator"] = generator(source.generator)
    if source.batch_size is None and source.batch_sampler is not None:
        batch_sampler = source.batch_sampler
        if type(batch_sampler) not in (utils.data.BatchSampler, StatefulBatchSampler):
            raise TypeError("custom validation batching must provide build_evaluation_dataloader(split)")
        kwargs["batch_sampler"] = type(batch_sampler)(sampler, batch_sampler.batch_size, batch_sampler.drop_last)
    else:
        kwargs.update(sampler=sampler, batch_size=source.batch_size, drop_last=source.drop_last)
    return type(source)(dataset, **kwargs)

train_step

Python

train_step(data: Any) -> tuple[Any, Tensor | None]

Run one training micro-step.

The default implementation runs forward → loss → metric update → backward → step for one micro-batch.

Called when: once per micro-batch by train_epoch/train_steps. The caller seeds the loop’s accumulation state before each invocation; this method consumes that state through backward() and step().

Precondition: self.model, self.optimizer, and self.criterion are bound; self.mode == RunnerMode.train.

Parameters:

Name	Type	Description	Default
data	Any	One micro-batch. The default unpacks data["input"] / data.get("target") for mappings, (data[0], data[1]) for non-string sequences, and (data, None) otherwise. Override train_step if your batch shape differs.	required

Returns:

Type	Description
Any	(pred, loss). pred is the model output (used by metrics.update).
Tensor | None	loss is the scalar loss returned to the caller for reduced logging.
tuple[Any, Tensor | None]	The default raises when criterion is missing or returns None;
tuple[Any, Tensor | None]	overrides may return (pred, None) to signal no loss available, in
tuple[Any, Tensor | None]	which case the caller skips loss bookkeeping.

Raises:

Type	Description
ValueError	self.model is not initialized, or criterion is missing or returned None.

Side effects: moves data to self.device, runs forward under train_context() (autocast + optional DDP no-sync), updates self.metrics when bound, then calls self.backward(loss) and self.step() to scale gradients, advance accumulation state, and flush the optimizer on accumulation boundaries.

Do not

• Zero gradients (optimizer_step does this on flush).
• Call self.optimizer.step() directly (use self.step()).
• Mutate train_state.global_step or train_state.micro_step.
• Implement gradient scaling here (override _backward() instead).
• Call save_checkpoint() (cadence is owned by the loop method).

Backend notes:

• DeepSpeedRunner inherits the default; _backward/step route through the DeepSpeed engine.
• ParallelRunner overrides this method when a pipeline schedule is set; the schedule owns micro-batching and loss reduction.

Source code in danling/runners/torch_runner.py

def train_step(self, data: Any) -> tuple[Any, torch.Tensor | None]:
    """
    Run one training micro-step.

    The default implementation runs forward → loss → metric update → backward
    → step for one micro-batch.

    **Called when:** once per micro-batch by `train_epoch`/`train_steps`. The
    caller seeds the loop's accumulation state before each invocation; this
    method consumes that state through `backward()` and `step()`.

    **Precondition:** `self.model`, `self.optimizer`, and `self.criterion`
    are bound; `self.mode == RunnerMode.train`.

    Args:
        data: One micro-batch. The default unpacks `data["input"]` /
            `data.get("target")` for mappings, `(data[0], data[1])` for
            non-string sequences, and `(data, None)` otherwise. Override
            `train_step` if your batch shape differs.

    Returns:
        `(pred, loss)`. `pred` is the model output (used by `metrics.update`).
        `loss` is the scalar loss returned to the caller for reduced logging.
        The default raises when `criterion` is missing or returns `None`;
        overrides may return `(pred, None)` to signal no loss available, in
        which case the caller skips loss bookkeeping.

    Raises:
        ValueError: `self.model` is not initialized, or `criterion` is missing
            or returned `None`.

    **Side effects:** moves `data` to `self.device`, runs forward under
    `train_context()` (autocast + optional DDP no-sync), updates
    `self.metrics` when bound, then calls `self.backward(loss)` and
    `self.step()` to scale gradients, advance accumulation state, and flush
    the optimizer on accumulation boundaries.

    !!! danger "Do not"
        - Zero gradients (`optimizer_step` does this on flush).
        - Call `self.optimizer.step()` directly (use `self.step()`).
        - Mutate `train_state.global_step` or `train_state.micro_step`.
        - Implement gradient scaling here (override `_backward()` instead).
        - Call `save_checkpoint()` (cadence is owned by the loop method).

    **Backend notes:**

    - `DeepSpeedRunner` inherits the default; `_backward`/`step` route
      through the DeepSpeed engine.
    - `ParallelRunner` overrides this method when a pipeline schedule is
      set; the schedule owns micro-batching and loss reduction.
    """
    data = self.to_device(data)
    with self.train_context():
        if isinstance(data, Mapping):
            inputs = data["input"]
            target = data.get("target")
        elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
            inputs = data[0]
            target = data[1] if len(data) > 1 else None
        else:
            inputs = data
            target = None

        if self.model is None:
            raise ValueError("cannot run train_step: model is not initialized")
        pred = self.model(**inputs) if isinstance(inputs, Mapping) else self.model(inputs)
        loss = self._compute_loss(pred, target)
        if loss is None:
            raise ValueError("cannot run train_step: criterion did not produce a loss")
        skipped = self.backward(loss)
        if not skipped and self.metrics is not None and pred is not None and target is not None:
            self.metrics.update(pred, target)
        self.step()
    return pred, loss.detach().new_zeros(()) if skipped else loss

backward

Python

backward(loss: Tensor) -> bool

Backpropagate one loss, or a synchronized zero when any replica is non-finite.

Returns:

Type	Description
bool	Whether this micro-batch was replaced with zero.

Source code in danling/runners/torch_runner.py

def backward(self, loss: torch.Tensor) -> bool:
    """Backpropagate one loss, or a synchronized zero when any replica is non-finite.

    Returns:
        Whether this micro-batch was replaced with zero.
    """

    zero_weight = self._pending_loss_normalizer == 0
    globally_nonfinite = False
    fault_tolerant = self.fault_tolerance is not None and self.fault_tolerance.enabled
    nonfinite = False
    if self.skip_nonfinite_loss or fault_tolerant:
        # Every rank votes, including an empty rank whose criterion may
        # legitimately be undefined. Zero weight makes that vote finite.
        nonfinite = not zero_weight and not bool(torch.isfinite(loss.detach()).all().item())
    if self.skip_nonfinite_loss:
        globally_nonfinite = self._sync_optimizer_skip_decision(nonfinite)
    if fault_tolerant and self.fault_tolerance is not None and (nonfinite or globally_nonfinite):
        self.fault_tolerance.record_nonfinite_loss()
    if zero_weight or globally_nonfinite:
        self._backward_zero_vjp(loss, preserve_missing=globally_nonfinite)
        return True
    self._backward(loss)
    return False

step

Python

step() -> None

Advance the accumulation state machine after one training micro-step.

Called when: train_step finishes backward for a micro-batch.

Side effects: increments train_state.micro_step and calls optimizer_step() only when the accumulation boundary is reached or the surrounding loop marks the current batch as the final flush in a partial window.

Do not

• Call this from evaluation/inference paths.
• Call optimizer_step() in addition to this method from the same micro-step.
• Adjust train_state.micro_step in train_step overrides.

Source code in danling/runners/torch_runner.py

def step(self) -> None:
    """
    Advance the accumulation state machine after one training micro-step.

    **Called when:** `train_step` finishes backward for a micro-batch.

    **Side effects:** increments `train_state.micro_step` and calls
    `optimizer_step()` only when the accumulation boundary is reached or
    the surrounding loop marks the current batch as the final flush in a
    partial window.

    !!! danger "Do not"
        - Call this from evaluation/inference paths.
        - Call `optimizer_step()` in addition to this method from the same
          micro-step.
        - Adjust `train_state.micro_step` in `train_step` overrides.
    """
    micro_steps = self.train_state.micro_step + 1
    self.train_state.micro_step = micro_steps
    if self._train_window_will_flush:
        self.optimizer_step()
        remainder = micro_steps % self.accum_steps
        if self.accum_steps > 1 and remainder != 0:
            self.train_state.micro_step += self.accum_steps - remainder
        return
    if self.accum_steps <= 1 or micro_steps % self.accum_steps == 0:
        self.optimizer_step()

optimizer_step

Python

optimizer_step() -> bool

Perform one backend optimizer update.

The default Torch implementation waits for checkpoint staging, applies accumulated-loss gradient scaling, optional grad clipping, non-finite grad skip logic, optimizer/scheduler stepping through OptimizerContainer, gradient zeroing, profiler advancement, and garbage-collection cadence.

Called when: step() reaches an accumulation boundary, or _flush_pending_optimizer_step() flushes a partial boundary before shutdown.

Returns:

Type	Description
bool	True when an optimizer update is applied, otherwise False.

Side effects: may update optimizer/scheduler state; increments train_state.global_step only when an update is actually applied.

Do not

• Increment global_step on skipped updates.
• Forget to zero gradients after a successful update or skipped non-finite update.
• Bypass checkpoint_manager.maybe_wait_for_staging().

Backend notes:

• DeepSpeedRunner overrides this hook because the DeepSpeed engine owns the concrete optimizer update.

Source code in danling/runners/torch_runner.py

def optimizer_step(self) -> bool:
    """
    Perform one backend optimizer update.

    The default Torch implementation waits for checkpoint staging, applies
    accumulated-loss gradient scaling, optional grad clipping, non-finite
    grad skip logic, optimizer/scheduler stepping through
    `OptimizerContainer`, gradient zeroing, profiler advancement, and
    garbage-collection cadence.

    **Called when:** `step()` reaches an accumulation boundary, or
    `_flush_pending_optimizer_step()` flushes a partial boundary before
    shutdown.

    Returns:
        `True` when an optimizer update is applied, otherwise `False`.

    **Side effects:** may update optimizer/scheduler state; increments
    `train_state.global_step` only when an update is actually applied.

    !!! danger "Do not"
        - Increment `global_step` on skipped updates.
        - Forget to zero gradients after a successful update or skipped
          non-finite update.
        - Bypass `checkpoint_manager.maybe_wait_for_staging()`.

    **Backend notes:**

    - `DeepSpeedRunner` overrides this hook because the DeepSpeed engine
      owns the concrete optimizer update.
    """
    if self.optimizer_container is None and self.optimizer is None:
        raise ValueError(
            "cannot perform optimizer step: no optimizer is configured; "
            "set `self.optimizer`, implement `build_optimizer()`, or override `optimizer_step()`"
        )

    self.validate_ema_update_contract()
    self._begin_optimizer_window()
    context = self.diagnostics.span("checkpoint.staging_wait") if self.diagnostics is not None else nullcontext()
    with context:
        self.checkpoint_manager.maybe_wait_for_staging()
    self._sync_fault_tolerant_gradients()
    if self.grad_scaler is not None and self.optimizer is not None:
        self.grad_scaler.unscale_(self.optimizer)
    grad_scale = self._gradient_scale_for_step()
    if grad_scale is not None:
        self._scale_optimizer_gradients(grad_scale)
    empty_window = self._accumulation_is_empty()
    max_grad_value = self.max_grad_value
    max_grad_norm = self.max_grad_norm
    skip_nonfinite_grad = self.skip_nonfinite_grad
    fault_tolerant = self.fault_tolerance is not None and self.fault_tolerance.enabled
    has_nonfinite_grad = False
    if skip_nonfinite_grad or self.grad_scaler is not None or fault_tolerant:
        has_nonfinite_grad = self._optimizer_has_nonfinite_grad()
        has_nonfinite_grad = self._sync_optimizer_skip_decision(has_nonfinite_grad)
    commit = not has_nonfinite_grad and not empty_window
    if fault_tolerant and self.fault_tolerance is not None:
        commit = self.fault_tolerance.should_commit(nonfinite=has_nonfinite_grad) and not empty_window
    if not commit:
        self._zero_optimizer_grad()
        self._update_grad_scaler(overflow=has_nonfinite_grad, committed=False)
        self._reset_accumulation_normalization()
        return False

    grad_norm = self._clip_optimizer_gradients(max_grad_value=max_grad_value, max_grad_norm=max_grad_norm)
    self._record_grad_norm(grad_norm)
    self._optimizer_step_and_zero_grad()
    self.update_ema()

    self._reset_accumulation_normalization()
    self.train_state.global_step += 1
    self._step_profiler()
    self._maybe_reduce_train_process_group_timeout()
    self.supervisor.maybe_collect_garbage(self.train_state.global_step, scope="train")
    return True

can_checkpoint_for_shutdown

Python

can_checkpoint_for_shutdown() -> bool

A fetched batch and its accumulation window must be consumed before saving.

Source code in danling/runners/torch_runner.py

def can_checkpoint_for_shutdown(self) -> bool:
    """A fetched batch and its accumulation window must be consumed before saving."""
    return getattr(self, "_train_checkpoint_ready", True)

synchronize_termination_signal

Python

synchronize_termination_signal(signum: int | None) -> int | None

Let one signaled rank stop its torchrun replica at a shared safe boundary.

Source code in danling/runners/torch_runner.py

def synchronize_termination_signal(self, signum: int | None) -> int | None:
    """Let one signaled rank stop its torchrun replica at a shared safe boundary."""
    if not (dist.is_available() and dist.is_initialized()):
        return signum
    # TorchFT replicas have independent default process groups. A shutdown
    # vote must not wait for its managed group or stop the other replicas.
    backend = str(dist.get_backend()).lower()
    device = torch.device("cuda", torch.cuda.current_device()) if "nccl" in backend else torch.device("cpu")
    pending = torch.tensor(0 if signum is None else signum, dtype=torch.int32, device=device)
    dist.all_reduce(pending, op=dist.ReduceOp.MAX)
    return int(pending.item()) or None

train

Python

train(train_splits: list[str] | None = None, evaluate_splits: list[str] | None = None) -> RoundDict

Run the full training workflow.

Selects epoch mode or step mode from self.is_step_mode, validates explicit split lists against the runner’s configured/inferred splits, and delegates to train_epochs or train_steps.

Called when: user code starts training.

Parameters:

Name	Type	Description	Default
train_splits	list[str] | None	Optional training splits. When None, use self.train_splits.	None
evaluate_splits	list[str] | None	Optional evaluation splits. When None, use self.evaluate_splits.	None

Returns:

Type	Description
RoundDict	Aggregated runner results (self.results).

Raises:

Type	Description
ValueError	no valid training split can be resolved.

Side effects: prints selected splits and runs the selected training loop. Checkpointing, result writing, scheduler stepping, and early stop are owned by the delegated loop method.

Source code in danling/runners/torch_runner.py

def train(
    self,
    train_splits: list[str] | None = None,
    evaluate_splits: list[str] | None = None,
) -> RoundDict:
    """
    Run the full training workflow.

    Selects epoch mode or step mode from `self.is_step_mode`, validates
    explicit split lists against the runner's configured/inferred splits,
    and delegates to `train_epochs` or `train_steps`.

    **Called when:** user code starts training.

    Args:
        train_splits: Optional training splits. When `None`, use `self.train_splits`.
        evaluate_splits: Optional evaluation splits. When `None`, use `self.evaluate_splits`.

    Returns:
        Aggregated runner results (`self.results`).

    Raises:
        ValueError: no valid training split can be resolved.

    **Side effects:** prints selected splits and runs the selected training
    loop. Checkpointing, result writing, scheduler stepping, and early stop
    are owned by the delegated loop method.
    """

    train_splits = self._resolve_requested_splits(train_splits, self.train_splits, kind="training")
    if not train_splits:
        raise ValueError("cannot start training: no valid training split was resolved")

    evaluate_splits = self._resolve_requested_splits(evaluate_splits, self.evaluate_splits, kind="evaluation")

    self.log_startup_summary(train_splits=train_splits, evaluate_splits=evaluate_splits)
    try:
        if self.is_step_mode:
            result = self.train_steps(train_splits=train_splits, evaluate_splits=evaluate_splits)
        else:
            result = self.train_epochs(train_splits=train_splits, evaluate_splits=evaluate_splits)
    except Exception as exc:
        self.log_crash_summary(exc)
        raise
    self.log_final_summary()
    return result

train_epochs

Python

train_epochs(train_splits: list[str] | None = None, evaluate_splits: list[str] | None = None) -> RoundDict

Run epoch-mode training until self.epochs is reached.

Each epoch runs all train splits, then all evaluation splits, advances epoch/metric schedulers, appends and writes results, saves periodic checkpoints on ckpt.interval, and refreshes the best checkpoint whenever the score improves.

Called when: train dispatches while config.epochs is set, or user code explicitly wants epoch-mode semantics.

Parameters:

Name	Type	Description	Default
train_splits	list[str] | None	Training splits for each epoch.	None
evaluate_splits	list[str] | None	Evaluation splits after each epoch.	None

Returns:

Type	Description
RoundDict	Aggregated runner results (self.results).

Raises:

Type	Description
ValueError	config.epochs is not set.

Source code in danling/runners/torch_runner.py

def train_epochs(
    self,
    train_splits: list[str] | None = None,
    evaluate_splits: list[str] | None = None,
) -> RoundDict:
    """
    Run epoch-mode training until `self.epochs` is reached.

    Each epoch runs all train splits, then all evaluation splits, advances
    epoch/metric schedulers, appends and writes results, saves periodic
    checkpoints on `ckpt.interval`, and refreshes the best checkpoint
    whenever the score improves.

    **Called when:** `train` dispatches while `config.epochs` is set, or
    user code explicitly wants epoch-mode semantics.

    Args:
        train_splits: Training splits for each epoch.
        evaluate_splits: Evaluation splits after each epoch.

    Returns:
        Aggregated runner results (`self.results`).

    Raises:
        ValueError: `config.epochs` is not set.
    """
    if train_splits is None:
        train_splits = self.train_splits
    if evaluate_splits is None:
        evaluate_splits = self.evaluate_splits

    total_epochs = self.epochs
    if total_epochs is None:
        raise ValueError("cannot run epoch-mode training: config.epochs is not set")
    # Even a completed resume needs healing and membership before final
    # collective checkpoint/export; no batch is fetched by this handshake.
    self._begin_optimizer_window()
    print(f"train: epoch mode start epoch={self.train_state.epoch} total_epochs={total_epochs}")
    periodic = self._training_evaluation_enabled
    early_stop_counter = 0
    patience = self.patience
    for epoch in range(self.train_state.epoch, total_epochs):
        self.supervisor.maybe_handle_termination_signal()
        self.train_state.epoch = epoch
        result = RoundDict()
        with self._training_evaluation_splits(evaluate_splits):
            for split in train_splits:
                result[split] = self.train_epoch(split)
                self.supervisor.maybe_handle_termination_signal()
        if periodic:
            self._maybe_evaluate_during_training(evaluate_splits, force=True, checkpoint=False)
            row = self.results.get(self.train_state.global_step, {})
            result.update({split: row[split] for split in evaluate_splits if split in row})
        else:
            for split in evaluate_splits:
                result[split] = self.evaluate_epoch(split)
                self.supervisor.maybe_handle_termination_signal()
        self._step_epoch_scheduler(result)
        index = self.train_state.global_step if periodic else epoch
        self.append_result(result, index=index)
        print(self.format_epoch_result(result, epochs=epoch, total_epochs=total_epochs))
        self.save_result()
        self.train_state.epoch = epoch + 1
        # Call every epoch: the manager persists periodic history on cadence and
        # publishes `best.pth` whenever the score improved this epoch.
        if periodic:
            self._save_evaluation_checkpoint(epochs=epoch, at_epoch_end=True)
        else:
            self.save_checkpoint(epochs=epoch)
        early_stop_counter = 0 if self.is_best else early_stop_counter + 1
        if early_stop_counter > patience:
            print("train: early-stop triggered")
            break
    self.save_checkpoint(last_step=True)
    self.save_model_checkpoint()
    return self.results

train_epoch

Python

train_epoch(split: str = 'train') -> RoundDict

Run one full dataloader pass for a training split.

This is the per-split epoch loop. It sets train mode, resets meters and train metrics, manages accumulation-window normalization, invokes train_step for each micro-batch, emits step logs, and records interval/epoch telemetry.

Called when: train_epochs processes one train split.

Parameters:

Name	Type	Description	Default
split	str	Training split name.	'train'

Returns:

Type	Description
RoundDict	Epoch-level metric mapping for this split.

Side effects: updates optimizer state through train_step, advances train_state.global_step on optimizer flushes, and writes step logs.

Do not

• Call this for evaluation data; use evaluate_epoch.
• Override this just to change one batch’s forward/loss logic; override train_step.
• Manually manage gradient zeroing inside train_step; this loop and optimizer_step own accumulation boundaries.
• Increment train_state.epoch; the surrounding train_epochs loop owns epoch progress.
• Save result or checkpoint aliases here; train_epochs owns epoch-level persistence.

See Also

train_steps: Step-mode counterpart that consumes splits against a global step budget instead of one epoch per split.

Source code in danling/runners/torch_runner.py

def train_epoch(self, split: str = "train") -> RoundDict:
    """
    Run one full dataloader pass for a training split.

    This is the per-split epoch loop. It sets train mode, resets meters and
    train metrics, manages accumulation-window normalization, invokes
    `train_step` for each micro-batch, emits step logs, and records
    interval/epoch telemetry.

    **Called when:** `train_epochs` processes one train split.

    Args:
        split: Training split name.

    Returns:
        Epoch-level metric mapping for this split.

    **Side effects:** updates optimizer state through `train_step`,
    advances `train_state.global_step` on optimizer flushes, and writes
    step logs.

    !!! danger "Do not"
        - Call this for evaluation data; use `evaluate_epoch`.
        - Override this just to change one batch's forward/loss logic;
          override `train_step`.
        - Manually manage gradient zeroing inside `train_step`; this loop
          and `optimizer_step` own accumulation boundaries.
        - Increment `train_state.epoch`; the surrounding `train_epochs`
          loop owns epoch progress.
        - Save result or checkpoint aliases here; `train_epochs` owns
          epoch-level persistence.

    See Also:
        [`train_steps`][danling.runners.TorchRunner.train_steps]:
            Step-mode counterpart that consumes splits against a global
            step budget instead of one epoch per split.
    """
    loader = self.dataloaders[split]
    loader_length = self._loader_length(loader)
    length = loader_length - 1 if loader_length is not None else None
    last_loss: torch.Tensor | None = None
    last_loss_n: float | None = None
    self._set_loader_epoch(loader, self.train_state.epoch)
    self.mode = RunnerMode.train
    self.split = split
    self.meters.reset()
    self.metrics = self.train_metrics
    if self.metrics is not None:
        self.metrics.reset()
    telemetry = LoopTelemetry(self, start_time=self.loop_time())
    self._reset_accumulation_normalization()
    if self.optimizer_container is not None:
        self.optimizer_container.zero_grad()
    elif self.optimizer is not None:
        self.optimizer.zero_grad()

    for iteration, data, will_flush in self._iter_train_batches(loader):
        self.supervisor.maybe_handle_termination_signal()
        step_before = self.train_state.global_step
        loss_n = self._get_loss_normalizer(data)
        self._pending_loss_normalizer = loss_n
        self._pending_loss_reduction = self._get_loss_reduction(data)
        self._train_window_will_flush = will_flush
        try:
            _, loss = self.train_step(data)
        finally:
            self._train_window_will_flush = False
            self._pending_loss_normalizer = None
            self._pending_loss_reduction = None

        loss = self._loss_for_reporting(loss, loss_n, data)
        self._train_checkpoint_ready = will_flush
        self.supervisor.mark_heartbeat_progress()
        self.supervisor.maybe_handle_termination_signal()
        current_time = self.loop_time()
        if self.scheduler is not None and hasattr(self.scheduler, "get_last_lr"):
            self.meters.lr.update(self.scheduler.get_last_lr()[0])
        if loss is not None:
            self._update_loss_meter(loss, loss_n)
        telemetry.observe(iteration=iteration, data=data, current_time=current_time)

        if self.train_state.global_step > step_before:
            splits = getattr(self, "_active_evaluation_splits", self.evaluate_splits)
            self._maybe_evaluate_during_training(splits)

        if self.log_interval > 0 and (
            (iteration > 0 and iteration % self.log_interval == 0) or iteration == length
        ):
            telemetry.emit_log(split=split, iteration=iteration, length=length, loss=loss, loss_n=loss_n)
        last_loss = loss
        last_loss_n = loss_n

    if (
        length is None
        and self.log_interval > 0
        and telemetry.last_iteration is not None
        and telemetry.last_iteration != telemetry.last_print_iteration
    ):
        assert telemetry.last_iteration is not None
        telemetry.emit_log(
            split=split,
            iteration=telemetry.last_iteration,
            length=length,
            loss=last_loss,
            loss_n=last_loss_n,
            reset_peak_stats=False,
        )
    result = self.get_epoch_result()
    telemetry.finalize_result(result, elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time)
    return result

train_steps

Python

train_steps(train_splits: list[str] | None = None, evaluate_splits: list[str] | None = None) -> RoundDict

Run step-mode training for the configured global step budget.

Step mode consumes train splits in sorted split order until train_state.global_step >= self.steps, then optionally evaluates configured evaluation splits with evaluate_steps.

Called when: train dispatches while config.epochs is unset, or user code explicitly wants a global-step budget.

Parameters:

Name	Type	Description	Default
train_splits	list[str] | None	Training splits to consume in order.	None
evaluate_splits	list[str] | None	Evaluation splits to run after training steps finish.	None

Returns:

Type	Description
RoundDict	Aggregated runner results (self.results).

Raises:

Type	Description
ValueError	total step budget cannot be resolved.

Side effects: updates epoch as an outer split-round counter, appends one result row indexed by global_step, writes result files, and saves the final checkpoint.

Do not

• Assume a split is consumed exactly once; step mode can resume a split iterator across outer rounds.
• Mutate train_state.global_step outside optimizer stepping.

See Also

train_epoch: Per-split epoch loop used by epoch-mode training.

Source code in danling/runners/torch_runner.py

def train_steps(
    self,
    train_splits: list[str] | None = None,
    evaluate_splits: list[str] | None = None,
) -> RoundDict:
    """
    Run step-mode training for the configured global step budget.

    Step mode consumes train splits in sorted split order until
    `train_state.global_step >= self.steps`, then optionally evaluates
    configured evaluation splits with `evaluate_steps`.

    **Called when:** `train` dispatches while `config.epochs` is unset, or
    user code explicitly wants a global-step budget.

    Args:
        train_splits: Training splits to consume in order.
        evaluate_splits: Evaluation splits to run after training steps finish.

    Returns:
        Aggregated runner results (`self.results`).

    Raises:
        ValueError: total step budget cannot be resolved.

    **Side effects:** updates epoch as an outer split-round counter,
    appends one result row indexed by `global_step`, writes result files,
    and saves the final checkpoint.

    !!! danger "Do not"
        - Assume a split is consumed exactly once; step mode can resume a
          split iterator across outer rounds.
        - Mutate `train_state.global_step` outside optimizer stepping.

    See Also:
        [`train_epoch`][danling.runners.TorchRunner.train_epoch]:
            Per-split epoch loop used by epoch-mode training.
    """
    if train_splits is None:
        train_splits = self.train_splits
    if evaluate_splits is None:
        evaluate_splits = self.evaluate_splits

    total_steps = self.steps
    if total_steps is None:
        raise ValueError("cannot run step-mode training: config.steps could not be resolved")
    self._begin_optimizer_window()
    print(f"train: step mode start global_step={self.train_state.global_step} steps={total_steps}")
    periodic = self._training_evaluation_enabled
    result = RoundDict()
    step_mode_iterators: dict[str, Iterator[tuple[int, Any, bool]] | None] = dict.fromkeys(train_splits)
    step_mode_sampler_epochs = self.train_state.sampler_epochs
    for split in train_splits:
        step_mode_sampler_epochs.setdefault(split, self.train_state.epoch)
    while self.train_state.global_step < total_steps:
        self.supervisor.maybe_handle_termination_signal()
        round_start_step = self.train_state.global_step
        plan_epoch = self.train_state.epoch
        restart_round = False
        round_finished_restored_split = False
        round_result = RoundDict()
        total_train_splits = len(train_splits)
        for split_index, split in enumerate(train_splits):
            if split_index < self.train_state.step_mode_split_index:
                continue
            self.supervisor.maybe_handle_termination_signal()
            self.mode = RunnerMode.train
            self.split = split
            remaining = total_steps - self.train_state.global_step
            if remaining <= 0:
                break
            loader = self.dataloaders[split]
            remaining_splits = total_train_splits - split_index
            split_steps = self._step_mode_split_budget(
                remaining_steps=remaining,
                remaining_splits=remaining_splits,
                loader=loader,
            )
            if split_steps <= 0:
                break
            start_global_step = self.train_state.global_step
            target_global_step = self.train_state.step_mode_target or start_global_step + split_steps
            # A checkpoint can be saved by the last update of a split,
            # before the loop clears its target and advances the round.
            round_finished_restored_split |= target_global_step <= start_global_step
            self.train_state.step_mode_split_index = split_index
            self.train_state.step_mode_target = target_global_step
            self.meters.reset()
            self.metrics = self.train_metrics
            if self.metrics is not None:
                self.metrics.reset()
            telemetry = LoopTelemetry(self, start_time=self.loop_time())
            self._reset_accumulation_normalization()
            if self.optimizer_container is not None:
                self.optimizer_container.zero_grad()
            elif self.optimizer is not None:
                self.optimizer.zero_grad()
            checkpoint_cadence = self.checkpoint_interval
            batch_iteration = -1

            while self.train_state.global_step < target_global_step:
                self.supervisor.maybe_handle_termination_signal()
                generation = getattr(self, "_dataloader_state_generation", 0)
                # Quorum can restore newer progress as well as model and
                # loader state. Revalidate this plan before any iterator
                # creation or fetch; the generator's begin is idempotent.
                self._begin_optimizer_window()
                if getattr(self, "_dataloader_state_generation", 0) != generation:
                    # Healing can reload only some local ranks, even when
                    # their shared training plan is unchanged. Invalidate
                    # local loader caches without changing collective loop
                    # ordering on the ranks that did not need recovery.
                    step_mode_iterators = dict.fromkeys(train_splits)
                    step_mode_sampler_epochs = self.train_state.sampler_epochs
                    for train_split in train_splits:
                        step_mode_sampler_epochs.setdefault(train_split, self.train_state.epoch)
                if (
                    self.train_state.global_step >= total_steps
                    or self.train_state.epoch != plan_epoch
                    or self.train_state.step_mode_split_index != split_index
                    or self.train_state.step_mode_target != target_global_step
                ):
                    restart_round = True
                    break
                if self.train_state.global_step >= target_global_step:
                    break
                batch: tuple[int, Any, bool] | None = None
                iterator = step_mode_iterators[split]
                recreated = 0
                while True:
                    if iterator is None:
                        # A restored cursor can already be at the end of
                        # its pass before this process creates an iterator.
                        # Allow that pass to finish, then try one fresh pass.
                        if recreated >= 2:
                            break
                        self._set_loader_epoch(loader, step_mode_sampler_epochs[split])
                        iterator = self._iter_train_batches(loader)
                        step_mode_iterators[split] = iterator
                        recreated += 1
                    try:
                        batch = next(iterator)
                        break
                    except StopIteration:
                        iterator = None
                        step_mode_iterators[split] = None
                        step_mode_sampler_epochs[split] += 1
                if batch is None:
                    break
                _, data, will_flush = batch
                batch_iteration += 1
                self.supervisor.maybe_handle_termination_signal()
                step_before = self.train_state.global_step
                loss_n = self._get_loss_normalizer(data)
                self._pending_loss_normalizer = loss_n
                self._pending_loss_reduction = self._get_loss_reduction(data)
                self._train_window_will_flush = will_flush
                try:
                    _, loss = self.train_step(data)
                finally:
                    self._train_window_will_flush = False
                    self._pending_loss_normalizer = None
                    self._pending_loss_reduction = None

                loss = self._loss_for_reporting(loss, loss_n, data)
                self._train_checkpoint_ready = will_flush
                self.supervisor.mark_heartbeat_progress()
                self.supervisor.maybe_handle_termination_signal()
                current_time = self.loop_time()
                if self.scheduler is not None and hasattr(self.scheduler, "get_last_lr"):
                    self.meters.lr.update(self.scheduler.get_last_lr()[0])
                if loss is not None:
                    self._update_loss_meter(loss, loss_n)
                telemetry.observe(iteration=batch_iteration, data=data, current_time=current_time)

                step_after = self.train_state.global_step
                evaluated = step_after > step_before and self._maybe_evaluate_during_training(evaluate_splits)
                if (
                    not evaluated
                    and checkpoint_cadence > 0
                    and step_after != step_before
                    and step_after % checkpoint_cadence == 0
                ):
                    self.save_checkpoint(save_best=False) if periodic else self.save_checkpoint()

                global_step = step_after if step_after != step_before else None
                is_log_step = (
                    self.log_interval > 0
                    and global_step is not None
                    and global_step > 0
                    and global_step % self.log_interval == 0
                )
                is_boundary_step = global_step in (target_global_step, total_steps)
                if self.log_interval > 0 and global_step is not None and (is_log_step or is_boundary_step):
                    telemetry.emit_log(
                        split=split,
                        iteration=batch_iteration,
                        length=total_steps,
                        loss=loss,
                        loss_n=loss_n,
                        display_iteration=global_step,
                    )

            if restart_round:
                break
            round_result[split] = self.get_epoch_result()
            telemetry.finalize_result(
                round_result[split], elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time
            )
            self.train_state.step_mode_split_index = split_index + 1
            self.train_state.step_mode_target = 0
            self.supervisor.maybe_handle_termination_signal()

        if restart_round:
            step_mode_iterators = dict.fromkeys(train_splits)
            step_mode_sampler_epochs = self.train_state.sampler_epochs
            for split in train_splits:
                step_mode_sampler_epochs.setdefault(split, self.train_state.epoch)
            continue
        if self.train_state.global_step == round_start_step and not round_finished_restored_split:
            remaining_steps = total_steps - self.train_state.global_step
            warn(
                f"step-mode training made no progress after one full split pass "
                f"(target={total_steps}, reached={self.train_state.global_step}, remaining={remaining_steps})",
                RuntimeWarning,
                stacklevel=2,
            )
            break
        self._step_epoch_scheduler(round_result)
        result = round_result
        self.train_state.epoch += 1
        self.train_state.step_mode_split_index = 0
        self.train_state.step_mode_target = 0
    remaining_steps = total_steps - self.train_state.global_step
    if remaining_steps > 0:
        warn(
            f"step-mode training finished with {remaining_steps} step(s) remaining "
            f"(target={total_steps}, reached={self.train_state.global_step})",
            RuntimeWarning,
            stacklevel=2,
        )
    if periodic:
        self._maybe_evaluate_during_training(evaluate_splits, force=True, checkpoint=False)
        row = self.results.get(self.train_state.global_step, {})
        result.update({split: row[split] for split in evaluate_splits if split in row})
    else:
        for split in evaluate_splits:
            result[split] = self.evaluate_steps(split=split)
    self.append_result(result, index=self.train_state.global_step)
    self.save_result()
    self.save_checkpoint(last_step=True)
    self.save_model_checkpoint()
    return self.results

evaluate_step

Python

evaluate_step(data: Any) -> tuple[Any, Tensor | None]

Run one evaluation micro-step.

The default implementation runs forward → optional loss → optional metric update under infer_context(). No backward pass and no optimizer step.

Called when: once per micro-batch by evaluate_epoch/evaluate_steps, which run under torch.inference_mode().

Precondition: at least one of self.model or self.ema is bound. self.mode == RunnerMode.evaluate. The default prefers self.ema over self.model when both are available.

Parameters:

Name	Type	Description	Default
data	Any	One micro-batch. The default unpacks data["input"] / data.get("target") for mappings, (data[0], data[1]) for non-string sequences, and (data, None) otherwise. Override evaluate_step if your batch shape differs.	required

Returns:

Type	Description
Any	(pred, loss). pred is the model output (used by metrics.update).
Tensor | None	loss is the scalar loss returned to the caller for reduced
tuple[Any, Tensor | None]	logging, or None when no criterion is set.

Raises:

Type	Description
ValueError	neither self.model nor self.ema is initialized.

Side effects: moves data to self.device, runs forward through self.ema or self.model under infer_context(), computes loss when criterion is set, and updates self.metrics when bound.

Do not

• Call self.backward(...) or self.step() (no optimizer here).
• Mutate train_state.global_step or train_state.micro_step.
• Switch the runner mode (the loop owns self.mode).
• Call save_checkpoint() (cadence is owned by training loops only).

Backend notes:

• ParallelRunner overrides this method when a pipeline schedule is set; the schedule owns micro-batching and pipeline-stage loss reduction.

Source code in danling/runners/torch_runner.py

def evaluate_step(self, data: Any) -> tuple[Any, torch.Tensor | None]:
    """
    Run one evaluation micro-step.

    The default implementation runs forward → optional loss → optional
    metric update under `infer_context()`. No backward pass and no optimizer step.

    **Called when:** once per micro-batch by `evaluate_epoch`/`evaluate_steps`,
    which run under `torch.inference_mode()`.

    **Precondition:** at least one of `self.model` or `self.ema` is bound.
    `self.mode == RunnerMode.evaluate`. The default prefers `self.ema` over
    `self.model` when both are available.

    Args:
        data: One micro-batch. The default unpacks `data["input"]` /
            `data.get("target")` for mappings, `(data[0], data[1])` for
            non-string sequences, and `(data, None)` otherwise. Override
            `evaluate_step` if your batch shape differs.

    Returns:
        `(pred, loss)`. `pred` is the model output (used by `metrics.update`).
        `loss` is the scalar loss returned to the caller for reduced
        logging, or `None` when no `criterion` is set.

    Raises:
        ValueError: neither `self.model` nor `self.ema` is initialized.

    **Side effects:** moves `data` to `self.device`, runs forward through
    `self.ema or self.model` under `infer_context()`, computes loss when
    `criterion` is set, and updates `self.metrics` when bound.

    !!! danger "Do not"
        - Call `self.backward(...)` or `self.step()` (no optimizer here).
        - Mutate `train_state.global_step` or `train_state.micro_step`.
        - Switch the runner mode (the loop owns `self.mode`).
        - Call `save_checkpoint()` (cadence is owned by training loops only).

    **Backend notes:**

    - `ParallelRunner` overrides this method when a pipeline schedule is
      set; the schedule owns micro-batching and pipeline-stage loss
      reduction.
    """
    data = self.to_device(data)
    if isinstance(data, Mapping):
        inputs = data["input"]
        target = data.get("target")
    elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
        inputs = data[0]
        target = data[1] if len(data) > 1 else None
    else:
        inputs = data
        target = None

    if self.model is None and self.ema is None:
        raise ValueError("cannot run evaluate_step: model is not initialized")
    model = self.ema or self.model
    with self.infer_context():
        pred = model(**inputs) if isinstance(inputs, Mapping) else model(inputs)
        loss = self._compute_loss(pred, target)

    if self.metrics is not None and pred is not None and target is not None:
        self.metrics.update(pred, target)

    return pred, loss

evaluate

Python

evaluate(evaluate_splits: list[str] | None = None) -> RoundDict

Run evaluation across splits with epoch-mode semantics.

Called when: user code explicitly evaluates a runner, or training code delegates to evaluation helpers.

Parameters:

Name	Type	Description	Default
evaluate_splits	list[str] | None	Optional evaluation splits. When None, use self.evaluate_splits.	None

Returns:

Type	Description
RoundDict	Mapping of split -> evaluation result for this call.

Raises:

Type	Description
ValueError	no valid evaluation split can be resolved.

Side effects: sets evaluation mode per split, prints a formatted aggregate result, and writes scalar outputs through evaluate_epoch.

Source code in danling/runners/torch_runner.py

def evaluate(self, evaluate_splits: list[str] | None = None) -> RoundDict:
    """
    Run evaluation across splits with epoch-mode semantics.

    **Called when:** user code explicitly evaluates a runner, or training
    code delegates to evaluation helpers.

    Args:
        evaluate_splits: Optional evaluation splits. When `None`, use `self.evaluate_splits`.

    Returns:
        Mapping of split -> evaluation result for this call.

    Raises:
        ValueError: no valid evaluation split can be resolved.

    **Side effects:** sets evaluation mode per split, prints a formatted
    aggregate result, and writes scalar outputs through `evaluate_epoch`.
    """

    evaluate_splits = self._resolve_requested_splits(evaluate_splits, self.evaluate_splits, kind="evaluation")
    if not evaluate_splits:
        raise ValueError("cannot start evaluation: no valid evaluation split was resolved")
    print("evaluate: start")
    print(f"evaluate: splits={evaluate_splits}")
    result = RoundDict()
    for split in evaluate_splits:
        result[split] = self.evaluate_epoch(split=split)
    display_epoch = self.train_state.epoch
    if self.epochs is not None and display_epoch > 0:
        display_epoch -= 1
    print(self.format_epoch_result(result, epochs=display_epoch))
    return result

evaluate_epoch

Python

evaluate_epoch(split: str = 'val') -> RoundDict

Run one full dataloader pass for an evaluation split.

Sets evaluation mode, resets meters/evaluation metrics, runs evaluate_step for every batch under inference mode, emits step logs, and writes the split result at the current epoch index.

Called when: evaluate or train_epochs evaluates a split.

Parameters:

Name	Type	Description	Default
split	str	Evaluation split name.	'val'

Returns:

Type	Description
RoundDict	Epoch-level metric mapping for this split.

Side effects: updates evaluation meters/metrics, emits logs, writes scalar results, and records telemetry. It does not update optimizer or training progress counters.

Source code in danling/runners/torch_runner.py

@torch.inference_mode()
def evaluate_epoch(self, split: str = "val") -> RoundDict:
    """
    Run one full dataloader pass for an evaluation split.

    Sets evaluation mode, resets meters/evaluation metrics, runs
    `evaluate_step` for every batch under inference mode, emits step logs,
    and writes the split result at the current epoch index.

    **Called when:** `evaluate` or `train_epochs` evaluates a split.

    Args:
        split: Evaluation split name.

    Returns:
        Epoch-level metric mapping for this split.

    **Side effects:** updates evaluation meters/metrics, emits logs, writes
    scalar results, and records telemetry. It does not update optimizer or
    training progress counters.
    """
    loader = self.dataloaders[split]
    loader_length = self._loader_length(loader)
    length = loader_length - 1 if loader_length is not None else None

    last_loss: torch.Tensor | None = None
    last_loss_n: float | None = None
    self.mode = RunnerMode.evaluate
    self.split = split
    self.meters = AverageMeters()
    self.metrics = self.evaluate_metrics
    if self.metrics is not None:
        self.metrics.reset()
    telemetry = LoopTelemetry(self, start_time=self.loop_time())
    consumed = 0
    for iteration, data in enumerate(loader):
        consumed = iteration + 1
        self.supervisor.maybe_handle_termination_signal()
        loss_n = self._get_loss_normalizer(data)
        if loss_n is not None and loss_n < 0:
            loss_n = None
        _, loss = self.evaluate_step(data)
        loss = self._loss_for_reporting(loss, loss_n, data)
        self.supervisor.mark_heartbeat_progress()
        self.supervisor.maybe_handle_termination_signal()
        current_time = self.loop_time()
        if loss is not None:
            self._update_loss_meter(loss, loss_n)
        telemetry.observe(
            iteration=iteration,
            data=data,
            current_time=current_time,
        )
        self.supervisor.maybe_collect_garbage(iteration + 1, scope=f"evaluate:{split}")

        if self.log_interval > 0 and (
            (iteration > 0 and iteration % self.log_interval == 0) or iteration == length
        ):
            telemetry.emit_log(split=split, iteration=iteration, length=length, loss=loss, loss_n=loss_n)
        last_loss = loss
        last_loss_n = loss_n

    if (
        length is None
        and self.log_interval > 0
        and telemetry.last_iteration is not None
        and telemetry.last_iteration != telemetry.last_print_iteration
    ):
        assert telemetry.last_iteration is not None
        telemetry.emit_log(
            split=split,
            iteration=telemetry.last_iteration,
            length=length,
            loss=last_loss,
            loss_n=last_loss_n,
            reset_peak_stats=False,
        )
    result = self.get_epoch_result()
    telemetry.finalize_result(
        result, elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time, steps=consumed
    )
    self.write_result(result, split, getattr(self, "_evaluation_result_index", self.train_state.epoch))
    return result

evaluate_steps

Python

evaluate_steps(split: str = 'val', steps: int | None = None) -> RoundDict

Run bounded evaluation steps on one split.

Used by step-mode training to evaluate a small fixed number of batches without requiring a full evaluation pass.

Called when: train_steps evaluates configured splits after the step budget finishes, or user code requests bounded evaluation.

Parameters:

Name	Type	Description	Default
split	str	Evaluation split name.	'val'
steps	int | None	Number of batches to evaluate. When None, defaults to max(self.steps // 20, 1).	None

Returns:

Type	Description
RoundDict	Step-bounded evaluation metrics.

Raises:

Type	Description
ValueError	step budget cannot be inferred, steps is negative, or the dataloader exhausts before the requested number of steps.

Side effects: writes scalar results at train_state.global_step.

Source code in danling/runners/torch_runner.py

@torch.inference_mode()
def evaluate_steps(self, split: str = "val", steps: int | None = None) -> RoundDict:
    """
    Run bounded evaluation steps on one split.

    Used by step-mode training to evaluate a small fixed number of batches
    without requiring a full evaluation pass.

    **Called when:** `train_steps` evaluates configured splits after the
    step budget finishes, or user code requests bounded evaluation.

    Args:
        split: Evaluation split name.
        steps: Number of batches to evaluate. When `None`, defaults to `max(self.steps // 20, 1)`.

    Returns:
        Step-bounded evaluation metrics.

    Raises:
        ValueError: step budget cannot be inferred, `steps` is negative, or
            the dataloader exhausts before the requested number of steps.

    **Side effects:** writes scalar results at `train_state.global_step`.
    """
    if steps is None:
        total_steps = self.steps
        if total_steps is None:
            raise ValueError("cannot infer evaluation steps: step budget is unavailable; pass `steps`")
        steps = max(total_steps // 20, 1)
    if steps < 0:
        raise ValueError(f"invalid steps: expected a non-negative value, got {steps}")
    loader = self.dataloaders[split]
    length = steps - 1

    self.mode = RunnerMode.evaluate
    self.split = split
    if steps == 0:
        self.meters = AverageMeters()
        self.metrics = self.evaluate_metrics
        if self.metrics is not None:
            self.metrics.reset()
        result = self.get_epoch_result()
        self.write_result(result, split, self.train_state.global_step)
        return result

    self.meters = AverageMeters()
    self.metrics = self.evaluate_metrics
    if self.metrics is not None:
        self.metrics.reset()
    telemetry = LoopTelemetry(self, start_time=self.loop_time())
    consumed = 0
    batches = loader if steps is None else islice(loader, steps)
    for iteration, data in enumerate(batches):
        consumed = iteration + 1
        self.supervisor.maybe_handle_termination_signal()
        loss_n = self._get_loss_normalizer(data)
        if loss_n is not None and loss_n < 0:
            loss_n = None
        _, loss = self.evaluate_step(data)
        loss = self._loss_for_reporting(loss, loss_n, data)
        self.supervisor.mark_heartbeat_progress()
        self.supervisor.maybe_handle_termination_signal()
        current_time = self.loop_time()
        if loss is not None:
            self._update_loss_meter(loss, loss_n)
        telemetry.observe(iteration=iteration, data=data, current_time=current_time)
        self.supervisor.maybe_collect_garbage(iteration + 1, scope=f"evaluate:{split}")

        if self.log_interval > 0 and (
            (iteration > 0 and iteration % self.log_interval == 0) or iteration == length
        ):
            telemetry.emit_log(split=split, iteration=iteration, length=length, loss=loss, loss_n=loss_n)

    if steps is not None and consumed < steps:
        raise ValueError(
            f"evaluate steps exhausted early on split '{split}': requested {steps} step(s), got {consumed}"
        )
    result = self.get_epoch_result()
    telemetry.finalize_result(
        result, elapsed_seconds=self.loop_time(sync=True) - telemetry.start_time, steps=consumed
    )
    self.write_result(result, split, self.train_state.global_step)
    return result

infer_step

Python

infer_step(data: Any) -> list[float]

Run one inference micro-step.

The default implementation runs forward through self.ema or self.model, detaches scalar-per-example predictions, squeezes the trailing dimension, moves them to CPU, and returns them as a Python list.

Called when: once per micro-batch by infer/_iter_infer_batches. The method is decorated with torch.inference_mode().

Precondition: at least one of self.model or self.ema is bound. self.mode == RunnerMode.infer.

Parameters:

Name	Type	Description	Default
data	Any	One micro-batch. The default unpacks data["input"] for mappings, data[0] for non-string sequences, and data itself otherwise. Override infer_step if your batch shape differs or you need to pass auxiliary tensors to the model.	required

Returns:

Type	Description
list[float]	List of CPU floats for scalar-per-example predictions. The
list[float]	default converts with pred.squeeze(-1).detach().cpu().tolist().
list[float]	Override if your model emits multi-dim tensors, mappings, or
list[float]	non-numeric outputs.

Raises:

Type	Description
ValueError	neither self.model nor self.ema is initialized.

Side effects: moves data to self.device, runs forward through self.ema or self.model under infer_context(), then converts the output to a CPU list.

Do not

• Compute or accumulate metrics (inference is metric-free).
• Mutate runner state counters.
• Return a torch.Tensor (callers expect list[float] for batched aggregation and streaming).
• Call self.backward(...) or self.step().

Backend notes:

• ParallelRunner overrides this method when a pipeline schedule is set; non-first-stage ranks pass data=None and the schedule routes activations through pipeline communication.

Source code in danling/runners/torch_runner.py

@torch.inference_mode()
def infer_step(self, data: Any) -> list[float]:
    """
    Run one inference micro-step.

    The default implementation runs forward through `self.ema or self.model`,
    detaches scalar-per-example predictions, squeezes the trailing
    dimension, moves them to CPU, and returns them as a Python list.

    **Called when:** once per micro-batch by `infer`/`_iter_infer_batches`.
    The method is decorated with `torch.inference_mode()`.

    **Precondition:** at least one of `self.model` or `self.ema` is bound.
    `self.mode == RunnerMode.infer`.

    Args:
        data: One micro-batch. The default unpacks `data["input"]` for
            mappings, `data[0]` for non-string sequences, and `data`
            itself otherwise. Override `infer_step` if your batch shape
            differs or you need to pass auxiliary tensors to the model.

    Returns:
        List of CPU floats for scalar-per-example predictions. The
        default converts with `pred.squeeze(-1).detach().cpu().tolist()`.
        Override if your model emits multi-dim tensors, mappings, or
        non-numeric outputs.

    Raises:
        ValueError: neither `self.model` nor `self.ema` is initialized.

    **Side effects:** moves `data` to `self.device`, runs forward through
    `self.ema or self.model` under `infer_context()`, then converts the
    output to a CPU list.

    !!! danger "Do not"
        - Compute or accumulate metrics (inference is metric-free).
        - Mutate runner state counters.
        - Return a `torch.Tensor` (callers expect `list[float]` for
          batched aggregation and streaming).
        - Call `self.backward(...)` or `self.step()`.

    **Backend notes:**

    - `ParallelRunner` overrides this method when a pipeline schedule is
      set; non-first-stage ranks pass `data=None` and the schedule routes
      activations through pipeline communication.
    """
    data = self.to_device(data)
    if isinstance(data, Mapping):
        inputs = data["input"]
    elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
        inputs = data[0]
    else:
        inputs = data

    if self.model is None and self.ema is None:
        raise ValueError("cannot run infer_step: model is not initialized")
    model = self.ema or self.model
    with self.infer_context():
        pred = model(**inputs) if isinstance(inputs, Mapping) else model(inputs)
    values = pred.squeeze(-1).detach().cpu().tolist()
    if isinstance(values, list):
        return values
    return [float(values)]

infer

Python

infer(split: str = 'infer', *, steps: int | None = None, stream: bool | None = None) -> list[float] | Iterator[list[float]]

Run inference on one split.

In non-stream mode this consumes all requested batches and returns a flattened Python list. In stream mode it returns an iterator of per-batch outputs and leaves consumption to the caller.

Called when: user code requests prediction-only execution.

Parameters:

Name	Type	Description	Default
split	str	Inference split name.	'infer'
steps	int | None	Optional max number of batches to consume.	None
stream	bool | None	True returns a generator of per-batch outputs, False returns a flattened list. When None, stream only for unsized loaders without explicit steps.	None

Returns:

Type	Description
list[float] | Iterator[list[float]]	Flattened predictions or a streaming iterator of batch predictions.

Raises:

Type	Description
ValueError	steps is negative, or non-stream inference is requested for an unsized loader without an explicit step count.

Side effects: sets inference mode/split. It does not update metrics or optimizer state.

Source code in danling/runners/torch_runner.py

def infer(
    self,
    split: str = "infer",
    *,
    steps: int | None = None,
    stream: bool | None = None,
) -> list[float] | Iterator[list[float]]:
    """
    Run inference on one split.

    In non-stream mode this consumes all requested batches and returns a
    flattened Python list. In stream mode it returns an iterator of
    per-batch outputs and leaves consumption to the caller.

    **Called when:** user code requests prediction-only execution.

    Args:
        split: Inference split name.
        steps: Optional max number of batches to consume.
        stream: `True` returns a generator of per-batch outputs, `False` returns a flattened list.
            When `None`, stream only for unsized loaders without explicit `steps`.

    Returns:
        Flattened predictions or a streaming iterator of batch predictions.

    Raises:
        ValueError: `steps` is negative, or non-stream inference is
            requested for an unsized loader without an explicit step count.

    **Side effects:** sets inference mode/split. It does not update metrics
    or optimizer state.
    """

    self.mode = RunnerMode.infer
    self.split = split
    loader = self.dataloaders[split]
    if steps is not None and steps < 0:
        raise ValueError(f"invalid steps: expected a non-negative value, got {steps}")

    loader_length = self._loader_length(loader)
    if stream is None:
        stream = steps is None and loader_length is None

    if not stream and loader_length is None and steps is None:
        raise ValueError("infer with stream=False requires `steps` for unsized loaders")

    iterator = self._iter_infer_batches(loader, steps=steps, split=split)
    if stream:
        return iterator

    total = steps if steps is not None else loader_length
    output: list = []
    for values in tqdm(iterator, total=total, disable=self.distributed and not self.is_main_process):
        output.extend(values)
    if self.distributed:
        output = self._gather_infer_predictions(loader, output)
    return output

state_dict

Python

state_dict(cls: type = dict) -> Mapping

Return the TorchRunner checkpoint payload.

Extends BaseRunner.state_dict with backend metadata plus EMA, optimizer, scheduler, and unwrapped model state.

Called when: checkpoint managers persist a TorchRunner checkpoint.

Parameters:

Name	Type	Description	Default
cls	type	Mapping factory used for nested payloads.	dict

Returns:

Type	Description
Mapping	Mapping containing base runner state and torch component state.

Side effects: snapshots Python/NumPy/Torch RNG state before exporting.

Source code in danling/runners/torch_runner.py

def state_dict(self, cls: type = dict) -> Mapping:
    """
    Return the TorchRunner checkpoint payload.

    Extends `BaseRunner.state_dict` with backend metadata plus EMA,
    optimizer, scheduler, and unwrapped model state.

    **Called when:** checkpoint managers persist a TorchRunner checkpoint.

    Args:
        cls: Mapping factory used for nested payloads.

    Returns:
        Mapping containing base runner state and torch component state.

    **Side effects:** snapshots Python/NumPy/Torch RNG state before
    exporting.
    """
    state = cls(super().state_dict(cls))
    state.update(self._export_checkpoint_metadata(cls))
    state.update(self._export_checkpoint_components(cls))
    if self.fault_tolerance is not None and self.fault_tolerance.enabled:
        state["fault_tolerance"] = self.fault_tolerance.state_dict()
    return state

load_state_dict

Python

load_state_dict(checkpoint: Mapping[str, Any]) -> None

Restore base runner state plus Torch RNG state.

Model, optimizer, scheduler, and dataloader components are restored by load_checkpoint; this method owns only runner/RNG state.

Source code in danling/runners/torch_runner.py

def load_state_dict(self, checkpoint: Mapping[str, Any]) -> None:
    """
    Restore base runner state plus Torch RNG state.

    Model, optimizer, scheduler, and dataloader components are restored by
    `load_checkpoint`; this method owns only runner/RNG state.
    """
    super().load_state_dict(checkpoint)
    if self.fault_tolerance is not None and "fault_tolerance" in checkpoint:
        self.fault_tolerance.load_state_dict(checkpoint["fault_tolerance"])
    if self.grad_scaler is not None:
        scaler_state = checkpoint.get("grad_scaler")
        if scaler_state is not None:
            self.grad_scaler.load_state_dict(self._require_checkpoint_component_state("grad_scaler", scaler_state))
    state_dict = checkpoint.get("state") or {}
    rng_state = state_dict.get("rng")
    if isinstance(rng_state, Mapping) and "torch_cpu" in rng_state and self.rng_state.torch_cpu is not None:
        torch.set_rng_state(self.rng_state.torch_cpu)
    if (
        torch.cuda.is_available()
        and isinstance(rng_state, Mapping)
        and "torch_cuda" in rng_state
        and self.rng_state.torch_cuda is not None
    ):
        torch.cuda.set_rng_state_all(self.rng_state.torch_cuda)

load_checkpoint

Python

load_checkpoint(checkpoint: Mapping | bytes | str | PathLike, *args, **kwargs) -> None

Load a full checkpoint and rebind optimizer/scheduler helpers.

This delegates component restore to BaseRunner.load_checkpoint, then rebuilds the OptimizerContainer so scheduler and optimizer state stay bound after restore.

Source code in danling/runners/torch_runner.py

def load_checkpoint(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> None:
    """
    Load a full checkpoint and rebind optimizer/scheduler helpers.

    This delegates component restore to `BaseRunner.load_checkpoint`, then
    rebuilds the `OptimizerContainer` so scheduler and optimizer state stay
    bound after restore.
    """
    super().load_checkpoint(checkpoint, *args, **kwargs)
    self._bind_optimizer_container()

read_checkpoint

Python

read_checkpoint(checkpoint: Mapping | bytes | str | PathLike, *args, **kwargs) -> Mapping[str, Any]

Read checkpoint payload from mapping/file/DCP directory input.

Source code in danling/runners/torch_runner.py

def read_checkpoint(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> Mapping[str, Any]:
    """Read checkpoint payload from mapping/file/DCP directory input."""
    if isinstance(checkpoint, Mapping):
        return checkpoint

    checkpoint_path = os.fsdecode(checkpoint)
    if str(self.config.get("ckpt.backend")).lower() == "dcp" and (
        checkpoint_path.endswith(".pointer") or not os.path.isfile(checkpoint_path)
    ):
        return self.checkpoint_manager.load_checkpoint(checkpoint)
    return super().read_checkpoint(checkpoint, *args, **kwargs)

read_config classmethod

Python

read_config(checkpoint: Mapping | bytes | str | PathLike, *args, **kwargs) -> RunnerConfig

Read runner config from checkpoint payload, including DCP directory inputs.

Source code in danling/runners/torch_runner.py

@classmethod
def read_config(
    cls,
    checkpoint: Mapping | bytes | str | os.PathLike,
    *args,
    **kwargs,
) -> RunnerConfig:
    """Read runner config from checkpoint payload, including DCP directory inputs."""
    if isinstance(checkpoint, Mapping):
        return super().read_config(checkpoint, *args, **kwargs)

    if TorchDistributedCheckpointManager.is_checkpoint_path(checkpoint):
        return TorchDistributedCheckpointManager.read_config(checkpoint)

    return super().read_config(checkpoint, *args, **kwargs)

close

Python

close(timeout: float | None = None) -> bool

Close runner resources.

Source code in danling/runners/torch_runner.py

def close(self, timeout: float | None = None) -> bool:
    """Close runner resources."""
    try:
        self._close_profiler()
        drained = super().close(timeout=timeout)
    except Exception:
        self._close_profiler()
        if self.diagnostics is not None:
            self.diagnostics.close()
        self.destroy_process_group()
        raise
    if not drained:
        return False
    if self.diagnostics is not None:
        self.diagnostics.close()
    self.destroy_process_group()
    return drained