Add single-node DDP training to shared SL

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/

agent_context/topics/supervised-learning/supervised-learning.md

@@ -19,15 +19,30 @@ 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
   showcase's default `train` command uses this same common recipe.
 
 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.

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 <https://github.com/pytorch/pytorch/blob/v2.7.0/torch/distributed/run.py>`_
+and `DistributedDataParallel
+<https://docs.pytorch.org/docs/2.7/generated/torch.nn.parallel.DistributedDataParallel.html>`_.
+
 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
 -------------------------------
 

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",
 ]

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