diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index d01eba53f1..49e2c33ab5 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -548,7 +548,7 @@ topics: - id: supervised-learning title: Supervised Learning aliases: [supervised learning, SL learning, 监督学习, 统一监督学习] - keywords: [train-sl, SupervisedBatch, LossResult, SupervisedObjective, SupervisedLearner, SupervisedTrainer, TensorBatchSource] + keywords: [train-sl, SupervisedBatch, LossResult, SupervisedObjective, SupervisedLearner, SupervisedTrainer, TensorBatchSource, distributed_run, torchrun, DDP, 单机多卡, 分布式监督学习] paths: [topics/supervised-learning/supervised-learning.md] source_of_truth: - embodichain/learning/sl/ diff --git a/agent_context/topics/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md index 982c189bd7..aa7f565b7b 100644 --- a/agent_context/topics/supervised-learning/supervised-learning.md +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -19,6 +19,11 @@ operations remain in `compute`; source storage belongs to data adapters. - `embodichain train-sl --config FILE [--resume FILE]` lazily resolves a `factory: module:function`. The factory receives the full model/data/objective/ optimizer/trainer/evaluation configuration and returns a `SupervisedTrainer`. + Under single-node `torchrun`, the workflow initializes a process group before + the factory and resolves `trainer.device` to CPU/Gloo or `cuda:LOCAL_RANK`/NCCL. + `distributed_run(device)` provides the same scoped setup for programmatic + callers; it cleans up only groups it creates. Launch commands and a complete + masked-regression factory live in `examples/learning/sl_ddp/`. - The [NeRD recipe](../neural-dynamics/neural-dynamics.md) is an installable factory at `embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`. The core does not import the model extension or example helpers. The NeRD @@ -26,8 +31,18 @@ operations remain in `compute`; source storage belongs to data adapters. Objectives own loss masks, normalization, integration and multi-step graph construction. The optimizer does not infer action semantics or detach histories. +Under DDP, `LossResult.loss` is a local mean and `weight` is its effective valid +supervision count; the learner forms the global weighted mean. Batch size is per +rank, the learning rate is unchanged, and rank-local sampling streams draw with +replacement, so samples may overlap. All ranks call fitting, saving and recovery; +only rank 0 writes artifacts or runs validation through the unwrapped model. Schema, action/frame/timestep and numerical-policy identifiers belong in each recipe's checkpoint signature. Tensor sources additionally bind content identity. +Distributed checkpoints retain each rank's RNG, sampler and buffers and require +the same world size for exact recovery. Inside rank 0's validation callback only, +a save uses the previously gathered committed-update snapshot without a new +collective; ordinary save calls require every rank to participate. + `learning/sl/_artifacts.py` owns run identities and output-history admission for both SL and historical NeRD trainers. A fresh run claims `run.json`; resuming an occupied directory requires the checkpoint's matching run identity. Checkpoints @@ -42,4 +57,6 @@ Acquisition and producer lifecycles remain application-owned. `tests/learning/sl/` checks independent masked updates, graph propagation, nonfinite update boundaries, owned samples, simulator-free imports and exact CPU recovery with dropout, optimizer and sampling state. Model-specific -qualification belongs to each model's acceptance recipe. +qualification belongs to each model's acceptance recipe. The distributed tests +exercise weighted global updates and fixed-world-size recovery with CPU/Gloo; +the example is validated by launch smoke runs rather than example-specific tests. diff --git a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst index 47aa97a49d..6251a7c10c 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -6,6 +6,89 @@ offline run lifecycle. Batches retain their tensor-tree structure and causal masks; source adapters, objectives and downstream execution define their semantics. Importing this package does not initialize simulation or reinforcement learning. +Single-node launch +------------------ + +Run the same recipe on one CPU or multiple workers from an installed repository +checkout. Each fresh run needs a separate, artifact-free ``trainer.output_dir``. +The example factory uses a small masked MLP regression corpus generated locally. + +.. code-block:: bash + + # One CPU worker. + python -m embodichain train-sl --config examples/learning/sl_ddp/cpu.yaml + + # Two CPU workers using Gloo; choose a fresh output directory in the recipe. + OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \ + -m embodichain train-sl --config examples/learning/sl_ddp/cpu.yaml + + # Two GPUs using NCCL, one worker per visible GPU. + CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=1 \ + torchrun --standalone --nnodes=1 --nproc-per-node=2 \ + -m embodichain train-sl --config examples/learning/sl_ddp/cuda.yaml + +For ``WORLD_SIZE > 1``, the workflow initializes the process group before invoking +the recipe factory. It resolves ``trainer.device`` to ``cpu`` or +``cuda:LOCAL_RANK``. ``trainer.batch_size`` is per worker; the global number of +drawn examples is ``batch_size * world_size``. Each worker samples independently +with replacement, so examples can overlap. The learning rate is not scaled +automatically. + +Objectives return a local mean loss and its valid supervision count in +``LossResult.weight``. The distributed learner differentiates the global +weighted mean ``sum(weight_rank * loss_rank) / sum(weight_rank)``. For masked +regression, that weight is the number of valid scalar targets, rather than the +number of padded tensor elements or an unweighted rank count. + +Only rank 0 writes training artifacts and runs validation. Validation receives +``learner.model`` as an unwrapped model so its forward pass does not require other +ranks. Each rank participates in ordinary ``fit()``, ``save()`` and ``resume()`` +calls. A save within rank 0's validation callback uses the already gathered +committed-update snapshot and does not start another collective. + +Distributed checkpoints preserve common optimizer/update state together with +each rank's RNG, source-sampling state and model buffers. Resume with the same +world size, recipe, source and training configuration: + +.. code-block:: bash + + OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \ + -m embodichain train-sl --config /tmp/sl-ddp-resume.yaml \ + --resume outputs/sl_ddp_cpu/step_100.pt + +In this example, copy ``cpu.yaml`` to ``/tmp/sl-ddp-resume.yaml`` and change only +``trainer.output_dir`` to a fresh destination. This preserves the completed +original run when demonstrating continuation from its older update 100 checkpoint. +An interrupted run can use its original directory when no later metrics would +be overwritten. See ``examples/learning/sl_ddp/README.md`` for the full recipe. + +Programmatic process-group scope +-------------------------------- + +The CLI owns process-group setup and cleanup automatically. Programmatic callers +can use ``distributed_run(device)`` around recipe construction and training; every +worker enters the scope. It yields the resolved ``torch.device`` and destroys +only a process group that it created. An already initialized group remains owned +by its caller. + +.. code-block:: python + + from pathlib import Path + import yaml + from embodichain.learning.sl import distributed_run + from examples.learning.sl_ddp.recipe import build_trainer + + config = yaml.safe_load(Path("examples/learning/sl_ddp/cpu.yaml").read_text()) + with distributed_run(config["trainer"]["device"]) as device: + config["trainer"]["device"] = str(device) + trainer = build_trainer(config) + trainer.fit() + +The corresponding PyTorch 2.7 references are +`torchrun `_ +and `DistributedDataParallel +`_. + embodichain.learning.sl ----------------------- @@ -21,6 +104,7 @@ embodichain.learning.sl TensorBatchSource SupervisedTrainer SupervisedTrainerCfg + distributed_run .. autoclass:: LossResult :members: @@ -46,6 +130,19 @@ embodichain.learning.sl .. autoclass:: SupervisedTrainerCfg :members: +.. autofunction:: distributed_run + +embodichain.learning.sl.distributed +----------------------------------- + +.. currentmodule:: embodichain.learning.sl.distributed + +.. autosummary:: + + distributed_run + +.. autofunction:: distributed_run + embodichain.learning.sl.learner ------------------------------- diff --git a/embodichain/learning/sl/__init__.py b/embodichain/learning/sl/__init__.py index d63ae13c3a..7bbe8974b3 100644 --- a/embodichain/learning/sl/__init__.py +++ b/embodichain/learning/sl/__init__.py @@ -19,6 +19,7 @@ from __future__ import annotations from .types import LossResult, SupervisedBatch, SupervisedObjective +from .distributed import distributed_run from .learner import SupervisedLearner from .sources import SupervisedSource, TensorBatchSource from .trainer import SupervisedTrainer, SupervisedTrainerCfg @@ -32,4 +33,5 @@ "TensorBatchSource", "SupervisedTrainer", "SupervisedTrainerCfg", + "distributed_run", ] diff --git a/embodichain/learning/sl/distributed.py b/embodichain/learning/sl/distributed.py new file mode 100644 index 0000000000..07c28dc773 --- /dev/null +++ b/embodichain/learning/sl/distributed.py @@ -0,0 +1,166 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Single-node process-group lifecycle and small SL control messages.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from datetime import timedelta +import math +import os +from typing import Any, TypeVar + +import torch +import torch.distributed as dist + +__all__ = ["distributed_run"] + +_T = TypeVar("_T") + + +@contextmanager +def distributed_run( + device: str | torch.device, *, timeout_seconds: float = 1800.0 +) -> Iterator[torch.device]: + """Bind one device and own a process group when launched by ``torchrun``. + + CUDA jobs use NCCL and ``cuda:LOCAL_RANK``; CPU jobs use Gloo. Ordinary + single-process calls need no process group. An existing group is borrowed + and is never destroyed here. All ranks must enter this context together. + + Args: + device: Recipe device. Under torchrun a CUDA index is replaced by the + local rank within ``CUDA_VISIBLE_DEVICES``. + timeout_seconds: Positive collective timeout, including rank-zero + validation time. It does not enable elastic world-size recovery. + + Yields: + The device to use before constructing a model and optimizer. + """ + if not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError("Distributed timeout must be finite and positive") + resolved = torch.device(device) + requested = int(os.environ.get("WORLD_SIZE", "1")) + if requested < 1: + raise ValueError("WORLD_SIZE must be positive") + if requested > 1: + if int(os.environ.get("LOCAL_WORLD_SIZE", str(requested))) != requested: + raise ValueError("SL torchrun currently supports a single node") + if resolved.type not in ("cpu", "cuda"): + raise ValueError("SL DDP supports CPU/Gloo or CUDA/NCCL") + if resolved.type == "cuda": + resolved = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + if resolved.type == "cuda": + resolved = torch.device( + "cuda", + ( + resolved.index + if resolved.index is not None + else torch.cuda.current_device() + ), + ) + torch.cuda.set_device(resolved) + owned = requested > 1 and not dist.is_initialized() + if owned: + dist.init_process_group( + backend="nccl" if resolved.type == "cuda" else "gloo", + timeout=timedelta(seconds=timeout_seconds), + ) + try: + yield resolved + finally: + if owned and dist.is_initialized(): + dist.destroy_process_group() + + +class _DistributedContext: + """Coordinate control boundaries; model gradients use native Torch DDP.""" + + def __init__(self, device: str | torch.device) -> None: + self.device = torch.device(device) + self.enabled = dist.is_initialized() and dist.get_world_size() > 1 + self.rank = dist.get_rank() if self.enabled else 0 + self.world_size = dist.get_world_size() if self.enabled else 1 + self.control_device = ( + self.device + if self.enabled and dist.get_backend() == "nccl" + else torch.device("cpu") + ) + + def check(self, error: Exception | None, stage: str) -> None: + """Propagate reachable local failures before the next collective phase.""" + if not self.enabled: + if error is not None: + raise error + return + failed = torch.tensor( + int(error is not None), dtype=torch.int32, device=self.control_device + ) + dist.all_reduce(failed, op=dist.ReduceOp.MAX) + if not failed.item(): + return + message = None if error is None else f"{type(error).__name__}: {error}" + messages: list[Any] = [None] * self.world_size + dist.all_gather_object(messages, message) + for rank, failure in enumerate(messages): + if failure is not None: + raise RuntimeError( + f"Distributed SL {stage} failed on rank {rank}: {failure}" + ) from error + + def run(self, function: Callable[[], _T], stage: str) -> _T: + error = None + result: Any = None + try: + result = function() + except Exception as caught: + error = caught + self.check(error, stage) + return result + + def root(self, function: Callable[[], _T], stage: str) -> _T: + if not self.enabled: + return function() + payload: list[Any] = [None] + if self.rank == 0: + try: + payload[0] = (function(), None) + except Exception as error: + payload[0] = (None, f"{type(error).__name__}: {error}") + dist.broadcast_object_list(payload, src=0) + value, failure = payload[0] + if failure is not None: + raise RuntimeError(f"Distributed SL {stage} failed on rank 0: {failure}") + return value + + def agree(self, value: Any, stage: str) -> None: + if not self.enabled: + return + values: list[Any] = [None] * self.world_size + dist.all_gather_object(values, value) + if any(item != values[0] for item in values[1:]): + raise ValueError(f"Distributed SL {stage} differs between ranks") + + def gather(self, value: Any) -> list[Any] | None: + if not self.enabled: + return [value] + # Keep the small control-state path on all-gather collectives, which + # are supported by both NCCL and Gloo. Non-root copies are temporary. + values: list[Any] = [None] * self.world_size + dist.all_gather_object(values, value) + return values if self.rank == 0 else None diff --git a/embodichain/learning/sl/learner.py b/embodichain/learning/sl/learner.py index 33af0e32a9..0c8c7e898c 100644 --- a/embodichain/learning/sl/learner.py +++ b/embodichain/learning/sl/learner.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Single-device optimizer updates shared by policy and dynamics fitting.""" +"""Single-device and distributed optimizer updates for policy and dynamics fitting.""" from __future__ import annotations @@ -24,13 +24,31 @@ from typing import Any import torch +import torch.distributed as dist from torch import nn +from torch.nn.parallel import DistributedDataParallel +from .distributed import _DistributedContext from .types import SupervisedBatch, SupervisedObjective __all__ = ["SupervisedLearner"] +class _ObjectiveModule(nn.Module): + """Wrap the complete graph while keeping model-specific APIs available.""" + + def __init__(self, model: nn.Module, objective: SupervisedObjective) -> None: + super().__init__() + self.model = model + object.__setattr__(self, "objective", objective) + + def forward( + self, batch: SupervisedBatch + ) -> tuple[torch.Tensor, Mapping[str, torch.Tensor | float], float]: + result = self.objective.compute_loss(self.model, batch) + return result.loss, result.metrics, result.weight + + class SupervisedLearner: """Own one optimizer while objectives own mathematical and graph semantics. @@ -43,6 +61,8 @@ class SupervisedLearner: objectives may explicitly leave autocast for FP32 state integration. clip_grad: Optional numerical-profile-specific clipping implementation. numerical_context: Context factory spanning forward, backward and clipping. + ddp_find_unused_parameters: Enable graph traversal for objectives that + omit trainable parameters. Applies only with an initialized group. Sampling, normalizer updates, learning-rate schedules and model publication belong to callers. An update is committed only after optimizer.step succeeds. @@ -60,6 +80,7 @@ def __init__( Callable[[Iterable[nn.Parameter], float], torch.Tensor] | None ) = None, numerical_context: Callable[[], AbstractContextManager[Any]] | None = None, + ddp_find_unused_parameters: bool = False, ) -> None: if not math.isfinite(max_grad_norm) or max_grad_norm < 0: raise ValueError("max_grad_norm must be finite and nonnegative") @@ -70,6 +91,22 @@ def __init__( self.clip_grad = clip_grad self.numerical_context = numerical_context or nullcontext self.num_updates = 0 + self.ddp_find_unused_parameters = ddp_find_unused_parameters + self._ddp: DistributedDataParallel | None = None + self._distributed_metrics: tuple[str, ...] | None = None + + def _configure_distributed(self) -> _DistributedContext: + context = _DistributedContext(next(self.model.parameters()).device) + if context.enabled and self._ddp is None: + self._ddp = DistributedDataParallel( + _ObjectiveModule(self.model, self.objective), + device_ids=( + [context.device.index] if context.device.type == "cuda" else None + ), + broadcast_buffers=False, + find_unused_parameters=self.ddp_find_unused_parameters, + ) + return context def update(self, batch: SupervisedBatch) -> dict[str, float | int]: """Backpropagate one objective and commit one finite parameter update. @@ -85,6 +122,9 @@ def update(self, batch: SupervisedBatch) -> dict[str, float | int]: FloatingPointError: Loss or gradient norm is not finite; optimizer parameters and the committed update counter are not advanced. """ + context = self._configure_distributed() + if context.enabled: + return self._distributed_update(batch, context) self.model.train() self.optimizer.zero_grad(set_to_none=True) parameter = next(self.model.parameters()) @@ -133,6 +173,96 @@ def update(self, batch: SupervisedBatch) -> dict[str, float | int]: ) return metrics + def _distributed_update( + self, batch: SupervisedBatch, context: _DistributedContext + ) -> dict[str, float | int]: + self.model.train() + self.optimizer.zero_grad(set_to_none=True) + autocast = ( + torch.autocast(context.device.type, dtype=self.amp_dtype) + if self.amp_dtype is not None + else nullcontext() + ) + with self.numerical_context(): + + def forward() -> tuple[torch.Tensor, dict[str, float], float]: + assert self._ddp is not None + with autocast: + loss, values, weight = self._ddp(batch) + if loss.ndim != 0 or not loss.requires_grad: + raise ValueError( + "A supervised objective must return a scalar loss with gradients" + ) + weight = float(weight) + if not math.isfinite(weight) or weight <= 0: + raise ValueError( + "Objective weight must be finite and positive on every rank" + ) + if not torch.isfinite(loss.detach()): + raise FloatingPointError("Non-finite supervised loss") + metrics = {name: float(value) for name, value in values.items()} + if not all(math.isfinite(value) for value in metrics.values()): + raise FloatingPointError("Non-finite supervised metric") + names = tuple(sorted(metrics)) + if ( + self._distributed_metrics is not None + and names != self._distributed_metrics + ): + raise ValueError("DDP objective metric names must remain constant") + return loss, metrics, weight + + loss, local_metrics, weight = context.run(forward, "objective") + names = tuple(sorted(local_metrics)) + if self._distributed_metrics is None: + context.agree(names, "objective metric names") + self._distributed_metrics = names + # Native DDP averages parameter gradients. Scale each local mean by + # W*w/sum(w) to obtain the mean over valid supervision, not ranks. + observations = torch.tensor( + [weight, float(loss.detach()) * weight] + + [local_metrics[name] * weight for name in names], + dtype=torch.float64, + device=context.control_device, + ) + dist.all_reduce(observations) + if not torch.isfinite(observations).all(): + raise FloatingPointError( + "Non-finite globally reduced supervised statistics" + ) + global_weight = float(observations[0]) + (loss * (context.world_size * weight / global_weight)).backward() + + def clip() -> float: + parameters = list(self.model.parameters()) + norm = ( + self.clip_grad(parameters, self.max_grad_norm) + if self.clip_grad is not None + else torch.nn.utils.clip_grad_norm_( + parameters, self.max_grad_norm or float("inf") + ) + ) + value = float(norm) + if not math.isfinite(value): + raise FloatingPointError("Non-finite supervised gradient norm") + return value + + norm = context.run(clip, "gradient clipping") + # A failed optimizer/communication operation aborts the job; no attempt + # is made to reuse a partially updated replica or broken process group. + context.run(self.optimizer.step, "optimizer update") + self.num_updates += 1 + reduced = observations.cpu().tolist() + metrics = { + name: reduced[index + 2] / global_weight for index, name in enumerate(names) + } + metrics.update( + loss=reduced[1] / global_weight, + grad_norm=norm, + weight=global_weight, + num_updates=self.num_updates, + ) + return metrics + def _optimizer_layout(self) -> list[list[str]]: names = { id(parameter): name for name, parameter in self.model.named_parameters() @@ -167,6 +297,12 @@ def load_state_dict(self, state: Mapping[str, Any]) -> None: Args: state: State produced by :meth:`state_dict`. """ + self._validate_state_dict(state) + self.model.load_state_dict(state["model"]) + self.optimizer.load_state_dict(state["optimizer"]) + self.num_updates = state["num_updates"] + + def _validate_state_dict(self, state: Mapping[str, Any]) -> None: current = self.model.state_dict() saved = state["model"] if current.keys() != saved.keys(): @@ -199,6 +335,3 @@ def load_state_dict(self, state: Mapping[str, Any]) -> None: for group, actual in zip(groups, self.optimizer.param_groups) ): raise ValueError("Supervised checkpoint optimizer parameter groups differ") - self.model.load_state_dict(saved) - self.optimizer.load_state_dict(state["optimizer"]) - self.num_updates = count diff --git a/embodichain/learning/sl/trainer.py b/embodichain/learning/sl/trainer.py index d3d46cc816..51069a65aa 100644 --- a/embodichain/learning/sl/trainer.py +++ b/embodichain/learning/sl/trainer.py @@ -28,6 +28,7 @@ import torch from embodichain.utils import configclass from ._artifacts import _RunArtifacts +from .distributed import _DistributedContext from .learner import SupervisedLearner from .sources import SupervisedSource @@ -36,11 +37,12 @@ @configclass class SupervisedTrainerCfg: - """Configure an offline, single-device supervised run. + """Configure an offline run, using one device per process under DDP. Args: total_updates: Absolute optimizer-update budget, including resumed work. - batch_size: Source examples per update. + batch_size: Source examples per rank and update. Global batch size under + DDP is this value times the process-group world size. device: Tensor destination; runtime adapters remain caller-owned. seed: Independent source-sampling seed. Initialize models before building the trainer using the recipe's explicitly chosen model seed. @@ -99,12 +101,62 @@ def __init__( json.dumps(dict(signature or {}), sort_keys=True, allow_nan=False) ) self.scheduler, self.evaluate = scheduler, evaluate - self.sampler = torch.Generator(device="cpu").manual_seed(self.cfg.seed) + self._context = _DistributedContext(self.cfg.device) + self.sampler = torch.Generator(device="cpu").manual_seed( + self.cfg.seed + self._context.rank + ) self.output = Path(self.cfg.output_dir) - self.output.mkdir(parents=True, exist_ok=True) self._artifacts = _RunArtifacts(self.output) self.learner.model.to(self.cfg.device) - self._validation_boundary: dict[str, Any] | None = None + self._distributed_prepared = False + self._evaluation_snapshot: list[Any] | None = None + + def _prepare_distributed(self) -> None: + if not self._context.enabled or self._distributed_prepared: + return + configuration = self.cfg.to_dict() + configuration["device"] = torch.device(self.cfg.device).type + configuration["output_dir"] = str(self.output.resolve()) + self._context.agree( + { + "config": configuration, + "signature": self.signature, + "source_id": self.source.source_id, + "objective": f"{type(self.learner.objective).__module__}.{type(self.learner.objective).__qualname__}", + "optimizer_type": f"{type(self.learner.optimizer).__module__}.{type(self.learner.optimizer).__qualname__}", + "optimizer_layout": self.learner._optimizer_layout(), + "optimizer_options": [ + { + key: ( + value.detach().cpu().tolist() + if isinstance(value, torch.Tensor) + else value + ) + for key, value in group.items() + if key != "params" + } + for group in self.learner.optimizer.param_groups + ], + "max_grad_norm": self.learner.max_grad_norm, + "amp_dtype": str(self.learner.amp_dtype), + "find_unused_parameters": self.learner.ddp_find_unused_parameters, + "scheduler_type": ( + None + if self.scheduler is None + else f"{type(self.scheduler).__module__}.{type(self.scheduler).__qualname__}" + ), + "validation": self.evaluate is not None, + }, + "recipe or dataset identity", + ) + self.learner._configure_distributed() + # Factories use a common seed for construction. After DDP synchronizes + # replicas, stochastic training and sampling use independent rank seeds. + seed = self.cfg.seed + self._context.rank + random.seed(seed) + np.random.seed(seed % 2**32) + torch.manual_seed(seed) + self._distributed_prepared = True def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: """Train up to the configured or an earlier committed update boundary. @@ -115,26 +167,37 @@ def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: Returns: Completed update count, final metrics and last-checkpoint path. """ + self._prepare_distributed() end = ( self.cfg.total_updates if stop_after is None else min(stop_after, self.cfg.total_updates) ) - if end < self.learner.num_updates: - raise ValueError("stop_after precedes the restored update cursor") - self._artifacts.claim() - self._artifacts.validate_history(self.learner.num_updates, key="num_updates") + self._context.agree((end, self.learner.num_updates), "update boundary") + + def admit_output() -> None: + if end < self.learner.num_updates: + raise ValueError("stop_after precedes the restored update cursor") + self._artifacts.claim() + self._artifacts.validate_history( + self.learner.num_updates, key="num_updates" + ) + + self._context.root(admit_output, "output admission") metrics = {} while self.learner.num_updates < end: candidate = torch.Generator(device="cpu") candidate.set_state(self.sampler.get_state()) - batch = self.source.sample(self.cfg.batch_size, generator=candidate).to( - self.cfg.device + batch = self._context.run( + lambda: self.source.sample(self.cfg.batch_size, generator=candidate).to( + self.cfg.device + ), + "sampling", ) metrics = self.learner.update(batch) self.sampler.set_state(candidate.get_state()) if self.scheduler is not None: - self.scheduler.step() + self._context.run(self.scheduler.step, "scheduler") step = self.learner.num_updates if ( self.evaluate is not None @@ -142,8 +205,12 @@ def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: and step % self.cfg.validation_interval == 0 ): metrics.update(self._evaluate_step()) - with (self.output / "metrics.jsonl").open("a") as handle: - handle.write(json.dumps(metrics, allow_nan=False) + "\n") + + def append_metrics() -> None: + with (self.output / "metrics.jsonl").open("a") as handle: + handle.write(json.dumps(metrics, allow_nan=False) + "\n") + + self._context.root(append_metrics, "metric logging") if ( self.cfg.checkpoint_interval and step % self.cfg.checkpoint_interval == 0 @@ -155,35 +222,56 @@ def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: "num_updates": self.learner.num_updates, "metrics": metrics, "checkpoint": str(path), + "world_size": self._context.world_size, + "global_batch_size": self.cfg.batch_size * self._context.world_size, } - def _evaluate_step(self) -> dict[str, float]: - boundary = { - "rng": self._rng_state(), + def _local_rank_state(self) -> dict[str, Any]: + return { "sampler": self.sampler.get_state(), + "rng": self._rng_state(), "buffers": { - name: value.detach().clone() - for name, value in self.learner.model.named_buffers() + name: buffer.detach().cpu().clone() + for name, buffer in self.learner.model.named_buffers() }, } - was_training = self.learner.model.training - self._validation_boundary = boundary - try: - self.learner.model.eval() - assert self.evaluate is not None - with torch.no_grad(): - return { - f"validation/{key}": float(value) - for key, value in self.evaluate(self.learner).items() - } - finally: - with torch.no_grad(): - for name, value in self.learner.model.named_buffers(): - value.copy_(boundary["buffers"][name]) - self.sampler.set_state(boundary["sampler"]) - self._restore_rng(boundary["rng"]) - self.learner.model.train(was_training) - self._validation_boundary = None + + def _gather_rank_states(self) -> list[Any] | None: + state = self._context.run(self._local_rank_state, "checkpoint rank state") + return self._context.gather(state) + + def _evaluate_step(self) -> dict[str, float]: + # Validation may publish a checkpoint from rank zero. Capture every + # rank's committed boundary first, so callback save() has no collectives. + snapshot = self._gather_rank_states() + + def evaluate() -> dict[str, float]: + rng = self._rng_state() + assert snapshot is not None + buffers = snapshot[0]["buffers"] + was_training = self.learner.model.training + self._evaluation_snapshot = snapshot + try: + self.learner.model.eval() + assert self.evaluate is not None + with torch.no_grad(): + metrics = { + f"validation/{key}": float(value) + for key, value in self.evaluate(self.learner).items() + } + if not all(math.isfinite(value) for value in metrics.values()): + raise FloatingPointError("Non-finite validation metric") + return metrics + finally: + self._evaluation_snapshot = None + with torch.no_grad(): + for name, value in self.learner.model.named_buffers(): + value.copy_(buffers[name]) + self.sampler.set_state(snapshot[0]["sampler"]) + self._restore_rng(rng) + self.learner.model.train(was_training) + + return self._context.root(evaluate, "validation") def _rng_state(self) -> dict[str, Any]: device = torch.device(self.cfg.device) @@ -206,12 +294,28 @@ def save(self, path: str | Path) -> None: Args: path: Local checkpoint destination, replaced atomically. + + All ranks call this method together, except inside the trainer's + rank-zero validation callback, where a captured boundary is available. """ + self._prepare_distributed() + if self._evaluation_snapshot is not None: + self._write_checkpoint(path, self._evaluation_snapshot) + return + self._context.agree(str(Path(path).resolve()), "checkpoint destination") + rank_states = self._gather_rank_states() if self._context.enabled else None + self._context.root( + lambda: self._write_checkpoint(path, rank_states), "checkpoint write" + ) + + def _write_checkpoint( + self, path: str | Path, rank_states: list[Any] | None + ) -> None: path = Path(path) self._artifacts.claim() path.parent.mkdir(parents=True, exist_ok=True) state = { - "version": 1, + "version": 2 if self._context.enabled else 1, "run_id": self._artifacts.run_id, "learner": self.learner.state_dict(), "sampler": self.sampler.get_state(), @@ -228,11 +332,18 @@ def save(self, path: str | Path) -> None: None if self.scheduler is None else self.scheduler.state_dict() ), } - if self._validation_boundary is not None: - boundary = self._validation_boundary - state["sampler"] = boundary["sampler"] - state["rng"] = boundary["rng"] - for name, value in boundary["buffers"].items(): + if self._context.enabled: + assert rank_states is not None + state["distributed"] = { + "world_size": self._context.world_size, + "rank_states": rank_states, + } + if rank_states is not None: + state["sampler"] = rank_states[0]["sampler"] + state["rng"] = rank_states[0]["rng"] + # An eval forward may maintain caches/counters. A callback save is + # still the pre-validation training boundary, including buffers. + for name, value in rank_states[0]["buffers"].items(): if name in state["learner"]["model"]: state["learner"]["model"][name] = value temporary = path.with_suffix(path.suffix + ".tmp") @@ -244,10 +355,54 @@ def resume(self, path: str | Path) -> None: Args: path: Trusted training checkpoint from this trainer. + + DDP recovery is collective and requires the original world size. """ - state = torch.load(path, map_location="cpu", weights_only=False) + self._context.agree(str(Path(path).resolve()), "resume checkpoint") + state = self._context.run( + lambda: torch.load(path, map_location="cpu", weights_only=False), + "checkpoint read", + ) + self._context.run( + lambda: self._validate_checkpoint(state), "checkpoint compatibility" + ) + + def admit_resume() -> str: + run_id = self._artifacts.validate_resume(state.get("run_id")) + self._artifacts.validate_history( + state["learner"]["num_updates"], key="num_updates" + ) + return run_id + + run_id = self._context.root(admit_resume, "resume output admission") + # Initial DDP broadcasts must happen before restoring rank-local buffers. + self._prepare_distributed() + + def restore() -> None: + self.learner.load_state_dict(state["learner"]) + if self.scheduler is not None: + self.scheduler.load_state_dict(state["scheduler"]) + if self._context.enabled: + local = state["distributed"]["rank_states"][self._context.rank] + with torch.no_grad(): + for name, buffer in self.learner.model.named_buffers(): + buffer.copy_(local["buffers"][name]) + else: + local = state + self.sampler.set_state(local["sampler"]) + self._restore_rng(local["rng"]) + + self._context.run(restore, "checkpoint restore") + self._artifacts.adopt_resume(run_id) + + def _validate_checkpoint(self, state: Mapping[str, Any]) -> None: + distributed = state.get("distributed", {}) + if distributed.get("world_size", 1) != self._context.world_size: + raise ValueError( + "Supervised checkpoint world size differs; exact recovery requires the same world size" + ) if ( - state.get("version") != 1 + state.get("version") != (2 if self._context.enabled else 1) or state["signature"] != self.signature or state["source_id"] != self.source.source_id ): @@ -256,6 +411,9 @@ def resume(self, path: str | Path) -> None: for key in ("output_dir",): current.pop(key, None) saved.pop(key, None) + if self._context.enabled: + current["device"] = torch.device(current["device"]).type + saved["device"] = torch.device(saved["device"]).type if current != saved or (state["scheduler"] is None) != (self.scheduler is None): raise ValueError("Supervised checkpoint training configuration differs") scheduler_type = ( @@ -265,13 +423,16 @@ def resume(self, path: str | Path) -> None: ) if state.get("scheduler_type") != scheduler_type: raise ValueError("Supervised checkpoint scheduler type differs") - run_id = self._artifacts.validate_resume(state.get("run_id")) - self._artifacts.validate_history( - state["learner"]["num_updates"], key="num_updates" - ) - self.learner.load_state_dict(state["learner"]) - if self.scheduler is not None: - self.scheduler.load_state_dict(state["scheduler"]) - self.sampler.set_state(state["sampler"]) - self._restore_rng(state["rng"]) - self._artifacts.adopt_resume(run_id) + self.learner._validate_state_dict(state["learner"]) + if self._context.enabled: + rank_states = distributed["rank_states"] + if len(rank_states) != self._context.world_size: + raise ValueError("Supervised checkpoint rank states differ") + buffers = dict(self.learner.model.named_buffers()) + local_buffers = rank_states[self._context.rank]["buffers"] + if buffers.keys() != local_buffers.keys() or any( + value.shape != local_buffers[name].shape + or value.dtype != local_buffers[name].dtype + for name, value in buffers.items() + ): + raise ValueError("Supervised checkpoint rank-local buffers differ") diff --git a/embodichain/learning/sl/types.py b/embodichain/learning/sl/types.py index 4d3a85ec74..00e58ebf25 100644 --- a/embodichain/learning/sl/types.py +++ b/embodichain/learning/sl/types.py @@ -84,8 +84,10 @@ class LossResult: Args: loss: Scalar connected to the parameters to optimize. metrics: Scalar diagnostics; they do not enter backward automatically. - weight: Effective observation count for evaluation aggregation. Loss - reduction and sequence weighting remain the objective's decision. + weight: Effective observation count. Under DDP, ``loss`` and metrics + must be local means over this positive count; the learner uses it + to produce a global weighted mean and the corresponding gradient. + Single-device loss reduction remains the objective's decision. """ loss: torch.Tensor diff --git a/embodichain/learning/sl/workflow.py b/embodichain/learning/sl/workflow.py index 9d6c68275f..fa6620fb64 100644 --- a/embodichain/learning/sl/workflow.py +++ b/embodichain/learning/sl/workflow.py @@ -20,11 +20,16 @@ import argparse from collections.abc import Sequence +import copy import importlib import json +import os from pathlib import Path from typing import Any +import torch.distributed as dist + +from .distributed import _DistributedContext, distributed_run from .trainer import SupervisedTrainer __all__ = ["train_from_config", "cli"] @@ -48,12 +53,37 @@ def train_from_config( module, separator, name = factory.partition(":") if not separator or not module or not name: raise ValueError("SL configuration requires factory: module:function") - trainer = getattr(importlib.import_module(module), name)(config) - if not isinstance(trainer, SupervisedTrainer): - raise TypeError("SL recipe factory must return a SupervisedTrainer") - if resume is not None: - trainer.resume(resume) - return trainer.fit() + prepared = copy.deepcopy(config) + device = prepared.get("trainer", {}).get("device", "cpu") + trainer: SupervisedTrainer | None = None + with distributed_run(device) as resolved: + context = _DistributedContext(resolved) + if context.enabled: + prepared["trainer"] = { + **prepared.get("trainer", {}), + "device": str(resolved), + } + + def build() -> SupervisedTrainer: + nonlocal trainer + candidate = getattr(importlib.import_module(module), name)(prepared) + if not isinstance(candidate, SupervisedTrainer): + raise TypeError("SL recipe factory must return a SupervisedTrainer") + trainer = candidate + return candidate + + try: + trainer = context.run(build, "recipe factory") + if resume is not None: + trainer.resume(resume) + return trainer.fit() + finally: + # Factories own acquisition, but this entry point owns the returned + # source's consumption lifetime, including background prefetch work. + if trainer is not None: + close = getattr(trainer.source, "close", None) + if close is not None: + close() def cli(argv: Sequence[str] | None = None) -> None: @@ -71,7 +101,12 @@ def cli(argv: Sequence[str] | None = None) -> None: config = yaml.safe_load(args.config.read_text()) if not isinstance(config, dict): raise ValueError("SL configuration must be a mapping") - print(json.dumps(train_from_config(config, resume=args.resume), allow_nan=False)) + rank = ( + dist.get_rank() if dist.is_initialized() else int(os.environ.get("RANK", "0")) + ) + result = train_from_config(config, resume=args.resume) + if rank == 0: + print(json.dumps(result, allow_nan=False)) if __name__ == "__main__": diff --git a/examples/learning/nerd/_support/_workflow.py b/examples/learning/nerd/_support/_workflow.py index c741ded2f9..2a8347221a 100644 --- a/examples/learning/nerd/_support/_workflow.py +++ b/examples/learning/nerd/_support/_workflow.py @@ -271,7 +271,9 @@ def run_export(checkpoint: str | Path, output: str | Path) -> dict[str, Any]: ] ): state = torch.load(checkpoint, map_location="cpu", weights_only=True) - if state.get("version") == 1 and "learner" in state: + # SL v2 adds rank-local recovery state; portable export still uses the + # committed rank-zero model with the same recipe and ordered feature layout. + if state.get("version") in (1, 2) and "learner" in state: signature = state["signature"] model_cfg = dict(signature["model"]) weights = state["learner"]["model"] diff --git a/examples/learning/sl_ddp/README.md b/examples/learning/sl_ddp/README.md new file mode 100644 index 0000000000..eee8bec66f --- /dev/null +++ b/examples/learning/sl_ddp/README.md @@ -0,0 +1,107 @@ +# Single-node distributed supervised learning + +This example fits a `4 → 32 → 2` MLP to a fixed synthetic regression corpus. +It uses ordinary PyTorch and the common SL source, objective, learner and trainer; +no dataset download or simulator is needed. Run from the repository checkout in +an installed EmbodiChain `py311` environment. + +The first target is a linear combination of inputs; the second adds a sine term. +Every example supervises the first scalar, while the second is masked according +to an input condition. The objective returns the local masked mean squared error +and `LossResult.weight = mask.sum()`. This makes the distributed update represent +`sum(weight_rank * loss_rank) / sum(weight_rank)`, including batches with different +numbers of valid targets. The independent fixed validation corpus is evaluated +on rank 0. + +## Launch + +Activate the environment and run two CPU workers using Gloo: + +```bash +conda activate py311 +OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \ + -m embodichain train-sl --config examples/learning/sl_ddp/cpu.yaml +``` + +For two GPUs, select the available devices and use the CUDA recipe. The workflow +maps each worker to `cuda:LOCAL_RANK` within `CUDA_VISIBLE_DEVICES` and uses NCCL: + +```bash +CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=1 \ + torchrun --standalone --nnodes=1 --nproc-per-node=2 \ + -m embodichain train-sl --config examples/learning/sl_ddp/cuda.yaml +``` + +The same factory also runs on one CPU without a distributed process group. For +a separate run, copy `cpu.yaml` to `/tmp/sl-single-cpu.yaml` and set its +`trainer.output_dir` to `outputs/sl_single_cpu`: + +```bash +python -m embodichain train-sl --config /tmp/sl-single-cpu.yaml +``` + +Each fresh run needs an artifact-free output directory. Rank 0 claims its durable +identity in `run.json` and writes `metrics.jsonl`, `step_100.pt`, `step_200.pt` and +`last.pt`; the validation +MSE appears every 50 updates as `validation/mse`. These launch commands illustrate +the execution path, without claiming a multi-GPU speedup for this tiny model. + +## Batch size, sampling and recovery + +`trainer.batch_size: 32` is per worker: two workers draw 64 examples per optimizer +update. Each worker draws independently with replacement from the same immutable +corpus, so samples can overlap across workers. The learning rate stays 0.001; +it is not automatically multiplied by the worker count. Changing world size +changes the effective batch and sampling sequence. + +Every worker participates in fitting, saving and restoring checkpoints, while +only rank 0 writes artifacts or runs validation. The checkpoint retains the common +optimizer/update state and every rank's RNG, sampler and model buffers. Exact +recovery requires the same world size, recipe, source identity and training +configuration. This example's fixed corpus/model construction is the same on +all ranks; the shared trainer supplies different sampling streams. + +To demonstrate continuation from update 100, copy the same recipe and change only +`trainer.output_dir`, for example to `outputs/sl_ddp_cpu_resumed`. Keep the 200-update +budget and all other settings. Then launch the same two workers: + +```bash +OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \ + -m embodichain train-sl --config /tmp/sl-ddp-resume.yaml \ + --resume outputs/sl_ddp_cpu/step_100.pt +``` + +A fresh destination preserves the original 200-update result when restoring its +older 100-update snapshot. To resume an interrupted run from its latest snapshot, +use the original directory only when `run.json` matches the checkpoint's run +identity and its complete metric history is not ahead of that snapshot. +Checkpoints made before run identities were added require a fresh destination; +their next save records a run identity. Checkpoints are trusted local PyTorch files. + +The installed NeRD factory uses this same launch path. After preparing the data +and checking the paths in its recipe, launch it with two available GPUs: + +```bash +CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=1 \ + torchrun --standalone --nnodes=1 --nproc-per-node=2 \ + -m embodichain train-sl \ + --config examples/learning/nerd/configs/cartpole/sl.train.yaml +``` + +Its factory is `embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`. +The existing precision, causal attention, normalization and loss-mask settings +remain recipe-owned. Each rank loads the dataset and keeps its own resident +cache, so budget memory per worker. Validation runs on rank 0 with its unwrapped +model. If a validation callback calls `trainer.save()`, that save uses the +committed state captured from all ranks before validation; calling `save()` +elsewhere requires every rank to participate. + +The CLI handles process-group setup and cleanup automatically. Programmatic +launchers can use `embodichain.learning.sl.distributed_run(device)` around model +and trainer construction; all ranks must enter that scope and participate in the +training calls. It yields the resolved `torch.device`, and destroys only a process +group that it created. A pre-existing group remains caller-owned. + +See the pinned-version PyTorch references for +[`torchrun` in 2.7](https://github.com/pytorch/pytorch/blob/v2.7.0/torch/distributed/run.py) +and [DistributedDataParallel 2.7](https://docs.pytorch.org/docs/2.7/generated/torch.nn.parallel.DistributedDataParallel.html). diff --git a/examples/learning/sl_ddp/cpu.yaml b/examples/learning/sl_ddp/cpu.yaml new file mode 100644 index 0000000000..ca3c50b94d --- /dev/null +++ b/examples/learning/sl_ddp/cpu.yaml @@ -0,0 +1,19 @@ +factory: examples.learning.sl_ddp.recipe:build_trainer +model: + hidden_dim: 32 + seed: 29 +data: + train_count: 1024 + validation_count: 128 + seed: 17 +optimizer: + lr: 0.001 + weight_decay: 0.0 +trainer: + total_updates: 200 + batch_size: 32 + device: cpu + seed: 101 + output_dir: outputs/sl_ddp_cpu + checkpoint_interval: 100 + validation_interval: 50 diff --git a/examples/learning/sl_ddp/cuda.yaml b/examples/learning/sl_ddp/cuda.yaml new file mode 100644 index 0000000000..6e19077009 --- /dev/null +++ b/examples/learning/sl_ddp/cuda.yaml @@ -0,0 +1,19 @@ +factory: examples.learning.sl_ddp.recipe:build_trainer +model: + hidden_dim: 32 + seed: 29 +data: + train_count: 1024 + validation_count: 128 + seed: 17 +optimizer: + lr: 0.001 + weight_decay: 0.0 +trainer: + total_updates: 200 + batch_size: 32 + device: cuda + seed: 101 + output_dir: outputs/sl_ddp_cuda + checkpoint_interval: 100 + validation_interval: 50 diff --git a/examples/learning/sl_ddp/recipe.py b/examples/learning/sl_ddp/recipe.py new file mode 100644 index 0000000000..3bc44e13ed --- /dev/null +++ b/examples/learning/sl_ddp/recipe.py @@ -0,0 +1,122 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""A fixed masked-regression recipe for single-device or distributed SL.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +from torch import nn + +from embodichain.learning.sl import ( + LossResult, + SupervisedBatch, + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, + TensorBatchSource, +) + +__all__ = ["build_trainer"] + + +class _MaskedMSE: + def compute_loss(self, model: nn.Module, batch: SupervisedBatch) -> LossResult: + """Return a local mean weighted by the number of valid scalar targets.""" + predictions = model(batch.inputs) + mask = batch.masks["targets"].to(dtype=predictions.dtype) + weight = mask.sum() + loss = ((predictions - batch.targets).square() * mask).sum() / weight + return LossResult(loss, {"mse": loss.detach()}, weight=float(weight.detach())) + + +def _corpus(size: int, *, seed: int) -> SupervisedBatch: + generator = torch.Generator(device="cpu").manual_seed(seed) + inputs = torch.randn(size, 4, generator=generator) + targets = torch.stack( + ( + inputs[:, 0] + 0.5 * inputs[:, 1] - 0.25 * inputs[:, 2], + inputs[:, 0].sin() + inputs[:, 3], + ), + dim=-1, + ) + # Every row retains one supervised scalar. The second varies by sample, + # so equally sized rank-local batches may have different valid weights. + mask = torch.stack((torch.ones(size, dtype=torch.bool), inputs[:, 0] > 0), dim=-1) + return SupervisedBatch( + inputs, + targets, + masks={"targets": mask}, + metadata={"features": 4, "targets": 2, "units": "dimensionless"}, + ) + + +def build_trainer(config: Mapping[str, Any]) -> SupervisedTrainer: + """Build a tiny recipe using the same factory for one or multiple workers. + + Args: + config: Mapping containing data counts/seeds, hidden width/model seed, + AdamW options and the shared trainer configuration. The SL workflow + resolves the trainer device before calling this factory. + + Returns: + A trainer with a fixed CPU corpus and an independent validation corpus. + """ + data = { + "train_count": 1024, + "validation_count": 128, + "seed": 17, + **config.get("data", {}), + } + model_config = {"hidden_dim": 32, "seed": 29, **config.get("model", {})} + optimizer_config = {"lr": 1e-3, "weight_decay": 0.0, **config.get("optimizer", {})} + cfg = SupervisedTrainerCfg(**config.get("trainer", {})) + training = _corpus(data["train_count"], seed=data["seed"]) + validation = _corpus(data["validation_count"], seed=data["seed"] + 1) + source = TensorBatchSource(training, source_id="masked-regression-v1") + torch.manual_seed(model_config["seed"]) + model = nn.Sequential( + nn.Linear(4, model_config["hidden_dim"]), + nn.Tanh(), + nn.Linear(model_config["hidden_dim"], 2), + ).to(cfg.device) + objective = _MaskedMSE() + learner = SupervisedLearner( + model, + objective, + torch.optim.AdamW(model.parameters(), **optimizer_config), + ) + + def evaluate(current: SupervisedLearner) -> Mapping[str, float]: + result = objective.compute_loss(current.model, validation.to(cfg.device)) + return {"mse": float(result.loss)} + + return SupervisedTrainer( + learner, + source, + cfg, + signature={ + "recipe": "masked-regression-v1", + "model": model_config, + "data": data, + "optimizer": optimizer_config, + "objective": "masked-mse-per-valid-scalar-v1", + }, + evaluate=evaluate, + ) diff --git a/tests/benchmark/motion_generation/test_embodiment_runtime.py b/tests/benchmark/motion_generation/test_embodiment_runtime.py index 5549b17e8b..120c81b502 100644 --- a/tests/benchmark/motion_generation/test_embodiment_runtime.py +++ b/tests/benchmark/motion_generation/test_embodiment_runtime.py @@ -253,7 +253,7 @@ def test_gripper_and_domain_overrides_cannot_change_component_commands( def test_official_ur5_component_uses_robot_catalog(monkeypatch, tmp_path: Path) -> None: from embodichain.lab.sim import cfg as cfg_module - from embodichain_robots import ur_robot + from embodichain_robots.universal_robots import ur_robot def local_asset(path: str) -> str: target = tmp_path / path diff --git a/tests/cli/test_list_robot.py b/tests/cli/test_list_robot.py index 648e73b6aa..ff65ecd0a0 100644 --- a/tests/cli/test_list_robot.py +++ b/tests/cli/test_list_robot.py @@ -248,8 +248,10 @@ def test_list_robot_wraps_long_catalog_metadata( def test_list_robot_prints_control_metadata_without_loading_robot( catalogs: list[Any], capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, ) -> None: """A provider's control summary is printed directly from its catalog.""" + monkeypatch.setenv("COLUMNS", "120") spec = { "class_path": "absent_robots:RobotCfg", "dof": 7, diff --git a/tests/learning/sl/_distributed_recipe.py b/tests/learning/sl/_distributed_recipe.py new file mode 100644 index 0000000000..ee0ad474df --- /dev/null +++ b/tests/learning/sl/_distributed_recipe.py @@ -0,0 +1,85 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Private CPU recipe for testing real torchrun lifecycle and resume.""" + +from __future__ import annotations + +import atexit +import json +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist + +from embodichain.learning.sl import ( + LossResult, + SupervisedBatch, + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, + TensorBatchSource, +) + + +class _Regression: + def compute_loss( + self, model: torch.nn.Module, batch: SupervisedBatch + ) -> LossResult: + error = (model(batch.inputs) - batch.targets).square() + return LossResult(error.mean(), weight=error.numel()) + + +class _BoundedTrainer(SupervisedTrainer): + boundary: int | None = None + + def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: + return super().fit( + stop_after=self.boundary if stop_after is None else stop_after + ) + + +def build(config: dict[str, Any]) -> SupervisedTrainer: + """Construct a tiny learner only after the production workflow initialized DDP.""" + initialized = dist.is_initialized() + assert initialized, "torchrun workflow must initialize the process group" + rank = dist.get_rank() + target = Path(config["test_audit_dir"]) / f"{config['test_phase']}-rank{rank}.json" + + def audit_exit() -> None: + target.write_text( + json.dumps( + { + "initialized_in_factory": initialized, + "initialized_at_exit": dist.is_initialized(), + } + ) + ) + + atexit.register(audit_exit) + torch.manual_seed(31) + model = torch.nn.Linear(2, 1) + values = torch.arange(16, dtype=torch.float32).view(8, 2) / 16 + source = TensorBatchSource(SupervisedBatch(values, values.sum(-1, keepdim=True))) + learner = SupervisedLearner( + model, _Regression(), torch.optim.Adam(model.parameters(), lr=0.02) + ) + trainer = _BoundedTrainer( + learner, source, SupervisedTrainerCfg(**config["trainer"]) + ) + trainer.boundary = config["test_stop_after"] + return trainer diff --git a/tests/learning/sl/test_distributed.py b/tests/learning/sl/test_distributed.py new file mode 100644 index 0000000000..bf6dda86c6 --- /dev/null +++ b/tests/learning/sl/test_distributed.py @@ -0,0 +1,555 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Bounded CPU Gloo coverage for shared-loss DDP and rank-local recovery.""" + +from __future__ import annotations + +import copy +from datetime import timedelta +import json +import os +from pathlib import Path +import random +import signal +import subprocess +import sys +import time +from typing import Any + +import numpy as np +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from embodichain.learning.sl import ( + LossResult, + SupervisedBatch, + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, + TensorBatchSource, +) + +pytestmark = [ + pytest.mark.slow, + pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="Requires CPU Gloo", + ), +] + + +class _Transition(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(2, 1, dtype=torch.float64) + self.step_size = 0.25 + + def forward(self, value: torch.Tensor) -> torch.Tensor: + return self.linear(value).tanh() + + +class _MaskedObjective: + def __init__(self, multistep: bool = False) -> None: + self.multistep = multistep + + def compute_loss( + self, model: torch.nn.Module, batch: SupervisedBatch + ) -> LossResult: + if self.multistep: + prediction = batch.inputs["state"] + for action in batch.inputs["actions"].unbind(1): + prediction = prediction + model.step_size * model( + torch.cat((prediction, action), -1) + ) + else: + prediction = model(batch.inputs) + errors = (prediction - batch.targets).square()[batch.masks["valid"]] + return LossResult(errors.mean(), {"mse": errors.mean()}, weight=errors.numel()) + + +def _weighted_worker(rank: int) -> None: + for multistep in (False, True): + torch.manual_seed(73) + model = _Transition() + reference = copy.deepcopy(model) + objective = _MaskedObjective(multistep) + values = torch.arange(12, dtype=torch.float64).view(6, 2) / 10 + targets = torch.linspace(-0.4, 0.6, 6, dtype=torch.float64).view(6, 1) + mask = torch.tensor([[True], [False], [False], [True], [True], [True]]) + if multistep: + inputs = { + "state": values[:, :1], + "actions": torch.arange(18, dtype=torch.float64).view(6, 3, 1) / 20, + } + local_inputs = { + key: value[rank * 3 : rank * 3 + 3] for key, value in inputs.items() + } + else: + inputs = values + local_inputs = values[rank * 3 : rank * 3 + 3] + global_batch = SupervisedBatch(inputs, targets, {"valid": mask}) + # The reference differentiates the independent concatenated-batch loss. + expected = objective.compute_loss(reference, global_batch) + expected_optimizer = torch.optim.SGD(reference.parameters(), lr=0.07) + expected.loss.backward() + expected_optimizer.step() + learner = SupervisedLearner( + model, objective, torch.optim.SGD(model.parameters(), lr=0.07) + ) + assert learner.model is model + report = learner.update( + SupervisedBatch( + local_inputs, + targets[rank * 3 : rank * 3 + 3], + {"valid": mask[rank * 3 : rank * 3 + 3]}, + ) + ) + assert report["weight"] == 4 + assert report["loss"] == pytest.approx(expected.loss.item(), abs=1e-14) + assert report["mse"] == pytest.approx(expected.loss.item(), abs=1e-14) + for actual, target in zip(model.parameters(), reference.parameters()): + torch.testing.assert_close(actual, target, rtol=1e-12, atol=1e-14) + copies = [torch.empty_like(actual) for _ in range(2)] + dist.all_gather(copies, actual.detach()) + torch.testing.assert_close(copies[0], copies[1], rtol=0, atol=0) + + +class _DropoutModel(torch.nn.Module): + def __init__(self, rank: int) -> None: + super().__init__() + self.rank = rank + self.network = torch.nn.Sequential( + torch.nn.Linear(2, 8), torch.nn.Dropout(0.3), torch.nn.Linear(8, 1) + ) + self.register_buffer("seen", torch.zeros((), dtype=torch.int64)) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + self.seen.add_(self.rank + 1) + return self.network(value) + + +class _Regression: + def compute_loss( + self, model: torch.nn.Module, batch: SupervisedBatch + ) -> LossResult: + error = (model(batch.inputs) - batch.targets).square() + return LossResult(error.mean(), weight=error.numel()) + + +class _StochasticSource: + def __init__(self) -> None: + values = torch.arange(40, dtype=torch.float32).view(20, 2) / 40 + self.source = TensorBatchSource( + SupervisedBatch(values, values.sum(-1, keepdim=True)) + ) + self.source_id = self.source.source_id + ":python-numpy-noise-v1" + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + batch = self.source.sample(batch_size, generator=generator) + offset = 0.01 * (random.random() + float(np.random.random())) + return SupervisedBatch(batch.inputs + offset, batch.targets) + + +def _build_recovery(path: Path, rank: int) -> SupervisedTrainer: + random.seed(101 + rank) + np.random.seed(201 + rank) + torch.manual_seed(301 + rank) + model = _DropoutModel(rank) + learner = SupervisedLearner( + model, _Regression(), torch.optim.Adam(model.parameters(), lr=0.01) + ) + cfg = SupervisedTrainerCfg( + total_updates=6, + batch_size=4, + seed=17, + output_dir=str(path), + checkpoint_interval=2, + validation_interval=2, + ) + trainer = SupervisedTrainer( + learner, + _StochasticSource(), + cfg, + signature={"regression": "rank-local recovery"}, + ) + + def evaluate(current: SupervisedLearner) -> dict[str, float]: + assert rank == 0, "Validation must run only on rank zero" + assert not current.model.training + current.model(torch.ones(1, 2)) + random.random() + np.random.random(3) + torch.rand(5) + torch.randint(10, (3,), generator=trainer.sampler) + # BC-style checkpoint publication inside a rank-zero callback must not hang. + trainer.save(path / "selected.pt") + with (path.parent / f"{path.name}-validation-rank{rank}.jsonl").open( + "a" + ) as handle: + handle.write(json.dumps({"step": current.num_updates}) + "\n") + return {"probe": float(current.num_updates)} + + trainer.evaluate = evaluate + return trainer + + +def _snapshot(trainer: SupervisedTrainer) -> dict[str, Any]: + return copy.deepcopy( + { + "learner": trainer.learner.state_dict(), + "sampler": trainer.sampler.get_state(), + "rng": { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + }, + } + ) + + +def _assert_tree_equal(actual: Any, expected: Any) -> None: + if isinstance(expected, torch.Tensor): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + elif isinstance(expected, np.ndarray): + np.testing.assert_array_equal(actual, expected) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys() + for key in expected: + _assert_tree_equal(actual[key], expected[key]) + elif isinstance(expected, (tuple, list)): + assert type(actual) is type(expected) + assert len(actual) == len(expected) + for value, target in zip(actual, expected): + _assert_tree_equal(value, target) + else: + assert actual == expected + + +def _recovery_worker(rank: int, directory: Path, stage: str) -> None: + trainer = _build_recovery(directory / stage, rank) + expected_sampler = torch.Generator().manual_seed(trainer.cfg.seed + rank) + torch.testing.assert_close( + trainer.sampler.get_state(), expected_sampler.get_state() + ) + if stage == "resumed": + trainer.resume(directory / "partial/selected.pt") + # Check rank-local buffers/RNG before a later forward could hide an error. + partial = torch.load(directory / f"partial-rank{rank}.pt", weights_only=False) + _assert_tree_equal(_snapshot(trainer), partial) + trainer.fit(stop_after=2 if stage == "partial" else None) + torch.save(_snapshot(trainer), directory / f"{stage}-rank{rank}.pt") + + +class _FailingSource: + source_id = "collaborative-failure-source-v1" + + def __init__(self, rank: int, fails: bool) -> None: + self.rank, self.fails = rank, fails + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + if self.fails and self.rank == 1: + raise ValueError("injected-source-failure") + return SupervisedBatch(torch.ones(batch_size, 2), torch.zeros(batch_size, 1)) + + +class _FailingObjective(_Regression): + def __init__(self, rank: int, failure: str) -> None: + self.rank, self.failure = rank, failure + + def compute_loss( + self, model: torch.nn.Module, batch: SupervisedBatch + ) -> LossResult: + if self.failure == "objective" and self.rank == 1: + raise ValueError("injected-objective-failure") + result = super().compute_loss(model, batch) + if self.rank == 1 and self.failure == "loss": + return LossResult(result.loss * float("nan")) + if self.rank == 1 and self.failure == "weight": + return LossResult(result.loss, weight=0) + return result + + +def _failure_worker(rank: int, directory: Path) -> None: + for phase in ("source", "objective", "loss", "weight", "validation"): + torch.manual_seed(11) + model = torch.nn.Linear(2, 1) + before = copy.deepcopy(model.state_dict()) + learner = SupervisedLearner( + model, + _FailingObjective(rank, phase), + torch.optim.Adam(model.parameters(), lr=0.01), + ) + + def evaluate(current: SupervisedLearner) -> dict[str, float]: + assert rank == 0 + raise ValueError("injected-validation-failure") + + trainer = SupervisedTrainer( + learner, + _FailingSource(rank, phase == "source"), + SupervisedTrainerCfg( + total_updates=2, + batch_size=2, + output_dir=str(directory / phase), + checkpoint_interval=0, + validation_interval=1 if phase == "validation" else 0, + ), + evaluate=evaluate if phase == "validation" else None, + ) + sampler_before = trainer.sampler.get_state().clone() + try: + trainer.fit() + except Exception as error: + message = { + "loss": "Non-finite supervised loss", + "weight": "Objective weight", + }.get(phase, f"injected-{phase}-failure") + assert message in str(error) + else: + raise AssertionError(f"Rank {rank} did not observe the {phase} failure") + assert learner.num_updates == (1 if phase == "validation" else 0) + if phase != "validation": + _assert_tree_equal(model.state_dict(), before) + torch.testing.assert_close(trainer.sampler.get_state(), sampler_before) + assert not learner.optimizer.state + dist.barrier() + + +def _artifact_worker(rank: int, directory: Path) -> None: + def build(output: Path) -> SupervisedTrainer: + trainer = _build_recovery(output, rank) + trainer.evaluate = None + trainer.cfg.validation_interval = 0 + return trainer + + source = build(directory / "source") + collision = build(source.output) + source.fit(stop_after=2) + original = { + name: (source.output / name).read_bytes() + for name in ("metrics.jsonl", "last.pt", "run.json") + } + with pytest.raises(RuntimeError, match="another run"): + collision.fit(stop_after=1) + for name, content in original.items(): + assert (source.output / name).read_bytes() == content + + foreign = build(directory / "foreign") + foreign.fit(stop_after=1) + candidate = build(foreign.output) + before = _snapshot(candidate) + with pytest.raises(RuntimeError, match="another run"): + candidate.resume(source.output / "last.pt") + _assert_tree_equal(_snapshot(candidate), before) + + for history in ("newer", "incomplete"): + if rank == 0: + (source.output / "metrics.jsonl").write_bytes( + original["metrics.jsonl"] + + ( + b'{"num_updates": 3}\n' + if history == "newer" + else b'{"num_updates":' + ) + ) + dist.barrier() + log = (source.output / "metrics.jsonl").read_bytes() + candidate = build(source.output) + before = _snapshot(candidate) + with pytest.raises(RuntimeError, match="history"): + candidate.resume(source.output / "last.pt") + _assert_tree_equal(_snapshot(candidate), before) + assert (source.output / "metrics.jsonl").read_bytes() == log + dist.barrier() + + # Previously published DDP checkpoints have no run ID. They still recover + # exactly, but must acquire their new identity in a fresh destination. + if rank == 0: + state = torch.load(source.output / "last.pt", weights_only=False) + state.pop("run_id") + torch.save(state, directory / "legacy.pt") + dist.barrier() + recovered = build(directory / "legacy-recovered") + recovered.resume(directory / "legacy.pt") + recovered.fit(stop_after=3) + saved = torch.load(recovered.output / "last.pt", weights_only=False) + assert ( + saved["run_id"] + == json.loads((recovered.output / "run.json").read_text())["run_id"] + ) + assert saved["learner"]["num_updates"] == 3 + + +def _distributed_worker( + rank: int, group_file: str, directory: str, scenario: str +) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{group_file}", + rank=rank, + world_size=2, + timeout=timedelta(seconds=45), + ) + try: + if scenario == "weighted": + _weighted_worker(rank) + elif scenario == "errors": + _failure_worker(rank, Path(directory)) + elif scenario == "artifacts": + _artifact_worker(rank, Path(directory)) + else: + _recovery_worker(rank, Path(directory), scenario) + finally: + dist.destroy_process_group() + + +def _spawn(scenario: str, directory: Path) -> None: + context = mp.spawn( + _distributed_worker, + args=(str(directory / f"group-{scenario}"), str(directory), scenario), + nprocs=2, + join=False, + ) + deadline = time.monotonic() + 100 + try: + while not context.join(timeout=1): + if time.monotonic() > deadline: + pytest.fail(f"Gloo {scenario} exceeded 100 seconds") + finally: + for process in context.processes: + if process.is_alive(): + process.terminate() + for process in context.processes: + process.join(timeout=5) + + +def test_unequal_weight_single_and_multistep_match_global_batch(tmp_path: Path) -> None: + _spawn("weighted", tmp_path) + + +def test_restart_preserves_rank_rng_sampling_buffers_and_single_writer( + tmp_path: Path, +) -> None: + for stage in ("full", "partial", "resumed"): + _spawn(stage, tmp_path) + for rank in range(2): + full = torch.load(tmp_path / f"full-rank{rank}.pt", weights_only=False) + resumed = torch.load(tmp_path / f"resumed-rank{rank}.pt", weights_only=False) + _assert_tree_equal(resumed, full) + full_lines = (tmp_path / "full/metrics.jsonl").read_text() + combined = (tmp_path / "partial/metrics.jsonl").read_text() + ( + tmp_path / "resumed/metrics.jsonl" + ).read_text() + assert combined == full_lines + records = [json.loads(line) for line in full_lines.splitlines()] + assert [record["num_updates"] for record in records] == list(range(1, 7)) + assert all(record["weight"] == 8 for record in records) + assert (tmp_path / "full/last.pt").is_file() + assert len((tmp_path / "full-validation-rank0.jsonl").read_text().splitlines()) == 3 + assert not list(tmp_path.glob("*-validation-rank1.jsonl")) + # A real two-rank checkpoint cannot silently become a one-rank run. + candidate = _build_recovery(tmp_path / "different-world-size", rank=0) + before = _snapshot(candidate) + with pytest.raises((ValueError, RuntimeError), match="[Ww]orld.?size"): + candidate.resume(tmp_path / "partial/selected.pt") + _assert_tree_equal(_snapshot(candidate), before) + + +def test_rank_failures_reach_all_workers_without_partial_update(tmp_path: Path) -> None: + _spawn("errors", tmp_path) + + +def test_output_identity_and_history_admission_reaches_all_ranks( + tmp_path: Path, +) -> None: + _spawn("artifacts", tmp_path) + + +def test_torchrun_cli_initializes_cleans_up_and_resumes(tmp_path: Path) -> None: + output = tmp_path / "run" + config = { + "factory": "_distributed_recipe:build", + "trainer": { + "total_updates": 4, + "batch_size": 3, + "seed": 23, + "device": "cpu", + "output_dir": str(output), + "checkpoint_interval": 0, + }, + "test_audit_dir": str(tmp_path), + "test_phase": "partial", + "test_stop_after": 2, + } + config_path = tmp_path / "recipe.json" + child_env = os.environ.copy() + child_env["CUDA_VISIBLE_DEVICES"] = "" + child_env["OMP_NUM_THREADS"] = "1" + child_env["PYTHONPATH"] = os.pathsep.join( + filter( + None, (str(Path(__file__).resolve().parent), child_env.get("PYTHONPATH")) + ) + ) + for phase in ("partial", "resumed"): + config["test_phase"] = phase + config["test_stop_after"] = 2 if phase == "partial" else None + config_path.write_text(json.dumps(config)) + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc_per_node=2", + "-m", + "embodichain", + "train-sl", + "--config", + str(config_path), + ] + if phase == "resumed": + command.extend(["--resume", str(output / "last.pt")]) + with subprocess.Popen( + command, + cwd=Path(__file__).resolve().parents[3], + env=child_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) as process: + try: + stdout, stderr = process.communicate(timeout=100) + except subprocess.TimeoutExpired: + # The new process group contains only this test's launcher/workers. + os.killpg(process.pid, signal.SIGKILL) + stdout, stderr = process.communicate() + pytest.fail(f"torchrun {phase} exceeded 100 seconds:\n{stdout}{stderr}") + assert process.returncode == 0, stdout + stderr + for rank in range(2): + audit = json.loads((tmp_path / f"{phase}-rank{rank}.json").read_text()) + assert audit == { + "initialized_in_factory": True, + "initialized_at_exit": False, + } + metrics = [ + json.loads(line) for line in (output / "metrics.jsonl").read_text().splitlines() + ] + assert [row["num_updates"] for row in metrics] == [1, 2, 3, 4] + assert all(row["weight"] == 6 for row in metrics) diff --git a/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py index c1681efbad..97bab43e33 100644 --- a/tests/learning/sl/test_trainer.py +++ b/tests/learning/sl/test_trainer.py @@ -214,6 +214,7 @@ def record_forward(model, inputs): count = 1 if callback_fails else 3 baseline.fit(stop_after=count) expected = copy.deepcopy(baseline.learner.model.state_dict()) + expected_sampler = baseline.sampler.get_state().clone() expected_random = (random.random(), np.random.rand(), torch.rand(4)) reset_rng() candidate = build(tmp_path / "candidate") @@ -225,6 +226,7 @@ def evaluate(learner): random.random() np.random.rand(5) torch.rand(9) + torch.randint(10, (3,), generator=candidate.sampler) candidate.save(candidate.output / "selected.pt") if callback_fails: raise RuntimeError("validation failed") @@ -237,6 +239,7 @@ def evaluate(learner): else: candidate.fit(stop_after=count) assert candidate.learner.model.training + torch.testing.assert_close(candidate.sampler.get_state(), expected_sampler) for key, value in expected.items(): torch.testing.assert_close( candidate.learner.model.state_dict()[key], value, rtol=0, atol=0 @@ -248,6 +251,7 @@ def evaluate(learner): count_forwards(resumed) resumed.cfg.validation_interval = 1 resumed.resume(candidate.output / "selected.pt") + torch.testing.assert_close(resumed.sampler.get_state(), expected_sampler) for key, value in expected.items(): torch.testing.assert_close( resumed.learner.model.state_dict()[key], value, rtol=0, atol=0