Codeseys's picture
Wave 13: serverless DiLoCo + replaysim normalization + 3 distillation losses + PRIME-RL + Monarch
b266c31
|
Raw
History Blame Contribute Delete
6.2 kB

Monarch actor mesh — design for hosting the framework's training topology

Status: Design + skeleton. Real Monarch integration is post-replication work (ADR-006 explicitly defers it to v0.2+).

ADR: 006

What Monarch is

Monarch (https://github.com/meta-pytorch/monarch, BSD-3) is Meta's actor- mesh runtime — a thin coordination layer over Python processes that lets you describe a training topology as a graph of typed actors, then run that topology on top of any cluster manager (k8s, Slurm, raw ssh).

Per ADR-006, Monarch is the only Meta PyTorch agentic-stack component that's actively shipping (v0.4.1 stable, v0.5 dev daily) and not paused. TorchForge, the original "agent" piece, is paused per its own repo banner.

Why Monarch fits the framework's design

The framework already has an N-actor topology even without Monarch:

  • Trainer (channel 1: GRPO; channel 2: SDPO; channel 3: trace-replay DPO)
  • Generator (rollout / vllm)
  • Rewarder (RLVR test runner / verifiers protocol)
  • N teachers (channel 3: external OpenRouter calls)
  • DiLoCo replicas (N copies of trainer, syncing via object store)

PRIME-RL gives us the trainer/generator/rewarder split for free. Monarch takes that further: each of those becomes a Monarch actor, and the framework gains:

  1. Heterogeneous executor support — actors run wherever Monarch's backend places them (Modal, k8s, on-prem cluster). Composes naturally with our ServerlessExecutor Protocol.
  2. Failure recovery — Monarch handles actor crashes + restarts; the framework's DiLoCo state is durable in object storage, so a restarted trainer replica can resume from the last outer round.
  3. Hot-swap of actor implementations — switch teacher backends from "OpenRouter" to "local vllm" by changing one Monarch actor binding.

Actor topology (proposed)

┌───────────────────────────────────────────────────────────────┐
│                  ComposerReplicationMesh                       │
│                                                                │
│   ┌──────────────┐   ┌──────────────┐   ┌──────────────────┐  │
│   │  Trainer × N │←─│  Generator    │←─│  Rewarder         │  │
│   │  (DiLoCo     │   │  (vllm)       │   │  (verifiers)      │  │
│   │   replicas)  │   └──────────────┘   └──────────────────┘  │
│   └──────┬───────┘                                              │
│          │                                                      │
│          │   Channel 2: same-model hint-conditioned forward    │
│          │   Channel 3: cross-model OpenRouter teachers        │
│          ▼                                                      │
│   ┌──────────────┐                                              │
│   │ TeacherPool  │ ── OpenRouter (Claude, GPT, DeepSeek, ...) │
│   │ (channel 3)  │                                              │
│   └──────────────┘                                              │
│                                                                │
│   ┌──────────────────────────────────────────────────────────┐ │
│   │  ObjectStore (s3://, hf://, file://)                     │ │
│   │  · DiLoCo pseudo-gradients (round_N/rank_R.pt)           │ │
│   │  · Replay datasets (NormalizedDPOPair JSONL)             │ │
│   └──────────────────────────────────────────────────────────┘ │
└────────────────────────────────────────────────────────────────┘

Mapping to Monarch primitives

from monarch import Actor, mesh, endpoint

class TrainerActor(Actor):
    """Hosts the GRPO trainer + composer 3-channel loss."""
    @endpoint
    async def train_outer_step(self, batch_id: int): ...

class GeneratorActor(Actor):
    """vllm rollout server — generates trajectories on demand."""
    @endpoint
    async def rollout(self, prompts: list[str]) -> list[str]: ...

class RewarderActor(Actor):
    """Runs verifiers protocol — RLVR-style test execution."""
    @endpoint
    async def score(self, completions: list[str]) -> list[float]: ...

class TeacherPoolActor(Actor):
    """Channel 3 — OpenRouter calls to N external teachers."""
    @endpoint
    async def replay(self, states: list[dict]) -> list[dict]: ...

# Topology
trainers   = mesh.spawn(TrainerActor, n=4, gpu="A100")
generator  = mesh.spawn(GeneratorActor, n=1, gpu="A100")
rewarder   = mesh.spawn(RewarderActor, n=1, gpu=None)
teachers   = mesh.spawn(TeacherPoolActor, n=1, gpu=None)

Status of this directory

  • monarch_actor_layout.md — this file (design)
  • actors.py — skeleton actor definitions; do not import without monarch installed
  • composer_mesh.py — composition glue; not yet implemented

Open questions (deferred to v0.2)

  • Does Monarch v0.5's Slurm backend hand-shake cleanly with HF Jobs? (HF Jobs runs each "job" as an independent container; Monarch wants to manage the lifecycle. Possible mismatch.)
  • Can the TrainerActor host the framework's ComposerReplicationTrainer unmodified, or does it need to be split into step_init / step_compute endpoints to fit Monarch's async actor model?

References