diff --git a/examples/tinker/moe/modal_moe_tinker.py b/examples/tinker/moe/modal_moe_tinker.py new file mode 100644 index 0000000000..134e324e40 --- /dev/null +++ b/examples/tinker/moe/modal_moe_tinker.py @@ -0,0 +1,385 @@ +"""Serve an MoE model behind SkyRL's Tinker API server on Modal, and submit client jobs to it. + + +demo -- one gpu container runs tinker server + gsm8k client + +serve -- long lived gpu container running tinker server, exposed through modal tunnel + +client -- cpu only container that runs the client against server's modal url + +engine-test -- isolation test: boot ONE plain vLLM server (no Tinker, no megatron, empty GPUs) and time it to first /health 200. Distinguishes 'vLLM+model is slow/stuck' from 'the colocated Tinker path breaks engine boot' + +Qwen3-30B-a3B-base recipe: TP2/EP8/ETP1 for training, 2 vllm engines x TP4 for inference + +Usage (from the SkyRL repo root): + # long-lived server; prints the tunnel URL, stays up until you stop it + modal run examples/tinker/moe/modal_moe_tinker.py --mode serve + + # then submit client job against running server + modal run examples/tinker/moe/modal_moe_tinker.py --mode client --base-url +""" + +from __future__ import annotations + +import json +import os +import pathlib +import shlex +from typing import Optional + +import modal + +REMOTE_REPO = "/root/SkyRL" +HF_CACHE = "/root/.cache/huggingface" +CKPT_DIR = "/root/ckpts/tinker_moe" +PORT = 8000 +BASE_URL = f"http://localhost:{PORT}" + +DEFAULT_BASE_MODEL = "Qwen/Qwen3-30B-A3B-Base" +GPU_SPEC = os.environ.get("SKYRL_TINKER_GPU", "H200:8") +NUM_GPUS = int(GPU_SPEC.split(":")[1]) if ":" in GPU_SPEC else 1 + + +DEFAULT_BACKEND_CONFIG: dict = { + "trainer.placement.colocate_all": True, + "trainer.placement.policy_num_gpus_per_node": NUM_GPUS, + + "trainer.policy.megatron_config.tensor_model_parallel_size": 2, + "trainer.policy.megatron_config.pipeline_model_parallel_size": 1, + "trainer.policy.megatron_config.expert_model_parallel_size": NUM_GPUS, + "trainer.policy.megatron_config.expert_tensor_parallel_size": 1, + "trainer.micro_forward_batch_size_per_gpu": 2, + "trainer.micro_train_batch_size_per_gpu": 1, + + "generator.inference_engine.num_engines": 2, + "generator.inference_engine.tensor_parallel_size": NUM_GPUS // 2, + "generator.inference_engine.backend": "vllm", + "generator.inference_engine.run_engines_locally": True, + "generator.inference_engine.weight_sync_backend": "nccl", + + "generator.inference_engine.gpu_memory_utilization": 0.5, + "generator.batched": True, + + "trainer.log_path": f"{CKPT_DIR}/logs", +} + + +def _is_repo_root(path: pathlib.Path) -> bool: + try: + return ( + path.exists() + and (path / "pyproject.toml").exists() + and (path / "examples").exists() + and (path / "skyrl").exists() + and (path / "skyrl-gym").exists() + ) + except OSError: + # e.g. PermissionError probing /root/SkyRL when running locally as a + # non-root user (Path.exists() raises on EACCES rather than returning + # False on older Pythons). + return False + + +def _find_repo_root() -> pathlib.Path: + candidates: list[pathlib.Path] = [] + env_repo_root = os.environ.get("SKYRL_REPO_ROOT") + if env_repo_root: + candidates.append(pathlib.Path(env_repo_root)) + candidates.append(pathlib.Path(REMOTE_REPO)) + for start in (pathlib.Path(__file__).resolve(), pathlib.Path.cwd().resolve()): + base = start if start.is_dir() else start.parent + candidates.extend([base, *base.parents]) + for candidate in dict.fromkeys(candidates): + if _is_repo_root(candidate): + return candidate + raise RuntimeError( + "Could not locate the SkyRL repo root. Set SKYRL_REPO_ROOT or run from inside the repository." + ) + + +REPO_ROOT = _find_repo_root() + +hf_volume = modal.Volume.from_name("skyrl-hf-cache", create_if_missing=True) +ckpt_volume = modal.Volume.from_name("skyrl-tinker-moe-ckpts", create_if_missing=True) + +wandb_secret = modal.Secret.from_name(os.environ.get("SKYRL_WANDB_SECRET", "wandb-secret")) + + +VENV_DIR = "/root/venvs/skyrl" + + +image = ( + modal.Image.from_registry("novaskyai/skyrl-train-ray-2.56.0-py3.12-cu12.8") + .env( + { + "HF_HOME": HF_CACHE, + "HF_HUB_ENABLE_HF_TRANSFER": "1", + "SKYRL_REPO_ROOT": REMOTE_REPO, + "UV_LINK_MODE": "copy", + } + ) + + .add_local_file(f"{REPO_ROOT}/pyproject.toml", f"{REMOTE_REPO}/pyproject.toml", copy=True) + .add_local_file(f"{REPO_ROOT}/uv.lock", f"{REMOTE_REPO}/uv.lock", copy=True) + .add_local_file(f"{REPO_ROOT}/skyrl-gym/pyproject.toml", f"{REMOTE_REPO}/skyrl-gym/pyproject.toml", copy=True) + .workdir(REMOTE_REPO) + .run_commands( + f"UV_PROJECT_ENVIRONMENT={VENV_DIR} uv sync --frozen --extra tinker --extra megatron " + "--no-install-project --no-install-package skyrl-gym", + gpu="any", + ) + # --- Code layer: runtime mount, no rebuild on edits. + .add_local_dir( + str(REPO_ROOT), + REMOTE_REPO, + ignore=[".venv", "**/__pycache__", ".git", "*.egg-info", ".pytest_cache"], + ) +) + +app = modal.App("skyrl-tinker-moe") + + +# --------------------------------------------------------------------------- +# Server helpers +# --------------------------------------------------------------------------- + + +def _pump_logs(stream, prefix: str) -> None: + for line in stream: + print(f"[{prefix}] {line}", end="") + + +def _start_server(env: dict, base_model: str, backend_config: dict): + import subprocess + import threading + + # Sync the venv exactly ONCE per container (installs the editable + # skyrl/skyrl-gym from the runtime mount and any packages the baked layer + # is missing), then run everything with UV_NO_SYNC=1. Without this, the + # engine child process and every Ray worker actor re-run `uv run` against + # the SAME shared venv concurrently; their syncs flip-flop the editable + # install path (repo mount vs ray working_dir), and an actor importing + # transformer-engine during a sibling's uninstall->install window dies + # with "Found empty `transformer-engine` meta package installed". + sync_cmd = ["uv", "sync", "--frozen", "--extra", "tinker", "--extra", "megatron"] + print(">>> [modal] pre-sync command:", " ".join(shlex.quote(c) for c in sync_cmd)) + subprocess.run(sync_cmd, cwd=REMOTE_REPO, env=env, check=True) + env = {**env, "UV_NO_SYNC": "1"} + + # NOT --isolated: the server must run in the stable UV_PROJECT_ENVIRONMENT + # venv so Ray workers' uv re-invocations resolve to the same environment. + cmd = [ + "uv", + "run", + "--extra", + "tinker", + "--extra", + "megatron", + "-m", + "skyrl.tinker.api", + "--base-model", + base_model, + "--backend", + "megatron", + "--port", + str(PORT), + "--backend-config", + json.dumps(backend_config), + "--checkpoints-base", + CKPT_DIR, + ] + print(">>> [modal] server command:", " ".join(shlex.quote(c) for c in cmd)) + server = subprocess.Popen( + cmd, + cwd=REMOTE_REPO, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + threading.Thread(target=_pump_logs, args=(server.stdout, "server"), daemon=True).start() + return server + + +def _kill_server(server) -> None: + import signal + import subprocess + + if server.poll() is not None: + return + print(">>> [modal] Stopping Tinker API server") + try: + os.killpg(os.getpgid(server.pid), signal.SIGTERM) + except ProcessLookupError: + return + try: + server.wait(timeout=60) + except subprocess.TimeoutExpired: + os.killpg(os.getpgid(server.pid), signal.SIGKILL) + server.wait(timeout=30) + + +def _wait_ready(server, timeout_s: int = 45 * 60) -> None: + """Poll /api/v1/healthz (all Tinker routes live under /api/v1) until the + server answers 200. The API answers quickly; the heavy model init happens + lazily on the first create_model/sample request.""" + import time + import urllib.error + import urllib.request + + deadline = time.time() + timeout_s + while time.time() < deadline: + if server.poll() is not None: + raise RuntimeError(f"Tinker server exited early with code {server.returncode}") + try: + with urllib.request.urlopen(f"{BASE_URL}/api/v1/healthz", timeout=5) as resp: + if resp.status == 200: + print(">>> [modal] Tinker API server is ready") + return + except (urllib.error.URLError, ConnectionError, TimeoutError): + pass + time.sleep(5) + raise RuntimeError(f"Tinker server not ready within {timeout_s}s") + + +def _merged_backend_config(overrides_json: Optional[str]) -> dict: + cfg = dict(DEFAULT_BACKEND_CONFIG) + if overrides_json: + cfg.update(json.loads(overrides_json)) + return cfg + + +def _server_env() -> dict: + env = os.environ.copy() + env["HOME"] = "/root" + env["TINKER_API_KEY"] = "tml-dummy" + env["UV_PROJECT_ENVIRONMENT"] = VENV_DIR + env.setdefault("SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S", "1800") + + env.setdefault("VLLM_LOGGING_LEVEL", "INFO") + env.setdefault("RAY_DEDUP_LOGS", "0") + return env + + +def _run_client( + base_url: str, + base_model: str, + client_args: str, + metrics_path: Optional[str] = None, +) -> None: + import subprocess + + + env = os.environ.copy() + env["HOME"] = "/root" + env["TINKER_API_KEY"] = "tml-dummy" + env["PYTHONUNBUFFERED"] = "1" # stream step prints live instead of block-buffering + + # run in baked venv + env["UV_PROJECT_ENVIRONMENT"] = VENV_DIR + cmd = [ + "uv", + "run", + "--no-sync", + "--with", + "datasets", + "--with", + "wandb", + "python", + "examples/tinker/moe/moe_smoke_client.py", + "--base-url", + base_url, + "--base-model", + base_model, + ] + if metrics_path and "--metrics-path" not in client_args: + cmd += ["--metrics-path", metrics_path] + if client_args.strip(): + cmd.extend(shlex.split(client_args)) + print(">>> [modal] client command:", " ".join(shlex.quote(c) for c in cmd)) + subprocess.run(cmd, cwd=REMOTE_REPO, env=env, check=True) + + +# --------------------------------------------------------------------------- +# Modal functions +# --------------------------------------------------------------------------- + + +@app.function( + image=image, + gpu=GPU_SPEC, + timeout=24 * 60 * 60, + volumes={HF_CACHE: hf_volume, "/root/ckpts": ckpt_volume}, +) +def serve(base_model: str = DEFAULT_BASE_MODEL, backend_config: Optional[str] = None) -> None: + """Run the Tinker server behind a Modal tunnel until the job is stopped.""" + import time + + os.makedirs(CKPT_DIR, exist_ok=True) + server = _start_server(_server_env(), base_model, _merged_backend_config(backend_config)) + try: + # expose public server (todo: modal.web_server? For this it's overkill i think) + with modal.forward(PORT) as tunnel: + print(f">>> [modal] Tinker tunnel URL (may still be warming up): {tunnel.url}") + _wait_ready(server) + print(">>> [modal] ================================================") + print(f">>> [modal] Tinker server ready: {tunnel.url}") + print(">>> [modal] TINKER_API_KEY=tml-dummy") + print(f">>> [modal] base_model={base_model}") + print(">>> [modal] Stop with: modal app stop skyrl-tinker-moe") + print(">>> [modal] ================================================") + while server.poll() is None: + time.sleep(10) + raise RuntimeError(f"Tinker server exited with code {server.returncode}") + finally: + _kill_server(server) + ckpt_volume.commit() + + +@app.function(image=image, timeout=24 * 60 * 60, secrets=[wandb_secret]) # CPU-only: HTTP calls +def client(base_url: str, base_model: str = DEFAULT_BASE_MODEL, client_args: str = "") -> None: + """Submit the smoke client as its own Modal job against a running server.""" + _run_client(base_url, base_model, client_args) + + +@app.function( + image=image, + gpu=GPU_SPEC, + timeout=6 * 60 * 60, + volumes={HF_CACHE: hf_volume, "/root/ckpts": ckpt_volume}, + secrets=[wandb_secret], +) +def demo(base_model: str = DEFAULT_BASE_MODEL, backend_config: Optional[str] = None, client_args: str = "") -> None: + """Server + client in one container under two processes""" + os.makedirs(CKPT_DIR, exist_ok=True) + server = _start_server(_server_env(), base_model, _merged_backend_config(backend_config)) + try: + _wait_ready(server) + # Metrics land on the ckpt volume so they survive the container. + _run_client(BASE_URL, base_model, client_args, metrics_path=f"{CKPT_DIR}/moe_gsm8k_metrics.jsonl") + finally: + _kill_server(server) + ckpt_volume.commit() + print(">>> [modal] demo finished successfully") + + + +@app.local_entrypoint() +def main( + mode: str = "demo", + base_model: str = DEFAULT_BASE_MODEL, + backend_config: Optional[str] = None, + base_url: Optional[str] = None, + client_args: str = "", +) -> None: + if mode == "demo": + demo.remote(base_model=base_model, backend_config=backend_config, client_args=client_args) + elif mode == "serve": + serve.remote(base_model=base_model, backend_config=backend_config) + elif mode == "client": + if not base_url: + raise ValueError("--mode client requires --base-url (the tunnel URL printed by --mode serve)") + client.remote(base_url=base_url, base_model=base_model, client_args=client_args) + else: + raise ValueError("choose demo serve or client") diff --git a/examples/tinker/moe/moe_smoke_client.py b/examples/tinker/moe/moe_smoke_client.py new file mode 100644 index 0000000000..eb5067d4a4 --- /dev/null +++ b/examples/tinker/moe/moe_smoke_client.py @@ -0,0 +1,333 @@ +"""LoRA GRPO run on GSM8K against SkyRL's Tinker server= + +Usage (server must already be running -- see modal_moe_tinker.py): + + TINKER_API_KEY=tml-dummy uv run --extra tinker --with datasets \ + python examples/tinker/moe/moe_smoke_client.py \ + --base-url http://localhost:8000 \ + --base-model Qwen/Qwen3-30B-A3B-Base \ + --num-steps 50 --metrics-path ./moe_gsm8k_metrics.jsonl + +For router replay, pass --enable-router-replay to echo moe routing matrices back into +forward_backward datums -- passes through client +""" + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import os +import random +import re +import time + +import datasets +import numpy as np +import tinker +from tinker import types + +INSTRUCTION = 'Let\'s think step by step and output the final answer after "####".' +STRICT_ANSWER_RE = re.compile(r"#### (\-?[0-9\.,]+)") + +LOGPROB_DIFF_KEYS = ( + "policy/rollout_train_logprobs_abs_diff_mean:mean", + "policy/rollout_train_logprobs_abs_diff_std:mean", + "policy/rollout_train_logprobs_abs_diff_max:max", +) + + +def tensor_data_int(values: list[int]) -> types.TensorData: + # dtype is required by the tinker SDK (TensorData.__init__ raises without it) + return types.TensorData(data=[int(v) for v in values], dtype="int64") + + +def tensor_data_float(values: list[float]) -> types.TensorData: + return types.TensorData(data=[float(v) for v in values], dtype="float32") + + +# --------------------------------------------------------------------------- +# Router replay (R3) support +# --------------------------------------------------------------------------- + + +def install_routing_matrix_capture() -> None: + # Patch tinker SDK to pass through routing matrix from server + from tinker.lib import api_future_impl + from tinker.types.sample_response import SampleResponse + + original = api_future_impl.deserialize_json_response + + def patched(result_dict, model_cls): + result = original(result_dict, model_cls) + if model_cls is SampleResponse and isinstance(result_dict, dict): + for seq_obj, seq_dict in zip(result.sequences, result_dict.get("sequences", [])): + # SampledSequence is a frozen dataclass; attach out-of-band. + object.__setattr__(seq_obj, "routing_matrix", seq_dict.get("routing_matrix")) + return result + + api_future_impl.deserialize_json_response = patched + + +def decode_routing_matrix(b64: str) -> np.ndarray: + # Decode vllm .npy: payload -> (num_forrwarded_tokens, layers, topk) + return np.load(io.BytesIO(base64.b64decode(b64))) + + +# mirror examples/train/gsm8k/gsm8k_dataset.py and examples/tinker/ppo/ppo_client.py + + +def extract_ground_truth(solution: str) -> str: + return solution.split("#### ")[1].replace(",", "").strip() + + +def compute_reward(response_text: str, ground_truth: str) -> float: + match = STRICT_ANSWER_RE.search(response_text) + if match is None: + return 0.0 + return 1.0 if match.group(1).replace(",", "").strip() == ground_truth else 0.0 + + +def load_gsm8k(tokenizer, max_prompt_length: int, seed: int) -> list[dict]: + """Load GSM8K train split straight from HF, chat-templated prompts.""" + ds = datasets.load_dataset("openai/gsm8k", "main", split="train") + records = [] + for row in ds: + messages = [{"role": "user", "content": f"{row['question']} {INSTRUCTION}"}] + # return_dict=False: transformers 5.x otherwise returns a BatchEncoding dict + tokens = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_dict=False + ) + if len(tokens) > max_prompt_length: + continue + records.append( + {"prompt_tokens": list(tokens), "ground_truth": extract_ground_truth(row["answer"])} + ) + random.Random(seed).shuffle(records) + print(f"Loaded {len(records)} GSM8K train examples") + return records + + +# --------------------------------------------------------------------------- +# GRPO datum construction +# --------------------------------------------------------------------------- + + +def build_datum( + prompt_tokens: list[int], + response_tokens: list[int], + sampling_logprobs: list[float], + advantage: float, + routing_matrix: np.ndarray | None = None, +) -> types.Datum: + + prefix = prompt_tokens + response_tokens[:-1] + n = len(response_tokens) + loss_fn_inputs = { + "target_tokens": tensor_data_int(response_tokens), + "weights": tensor_data_float([1.0] * n), + "logprobs": tensor_data_float(sampling_logprobs), + "advantages": tensor_data_float([advantage] * n), + } + if routing_matrix is not None: + # vllm covers all forwarded tokens, should be that (prompt + generated - 1) == len(prefix) + if routing_matrix.shape[0] != len(prefix): + raise ValueError( + f"routing_matrix covers {routing_matrix.shape[0]} tokens, expected {len(prefix)} (model_input)" + ) + loss_fn_inputs["routing_matrix"] = types.TensorData.from_numpy(routing_matrix.astype(np.int64)) + return types.Datum( + model_input=types.ModelInput.from_ints(prefix), + loss_fn_inputs=loss_fn_inputs, + ) + + +def grpo_advantages(rewards: list[float], norm_by_std: bool) -> list[float]: + mean = sum(rewards) / len(rewards) + centered = [r - mean for r in rewards] + if not norm_by_std: + return centered + var = sum(c * c for c in centered) / len(centered) + std = var**0.5 + return [c / (std + 1e-6) for c in centered] + + +# --------------------------------------------------------------------------- +# Main loop +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://localhost:8000") + parser.add_argument("--base-model", default="Qwen/Qwen3-30B-A3B-Base") + parser.add_argument("--api-key", default=os.environ.get("TINKER_API_KEY", "tml-dummy")) + parser.add_argument("--lora-rank", type=int, default=32) + parser.add_argument("--num-steps", type=int, default=50) + parser.add_argument("--batch-size", type=int, default=8, help="prompts per step") + parser.add_argument("--group-size", type=int, default=8, help="samples per prompt (GRPO group)") + parser.add_argument("--max-prompt-length", type=int, default=1024) + parser.add_argument("--max-generate-length", type=int, default=1024) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--lr", type=float, default=4e-5) + parser.add_argument("--loss-fn", default="importance_sampling", choices=["importance_sampling", "ppo"]) + parser.add_argument("--no-norm-by-std", action="store_true", help="GRPO without std normalization") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--metrics-path", default="./moe_gsm8k_metrics.jsonl") + parser.add_argument( + "--enable-router-replay", + action="store_true", + help="Echo MoE routing matrices from sample responses into forward_backward datums. " + "Requires the server to run with enable_return_routed_experts + moe_enable_routing_replay.", + ) + args = parser.parse_args() + + if args.enable_router_replay: + install_routing_matrix_capture() + + print(f"Connecting to {args.base_url} (model={args.base_model}, lora_rank={args.lora_rank})") + service_client = tinker.ServiceClient(base_url=args.base_url, api_key=args.api_key) + training_client = service_client.create_lora_training_client( + base_model=args.base_model, + rank=args.lora_rank, + seed=args.seed, + train_mlp=True, + train_attn=True, + ) + tokenizer = training_client.get_tokenizer() + records = load_gsm8k(tokenizer, args.max_prompt_length, args.seed) + + # The backend applies only the lr; the other Adam fields satisfy the schema. + adam = types.AdamParams(learning_rate=args.lr, beta1=0.9, beta2=0.999, eps=1.0e-8) + loss_fn_config = ( + {"clip_low_threshold": 0.8, "clip_high_threshold": 1.2} if args.loss_fn == "ppo" else None + ) + + metrics_file = open(args.metrics_path, "a", buffering=1) + print(f"Writing per-step metrics to {args.metrics_path}") + + + wandb_run = None + if os.environ.get("WANDB_API_KEY"): + try: + import wandb + + wandb_run = wandb.init( + project=os.environ.get("WANDB_PROJECT", "skyrl-tinker-moe"), + name=os.environ.get("WANDB_RUN_NAME", f"moe-gsm8k-lora-r{args.lora_rank}"), + config=vars(args), + ) + print(f"wandb logging enabled: {wandb_run.url}") + except Exception as e: + print(f"wandb logging disabled ({e})") + else: + print("wandb logging disabled (WANDB_API_KEY not set)") + + for step in range(args.num_steps): + step_start = time.time() + batch = [records[(step * args.batch_size + i) % len(records)] for i in range(args.batch_size)] + + + sampling_client = training_client.save_weights_and_get_sampling_client() + futures = [] + for offset, rec in enumerate(batch): + params = types.SamplingParams( + max_tokens=args.max_generate_length, + temperature=args.temperature, + seed=args.seed + step * 10_000 + offset, + ) + futures.append( + sampling_client.sample( + prompt=types.ModelInput.from_ints(rec["prompt_tokens"]), + num_samples=args.group_size, + sampling_params=params, + ) + ) + + + data: list[types.Datum] = [] + all_rewards: list[float] = [] + skipped = 0 + missing_routing = 0 + for rec, future in zip(batch, futures): + result = future.result() + group = [] + for seq in result.sequences: + tokens = list(seq.tokens) + logprobs = list(seq.logprobs or []) + if not tokens or len(logprobs) != len(tokens): + skipped += 1 + continue + routing_matrix = None + if args.enable_router_replay: + raw = getattr(seq, "routing_matrix", None) + if raw is None: + + missing_routing += 1 + skipped += 1 + continue + routing_matrix = decode_routing_matrix(raw) + text = tokenizer.decode(tokens, skip_special_tokens=True) + group.append((tokens, logprobs, routing_matrix, compute_reward(text, rec["ground_truth"]))) + if len(group) < 2: + skipped += len(group) + continue + advantages = grpo_advantages([g[3] for g in group], norm_by_std=not args.no_norm_by_std) + for (tokens, logprobs, routing_matrix, reward), adv in zip(group, advantages): + all_rewards.append(reward) + data.append(build_datum(rec["prompt_tokens"], tokens, logprobs, adv, routing_matrix)) + + if args.enable_router_replay and missing_routing: + print( + f"step {step}: WARNING {missing_routing} sequences lacked routing_matrix -- " + "is the server running with generator.inference_engine.enable_return_routed_experts=true?" + ) + + if not data: + print(f"step {step}: no usable samples (skipped={skipped}); skipping update") + continue + + + fb = training_client.forward_backward(data, args.loss_fn, loss_fn_config).result() + training_client.optim_step(adam).result() + + + server_metrics = dict(getattr(fb, "metrics", {}) or {}) + row = { + "step": step, + "reward_mean": sum(all_rewards) / len(all_rewards), + "num_datums": len(data), + "skipped": skipped, + "loss": server_metrics.get("total_loss:sum"), + "step_time_s": round(time.time() - step_start, 1), + } + for key in LOGPROB_DIFF_KEYS: + short = key.split("/")[-1].split(":")[0] # e.g. rollout_train_logprobs_abs_diff_mean + row[short] = server_metrics.get(key) + metrics_file.write(json.dumps(row) + "\n") + if wandb_run is not None: + wandb_run.log({k: v for k, v in row.items() if v is not None}, step=step) + + diff_mean = row.get("rollout_train_logprobs_abs_diff_mean") + diff_max = row.get("rollout_train_logprobs_abs_diff_max") + print( + f"step {step:3d} | reward {row['reward_mean']:.3f} | datums {len(data):3d} | " + f"logprob_diff mean={diff_mean if diff_mean is not None else 'n/a'} " + f"max={diff_max if diff_max is not None else 'n/a'} | {row['step_time_s']}s" + ) + if step == 0 and diff_mean is None: + print( + "WARNING: server metrics did not include " + "policy/rollout_train_logprobs_abs_diff_* — check that datums carry " + "sampling logprobs and the server is the skyrl_train backend." + ) + + metrics_file.close() + if wandb_run is not None: + wandb_run.finish() + print("run complete") + + +if __name__ == "__main__": + main() diff --git a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py index b2d1e4e812..23561ac232 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -722,6 +722,10 @@ async def sample( "tokens": choice["token_ids"], "logprobs": seq_logprobs, "stop_reason": choice.get("finish_reason"), + # MoE routing indices (base64 .npy), present when the server + # runs with enable_return_routed_experts. Passed through + # verbatim; the client decodes. + "routing_matrix": choice.get("routed_experts"), } ) diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index 885ec362b8..ced13cefc8 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -37,6 +37,7 @@ from skyrl.env_vars import SKYRL_RAY_PG_TIMEOUT_IN_S from skyrl.tinker import types from skyrl.train.config import SkyRLTrainConfig, get_config_as_yaml_str +from skyrl.train.dataset.preprocess import pad_rollout_expert_indices from skyrl.train.utils.utils import ( ResolvedPlacementGroup, get_ray_pg_ready_with_timeout, @@ -194,6 +195,7 @@ def _split_model_pass_batch_by_model_id( "all_advantages", "all_values", "all_returns", + "all_routing_matrices", "all_model_ids", "all_loss_fns", "all_loss_fn_configs", @@ -668,6 +670,23 @@ def _to_training_batch(self, prepared_batch: types.PreparedModelPassBatch, role: batch_dict["rollout_logprobs"] = action_log_probs_tensor if has_advantages: batch_dict["advantages"] = torch.tensor(advantages_list, dtype=torch.float32) + + # MoE router replay: routing indices cover exactly the forwarded tokens + # (model_input = pre-shifted prompt + response[:-1]), so the final + # sequence position is left zero by the padding helper. Alignment is + # sequence-start against attention_mask, matching the classic path. + routing_matrices = prepared_batch.all_routing_matrices + if any(rm is not None for rm in routing_matrices): + if not all(rm is not None for rm in routing_matrices): + raise ValueError( + "forward_backward batch mixes datums with and without routing_matrix; " + "router replay is all-or-nothing per batch" + ) + batch_dict["rollout_expert_indices"] = pad_rollout_expert_indices( + [rm.to_numpy() for rm in routing_matrices], + row_lengths=[len(seq) for seq in full_sequences], + max_total=max_seq_len, + ) if role == "critic": if has_values != has_returns: raise ValueError("Critic batches must provide both values and returns, or neither") @@ -1049,8 +1068,16 @@ def _sample_with_remote_client( # save_sampler_checkpoint via load_lora_adapter). Single-tenant / # FFT path falls back to resolve_policy_model_name(cfg). fallback_model_name = resolve_policy_model_name(self._cfg) + # Adapters exist on vLLM only when the sync takes the in-place LoRA path. + # With megatron merge_lora (the default), LoRA weights are merged and + # broadcast under the base model name -- routing by adapter name would + # 404. Mirror the gate the dispatch uses in broadcast_to_inference_engines. + serves_lora_adapters = self._base_lora_signature is not None and not ( + self._cfg.trainer.strategy == "megatron" + and self._cfg.trainer.policy.megatron_config.lora_config.merge_lora + ) per_request_models = [ - mid if (self._base_lora_signature is not None and mid in self._model_ids_to_role) else fallback_model_name + mid if (serves_lora_adapters and mid in self._model_ids_to_role) else fallback_model_name for mid in prepared_batch.all_model_ids ] @@ -1090,9 +1117,9 @@ def _aggregate_sample_results( logger.info(f"Aggregating sample results for {len(sample_outputs)} samples") def _extract_sequences(output): - """Yield (tokens, logprobs, stop_reason) from a single sample output.""" + """Yield (tokens, logprobs, stop_reason, routing_matrix) from a single sample output.""" for seq in output["sequences"]: - yield seq["tokens"], seq.get("logprobs"), seq.get("stop_reason") + yield seq["tokens"], seq.get("logprobs"), seq.get("stop_reason"), seq.get("routing_matrix") results = {} for request_id, model_id, start_idx, end_idx, prompt_logprobs_requested in prepared_batch.request_batch_slices: @@ -1115,7 +1142,7 @@ def _extract_sequences(output): logger.error(error_msg) break - for tokens, logprobs_raw, stop_reason_raw in _extract_sequences(output): + for tokens, logprobs_raw, stop_reason_raw, routing_matrix in _extract_sequences(output): # Map vLLM stop reason to Tinker format stop_reason = "stop" if stop_reason_raw in ("stop", "stop_token") else "length" logprobs = logprobs_raw or [] @@ -1130,6 +1157,7 @@ def _extract_sequences(output): tokens=tokens, logprobs=logprobs, stop_reason=stop_reason, + routing_matrix=routing_matrix, ) ) diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index c135f05f03..e6a3ad9433 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -415,9 +415,14 @@ def to_types(self) -> types.ModelInput: class TensorData(BaseModel): data: list[int] | list[float] + # Optional dtype/shape, matching the tinker SDK's TensorData wire format + # (the SDK flattens multi-dimensional arrays into `data` and sends the + # shape alongside). + dtype: Literal["int64", "float32"] | None = None + shape: list[int] | None = None def to_types(self) -> types.TensorData: - return types.TensorData(data=self.data) + return types.TensorData(data=self.data, dtype=self.dtype, shape=self.shape) class Datum(BaseModel): @@ -440,6 +445,7 @@ def to_types(self) -> types.Datum: logprobs=inp["logprobs"].to_types() if "logprobs" in inp else types.TensorData(data=[]), values=inp["values"].to_types() if "values" in inp else types.TensorData(data=[]), returns=inp["returns"].to_types() if "returns" in inp else types.TensorData(data=[]), + routing_matrix=inp["routing_matrix"].to_types() if "routing_matrix" in inp else None, ), model_input=self.model_input.to_types(), ) diff --git a/skyrl/tinker/engine.py b/skyrl/tinker/engine.py index bab1998d1a..e29b396831 100644 --- a/skyrl/tinker/engine.py +++ b/skyrl/tinker/engine.py @@ -126,6 +126,7 @@ def prepare_model_pass_batch( all_advantages = [] all_values = [] all_returns = [] + all_routing_matrices = [] all_loss_fns = [] all_loss_fn_configs = [] request_batch_slices = [] @@ -145,6 +146,7 @@ def prepare_model_pass_batch( all_advantages.append(loss_fn_inputs.advantages.data) all_values.append(loss_fn_inputs.values.data) all_returns.append(loss_fn_inputs.returns.data) + all_routing_matrices.append(loss_fn_inputs.routing_matrix) all_model_ids.append(model_id) all_loss_fns.append(request_data.loss_fn) all_loss_fn_configs.append(request_data.loss_fn_config) @@ -159,6 +161,7 @@ def prepare_model_pass_batch( all_advantages=all_advantages, all_values=all_values, all_returns=all_returns, + all_routing_matrices=all_routing_matrices, all_model_ids=all_model_ids, all_loss_fns=all_loss_fns, all_loss_fn_configs=all_loss_fn_configs, diff --git a/skyrl/tinker/types.py b/skyrl/tinker/types.py index aa22c657fd..640174ea20 100644 --- a/skyrl/tinker/types.py +++ b/skyrl/tinker/types.py @@ -9,6 +9,7 @@ from typing import Annotated, Any, Literal, TypedDict from urllib.parse import urlparse +import numpy as np from pydantic import Base64Bytes, BaseModel, Discriminator, Field @@ -142,6 +143,14 @@ class RenderedModelInput(BaseModel): class TensorData(BaseModel): data: list[int] | list[float] + dtype: Literal["int64", "float32"] | None = None + shape: list[int] | None = None + """Optional dtype/shape, matching the tinker SDK's TensorData wire format. + ``data`` is always the flattened values; ``shape`` restores dimensionality.""" + + def to_numpy(self) -> "np.ndarray": + arr = np.asarray(self.data) + return arr.reshape(self.shape) if self.shape is not None else arr class LossFnInputs(BaseModel): @@ -151,6 +160,9 @@ class LossFnInputs(BaseModel): logprobs: TensorData values: TensorData = Field(default_factory=lambda: TensorData(data=[])) returns: TensorData = Field(default_factory=lambda: TensorData(data=[])) + routing_matrix: TensorData | None = None + """MoE rollout routing indices for replay, shape (num_input_tokens, num_moe_layers, topk). + Covers every forwarded token of ``model_input`` (the pre-shifted prompt + response[:-1]).""" class Datum(BaseModel): @@ -264,6 +276,11 @@ class GeneratedSequence(BaseModel): stop_reason: Literal["length", "stop"] tokens: list[int] logprobs: list[float] + routing_matrix: str | None = None + """MoE expert indices for every forwarded token (prompt + generated - 1), as + base64-encoded ``.npy`` bytes passed through from vLLM. Present only when the + server runs with ``enable_return_routed_experts``. Decode with + ``np.load(io.BytesIO(base64.b64decode(s)))``.""" class SampleOutput(BaseModel): @@ -295,6 +312,7 @@ class PreparedModelPassBatch(BaseModel): all_advantages: list[list[float]] all_values: list[list[float]] all_returns: list[list[float]] + all_routing_matrices: list[TensorData | None] = Field(default_factory=list) # Per-example scalars all_model_ids: list[str] diff --git a/skyrl/train/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py index 2b34521d1c..a1af18981b 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -29,6 +29,43 @@ def _verify_inputs( ) +def pad_rollout_expert_indices(per_sample_indices, row_lengths, max_total): + """Pad per-sample MoE routing indices into a ``[B, max_total, layers, topk]`` tensor. + + Each sample's rows are written starting at that row's left-pad offset + (sequence-start alignment against ``attention_mask`` — unlike the + response-aligned loss tensors), with uncovered positions left zero. + Downcast to uint8/int16 when the expert count allows. + + Args: + per_sample_indices: per-sample ``[seq][layers][topk]`` nested lists or + arrays; falsy/None entries contribute all-zero rows. + row_lengths: real (unpadded) sequence length per row. + max_total: padded sequence length of the batch. + + Returns: + Integer tensor, or None if no sample carries indices. + """ + first_non_empty = next((x for x in per_sample_indices if x is not None and len(x) > 0), None) + if first_non_empty is None: + return None + num_layers = len(first_non_empty[0]) + topk = len(first_non_empty[0][0]) if num_layers > 0 else 0 + padded = torch.zeros(len(per_sample_indices), max_total, num_layers, topk, dtype=torch.int32) + for i, sample_indices in enumerate(per_sample_indices): + if sample_indices is not None and len(sample_indices) > 0: + left_pad = max_total - row_lengths[i] + n = min(len(sample_indices), max_total - left_pad) + padded[i, left_pad : left_pad + n] = torch.as_tensor(sample_indices[:n], dtype=torch.int32) + + # downcast to uint8 if possible, otherwise int16 to save memory + if padded.max().item() < 2**8: + return padded.to(torch.uint8) + if padded.max().item() < 2**15: + return padded.to(torch.int16) + return padded + + def convert_prompts_responses_to_batch_tensors( tokenizer: AutoTokenizer, prompts: List[List[int]], @@ -161,23 +198,10 @@ def convert_prompts_responses_to_batch_tensors( rollout_expert_indices_tensor = None if rollout_expert_indices: - first_non_empty = next((x for x in rollout_expert_indices if x), None) - if first_non_empty: - num_layers = len(first_non_empty[0]) - topk = len(first_non_empty[0][0]) if num_layers > 0 else 0 - padded = torch.zeros(len(rollout_expert_indices), max_total, num_layers, topk, dtype=torch.int32) - for i, sample_indices in enumerate(rollout_expert_indices): - if sample_indices: - left_pad = max_total - (prompt_token_lens[i] + response_token_lens[i]) - n = min(len(sample_indices), max_total - left_pad) - padded[i, left_pad : left_pad + n] = torch.tensor(sample_indices[:n], dtype=torch.int32) - rollout_expert_indices_tensor = padded - - # downcast to uint8 if possible, otherwise int16 to save memory - if rollout_expert_indices_tensor.max().item() < 2**8: - rollout_expert_indices_tensor = rollout_expert_indices_tensor.to(torch.uint8) - elif rollout_expert_indices_tensor.max().item() < 2**15: - rollout_expert_indices_tensor = rollout_expert_indices_tensor.to(torch.int16) + row_lengths = [p + r for p, r in zip(prompt_token_lens, response_token_lens)] + rollout_expert_indices_tensor = pad_rollout_expert_indices( + rollout_expert_indices, row_lengths=row_lengths, max_total=max_total + ) return ( sequences, diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 777704a3a3..930cd5316c 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -785,6 +785,12 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: "UV_LINK_MODE", "UV_PYTHON", "UV_OFFLINE", + # UV_NO_SYNC=1 keeps the workers' `uv run` re-invocations from re-syncing a + # shared venv concurrently. Without it, each worker's sync flip-flops the + # editable install (driver path vs ray working_dir path), and an actor that + # imports transformer-engine during a sibling's uninstall->install window + # dies with "Found empty `transformer-engine` meta package installed". + "UV_NO_SYNC", "PYTORCH_CUDA_ALLOC_CONF", # Debug/trace knobs — forwarded so they reach the worker actors, not just the driver. "CUDA_LAUNCH_BLOCKING",