diff --git a/Dockerfile.cuda.runtime b/Dockerfile.cuda.runtime new file mode 100644 index 0000000000..38499fd90a --- /dev/null +++ b/Dockerfile.cuda.runtime @@ -0,0 +1,189 @@ +# Multi-stage Dockerfile for prime-rl with NVRTC support on GB200 (sm_100a). +# +# WHY THIS EXISTS (separate from Dockerfile.cuda): +# The original Dockerfile.cuda uses `python:3.12-slim` as the runtime base. +# That image has no CUDA toolkit, so tilelang's JIT path (which compiles the +# sparse-MLA kernels at runtime via NVRTC) fails with: +# +# atomic.h(7): catastrophic error: cannot open source file "cuda/atomic" +# +# `cuda/atomic` is a libcudacxx (CCCL) header. It is shipped only by the +# CUDA dev/devel toolkit, not by the `nvidia-cuda-*` pip wheels. Without it +# NVRTC cannot compile any kernel that pulls in tilelang's atomic.h. +# +# This Dockerfile uses NVIDIA's cuda-dl-base (devel) image for both stages, +# so the runtime image carries the libcudacxx / CCCL headers tilelang needs +# and the wheels in /app/.venv keep their CUDA 12.x ABI. +# +# BASE IMAGE: +# nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04 +# - matches dynamo/container/context.yaml `prime-rl.cuda12.9` entry +# - ships Python 3.12 by default (Ubuntu 24.04) +# - includes libcudacxx, CCCL, cuDNN, nvcc, and full CUDA dev headers +# - forward-compatible with the CUDA 12.8 wheels pinned in uv.lock +# +# USAGE: +# docker buildx build --platform linux/arm64 \ +# --build-arg TARGETARCH=arm64 \ +# -f Dockerfile.cuda.runtime -t . + +ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base +ARG BASE_IMAGE_TAG=25.06-cuda12.9-devel-ubuntu24.04 + +############################ +##### Build stage ########## +############################ +FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS builder +LABEL maintainer="prime intellect" +LABEL repository="prime-rl" + +# Set en_US.UTF-8 locale by default +RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment + +# CUDA_HOME / PATH from base image are already correct (/usr/local/cuda), but +# pin them explicitly so downstream tooling (tilelang, flash-attn) sees them. +ENV CUDA_HOME=/usr/local/cuda +ENV PATH=$PATH:/usr/local/cuda/bin + +# Install build tooling. +ARG DEBIAN_FRONTEND=noninteractive +ENV DEBIAN_FRONTEND=noninteractive +ENV TZ=Etc/UTC +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + curl \ + sudo \ + git \ + ninja-build \ + && apt-get clean autoclean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* + +# Install uv. +ADD https://astral.sh/uv/install.sh /uv-installer.sh +RUN INSTALLER_NO_MODIFY_PATH=1 UV_INSTALL_DIR="/usr/local/bin" sh /uv-installer.sh && rm /uv-installer.sh +ENV PATH="/usr/local/bin:$PATH" +ENV UV_PYTHON_INSTALL_DIR="/usr/local/share/uv/python" +ENV UV_CACHE_DIR="/usr/local/share/uv/cache" + +# Install Python dependencies (gradual copies help with caching). +WORKDIR /app + +ENV UV_COMPILE_BYTECODE=1 UV_LINK_MODE=copy + +COPY pyproject.toml /app/pyproject.toml +COPY uv.lock /app/uv.lock +COPY README.md /app/README.md +COPY src/ /app/src/ +COPY packages/ /app/packages/ +COPY deps/ /app/deps/ +COPY configs /app/configs +COPY examples /app/examples +COPY benchmarks/scripts /app/benchmarks/scripts + +RUN --mount=type=cache,target=/app/.cache/uv \ + uv sync --extra flash-attn --extra flash-attn-3 --extra flash-attn-cute --extra envs --extra gpt-oss --group mamba-ssm --locked --no-dev + +# arm64: build flash-attn + DeepGEMM from source. +ARG TARGETARCH +COPY scripts/docker-arm64-post-install.sh /app/scripts/docker-arm64-post-install.sh +COPY scripts/install_deep_gemm.sh /app/scripts/install_deep_gemm.sh +RUN if [ "$TARGETARCH" = "arm64" ]; then /app/scripts/docker-arm64-post-install.sh; fi + +# vLLM PR #39366 (two-phase DP pause) is native in vLLM 0.22 — no patch needed +# (the rl-sdk-4 merge bumped vLLM 0.20.2 -> 0.22 and dropped this patch). + +############################ +##### Runtime stage ######## +############################ +# Same image so libcudacxx, CCCL headers, and the system Python 3.12 are all +# present at runtime — the original Dockerfile.cuda switched to python:3.12-slim +# here and lost the CUDA dev headers, which broke tilelang's NVRTC backend. +FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} + +ENV CUDA_HOME=/usr/local/cuda +ENV PATH=/usr/local/cuda/bin:$PATH + +ARG DEBIAN_FRONTEND=noninteractive +ENV DEBIAN_FRONTEND=noninteractive +ENV TZ=Etc/UTC + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + wget \ + clang \ + tmux \ + iperf \ + openssh-server \ + git \ + git-lfs \ + gpg \ + sudo \ + iputils-ping \ + net-tools \ + curl \ + vim \ + libibverbs1 \ + ibverbs-providers \ + python3.12 \ + python3.12-venv \ + && apt-get clean autoclean \ + && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* + +# Ensure `python` / `python3` point at 3.12 (ubuntu24.04 ships 3.12 already, +# but the symlinks aren't created by default). +RUN ln -sf /usr/bin/python3.12 /usr/local/bin/python \ + && ln -sf /usr/bin/python3.12 /usr/local/bin/python3 \ + && ln -sf /usr/bin/python3.12 /usr/local/bin/python3.12 + +ARG USER_ID=1000 +ARG GROUP_ID=1000 +# Ubuntu 24.04 ships a default `ubuntu` user at uid 1000; remove it so the +# explicit appuser keeps uid/gid 1000 (matches Dockerfile.cuda + k8s manifests). +RUN userdel -r ubuntu 2>/dev/null || true \ + && groupadd --gid $GROUP_ID appuser \ + && useradd --uid $USER_ID --gid appuser --create-home --shell /bin/bash appuser \ + && usermod -aG sudo appuser \ + && echo "appuser ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers + +# Install uv for development use. +ADD https://astral.sh/uv/install.sh /uv-installer.sh +RUN INSTALLER_NO_MODIFY_PATH=1 UV_INSTALL_DIR="/usr/local/bin" sh /uv-installer.sh && rm /uv-installer.sh + +USER appuser +ENV PATH="/usr/local/bin:$PATH" +WORKDIR /app +# Copy the application + venv from the builder. +COPY --from=builder --chown=appuser:appuser /app /app + +# Copy and set up entrypoint script. +COPY --chown=appuser:appuser scripts/docker-entrypoint.sh /app/docker-entrypoint.sh +RUN chmod +x /app/docker-entrypoint.sh + +# Repoint venv Python symlinks at the runtime-stage interpreter (the builder +# used a uv-managed Python that does not exist here). +RUN rm /app/.venv/bin/python && ln -s /usr/bin/python3.12 /app/.venv/bin/python +RUN rm /app/.venv/bin/python3 && ln -s /usr/bin/python3.12 /app/.venv/bin/python3 +RUN rm /app/.venv/bin/python3.12 && ln -s /usr/bin/python3.12 /app/.venv/bin/python3.12 + +# python3.12-dev: Python.h headers for vLLM's Triton CudaUtils JIT-compile at runtime. +# Required by vLLM serving (dynamo.vllm / inference); absent from the cuda-dl-base +# runtime, which crashed inference with `fatal error: Python.h: No such file or directory`. +USER root +RUN apt-get update && apt-get install -y --no-install-recommends python3.12-dev \ + && apt-get clean autoclean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* +USER appuser + +# Place executables in the environment at the front of the path. +ENV PATH="/app/.venv/bin:$PATH" + +# HuggingFace Hub timeouts (defaults are 10s which causes issues on slow networks). +ENV HF_HUB_ETAG_TIMEOUT=500 +ENV HF_HUB_DOWNLOAD_TIMEOUT=300 + +# Enable FP8 grouped-GEMM kernels in vLLM MoE layers (requires DeepGEMM, built +# during the arm64 post-install step above). +ENV VLLM_USE_DEEP_GEMM=1 +ENV VLLM_MOE_USE_DEEP_GEMM=1 + +# Use entrypoint for setup (ulimit, etc) but default to sleep infinity for K8s. +ENTRYPOINT ["/app/docker-entrypoint.sh"] +CMD ["sleep", "infinity"] diff --git a/k8s/dynamo-deploy/dynamo-dgd.yaml b/k8s/dynamo-deploy/dynamo-dgd.yaml new file mode 100644 index 0000000000..9791d5d8d9 --- /dev/null +++ b/k8s/dynamo-deploy/dynamo-dgd.yaml @@ -0,0 +1,90 @@ +# Example DynamoGraphDeployment for serving inference to prime-rl. +# kubectl apply -f dynamo-dgd.yaml -n +# +# Replace ``, `/dynamo:`, +# ``, and the model name below for your environment. +# Requires DYN_ENABLE_RL=true on the Dynamo runtime so /v1/rl/* endpoints are +# served natively. +apiVersion: nvidia.com/v1alpha1 +kind: DynamoGraphDeployment +metadata: + name: prime-rl-dynamo + namespace: +spec: + backendFramework: vllm + pvcs: + - create: false + name: model-cache + services: + Frontend: + componentType: frontend + + extraPodSpec: + imagePullSecrets: + - name: + mainContainer: + image: /dynamo: + startupProbe: + failureThreshold: 360 + httpGet: + path: /health + port: 8000 + periodSeconds: 10 + timeoutSeconds: 5 + replicas: 1 + volumeMounts: + - mountPoint: /model-cache + name: model-cache + VllmWorker: + componentType: worker + + envFromSecret: hf-token-secret + extraPodSpec: + imagePullSecrets: + - name: + mainContainer: + command: + - python3 + - -m + - dynamo.vllm + args: + - --model + - Qwen/Qwen3-4B-Instruct-2507 + - --served-model-name + - Qwen/Qwen3-4B-Instruct-2507 + - --tensor-parallel-size + - "1" + - --max-model-len + - "2048" + - --max-num-seqs + - "64" + - --gpu-memory-utilization + - "0.90" + - --enforce-eager + env: + - name: HF_HOME + value: /model-cache/huggingface + image: /dynamo: + startupProbe: + failureThreshold: 360 + httpGet: + path: /health + port: 9090 + periodSeconds: 10 + timeoutSeconds: 10 + workingDir: /workspace/examples/backends/vllm + tolerations: + - effect: NoSchedule + key: nvidia.com/gpu + operator: Exists + replicas: 1 + resources: + limits: + gpu: "1" + requests: + gpu: "1" + sharedMemory: + size: 16Gi + volumeMounts: + - mountPoint: /model-cache + name: model-cache diff --git a/k8s/dynamo-deploy/prime-rl-configs.yaml b/k8s/dynamo-deploy/prime-rl-configs.yaml new file mode 100644 index 0000000000..11a8cc8c53 --- /dev/null +++ b/k8s/dynamo-deploy/prime-rl-configs.yaml @@ -0,0 +1,60 @@ +# Example ConfigMap mounted at /configs in the orchestrator and trainer pods. +# kubectl apply -f prime-rl-configs.yaml -n +# +# Replace `` below with your namespace, and adjust the +# Dynamo frontend service hostname if it differs in your cluster. +apiVersion: v1 +kind: ConfigMap +metadata: + name: prime-rl-configs + namespace: +data: + orch.toml: | + max_steps = 20 + seq_len = 2048 + batch_size = 64 + rollouts_per_example = 4 + use_token_client = false + + [wandb] + project = "prime-rl-dynamo-k8s" + name = "dynamo-smoke-qwen3-4b" + + [model] + name = "Qwen/Qwen3-4B-Instruct-2507" + + [sampling] + max_tokens = 256 + + [[env]] + id = "math-env" + name = "hendrycks-math" + args = { dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", math_verify_max_workers = 32, math_verify_timeout = 60 } + + [buffer] + easy_threshold = 1.0 + hard_threshold = 0.0 + + [client] + # Use the Dynamo admin/rollout transport (/engine/* + nvext), not the + # default vLLM frontend — required for the Dynamo /v1/rl/* admin routes below. + backend = "dynamo" + base_url = ["http://prime-rl-dynamo-frontend..svc.cluster.local:8000/v1"] + # Admin endpoints (/v1/rl/*) are served natively by the Dynamo Rust frontend + # (DYN_ENABLE_RL=true). No separate admin-stub service needed. + admin_base_url = ["http://prime-rl-dynamo-frontend..svc.cluster.local:8000/v1/rl"] + skip_model_check = true + + train.toml: | + max_steps = 20 + + [model] + name = "Qwen/Qwen3-4B-Instruct-2507" + seq_len = 2048 + + [wandb] + project = "prime-rl-dynamo-k8s" + name = "dynamo-smoke-qwen3-4b-trainer" + + [optim] + lr = 3e-6 diff --git a/k8s/dynamo-deploy/prime-rl-values.yaml b/k8s/dynamo-deploy/prime-rl-values.yaml new file mode 100644 index 0000000000..0b97ef87e2 --- /dev/null +++ b/k8s/dynamo-deploy/prime-rl-values.yaml @@ -0,0 +1,114 @@ +# Example Helm values for prime-rl when Dynamo serves inference. +# +# helm install prime-rl k8s/prime-rl -f k8s/dynamo-deploy/prime-rl-values.yaml -n +# +# Dynamo serves inference, so prime-rl's own inference component is disabled. +# Orchestrator and trainer point at the Dynamo frontend (deploy with dynamo-dgd.yaml). +# +# Replace the following placeholders for your environment: +# - namespace +# - image.repository / image.tag (your prime-rl image) +# - storage.storageClassName (your cluster's RWX storage class) +# - orchestrator/trainer.imagePullSecrets[*].name +# - secrets.name (HF token / W&B key) + +namespace: + +image: + repository: /prime-rl + tag: "latest" + pullPolicy: IfNotPresent + +storage: + enabled: true + storageClassName: + accessModes: + - ReadWriteMany + size: 100Gi + mountPath: /data + +config: + example: "prime-rl-dynamo" + secrets: + enabled: true + name: hf-token-secret + +# Orchestrator: CPU only, talks to Dynamo for inference +orchestrator: + enabled: true + replicas: 1 + autoStart: true + command: >- + BENCH_API_KEY=EMPTY + uv run orchestrator + @ /configs/orch.toml + --output-dir /data/outputs/run_default + --max-concurrent 16 + ; echo "orchestrator exited: $?" ; sleep infinity + resources: + requests: + memory: "8Gi" + cpu: "4" + env: + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: HF_TOKEN + optional: true + - name: WANDB_API_KEY + valueFrom: + secretKeyRef: + name: wandb-secret + key: WANDB_API_KEY + configMap: prime-rl-configs + imagePullSecrets: + - name: + nodeSelector: {} + +# Inference: DISABLED (Dynamo DGD replaces this) +inference: + enabled: false + +# Trainer: 1 GPU is sufficient for small models (e.g. Qwen3-4B GRPO) +trainer: + enabled: true + replicas: 1 + autoStart: true + command: >- + uv run torchrun + --nproc-per-node=1 + --rdzv-endpoint=localhost:29510 + --rdzv-id=smoke_$(date +%s) + -m prime_rl.trainer.rl.train + @ /configs/train.toml + --output-dir /data/outputs + ; echo "trainer exited: $?" ; sleep infinity + pytorchCudaAllocConf: "expandable_segments:True" + gpu: + enabled: true + count: 1 + resources: + requests: + memory: "32Gi" + cpu: "8" + env: + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: HF_TOKEN + optional: true + - name: WANDB_API_KEY + valueFrom: + secretKeyRef: + name: wandb-secret + key: WANDB_API_KEY + configMap: prime-rl-configs + runtimeClassName: nvidia + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + imagePullSecrets: + - name: diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index 1ed16b9f5e..f3c01a9ab7 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -26,6 +26,14 @@ spec: nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.orchestrator.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.orchestrator.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: prime-rl-orchestrator image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" @@ -33,7 +41,7 @@ spec: {{- if .Values.orchestrator.autoStart }} command: ["/bin/bash", "-c"] args: - - {{ .Values.orchestrator.command }} + - {{ .Values.orchestrator.command | quote }} {{- else }} command: ["sleep", "infinity"] {{- end }} @@ -91,16 +99,29 @@ spec: {{- end }} resources: {{- toYaml .Values.orchestrator.resources | nindent 10 }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.orchestrator.configMap }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.orchestrator.configMap }} + - name: rl-configs + mountPath: /configs + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.orchestrator.configMap }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ .Values.storage.existingClaim | default (printf "%s-shared-data" .Release.Name) }} + {{- end }} + {{- if .Values.orchestrator.configMap }} + - name: rl-configs + configMap: + name: {{ .Values.orchestrator.configMap }} + {{- end }} {{- end }} {{- end }} --- @@ -132,6 +153,19 @@ spec: {{- if .Values.inference.runtimeClassName }} runtimeClassName: {{ .Values.inference.runtimeClassName }} {{- end }} + {{- with .Values.inference.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.inference.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- if .Values.inference.resourceClaim.enabled }} + resourceClaims: + - name: {{ .Values.inference.resourceClaim.name }} + resourceClaimTemplateName: {{ .Values.inference.resourceClaim.templateName }} + {{- end }} containers: - name: prime-rl-inference image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" @@ -139,7 +173,7 @@ spec: {{- if .Values.inference.autoStart }} command: ["/bin/bash", "-c"] args: - - {{ .Values.inference.command }} + - {{ .Values.inference.command | quote }} {{- else }} command: ["sleep", "infinity"] {{- end }} @@ -184,19 +218,23 @@ spec: {{- end }} resources: requests: - {{- if .Values.inference.gpu.enabled }} + {{- if and .Values.inference.gpu.enabled (not .Values.inference.resourceClaim.enabled) }} nvidia.com/gpu: {{ .Values.inference.gpu.count }} {{- end }} memory: {{ .Values.inference.resources.requests.memory }} cpu: {{ .Values.inference.resources.requests.cpu }} limits: - {{- if .Values.inference.gpu.enabled }} + {{- if and .Values.inference.gpu.enabled (not .Values.inference.resourceClaim.enabled) }} nvidia.com/gpu: {{ .Values.inference.gpu.count }} {{- end }} {{- if .Values.inference.resources.limits }} memory: {{ .Values.inference.resources.limits.memory }} cpu: {{ .Values.inference.resources.limits.cpu }} {{- end }} + {{- if .Values.inference.resourceClaim.enabled }} + claims: + - name: {{ .Values.inference.resourceClaim.name }} + {{- end }} {{- if and .Values.inference.probes .Values.inference.probes.enabled }} startupProbe: httpGet: @@ -220,16 +258,29 @@ spec: failureThreshold: {{ .Values.inference.probes.readiness.failureThreshold }} timeoutSeconds: {{ .Values.inference.probes.readiness.timeoutSeconds }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.inference.configMap }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.inference.configMap }} + - name: rl-configs + mountPath: /configs + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.inference.configMap }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ .Values.storage.existingClaim | default (printf "%s-shared-data" .Release.Name) }} + {{- end }} + {{- if .Values.inference.configMap }} + - name: rl-configs + configMap: + name: {{ .Values.inference.configMap }} + {{- end }} {{- end }} {{- end }} --- @@ -261,6 +312,19 @@ spec: {{- if .Values.trainer.runtimeClassName }} runtimeClassName: {{ .Values.trainer.runtimeClassName }} {{- end }} + {{- with .Values.trainer.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.trainer.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- if .Values.trainer.resourceClaim.enabled }} + resourceClaims: + - name: {{ .Values.trainer.resourceClaim.name }} + resourceClaimTemplateName: {{ .Values.trainer.resourceClaim.templateName }} + {{- end }} containers: - name: prime-rl-trainer image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" @@ -268,7 +332,7 @@ spec: {{- if .Values.trainer.autoStart }} command: ["/bin/bash", "-c"] args: - - {{ .Values.trainer.command }} + - {{ .Values.trainer.command | quote }} {{- else }} command: ["sleep", "infinity"] {{- end }} @@ -330,19 +394,23 @@ spec: {{- end }} resources: requests: - {{- if .Values.trainer.gpu.enabled }} + {{- if and .Values.trainer.gpu.enabled (not .Values.trainer.resourceClaim.enabled) }} nvidia.com/gpu: {{ .Values.trainer.gpu.count }} {{- end }} memory: {{ .Values.trainer.resources.requests.memory }} cpu: {{ .Values.trainer.resources.requests.cpu }} limits: - {{- if .Values.trainer.gpu.enabled }} + {{- if and .Values.trainer.gpu.enabled (not .Values.trainer.resourceClaim.enabled) }} nvidia.com/gpu: {{ .Values.trainer.gpu.count }} {{- end }} {{- if .Values.trainer.resources.limits }} memory: {{ .Values.trainer.resources.limits.memory }} cpu: {{ .Values.trainer.resources.limits.cpu }} {{- end }} + {{- if .Values.trainer.resourceClaim.enabled }} + claims: + - name: {{ .Values.trainer.resourceClaim.name }} + {{- end }} {{- if and .Values.trainer.probes .Values.trainer.probes.enabled }} startupProbe: httpGet: @@ -366,15 +434,28 @@ spec: failureThreshold: {{ .Values.trainer.probes.readiness.failureThreshold }} timeoutSeconds: {{ .Values.trainer.probes.readiness.timeoutSeconds }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.trainer.configMap }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.trainer.configMap }} + - name: rl-configs + mountPath: /configs + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.trainer.configMap }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ .Values.storage.existingClaim | default (printf "%s-shared-data" .Release.Name) }} + {{- end }} + {{- if .Values.trainer.configMap }} + - name: rl-configs + configMap: + name: {{ .Values.trainer.configMap }} + {{- end }} {{- end }} {{- end }} diff --git a/k8s/prime-rl/templates/pvc.yaml b/k8s/prime-rl/templates/pvc.yaml index 7afd6ff386..a3e37a0f54 100644 --- a/k8s/prime-rl/templates/pvc.yaml +++ b/k8s/prime-rl/templates/pvc.yaml @@ -1,4 +1,9 @@ -{{- if .Values.storage.enabled }} +{{- /* + Only create a PVC when storage is enabled AND the user has not supplied an + existing claim. When `storage.existingClaim` is set the chart mounts that + claim instead — e.g. the cluster-shared `shared-model-cache` PVC. +*/ -}} +{{- if and .Values.storage.enabled (not .Values.storage.existingClaim) }} apiVersion: v1 kind: PersistentVolumeClaim metadata: diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index 7fc3b48989..1905c6ddab 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -14,7 +14,14 @@ image: # Shared storage configuration storage: enabled: true - # PVC name will be automatically set to {{ .Release.Name }}-shared-data + # If existingClaim is set, the chart will not create a PVC and will instead + # mount the named claim. Useful for cluster-wide model caches (e.g. the + # auto-provisioned `shared-model-cache` PVC documented at + # work/ops/shared-model-cache.md) so models downloaded by other namespaces + # are reused with zero re-download. + existingClaim: "" + # Used only when existingClaim is empty (chart creates the PVC). + # Resulting PVC name: {{ .Release.Name }}-shared-data storageClassName: nfs accessModes: - ReadWriteMany @@ -50,6 +57,19 @@ orchestrator: nodeSelector: {} # nvidia.com/gpu.present: "true" # Orchestrator doesn't need GPUs + tolerations: [] + # - key: "nvidia.com/gpu" + # operator: "Exists" + # effect: "NoSchedule" + + imagePullSecrets: [] + # - name: regcred + + # Optional: name of an existing ConfigMap to mount at /configs in the + # orchestrator container. Useful for shipping TOML configs without rebuilding + # the image. Leave empty to skip. + configMap: "" + # Inference component inference: enabled: true @@ -63,6 +83,18 @@ inference: enabled: true count: 1 + # Dynamic Resource Allocation (DRA) — see trainer.resourceClaim for semantics. + resourceClaim: + enabled: false + templateName: "" + name: "gpus" + + # Optional: name of an existing ConfigMap to mount at /configs in the + # inference container. Symmetric with orchestrator.configMap and + # trainer.configMap. Lets a native vLLM inference server read its + # infer.toml from a ConfigMap. + configMap: "" + resources: requests: memory: "4Gi" @@ -75,6 +107,9 @@ inference: runtimeClassName: nvidia + tolerations: [] + imagePullSecrets: [] + # Health probes for inference server (/health for startup/readiness, /liveness for engine liveness) probes: enabled: false @@ -108,6 +143,22 @@ trainer: enabled: true count: 1 + # Dynamic Resource Allocation (DRA) — for DRA-only node pools that reject the + # legacy `nvidia.com/gpu` device-plugin resource. When `resourceClaim.enabled` + # is true the chart will NOT emit `nvidia.com/gpu` requests/limits and will + # instead add a `spec.resourceClaims` entry to the pod and a `resources.claims` + # reference inside the container. + # templateName: name of a pre-existing ResourceClaimTemplate in the same ns + # name: logical claim name used to link pod <-> container + resourceClaim: + enabled: false + templateName: "" + name: "gpus" + + # Optional: name of an existing ConfigMap to mount at /configs in the + # trainer container. Symmetric with `orchestrator.configMap`. + configMap: "" + resources: requests: memory: "4Gi" @@ -123,6 +174,9 @@ trainer: runtimeClassName: nvidia + tolerations: [] + imagePullSecrets: [] + # Health probes for trainer (requires metrics_server config) probes: enabled: false diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 9579259366..203cc37f83 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -342,6 +342,13 @@ class InferenceConfig(BaseConfig): enable_return_routed_experts: bool = False """Return routed experts in responses. Forwarded as ``--enable-return-routed-experts``.""" + moe_backend: str = "auto" + """MoE kernel backend, forwarded as vLLM ``--moe-backend``. The default fused + backends (e.g. FlashInfer TRTLLM) compute gating+routing+experts in one kernel and + bypass the Python ``BaseRouter`` that the routed-experts capture hook binds to, so + capture yields all-zero routing. Use a non-fused backend (``triton``) when + ``enable_return_routed_experts``/router replay is on so real routing is captured.""" + enable_fp32_lm_head: bool = True """Run the lm_head projection in fp32 via a native bf16×bf16 → fp32 GEMM (``torch.mm`` with ``out_dtype=torch.float32``). Stabilizes logprob precision under FP8/bf16 inference, matching SGLang's ``--enable-fp32-lm-head``. Implemented as a monkey-patch over vLLM's LogitsProcessor, activated by setting ``additional_config["fp32_lm_head"] = True`` on the vLLM config.""" @@ -503,6 +510,7 @@ def to_vllm(self) -> Namespace: "gpu_memory_utilization": "gpu_memory_utilization", "api_server_count": "api_server_count", "enable_return_routed_experts": "enable_return_routed_experts", + "moe_backend": "moe_backend", "enable_expert_parallel": "enable_expert_parallel", "all2all_backend": "all2all_backend", "enable_eplb": "enable_eplb", diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index be5fe249f3..a3c28957b9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -907,5 +907,10 @@ def resolve_env_config(self): if is_vllm: env.sampling.extra_body.setdefault("top_k", -1) env.sampling.extra_body.setdefault("min_p", 0.0) - env.sampling.extra_body.setdefault("return_token_ids", True) + # DRA-TITO-PATCH: Dynamo strict-rejects `return_token_ids`. + # Token IDs come back via `nvext.engine_data.completion_token_ids` + # (PR #8119 channel) on the rl-sdk-2 backend, so the client does + # not need this vLLM-native field set. Stripped here at config + # validate time so it never reaches the inference server. + # env.sampling.extra_body.setdefault("return_token_ids", True) return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index dab46a9ce1..6aa3cd1ff7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -428,9 +428,19 @@ def auto_setup_router_replay(self): stacklevel=2, ) self.inference.enable_return_routed_experts = True + # Fused MoE backends (FlashInfer/TRTLLM) bypass the BaseRouter + # capture hook, yielding all-zero routing. Force a non-fused + # backend for capture unless the user set one explicitly. + if self.inference.moe_backend == "auto": + warnings.warn( + "Router replay is enabled; setting inference.moe_backend='triton' " + "(fused MoE backends bypass routed-experts capture).", + stacklevel=2, + ) + self.inference.moe_backend = "triton" else: warnings.warn( - "Router replay is enabled, but inference is not configured. When manually starting the inference server, make sure to pass `--enable-return-routed-experts` to the vLLM server.", + "Router replay is enabled, but inference is not configured. When manually starting the inference server, make sure to pass `--enable-return-routed-experts` and `--moe-backend triton` to the vLLM server.", stacklevel=2, ) return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index ff311f145d..9c16c410bd 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -136,6 +136,12 @@ class ClientConfig(BaseConfig): elastic: ElasticConfig | None = None """Elastic inference pool config for DNS-based service discovery. When set, ``base_url`` is ignored and inference servers are discovered dynamically via DNS.""" + rl_base_url: list[str] | None = None + """Dynamo RL worker discovery base URLs. Used only for backend='dynamo' when admin_base_url is unset. These URLs point at the Dynamo RL discovery listener (DYN_RL_PORT, default 8001), which serves GET /v1/rl/workers. If unset, prime-rl derives the discovery URL from base_url by replacing the port with DYN_RL_PORT or 8001.""" + + backend: Literal["vllm", "dynamo"] = "vllm" + """Inference backend selector. Picks the AdminAPI implementation used for pause/resume/update_weights/load_lora_adapter/list_models. Default 'vllm' matches prime-rl's bundled vLLM frontend. 'dynamo' targets NVIDIA Dynamo's worker /engine/* admin routes on admin_base_url and routes /v1/models to the OpenAI-compat base_url.""" + router_url: str | None = None """vllm-router URL for load-aware inference routing. With elastic mode, inference requests go through the router while admin ops still hit discovered pods directly.""" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 00f4e07deb..00a5325b7f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -457,12 +457,25 @@ class BaseWeightBroadcastConfig(BaseConfig): class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig): type: Literal["filesystem"] = "filesystem" - - save_sharded: bool = True - """Save the weight checkpoint in sharded format.""" - - save_format: Literal["safetensors", "torch"] = "safetensors" - """Weight checkpoint serialization format.""" + save_sharded: Annotated[bool, Field(description="Whether to save the weight checkpoint in sharded format.")] = True + save_format: Annotated[ + Literal["safetensors", "torch"], Field(description="The format to save the weight checkpoint in.") + ] = "safetensors" + keep_recent: Annotated[ + int | None, + Field( + ge=0, + description=( + "Number of recent broadcast step directories to keep on disk. At trainer step N, " + "step_{N - keep_recent - 1} is deleted. If None (default), falls back to " + "trainer.max_async_level so retention matches the staleness window. Set this above " + "max_async_level for LoRA + filesystem broadcast: vLLM lazy-loads the adapter inside " + "_prepare_inputs, so a generate request admitted under an older adapter can page that " + "dir in *after* the orchestrator has moved on. A +1 or +2 buffer over max_async_level " + "is typical." + ), + ), + ] = None class NCCLWeightBroadcastConfig(BaseWeightBroadcastConfig): diff --git a/scripts/docker-arm64-post-install.sh b/scripts/docker-arm64-post-install.sh index 55f85a3a03..12a71ad987 100755 --- a/scripts/docker-arm64-post-install.sh +++ b/scripts/docker-arm64-post-install.sh @@ -42,3 +42,7 @@ echo " target venv: $VENV_PATH" echo "=== reinstalling flash-attn-cute (flash-attn overwrites it with a stub) ===" uv pip install --python "$VENV_PATH/bin/python" --reinstall --no-deps \ "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@96bd151#subdirectory=flash_attn/cute" + +echo "=== building DeepGEMM from source (FP8 grouped-GEMM for SM100 / GB200) ===" +# Runs from /app (WORKDIR), so uv auto-detects the project venv at /app/.venv. +bash /app/scripts/install_deep_gemm.sh diff --git a/src/prime_rl/_compat.py b/src/prime_rl/_compat.py index 268c38376b..1ab251a777 100644 --- a/src/prime_rl/_compat.py +++ b/src/prime_rl/_compat.py @@ -35,7 +35,9 @@ # --------------------------------------------------------------------------- try: import transformers.integrations.hub_kernels as _hub_kernels -except ImportError: + # AttributeError: module exists in some builds but lazy_load_kernel was added in 5.5 + _ = _hub_kernels.lazy_load_kernel +except (ImportError, AttributeError): _hub_kernels = None # transformers < 5.5, no patch needed if _hub_kernels is not None: diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py index 8190f210b0..7d85153c47 100644 --- a/src/prime_rl/inference/patches.py +++ b/src/prime_rl/inference/patches.py @@ -18,6 +18,9 @@ def transformers_v5_compat(): _patch_qwen35_lora() _patch_lora_key_prefix() monkey_patch_deep_gemm_silu_mul_quant_int64() + monkey_patch_deep_gemm_silu_mul_quant_packed_int64() + # monkey_patch_dp_engine_core_pause_resume_deadlock() # DISABLED: use vLLM PR #39366-native two-phase pause (avoid double pause/resume fix) + monkey_patch_fp32_lm_head() monkey_patch_vllm_padded_input_scrub() monkey_patch_return_routed_experts_with_nixl_connector() @@ -207,6 +210,207 @@ def monkey_patch_deep_gemm_silu_mul_quant_int64(): logger.warning("Enabled int64-addressing Triton patch for vLLM DeepGEMM SiLU/mul FP8 quant.") +@triton.jit +def _silu_mul_quant_fp8_packed_kernel_int64( + input_ptr, + output_q_ptr, + output_scale_ptr, + M: tl.int64, + input_stride_m: tl.int64, + output_q_stride_m: tl.int64, + output_scale_stride_k: tl.int64, + clamp_limit, + N: tl.constexpr, + NUM_GROUPS: tl.constexpr, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + GROUP_SIZE: tl.constexpr, + BLOCK_M: tl.constexpr, + HAS_CLAMP: tl.constexpr, +): + """int64-addressing variant of vLLM's _silu_mul_quant_fp8_packed_kernel. + + Mirrors `vllm/model_executor/layers/quantization/utils/fp8_utils.py:152` + but casts row/column offsets to tl.int64 so the address arithmetic in + `base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m` does not + overflow when M * input_stride_m exceeds 2**31 (e.g. Qwen3-235B profile_run + on EP=2 × DP=4 layouts, 128 experts × 6144 fused MoE dim). + """ + N_2: tl.constexpr = N // 2 + + # Grid layout: dim0 = M-blocks (large, up to 2**31-1 on Blackwell), + # dim1 = packed groups (small, ~3-6). The upstream kernel placed the + # large M dim on grid Y, which is capped at 65,535 and produced + # `Triton Error [CUDA]: invalid argument` on Qwen3-235B profile_run. + pid_m = tl.program_id(0) + pid_pack = tl.program_id(1) + m_offset = (pid_m * BLOCK_M).to(tl.int64) + + if m_offset >= M: + return + + offs_m = tl.arange(0, BLOCK_M).to(tl.int64) + offs_n = tl.arange(0, GROUP_SIZE).to(tl.int64) + row_mask = (m_offset + offs_m) < M + + base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m + base_out_offset = (m_offset + offs_m[:, None]) * output_q_stride_m + + packed_scale = tl.zeros((BLOCK_M,), dtype=tl.int32) + + for pack_idx in tl.static_range(4): + group_id = pid_pack * 4 + pack_idx + + if group_id < NUM_GROUPS: + n_offset = (group_id * GROUP_SIZE).to(tl.int64) + + act_ptrs = input_ptr + base_row_offset + n_offset + offs_n[None, :] + act_in = tl.load(act_ptrs, mask=row_mask[:, None], other=0.0) + + mul_ptrs = act_ptrs + N_2 + mul_in = tl.load(mul_ptrs, mask=row_mask[:, None], other=0.0) + + act_f32 = act_in.to(tl.float32) + mul_f32 = mul_in.to(tl.float32) + + if HAS_CLAMP: + act_f32 = tl.minimum(act_f32, clamp_limit) + mul_f32 = tl.clamp(mul_f32, -clamp_limit, clamp_limit) + + y = (act_f32 / (1.0 + tl.exp(-act_f32))) * mul_f32 + # Round through bf16 to match unfused precision path + y = y.to(tl.bfloat16).to(tl.float32) + + absmax = tl.max(tl.abs(y), axis=1) + + scale_raw = tl.maximum(absmax / fp8_max, 1e-10) + exponent = tl.ceil(tl.log2(scale_raw)) + scale = tl.math.exp2(exponent) + + y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max) + + out_q_ptrs = output_q_ptr + base_out_offset + n_offset + offs_n[None, :] + tl.store( + out_q_ptrs, + y_q.to(output_q_ptr.dtype.element_ty), + mask=row_mask[:, None], + ) + + exponent_biased = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32) + packed_scale = packed_scale | (exponent_biased << (pack_idx * 8)) + + scale_ptrs = output_scale_ptr + pid_pack.to(tl.int64) * output_scale_stride_k + m_offset + offs_m + tl.store(scale_ptrs, packed_scale, mask=row_mask) + + +def silu_mul_quant_fp8_packed_triton_int64( + input: torch.Tensor, + group_size: int = 128, + output_q: torch.Tensor | None = None, + clamp_limit: float | None = None, +): + """int64-addressing variant of vLLM's silu_mul_quant_fp8_packed_triton. + + Same semantics and return shape as the upstream function, but launches the + int64 kernel above. Required for Qwen3-235B FP8 inference with DeepGEMM + enabled — the upstream packed kernel overflows on row offsets at profile_run + shapes (see issues.md Issue 7 in work/bis-dev/may-26/01-qwen-235b-whiteffiber). + """ + assert input.dim() == 2 + assert input.is_contiguous() + + M, N = input.shape + N_2 = N // 2 + + assert N_2 % group_size == 0 + + fp8_dtype = torch.float8_e4m3fn + finfo = torch.finfo(fp8_dtype) + fp8_min, fp8_max = finfo.min, finfo.max + + num_groups_per_row = N_2 // group_size + num_packed_groups = (num_groups_per_row + 3) // 4 + tma_aligned_M = ((M + 3) // 4) * 4 + + if output_q is None: + output_q = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device) + + output_scale_packed = torch.zeros( + (num_packed_groups, tma_aligned_M), + dtype=torch.int32, + device=input.device, + ).T[:M, :] + + BLOCK_M = 8 + # gridX gets the large M dim (Blackwell cap is 2**31-1); gridY gets + # the small packed-group dim (Blackwell cap is 65,535). Swapping vs. + # upstream avoids "invalid argument" when (M+7)//8 > 65,535 + # (Qwen3-235B with 128 experts on EP=2 trips this in profile_run). + grid = ((M + BLOCK_M - 1) // BLOCK_M, num_packed_groups) + + num_warps = max(4, group_size // 32) + num_stages = 2 + + has_clamp = clamp_limit is not None + _silu_mul_quant_fp8_packed_kernel_int64[grid]( + input, + output_q, + output_scale_packed, + M, + input.stride(0), + output_q.stride(0), + output_scale_packed.stride(1), + clamp_limit if has_clamp else 0.0, + N=N, + NUM_GROUPS=num_groups_per_row, + fp8_min=fp8_min, + fp8_max=fp8_max, + GROUP_SIZE=group_size, + BLOCK_M=BLOCK_M, + HAS_CLAMP=has_clamp, + num_warps=num_warps, + num_stages=num_stages, + ) + + return output_q, output_scale_packed + + +def monkey_patch_deep_gemm_silu_mul_quant_packed_int64(): + """Replace vLLM's silu_mul_quant_fp8_packed_triton with our int64 variant. + + The upstream packed kernel in vLLM 0.21.0 uses int32 arithmetic for row + offsets (`(m_offset + offs_m[:, None]) * input_stride_m`). For Qwen3-235B + FP8 inference with `VLLM_USE_DEEP_GEMM=1`, profile_run hits M*stride > 2**31 + on the 128-expert layout and Triton emits CUDA "invalid argument". + + We patch: + 1. fp8_utils.silu_mul_quant_fp8_packed_triton (the public wrapper) + 2. deep_gemm_moe.fused_silu_mul_fp8_quant_packed (the alias inside the + DeepGEMM MoE experts module — imported by name at module load) + """ + import sys + + from vllm.logger import init_logger + from vllm.model_executor.layers.quantization.utils import fp8_utils + + logger = init_logger(__name__) + + fp8_utils.silu_mul_quant_fp8_packed_triton = silu_mul_quant_fp8_packed_triton_int64 + + deep_gemm_moe_module = sys.modules.get("vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe") + if deep_gemm_moe_module is not None: + # deep_gemm_moe.py re-exports the wrapper under the name + # `fused_silu_mul_fp8_quant_packed` (see line 208 in the upstream module). + if hasattr(deep_gemm_moe_module, "fused_silu_mul_fp8_quant_packed"): + deep_gemm_moe_module.fused_silu_mul_fp8_quant_packed = silu_mul_quant_fp8_packed_triton_int64 + if hasattr(deep_gemm_moe_module, "silu_mul_quant_fp8_packed_triton"): + deep_gemm_moe_module.silu_mul_quant_fp8_packed_triton = silu_mul_quant_fp8_packed_triton_int64 + + logger.warning( + "Enabled int64-addressing Triton patch for vLLM DeepGEMM packed SiLU/mul FP8 quant." + ) + + def _patch_qwen35_lora(): """Fix Qwen3.5 LoRA: align packed_modules_mapping with output_sizes. @@ -803,6 +1007,10 @@ def monkey_patch_fp32_lm_head(): logger = init_logger(__name__) + if getattr(LogitsProcessor, "_prime_rl_fp32_lm_head_patch_installed", False): + logger.debug("fp32 lm_head patch already installed; skipping.") + return + _original_init = LogitsProcessor.__init__ _original_get_logits = LogitsProcessor._get_logits @@ -835,4 +1043,5 @@ def _patched_get_logits(self, hidden_states, lm_head, embedding_bias): LogitsProcessor.__init__ = _patched_init LogitsProcessor._get_logits = _patched_get_logits + LogitsProcessor._prime_rl_fp32_lm_head_patch_installed = True logger.info("Installed fp32 lm_head patch (native out_dtype=fp32 mm).") diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py index e9ef70b049..0bbd625f1a 100644 --- a/src/prime_rl/inference/vllm/routed_experts.py +++ b/src/prime_rl/inference/vllm/routed_experts.py @@ -17,13 +17,22 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A assert np.issubdtype(array.dtype, np.integer) if array.size: assert array.min() >= 0 - assert array.max() <= np.iinfo(np.uint8).max - - compact = np.ascontiguousarray(array.astype(np.uint8, copy=False)) + # Preserve vLLM's per-MODEL dtype (RoutedExpertsManager uses uint8 for + # <=256 experts, uint16 otherwise), so every sample of a model shares one + # dtype. Re-narrowing per-sample on the observed max id would emit uint8 + # for one sample and uint16 for another of the SAME >256-expert model, + # which the trainer's same-dtype routed-experts packing rejects. Only an + # unexpectedly wide capture dtype (e.g. an int64 buffer) is capped to int32 + # (still consistent per model). The dtype rides the payload. + if array.dtype in (np.uint8, np.uint16, np.int16, np.int32): + compact = np.ascontiguousarray(array) + else: + compact = np.ascontiguousarray(array.astype(np.int32, copy=False)) return { "data": pybase64.b64encode(memoryview(compact)).decode("ascii"), "shape": list(compact.shape), "start": start, + "dtype": compact.dtype.name, } diff --git a/src/prime_rl/inference/vllm/worker/nccl.py b/src/prime_rl/inference/vllm/worker/nccl.py index 95d35b583c..1acca94dcd 100644 --- a/src/prime_rl/inference/vllm/worker/nccl.py +++ b/src/prime_rl/inference/vllm/worker/nccl.py @@ -139,16 +139,81 @@ def init_broadcaster( inference_world_size: Total number of inference GPUs across all servers. """ self.quantize_in_weight_transfer = quantize_in_weight_transfer - # Use the worker's device index directly as the local rank. - # The previous dp_group-based computation broke in vLLM v1 multiprocess - # DP mode where each worker is a separate process with a singleton - # DP group (rank_in_group is always 0). + + # ===================================================================== + # PATCHED — DP+TP-aware rank computation (Option B: parallel_config + # dispatch). Overrides the image-baked nccl.py via ConfigMap subPath mount. + # + # The vLLM Worker exposes `self.rank` (executor-wide rank) and + # `self.device.index` (local GPU index 0..local_world_size-1). Neither + # alone correctly identifies the GPU's position across BOTH TP and DP + # multinode topologies: + # + # - TP/PP multinode (one executor spans pods): + # self.rank is unique 0..world_size-1 + # self.device.index repeats per node (0..3 on each node) + # --> use self.rank + # + # - DP-only (TP=1) — e.g. vLLM `--data-parallel-size N --tp 1`: + # self.rank is 0 for every subprocess + # self.device.index is the LOCAL gpu index (0..size_local-1); for a + # MULTINODE-DP follower it does NOT equal the global DP rank + # --> use parallel_config.data_parallel_rank (global DP rank) + # + # Dispatch on parallel_config (always available on a vLLM Worker via + # vllm_config or as a direct attribute, depending on vLLM version). + # ===================================================================== + pc = None + for attr_path in ("vllm_config.parallel_config", "parallel_config"): + obj = self + try: + for part in attr_path.split("."): + obj = getattr(obj, part) + pc = obj + break + except AttributeError: + continue + + tp_size = getattr(pc, "tensor_parallel_size", 1) if pc is not None else 1 + pp_size = getattr(pc, "pipeline_parallel_size", 1) if pc is not None else 1 + worker_rank = getattr(self, "rank", 0) local_rank = self.device.index - global_rank_inference = rank_offset + local_rank + dp_rank = getattr(pc, "data_parallel_rank", None) if pc is not None else None + + if tp_size > 1 or pp_size > 1: + # Executor spans multiple ranks (TP or PP); self.rank is canonical. + effective_rank = worker_rank + rank_source = "self.rank (TP/PP > 1)" + elif dp_rank is not None: + # Pure DP: use the GLOBAL data-parallel rank. Correct for BOTH single-node + # DP (dp_rank == device.index) AND MULTINODE DP, where a follower node's + # device.index is the LOCAL gpu index (0..size_local-1) and does NOT equal + # the global DP rank (e.g. follower DP ranks 4..7 have device.index 0..3). + # Using device.index there collides with the leader's ranks 0..3 and the + # NCCL broadcast group never forms (Bootstrap "rank N already checked in"). + effective_rank = local_rank # DGD per-pod-offset override (issue #7); proper fix = per-engine offsets in orch + rank_source = "self.device.index (DGD per-pod-offset, issue #7)" + else: + # Fallback (older vLLM lacking data_parallel_rank): single-node DP only, + # where device.index uniquely identifies the GPU within the pod (0..N-1). + effective_rank = local_rank + rank_source = "self.device.index (DP-only fallback)" + + global_rank_inference = rank_offset + effective_rank logger.info( - f"Worker [local_rank={local_rank} rank_offset={rank_offset}] " - f"-> [global_rank={global_rank_inference} inference_world_size={inference_world_size}]" + "Worker [tp_size=%s pp_size=%s local_rank=%s worker_rank=%s " + "effective_rank=%s (%s) rank_offset=%s] " + "-> [global_rank=%s inference_world_size=%s] (dp-tp-aware-patch)", + tp_size, + pp_size, + local_rank, + worker_rank, + effective_rank, + rank_source, + rank_offset, + global_rank_inference, + inference_world_size, ) self.nccl_broadcast_receiver = NCCLWeightBroadcastReceiver( diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py index 7b448f021a..efe6f0a7c9 100644 --- a/src/prime_rl/inference/vllm/worker/weight_transfer.py +++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py @@ -84,12 +84,121 @@ def build_expert_map(model: Module) -> dict[str, torch.Tensor]: return source_indices_by_module +def _try_e8m0_scale_conversion( + name: str, + received_scale: torch.Tensor, + param: torch.Tensor, + params: dict[str, torch.Tensor], + updated_weights: set[str], +) -> bool: + """Convert a trainer-format blockwise FP8 scale to vLLM's E8M0 TMA layout. + + When DeepGEMM E8M0 is active (GB200 default), vLLM stores weight_scale_inv + tensors in a TMA-aligned [N, K/512] packed UE8M0 layout. The trainer sends + [N/128, K/128] float32 blockwise scales. + + This delegates to vLLM's own unified helper + ``deepgemm_post_process_fp8_weight_block``, which is the *exact* call vLLM + invokes at initial model load (via ``DeepGemmFp8BlockScaledMMKernel`` for + dense linears and ``prepare_fp8_moe_layer_for_deepgemm`` for MoE). The + helper: + + 1. Re-quantises the FP8 weight to UE8M0 (power-of-two) scales in-place + (``requant_weight_ue8m0_inplace``). + 2. Repacks the scale tensor to the TMA-aligned ``[..., N, K/512]`` layout + (``transform_sf_into_required_layout`` with recipe ``(1, 128, 128)``). + 3. Handles 2D (dense) vs 3D (MoE) dispatch internally. + + Reusing vLLM's wrapper guarantees the broadcast path produces the same + in-memory layout as vLLM's own startup load path — drift-free by + construction. + + Returns True on success, False if conversion is not applicable or fails. + The caller is responsible for raising shape-mismatch errors for False returns. + + ORDERING REQUIREMENT: the corresponding FP8 weight param must have been + updated (copied from the received tensor) before this function is called, + i.e. the weight must appear before its weight_scale_inv in the state iterator. + """ + # Only handle weight_scale_inv tensors. + if not name.endswith("weight_scale_inv"): + return False + + # Derive the corresponding weight parameter name. + # e.g. "...qkv_proj.weight_scale_inv" -> "...qkv_proj.weight" + # "...w13_weight_scale_inv" -> "...w13_weight" + weight_name = name[: -len("weight_scale_inv")] + "weight" + if weight_name not in params: + return False + + # The weight must have been updated in this broadcast pass so that + # requant_weight_ue8m0_inplace dequantises the *new* FP8 values. + if weight_name not in updated_weights: + logger.warning( + "E8M0 scale conversion for %s: weight %s was not updated in this pass " + "(scale arrived before weight). Skipping conversion.", + name, + weight_name, + ) + return False + + try: + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + deepgemm_post_process_fp8_weight_block, + ) + except ImportError as exc: + logger.debug("E8M0 conversion unavailable: %s", exc) + return False + + weight_param = params[weight_name] + + try: + # Single unified call: re-quantises weight in-place with UE8M0 scales + # AND repacks the scale tensor to the TMA-aligned [..., N, K/512] layout + # that vLLM stores. Handles 2D (dense) and 3D (MoE) dispatch internally. + # The first return value is the same underlying storage as weight_param + # (modified in-place); we only need the new scale tensor. + _wq, dg_scale = deepgemm_post_process_fp8_weight_block( + wq=weight_param, + ws=received_scale, + quant_block_shape=(128, 128), + use_e8m0=True, + ) + param.copy_(dg_scale.view_as(param)) + logger.debug("E8M0 scale conversion applied for %s", name) + return True + + except Exception as exc: + logger.warning("E8M0 scale conversion failed for %s: %s", name, exc) + return False + + @torch.no_grad() def load_weights_kernel(model: Module, state_iter: Generator[tuple[str, torch.Tensor], None, None]) -> None: - """Load vLLM kernel-format tensors using in-place copy_ updates.""" + """Load vLLM kernel-format tensors using in-place copy_ updates. + + Handles the GB200 DeepGEMM E8M0 scale mismatch: the trainer broadcasts + blockwise [N/128, K/128] float32 scales, while vLLM with E8M0 enabled + stores weight_scale_inv in [N, K/512] TMA-aligned packed UE8M0 layout. + When is_deep_gemm_e8m0_used() is True, scale tensors with a shape mismatch + are automatically converted via requant_weight_ue8m0_inplace + + transform_sf_into_required_layout. + """ params = dict(model.named_parameters()) expert_source_indices = build_expert_map(model) + # Detect E8M0 once (cached call, cheap). + try: + from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used + _use_e8m0 = is_deep_gemm_e8m0_used() + except ImportError: + _use_e8m0 = False + + # Track which weight params have been updated in this call so we can + # safely call requant_weight_ue8m0_inplace on them (it must dequantise + # the *new* FP8 values, not the stale ones from the initial model load). + updated_weights: set[str] = set() + loaded = 0 skipped: list[str] = [] shape_mismatches: list[str] = [] @@ -108,11 +217,19 @@ def load_weights_kernel(model: Module, state_iter: Generator[tuple[str, torch.Te break if param.shape != tensor.shape: + # Attempt E8M0 scale conversion before declaring a mismatch. + if _use_e8m0 and _try_e8m0_scale_conversion( + name, tensor, param, params, updated_weights + ): + loaded += 1 + continue + shape_mismatches.append(f"{name}: param={list(param.shape)} != received={list(tensor.shape)}") continue param.copy_(tensor) loaded += 1 + updated_weights.add(name) if shape_mismatches: raise ValueError(f"Kernel weight transfer had {len(shape_mismatches)} shape mismatches: {shape_mismatches}") diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index d685b519b6..8b21c7fa4d 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -314,6 +314,10 @@ async def setup(self) -> None: await self.inference_metrics.start() get_logger().info(f"Initializing weight broadcast ({config.weight_broadcast})") + # Propagate the configured broadcast type to the admin API. The Dynamo + # backend gates its NCCL init path on ``_weight_broadcast_type``. + if hasattr(self.student_inference._admin_api, "_weight_broadcast_type"): + self.student_inference._admin_api._weight_broadcast_type = config.weight_broadcast.type if config.weight_broadcast.type == "nccl": await init_nccl_broadcast( self.student_inference.admin_clients, @@ -322,6 +326,7 @@ async def setup(self) -> None: config.weight_broadcast.timeout, inference_world_size=config.weight_broadcast.inference_world_size, quantize_in_weight_transfer=config.weight_broadcast.quantize_in_weight_transfer, + admin=self.student_inference._admin_api, ) get_logger().info(f"Initializing training batch sender ({config.rollout_transport})") diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 3e8431c12a..f5e11984cd 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -247,9 +247,20 @@ def prepare_step_tokens(step: vf.TrajectoryStep, step_idx: int) -> dict[str, Any routed_experts_start = None if routed_experts_payload is not None: decoded_routed_experts = pybase64.b64decode_as_bytearray(routed_experts_payload["data"]) - routed_experts = np.frombuffer(decoded_routed_experts, dtype=np.uint8).reshape( + # dtype rides the payload so >256-expert MoEs (uint16) and + # int32 captures decode correctly; default uint8 for payloads + # serialized before the dtype field existed. + re_dtype = np.dtype(routed_experts_payload.get("dtype", "uint8")) + routed_experts = np.frombuffer(decoded_routed_experts, dtype=re_dtype).reshape( routed_experts_payload["shape"] ) + # The trainer (batch.py row sizing + tensor decode) supports + # uint8/int16/int32 (all torch-safe). Normalize a uint16 wire + # dtype (vLLM/dynamo >256-expert capture) to int32 so it never + # reaches the trainer; expert ids are non-negative so this is + # lossless. + if routed_experts.dtype == np.uint16: + routed_experts = routed_experts.astype(np.int32) routed_experts_start = routed_experts_payload["start"] return { diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 5675ba3f34..c056f12cd7 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -4,16 +4,20 @@ from concurrent.futures import ThreadPoolExecutor from itertools import cycle from pathlib import Path +from typing import Any import orjson +import pandas as pd import verifiers as vf +from rich.console import Console +from rich.table import Table from verifiers.utils.client_utils import setup_openai_client from verifiers.utils.save_utils import make_serializable from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.transport import TrainingSample from prime_rl.utils.client import setup_inference_pool -from prime_rl.utils.logger import InterceptHandler, get_logger +from prime_rl.utils.logger import InterceptHandler, format_time, get_logger from prime_rl.utils.utils import ( get_broadcast_dir, get_ckpt_dir, @@ -96,58 +100,177 @@ def set_default_executor(max_workers: int = 64) -> None: asyncio.get_event_loop().set_default_executor(ThreadPoolExecutor(max_workers=max_workers)) +def print_benchmark(history: dict[str, list[Any]]) -> None: + """ + Print benchmark results as rich table. Shows formatted step time values. + First N rows show the per-step values, and the last row shows the mean, + std, min, and max values. + """ + history.pop("step") + assert all(len(v) for v in history.values()), "All metrics must have logged the same number of steps" + + # Turn metric history into pd.DataFrame + df = pd.DataFrame(dict(history.items())) + columns = { + "time/step": "Step Time", + } + df = df.rename(columns=columns) + df = df[list(columns.values())] + df = df.iloc[1:] # Exclude first row + + # Setup console + console = Console() + table = Table(title="Benchmark") + + # Add columns + table.add_column("Step", justify="right") + for col in df.columns: + table.add_column(col, justify="center", style="magenta") + + # Add formatted rows + formatted_df = pd.DataFrame(columns=df.columns) + formatted_df["Step Time"] = df["Step Time"].apply(format_time) + for step, row in formatted_df.iterrows(): + table.add_row(*([str(step)] + [str(x) for x in row])) + + # Separator + num_table_columns = 1 + len(df.columns) + table.add_row(*([""] * num_table_columns)) + + # Add row for formatted, aggregated statistics + mean_df = df.describe().loc[["mean", "std", "min", "max"], :] + formatted_mean_df = pd.DataFrame(columns=mean_df.columns) + formatted_mean_df["Step Time"] = mean_df["Step Time"].apply(format_time) + mean_row = ["Overall"] + formatted_mean_df.T.apply( + lambda row: f"{row['mean']} ± {row['std']} [{row['min']}, {row['max']}]", axis=1 + ).tolist() + table.add_row(*mean_row) + + # Display table + console.print(table) + + +def _flatten_prompt_logprobs(raw: list[Any] | None) -> list[float]: + """Shared flattener used by both transports. + + ``prompt_logprobs[i]`` is a ``{token_id: Logprob}`` dict for tokens the + engine could score, or ``None`` for the leading token which has no + preceding context. Flatten to ``list[float]`` with 0.0 in the unscored + slot. Accepts both vLLM's typed ``Logprob`` objects and dynamo's + ``PromptLogprobEntry`` dict shape (`{logprob, rank?, decoded_token?}`). + """ + flat: list[float] = [] + for entry in raw or []: + if not entry: + flat.append(0.0) + continue + first = next(iter(entry.values())) + lp = first.logprob if hasattr(first, "logprob") else first.get("logprob") + flat.append(float(lp) if lp is not None else 0.0) + return flat + + +async def _compute_teacher_logprobs_vllm( + client_config: vf.ClientConfig, model_name: str, sample: TrainingSample +) -> list[float]: + """Legacy path: prime-rl's vLLM sidecar ``/inference/v1/generate``.""" + import httpx + from vllm.entrypoints.serve.disagg.protocol import GenerateResponse + + client = setup_openai_client(client_config) + # Two escape hatches from ``AsyncOpenAI.post``: + # 1. URL — ``/inference/v1/generate`` is mounted at server root, not + # under ``/v1``. Pass an absolute URL so the SDK's ``_prepare_url`` + # skips the base-url merge. + # 2. Parse — vLLM's ``GenerateResponse`` isn't an ``openai.BaseModel``. + # Use ``cast_to=httpx.Response`` and validate the body ourselves. + base = str(client.base_url).rstrip("/").removesuffix("/v1") + http_response = await client.post( + f"{base}/inference/v1/generate", + cast_to=httpx.Response, + body={ + "model": model_name, + "token_ids": list(sample.prompt_ids) + list(sample.completion_ids), + "sampling_params": { + "max_tokens": 1, + "temperature": 1.0, + "top_p": 1.0, + "prompt_logprobs": 1, + }, + }, + ) + response = GenerateResponse.model_validate_json(http_response.content) + return _flatten_prompt_logprobs(response.prompt_logprobs) + + +async def _compute_teacher_logprobs_dynamo( + client_config: vf.ClientConfig, model_name: str, sample: TrainingSample +) -> list[float]: + """rl-sdk-2 path: dynamo via ``/v1/chat/completions`` with nvext envelope. + + Wire shape (per plan.md A1+A3 — already implemented on rl-sdk-2): + - top-level ``prompt_logprobs: 1`` (CommonExt sampling param) + - ``nvext.token_data`` carries pre-tokenized prompt + - ``nvext.extra_fields = ["prompt_logprobs"]`` opts into the response + field; dynamo emits ``response.nvext.prompt_logprobs`` shaped as + ``[None | {token_id: {logprob, rank?, decoded_token?}}]``, which the + shared flattener consumes unchanged. + + Required engine adapter support: vLLM worker must populate + ``LLMEngineOutput.prompt_logprobs`` when ``SamplingParams.prompt_logprobs`` + is set. Without that (A10), the response payload is None. + """ + client = setup_openai_client(client_config) + token_ids = list(sample.prompt_ids) + list(sample.completion_ids) + body = { + "model": model_name, + "messages": [], + "max_completion_tokens": 1, + "temperature": 1.0, + "top_p": 1.0, + "prompt_logprobs": 1, + "nvext": { + "token_data": token_ids, + "extra_fields": ["prompt_logprobs"], + }, + } + # Dynamo's response is just a standard chat-completion JSON with an extra + # ``nvext`` field. Use ``cast_to=httpx.Response`` so we can read the raw + # body and pluck ``nvext.prompt_logprobs`` ourselves — the OpenAI SDK + # response models drop unknown fields. + import httpx as _httpx + + http_response = await client.post( + "/chat/completions", + cast_to=_httpx.Response, + body=body, + ) + payload = http_response.json() + nvext_resp = (payload or {}).get("nvext") or {} + raw = nvext_resp.get("prompt_logprobs") + return _flatten_prompt_logprobs(raw) + + async def compute_teacher_logprobs( clients: list[vf.ClientConfig], model_name: str, samples: list[TrainingSample], ) -> list[list[float]]: - """Compute teacher model logprobs for a batch of training samples via prefill.""" - import httpx - from vllm.entrypoints.serve.disagg.protocol import GenerateResponse + """Compute teacher model logprobs for a batch of training samples via prefill. + + Dispatches to the vLLM-sidecar or dynamo-nvext path based on the + per-client ``renderer_transport``: + + - ``vllm_generate`` (default): POST ``/inference/v1/generate`` + - ``dynamo_chat`` : POST ``/v1/chat/completions`` with nvext + + Both flatten to ``list[float]`` via the shared helper. + """ async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample) -> list[float]: - client = setup_openai_client(client_config) - - # Two escape hatches from ``AsyncOpenAI.post``: - # 1. URL — ``/inference/v1/generate`` is mounted at server root, not - # under ``/v1``. Pass an absolute URL so the SDK's - # ``_prepare_url`` skips the base-url merge (it short-circuits - # when the path passes ``httpx.URL.is_relative_url`` as False). - # 2. Parse — vLLM's ``GenerateResponse`` is a plain - # ``pydantic.BaseModel`` and the SDK's parse layer rejects any - # ``cast_to`` that doesn't subclass ``openai.BaseModel``. Use - # ``cast_to=httpx.Response`` so the SDK still builds the request - # (preserving ``auth_headers``, retries, timeouts, idempotency - # keys) and just hands us the raw response to validate ourselves. - base = str(client.base_url).rstrip("/").removesuffix("/v1") - http_response = await client.post( - f"{base}/inference/v1/generate", - cast_to=httpx.Response, - body={ - "model": model_name, - "token_ids": list(sample.prompt_ids) + list(sample.completion_ids), - "sampling_params": { - "max_tokens": 1, - "temperature": 1.0, - "top_p": 1.0, - "prompt_logprobs": 1, - }, - }, - ) - response = GenerateResponse.model_validate_json(http_response.content) - # ``prompt_logprobs[i]`` is a ``{token_id: Logprob}`` dict for tokens - # the engine could score, or ``None`` for the leading token which has - # no preceding context. Flatten to ``list[float]`` with 0.0 in the - # unscored slot. - flat: list[float] = [] - for entry in response.prompt_logprobs or []: - if not entry: - flat.append(0.0) - continue - first = next(iter(entry.values())) - lp = first.logprob if hasattr(first, "logprob") else first.get("logprob") - flat.append(float(lp) if lp is not None else 0.0) - return flat + if client_config.renderer_transport == "dynamo_chat": + return await _compute_teacher_logprobs_dynamo(client_config, model_name, sample) + return await _compute_teacher_logprobs_vllm(client_config, model_name, sample) return await asyncio.gather(*[_compute_single(client, sample) for client, sample in zip(cycle(clients), samples)]) diff --git a/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py index 85025821a1..ea5cf3c6f8 100644 --- a/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py @@ -7,6 +7,18 @@ def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: return max(int(i.split(".")[2]) for i in state_dict.keys() if "model.layers." in i) + 1 +def _expert_indices(state_dict: dict[str, Tensor], layer_idx: int) -> list[int]: + prefix = f"model.layers.{layer_idx}.mlp.experts." + indices: set[int] = set() + for key in state_dict: + if not key.startswith(prefix): + continue + candidate = key[len(prefix) :].split(".", 1)[0] + if candidate.isdigit(): + indices.add(int(candidate)) + return sorted(indices) + + def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): i = layer_idx @@ -35,24 +47,29 @@ def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): del state_dict[f"model.layers.{i}.mlp.experts.down_proj"] else: # Old per-expert format - num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3 + expert_indices = _expert_indices(state_dict, i) + num_experts = len(expert_indices) if num_experts == 0: return - dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape + first_expert = expert_indices[0] + dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].shape w1 = torch.empty( - (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype + (num_experts, moe_dim, dim), + dtype=state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype, ) # Gate w2 = torch.empty( - (num_experts, dim, moe_dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype + (num_experts, dim, moe_dim), + dtype=state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype, ) # Down w3 = torch.empty( - (num_experts, moe_dim, dim), dtype=state_dict[f"model.layers.{i}.mlp.experts.1.down_proj.weight"].dtype + (num_experts, moe_dim, dim), + dtype=state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype, ) # Up - for j in range(num_experts): - w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) - w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) - w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) + for expert_pos, j in enumerate(expert_indices): + w1[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) + w2[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) + w3[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py index e8e5b311b8..4d31a3cabc 100644 --- a/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py +++ b/src/prime_rl/trainer/models/glm_moe_dsa/converting_glm_moe_dsa.py @@ -13,6 +13,18 @@ def _is_moe_layer(state_dict: dict[str, Tensor], layer_idx: int) -> bool: return f"model.layers.{layer_idx}.mlp.gate.weight" in state_dict +def _expert_indices(state_dict: dict[str, Tensor], layer_idx: int) -> list[int]: + prefix = f"model.layers.{layer_idx}.mlp.experts." + indices: set[int] = set() + for key in state_dict: + if not key.startswith(prefix): + continue + candidate = key[len(prefix) :].split(".", 1)[0] + if candidate.isdigit(): + indices.add(int(candidate)) + return sorted(indices) + + def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): i = layer_idx @@ -38,19 +50,21 @@ def convert_hf_layer_to_tt(state_dict: dict[str, Tensor], layer_idx: int): del state_dict[f"model.layers.{i}.mlp.experts.gate_up_proj"] del state_dict[f"model.layers.{i}.mlp.experts.down_proj"] else: - num_experts = len([j for j in state_dict.keys() if f"model.layers.{i}.mlp.experts" in j]) // 3 + expert_indices = _expert_indices(state_dict, i) + num_experts = len(expert_indices) if num_experts == 0: return - dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].shape - dtype = state_dict[f"model.layers.{i}.mlp.experts.0.down_proj.weight"].dtype + first_expert = expert_indices[0] + dim, moe_dim = state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].shape + dtype = state_dict[f"model.layers.{i}.mlp.experts.{first_expert}.down_proj.weight"].dtype w1 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) w2 = torch.empty((num_experts, dim, moe_dim), dtype=dtype) w3 = torch.empty((num_experts, moe_dim, dim), dtype=dtype) - for j in range(num_experts): - w1[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) - w2[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) - w3[j].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) + for expert_pos, j in enumerate(expert_indices): + w1[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"]) + w2[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"]) + w3[expert_pos].copy_(state_dict[f"model.layers.{i}.mlp.experts.{j}.up_proj.weight"]) del state_dict[f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight"] del state_dict[f"model.layers.{i}.mlp.experts.{j}.down_proj.weight"] diff --git a/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py index 48ccc5b941..05362de1d3 100644 --- a/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py +++ b/src/prime_rl/trainer/models/qwen3_moe/converting_qwen3_moe.py @@ -1,6 +1,8 @@ import torch from torch import Tensor +from prime_rl.trainer.models.fp8 import quantize_to_fp8_blockwise + def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: """Get the maximum number of layers in the model.""" @@ -98,3 +100,104 @@ def convert_tt_to_hf_moe(state_dict: dict[str, Tensor]): num_layers = get_max_layer_num(state_dict) for i in range(num_layers): convert_tt_layer_to_hf(state_dict, i) + + +def _emit_weight(out: dict[str, Tensor], name: str, tensor: Tensor, quantize_fp8: bool) -> None: + """Emit a 2D projection weight, optionally FP8-quantized with block scales.""" + if quantize_fp8: + fp8_weight, scale = quantize_to_fp8_blockwise(tensor) + out[name] = fp8_weight + out[name.removesuffix(".weight") + ".weight_scale_inv"] = scale + else: + out[name] = tensor + + +def _emit_moe_experts( + out: dict[str, Tensor], + prefix: str, + w1: Tensor, + w2: Tensor, + w3: Tensor, + quantize_fp8: bool, +) -> None: + """Emit fused MoE expert weights in vLLM FusedMoE W13 layout. + + vLLM's DEEPGEMM FP8 MoE path keeps weights in `[gate; up]` (W13) order; + FLASHINFER paths flip to W31 during `process_weights_after_loading`, so + inference must select the DEEPGEMM backend (`VLLM_USE_DEEP_GEMM=1`) for + RL weight reloads to stay valid across broadcasts. + """ + w13 = torch.cat([w1, w3], dim=1) + if not quantize_fp8: + out[f"{prefix}.mlp.experts.w13_weight"] = w13 + out[f"{prefix}.mlp.experts.w2_weight"] = w2 + return + + num_experts = w1.shape[0] + w13_fp8: list[Tensor] = [] + w13_scales: list[Tensor] = [] + w2_fp8: list[Tensor] = [] + w2_scales: list[Tensor] = [] + for expert_idx in range(num_experts): + q13, s13 = quantize_to_fp8_blockwise(w13[expert_idx]) + q2, s2 = quantize_to_fp8_blockwise(w2[expert_idx]) + w13_fp8.append(q13) + w13_scales.append(s13) + w2_fp8.append(q2) + w2_scales.append(s2) + + out[f"{prefix}.mlp.experts.w13_weight"] = torch.stack(w13_fp8) + out[f"{prefix}.mlp.experts.w13_weight_scale_inv"] = torch.stack(w13_scales) + out[f"{prefix}.mlp.experts.w2_weight"] = torch.stack(w2_fp8) + out[f"{prefix}.mlp.experts.w2_weight_scale_inv"] = torch.stack(w2_scales) + + +def _is_dense_layer(state_dict: dict[str, Tensor], layer_idx: int) -> bool: + """Dense MLP layers have a fused gate_proj; MoE layers route through mlp.router.""" + return f"model.layers.{layer_idx}.mlp.gate_proj.weight" in state_dict + + +def convert_tt_layer_to_vllm_kernel( + state_dict: dict[str, Tensor], + layer_idx: int, + quantize_fp8: bool = False, +) -> dict[str, Tensor]: + """Convert a single Qwen3MoE layer from PrimeRL format to vLLM kernel format.""" + prefix = f"model.layers.{layer_idx}" + out: dict[str, Tensor] = { + f"{prefix}.input_layernorm.weight": state_dict[f"{prefix}.input_layernorm.weight"], + f"{prefix}.post_attention_layernorm.weight": state_dict[f"{prefix}.post_attention_layernorm.weight"], + f"{prefix}.self_attn.q_norm.weight": state_dict[f"{prefix}.self_attn.q_norm.weight"], + f"{prefix}.self_attn.k_norm.weight": state_dict[f"{prefix}.self_attn.k_norm.weight"], + } + + qkv = torch.cat( + [ + state_dict[f"{prefix}.self_attn.q_proj.weight"], + state_dict[f"{prefix}.self_attn.k_proj.weight"], + state_dict[f"{prefix}.self_attn.v_proj.weight"], + ], + dim=0, + ) + _emit_weight(out, f"{prefix}.self_attn.qkv_proj.weight", qkv, quantize_fp8) + _emit_weight(out, f"{prefix}.self_attn.o_proj.weight", state_dict[f"{prefix}.self_attn.o_proj.weight"], quantize_fp8) + + if _is_dense_layer(state_dict, layer_idx): + gate_up = torch.cat( + [state_dict[f"{prefix}.mlp.gate_proj.weight"], state_dict[f"{prefix}.mlp.up_proj.weight"]], + dim=0, + ) + _emit_weight(out, f"{prefix}.mlp.gate_up_proj.weight", gate_up, quantize_fp8) + _emit_weight(out, f"{prefix}.mlp.down_proj.weight", state_dict[f"{prefix}.mlp.down_proj.weight"], quantize_fp8) + else: + out[f"{prefix}.mlp.gate.weight"] = state_dict[f"{prefix}.mlp.router.gate.weight"] + _emit_moe_experts( + out, + prefix, + state_dict[f"{prefix}.mlp.experts.w1"], + state_dict[f"{prefix}.mlp.experts.w2"], + state_dict[f"{prefix}.mlp.experts.w3"], + quantize_fp8, + ) + + return out diff --git a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py index 25d83f9436..59c18efcc6 100644 --- a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py @@ -38,6 +38,7 @@ convert_hf_layer_to_tt, convert_hf_to_tt_moe, convert_tt_layer_to_hf, + convert_tt_layer_to_vllm_kernel, convert_tt_to_hf_moe, ) from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids @@ -170,6 +171,13 @@ def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) - convert_hf_layer_to_tt(state_dict, layer_idx) return state_dict + @classmethod + def convert_layer_to_vllm_kernel( + cls, state_dict: dict[str, Tensor], layer_idx: int, quantize_fp8: bool = False + ) -> dict[str, Tensor]: + """Convert a single layer's weights from PrimeRL format to vLLM FusedMoE/linear kernel format.""" + return convert_tt_layer_to_vllm_kernel(state_dict, layer_idx, quantize_fp8=quantize_fp8) + @auto_docstring class Qwen3MoeModel(Qwen3MoePreTrainedModel): diff --git a/src/prime_rl/trainer/rl/broadcast/filesystem.py b/src/prime_rl/trainer/rl/broadcast/filesystem.py index e8e2d68db9..35a8a0c89a 100644 --- a/src/prime_rl/trainer/rl/broadcast/filesystem.py +++ b/src/prime_rl/trainer/rl/broadcast/filesystem.py @@ -1,3 +1,4 @@ +import os import shutil import time from pathlib import Path @@ -104,14 +105,43 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: self.logger.debug(f"Weights broadcasted in {time.perf_counter() - start_time:.2f}s") def _notify_orchestrator(self, save_dir: Path): - """Notify the orchestrator that the weights have been broadcast by writing a 'STABLE' file to a shared filesystem.""" - stable_file = save_dir / "STABLE" - stable_file.touch() - - def maybe_clean(self, interval_to_keep: int | None): + """Notify the orchestrator that the weights have been broadcast by writing a 'STABLE' file to a shared filesystem. + + On shared filesystems (NFS / CephFS / PVC-backed volumes), a write on one + node is not immediately visible on another. Without explicit flushing, the + orchestrator can see the STABLE sentinel before the adapter files + (``adapter_model.safetensors``, ``adapter_config.json``) are visible to + the inference worker — causing ``LoRAAdapterNotFoundError`` when + ``/v1/rl/load_lora_adapter`` is called. We fsync the adapter files and + the directory's dentries before touching STABLE, then fsync the + directory again so STABLE itself is durable. + """ + for f in save_dir.iterdir(): + if f.is_file(): + _fsync_path(f) + _fsync_path(save_dir) + (save_dir / "STABLE").touch() + _fsync_path(save_dir) + + def maybe_clean(self, retention: int, interval_to_keep: int | None): for idx in self.multi_run_manager.used_idxs: maybe_clean( get_broadcast_dir(self.multi_run_manager.get_run_dir(idx)), self.multi_run_manager.progress[idx].step, + retention, interval_to_keep, ) + + +def _fsync_path(path: Path) -> None: + """Best-effort fsync of a file or directory. Silent on OSError because the + sync is an availability-of-data hint to the kernel, not a correctness gate + on its own — STABLE ordering is what matters.""" + try: + fd = os.open(str(path), os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + except OSError: + pass diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py index 13a887308a..8412e5dd80 100644 --- a/src/prime_rl/trainer/rl/broadcast/nccl.py +++ b/src/prime_rl/trainer/rl/broadcast/nccl.py @@ -4,6 +4,7 @@ from typing import Callable, Generator, cast import torch +import torch.distributed as dist import torch.nn as nn from torch import Tensor from torch.distributed.tensor import DTensor @@ -156,10 +157,21 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: preprocess_fn = preprocess_layer_checkpoint for layer_id, layer_state_dict in filter_state_dict_by_layers(state_dict, num_layers, layer_prefix): + # _resolve_dtensors triggers DTensor.full_tensor() all_gathers on ALL trainer + # ranks (FSDP/EP collectives). Every rank must participate. layer_state_dict = self._resolve_dtensors(layer_state_dict) layer_state_dict = preprocess_fn(model, layer_state_dict, layer_id) if self.world.is_master: broadcast_state_dict(layer_state_dict, self.communicator) + # Ensure the inference NCCL sends complete on the master's CUDA stream + # before we release non-master ranks to start the next layer's all_gather. + torch.cuda.synchronize() + # Barrier: prevents ranks 1-N from advancing to the next layer's + # _resolve_dtensors (and enqueuing the next FSDP/EP all_gather) before + # rank 0 finishes broadcasting the current layer to inference workers. + # Without this, ranks 1-N enqueue unmatched all_gathers and the NCCL + # watchdog fires after 600 s, killing all trainer processes. + dist.barrier() def _resolve_dtensors(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: for key, value in list(state_dict.items()): diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 83afa666dc..dcf2aecb98 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -279,7 +279,14 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: # Clean up old broadcast directories (unless at ckpt interval if using filesystem weight broadcast) if config.weight_broadcast.type == "filesystem": interval_to_keep = config.ckpt and config.ckpt.interval - weight_broadcast.maybe_clean(interval_to_keep) + # Fallback retention when keep_recent unset: 2 (matches main's + # original step-2 cleanup; trainer no longer carries max_async_level). + retention = ( + config.weight_broadcast.keep_recent + if config.weight_broadcast.keep_recent is not None + else 2 + ) + weight_broadcast.maybe_clean(retention, interval_to_keep) else: broadcast_weights_time = 0 # Usually the broadcast will set this. If broadcast is skipped, we need to reset this here. diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 9a8a86a526..10aebd0bce 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -450,18 +450,18 @@ def step(self): self.step_num += 1 -def maybe_clean(path: Path, step: int, interval_to_keep: int | None) -> None: - """Delete the broadcast dir from 2 trainer steps ago. +def maybe_clean(path: Path, step: int, retention: int, interval_to_keep: int | None) -> None: + """Delete the broadcast step directory that falls outside the retention window. - With a 1-step async barrier, the orchestrator at trainer step ``step`` is still consuming the - ckpt from ``step - 1``; ``step - 2`` is therefore safe to remove unless it falls on a - checkpoint interval that we want to preserve. + At trainer step N, deletes step_{N - retention - 1}. Caller picks whether + `retention` is the orchestrator's staleness bound (max_async_level) or a + wider buffer (weight_broadcast.keep_recent). """ logger = get_logger() - candidate_step = max(step - 2, 0) - candidate_path = get_step_path(path, candidate_step) - if interval_to_keep and candidate_step % interval_to_keep == 0: - logger.debug(f"Keeping path {candidate_path} (on ckpt interval)") - return - logger.debug(f"Removing path {candidate_path}") - shutil.rmtree(candidate_path, ignore_errors=True) + step = max(step - (retention + 1), 0) + candidate_path_to_delete = get_step_path(path, step) + keep = bool(interval_to_keep and step % interval_to_keep == 0) + logger.debug(f"Considering deleting path {candidate_path_to_delete}") + if not keep: + logger.debug(f"Removing path {candidate_path_to_delete}") + shutil.rmtree(candidate_path_to_delete, ignore_errors=True) diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index b9ee8f4b9d..fd608a92ea 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -6,6 +6,7 @@ from itertools import cycle from pathlib import Path from typing import Protocol, runtime_checkable +from urllib.parse import urlsplit, urlunsplit import httpx import verifiers as vf @@ -28,6 +29,237 @@ def client_identity(client: vf.ClientConfig) -> ClientIdentity: return (client.api_base_url, client.extra_headers.get("X-data-parallel-rank")) +class AdminAPI(Protocol): + """Admin endpoints for an inference backend. + + Per-method: construct one HTTP call. Per-server parallelism, retry, and + raise-for-status policy live in the caller. + """ + + async def health(self, client: AsyncClient) -> None: ... + async def list_models(self, client: AsyncClient) -> list[dict]: ... + async def pause(self, client: AsyncClient) -> None: ... + async def resume(self, client: AsyncClient) -> None: ... + async def update_weights(self, client: AsyncClient, weight_dir: str | None) -> None: ... + async def load_lora_adapter( + self, + client: AsyncClient, + lora_name: str, + lora_path: str, + *, + timeout: httpx.Timeout, + ) -> None: ... + async def init_broadcaster( + self, + client: AsyncClient, + *, + host: str, + port: int, + rank_offset: int, + inference_world_size: int, + timeout: int, + quantize_in_weight_transfer: bool, + ) -> None: ... + + +class VLLMAdminAPI: + """vLLM admin endpoints.""" + + async def health(self, client: AsyncClient) -> None: + # No raise_for_status: any HTTP response means the server is up. + # Only transport errors mean "not ready yet" (caller retries). + await client.get("/health") + + async def list_models(self, client: AsyncClient) -> list[dict]: + response = await client.get("/v1/models") + return response.json()["data"] + + async def pause(self, client: AsyncClient) -> None: + response = await client.post("/pause", params={"mode": "keep", "clear_cache": "false"}) + response.raise_for_status() + + async def resume(self, client: AsyncClient) -> None: + response = await client.post("/resume") + response.raise_for_status() + + async def update_weights(self, client: AsyncClient, weight_dir: str | None) -> None: + response = await client.post("/update_weights", json={"weight_dir": weight_dir}) + response.raise_for_status() + + async def load_lora_adapter( + self, + client: AsyncClient, + lora_name: str, + lora_path: str, + *, + timeout: httpx.Timeout, + ) -> None: + response = await client.post( + "/load_lora_adapter", + json={"lora_name": lora_name, "lora_path": lora_path}, + timeout=timeout, + ) + response.raise_for_status() + + async def init_broadcaster( + self, + client: AsyncClient, + *, + host: str, + port: int, + rank_offset: int, + inference_world_size: int, + timeout: int, + quantize_in_weight_transfer: bool, + ) -> None: + response = await client.post( + "/init_broadcaster", + json={ + "host": host, + "port": port, + "rank_offset": rank_offset, + "inference_world_size": inference_world_size, + "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, + }, + ) + response.raise_for_status() + + +class DynamoAdminAPI(VLLMAdminAPI): + """NVIDIA Dynamo worker admin endpoints via ``POST /engine/``. + + Each Dynamo worker exposes engine routes on its system status server + (``DYN_SYSTEM_PORT``, default 8081). Multi-worker deployments are handled by + iterating over ``admin_clients``. + + Args: + engine_rpc: The ``collective_rpc`` target forwarded by + ``update_weights_from_disk``. Use ``"reload_weights"`` for plain + vLLM / dynamo.vllm without a worker extension (default). Use + ``"update_weights_from_path"`` only when + FileSystemWeightUpdateWorker / NCCLWeightUpdateWorker is loaded via + ``--worker-extension-cls``. + """ + + def __init__(self, engine_rpc: str = "reload_weights", weight_broadcast_type: str = "filesystem") -> None: + self._engine_rpc = engine_rpc + # Determines which engine method is called per step: "update_weights_from_distributed" + # for NCCL (trainer broadcasts; worker just needs to receive) vs + # "update_weights_from_disk" for filesystem. Set externally by the orchestrator + # once weight_broadcast config is resolved. Defaults to filesystem (run #35 behaviour). + self._weight_broadcast_type = weight_broadcast_type + + async def health(self, client: AsyncClient) -> None: + await client.get("/health") + + async def _post_engine( + self, + client: AsyncClient, + method: str, + body: dict | None = None, + *, + timeout: httpx.Timeout | None = None, + ) -> dict: + response = await client.post(f"/engine/{method}", json=body or {}, timeout=timeout) + response.raise_for_status() + data = response.json() + if isinstance(data, dict) and data.get("status") == "error": + raise RuntimeError(data.get("message", f"Dynamo /engine/{method} failed")) + return data + + async def pause(self, client: AsyncClient) -> None: + await self._post_engine(client, "pause_generation", {"mode": "keep", "clear_cache": False}) + + async def resume(self, client: AsyncClient) -> None: + await self._post_engine(client, "resume_generation") + + async def update_weights(self, client: AsyncClient, weight_dir: str | None) -> None: + if weight_dir is None: + return + if self._weight_broadcast_type == "nccl": + # NCCL path: trainer has already broadcast weights via the NCCL group; + # this RPC tells the inference worker to call receive_state_dict(). + # NCCLWeightUpdateWorker exposes "update_weights_from_path", not "reload_weights". + await self._post_engine( + client, + "update_weights_from_distributed", + { + "weight_version": Path(weight_dir).name, + "weight_dir": weight_dir, + "engine_rpc": "update_weights_from_path", + }, + timeout=httpx.Timeout(180.0), + ) + else: + # Resolve to absolute path so the inference worker (which may run in a + # different working directory) can find the checkpoint on the shared NFS. + abs_path = str(Path(weight_dir).resolve()) + await self._post_engine( + client, + "update_weights_from_disk", + { + "model_path": abs_path, + "weight_version": Path(weight_dir).name, + "engine_rpc": self._engine_rpc, + }, + timeout=httpx.Timeout(180.0), + ) + + async def load_lora_adapter( + self, + client: AsyncClient, + lora_name: str, + lora_path: str, + *, + timeout: httpx.Timeout, + ) -> None: + await self._post_engine( + client, + "load_lora", + { + "lora_name": lora_name, + "source": {"uri": Path(lora_path).absolute().as_uri()}, + }, + timeout=timeout, + ) + + async def init_broadcaster( + self, + client: AsyncClient, + *, + host: str, + port: int, + rank_offset: int, + inference_world_size: int, + timeout: int, + quantize_in_weight_transfer: bool, + ) -> None: + await self._post_engine( + client, + "init_weights_update_group", + { + "host": host, + "port": port, + "rank_offset": rank_offset, + "inference_world_size": inference_world_size, + "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, + "engine_rpc": "init_broadcaster", + }, + ) + + +def setup_admin_api(client_config: ClientConfig) -> AdminAPI: + """Pick the AdminAPI implementation that matches ``client_config.backend``.""" + if client_config.backend == "dynamo": + return DynamoAdminAPI() + return VLLMAdminAPI() + + +_DEFAULT_ADMIN: AdminAPI = VLLMAdminAPI() + + @runtime_checkable class InferencePool(Protocol): """Protocol for inference pools (static or elastic).""" @@ -103,6 +335,12 @@ def __init__( ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) self._admin_clients = setup_admin_clients(client_config) + self._model_clients = ( + setup_admin_clients(client_config, use_admin_base_url=False) + if client_config.backend == "dynamo" or client_config.admin_base_url + else self._admin_clients + ) + self._admin_api = setup_admin_api(client_config) self._skip_model_check = client_config.skip_model_check self._wait_for_ready_timeout = client_config.wait_for_ready_timeout self._eval_cycle = cycle(self._eval_clients) @@ -133,12 +371,16 @@ async def select_train_client(self, load: Mapping[ClientIdentity, int]) -> vf.Cl async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: await check_health( - self._admin_clients, timeout=timeout if timeout is not None else self._wait_for_ready_timeout + self._admin_clients, + timeout=timeout if timeout is not None else self._wait_for_ready_timeout, + admin=self._admin_api, + ) + await maybe_check_has_model( + self._model_clients, model_name, skip_model_check=self._skip_model_check, admin=self._admin_api ) - await maybe_check_has_model(self._admin_clients, model_name, skip_model_check=self._skip_model_check) async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: - await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step) + await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step, admin=self._admin_api) def get_metrics(self) -> dict[str, float]: return {} @@ -185,6 +427,13 @@ def setup_clients( renderer_model_name: str | None = None, pool_size: int | None = None, ) -> list[vf.ClientConfig]: + # Pick the verifiers wire-shape selector based on client_config.backend. + # When backend == "dynamo", both RendererClient and + # OpenAIChatCompletionsTokenClient route through Dynamo's nvext path: + # - request: nvext.token_data carries pre-tokenized prompt + # - response: nvext.engine_data carries completion_token_ids + logprobs + # Default backend keeps the legacy vLLM TITO surface. + renderer_transport = "dynamo_chat" if client_config.backend == "dynamo" else "vllm_generate" clients = [] client_idx = 0 # Only forward the renderer config when the client actually uses a @@ -208,6 +457,9 @@ def setup_clients( vf.ClientConfig( client_idx=client_idx, client_type=client_type, + # Dynamo backend routes both renderer and token clients through + # the nvext path; default backend keeps the legacy vLLM TITO surface. + renderer_transport=renderer_transport, api_base_url=base_url, api_key_var=client_config.api_key_var, timeout=client_config.timeout, @@ -224,14 +476,21 @@ def setup_clients( return clients -def setup_admin_clients(client_config: ClientConfig) -> list[AsyncClient]: +def setup_admin_clients(client_config: ClientConfig, *, use_admin_base_url: bool = True) -> list[AsyncClient]: """Create dedicated admin clients for weight update operations. Uses a separate connection pool to avoid queueing behind streaming requests. - When admin_base_url is set, uses those URLs instead of base_url, allowing - weight updates to bypass routers in disaggregated P/D deployments. + When admin_base_url is set and use_admin_base_url is true, uses those URLs + instead of base_url, allowing weight updates to bypass routers in + disaggregated P/D deployments. For Dynamo, if admin_base_url is unset, + discover worker-advertised system URLs from GET /v1/rl/workers. """ - urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url + if use_admin_base_url and client_config.admin_base_url: + urls = client_config.admin_base_url + elif use_admin_base_url and client_config.backend == "dynamo": + urls = discover_dynamo_admin_base_urls(client_config) + else: + urls = client_config.base_url def _setup_admin_client(base_url: str) -> httpx.AsyncClient: env_headers = { @@ -255,23 +514,84 @@ def _setup_admin_client(base_url: str) -> httpx.AsyncClient: return [_setup_admin_client(base_url) for base_url in urls] +def discover_dynamo_admin_base_urls(client_config: ClientConfig) -> list[str]: + urls: list[str] = [] + headers = client_config.headers.copy() + api_key = os.getenv(client_config.api_key_var, "EMPTY") + if api_key and api_key != "EMPTY": + headers["Authorization"] = f"Bearer {api_key}" + + for base_url in _dynamo_rl_discovery_base_urls(client_config): + discovery_base = base_url.rstrip("/").removesuffix("/v1") + with httpx.Client( + base_url=discovery_base, + headers=headers, + timeout=httpx.Timeout(connect=client_config.connect_timeout, read=30.0, write=30.0, pool=10.0), + ) as client: + response = client.get("/v1/rl/workers") + response.raise_for_status() + for worker in response.json().get("workers", []): + system_url = worker.get("system_url") + if system_url: + urls.append(system_url) + + deduped = list(dict.fromkeys(urls)) + if not deduped: + raise ValueError( + "Dynamo backend did not discover any worker system URLs from /v1/rl/workers. " + "Set client.admin_base_url explicitly, set client.rl_base_url to the Dynamo " + "RL discovery listener, and make sure Dynamo workers run with DYN_ENABLE_RL " + "and a system status server enabled." + ) + return deduped + + +def _dynamo_rl_discovery_base_urls(client_config: ClientConfig) -> list[str]: + configured = getattr(client_config, "rl_base_url", None) + if configured: + return configured + + rl_port = int(os.getenv("DYN_RL_PORT", "8001")) + return [_replace_url_port(base_url, rl_port) for base_url in client_config.base_url] + + +def _replace_url_port(base_url: str, port: int) -> str: + parsed = urlsplit(base_url.rstrip("/").removesuffix("/v1")) + scheme = parsed.scheme or "http" + host = parsed.hostname or parsed.netloc + if not host: + raise ValueError(f"Cannot derive Dynamo RL discovery URL from base_url={base_url!r}") + if ":" in host and not host.startswith("["): + host = f"[{host}]" + netloc = f"{host}:{port}" + return urlunsplit((scheme, netloc, "", "", "")) + + async def maybe_check_has_model( - admin_clients: list[AsyncClient], model_name: str, skip_model_check: bool = False + admin_clients: list[AsyncClient], + model_name: str, + skip_model_check: bool = False, + *, + admin: AdminAPI = _DEFAULT_ADMIN, ) -> None: if skip_model_check: return logger = get_logger() logger.debug(f"Checking if model {model_name} is in the inference pool") - results = await asyncio.gather(*[admin_client.get("/v1/models") for admin_client in admin_clients]) - for admin_client, result in zip(admin_clients, results): - models = result.json()["data"] + results = await asyncio.gather(*[admin.list_models(admin_client) for admin_client in admin_clients]) + for admin_client, models in zip(admin_clients, results): if not any(model["id"] == model_name for model in models): raise ValueError(f"Model {model_name} was not found in the inference pool on {admin_client.base_url}") logger.debug(f"Model {model_name} was found in the inference pool") async def check_health( - admin_clients: list[AsyncClient], interval: int = 1, log_interval: int = 10, timeout: int = 1800 + admin_clients: list[AsyncClient], + interval: int = 1, + log_interval: int = 10, + timeout: int = 1800, + *, + admin: AdminAPI = _DEFAULT_ADMIN, ) -> None: logger = get_logger() @@ -280,7 +600,7 @@ async def _check_health(admin_client: AsyncClient) -> None: logger.debug("Starting pinging /health to check health") while wait_time < timeout: try: - await admin_client.get("/health") + await admin.health(admin_client) logger.debug(f"Inference pool is ready after {wait_time} seconds") return except NotFoundError: @@ -303,105 +623,44 @@ async def _check_health(admin_client: AsyncClient) -> None: NCCL_READY_MARKER = "NCCL_READY" -def _is_retryable_pause_error(exception: BaseException) -> bool: - """Check if an exception should trigger a retry for pausing engines.""" - if isinstance(exception, httpx.HTTPStatusError): - # Retry on transient server errors (5xx, e.g. engine briefly unresponsive); - # client errors (4xx) won't fix themselves on retry. - return exception.response.status_code >= 500 - # Retry on transport-level failures (timeouts, connection resets, etc.) so the - # per-attempt read timeout below turns a stuck server into a bounded retry loop - # instead of hanging forever on the global timeout=None admin client. - if isinstance(exception, (httpx.TimeoutException, httpx.TransportError)): - return True - return False - - -# Per-attempt and total bounds for `/pause`. Pausing drains in-flight requests -# (mode="keep"), so a single attempt can legitimately take a while, but the global -# admin AsyncClient uses `timeout=None`, so a stuck server would hang the weight -# update forever. `_READ_TIMEOUT` converts a hang into a TimeoutException so -# tenacity retries; `_TOTAL` is the wall-clock budget across all retries. -PAUSE_READ_TIMEOUT_S = 120.0 -PAUSE_TOTAL_TIMEOUT_S = 300.0 - - -async def _pause_engines(admin_clients: list[AsyncClient], *, step: int) -> None: - """Pause all inference engines, waiting for in-flight requests to drain.""" - logger = get_logger() - logger.info(f"Updating policy in-flight to v{step}") - - @retry( - retry=retry_if_exception(_is_retryable_pause_error), - stop=stop_after_delay(PAUSE_TOTAL_TIMEOUT_S) | stop_after_attempt(10), - wait=wait_exponential(multiplier=1, min=1, max=10), - reraise=True, - ) - async def _pause(client: AsyncClient) -> None: - response = await client.post( - "/pause", - params={"mode": "keep", "clear_cache": "false"}, - timeout=httpx.Timeout(connect=10.0, read=PAUSE_READ_TIMEOUT_S, write=60.0, pool=10.0), - ) - response.raise_for_status() - - await asyncio.gather(*[_pause(client) for client in admin_clients]) - logger.debug("All inference engines paused") - - -async def _resume_engines(admin_clients: list[AsyncClient]) -> None: - """Resume all inference engines after weight update.""" - logger = get_logger() - - async def _resume(client: AsyncClient) -> None: - response = await client.post("/resume") - response.raise_for_status() - - await asyncio.gather(*[_resume(client) for client in admin_clients]) - logger.debug("All inference engines resumed") - - async def update_weights( admin_clients: list[AsyncClient], weight_dir: Path | None, lora_name: str | None = None, step: int = 0, + *, + admin: AdminAPI = _DEFAULT_ADMIN, ) -> None: """Update weights on static inference servers. - Pauses all engines first to drain in-flight requests, then performs the - weight update, then resumes. This ensures all DP workers are idle and can - participate in the collective weight transfer. - - Note: The server-side /update_weights endpoint automatically resets the prefix cache - to invalidate any cached KV states computed with the old weights. + Pauses all engines to drain in-flight requests, performs the weight update, + then resumes. Ensures all DP workers are idle and can participate in the + collective weight transfer. The server-side ``/update_weights`` endpoint + resets the prefix cache to invalidate any KV states computed with the old + weights. """ logger = get_logger() - weight_dir_posix = weight_dir.as_posix() if weight_dir is not None else None - if lora_name is not None and weight_dir is not None: - await load_lora_adapter(admin_clients, lora_name, weight_dir) - else: - - async def _update_weights(admin_client: AsyncClient, weight_dir: str | None) -> None: - response = await admin_client.post("/update_weights", json={"weight_dir": weight_dir}) - response.raise_for_status() + await load_lora_adapter(admin_clients, lora_name, weight_dir, admin=admin) + return - # Pause engines so all DP workers drain in-flight work and can join the NCCL broadcast - await _pause_engines(admin_clients, step=step) + weight_dir_posix = weight_dir.as_posix() if weight_dir is not None else None - try: - # Create ready marker before servers enter receive path (used by NCCL broadcast) - if weight_dir is not None: - nccl_ready_file = weight_dir / NCCL_READY_MARKER - nccl_ready_file.parent.mkdir(parents=True, exist_ok=True) - nccl_ready_file.touch() - logger.debug(f"Created NCCL_READY marker at {nccl_ready_file}") + logger.info("Pausing inference engines for weight update") + await asyncio.gather(*[admin.pause(c) for c in admin_clients]) + try: + # NCCL_READY marker is created before servers enter the receive path + if weight_dir is not None: + nccl_ready_file = weight_dir / NCCL_READY_MARKER + nccl_ready_file.parent.mkdir(parents=True, exist_ok=True) + nccl_ready_file.touch() + logger.debug(f"Created NCCL_READY marker at {nccl_ready_file}") - await asyncio.gather(*[_update_weights(admin_client, weight_dir_posix) for admin_client in admin_clients]) - finally: - await _resume_engines(admin_clients) + await asyncio.gather(*[admin.update_weights(c, weight_dir_posix) for c in admin_clients]) + finally: + await asyncio.gather(*[admin.resume(c) for c in admin_clients]) + logger.info("Inference engines resumed") def _is_retryable_lora_error(exception: BaseException) -> bool: @@ -428,7 +687,13 @@ def _is_retryable_lora_error(exception: BaseException) -> bool: LORA_LOAD_TOTAL_TIMEOUT_S = 120.0 -async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lora_path: Path) -> None: +async def load_lora_adapter( + admin_clients: list[AsyncClient], + lora_name: str, + lora_path: Path, + *, + admin: AdminAPI = _DEFAULT_ADMIN, +) -> None: """Make a HTTP post request to the vLLM server to load a LoRA adapter. Uses our wrapper endpoint that also resets the prefix cache to invalidate @@ -439,6 +704,7 @@ async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lo """ logger = get_logger() lora_path_posix = lora_path.as_posix() + per_attempt_timeout = httpx.Timeout(connect=10.0, read=LORA_LOAD_READ_TIMEOUT_S, write=60.0, pool=10.0) @retry( retry=retry_if_exception(_is_retryable_lora_error), @@ -448,29 +714,11 @@ async def load_lora_adapter(admin_clients: list[AsyncClient], lora_name: str, lo ) async def _load_lora_adapter(admin_client: AsyncClient) -> None: logger.debug(f"Sending request to load LoRA adapter {lora_name} from {lora_path}") - response = await admin_client.post( - "/load_lora_adapter", - json={"lora_name": lora_name, "lora_path": lora_path_posix}, - timeout=httpx.Timeout(connect=10.0, read=LORA_LOAD_READ_TIMEOUT_S, write=60.0, pool=10.0), - ) - response.raise_for_status() + await admin.load_lora_adapter(admin_client, lora_name, lora_path_posix, timeout=per_attempt_timeout) await asyncio.gather(*[_load_lora_adapter(admin_client) for admin_client in admin_clients]) -async def unload_lora_adapter(admin_clients: list[AsyncClient], lora_name: str) -> None: - """Make a HTTP post request to the vLLM server to unload a LoRA adapter.""" - logger = get_logger() - - async def _unload_lora_adapter(admin_client: AsyncClient) -> None: - logger.debug(f"Sending request to unload LoRA adapter {lora_name}") - await admin_client.post("/v1/unload_lora_adapter", json={"lora_name": lora_name}) - # TODO: The first one can fail, but subsequent ones should succeed. - # response.raise_for_status() - - await asyncio.gather(*[_unload_lora_adapter(admin_client) for admin_client in admin_clients]) - - async def init_nccl_broadcast( admin_clients: list[AsyncClient], host: str, @@ -478,6 +726,8 @@ async def init_nccl_broadcast( timeout: int, inference_world_size: int | None = None, quantize_in_weight_transfer: bool = False, + *, + admin: AdminAPI = _DEFAULT_ADMIN, ) -> None: """Initialize NCCL broadcast on all inference servers. @@ -502,18 +752,15 @@ async def init_nccl_broadcast( async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> None: try: - response = await admin_client.post( - "/init_broadcaster", - json={ - "host": host, - "port": port, - "rank_offset": rank_offset, - "inference_world_size": inference_world_size, - "timeout": timeout, - "quantize_in_weight_transfer": quantize_in_weight_transfer, - }, + await admin.init_broadcaster( + admin_client, + host=host, + port=port, + rank_offset=rank_offset, + inference_world_size=inference_world_size, + timeout=timeout, + quantize_in_weight_transfer=quantize_in_weight_transfer, ) - response.raise_for_status() except httpx.HTTPStatusError as e: if e.response.status_code == 404: logger.warning("The route /init_broadcaster does not exist. Skipping NCCL broadcast initialization.") diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index 951b3673c1..8cee2e4cca 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -22,7 +22,14 @@ from renderers import RendererConfig from prime_rl.configs.shared import ClientConfig -from prime_rl.utils.client import ClientIdentity, client_identity, load_lora_adapter, setup_admin_clients, setup_clients +from prime_rl.utils.client import ( + ClientIdentity, + client_identity, + load_lora_adapter, + setup_admin_api, + setup_admin_clients, + setup_clients, +) from prime_rl.utils.logger import get_logger # --- Shared discovery functions --- @@ -127,6 +134,7 @@ def __init__( self._servers: dict[str, ServerState] = {} self._admin_clients: dict[str, AsyncClient] = {} + self._admin_api = setup_admin_api(client_config) self._lock = asyncio.Lock() self._desired: AdapterState = AdapterState() @@ -334,7 +342,9 @@ async def _sync_server_adapter(self, ip: str) -> bool: if self._desired.name and self._desired.path: try: self.logger.debug(f"Loading adapter {self._desired.name} on {ip}") - await load_lora_adapter([self._admin_clients[ip]], self._desired.name, self._desired.path) + await load_lora_adapter( + [self._admin_clients[ip]], self._desired.name, self._desired.path, admin=self._admin_api + ) except Exception as e: server.status = "unhealthy" server.sync_failures += 1 @@ -369,12 +379,8 @@ async def _check_server_health(self, admin_client: AsyncClient, ip: str) -> bool return False try: - response = await admin_client.get("/v1/models") - response.raise_for_status() - data = response.json() - models = [m.get("id") for m in data.get("data", [])] - - if self.base_model_name not in models: + models = await self._admin_api.list_models(admin_client) + if self.base_model_name not in [m.get("id") for m in models]: self.logger.debug(f"Server {ip} does not have base model {self.base_model_name}") return False except Exception as e: diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 69325e6c4a..002680bd0c 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -6,7 +6,14 @@ import verifiers as vf from prime_rl.configs.shared import ClientConfig -from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients +from prime_rl.utils.client import ( + DynamoAdminAPI, + _dynamo_rl_discovery_base_urls, + _is_retryable_lora_error, + discover_dynamo_admin_base_urls, + load_lora_adapter, + setup_clients, +) def test_is_retryable_lora_error_returns_true_for_404(): @@ -49,6 +56,32 @@ def test_load_lora_adapter_succeeds_on_first_attempt(): ) +def test_dynamo_load_lora_adapter_uses_existing_lora_engine_route(): + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = {"status": "success"} + mock_client.post.return_value = mock_response + + asyncio.run( + DynamoAdminAPI().load_lora_adapter( + mock_client, + "test-lora", + "/test/path", + timeout=httpx.Timeout(connect=10.0, read=30.0, write=60.0, pool=10.0), + ) + ) + + mock_client.post.assert_called_once_with( + "/engine/load_lora", + json={ + "lora_name": "test-lora", + "source": {"uri": Path("/test/path").absolute().as_uri()}, + }, + timeout=httpx.Timeout(connect=10.0, read=30.0, write=60.0, pool=10.0), + ) + + def test_setup_clients_assigns_renderer_and_dp_rank_headers(): from renderers import Qwen3VLRendererConfig @@ -117,3 +150,70 @@ def test_setup_clients_preserves_chat_client_defaults(): extra_headers_from_state={}, ) ] + + +def test_dynamo_rl_discovery_base_urls_derive_from_base_url(monkeypatch): + monkeypatch.setenv("DYN_RL_PORT", "18001") + client_config = ClientConfig( + base_url=["http://frontend.local:8000/v1"], + backend="dynamo", + ) + + assert _dynamo_rl_discovery_base_urls(client_config) == ["http://frontend.local:18001"] + + +def test_dynamo_rl_discovery_base_urls_honor_explicit_config(): + client_config = ClientConfig( + base_url=["http://frontend.local:8000/v1"], + backend="dynamo", + rl_base_url=["http://frontend.local:8001/v1"], + ) + + assert _dynamo_rl_discovery_base_urls(client_config) == ["http://frontend.local:8001/v1"] + + +def test_discover_dynamo_admin_base_urls_reads_workers(monkeypatch): + calls = [] + + class FakeResponse: + def raise_for_status(self): + pass + + def json(self): + return { + "workers": [ + {"system_url": "http://worker-0:8081"}, + {"system_url": "http://worker-0:8081"}, + {"system_url": "http://worker-1:8081"}, + {}, + ] + } + + class FakeClient: + def __init__(self, *, base_url, headers, timeout): + calls.append(("init", base_url, headers, timeout)) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get(self, path): + calls.append(("get", path)) + return FakeResponse() + + monkeypatch.setattr("prime_rl.utils.client.httpx.Client", FakeClient) + + client_config = ClientConfig( + base_url=["http://frontend.local:8000/v1"], + backend="dynamo", + rl_base_url=["http://frontend.local:8001/v1"], + ) + + assert discover_dynamo_admin_base_urls(client_config) == [ + "http://worker-0:8081", + "http://worker-1:8081", + ] + assert calls[0][0:2] == ("init", "http://frontend.local:8001") + assert ("get", "/v1/rl/workers") in calls diff --git a/tools/dynamo/configs/smoke_rl.toml b/tools/dynamo/configs/smoke_rl.toml new file mode 100644 index 0000000000..b2c1d77342 --- /dev/null +++ b/tools/dynamo/configs/smoke_rl.toml @@ -0,0 +1,47 @@ +# Prime-RL smoke test: Dynamo as inference backend (short -- 5 steps) +# Model: PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT +# GPU layout: GPU 0 = Dynamo, GPU 1 = trainer +# +# Usage: uv run orchestrator @ tools/dynamo/configs/smoke_rl.toml --output-dir /tmp/.../run_default + +max_steps = 5 +seq_len = 512 +batch_size = 16 +rollouts_per_example = 4 +# Disable TITO: Dynamo injects token_ids via nvext (DYN_ENABLE_RL=true) rather than +# the TITO /v1/chat/completions/tokens path. Keep false for all Dynamo deployments. +use_token_client = false + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +offline = true + +[sampling] +max_tokens = 64 + +[[filters]] +type = "gibberish" +enforce = false + +[[filters]] +type = "repetition" +enforce = false + +[[filters]] +type = "zero_advantage" +enforce = false + +[[env]] +id = "reverse-text" + +[client] +# Point at Dynamo instead of prime-rl's own vLLM +base_url = ["http://localhost:8000/v1"] +# Discover worker-advertised system URLs from Dynamo's dedicated RL listener, +# then call each worker's /engine/* admin routes directly. +backend = "dynamo" +rl_base_url = ["http://localhost:8001/v1"] +# Dynamo may serve model under a different name -- skip the /v1/models check +skip_model_check = true diff --git a/tools/dynamo/configs/smoke_rl_long.toml b/tools/dynamo/configs/smoke_rl_long.toml new file mode 100644 index 0000000000..d21c8f9cb9 --- /dev/null +++ b/tools/dynamo/configs/smoke_rl_long.toml @@ -0,0 +1,40 @@ +# Prime-RL LONG smoke test: Dynamo as inference backend (20 steps) +# Enough steps to observe training dynamics over more iterations. +# +# Usage: uv run orchestrator @ tools/dynamo/configs/smoke_rl_long.toml --output-dir /tmp/.../run_default + +max_steps = 20 +seq_len = 512 +batch_size = 16 +rollouts_per_example = 4 +use_token_client = false + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +offline = true + +[sampling] +max_tokens = 64 + +[[filters]] +type = "gibberish" +enforce = false + +[[filters]] +type = "repetition" +enforce = false + +[[filters]] +type = "zero_advantage" +enforce = false + +[[env]] +id = "reverse-text" + +[client] +base_url = ["http://localhost:8000/v1"] +backend = "dynamo" +rl_base_url = ["http://localhost:8001/v1"] +skip_model_check = true diff --git a/tools/dynamo/configs/smoke_trainer.toml b/tools/dynamo/configs/smoke_trainer.toml new file mode 100644 index 0000000000..aff75cd895 --- /dev/null +++ b/tools/dynamo/configs/smoke_trainer.toml @@ -0,0 +1,20 @@ +# Prime-RL trainer config for Dynamo smoke test (short -- 5 steps) +# Single GPU, no FSDP. Reads rollouts written by the orchestrator process. +# +# Usage: uv run torchrun --nproc-per-node=1 --rdzv-endpoint=localhost:29510 \ +# -m prime_rl.trainer.rl.train @ tools/dynamo/configs/smoke_trainer.toml \ +# --output-dir /tmp/dynamo_smoke_outputs + +max_steps = 5 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" +seq_len = 512 + +[wandb] +offline = true + +[optim] +lr = 3e-6 + +[ckpt] diff --git a/tools/dynamo/configs/smoke_trainer_long.toml b/tools/dynamo/configs/smoke_trainer_long.toml new file mode 100644 index 0000000000..82768c456e --- /dev/null +++ b/tools/dynamo/configs/smoke_trainer_long.toml @@ -0,0 +1,19 @@ +# Prime-RL trainer config for long smoke test (20 steps) +# +# Usage: uv run torchrun --nproc-per-node=1 --rdzv-endpoint=localhost:29510 \ +# -m prime_rl.trainer.rl.train @ tools/dynamo/configs/smoke_trainer_long.toml \ +# --output-dir /tmp/dynamo_smoke_long + +max_steps = 20 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" +seq_len = 512 + +[wandb] +offline = true + +[optim] +lr = 3e-6 + +[ckpt] diff --git a/tools/dynamo/run_dynamo.sh b/tools/dynamo/run_dynamo.sh new file mode 100755 index 0000000000..6a27580cec --- /dev/null +++ b/tools/dynamo/run_dynamo.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# Launch Dynamo frontend + vLLM worker for smoke testing with prime-rl. +# +# Prerequisites: +# - Dynamo virtualenv at $DYNAMO_VENV (default: ~/dev/dynamo/dynamo) +# - etcd + NATS running (dynamo depends on them) +# +# Usage: +# ./tools/dynamo/run_dynamo.sh +# CUDA_VISIBLE_DEVICES=0 ./tools/dynamo/run_dynamo.sh +# DYNAMO_MODEL=my-org/my-model ./tools/dynamo/run_dynamo.sh + +set -euo pipefail + +DYNAMO_VENV="${DYNAMO_VENV:-/home/biswaranjanp/dev/rl/dynamo/.venv}" +DYNAMO_MODEL="${DYNAMO_MODEL:-PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT}" +DYNAMO_GPU_MEMORY_UTILIZATION="${DYNAMO_GPU_MEMORY_UTILIZATION:-0.50}" +CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +LOG_DIR="${LOG_DIR:-/tmp}" +PRIME_RL_SRC="${PRIME_RL_SRC:-/home/biswaranjanp/dev/rl/prime-rl/src}" + +export CUDA_VISIBLE_DEVICES + +source "$DYNAMO_VENV/bin/activate" + +echo "[dynamo-smoke] Starting frontend..." +DYN_ENABLE_RL=true \ +DYN_RL_PORT="${DYN_RL_PORT:-8001}" \ +python -m dynamo.frontend > "$LOG_DIR/dynamo_frontend.log" 2>&1 & +FRONTEND_PID=$! +echo "[dynamo-smoke] Frontend PID: $FRONTEND_PID (log: $LOG_DIR/dynamo_frontend.log)" + +sleep 5 + +echo "[dynamo-smoke] Starting vLLM worker (model: $DYNAMO_MODEL)..." +PYTHONPATH="${PRIME_RL_SRC}${PYTHONPATH:+:$PYTHONPATH}" \ +DYN_SYSTEM_PORT="${DYN_SYSTEM_PORT:-8081}" \ +python -m dynamo.vllm \ + --model "$DYNAMO_MODEL" \ + --enforce-eager \ + --max-model-len 2048 \ + --max-num-seqs 32 \ + --gpu-memory-utilization "$DYNAMO_GPU_MEMORY_UTILIZATION" \ + --enable-rl \ + --worker-extension-cls prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker \ + > "$LOG_DIR/dynamo_vllm.log" 2>&1 & +WORKER_PID=$! +echo "[dynamo-smoke] Worker PID: $WORKER_PID (log: $LOG_DIR/dynamo_vllm.log)" + +echo "[dynamo-smoke] Both processes launched." +echo "[dynamo-smoke] Frontend log: $LOG_DIR/dynamo_frontend.log" +echo "[dynamo-smoke] Worker log: $LOG_DIR/dynamo_vllm.log" +wait diff --git a/tools/dynamo/run_smoke_test.sh b/tools/dynamo/run_smoke_test.sh new file mode 100755 index 0000000000..16e5af1548 --- /dev/null +++ b/tools/dynamo/run_smoke_test.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# Full Dynamo + prime-rl smoke test: orchestrator + trainer. +# +# GPU 0: Dynamo inference (must already be running -- see run_dynamo.sh) +# GPU 1: prime-rl trainer (single GPU, uses torchrun) +# +# Directory structure created: +# $OUTPUT_DIR/ <- trainer output_dir +# $OUTPUT_DIR/run_default/ <- orchestrator output_dir (run_* naming required) +# +# Usage: +# # Short run (5 steps, default): +# ./tools/dynamo/run_smoke_test.sh +# +# # Long run (20 steps): +# ./tools/dynamo/run_smoke_test.sh --long +# +# # Custom output dir: +# OUTPUT_DIR=/tmp/my_run ./tools/dynamo/run_smoke_test.sh + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)" +CONFIG_DIR="$SCRIPT_DIR/configs" + +# Defaults +OUTPUT_DIR="${OUTPUT_DIR:-/tmp/dynamo_smoke_outputs}" +TRAINER_GPU="${TRAINER_GPU:-1}" +RDZV_PORT="${RDZV_PORT:-29510}" +SMOKE_MODEL="${SMOKE_MODEL:-}" +MODEL_ARGS=() +if [[ -n "$SMOKE_MODEL" ]]; then + MODEL_ARGS=(--model.name "$SMOKE_MODEL" --tokenizer.name "$SMOKE_MODEL") +fi + +# Pick short vs long config +if [[ "${1:-}" == "--long" ]]; then + ORCH_CONFIG="$CONFIG_DIR/smoke_rl_long.toml" + TRAINER_CONFIG="$CONFIG_DIR/smoke_trainer_long.toml" + ORCH_LOG="${OUTPUT_DIR}/smoke_long_orchestrator.log" + TRAINER_LOG="${OUTPUT_DIR}/smoke_long_trainer.log" + echo "[smoke] Using LONG configs (20 steps)" +else + ORCH_CONFIG="$CONFIG_DIR/smoke_rl.toml" + TRAINER_CONFIG="$CONFIG_DIR/smoke_trainer.toml" + ORCH_LOG="${OUTPUT_DIR}/smoke_orchestrator.log" + TRAINER_LOG="${OUTPUT_DIR}/smoke_trainer.log" + echo "[smoke] Using SHORT configs (5 steps)" +fi + +ORCH_DIR="$OUTPUT_DIR/run_default" +mkdir -p "$ORCH_DIR" + +cd "$REPO_DIR" + +# Start orchestrator (no GPU needed -- just API calls to Dynamo) +echo "[smoke] Starting orchestrator (output: $ORCH_DIR)..." +BENCH_API_KEY=EMPTY CUDA_VISIBLE_DEVICES="" \ + uv run orchestrator \ + @ "$ORCH_CONFIG" \ + --output-dir "$ORCH_DIR" \ + "${MODEL_ARGS[@]}" \ + > "$ORCH_LOG" 2>&1 & +ORCH_PID=$! +echo "[smoke] Orchestrator PID: $ORCH_PID (log: $ORCH_LOG)" + +# Give orchestrator time to write first batch +sleep 12 + +# Start trainer on GPU via uv run torchrun (uses project venv Python, needed for torch.distributed) +echo "[smoke] Starting trainer (output: $OUTPUT_DIR)..." +CUDA_VISIBLE_DEVICES="$TRAINER_GPU" uv run torchrun \ + --nproc-per-node=1 \ + --rdzv-endpoint="localhost:$RDZV_PORT" \ + --rdzv-id="smoke_$(date +%s)" \ + -m prime_rl.trainer.rl.train \ + @ "$TRAINER_CONFIG" \ + --output-dir "$OUTPUT_DIR" \ + "${MODEL_ARGS[@]}" \ + > "$TRAINER_LOG" 2>&1 & +TRAINER_PID=$! +echo "[smoke] Trainer PID: $TRAINER_PID (log: $TRAINER_LOG)" + +echo "[smoke] Waiting for both processes..." + +ORCH_EXIT=0 +TRAINER_EXIT=0 +wait $ORCH_PID || ORCH_EXIT=$? +wait $TRAINER_PID || TRAINER_EXIT=$? + +echo "" +echo "[smoke] Orchestrator exit: $ORCH_EXIT" +echo "[smoke] Trainer exit: $TRAINER_EXIT" + +if [[ $ORCH_EXIT -eq 0 && $TRAINER_EXIT -eq 0 ]]; then + echo "[smoke] Smoke test PASSED." +else + echo "[smoke] Smoke test FAILED." + exit 1 +fi