.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated/examples/grpo/kubernetes_grpo.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_examples_grpo_kubernetes_grpo.py: GRPO on Kubernetes (Qwen3.5-0.8B-Base on GSM8K) =============================================== This tutorial runs a GRPO-style RL fine-tuning loop on Kubernetes using Monarch's ``KubernetesJob`` with two heterogeneous ``MonarchMesh`` CRDs: - a ``learner`` mesh on one pod with 2 GPUs; the trainable policy is split across both via ``device_map="auto"`` so the model + AdamW state + activations fit comfortably; - a multi-pod ``generator`` mesh that rolls out G=6 completions per prompt and syncs updated policy weights back from the learner over ``RDMABuffer``. Prompts come from ``openai/gsm8k``. The model is asked via a system prompt to emit its final integer inside ``...`` tags. The reward is ``1.0`` for exact-string match of the extracted answer with the gold, ``0.0`` otherwise. The algorithm is intentionally minimal -- token-level REINFORCE on group-relative advantages, no KL anchor, no PPO clipping, no format shaping. The point is to show how to wire a Monarch actor mesh for distributed RL on Kubernetes (two heterogeneous meshes, replay-buffer actor, scorer actor, RDMA weight sync, async rollout / train). For KL-anchored production GRPO see the :doc:`grpo_actor` sibling example. Prerequisites ------------- - A Kubernetes cluster with GPU worker nodes; the learner pod requests ``nvidia.com/gpu: 2`` and each generator pod requests as many as ``--gpus_per_generator`` (default 4). - ``nvidia.com/gpu`` device plugin installed. - Monarch operator and ``MonarchMesh`` CRD installed (see https://github.com/meta-pytorch/monarch-kubernetes). - ``kubectl`` configured for the cluster. - Internet access from worker pods (for downloading Qwen3.5-0.8B-Base and pip installing ``transformers``). No HF token is required -- Qwen3.5 is Apache 2.0 and ungated. - The controller pod installs ``datasets`` at startup (see ``manifests/grpo_provision.yaml``); the worker pods do not need it because prompts are loaded in ``main()`` on the controller. RDMA transport caveat --------------------- This example calls ``monarch.configure(rdma_allow_tcp_fallback=True)`` so cross-pod ``RDMABuffer`` reads fall back to TCP on clusters without an RDMA CNI / device plugin. This is demonstrative, not performance-tuned. On a cluster with ibverbs plumbing, drop the ``configure()`` call and add the appropriate device resource requests to ``build_pod_template``. Pod startup time ---------------- Cold worker pods need a few minutes before they are ready: - ``pip install transformers tokenizers accelerate`` - Downloading the model weights from Hugging Face ``KubernetesJob(timeout=600)`` gives 10 minutes for the meshes to come up. Running the example ------------------- :: kubectl apply -f manifests/grpo_provision.yaml kubectl wait --for=condition=Ready pod/grpo-controller -n monarch-tests kubectl cp kubernetes_grpo.py monarch-tests/grpo-controller:/tmp/kubernetes_grpo.py kubectl exec -it grpo-controller -n monarch-tests -- python /tmp/kubernetes_grpo.py kubectl delete -f manifests/grpo_provision.yaml Each step prints one ``[Step NNN] loss=... reward=... gn=...`` line. ``reward`` is the mean correctness of the sampled replay-buffer batch; ``gn`` is the pre-clip grad norm. Steps with ``reward=0.000`` or ``reward=1.000`` show ``gn=0.000``: every completion in the group got the same score, so the group-relative advantage is zero and the gradient vanishes. Mixing in non-degenerate groups is what drives learning. The headline "did training work?" signal is the held-out GSM8K test accuracy: a baseline pass before step 0 and a final pass after the training loop, both sharded greedy-decode across the generator mesh. A representative 100-step run on the reference config (1 learner pod x 2 GPUs, 2 generator pods x 4 GPUs, and 512 eval prompts) ends up looking like this: :: ============================================================ Training summary ------------------------------------------------------------ Baseline acc: 0.375 (192/512) Final acc: 0.562 (288/512) Delta: +0.188 ============================================================ Larger Qwen3.5-Base policies show the same end-to-end story with bigger deltas. Bumping ``MODEL_NAME`` to ``Qwen/Qwen3.5-2B-Base`` (no other changes) on the same config produces: :: ============================================================ Training summary ------------------------------------------------------------ Baseline acc: 0.494 (253/512) Final acc: 0.742 (380/512) Delta: +0.248 ============================================================ The hyperparameters are reasonable defaults but this script is not tuned for SOTA GSM8K accuracy -- the point is to show how to wire a Monarch actor mesh for distributed RL on Kubernetes (two heterogeneous meshes, replay-buffer actor, scorer actor, RDMA weight sync, async rollout / train). Running the example on GKE ------------------- Follow the instructions in `Allocate network resources by using GKE managed DRANET `_ to create a GKE cluster with RDMA enabled. **NOTE:** The gpus_per_generator must match the generator resource claim template device count in ``resource_claim_templates.yaml``. :: kubectl apply -f manifests/grpo_provision.yaml kubectl apply -f manifests/gke/resource_claim_templates.yaml kubectl wait --for=condition=Ready pod/grpo-controller -n monarch-tests kubectl cp kubernetes_grpo.py monarch-tests/grpo-controller:/tmp/kubernetes_grpo.py kubectl exec -it grpo-controller -n monarch-tests -- python /tmp/kubernetes_grpo.py --gpus_per_generator 3 --learner_resource_claim_template "learner-rdma" --generator_resource_claim_template "generator-rdma" kubectl delete -f manifests/gke/resource_claim_templates.yaml kubectl delete -f manifests/grpo_provision.yaml .. GENERATED FROM PYTHON SOURCE LINES 142-144 Imports ------- .. GENERATED FROM PYTHON SOURCE LINES 144-179 .. code-block:: default import argparse import asyncio import math import random import re import textwrap import time from dataclasses import dataclass from enum import Enum from typing import Any, Dict, List, Optional, Tuple import monarch import torch import torch.nn as nn import torch.optim as optim from kubernetes.client import ( CoreV1ResourceClaim, V1Container, V1EmptyDirVolumeSource, V1EnvVar, V1PodResourceClaim, V1PodSpec, V1PodTemplateSpec, V1Probe, V1ResourceRequirements, V1TCPSocketAction, V1Volume, V1VolumeMount, ) from monarch._src.job.kubernetes import _WORKER_BOOTSTRAP_SCRIPT from monarch.actor import Actor, current_rank, current_size, endpoint from monarch.job.kubernetes import KubernetesJob from monarch.rdma import RDMABuffer .. GENERATED FROM PYTHON SOURCE LINES 180-184 Cross-pod RDMABuffer reads use TCP fallback by default on clusters without an RDMA CNI. Also extend supervision / spawn-idle timeouts so large model loads inside actor __init__ (which block the event loop for tens of seconds) don't get killed by the default 60s liveness watchdog. .. GENERATED FROM PYTHON SOURCE LINES 184-192 .. code-block:: default monarch.configure( rdma_allow_tcp_fallback=True, actor_spawn_max_idle="10m", get_proc_state_max_idle="10m", supervision_watchdog_timeout="10m", enable_log_forwarding=True, ) .. GENERATED FROM PYTHON SOURCE LINES 193-195 Constants --------- .. GENERATED FROM PYTHON SOURCE LINES 195-203 .. code-block:: default MODEL_NAME = "Qwen/Qwen3.5-0.8B-Base" G = 6 # GRPO group size (completions sampled per prompt) MAX_NEW_TOKENS = 256 MONARCH_NAMESPACE = "monarch-tests" MONARCH_IMAGE = "ghcr.io/meta-pytorch/monarch:latest" .. GENERATED FROM PYTHON SOURCE LINES 204-212 Task-specific: GSM8K reward, prompts, and dataset ------------------------------------------------- Everything below in this section is GSM8K-specific. To adapt this example to another task, replace this block: define a new prompt template, parsing helpers, dataset loader, and a class with a ``score`` method matching the ``GSM8KScorePolicy.score`` signature. Nothing in the actor topology, RL loss, or RDMA weight sync below knows about GSM8K. .. GENERATED FROM PYTHON SOURCE LINES 212-324 .. code-block:: default SYSTEM_PROMPT = ( "Solve the problem and put your final integer answer inside " "... tags." ) # Matches the GSM8K dataset's native '#### ' ground-truth marker, used # only to pull the gold answer out of ``ex["answer"]`` at dataset load time. GSM8K_GT_PATTERN = re.compile(r"####\s*(-?\d+)") def load_gsm8k_prompts(split: str, num_prompts: int) -> List[Tuple[str, str]]: """Load the first ``num_prompts`` problems from GSM8K ``split``. Returns a list of (prompt_text, ground_truth_string) tuples. Pass ``num_prompts=0`` to use the entire split. .. warning:: Downloading the dataset from the Hugging Face Hub at job launch is for DEMONSTRATION ONLY. In production, stage the dataset onto a persistent volume (e.g. a PVC or an object-store-backed mount) and point ``datasets.load_dataset`` at the local path. This avoids repeated downloads, removes a dependency on outbound network access, and gives reproducible, auditable inputs across runs. """ from datasets import load_dataset ds = load_dataset("openai/gsm8k", "main", split=split) if num_prompts > 0: ds = ds.select(range(min(num_prompts, len(ds)))) out: List[Tuple[str, str]] = [] for ex in ds: m = GSM8K_GT_PATTERN.search(ex["answer"]) if m is None: continue out.append((ex["question"], m.group(1))) return out def extract_xml_answer(text: str) -> Optional[str]: """Return the content of the last ``...`` block, stripped. Returns ``None`` if no well-formed block exists or the content is empty. """ parts = text.split("") if len(parts) < 2: return None last = parts[-1] if "" not in last: return None answer = last.split("")[0].strip() return answer or None def is_correct(generated_text: str, ground_truth: str) -> bool: """True when the answer extracted from ``generated_text`` matches the gold.""" return extract_xml_answer(generated_text) == ground_truth def format_prompt(tokenizer: Any, prompt_text: str) -> str: """Wrap ``prompt_text`` in the GSM8K system prompt + chat template.""" messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt_text}, ] try: return tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) except (AttributeError, ValueError): return f"{SYSTEM_PROMPT}\n\nQuestion:\n{prompt_text}\n\n\n" class GSM8KScorePolicy: """Rule-based reward for GSM8K: 1.0 for exact answer match, else 0.0. Plug a different task in by writing a class with a ``score`` method matching the signature here: ``score(generations, gen_mask, ground_truth) -> rewards``. """ def __init__(self, model_name: str) -> None: # Lazy tokenizer so this object pickles cheaply into the Scorer proc. self.model_name = model_name self._tokenizer: Any = None def _ensure_tokenizer(self) -> Any: if self._tokenizer is None: from transformers.models.auto.tokenization_auto import AutoTokenizer self._tokenizer = AutoTokenizer.from_pretrained(self.model_name) return self._tokenizer def score( self, generations: torch.Tensor, gen_mask: torch.Tensor, ground_truth: str, ) -> torch.Tensor: """Return per-completion total reward over the [G] completions.""" tokenizer = self._ensure_tokenizer() n = generations.shape[0] token_lists = [generations[i][gen_mask[i].bool()].tolist() for i in range(n)] texts = tokenizer.batch_decode(token_lists, skip_special_tokens=True) rewards = torch.zeros(n) for i, text in enumerate(texts): rewards[i] = 1.0 if extract_xml_answer(text) == ground_truth else 0.0 return rewards .. GENERATED FROM PYTHON SOURCE LINES 325-327 Generic helpers --------------- .. GENERATED FROM PYTHON SOURCE LINES 327-356 .. code-block:: default def gather_logprobs( sliced_logits: torch.Tensor, target_ids: torch.Tensor ) -> torch.Tensor: """Memory-efficient per-token log-probability of ``target_ids``. Computes ``target_logit - logsumexp(all_logits)`` rather than materialising a full ``[B, L, V]`` fp32 ``log_softmax`` tensor, which saves several GB of fp32 activations at typical batch/sequence/vocab sizes. """ target_logits = sliced_logits.gather(-1, target_ids.unsqueeze(-1)).squeeze(-1) lse = torch.logsumexp(sliced_logits.float(), dim=-1) return target_logits - lse def select_runtime_device() -> str: return "cuda" if torch.cuda.is_available() else "cpu" def select_torch_dtype() -> torch.dtype: return torch.bfloat16 if torch.cuda.is_available() else torch.float32 def pad_trailing(x: torch.Tensor, pad_len: int, fill: float) -> torch.Tensor: """Right-pad a ``[G, L]`` tensor along dim 1 with ``fill`` to width ``L + pad_len``.""" return torch.nn.functional.pad(x, (0, pad_len), value=fill) .. GENERATED FROM PYTHON SOURCE LINES 357-359 Data structures --------------- .. GENERATED FROM PYTHON SOURCE LINES 359-398 .. code-block:: default @dataclass class TrajectorySlice: """One rollout: a prompt + G sampled completions. Attributes: policy_version: Policy version that generated this slice. prompt_ids: Tokenized prompt, shape [prompt_len]. generations: Generated token ids, shape [G, MAX_NEW_TOKENS]. Padded after each sequence's first EOS. gen_mask: 1.0 for real generated tokens, 0.0 after first EOS, shape [G, MAX_NEW_TOKENS]. Masks the loss. rewards: Per-completion total reward, shape [G]. Zero until the score policy fills it in. ground_truth: Correct answer string passed to the score policy. """ policy_version: int prompt_ids: torch.Tensor generations: torch.Tensor gen_mask: torch.Tensor rewards: torch.Tensor ground_truth: str @dataclass class PolicyBatch: """Tensor bundle for one learner-side gradient step.""" input_ids: torch.Tensor attention_mask: torch.Tensor generations: torch.Tensor gen_mask: torch.Tensor advantages: torch.Tensor row_idx: torch.Tensor col_idx: torch.Tensor .. GENERATED FROM PYTHON SOURCE LINES 399-401 Actors ------ .. GENERATED FROM PYTHON SOURCE LINES 401-869 .. code-block:: default class ReplayBuffer(Actor): """Storage for scored slices with policy-version-weighted sampling.""" def __init__(self) -> None: self.storage: List[Tuple[int, TrajectorySlice]] = [] self.storage_event = asyncio.Event() @endpoint async def put(self, slice_: TrajectorySlice) -> None: # Keep storage to one policy version: drop older, ignore stale. # Sampling stale slices would bias the on-policy gradient estimator. v = slice_.policy_version if self.storage: cur = self.storage[0][0] if v > cur: self.storage.clear() self.storage_event.clear() elif v < cur: return self.storage.append((v, slice_)) self.storage_event.set() async def _wait_for_storage(self) -> None: if not self.storage: await self.storage_event.wait() @endpoint async def sample_from(self, k: int) -> List[TrajectorySlice]: try: await asyncio.wait_for(self._wait_for_storage(), timeout=30.0) except asyncio.TimeoutError: raise RuntimeError("Timeout waiting for ReplayBuffer to populate") fresh = [s for _, s in self.storage] chosen = random.choices(range(len(fresh)), k=k) return [fresh[i] for i in chosen] class Scorer(Actor): """Generic reward dispatcher. Owns the trajectory queue. Generators call ``put`` to enqueue completions; ``run`` consumes them, delegates scoring to ``score_policy``, and forwards the scored slice to the replay buffer. ``score_policy`` is any object exposing ``score(generations, gen_mask, ground_truth) -> rewards``. """ def __init__(self, score_policy: Any, replay_buffer: Any) -> None: self.score_policy = score_policy self.replay_buffer = replay_buffer self.queue: asyncio.Queue[TrajectorySlice] = asyncio.Queue() self.running = False @endpoint async def put(self, slice_: TrajectorySlice) -> None: await self.queue.put(slice_) async def _score_slice(self, slice_: TrajectorySlice) -> None: rewards = self.score_policy.score( slice_.generations, slice_.gen_mask, slice_.ground_truth ) scored = TrajectorySlice( policy_version=slice_.policy_version, prompt_ids=slice_.prompt_ids, generations=slice_.generations, gen_mask=slice_.gen_mask, rewards=rewards, ground_truth=slice_.ground_truth, ) await self.replay_buffer.put.call(scored) @endpoint async def run(self) -> None: if self.running: return self.running = True try: while self.running: try: # 10s wait_for bounds shutdown latency after scorer.shutdown(). slice_ = await asyncio.wait_for( self.queue.get(), timeout=10.0, ) except asyncio.TimeoutError: continue await self._score_slice(slice_) finally: self.running = False @endpoint async def shutdown(self) -> None: self.running = False class GeneratorState(Enum): READY_TO_GENERATE = "READY_TO_GENERATE" READY_TO_UPDATE = "READY_TO_UPDATE" class Generator(Actor): """Generates G completions per prompt and ships them to the Scorer. Holds an inference-only copy of the policy model. Between rollouts, syncs updated weights from the Learner via RDMABuffer. """ def __init__( self, model_name: str, weight_buffers: Dict[str, RDMABuffer], scorer: Any, ) -> None: from transformers.models.auto.modeling_auto import AutoModelForCausalLM from transformers.models.auto.tokenization_auto import AutoTokenizer if torch.cuda.is_available(): torch.cuda.set_device(current_rank()["gpus"]) self.device = select_runtime_device() self.model_dtype = select_torch_dtype() self.tokenizer = AutoTokenizer.from_pretrained(model_name) if self.tokenizer.pad_token_id is None: self.tokenizer.pad_token_id = self.tokenizer.eos_token_id self.model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=self.model_dtype, attn_implementation="eager", trust_remote_code=True, ).to(self.device) self.model.eval() self.weight_buffers = weight_buffers # Cache uint8 views of the live policy weights once. Storage is stable # across in-place writes (RDMABuffer.read_into), so the same view # always points at the right bytes -- no need to rebuild from # state_dict() every weight sync. Iterate named_parameters() to mirror # the learner-side export and avoid hitting registered buffers (e.g. # rotary_emb.inv_freq) that may not support .view(torch.uint8). self._param_views: Dict[str, torch.Tensor] = { n: p.data.view(torch.uint8).flatten() for n, p in self.model.named_parameters() if n in weight_buffers } self.scorer = scorer self.state = GeneratorState.READY_TO_GENERATE self.cond = asyncio.Condition() self.policy_version = 0 @endpoint async def generate(self, prompt_text: str, ground_truth: str) -> None: async with self.cond: await self.cond.wait_for( lambda: self.state == GeneratorState.READY_TO_GENERATE ) templated = format_prompt(self.tokenizer, prompt_text) prompt_ids = ( self.tokenizer(templated, return_tensors="pt") .input_ids[0] .to(self.device) ) input_ids = prompt_ids.unsqueeze(0).repeat(G, 1) attention_mask = torch.ones_like(input_ids) with torch.no_grad(): out = self.model.generate( input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=1.0, pad_token_id=self.tokenizer.pad_token_id, return_dict_in_generate=True, ) gen_tokens = out.sequences[:, input_ids.shape[1] :] # noqa: E203 eos_id = self.tokenizer.eos_token_id post_eos = (gen_tokens == eos_id).cumsum(dim=-1) > 1 gen_mask = (~post_eos).to(torch.float32) # HuggingFace generate() returns sequences of length max_over_group, # which can be shorter than MAX_NEW_TOKENS when every sampled # completion emits EOS early. The Learner downstream reshapes # generations as [-1, MAX_NEW_TOKENS] and requires this exact # trailing dim, so pad out to the full budget here. gl_raw = gen_tokens.shape[1] if gl_raw < MAX_NEW_TOKENS: pad_len = MAX_NEW_TOKENS - gl_raw gen_tokens = pad_trailing( gen_tokens, pad_len, self.tokenizer.pad_token_id ) gen_mask = pad_trailing(gen_mask, pad_len, 0.0) slice_ = TrajectorySlice( policy_version=self.policy_version, prompt_ids=prompt_ids.cpu(), generations=gen_tokens.cpu(), gen_mask=gen_mask.cpu(), rewards=torch.zeros(G), ground_truth=ground_truth, ) await self.scorer.put.call(slice_) async with self.cond: self.state = GeneratorState.READY_TO_UPDATE self.cond.notify_all() @endpoint async def update(self, version: int) -> None: async with self.cond: total_bytes = sum(v.numel() for v in self._param_views.values()) t0 = time.monotonic() await asyncio.gather( *( buf.read_into(self._param_views[n], timeout=30) for n, buf in self.weight_buffers.items() ) ) dt = time.monotonic() - t0 gb = total_bytes / 1e9 print( f"[Generator] weight sync v{version}: " f"{gb:.3f} GB in {dt:.2f}s = {gb / dt if dt > 0 else 0:.2f} GB/s " f"({len(self.weight_buffers)} buffers)", flush=True, ) self.policy_version = version self.state = GeneratorState.READY_TO_GENERATE self.cond.notify_all() @endpoint async def eval_shard(self, eval_prompts: List[Tuple[str, str]]) -> Tuple[int, int]: """Greedy-decode this proc's shard of ``eval_prompts``. Returns ``(correct, total)`` for this shard. ``.call()`` from the caller gathers one such tuple per generator proc into a ``ValueMesh``, which the caller sums to get the totals. Sharding uses ``current_rank().rank`` as a flat global rank across the full ``hosts * gpus_per_host`` generator mesh, so eval wall-clock scales with ``1 / world_size``. """ rank = current_rank().rank world_size = math.prod(int(v) for v in current_size().values()) my_shard = eval_prompts[rank::world_size] # Eval runs outside the rollout/update protocol (only between # training phases, with no concurrent ``generate`` or ``update`` in # flight), so it doesn't gate on ``state``. The lock is defensive, # to serialize with any unexpected concurrent endpoint call. async with self.cond: self.model.eval() correct = 0 for prompt_text, ground_truth in my_shard: templated = format_prompt(self.tokenizer, prompt_text) input_ids = self.tokenizer(templated, return_tensors="pt").input_ids.to( self.device ) with torch.no_grad(): out = self.model.generate( input_ids=input_ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, pad_token_id=self.tokenizer.pad_token_id, ) gen = out[0, input_ids.shape[1] :] # noqa: E203 text = self.tokenizer.decode(gen, skip_special_tokens=True) if is_correct(text, ground_truth): correct += 1 return correct, len(my_shard) class Learner(Actor): """Token-level GRPO update: REINFORCE on group-relative advantages.""" def __init__(self, model_name: str, replay_buffer: Any) -> None: # Use direct submodule imports rather than `from transformers import # AutoModelForCausalLM` because transformers 5.x uses a lazy-loading # __init__.py that intermittently fails in Monarch actor subprocesses. from transformers.models.auto.modeling_auto import AutoModelForCausalLM from transformers.models.auto.tokenization_auto import AutoTokenizer # Input tensors live on cuda:0; ``device_map="auto"`` puts the # embedding + first layers there, so this avoids a copy on the # forward path. Backward auto-shards across whichever devices hold # each layer's parameters. self.device = "cuda:0" if torch.cuda.is_available() else "cpu" self.model_dtype = select_torch_dtype() self.tokenizer = AutoTokenizer.from_pretrained(model_name) if self.tokenizer.pad_token_id is None: self.tokenizer.pad_token_id = self.tokenizer.eos_token_id # ``device_map="auto"`` splits the trainable policy across all # visible GPUs (2 on the learner pod) so the model + AdamW state # + activations fit comfortably. Eager attention avoids flash_attn # dependencies; gradient checkpointing trades compute for activation # memory. device_map = "auto" if torch.cuda.is_available() else None self.model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=self.model_dtype, attn_implementation="eager", trust_remote_code=True, device_map=device_map, ) self.model.gradient_checkpointing_enable() self.optim = optim.AdamW(self.model.parameters(), lr=5e-6) self.grad_clip_norm = 1.0 self.policy_version = 0 self.replay_buffer = replay_buffer self.batch_size = 1 self.generators: Optional[Any] = None # Pins the underlying ``p.data`` tensor storage for the lifetime of # the RDMABuffer registrations; see ``weights_handle``. self._pinned_weight_buffers: Dict[str, Tuple[torch.Tensor, RDMABuffer]] = {} @endpoint async def init_generators(self, generators: Any) -> None: self.generators = generators @endpoint async def weights_handle(self) -> Dict[str, RDMABuffer]: # Iterate parameters (not state_dict) so non-param buffers like # rotary_emb.inv_freq -- which may not support .view(torch.uint8) # -- are skipped. The local ``(p.data, buf)`` pair pins the tensor # storage so the RDMA registration stays valid; only the buffer is # shipped across the mesh boundary. Idempotent so repeat calls # return the same buffers without re-registering. if not self._pinned_weight_buffers: self._pinned_weight_buffers = { n: (p.data, RDMABuffer(p.data.view(torch.uint8).flatten())) for n, p in self.model.named_parameters() } return {n: buf for n, (_, buf) in self._pinned_weight_buffers.items()} def _compute_advantages(self, rewards: torch.Tensor) -> torch.Tensor: """GRPO-paper per-prompt advantage normalisation. For each prompt group of ``G`` completions, subtract the group mean and divide by the group standard deviation. """ batch_size = rewards.shape[0] // G reshaped = rewards.view(batch_size, G) means = reshaped.mean(dim=1, keepdim=True) stds = reshaped.std(dim=1, keepdim=True) advs = (reshaped - means) / (stds + 1e-4) return advs.reshape(-1) def _build_policy_batch( self, prompt_ids_list: List[torch.Tensor], generations: torch.Tensor, gen_mask: torch.Tensor, advantages: torch.Tensor, ) -> PolicyBatch: device = self.device pad_id = self.tokenizer.pad_token_id rows: List[torch.Tensor] = [] prompt_lens: List[int] = [] for i, prompt_ids in enumerate(prompt_ids_list): prompt_len = int(prompt_ids.shape[0]) for g in range(G): rows.append(torch.cat([prompt_ids, generations[i * G + g]])) prompt_lens.append(prompt_len) max_len = max(row.shape[0] for row in rows) batch_size = len(rows) input_ids = torch.full( (batch_size, max_len), pad_id, dtype=torch.long, device=device ) attention_mask = torch.zeros( (batch_size, max_len), dtype=torch.long, device=device ) for i, row in enumerate(rows): row_on_device = row.to(device) input_ids[i, : row.shape[0]] = row_on_device attention_mask[i, : row.shape[0]] = 1 generations_dev = generations.to(device) gen_len = generations_dev.shape[1] row_idx = torch.arange(batch_size, device=device).unsqueeze(1) prompt_len_tensor = torch.tensor(prompt_lens, device=device) col_idx = (prompt_len_tensor - 1).unsqueeze(1) + torch.arange( gen_len, device=device ) return PolicyBatch( input_ids=input_ids, attention_mask=attention_mask, generations=generations_dev, gen_mask=gen_mask.to(device), advantages=advantages.to(device), row_idx=row_idx, col_idx=col_idx, ) def _apply_policy_update(self, batch: PolicyBatch) -> Tuple[torch.Tensor, float]: """One REINFORCE update. Returns ``(loss, grad_norm)``.""" # ``device_map="auto"`` may place the lm_head on cuda:1; move the # logits back to ``self.device`` so row_idx/col_idx (cuda:0) can # index them. ``.to()`` is differentiable, so backward flows back # to cuda:1 transparently. policy_logits = self.model( batch.input_ids, attention_mask=batch.attention_mask ).logits.to(self.device) policy_selected = policy_logits[batch.row_idx, batch.col_idx] new_logps = gather_logprobs(policy_selected, batch.generations) token_count = batch.gen_mask.sum().clamp(min=1.0) advantages_b = batch.advantages.unsqueeze(-1) loss = -(advantages_b * new_logps * batch.gen_mask).sum() / token_count self.optim.zero_grad() loss.backward() grad_norm = nn.utils.clip_grad_norm_( self.model.parameters(), self.grad_clip_norm ) self.optim.step() self.policy_version += 1 return loss.detach(), float(grad_norm) @endpoint async def step(self) -> Dict[str, float]: """Run one gradient update. Returns per-step training metrics.""" # The weight push and the replay-buffer sample are independent; # overlap them so the RDMA round-trip hides behind the local fetch. sample_coro = self.replay_buffer.sample_from.call_one(self.batch_size) if self.generators is not None: _, slices = await asyncio.gather( self.generators.update.call(self.policy_version), sample_coro, ) else: slices = await sample_coro prompt_ids_list = [s.prompt_ids for s in slices] generations = torch.stack([s.generations for s in slices]).view( -1, MAX_NEW_TOKENS ) gen_mask = torch.cat([s.gen_mask for s in slices]) rewards = torch.cat([s.rewards for s in slices]) advs = self._compute_advantages(rewards) batch = self._build_policy_batch(prompt_ids_list, generations, gen_mask, advs) loss, grad_norm = self._apply_policy_update(batch) return { "loss": float(loss), "reward_mean": float(rewards.mean()), "grad_norm": grad_norm, } @endpoint async def sync_generators(self) -> None: """Push the learner's current policy weights to every generator. Eval runs on the generator mesh so it can shard the eval set and decode ``world_size``-way in parallel. ``step`` syncs generators once per gradient update, before the gradient itself runs, so by the end of training the generators carry the pre-final-gradient policy and are one step stale; this method is the explicit flush, called immediately before a sharded eval on the final policy. """ if self.generators is not None: await self.generators.update.call(self.policy_version) .. GENERATED FROM PYTHON SOURCE LINES 870-872 Pod spec -------- .. GENERATED FROM PYTHON SOURCE LINES 872-990 .. code-block:: default # Installing Python dependencies at container startup is for DEMONSTRATION # ONLY. In production, bake ``transformers``, ``tokenizers``, and # ``accelerate`` into the worker container image. PIP_INSTALL = textwrap.dedent("""\ import subprocess, sys # --break-system-packages is required on the Debian-based monarch image # (PEP 668); without it pip refuses to install into the system Python. subprocess.check_call([ sys.executable, "-m", "pip", "install", "--quiet", "--break-system-packages", "--upgrade", "transformers", "tokenizers", "accelerate", ]) """) def build_pod_template( gpus: int, resource_claim_template: Optional[str] = None ) -> V1PodTemplateSpec: """Shared pod template for learner and generator meshes. Prepends a pip-install prefix to the Monarch worker bootstrap so that ``transformers``, ``tokenizers``, and ``accelerate`` are present before ``run_worker_loop_forever`` is imported. Sets: - ``MONARCH_PORT=26600`` (matches the default label selector). - ``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`` (required by RDMABuffer's CUDA allocator check). - ``HF_HOME=/tmp/hf_cache`` so the Hugging Face cache lands on rootfs ephemeral storage. Mounts ``/dev/shm`` as a 16 GiB memory-backed emptyDir for shared memory used by NCCL / multi-worker inference. """ bootstrap = PIP_INSTALL + _WORKER_BOOTSTRAP_SCRIPT resources = None node_selector = None env = [ V1EnvVar(name="MONARCH_PORT", value="26600"), # /tmp/hf_cache forces the Hugging Face cache onto # rootfs ephemeral storage for DEMONSTRATION ONLY -- # every pod redownloads the full model at startup. In # production, stage model weights onto a persistent # volume (e.g. a read-only PVC populated by an image # pre-puller or an init container) and point HF_HOME # there, so worker pods start instantly and do not # hammer the HF Hub. V1EnvVar(name="HF_HOME", value="/tmp/hf_cache"), ] resource_claims = [] claims = [] if gpus > 0: if resource_claim_template: resource_claims = [ V1PodResourceClaim( name="rdma", resource_claim_template_name=resource_claim_template ) ] claims = [CoreV1ResourceClaim(name="rdma")] gpu_resources = {"nvidia.com/gpu": str(gpus)} resources = V1ResourceRequirements( limits=gpu_resources, requests=gpu_resources, claims=claims if claims else None, ) env.insert( 1, V1EnvVar( name="PYTORCH_CUDA_ALLOC_CONF", value="expandable_segments:True", ), ) return V1PodTemplateSpec( spec=V1PodSpec( containers=[ V1Container( name="worker", image=MONARCH_IMAGE, command=["python", "-u", "-c", bootstrap], env=env, resources=resources, # Mark pod Ready only once the Monarch worker loop is actually # listening on :26600. Without this, K8s reports Ready as soon # as the container starts -- which is during pip install, long # before run_worker_loop_forever binds the port. The controller # then connects too early and the 30s message delivery timeout # fires before the worker is reachable. readiness_probe=V1Probe( tcp_socket=V1TCPSocketAction(port=26600), initial_delay_seconds=30, period_seconds=5, timeout_seconds=5, failure_threshold=60, ), volume_mounts=[ V1VolumeMount(name="dshm", mount_path="/dev/shm"), ], ) ], node_selector=node_selector, resource_claims=resource_claims if resource_claims else None, volumes=[ V1Volume( name="dshm", empty_dir=V1EmptyDirVolumeSource( medium="Memory", size_limit="16Gi" ), ) ], ), ) .. GENERATED FROM PYTHON SOURCE LINES 991-993 Orchestration ------------- .. GENERATED FROM PYTHON SOURCE LINES 993-1169 .. code-block:: default async def run_sharded_eval( learner: Any, generators: Any, prompts_for_eval: List[Tuple[str, str]], sync_first: bool, ) -> Tuple[int, int]: """Sharded greedy-decode eval across the generator mesh. When ``sync_first`` is True, push the learner's current policy weights to every generator first; otherwise rely on the caller to know the meshes are already aligned (e.g. baseline eval immediately after init). Each generator proc decodes ``prompts_for_eval[rank::world_size]``; this function sums the per-shard ``(correct, total)`` tuples. """ if sync_first: await learner.sync_generators.call_one() value_mesh = await generators.eval_shard.call(prompts_for_eval) correct = 0 total = 0 for k, n in value_mesh.values(): correct += k total += n return correct, total async def main( model_name: str, num_generator_hosts: int, gpus_per_generator: int, training_steps: int, namespace: str, dataset_split: str, num_prompts: int, eval_size: int, learner_resource_claim_template: Optional[str] = None, generator_resource_claim_template: Optional[str] = None, ) -> None: """Run GRPO fine-tuning across the learner and generator meshes.""" prompts = load_gsm8k_prompts(split=dataset_split, num_prompts=num_prompts) eval_prompts = load_gsm8k_prompts(split="test", num_prompts=eval_size) print("=" * 60) print(f"Kubernetes GRPO ({model_name})") print( f"Config: 1 learner pod x 2 GPUs | " f"{num_generator_hosts} generator pods x {gpus_per_generator} GPU(s)" ) print(f"Namespace: {namespace} | Training steps: {training_steps}") print(f"Dataset: openai/gsm8k[{dataset_split}] ({len(prompts)} prompts)") print(f"Eval: openai/gsm8k[test] ({len(eval_prompts)} prompts)") if learner_resource_claim_template: print(f"Learner Resource Claim Template: {learner_resource_claim_template}") if generator_resource_claim_template: print(f"Generator Resource Claim Template: {generator_resource_claim_template}") print("=" * 60) # 600s timeout lets the meshes finish cold-start pip install + model # download before state() gives up. k8s_job = KubernetesJob(namespace=namespace, timeout=600) # Learner pod requests 2 GPUs; ``device_map="auto"`` inside the actor # spreads the trainable policy across both. ``spawn_procs({"gpus": 1})`` # below still spawns a single learner proc that sees both GPUs in its # CUDA namespace. To fine-tune a larger base model, bump this to 4 # (e.g. for a 4B-parameter policy, the static memory of params + bf16 # gradients + fp32 AdamW state alone is ~48 GB, requiring 4 x 22 GB GPUs). k8s_job.add_mesh( "learner", num_replicas=1, pod_template=build_pod_template( gpus=2, resource_claim_template=learner_resource_claim_template ), ) k8s_job.add_mesh( "generator", num_replicas=num_generator_hosts, pod_template=build_pod_template( gpus=gpus_per_generator, resource_claim_template=generator_resource_claim_template, ), ) learner_mesh = None gen_mesh = None try: job_state = k8s_job.state() learner_mesh = job_state.learner.spawn_procs({"gpus": 1}) gen_mesh = job_state.generator.spawn_procs({"gpus": gpus_per_generator}) await asyncio.gather( learner_mesh.logging_option(stream_to_client=True), gen_mesh.logging_option(stream_to_client=True), ) score_policy = GSM8KScorePolicy(model_name) replay_buf = learner_mesh.spawn("rb", ReplayBuffer) learner = learner_mesh.spawn("learner", Learner, model_name, replay_buf) scorer = learner_mesh.spawn("scorer", Scorer, score_policy, replay_buf) wb = await learner.weights_handle.call_one() generators = gen_mesh.spawn("generator", Generator, model_name, wb, scorer) await learner.init_generators.call(generators) # Seed the replay buffer with one scored slice before training so # step 0's ``sample_from`` doesn't race with its concurrent rollout. seed_prompt, seed_gt = prompts[0] await generators.generate.call(seed_prompt, seed_gt) scorer_run_future = scorer.run.call_one() # Baseline eval before any training. Generators and learner are # both freshly loaded from the same HF checkpoint, so skip the sync. baseline_correct, baseline_total = await run_sharded_eval( learner, generators, eval_prompts, sync_first=False ) baseline_acc = baseline_correct / max(1, baseline_total) print( f"[Baseline] acc={baseline_acc:.3f} ({baseline_correct}/{baseline_total})", flush=True, ) for step in range(training_steps): idx = step % len(prompts) prompt, gt = prompts[idx] _, stats = await asyncio.gather( generators.generate.call(prompt, gt), learner.step.call_one(), ) print( f"[Step {step:03d}] " f"loss={stats['loss']:+.4f} " f"reward={stats['reward_mean']:.3f} " f"gn={stats['grad_norm']:.3f}", flush=True, ) final_correct, final_total = await run_sharded_eval( learner, generators, eval_prompts, sync_first=True ) final_acc = final_correct / max(1, final_total) print("", flush=True) print("=" * 60, flush=True) print("Training summary", flush=True) print("-" * 60, flush=True) print( f"Baseline acc: {baseline_acc:.3f} ({baseline_correct}/{baseline_total})", flush=True, ) print( f"Final acc: {final_acc:.3f} ({final_correct}/{final_total})", flush=True, ) print(f"Delta: {final_acc - baseline_acc:+.3f}", flush=True) print("=" * 60, flush=True) # Shutdown latency bounded by the 10s wait_for in Scorer.run(). print("Stopping scorer...", flush=True) await scorer.shutdown.call_one() await scorer_run_future print("Training complete.", flush=True) finally: # Best-effort per-mesh cleanup; one failure must not skip the other. if learner_mesh is not None: try: learner_mesh.stop().get() except Exception as e: print(f"[cleanup] learner_mesh.stop() failed: {e}") if gen_mesh is not None: try: gen_mesh.stop().get() except Exception as e: print(f"[cleanup] gen_mesh.stop() failed: {e}") # Unconditional so MonarchMesh CRDs and pods are always torn down. k8s_job.kill() .. GENERATED FROM PYTHON SOURCE LINES 1170-1172 CLI --- .. GENERATED FROM PYTHON SOURCE LINES 1172-1252 .. code-block:: default if __name__ == "__main__": parser = argparse.ArgumentParser( description="Run GRPO fine-tuning on Kubernetes with Qwen3.5-0.8B-Base" ) parser.add_argument( "--model_name", type=str, default=MODEL_NAME, help="HF model used for learner and generator actors", ) parser.add_argument( "--num_generator_hosts", type=int, default=2, help="Number of generator worker pods", ) parser.add_argument( "--gpus_per_generator", type=int, default=4, help="GPUs per generator pod; set 0 for CPU", ) parser.add_argument( "--training_steps", type=int, default=100, help="Number of gradient update iterations", ) parser.add_argument( "--namespace", type=str, default=MONARCH_NAMESPACE, help="Kubernetes namespace for the MonarchMesh CRDs", ) parser.add_argument( "--dataset_split", type=str, default="train", help="GSM8K split to load (train or test)", ) parser.add_argument( "--num_prompts", type=int, default=1000, help="Number of GSM8K prompts to load; 0 means the full split", ) parser.add_argument( "--eval_size", type=int, default=512, help="Number of held-out GSM8K test prompts used for eval.", ) parser.add_argument( "--learner_resource_claim_template", type=str, default=None, help="Name of the ResourceClaimTemplate for Learner RDMA (e.g. 'two-rdma')", ) parser.add_argument( "--generator_resource_claim_template", type=str, default=None, help="Name of the ResourceClaimTemplate for Generator RDMA (e.g. 'two-rdma')", ) args = parser.parse_args() asyncio.run( main( model_name=args.model_name, num_generator_hosts=args.num_generator_hosts, gpus_per_generator=args.gpus_per_generator, training_steps=args.training_steps, namespace=args.namespace, dataset_split=args.dataset_split, num_prompts=args.num_prompts, eval_size=args.eval_size, learner_resource_claim_template=args.learner_resource_claim_template, generator_resource_claim_template=args.generator_resource_claim_template, ) ) .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.000 seconds) .. _sphx_glr_download_generated_examples_grpo_kubernetes_grpo.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: kubernetes_grpo.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: kubernetes_grpo.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_