Skip to content

Add single-node DDP training to shared SL - #864

Merged
yuecideng merged 38 commits into
mainfrom
codex/sl-ddp
Oct 26, 2026
Merged

yuecideng merged 38 commits into
mainfrom
codex/sl-ddp

Conversation

@yuecideng

@yuecideng yuecideng commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Description

Run shared SL recipes under single-node torchrun, with one CUDA/NCCL or CPU/Gloo worker per device. This branch has normally merged main@c5b1fdfc2, which contains the merged #855 foundation. it retains that foundation's durable run identity, append-only metrics, validation snapshots and checkpoint compatibility checks.

DDP wraps the objective computation while learner.model remains the original model. Local mean losses are weighted by their effective valid-supervision counts before gradient averaging. Only rank 0 validates or writes artifacts; checkpoints preserve every rank's RNG, sampler and model buffers. NeRD showcase export now accepts both SL v1 and distributed SL v2 checkpoints, preserving safe loading, dtype, normalization and ordered inputs.

Launch and contracts

conda activate py311
CUDA_VISIBLE_DEVICES=1,6 OMP_NUM_THREADS=1 torchrun \
  --standalone --nnodes=1 --nproc-per-node=2 \
  -m embodichain train-sl --config examples/learning/sl_ddp/cuda.yaml

Use cpu.yaml without visible CUDA devices for Gloo. Model recipes use the same entry point, including embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer. Append --resume PATH with the same world size/configuration and a compatible destination. See the launch guide.

Batch size is per worker, sampling uses independent replacement streams and learning rate is not automatically scaled. Exact recovery requires the same world size. All ranks participate in ordinary fitting/saving/recovery; rank-zero validation callbacks save the previously gathered committed boundary without another collective. Rank-local admission, sampling, objective, finite-value, validation and I/O errors are propagated to all workers. Fatal device/communication failures require restarting the job.

The checkpoint control-state transport temporarily uses O(world-size × maximum serialized rank-state size) storage per rank. It carries RNG, sampler and registered buffers, rather than parameters, optimizer state, images or corpora. This is a bounded compatibility tradeoff, not a universal speedup claim.

Main synchronization

Main synchronization at 8dae1b392: 88 relevant files remain byte/AST-identical to the prior qualified DDP source; main's vendor catalogs, CI, package and release checks are preserved. The new main integration passes 28 SL/Gloo checks, 94 metadata/context checks and 135 catalog/CLI checks; these groups are reported separately. A fresh two-GPU NCCL four-update run and recovery from update two match exactly. The prior 12-update interruption and NeRD v1/v2 export evidence remains source-bound and applicable to unchanged execution files.

Validation

Current branch source: 8dae1b3920677ff0be7c7fb262cd64b9a2896b13; prior qualified source: ba89775e387b7f89a9bde5bd36070ac8146ca666; normal base merge: a1a77ceaa.

  • 70 focused CPU/Gloo/context checks pass. Independent concatenated-batch oracles cover unequal valid counts and multi-step gradients. Recovery covers dropout, optimizer, rank buffers, sampler/RNG, callback checkpoints and root-only output ownership, including foreign runs and invalid/ahead histories.
  • Actual two-worker NCCL on GPUs 1/6: a continuous 12-update run matches a run terminated after update 6 and recovered to 12. Model, optimizer, all Python/NumPy/CPU/CUDA RNG, sampler, every rank's buffers and remaining log are identical. Device selection and normal process-group cleanup pass.
  • Actual NeRD v2 two-GPU and v1 single-device four-update checkpoints both export through the showcase and run inference. Portable weights, dtype and ordered feature layout match exactly; prediction difference is zero.
  • Combined with Add robot world models through the model extension and SL #859 and Add RigidFormer dynamics models and SL showcase #858, five production cases (state/visual BC interfaces, NeRD, three-step RWM and uneven-object RigidFormer) pass two real Gloo updates against independent full-batch updates. Replicas match exactly; maximum parameter deviation is 2.98e-8. The BC checks use its frozen draft source and do not change or qualify its single-process formal acceptance recipe.
  • The final combined wheel, installed after deleting its build source, passes model imports, CLI launches and two-worker Gloo interruption/recovery from an unrelated directory with simulator imports blocked.
  • Black 26.10.1 black ., context validation and API coverage 3170/3170 pass. Independent HPC/reliability/security review found no unresolved actionable defect.

Combined affected checks and package evidence are recorded separately from branch-only checks; overlapping counts are not added. No new long training/convergence matrix or speedup measurement is claimed.

CI integration repairs

Current review head: f372f2079ba94b79d928fc05f0fd114d126d29dc. The unchanged main CI exposed two tests left behind by the vendor migration: narrow-terminal label wrapping and a retired UR mock import. The exact two-file maintenance commit b1d8ab415 is reviewed separately in #878 and included here; semantic assertions, CI selection/gates and production behavior are unchanged. Its original two failures reproduce and the corrected related modules pass 42 tests.

Dependencies and review

Type of change

  • New feature and recovery/export compatibility fixes.

Checklist

  • Run black . with Black 26.10.1.
  • Review affected context and document DDP lifecycle/launch contracts.
  • Update public API documentation and pass the coverage gate.
  • Reuse independent behavioral oracles and add necessary collective-admission coverage.
  • Validate CPU/Gloo, actual CUDA/NCCL, v1/v2 export and installed-package behavior.
  • Preserve review-required merge control.

Base automatically changed from codex/sl-foundation to codex/nerd-model-data-training October 25, 2026 07:17
@greptile-apps

greptile-apps Bot commented Oct 25, 2026 •

Copy link
Copy Markdown

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The changes since the previous review appear safe to merge.

Summary

Adds single-node DDP training to shared SL.

  • Shared SL recipes train across single-node workers with weighted updates.
  • Distributed runs validate on rank 0 and save each worker’s recovery state.
  • NeRD export reads both single-device and distributed SL checkpoints.
  • Robot-list and runtime tests use stable imports and terminal width.

No actionable new issues or repository-rule violations were found. No previous review threads were supplied. This review used code inspection; no tests were run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[torchrun workers] --> B[Initialize process group]
    B --> C[Build recipe on each worker]
    C --> D[Sample local batches]
    D --> E[Combine weighted gradients]
    E --> F[Commit optimizer update]
    F --> G[Gather worker recovery state]
    G --> H[Rank 0 validates and writes artifacts]
    H --> D
Loading

Reviews (5) · Last reviewed commit: "Include vendor catalog test compatibilit..." · Reviewed by Greptile

@yuecideng yuecideng added feature New user-facing capabilities. performance Performance, memory, throughput and synchronization-cost improvements. sl Supervised learning updates, training, recovery and behavior cloning workflows. labels Oct 25, 2026
Base automatically changed from codex/nerd-model-data-training to main October 25, 2026 18:07
@yuecideng yuecideng mentioned this pull request Oct 25, 2026
7 tasks done
@yuecideng
yuecideng merged commit f3a399e into main Oct 26, 2026
9 checks passed
@yuecideng
yuecideng deleted the codex/sl-ddp branch October 26, 2026 01:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature New user-facing capabilities. performance Performance, memory, throughput and synchronization-cost improvements. sl Supervised learning updates, training, recovery and behavior cloning workflows.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant