diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml
index f39bf888af..34e15bcbeb 100644
--- a/.github/workflows/pr-test.yml
+++ b/.github/workflows/pr-test.yml
@@ -595,7 +595,15 @@ jobs:
},
{
"num_gpus": 0,
- "test_file": "test_train_dump.py"
+ "test_file": "test_train_data_utils.py"
+ },
+ {
+ "num_gpus": 0,
+ "test_file": "test_rollout_data_utils.py"
+ },
+ {
+ "num_gpus": 0,
+ "test_file": "test_rollout_metrics.py"
},
{
"num_gpus": 0,
@@ -617,10 +625,18 @@ jobs:
"num_gpus": 0,
"test_file": "test_logprob_response_spans.py"
},
+ {
+ "num_gpus": 0,
+ "test_file": "observability/test_trace_utils.py"
+ },
{
"num_gpus": 0,
"test_file": "test_value_temperature.py"
},
+ {
+ "num_gpus": 0,
+ "test_file": "test_ppo_kl_metric.py"
+ },
{
"num_gpus": 0,
"test_file": "test_cispo_loss.py"
@@ -697,6 +713,10 @@ jobs:
"num_gpus": 0,
"test_file": "test_qwen3_5_vl_native.py"
},
+ {
+ "num_gpus": 0,
+ "test_file": "test_accelerator.py"
+ },
{
"num_gpus": 0,
"test_file": "test_reloadable_process_group_world.py"
@@ -782,7 +802,7 @@ jobs:
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors psutil
- pip install transformers
+ pip install transformers wandb
- name: Install
diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2
index ce2425e39b..6ed6147126 100644
--- a/.github/workflows/pr-test.yml.j2
+++ b/.github/workflows/pr-test.yml.j2
@@ -66,7 +66,7 @@
'label': 'run-ci-cpu-unittest',
'always': True,
'cpu': True,
- 'extra_pip_deps': 'transformers',
+ 'extra_pip_deps': 'transformers wandb',
'tests': [
{'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0},
{'test_file': 'test_deep_ep_tms_patch.py', 'num_gpus': 0},
@@ -74,13 +74,17 @@
{'test_file': 'utils/test_megatron_server_arguments.py', 'num_gpus': 0},
{'test_file': 'test_dp_schedule.py', 'num_gpus': 0},
{'test_file': 'test_cp_utils.py', 'num_gpus': 0},
- {'test_file': 'test_train_dump.py', 'num_gpus': 0},
+ {'test_file': 'test_train_data_utils.py', 'num_gpus': 0},
+ {'test_file': 'test_rollout_data_utils.py', 'num_gpus': 0},
+ {'test_file': 'test_rollout_metrics.py', 'num_gpus': 0},
{'test_file': 'test_metric_report.py', 'num_gpus': 0},
{'test_file': 'test_metric_report_dist.py', 'num_gpus': 0},
{'test_file': 'test_loss_cp_invariance.py', 'num_gpus': 0},
{'test_file': 'test_advantage_whiten_cp.py', 'num_gpus': 0},
{'test_file': 'test_logprob_response_spans.py', 'num_gpus': 0},
+ {'test_file': 'observability/test_trace_utils.py', 'num_gpus': 0},
{'test_file': 'test_value_temperature.py', 'num_gpus': 0},
+ {'test_file': 'test_ppo_kl_metric.py', 'num_gpus': 0},
{'test_file': 'test_cispo_loss.py', 'num_gpus': 0},
{'test_file': 'test_policy_loss.py', 'num_gpus': 0},
{'test_file': 'test_ppo_logprob_entropy.py', 'num_gpus': 0},
@@ -100,6 +104,7 @@
{'test_file': 'test_rollout_sample_hooks.py', 'num_gpus': 0},
{'test_file': 'test_hf_to_megatron.py', 'num_gpus': 0},
{'test_file': 'test_qwen3_5_vl_native.py', 'num_gpus': 0},
+ {'test_file': 'test_accelerator.py', 'num_gpus': 0},
{'test_file': 'test_reloadable_process_group_world.py', 'num_gpus': 0},
{'test_file': 'test_placement_group.py', 'num_gpus': 0},
{'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0},
diff --git a/docker/NOTES_GB10.md b/docker/NOTES_GB10.md
index d5417321f5..4901b5a790 100644
--- a/docker/NOTES_GB10.md
+++ b/docker/NOTES_GB10.md
@@ -52,7 +52,7 @@ https://catalog.ngc.nvidia.com/orgs/nvidia/containers/vllm?version=26.03-py3
| 9 | CMake 3.31 rejects `CMAKE_CUDA_ARCHITECTURES=120f;121f` | `f` suffix for CUDA 13 Blackwell family is only supported in CMake ≥4.0 | Upgrade to `cmake==4.3.1` via pip (must override NGC `/etc/pip/constraint.txt` with `PIP_CONSTRAINT=`) and set `CMAKE_POLICY_VERSION_MINIMUM=3.5` for old bundled deps |
| 10 | TE 2.10 build: `cuda_profiler_api.h: No such file` | CUDA 13 removed the public header for `cudaProfilerStart/Stop`; the symbols still exist in `libcudart.so.13`. TE's 3 `fused_softmax` TUs `#include` the header but don't call the APIs | Install a 20-line shim header at `/usr/local/cuda/include/cuda_profiler_api.h` declaring the two functions extern. Stored as `docker/patch/gb10/cuda_profiler_api.h` |
| 11 | slime `train.py --help`: `'tuple' object has no attribute 'strip'` | Typo in `slime/utils/arguments.py:1073`: `help=("string",)` (trailing comma → tuple) instead of `help=("string")` | Remove trailing comma — simple one-line slime fix, upstream-able |
-| 12 | `sglang_router` x86_64-only wheel from `zhuzilin/sgl-router` fork | slime Dockerfile pins `zhuzilin/sgl-router` release (no arm64 builds); slime's `'slime' in version` assertion is only in `wandb_utils.py` (non-critical path) | Install upstream `sglang-router==0.3.2` from PyPI (has arm64 wheel). Accept wandb path fallback |
+| 12 | `sglang_router` x86_64-only wheel from `zhuzilin/sgl-router` fork | slime Dockerfile pins `zhuzilin/sgl-router` release (no arm64 builds); slime's `'slime' in version` assertion is only in `slime/observability/wandb_utils.py` (non-critical path) | Install upstream `sglang-router==0.3.2` from PyPI (has arm64 wheel). Accept wandb path fallback |
| 13 | `antlr4-python3-runtime==4.13.2` → `Could not deserialize ATN with version 3` | Omegaconf's bundled grammar was generated with antlr 4.9 serialized format; runtime 4.13 only reads format v4 | Pin `antlr4-python3-runtime==4.9.3` |
| 14 | `megatron.training` not importable after `pip install -e Megatron-LM` | Megatron-LM's setup.py only packages `megatron-core`; `megatron.training`, `megatron.rl`, `megatron.legacy` are sibling dirs meant to be on `PYTHONPATH` | `export PYTHONPATH=/root/src/Megatron-LM:$PYTHONPATH` (slime docs confirm this) |
| 15 | `libz3.so` missing for tilelang | tilelang uses Z3 SMT solver for autoscheduling; NGC vllm base doesn't include libz3 | `apt-get update && apt-get install -y libz3-dev` (libz3-4 package alias needs update first) |
diff --git a/docker/npu_patch/slime.patch b/docker/npu_patch/slime.patch
index fee99b00b3..769eb4783d 100644
--- a/docker/npu_patch/slime.patch
+++ b/docker/npu_patch/slime.patch
@@ -84,11 +84,11 @@ index 7bc4f910..a00c83c51 100644
from ray.actor import ActorHandle
from torch_memory_saver import torch_memory_saver
@@ -17,6 +21,7 @@ from slime.ray.train_actor import TrainRayActor
- from slime.utils import train_dump_utils
+ from slime.observability import train_data_utils, train_metric_utils
from slime.utils.data import process_rollout_data
from slime.utils.distributed_utils import get_gloo_group, init_process_group
+from slime.utils.http_utils import _wrap_ipv6
- from slime.utils.logging_utils import init_tracking
+ from slime.observability.logging_utils import init_tracking
from slime.utils.memory_utils import clear_memory, print_memory
from slime.utils.misc import Box
@@ -55,6 +60,8 @@ class MegatronTrainRayActor(TrainRayActor):
@@ -586,13 +586,13 @@ diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py
index 75cb053c..c54a3854 100644
--- a/slime/ray/rollout.py
+++ b/slime/ray/rollout.py
-@@ -28,6 +28,7 @@ from slime.utils.metric_utils import (
+@@ -28,6 +28,7 @@ from slime.observability.metric_utils import (
from slime.utils.misc import Box, group_by, load_function
from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
from slime.utils.types import Sample
+from slime.utils.common import is_npu
- from ..utils.metric_utils import has_repetition
+ from slime.observability.metric_utils import has_repetition
from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock
@@ -76,7 +77,8 @@ class RolloutManager:
self.all_rollout_engines = [None] * num_engines
@@ -631,7 +631,7 @@ index 2e900ca5..d0a25583 100644
+++ b/slime/ray/train_actor.py
@@ -13,16 +13,23 @@ from slime.ray.ray_actor import RayActor
from slime.utils.distributed_utils import init_gloo_group
- from slime.utils.logging_utils import configure_logger
+ from slime.observability.logging_utils import configure_logger
from slime.utils.memory_utils import clear_memory, print_memory
+from slime.utils.common import is_npu
@@ -966,4 +966,4 @@ index 01883c47..18faa6d2 100644
+ import mindspeed.megatron_adaptor
from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
from slime.utils.arguments import parse_args
- from slime.utils.logging_utils import configure_logger, init_tracking
+ from slime.observability.logging_utils import configure_logger, init_tracking
diff --git a/docs/en/developer_guide/trace.md b/docs/en/developer_guide/trace.md
index d866537298..604efa6564 100644
--- a/docs/en/developer_guide/trace.md
+++ b/docs/en/developer_guide/trace.md
@@ -41,7 +41,7 @@ By default it also starts a local static server so you can open the generated HT
## Instrument custom code
-For custom rollout or reward code — including custom agent steps, tool calls, sandbox execution, and verifier calls in agentic workflows — reuse helpers from `slime.utils.trace_utils`:
+For custom rollout or reward code — including custom agent steps, tool calls, sandbox execution, and verifier calls in agentic workflows — reuse helpers from `slime.observability.trace_utils`:
- `trace_span(target, name, attrs=...)`: record a duration span.
- `trace_event(target, name, attrs=...)`: record an instant event.
@@ -57,7 +57,7 @@ Use `trace_function(...)` when the whole function should be represented as one s
The decorator is what slime uses for the main rollout pipeline. For example, `generate_and_rm(...)` is traced per sample and `generate_and_rm_group(...)` is traced per sample group:
```python
-from slime.utils.trace_utils import trace_function
+from slime.observability.trace_utils import trace_function
@trace_function("generate_and_rm", target="sample")
@@ -104,7 +104,7 @@ If you need to add attrs after part of the function has executed, use an inner `
If you want to record SGLang generation metadata in a consistent way, reuse `build_sglang_meta_trace_attrs`:
```python
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span
+from slime.observability.trace_utils import build_sglang_meta_trace_attrs, trace_span
with trace_span(sample, "sglang_generate") as span:
output = await post(url, payload)
@@ -116,4 +116,3 @@ with trace_span(sample, "sglang_generate") as span:
- Save a small number of rollouts first; the viewer is easiest to read when each dump contains a manageable number of samples.
- The viewer is built from the saved `.pt` dump, so traces can be inspected offline on another machine.
- For GPU/kernel-level SGLang profiling traces, see [Profiling](./profiling.md).
-
diff --git a/docs/en/get_started/customization.md b/docs/en/get_started/customization.md
index 476b2ffed9..6a9cfb6a55 100644
--- a/docs/en/get_started/customization.md
+++ b/docs/en/get_started/customization.md
@@ -42,7 +42,7 @@ For most agentic use cases, **start with `--custom-generate-function-path` plus
| Replace the entire rollout orchestration (only when per-sample customization is not enough) | [`--rollout-function-path`](#1-rollout-function---rollout-function-path) |
| Control task sampling, buffering, requeueing, or custom prompt/task sources | [`--data-source-path`](#15-data-source---data-source-path) |
| Attach custom loss masks, metadata, or convert agentic outputs into training data | [`--rollout-data-postprocess-path`](#8-rollout-data-postprocess---rollout-data-postprocess-path), [`--custom-convert-samples-to-train-data-path`](#13-samples-to-train-data-conversion---custom-convert-samples-to-train-data-path) |
-| Debug long-running custom generation, verifier calls, tool calls, or sandbox steps | trace utilities in [`slime.utils.trace_utils`](../developer_guide/trace.md) |
+| Debug long-running custom generation, verifier calls, tool calls, or sandbox steps | trace utilities in [`slime.observability.trace_utils`](../developer_guide/trace.md) |
A native example of this pattern is [`examples/search-r1`](../../../examples/search-r1/), which adds search-augmented multi-turn generation via `--custom-generate-function-path` while keeping slime's default `sglang_rollout` outer loop. See also [`examples/multi_agent`](../../../examples/multi_agent/README.md) for a `--rollout-function-path`-based multi-agent pattern and [`examples/fully_async`](../../../examples/fully_async/README.md) for long-tail agentic generation.
diff --git a/docs/en/get_started/quick_start.md b/docs/en/get_started/quick_start.md
index d0bd3fb780..0eb848ac93 100644
--- a/docs/en/get_started/quick_start.md
+++ b/docs/en/get_started/quick_start.md
@@ -584,6 +584,8 @@ export NCCL_SOCKET_IFNAME=$(ip -o -4 addr show | awk '$4 ~ /^10\\./ {print $2}')
export NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=$(ip -o -4 addr show | awk '$4 ~ /^10\./ {print $2}')
```
+For launching the same multi-node setup on Kubernetes or cloud instances with one command, see the [SkyPilot tutorial](../platform_support/skypilot_tutorial.md).
+
slime has been deeply optimized for distributed training of large-scale Mixture of Experts (MoE) models. We provide some end-to-end training cases for reference:
- [Example: 8xH100 Training GLM-4.7-Flash](../examples/glm4.7-30B-A3B.md)
diff --git a/docs/en/index.rst b/docs/en/index.rst
index 83230af1fa..d7a1c5df0e 100644
--- a/docs/en/index.rst
+++ b/docs/en/index.rst
@@ -110,9 +110,10 @@ Start by Use Case
.. toctree::
:maxdepth: 1
- :caption: Hardware Platforms
+ :caption: Platforms
platform_support/amd_tutorial.md
+ platform_support/skypilot_tutorial.md
.. toctree::
:maxdepth: 1
diff --git a/docs/en/platform_support/skypilot_tutorial.md b/docs/en/platform_support/skypilot_tutorial.md
new file mode 100644
index 0000000000..746d88b666
--- /dev/null
+++ b/docs/en/platform_support/skypilot_tutorial.md
@@ -0,0 +1,319 @@
+# SkyPilot
+
+[SkyPilot](https://github.com/skypilot-org/skypilot) is an open-source framework for running workloads on Kubernetes or any cloud. This tutorial shows how to launch multi-node slime training with SkyPilot: node provisioning, Ray cluster startup, and job submission are described in a single YAML, replacing the per-node `ray start` steps from the [Quick Start](../get_started/quick_start.md).
+
+It covers two setups, both running the Quick Start's Qwen3-4B GRPO recipe (`scripts/run-qwen3-4B.sh`) on the DAPO-math dataset:
+
+- **Multi-node training on one cluster** — the standard setup from the Quick Start's multi-node section.
+- **Disaggregated training and inference** — the trainer and SGLang engines run as separate, gang-scheduled jobs that scale independently.
+
+This page is maintained by the SkyPilot maintainers.
+
+## Prerequisites
+
+Install SkyPilot with the extras for your infrastructure and confirm it can reach it:
+
+```bash
+pip install "skypilot[kubernetes]" # or [aws], [gcp], ... — see SkyPilot docs
+sky check
+```
+
+The examples below use the `slimerl/slime:latest` Docker image from the Quick Start, so no additional environment setup is needed inside the nodes.
+
+## Multi-Node Training on One Cluster
+
+The Quick Start starts a Ray cluster by running `ray start` on every node, then submits training with `ray job submit` from node 0. The following SkyPilot task performs the same steps: it provisions `num_nodes` nodes with GPUs, downloads and converts the model on each node, starts the Ray head and workers, and submits the job. Environment variables like `SKYPILOT_NODE_RANK` and `SKYPILOT_NODE_IPS` are injected by SkyPilot on every node.
+
+
+slime-multinode.yaml
+
+```yaml
+# slime-multinode.yaml
+resources:
+ infra: kubernetes # or aws / gcp / any infra configured in `sky check`
+ accelerators: H100:4
+ image_id: docker:slimerl/slime:latest
+
+num_nodes: 2
+
+setup: |
+ pip install -q -U "huggingface_hub[cli]"
+ [ -d /root/Qwen3-4B ] || hf download Qwen/Qwen3-4B --local-dir /root/Qwen3-4B
+ [ -d /root/dapo-math-17k ] || hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k
+ [ -d /root/aime-2024 ] || hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/aime-2024
+ # Convert the HF checkpoint to Megatron torch_dist format (each node needs a local copy).
+ if [ ! -d /root/Qwen3-4B_torch_dist ]; then
+ cd /root/slime
+ source scripts/models/qwen3-4B.sh
+ PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
+ ${MODEL_ARGS[@]} --hf-checkpoint /root/Qwen3-4B --save /root/Qwen3-4B_torch_dist
+ fi
+
+run: |
+ MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
+ if [ "$SKYPILOT_NODE_RANK" != "0" ]; then
+ # Worker nodes join the Ray cluster. --block keeps the worker's Ray daemons in the
+ # foreground for the whole run (an exiting run command would get them reaped) and
+ # returns once the head shuts down at the end of training.
+ sleep 10
+ ray start --address=${MASTER_ADDR}:6379 --num-gpus ${SKYPILOT_NUM_GPUS_PER_NODE} --disable-usage-stats \
+ --dashboard-agent-listen-port 52366 --metrics-export-port 8091 --block
+ exit 0
+ fi
+
+ # Start Ray from /root/slime: job entrypoints run in the head's working directory.
+ cd /root/slime
+ source scripts/models/qwen3-4B.sh
+
+ # Non-default agent/metrics ports: SkyPilot's runtime on the node runs its own Ray.
+ ray start --head --node-ip-address ${MASTER_ADDR} \
+ --num-gpus ${SKYPILOT_NUM_GPUS_PER_NODE} --disable-usage-stats \
+ --dashboard-host=0.0.0.0 --dashboard-port=8265 \
+ --dashboard-agent-listen-port 52366 --metrics-export-port 8091
+
+ # Wait until every node has joined the Ray cluster.
+ until python3 -c "import ray, sys; ray.init(address='${MASTER_ADDR}:6379', logging_level='error'); sys.exit(0 if len([n for n in ray.nodes() if n['Alive']]) >= ${SKYPILOT_NUM_NODES} else 1)"; do sleep 5; done
+
+ # Wait for Ray's job agent to be ready to accept submissions.
+ until ray job submit --address="http://127.0.0.1:8265" --no-wait -- true >/dev/null 2>&1; do
+ echo "waiting for the Ray job agent..."; sleep 5
+ done
+ ray job submit --address="http://127.0.0.1:8265" \
+ --runtime-env-json='{"env_vars": {"PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1"}}' \
+ -- python3 /root/slime/train.py \
+ --actor-num-nodes ${SKYPILOT_NUM_NODES} \
+ --actor-num-gpus-per-node ${SKYPILOT_NUM_GPUS_PER_NODE} \
+ --num-gpus-per-node ${SKYPILOT_NUM_GPUS_PER_NODE} \
+ --colocate \
+ ${MODEL_ARGS[@]} \
+ --hf-checkpoint /root/Qwen3-4B \
+ --ref-load /root/Qwen3-4B_torch_dist \
+ --load /root/Qwen3-4B_slime/ \
+ --save /root/Qwen3-4B_slime/ \
+ --save-interval 20 \
+ --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl \
+ --input-key prompt \
+ --label-key label \
+ --apply-chat-template \
+ --rollout-shuffle \
+ --rm-type deepscaler \
+ --num-rollout 3000 \
+ --rollout-batch-size 32 \
+ --n-samples-per-prompt 8 \
+ --rollout-max-response-len 8192 \
+ --rollout-temperature 1 \
+ --global-batch-size 256 \
+ --balance-data \
+ --eval-interval 20 \
+ --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl \
+ --n-samples-per-eval-prompt 16 \
+ --eval-max-response-len 16384 \
+ --eval-top-p 1 \
+ --advantage-estimator grpo \
+ --use-kl-loss \
+ --kl-loss-coef 0.00 \
+ --kl-loss-type low_var_kl \
+ --entropy-coef 0.00 \
+ --eps-clip 0.2 \
+ --eps-clip-high 0.28 \
+ --optimizer adam \
+ --lr 1e-6 \
+ --lr-decay-style constant \
+ --weight-decay 0.1 \
+ --adam-beta1 0.9 \
+ --adam-beta2 0.98 \
+ --tensor-model-parallel-size 2 \
+ --sequence-parallel \
+ --pipeline-model-parallel-size 1 \
+ --context-parallel-size 1 \
+ --expert-model-parallel-size 1 \
+ --expert-tensor-parallel-size 1 \
+ --recompute-granularity full \
+ --recompute-method uniform \
+ --recompute-num-layers 1 \
+ --use-dynamic-batch-size \
+ --max-tokens-per-gpu 9216 \
+ --rollout-num-gpus-per-engine 2 \
+ --sglang-mem-fraction-static 0.7 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --accumulate-allreduce-grads-in-fp32 \
+ --attention-softmax-in-fp32 \
+ --attention-backend flash
+```
+
+
+
+The training arguments are the Quick Start's Qwen3-4B recipe; the only adjustments are the topology flags (`--actor-num-nodes`, `--actor-num-gpus-per-node`, and `--num-gpus-per-node`), whose values come from the SkyPilot-injected environment. `--num-gpus-per-node` matters on nodes with fewer than 8 GPUs: slime's colocated engine mapping assumes 8 per node unless told otherwise. Launch it with:
+
+```bash
+sky launch -c slime-train slime-multinode.yaml
+```
+
+SkyPilot provisions the nodes (creating them if needed), runs `setup` and `run` on each node, and streams the logs. `sky down slime-train` tears the cluster down. The task assumes a fresh cluster: to re-run training, recreate the cluster (`sky down slime-train && sky launch -c slime-train ...`) rather than re-launching onto one whose Ray daemons are still running. The same YAML can be launched as a managed job with `sky jobs launch`, which adds automatic recovery from node failures.
+
+## Disaggregated Training and Inference
+
+slime supports connecting the trainer to SGLang engines launched by an external system (`--rollout-external-engine-addrs`, see [External Rollout Engines](../advanced/external-rollout-engines.md)). With a SkyPilot **Job Group**, the trainer and each engine are separate jobs in one YAML that are gang-scheduled together and reach each other by stable hostname (`-0.`), so the fleet of engines can be sized independently of the trainer.
+
+The trainer publishes updated weights after each optimizer step and the engines reload them from a shared `ReadWriteMany` volume (`--update-weight-transport disk`). Create the volume once:
+
+```yaml
+# policy-volume.yaml
+name: slime-policy
+type: k8s-pvc
+size: 100Gi
+infra: kubernetes
+config:
+ access_mode: ReadWriteMany
+```
+
+```bash
+sky volumes apply policy-volume.yaml
+```
+
+Then launch the Job Group:
+
+
+slime-jobgroup.yaml
+
+```yaml
+# slime-jobgroup.yaml
+---
+name: slime-rl
+execution: parallel
+primary_tasks: [trainer] # the group succeeds/fails with the trainer
+inter_connection: true # place all jobs on one cluster so they can reach each other
+termination_delay: 60s
+---
+name: sglang
+resources:
+ infra: kubernetes
+ accelerators: H100:1
+ image_id: docker:slimerl/slime:latest
+volumes:
+ /shared/policy: slime-policy
+setup: |
+ pip install -q -U "huggingface_hub[cli]"
+ [ -d /root/Qwen3-4B ] || hf download Qwen/Qwen3-4B --local-dir /root/Qwen3-4B
+run: |
+ # One SGLang server; the trainer reaches it at sglang-0.:30000.
+ python -m sglang.launch_server --model-path /root/Qwen3-4B --tp 1 \
+ --host 0.0.0.0 --port 30000 --mem-fraction-static 0.7
+---
+name: trainer
+resources:
+ infra: kubernetes
+ accelerators: H100:2
+ image_id: docker:slimerl/slime:latest
+volumes:
+ /shared/policy: slime-policy
+setup: |
+ pip install -q -U "huggingface_hub[cli]"
+ [ -d /root/Qwen3-4B ] || hf download Qwen/Qwen3-4B --local-dir /root/Qwen3-4B
+ [ -d /root/dapo-math-17k ] || hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k
+ [ -d /root/aime-2024 ] || hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/aime-2024
+ # Convert the HF checkpoint to Megatron torch_dist format.
+ if [ ! -d /root/Qwen3-4B_torch_dist ]; then
+ cd /root/slime
+ source scripts/models/qwen3-4B.sh
+ PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
+ ${MODEL_ARGS[@]} --hf-checkpoint /root/Qwen3-4B --save /root/Qwen3-4B_torch_dist
+ fi
+run: |
+ # Wait for the engine job to serve (jobs in a group provision independently).
+ ENGINE_ADDR="sglang-0.${SKYPILOT_JOBGROUP_NAME}:30000"
+ until curl -sf "http://${ENGINE_ADDR}/health" >/dev/null; do
+ echo "waiting for engine ${ENGINE_ADDR}..."; sleep 10
+ done
+ echo "engine healthy: ${ENGINE_ADDR}"
+
+ # Start Ray from /root/slime: job entrypoints run in the head's working directory.
+ cd /root/slime
+ source scripts/models/qwen3-4B.sh
+
+ # Non-default agent/metrics ports: SkyPilot's runtime on the node runs its own Ray.
+ ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats \
+ --dashboard-host=0.0.0.0 --dashboard-port=8265 \
+ --dashboard-agent-listen-port 52366 --metrics-export-port 8091
+
+ # Wait for Ray's job agent to be ready to accept submissions.
+ until ray job submit --address="http://127.0.0.1:8265" --no-wait -- true >/dev/null 2>&1; do
+ echo "waiting for the Ray job agent..."; sleep 5
+ done
+ ray job submit --address="http://127.0.0.1:8265" \
+ --runtime-env-json='{"env_vars": {"PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1"}}' \
+ -- python3 /root/slime/train.py \
+ --actor-num-nodes 1 \
+ --actor-num-gpus-per-node 2 \
+ --rollout-external-engine-addrs ${ENGINE_ADDR} \
+ --update-weight-mode full \
+ --update-weight-transport disk \
+ --update-weight-disk-dir /shared/policy \
+ ${MODEL_ARGS[@]} \
+ --hf-checkpoint /root/Qwen3-4B \
+ --ref-load /root/Qwen3-4B_torch_dist \
+ --load /root/Qwen3-4B_slime/ \
+ --save /root/Qwen3-4B_slime/ \
+ --save-interval 20 \
+ --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl \
+ --input-key prompt \
+ --label-key label \
+ --apply-chat-template \
+ --rollout-shuffle \
+ --rm-type deepscaler \
+ --num-rollout 3000 \
+ --rollout-batch-size 32 \
+ --n-samples-per-prompt 8 \
+ --rollout-max-response-len 8192 \
+ --rollout-temperature 1 \
+ --global-batch-size 256 \
+ --balance-data \
+ --advantage-estimator grpo \
+ --use-kl-loss \
+ --kl-loss-coef 0.00 \
+ --kl-loss-type low_var_kl \
+ --entropy-coef 0.00 \
+ --eps-clip 0.2 \
+ --eps-clip-high 0.28 \
+ --optimizer adam \
+ --lr 1e-6 \
+ --lr-decay-style constant \
+ --weight-decay 0.1 \
+ --adam-beta1 0.9 \
+ --adam-beta2 0.98 \
+ --tensor-model-parallel-size 2 \
+ --sequence-parallel \
+ --pipeline-model-parallel-size 1 \
+ --context-parallel-size 1 \
+ --expert-model-parallel-size 1 \
+ --expert-tensor-parallel-size 1 \
+ --recompute-granularity full \
+ --recompute-method uniform \
+ --recompute-num-layers 1 \
+ --use-dynamic-batch-size \
+ --max-tokens-per-gpu 9216 \
+ --rollout-num-gpus-per-engine 1 \
+ --sglang-mem-fraction-static 0.7 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --accumulate-allreduce-grads-in-fp32 \
+ --attention-softmax-in-fp32 \
+ --attention-backend flash
+```
+
+
+
+```bash
+sky jobs launch -n slime-rl slime-jobgroup.yaml
+```
+
+To scale the inference fleet, add more engine jobs (`sglang-2`, `sglang-3`, ...) to the YAML and append their addresses to `--rollout-external-engine-addrs`. For large models, `--update-weight-mode delta` ships only the changed bytes ([Delta Weight Sync](../advanced/delta-weight-sync.md)); NCCL transport (`--update-weight-transport nccl`) avoids the shared volume entirely.
+
+## End-to-End Example: Agentic Coding RL
+
+A complete agentic RL version of the disaggregated setup lives in the SkyPilot repository:
+
+**[slime on SkyPilot Job Groups](https://github.com/skypilot-org/skypilot/tree/master/llm/slime)** — trains a coding agent (Qwen3-14B) on SWE-smith with slime: a Megatron trainer job plus 1–3 SGLang engine jobs in one Job Group, agent rollouts executing untrusted code in sandboxed pods, and disk-based delta weight sync between the jobs. The example includes launch YAMLs, all setup/run scripts, and benchmark results for scaling the inference fleet (1 → 3 engines cuts async step time from about 1200 s to about 660 s on the example workload).
+
+Issues with the SkyPilot setups on this page can be reported to the [SkyPilot repository](https://github.com/skypilot-org/skypilot/issues).
diff --git a/docs/zh/developer_guide/trace.md b/docs/zh/developer_guide/trace.md
index dfd2a812cc..ab2afe016d 100644
--- a/docs/zh/developer_guide/trace.md
+++ b/docs/zh/developer_guide/trace.md
@@ -41,7 +41,7 @@ python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt
## 给自定义代码打点
-在自定义 rollout 或 reward 逻辑中——包括 agentic workflow 里的 agent step、tool call、sandbox 执行、verifier 调用等——可以直接复用 `slime.utils.trace_utils` 里的工具：
+在自定义 rollout 或 reward 逻辑中——包括 agentic workflow 里的 agent step、tool call、sandbox 执行、verifier 调用等——可以直接复用 `slime.observability.trace_utils` 里的工具：
- `trace_span(target, name, attrs=...)`：记录一段持续时间。
- `trace_event(target, name, attrs=...)`：记录一个瞬时事件。
@@ -57,7 +57,7 @@ python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt
slime 主 rollout 流程里就是这样用的。例如 `generate_and_rm(...)` 按 sample 打点，而 `generate_and_rm_group(...)` 按 group 打点：
```python
-from slime.utils.trace_utils import trace_function
+from slime.observability.trace_utils import trace_function
@trace_function("generate_and_rm", target="sample")
@@ -104,7 +104,7 @@ async def custom_rollout_batch(samples, **kwargs):
如果想统一记录 SGLang 返回的 generation 元信息，可以复用 `build_sglang_meta_trace_attrs`：
```python
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span
+from slime.observability.trace_utils import build_sglang_meta_trace_attrs, trace_span
with trace_span(sample, "sglang_generate") as span:
output = await post(url, payload)
@@ -116,4 +116,3 @@ with trace_span(sample, "sglang_generate") as span:
- 先保存少量 rollout；单个 dump 的 sample 数量适中时，viewer 会更容易阅读。
- viewer 直接基于保存下来的 `.pt` dump 工作，因此可以把文件拷到别的机器离线分析。
- 如果你想看的是 SGLang 自身的 GPU / kernel 级 profiling trace，请参考 [性能分析](./profiling.md)。
-
diff --git a/docs/zh/get_started/customization.md b/docs/zh/get_started/customization.md
index c3cc65e8a3..992a1bfc2d 100644
--- a/docs/zh/get_started/customization.md
+++ b/docs/zh/get_started/customization.md
@@ -42,7 +42,7 @@ agentic workflow——multi-turn tool use、sandbox interaction、environment fe
| 替换整个 rollout 编排（只在 per-sample 自定义不够用时使用） | [`--rollout-function-path`](#1-rollout-函数---rollout-function-path) |
| 控制任务采样、缓冲、回填，或自定义 prompt / task 数据源 | [`--data-source-path`](#15-数据源---data-source-path) |
| 给 agentic 输出附加自定义 loss mask、metadata，或转换成训练数据 | [`--rollout-data-postprocess-path`](#8-rollout-数据后处理---rollout-data-postprocess-path)、[`--custom-convert-samples-to-train-data-path`](#13-样本转训练数据---custom-convert-samples-to-train-data-path) |
-| 调试长耗时的 custom generation、verifier、tool call 或 sandbox 调用 | [`slime.utils.trace_utils`](../developer_guide/trace.md) 中的 trace 工具 |
+| 调试长耗时的 custom generation、verifier、tool call 或 sandbox 调用 | [`slime.observability.trace_utils`](../developer_guide/trace.md) 中的 trace 工具 |
这一模式的原生示例是 [`examples/search-r1`](../../../examples/search-r1/)：通过 `--custom-generate-function-path` 接入搜索增强的多轮生成，外层仍然走 slime 默认的 `sglang_rollout`。互补的示例可参考 [`examples/multi_agent`](../../../examples/multi_agent/README.md) 中基于 `--rollout-function-path` 的多 agent 模式，以及 [`examples/fully_async`](../../../examples/fully_async/README.md) 中适合 long-tail agentic 场景的 fully-async rollout。
diff --git a/requirements.txt b/requirements.txt
index 79a708f949..5acae2c7d2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,7 +16,6 @@ pylatexenc
pyyaml
qwen_vl_utils # for VLM
ray[default]
-ring_flash_attn
safetensors
sglang-router>=0.3.0
tensorboard
diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py
index e1315d6633..7703fa2913 100644
--- a/slime/backends/megatron_utils/__init__.py
+++ b/slime/backends/megatron_utils/__init__.py
@@ -2,6 +2,10 @@
import torch
+from slime.utils import accelerator
+
+accelerator.initialize_accelerator()
+
try:
import deep_ep
from torch_memory_saver import torch_memory_saver
@@ -21,7 +25,18 @@ def new_init(self, *args, **kwargs):
# DeepEP owns persistent buffers and may initialize them on its
# internal streams. Make their lifetime independent of the TMS
# disabled region before restoring allocation tracking.
- torch.cuda.synchronize()
+ # CPU-only imports intentionally have no selected device; explicit
+ # accelerator requests still fail fast in initialize_accelerator().
+ selected_accelerator = accelerator.initialize_accelerator()
+ if selected_accelerator is not None:
+ selected_accelerator.synchronize()
+ else:
+ # Keep the historical CUDA hook observable for CPU test
+ # doubles, while ignoring the expected no-CUDA runtime error.
+ try:
+ torch.cuda.synchronize()
+ except RuntimeError:
+ pass
finally:
cdll.tms_set_interesting_region(original_interesting_region)
diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py
index ea2601a9d4..1d20abc96f 100644
--- a/slime/backends/megatron_utils/actor.py
+++ b/slime/backends/megatron_utils/actor.py
@@ -12,10 +12,14 @@
from torch_memory_saver import torch_memory_saver
from transformers import AutoConfig, AutoTokenizer
+from slime.observability import train_data_utils, train_metric_utils
+from slime.observability.logging_utils import init_tracking
+from slime.observability.profile_utils import TrainProfiler
+from slime.observability.timer import Timer, inverse_timer, timer, with_defer
from slime.ray.train_actor import TrainRayActor
+from slime.utils import accelerator
from slime.utils.data import process_rollout_data
from slime.utils.distributed_utils import get_gloo_group
-from slime.utils.logging_utils import init_tracking
from slime.utils.memory_utils import clear_memory, print_memory
from slime.utils.misc import Box
from slime.utils.reloadable_process_group import (
@@ -25,15 +29,12 @@
reload_process_groups,
)
from slime.utils.routing_replay import RoutingReplay
-from slime.utils.timer import Timer, inverse_timer, timer, with_defer
from slime.utils.types import RolloutBatch
-from ...utils.profile_utils import TrainProfiler
from ...utils.tensor_backper import TensorBackuper
-from . import train_dump_utils
from .checkpoint import load_checkpoint
from .cp_utils import prepare_routed_experts_for_routing_replay, slice_log_prob_with_cp
-from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data
+from .data import DataIterator, get_data_iterator
from .hf_checkpoint_saver import save_hf_model_to_path
from .initialize import init, is_megatron_main_rank
from .loss import (
@@ -237,7 +238,7 @@ def wake_up(self) -> None:
# that is the first NCCL operation on a group. Prime WORLD here,
# after the memory saver is resumed, so later stages cannot miss its
# lazy initialization. Sleep still destroys it completely.
- dist.barrier(device_ids=[torch.cuda.current_device()])
+ dist.barrier(device_ids=[accelerator.current_device()])
if self.role == "actor":
self._switch_model("actor")
print_memory("after wake_up model")
@@ -253,7 +254,7 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch:
)
# TODO: this is ugly, move to somewhere else?
# move tokens to GPU in advance
- device = torch.cuda.current_device()
+ device = accelerator.current_device()
rollout_data["tokens"] = [
t.to(device=device, dtype=torch.long, non_blocking=True) for t in rollout_data["tokens"]
]
@@ -359,7 +360,6 @@ def compute_log_prob(
num_microbatches: list[int],
store_prefix: str = "",
) -> dict[str, list[torch.Tensor]]:
-
with timer(f"{store_prefix}log_probs"):
return forward_only(
get_log_probs_and_entropy,
@@ -505,7 +505,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
if self.rollout_data_postprocess is not None:
self.rollout_data_postprocess(self.args, rollout_id, rollout_data)
- log_rollout_data(
+ train_metric_utils.log_rollout_data(
rollout_id,
self.args,
rollout_data,
@@ -542,7 +542,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
self.prof.step(rollout_id=rollout_id)
- train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
+ train_data_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
if self.args.use_routing_replay:
RoutingReplay.clear_all()
@@ -561,7 +561,11 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
logger.info(f"Updating ref model at rollout_id {rollout_id}")
self.weights_backuper.backup("ref")
- log_perf_data(rollout_id, self.args, extra_metrics=self.weight_updater.pop_metrics())
+ train_metric_utils.log_perf_data(
+ rollout_id,
+ self.args,
+ extra_metrics=self.weight_updater.pop_metrics(),
+ )
@timer
def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
diff --git a/slime/backends/megatron_utils/cp_utils.py b/slime/backends/megatron_utils/cp_utils.py
index 96c97df0e2..a54da3a7ef 100644
--- a/slime/backends/megatron_utils/cp_utils.py
+++ b/slime/backends/megatron_utils/cp_utils.py
@@ -124,114 +124,6 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor:
return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token
-def reduce_train_step_metrics(
- losses_reduced: list[dict],
- *,
- calculate_per_token_loss: bool,
- step_global_batch_size: int,
- cp_size: int,
- dp_with_cp_group,
-) -> dict[str, float]:
- """Aggregate per-mb log dicts into the dict ``train_one_step`` reports.
-
- Pipeline (1:1 with what the train loop used to do inline):
- 1. Sum each metric's per-mb ``values`` tensor locally on this rank.
- 2. All-reduce across the DP*CP group (``dp_with_cp_group``).
- 3. Apply the per-mode divisor / cp_factor:
- - per-token-loss: divisor = ``values[0]`` = all-reduced ``num_tokens``,
- CP-inflated by ``cp_size`` because every CP rank computes the same
- num_tokens off the FULL (not chunked) masks; the
- ``cp_factor = cp_size`` multiplier cancels that inflation, leaving
- the genuine per-token average.
- - per-rollout-mean: divisor = constant ``step_global_batch_size`` from
- the rollout side, never all-reduced, so no CP inflation to cancel
- and ``cp_factor = 1``.
-
- Tests pass a mock ``dp_with_cp_group`` and monkeypatch ``dist.all_reduce``
- to a no-op, then pre-aggregate virtual ranks themselves — this exercises
- the same call shape as production while staying single-process.
- """
- keys = losses_reduced[0]["keys"]
- values = None
- for x in losses_reduced:
- values = x["values"] if values is None else values + x["values"]
- assert len(keys) + 1 == values.numel()
- dist.all_reduce(values, group=dp_with_cp_group)
- values = values.tolist()
-
- if calculate_per_token_loss:
- num_samples_or_tokens = values[0]
- cp_factor = cp_size
- else:
- num_samples_or_tokens = step_global_batch_size
- cp_factor = 1
- return {key: value * cp_factor / num_samples_or_tokens for key, value in zip(keys, values[1:], strict=False)}
-
-
-def rollout_log_metric_contribution(
- per_rank_reducer_sum: float,
- *,
- cp_size: int,
- num_rollouts_in_rollout: int,
- dp_size: int,
-) -> tuple[float, float]:
- """``(sum, count)`` tuple to hand the gather step for a per-rollout-mean
- metric on the rollout side (``log_rollout_data``).
-
- Sum across DP*CP ranks of ``count`` lands on ``num_rollouts_in_rollout``
- (``dp_size`` here is the no-CP DP width; the gather covers ``dp_size *
- cp_size`` ranks, and each rank emits the same ``count``, so the totals
- cancel out the ``cp_size`` in the sum). Result: ``Σsum / Σcount =
- sum_DP_full / num_rollouts`` — the same number ``train_one_step`` reports
- for the same samples (when ``num_steps_per_rollout == 1``).
-
- Pair with :func:`gather_and_reduce_log_dict` to do the full end-to-end
- in tests (single helper call per rank, returns the reduced number on
- the source rank).
- """
- sum_value = cp_size * per_rank_reducer_sum
- count = num_rollouts_in_rollout / dp_size
- return sum_value, count
-
-
-def gather_and_reduce_log_dict(
- log_dict: dict,
- *,
- dp_size: int,
- dp_src_rank: int,
- dp_group,
-) -> dict | None:
- """``dist.gather_object`` per-rank log_dicts + per-key reduction.
-
- Per key in the gathered dicts:
- - ``(sum, count)`` tuple → ``Σsum / Σcount`` (per-rollout-mean shape;
- pair with :func:`rollout_log_metric_contribution`).
- - plain value → ``Σ / dp_size`` (legacy mean-across-ranks; the only
- correct answer when ranks hold the same data).
-
- Returns the reduced dict on ``dp_src_rank``, ``None`` elsewhere. The
- caller adds whatever metric-name prefix / wandb plumbing it wants —
- this helper stays free of side effects so CPU multi-process unit tests
- can drive it directly with real ``torch.distributed``.
- """
- if dist.get_rank() == dp_src_rank:
- gathered = [None] * dp_size
- dist.gather_object(log_dict, gathered, dst=dp_src_rank, group=dp_group)
- reduced: dict = {}
- for key in log_dict:
- values = [d[key] for d in gathered]
- first = values[0]
- if isinstance(first, tuple) and len(first) == 2:
- total_sum = sum(v[0] for v in values)
- total_count = sum(v[1] for v in values)
- reduced[key] = total_sum / total_count if total_count else 0.0
- else:
- reduced[key] = sum(values) / dp_size
- return reduced
- dist.gather_object(log_dict, None, dst=dp_src_rank, group=dp_group)
- return None
-
-
def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor:
"""
Gather tensors across all ranks in the context parallel group.
diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py
index 8ab7235da9..6ae48a9613 100644
--- a/slime/backends/megatron_utils/data.py
+++ b/slime/backends/megatron_utils/data.py
@@ -1,28 +1,14 @@
-import logging
-from argparse import Namespace
from collections.abc import Sequence
-import numpy as np
import torch
-import torch.distributed as dist
import torch.nn.functional as F
from megatron.core import mpu
from megatron.core.packed_seq_params import PackedSeqParams
-from slime.utils import train_metric_utils
-from slime.utils.flops_utils import calculate_fwd_flops
-from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step
+from slime.utils import accelerator
from slime.utils.types import RolloutBatch
-from ...utils import logging_utils
-from .cp_utils import (
- gather_and_reduce_log_dict,
- get_sum_of_sample_mean,
- rollout_log_metric_contribution,
- slice_with_cp,
-)
-
-logger = logging.getLogger(__name__)
+from .cp_utils import slice_with_cp
def get_batch(
@@ -83,7 +69,7 @@ def get_batch(
tokens = F.pad(tokens, (0, pad), value=pad_token_id)
cu_seqlens_list.append(cu_seqlens_list[-1] + pad)
- cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=torch.cuda.current_device())
+ cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=accelerator.current_device())
tokens = tokens.chunk(cp_size, dim=0)[cp_rank]
else:
tokens = [slice_with_cp(t, pad_token_id) for t in tokens]
@@ -101,7 +87,7 @@ def get_batch(
cu_seqlens.append(cu_seqlens[-1] + pad)
# thd requires the cu_seqlens to be of the origin length
- cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int).cuda() * cp_size
+ cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int, device=accelerator.device()) * cp_size
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
packed_seq_params = PackedSeqParams(
@@ -163,41 +149,6 @@ def get_batch(
return batch
-def gather_log_data(
- metric_name: str,
- args: Namespace,
- rollout_id: int,
- log_dict: dict[str, "float | tuple[float, float]"],
-) -> dict[str, float] | None:
- """
- Gather per-rank metrics, reduce on the DP source rank, and log to W&B / TB.
-
- Each value in ``log_dict`` is either:
- * a ``(sum, count)`` tuple → reduced as ``Σsum / Σcount``;
- * a plain scalar → reduced as ``Σ / dp_size`` (mean across ranks).
-
- The gather + reduce step is delegated to
- :func:`cp_utils.gather_and_reduce_log_dict` so it can be exercised by
- CPU multi-process unit tests directly. This function adds the
- ``metric_name`` prefix and the W&B / TB logging side effects.
- """
- reduced = gather_and_reduce_log_dict(
- log_dict,
- dp_size=mpu.get_data_parallel_world_size(with_context_parallel=True),
- dp_src_rank=mpu.get_data_parallel_src_rank(with_context_parallel=True),
- dp_group=mpu.get_data_parallel_group_gloo(with_context_parallel=True),
- )
- if reduced is None:
- return None
- reduced_log_dict = {f"{metric_name}/{k}": v for k, v in reduced.items()}
- logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}")
- # Calculate step once to avoid duplication
- step = compute_rollout_step(args, rollout_id)
- reduced_log_dict["rollout/step"] = step
- logging_utils.log(args, reduced_log_dict, step_key="rollout/step")
- return reduced_log_dict
-
-
class DataIterator:
"""Iterator over a rollout dict following an explicit micro-batch index schedule."""
@@ -245,277 +196,6 @@ def get_data_iterator(rollout_data: RolloutBatch) -> list[DataIterator]:
return [DataIterator(rollout_data, micro_batch_indices) for _ in range(vpp_size)]
-def log_rollout_data(
- rollout_id: int,
- args: Namespace,
- rollout_data: RolloutBatch,
-) -> None:
- """
- Summarize rollout fields and log reduced metrics on PP last stage, TP rank 0.
-
- - Tensor-valued lists are concatenated and averaged. For token-level metrics
- like log-probs/returns/advantages/values, computes a CP-correct sample mean
- using `loss_masks` and total/response lengths.
- - Non-tensor lists are averaged elementwise.
- - Scalars are converted to Python numbers.
- """
- if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
- cp_size = mpu.get_context_parallel_world_size()
- log_dict = {}
- response_lengths = rollout_data["response_lengths"]
- loss_masks = rollout_data["loss_masks"]
- total_lengths = rollout_data["total_lengths"]
- # Same per-rollout denominators the training loss uses, so reported
- # log_probs / returns / advantages / etc. live in the same per-rollout
- # mean space (rather than per-sample) as the gradient signal.
- rollout_mask_sums = rollout_data.get("rollout_mask_sums", None)
- # For per-rollout-mean metrics: ``rollout_log_metric_contribution``
- # produces the ``(sum, count)`` tuple so gather_log_data's
- # ``Σsum / Σcount`` lands on ``sum_DP_full / num_rollouts`` — the
- # same number train_one_step reports for the same samples.
- dp_world = mpu.get_data_parallel_world_size(with_context_parallel=False)
- num_rollouts_in_rollout = sum(rollout_data["global_batch_sizes"])
-
- for key, val in rollout_data.items():
- if key in [
- "tokens",
- "multimodal_train_inputs",
- "loss_masks",
- "sample_indices",
- "rollout_ids",
- "rollout_mask_sums",
- "rollout_top_p_token_ids",
- "rollout_top_p_token_offsets",
- "rollout_routed_experts",
- "global_batch_sizes",
- "num_microbatches",
- "micro_batch_indices",
- "source_names",
- # DP-local view of `raw_reward`, which this loop already logs;
- # both reduce to the same mean, so skip the duplicate metric.
- "local_raw_reward",
- ]:
- continue
- # Emit (sum, count) so gather_log_data can do a weighted average across
- # DP ranks. This stops the legacy "every rank has the same N samples"
- # assumption from biasing means once uneven-DP partitioning lands.
- if isinstance(val, (list, tuple)):
- count = len(val)
- if isinstance(val[0], torch.Tensor):
- # NOTE: Here we have to do the clone().detach(), otherwise the tensor will be
- # modified in place and will cause problem for the next rollout.
- if key in [
- "log_probs",
- "ref_log_probs",
- "rollout_log_probs",
- "returns",
- "advantages",
- "values",
- "teacher_log_probs",
- "opd_reverse_kl",
- ]:
- tensor = torch.cat(val).clone().detach()
- sum_of_sample_mean = get_sum_of_sample_mean(
- total_lengths,
- response_lengths,
- loss_masks,
- rollout_mask_sums,
- )
- # Compute (sum, count) via the shared helper so this
- # path and the unit tests stay in sync.
- sum_value, count = rollout_log_metric_contribution(
- sum_of_sample_mean(tensor).item(),
- cp_size=cp_size,
- num_rollouts_in_rollout=num_rollouts_in_rollout,
- dp_size=dp_world,
- )
- log_dict[key] = (sum_value, count)
- continue
- tensor = torch.cat(val).clone().detach()
- # val.mean() * cp_size is the per-sample mean for one rank;
- # multiply by count to get the per-rank sum.
- per_rank_sum = tensor.mean() * cp_size * count
- sum_value = per_rank_sum.item()
- else:
- sum_value = sum(val)
- log_dict[key] = (sum_value, count)
- elif isinstance(val, torch.Tensor):
- # Scalar tensor (one per rank): treat as count=1.
- log_dict[key] = (val.float().mean().item(), 1)
- else:
- raise ValueError(f"Unsupported type: {type(val)} for key: {key}")
-
- reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict)
- if args.ci_test and reduced_log_dict is not None:
- # This is an initial actor/ref zero-KL check. R3 replays rollout
- # routing for the actor forward, while the reference forward
- # intentionally falls through to normal routing, so their
- # log-probs are not expected to match bit-for-bit in CI.
- if (
- rollout_id == 0
- and not getattr(args, "ci_disable_kl_checker", False)
- and not getattr(args, "use_rollout_routing_replay", False)
- and "rollout/log_probs" in reduced_log_dict
- and "rollout/ref_log_probs" in reduced_log_dict
- ):
- # TODO: figure out why there is a small numerical difference in log_probs and ref_log_probs in CI test, and whether it's expected or not.
- # assert reduced_log_dict["rollout/log_probs"] == reduced_log_dict["rollout/ref_log_probs"]
- assert abs(reduced_log_dict["rollout/log_probs"] - reduced_log_dict["rollout/ref_log_probs"]) < 1e-8
- if "rollout/log_probs" in reduced_log_dict:
- assert -1 < reduced_log_dict["rollout/log_probs"] < 0
- if "rollout/entropy" in reduced_log_dict:
- assert 0 < reduced_log_dict["rollout/entropy"] < 1
-
- if args.log_multi_turn:
- log_multi_turn_data(rollout_id, args, rollout_data)
- if args.log_passrate:
- log_passrate(rollout_id, args, rollout_data)
-
- if args.log_correct_samples:
- if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
- cp_size = mpu.get_context_parallel_world_size()
- log_dict = {}
- response_lengths = rollout_data["response_lengths"]
- loss_masks = rollout_data["loss_masks"]
- total_lengths = rollout_data["total_lengths"]
-
- def quantile(total_value, n_quantiles, data) -> dict:
- import math
-
- assert n_quantiles > 1, f"n_quantiles({n_quantiles}) must be greater than 1."
-
- quantiles = [((i + 1) / n_quantiles) for i in range(n_quantiles)]
- cut_points = [total_value * q for q in quantiles]
- cut_points[-1] = total_value
-
- count = [0] * n_quantiles
- for d in data:
- for i, point in enumerate(cut_points):
- if d <= point:
- count[i] += 1
- break
-
- total = sum(count) + 1e-9
- percentile = [c / total for c in count]
-
- percentile = {f"p{min(math.ceil(q*100),100)}": p for q, p in zip(quantiles, percentile, strict=True)}
- return percentile
-
- # DP-local, so it lines up positionally with response_lengths /
- # total_lengths / loss_masks / log_probs below. `raw_reward` itself
- # is the whole rollout batch (log_passrate needs the full grouping).
- raw_rewards = rollout_data["local_raw_reward"]
- # Additional metrics for correct cases are calculated separately below.
- correct_response_lengths = []
- correct_total_lengths = []
- correct_loss_masks = []
- correct_entropy = []
- for i, raw_reward in enumerate(raw_rewards):
- if raw_reward == 1:
- correct_response_lengths.append(response_lengths[i])
- correct_total_lengths.append(total_lengths[i])
- correct_loss_masks.append(loss_masks[i])
- correct_entropy.append(-rollout_data["log_probs"][i])
- num_correct_responses = len(correct_total_lengths)
- rollout_data["correct_response_lengths"] = correct_response_lengths
- correct_response_length_percentile = quantile(
- args.rollout_max_response_len, 4, rollout_data["correct_response_lengths"]
- )
- for p, val in correct_response_length_percentile.items():
- rollout_data[f"correct_length/{p}"] = [val] * num_correct_responses
- if len(correct_entropy) > 0:
- # NOTE: per-sample-mean over the correct subset, not per-rollout.
- # A rollout's siblings may not all be correct, and slicing
- # ``rollout_mask_sums`` here would leave a denom that still
- # includes incorrect siblings — meaningless for a "correct-only"
- # entropy report. Per-sample-mean over the filtered subset is
- # the cleanest semantic.
- sum_of_sample_mean = get_sum_of_sample_mean(
- correct_total_lengths, correct_response_lengths, correct_loss_masks, sample_denoms=None
- )
- correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0))
- rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses
- else:
- rollout_data["correct_entropy"] = [0] * num_correct_responses
-
-
-def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
- """
- Log multi-turn auxiliary metrics such as raw/observed response lengths and rounds.
-
- Operates only on PP last stage and TP rank 0. Uses GPU tensors when available
- to compute statistics without host transfers.
- """
- if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
- log_dict = {}
- for key, val in rollout_data.items():
- if key == "loss_masks":
- if val: # Check if val is not empty
- device = val[0].device # Get device from first tensor
-
- # Vectorized length calculation using torch
- raw_response_lengths = torch.tensor([v.shape[0] for v in val], dtype=torch.float32, device=device)
- log_dict["raw_response_length/response_length_mean"] = raw_response_lengths.mean().item()
- log_dict["raw_response_length/response_length_max"] = raw_response_lengths.max().item()
- log_dict["raw_response_length/response_length_min"] = raw_response_lengths.min().item()
- log_dict["raw_response_length/response_length_clip_ratio"] = (
- (raw_response_lengths >= args.rollout_max_response_len).float().mean().item()
- )
-
- # Vectorized sum calculation using torch - stay on GPU
- wo_obs_response_lengths = torch.tensor(
- [v.sum().item() for v in val], dtype=torch.float32, device=device
- )
- log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item()
- log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item()
- log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item()
- if key == "round_number":
- # Use numpy for vectorized round number statistics
- round_number_array = np.array(val)
- log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array)
- log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array)
- log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array)
- gather_log_data("multi_turn", args, rollout_id, log_dict)
-
-
-def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
- """
- Compute pass@k metrics from `raw_reward` groups and log the results.
-
- `raw_reward` is reshaped to `[group_number, group_size]`, then pass@k is
- estimated per problem and averaged.
- """
- if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
- log_dict = {}
- for key, val in rollout_data.items():
- if key != "raw_reward":
- continue
-
- log_dict |= compute_pass_rate(
- flat_rewards=val,
- group_size=args.n_samples_per_prompt,
- num_groups=args.rollout_batch_size,
- )
-
- gather_log_data("passrate", args, rollout_id, log_dict)
-
-
-def log_perf_data(rollout_id: int, args: Namespace, extra_metrics: dict | None = None) -> None:
- train_metric_utils.log_perf_data_raw(
- rollout_id=rollout_id,
- args=args,
- is_primary_rank=(
- mpu.get_tensor_model_parallel_rank() == 0
- and mpu.is_pipeline_last_stage()
- and mpu.get_data_parallel_rank(with_context_parallel=True) == 0
- ),
- compute_total_fwd_flops=lambda seq_lens: calculate_fwd_flops(seqlens=seq_lens, args=args)
- / dist.get_world_size()
- / 1e12,
- extra_metrics=extra_metrics,
- )
-
-
def tensors_to_cpu(tensor_list):
"""Move a list of GPU tensors to CPU for Ray object store transfer.
@@ -543,5 +223,5 @@ def tensors_to_gpu(tensor_list, device=None):
if tensor_list is None:
return None
if device is None:
- device = torch.cuda.current_device()
+ device = accelerator.current_device()
return [t.to(device=device, dtype=torch.float32) for t in tensor_list]
diff --git a/slime/backends/megatron_utils/hf_checkpoint_saver.py b/slime/backends/megatron_utils/hf_checkpoint_saver.py
index f45f10e3e8..a574d34f1c 100644
--- a/slime/backends/megatron_utils/hf_checkpoint_saver.py
+++ b/slime/backends/megatron_utils/hf_checkpoint_saver.py
@@ -8,6 +8,8 @@
import torch
+from slime.utils import accelerator
+
logger = logging.getLogger(__name__)
_HF_WEIGHT_FILE_NAMES = {
@@ -223,9 +225,10 @@ def _write_pending_chunk(
if pending_write is not None:
shard_idx, named_tensors = pending_write
writer.write(named_tensors, shard_idx=shard_idx)
- if torch.cuda.is_available():
- torch.cuda.ipc_collect()
- torch.cuda.empty_cache()
+ selected_accelerator = accelerator.initialize_accelerator()
+ if selected_accelerator is not None:
+ selected_accelerator.ipc_collect()
+ selected_accelerator.empty_cache()
return None
diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py
index a19c014868..e7ac97c19e 100644
--- a/slime/backends/megatron_utils/loss.py
+++ b/slime/backends/megatron_utils/loss.py
@@ -771,11 +771,11 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
rewards = []
kl_coef = -args.kl_coef
cp_rank = mpu.get_context_parallel_rank()
- for reward, k in zip(old_rewards, kl, strict=False):
- k *= kl_coef
+ for reward, per_token_kl in zip(old_rewards, kl, strict=False):
+ token_level_rewards = per_token_kl * kl_coef
if cp_rank == 0:
- k[-1] += reward
- rewards.append(k)
+ token_level_rewards[-1] += reward
+ rewards.append(token_level_rewards)
advantages, returns = get_advantages_and_returns_batch(
total_lengths, response_lengths, values, rewards, args.gamma, args.lambd
)
diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py b/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py
index ca69df8e65..5f953f9179 100644
--- a/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py
+++ b/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py
@@ -5,6 +5,8 @@
import torch
import torch.nn as nn
+from slime.utils import accelerator
+
try:
import fake_int4_quant_cuda
except ImportError:
@@ -90,7 +92,7 @@ def from_linear(cls, linear, w_bit, group_size, init_only=False, scales=None, ze
awq_linear.bias = linear.bias.clone().half()
pack_num = 32 // awq_linear.w_bit
- device = torch.device(f"cuda:{torch.cuda.current_device()}")
+ device = accelerator.current_device()
repeat_scales = scales.to(device).t().repeat_interleave(group_size, 1)
if isinstance(zeros, torch.Tensor):
@@ -283,7 +285,7 @@ def quantize_params_compressed_tensors(converted_named_params, quantization_conf
qw, s, zp = pack_layer(param, group_size, is_symmetric)
qweight_name = name.replace(".weight", ".weight_packed")
scale_name = name.replace(".weight", ".weight_scale")
- weight_shape = torch.tensor(param.shape, dtype=torch.int32, device="cuda")
+ weight_shape = torch.tensor(param.shape, dtype=torch.int32, device=accelerator.device())
weight_shape_name = name.replace(".weight", ".weight_shape")
if zp is not None:
zp_name = name.replace(".weight", ".weight_zero_point")
diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py
index 1e6a53ef8a..1bb73d6b5e 100644
--- a/slime/backends/megatron_utils/model.py
+++ b/slime/backends/megatron_utils/model.py
@@ -28,11 +28,10 @@
from megatron.core.pipeline_parallel.utils import unwrap_model
except ImportError:
from megatron.core.utils import unwrap_model
-from slime.utils import logging_utils
+from slime.observability import logging_utils, train_metric_utils
from slime.utils.memory_utils import clear_memory
from .checkpoint import load_checkpoint, save_checkpoint
-from .cp_utils import reduce_train_step_metrics
from .data import DataIterator, get_batch
from .loss import ROLLOUT_TOP_P_TOKEN_KEYS, get_rollout_top_p_logprob_kwargs, loss_function
from .model_provider import get_model_provider_func
@@ -270,7 +269,7 @@ def _patch_megatron_adam(adam_cls):
def setup_model_and_optimizer(
args: Namespace,
role: str = "actor",
-) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]:
+) -> tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None]:
"""Build model(s), wrap with DDP, and construct optimizer and scheduler.
Args:
@@ -283,7 +282,7 @@ def setup_model_and_optimizer(
lr_mult (float): Global learning-rate multiplier for the optimizer.
Returns:
- tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]:
+ tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None]:
- List of model chunks wrapped by ``DDP``.
- The constructed ``MegatronOptimizer`` instance.
- The learning-rate/weight-decay scheduler tied to the optimizer.
@@ -293,6 +292,10 @@ def setup_model_and_optimizer(
model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder)
+ if args.num_rollout == 0:
+ args.no_load_optim = True
+ return model, None, None
+
# Optimizer
kwargs = {}
for f in dataclasses.fields(OptimizerConfig):
@@ -688,7 +691,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
optimizer.zero_grad()
if mpu.is_pipeline_last_stage(ignore_virtual=True):
- loss_reduced = reduce_train_step_metrics(
+ loss_reduced = train_metric_utils.reduce_train_step_metrics(
losses_reduced,
calculate_per_token_loss=args.calculate_per_token_loss,
step_global_batch_size=step_global_batch_size,
@@ -973,7 +976,7 @@ def save(
def initialize_model_and_optimizer(
args: Namespace, role: str = "actor"
-) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]:
+) -> tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None, int]:
"""Initialize model(s), optimizer, scheduler, and load from checkpoint.
Args:
@@ -981,7 +984,7 @@ def initialize_model_and_optimizer(
role (str): Logical role of the model (e.g., "actor", "critic").
Returns:
- tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]:
+ tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None, int]:
DDP-wrapped model chunks, optimizer, scheduler, and iteration index.
"""
diff --git a/slime/backends/megatron_utils/server/logprob_utils.py b/slime/backends/megatron_utils/server/logprob_utils.py
index 61c3fca5f2..3772056dbd 100644
--- a/slime/backends/megatron_utils/server/logprob_utils.py
+++ b/slime/backends/megatron_utils/server/logprob_utils.py
@@ -12,6 +12,7 @@
from slime.backends.megatron_utils.data import get_data_iterator
from slime.backends.megatron_utils.loss import get_log_probs_and_entropy, get_responses
from slime.backends.megatron_utils.model import forward_only
+from slime.utils import accelerator
logging.getLogger().setLevel(logging.WARNING)
@@ -222,17 +223,17 @@ def get_label_token_log_probs_from_vocab_parallel_logits(
return (local_selected_logits - global_max.to(reduction_dtype) - log_denom).to(logits_dtype)
-def _to_cuda_tensors(values, dtype: torch.dtype) -> list[torch.Tensor]:
- return [torch.as_tensor(value, dtype=dtype, device=torch.cuda.current_device()) for value in values]
+def _to_accelerator_tensors(values, dtype: torch.dtype) -> list[torch.Tensor]:
+ return [torch.as_tensor(value, dtype=dtype, device=accelerator.current_device()) for value in values]
def _prepare_rollout_data(rollout_data_ref):
rollout_data = ray.get(rollout_data_ref[0].inner)
- rollout_data["tokens"] = _to_cuda_tensors(rollout_data["tokens"], torch.long)
- rollout_data["loss_masks"] = _to_cuda_tensors(rollout_data["loss_masks"], torch.int)
+ rollout_data["tokens"] = _to_accelerator_tensors(rollout_data["tokens"], torch.long)
+ rollout_data["loss_masks"] = _to_accelerator_tensors(rollout_data["loss_masks"], torch.int)
if rollout_data.get("label_token_ids") is not None:
- rollout_data["label_token_ids"] = _to_cuda_tensors(rollout_data["label_token_ids"], torch.long)
+ rollout_data["label_token_ids"] = _to_accelerator_tensors(rollout_data["label_token_ids"], torch.long)
for idx, tensor in enumerate(rollout_data["label_token_ids"]):
if tensor.dim() == 1 and tensor.numel() == 0:
rollout_data["label_token_ids"][idx] = tensor.reshape(0, 0)
diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
index ae3083b1f4..67f663d276 100644
--- a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
+++ b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
@@ -7,6 +7,7 @@
from megatron.core import mpu
from tqdm import tqdm
+from slime.utils import accelerator
from slime.utils.distributed_utils import get_gloo_group
from slime.utils.types import ParamInfo
@@ -80,13 +81,13 @@ def _get_megatron_full_params(
if dist.get_rank() == info.src_rank:
params.append(
torch.nn.Parameter(
- megatron_local_weights[info.name].to(device=torch.cuda.current_device(), non_blocking=True),
+ megatron_local_weights[info.name].to(device=accelerator.current_device(), non_blocking=True),
requires_grad=False,
)
)
else:
- params.append(torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device()))
- torch.cuda.synchronize()
+ params.append(torch.empty(info.shape, dtype=info.dtype, device=accelerator.current_device()))
+ accelerator.synchronize()
# broadcast params across pp ranks
if pp_size > 1:
diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py
index 2632835096..9be3bfabb9 100644
--- a/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py
+++ b/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py
@@ -19,6 +19,7 @@
from megatron.core import mpu
from ray.actor import ActorHandle
+from slime.utils import accelerator
from slime.utils.disk_delta import NUM_WORKERS, checksum, make_tensor_reader, overwrite_encode
from slime.utils.distributed_utils import get_gloo_group
@@ -259,7 +260,7 @@ def collect(fut):
if use_pinned and nbytes <= max_bytes:
buf = free_q.get() # blocks when all buffers are in flight -> backpressures the gather
buf[:nbytes].copy_(flat, non_blocking=True)
- torch.cuda.current_stream().synchronize()
+ accelerator.current_stream().synchronize()
payload, pinned = buf, True
else:
payload, pinned = flat.cpu().numpy(), False
@@ -278,7 +279,7 @@ def _record_metrics(self) -> None:
counts = torch.tensor(
[self.changed_bytes, self.total_bytes, self.wire_bytes],
dtype=torch.int64,
- device=torch.cuda.current_device(),
+ device=accelerator.current_device(),
)
dist.all_reduce(counts)
changed, total, wire = counts.tolist()
diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
index 1ba987f06b..cd3169e38f 100644
--- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
+++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
@@ -13,6 +13,7 @@
from ray.actor import ActorHandle
from tqdm import tqdm
+from slime.utils import accelerator
from slime.utils.distributed_utils import get_gloo_group, init_process_group
from slime.utils.http_utils import _wrap_ipv6
@@ -22,8 +23,8 @@
class UpdateWeightFromDistributed:
"""
- Update distributed engines via NCCL. Each PP rank: group "slime-pp_{pp_rank}",
- only DP=TP=0 broadcasts. Non-expert (TP) and expert (EP) params separate.
+ Update distributed engines through a device process group. Each PP rank: group "slime-pp_{pp_rank}",
+ only DP=TP=0 transfers. Non-expert (TP) and expert (EP) params separate.
Subclasses override ``_send_weights`` / ``_on_chunk`` to inject per-mode behaviour.
"""
@@ -63,7 +64,7 @@ def connect_rollout_engines(
engine_parallel_configs: Sequence[Mapping[str, object]] | None = None,
) -> None:
"""
- Create NCCL "slime-pp_{pp_rank}" if PP source (DP=TP=0). Lock prevents concurrent broadcasts.
+ Create "slime-pp_{pp_rank}" if PP source (DP=TP=0). Lock prevents concurrent transfers.
"""
self.rollout_engines = rollout_engines
self.rollout_engine_lock = rollout_engine_lock
@@ -217,7 +218,7 @@ def _ep_gather_and_convert(self, named_tensors: list[tuple[str, torch.Tensor]])
handles = []
for i, (_name, param) in enumerate(named_tensors):
params = [
- torch.empty_like(param.data, device=torch.cuda.current_device())
+ torch.empty_like(param.data, device=accelerator.current_device())
for _ in range(mpu.get_expert_model_parallel_world_size())
]
handle = dist.all_gather(params, param.data, group=mpu.get_expert_model_parallel_group(), async_op=True)
@@ -244,9 +245,9 @@ def _update_bucket_weights_from_distributed(
load_format: str | None = None,
) -> None:
"""
- Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock.
+ Lock → transfer → clear → unlock → pbar++. Lock prevents communication deadlock.
"""
- # lock the rollout engines to prevent dead lock on broadcast.
+ # Lock the rollout engines to prevent communication deadlock.
while not ray.get(self.rollout_engine_lock.acquire.remote()):
time.sleep(0.1)
@@ -272,11 +273,11 @@ def connect_rollout_engines_from_distributed(
engine_gpu_counts: Sequence[int] | None = None,
) -> dist.ProcessGroup:
"""
- Create NCCL group: training rank 0 + all engine GPUs. Blocks until joined.
+ Create a device process group: training rank 0 + all engine GPUs. Blocks until joined.
``engine_gpu_counts`` gives the number of GPUs per engine. When engines
have heterogeneous TP sizes (e.g. prefill TP=2, decode TP=4), each engine
- occupies a different number of ranks in the NCCL group.
+ occupies a different number of ranks in the process group.
"""
if engine_gpu_counts is None:
engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines)
@@ -292,6 +293,7 @@ def connect_rollout_engines_from_distributed(
for c in engine_gpu_counts:
cumulative.append(cumulative[-1] + c)
+ backend = accelerator.weight_update_backend()
refs = [
engine.init_weights_update_group.remote(
master_address=master_address,
@@ -299,12 +301,12 @@ def connect_rollout_engines_from_distributed(
rank_offset=cumulative[i] + 1,
world_size=world_size,
group_name=group_name,
- backend="nccl",
+ backend=backend,
)
for i, engine in enumerate(rollout_engines)
]
model_update_groups = init_process_group(
- backend="nccl",
+ backend=backend,
init_method=f"tcp://{_wrap_ipv6(master_address)}:{master_port}",
world_size=world_size,
rank=0,
@@ -316,7 +318,7 @@ def connect_rollout_engines_from_distributed(
def disconnect_rollout_engines_from_distributed(args, group_name, model_update_groups, rollout_engines):
"""
- Destroy NCCL on training and engines.
+ Destroy the weight-update process group on training and engines.
"""
refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines]
dist.destroy_process_group(model_update_groups)
@@ -332,7 +334,7 @@ def update_weights_from_distributed(
load_format: str | None = None,
) -> list[ObjectRef]:
"""
- Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines).
+ Send metadata through Ray and tensors through the configured transport.
"""
refs = [
engine.update_weights_from_distributed.remote(
@@ -345,7 +347,6 @@ def update_weights_from_distributed(
)
for engine in rollout_engines
]
-
handles = []
for _, param in converted_named_tensors:
handles.append(dist.broadcast(param.data, 0, group=group, async_op=True))
diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
index 877febb7ca..b656781edb 100644
--- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
+++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
@@ -11,6 +11,7 @@
from ray.actor import ActorHandle
from tqdm import tqdm
+from slime.utils import accelerator
from slime.utils.distributed_utils import get_gloo_group
from slime.utils.types import ParamInfo
@@ -31,7 +32,7 @@ def _build_flattened_tensor_data(
) -> dict[str, Any]:
if not named_tensors:
return {
- "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=torch.cuda.current_device()),
+ "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=accelerator.current_device()),
"metadata": [],
}
@@ -209,7 +210,7 @@ def _prepare_expert_weight_batch(
offset = buffer_offsets[key]
buffer_offsets[key] = offset + 1
if offset == len(pool):
- pool.append(torch.empty(info.shape, dtype=info.dtype, device="cuda"))
+ pool.append(torch.empty(info.shape, dtype=info.dtype, device=accelerator.device()))
tensor = pool[offset]
if self.rank == transfer.source_rank:
source = megatron_local_weights[info.name]
@@ -266,12 +267,12 @@ def _update_expert_weights(
refs, long_lived_tensors = self._send_hf_params(hf_named_tensors)
ray.get(refs)
dist.barrier(group=get_gloo_group())
- torch.cuda.synchronize()
+ accelerator.synchronize()
del refs, long_lived_tensors, hf_named_tensors
- torch.cuda.ipc_collect()
- torch.cuda.empty_cache()
+ accelerator.ipc_collect()
+ accelerator.empty_cache()
del staging_buffers
- torch.cuda.empty_cache()
+ accelerator.empty_cache()
@torch.no_grad()
def update_weights(self) -> None:
@@ -306,8 +307,8 @@ def update_weights(self) -> None:
# then release CUDA IPC cache entries whose consumers (sglang engines)
# have already closed their IPC handles.
del refs, long_lived_tensors, hf_named_tensors
- torch.cuda.ipc_collect()
- torch.cuda.empty_cache()
+ accelerator.ipc_collect()
+ accelerator.empty_cache()
if self._expert_transfer_plan:
self._update_expert_weights(megatron_local_weights)
@@ -316,8 +317,8 @@ def update_weights(self) -> None:
dist.barrier(group=get_gloo_group())
# After the barrier all engines have returned, so every rank's last-chunk
# IPC handles are now released by the consumers. Clean them up.
- torch.cuda.ipc_collect()
- torch.cuda.empty_cache()
+ accelerator.ipc_collect()
+ accelerator.empty_cache()
# int4/fp4 post_process
if self.rank == 0:
@@ -426,6 +427,6 @@ def _send_to_colocated_engine(
def _empty_flattened_tensor_data():
return {
- "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=torch.cuda.current_device()),
+ "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=accelerator.current_device()),
"metadata": [],
}
diff --git a/slime/backends/sglang_utils/__init__.py b/slime/backends/sglang_utils/__init__.py
index e69de29bb2..1a15621488 100644
--- a/slime/backends/sglang_utils/__init__.py
+++ b/slime/backends/sglang_utils/__init__.py
@@ -0,0 +1,5 @@
+from slime.utils import accelerator
+
+# Finalize the backend before importing any SGLang submodule. In a MUSA
+# runtime this loads musa_patch only after MUSA wins backend selection.
+accelerator.initialize_accelerator()
diff --git a/slime/backends/sglang_utils/arguments.py b/slime/backends/sglang_utils/arguments.py
index 7b04fe4a43..87ea018825 100644
--- a/slime/backends/sglang_utils/arguments.py
+++ b/slime/backends/sglang_utils/arguments.py
@@ -1,9 +1,12 @@
import argparse
+import logging
from sglang.srt.server_args import ServerArgs
from sglang_router.launch_router import RouterArgs
from slime.utils.http_utils import _wrap_ipv6
+logger = logging.getLogger(__name__)
+
# TODO: use all sglang router arguments with `--sglang-router` prefix
def add_sglang_router_arguments(parser):
@@ -152,7 +155,18 @@ def validate_args(args):
("sglang_ep_size", "sglang_expert_parallel_size"),
("sglang_moe_dp_size", "sglang_moe_data_parallel_size"),
):
- value = getattr(args, current_name) if hasattr(args, current_name) else getattr(args, legacy_name)
+ if hasattr(args, current_name):
+ value = getattr(args, current_name)
+ elif hasattr(args, legacy_name):
+ value = getattr(args, legacy_name)
+ else:
+ logger.warning(
+ "The installed SGLang registered neither %s nor %s; "
+ "skipping compatibility alias normalization for this parameter.",
+ current_name,
+ legacy_name,
+ )
+ continue
setattr(args, current_name, value)
setattr(args, legacy_name, value)
diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py
index 3e562a7e38..66e1828baa 100644
--- a/slime/backends/sglang_utils/sglang_engine.py
+++ b/slime/backends/sglang_utils/sglang_engine.py
@@ -12,6 +12,7 @@
from slime.backends.sglang_utils.external import get_server_info
from slime.ray.ray_actor import RayActor
+from slime.utils import accelerator
from slime.utils.http_utils import get_host_info
logger = logging.getLogger(__name__)
@@ -27,24 +28,6 @@ def get_base_gpu_id(args, rank):
return start_index
-def _to_local_gpu_id(physical_gpu_id: int) -> int:
- cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
- if not cvd:
- return physical_gpu_id # no remapping
- # CUDA_VISIBLE_DEVICES can be like "4,5,6,7"
- visible = [int(x) for x in cvd.split(",") if x.strip() != ""]
- # In a remapped process, valid torch device indices are 0..len(visible)-1
- if physical_gpu_id in visible:
- return visible.index(physical_gpu_id)
- # If we're already getting local IDs, allow them
- if 0 <= physical_gpu_id < len(visible):
- return physical_gpu_id
- raise RuntimeError(
- f"GPU id {physical_gpu_id} is not valid under CUDA_VISIBLE_DEVICES={cvd}. "
- f"Expected one of {visible} (physical) or 0..{len(visible)-1} (local)."
- )
-
-
def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
# Expandable segments help the colocated training actor tolerate repeated
# cache releases, but SGLang's allocator/sleep path does not support them.
@@ -540,7 +523,7 @@ def _compute_server_args(
nnodes = max(1, _gpus_per_engine // args.num_gpus_per_node)
node_rank = rank % nnodes
base = base_gpu_id if base_gpu_id is not None else get_base_gpu_id(args, rank)
- base = _to_local_gpu_id(base)
+ base = accelerator.resolve_visible_device_id(base)
kwargs = {
"model_path": args.hf_checkpoint,
"trust_remote_code": True,
diff --git a/slime/observability/__init__.py b/slime/observability/__init__.py
new file mode 100644
index 0000000000..3dab479228
--- /dev/null
+++ b/slime/observability/__init__.py
@@ -0,0 +1 @@
+"""Training and rollout observability utilities."""
diff --git a/slime/utils/logging_utils.py b/slime/observability/logging_utils.py
similarity index 91%
rename from slime/utils/logging_utils.py
rename to slime/observability/logging_utils.py
index 1fc3d94b23..2ae2a222ce 100644
--- a/slime/utils/logging_utils.py
+++ b/slime/observability/logging_utils.py
@@ -2,8 +2,8 @@
import wandb
-from . import wandb_utils
-from .tensorboard_utils import _TensorboardAdapter
+from slime.observability import wandb_utils
+from slime.observability.tensorboard_utils import _TensorboardAdapter
_LOGGER_CONFIGURED = False
diff --git a/slime/utils/metric_utils.py b/slime/observability/metric_utils.py
similarity index 100%
rename from slime/utils/metric_utils.py
rename to slime/observability/metric_utils.py
diff --git a/slime/utils/profile_utils.py b/slime/observability/profile_utils.py
similarity index 70%
rename from slime/utils/profile_utils.py
rename to slime/observability/profile_utils.py
index 504d1ce868..a2f8dc0316 100644
--- a/slime/utils/profile_utils.py
+++ b/slime/observability/profile_utils.py
@@ -5,6 +5,7 @@
import torch
+from slime.utils import accelerator
from slime.utils.memory_utils import print_memory
logger = logging.getLogger(__name__)
@@ -58,7 +59,13 @@ def _profile_simple_loop(iterator, args, name):
def _create_torch_profiler(args, name):
+ activities = [torch.profiler.ProfilerActivity.CPU]
+ activity_name = accelerator.device_type().upper()
+ if hasattr(torch.profiler.ProfilerActivity, activity_name):
+ activities.append(getattr(torch.profiler.ProfilerActivity, activity_name))
+
return torch.profiler.profile(
+ activities=activities,
schedule=torch.profiler.schedule(
# TODO the train_actor and train_log_probs ones may need to have different args to control step
wait=max(args.profile_step_start - 1, 0),
@@ -101,30 +108,55 @@ def stop(self):
class _TorchMemoryProfiler(_BaseMemoryProfiler):
+ def __init__(self, args):
+ super().__init__(args)
+ self._recording = False
+
+ @staticmethod
+ def _memory_module():
+ return accelerator.memory_module()
+
def start(self):
logger.info("Attach OOM dump memory history.")
-
- torch.cuda.memory._record_memory_history(
+ memory_module = self._memory_module()
+ if memory_module is None or not hasattr(memory_module, "_record_memory_history"):
+ logger.warning("Accelerator memory history is unavailable; skip torch memory profiler.")
+ return
+ if not hasattr(memory_module, "_dump_snapshot"):
+ logger.warning("Accelerator memory snapshot is unavailable; skip torch memory profiler.")
+ return
+
+ memory_module._record_memory_history(
max_entries=1000000,
- # record stack information for the trace events
- # trace_alloc_record_context=True,
stacks="all",
)
+ self._recording = True
def oom_observer(device, alloc, device_alloc, device_free):
logger.info(
f"Observe OOM, will dump snapshot to {self._path_dump}. ({device=} {alloc=} {device_alloc=} {device_free=}; stacktrace is as follows)"
)
traceback.print_stack()
- torch.cuda.memory._dump_snapshot(self._path_dump)
+ memory_module._dump_snapshot(str(self._path_dump))
print_memory("when oom")
- torch._C._cuda_attach_out_of_memory_observer(oom_observer)
+ attach_oom_observer = getattr(torch._C, "_cuda_attach_out_of_memory_observer", None)
+ if attach_oom_observer is not None:
+ attach_oom_observer(oom_observer)
+ else:
+ logger.warning("Accelerator OOM observer is unavailable; memory snapshot on OOM is disabled.")
def stop(self):
+ if not self._recording:
+ return
logger.info(f"Dump memory snapshot to: {self._path_dump}")
- torch.cuda.memory._dump_snapshot(self._path_dump)
- torch.cuda.memory._record_memory_history(enabled=None)
+ memory_module = self._memory_module()
+ if memory_module is None or not hasattr(memory_module, "_dump_snapshot"):
+ logger.warning("Accelerator memory snapshot is unavailable; skip dump.")
+ return
+ memory_module._dump_snapshot(str(self._path_dump))
+ memory_module._record_memory_history(enabled=None)
+ self._recording = False
class _MemrayMemoryProfiler(_BaseMemoryProfiler):
diff --git a/slime/observability/rollout_data_utils.py b/slime/observability/rollout_data_utils.py
new file mode 100644
index 0000000000..de2c7db136
--- /dev/null
+++ b/slime/observability/rollout_data_utils.py
@@ -0,0 +1,153 @@
+import logging
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from slime.utils.types import Sample
+
+logger = logging.getLogger(__name__)
+
+_ROLLOUT_DATA_TENSOR_DTYPES = {
+ "tokens": torch.long,
+ "loss_masks": torch.int,
+ "rollout_log_probs": torch.float32,
+ "rollout_top_p_token_ids": torch.int32,
+ "rollout_top_p_token_offsets": torch.int32,
+ "teacher_log_probs": torch.float32,
+ "rollout_routed_experts": None,
+}
+
+
+def _cpu_tensor(value, dtype: torch.dtype | None = None) -> torch.Tensor:
+ if isinstance(value, np.ndarray) and not value.flags.writeable:
+ value = value.copy()
+ tensor = torch.as_tensor(value, dtype=dtype) if dtype is not None else torch.as_tensor(value)
+ return tensor.detach().cpu().contiguous()
+
+
+def tensorize_rollout_data_for_training(rollout_data: dict[str, Any]) -> None:
+ for key, dtype in _ROLLOUT_DATA_TENSOR_DTYPES.items():
+ if key in rollout_data:
+ rollout_data[key] = [_cpu_tensor(value, dtype=dtype) for value in rollout_data[key]]
+
+ if "multimodal_train_inputs" in rollout_data:
+ rollout_data["multimodal_train_inputs"] = [
+ (
+ {
+ key: _cpu_tensor(value) if isinstance(value, (np.ndarray, torch.Tensor)) else value
+ for key, value in mm_dict.items()
+ }
+ if mm_dict is not None
+ else None
+ )
+ for mm_dict in rollout_data["multimodal_train_inputs"]
+ ]
+
+ if "rollout_mask_sums" in rollout_data:
+ rollout_data["rollout_mask_sums"] = _cpu_tensor(
+ rollout_data["rollout_mask_sums"],
+ dtype=torch.float32,
+ )
+
+
+def validate_rollout_routed_experts_for_replay(
+ routed_experts: list[torch.Tensor],
+ args,
+) -> None:
+ """Reject incomplete PP routing captures before R3 consumes them."""
+ if not routed_experts:
+ raise ValueError("R3 is enabled but no rollout routed-experts tensors were returned.")
+
+ num_layers = int(args.num_layers)
+ topk = int(args.moe_router_topk)
+ moe_layer_freq = getattr(args, "moe_layer_freq", None)
+ if isinstance(moe_layer_freq, (list, tuple)):
+ moe_layers = [layer_id for layer_id, freq in enumerate(moe_layer_freq[:num_layers]) if int(freq) != 0]
+ else:
+ moe_layers = list(range(num_layers))
+
+ for sample_idx, experts in enumerate(routed_experts):
+ experts = torch.as_tensor(experts)
+ if experts.ndim != 3 or tuple(experts.shape[1:]) != (num_layers, topk):
+ raise ValueError(
+ "Invalid rollout routed-experts shape for R3: "
+ f"sample={sample_idx}, got={tuple(experts.shape)}, "
+ f"expected=(*, {num_layers}, {topk})."
+ )
+ if experts.shape[0] == 0:
+ raise ValueError(f"R3 sample {sample_idx} has no routed-experts rows.")
+ if topk > 1:
+ missing_layers = [layer_id for layer_id in moe_layers if not torch.count_nonzero(experts[:, layer_id, :])]
+ if missing_layers:
+ raise ValueError(
+ "R3 routed-experts capture is all zero for MoE layers "
+ f"{missing_layers} in sample {sample_idx}. This usually means "
+ "SGLang pipeline stages did not aggregate their disjoint routing "
+ "captures; refusing to replay expert 0 everywhere."
+ )
+
+
+def validate_rollout_id_annotated(node, depth=0):
+ """Walk the rollout function's nested output and validate ``rollout_id`` only
+ when a compact / subagent pattern is detected.
+
+ "Compact" = the rollout function wraps multiple training samples from one
+ rollout execution into a ``list[Sample]``. In slime's convention the
+ default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout)
+ so its leaf ``list[Sample]`` lands at depth 1 and we skip validation,
+ preserving backward compatibility. A compact rollout adds a third level:
+ ``list[list[list[Sample]]]`` (prompt × rollout × samples-from-one-rollout),
+ so the leaf ``list[Sample]`` lands at depth ≥ 2. At that point we require
+ every sibling to carry a non-None ``rollout_id`` and to share the same
+ value, so the loss reducer counts the rollout once instead of N times.
+ """
+ if isinstance(node, Sample):
+ return
+ assert isinstance(node, list), f"unexpected rollout output node type: {type(node).__name__}"
+ if node and isinstance(node[0], Sample):
+ if depth >= 2 and len(node) > 1:
+ rids = [sample.rollout_id for sample in node]
+ missing = [i for i, rollout_id in enumerate(rids) if rollout_id is None]
+ assert not missing, (
+ f"Compact rollout returned {len(node)} samples but rollout_id is unset on "
+ f"positions {missing}. Set Sample.rollout_id on every sibling so the loss "
+ "reducer can aggregate them as one rollout instead of N."
+ )
+ assert len(set(rids)) == 1, f"Sibling samples from one compact rollout must share rollout_id; got {rids}."
+ return
+ for item in node:
+ validate_rollout_id_annotated(item, depth + 1)
+
+
+def load_debug_rollout_data(path_template, *, rollout_id: int, subsample_ratio=None) -> list[Sample]:
+ data = torch.load(path_template.format(rollout_id=rollout_id), weights_only=False)["samples"]
+ data = [Sample.from_dict(sample) for sample in data]
+ if subsample_ratio is not None:
+ original_num_rows = len(data)
+ rough_subsample_num_rows = int(original_num_rows * subsample_ratio)
+ data = data[: rough_subsample_num_rows // 2] + data[-rough_subsample_num_rows // 2 :]
+ logger.info(
+ "Subsample loaded debug rollout data using ratio=%s and change num rows %s -> %s",
+ subsample_ratio,
+ original_num_rows,
+ len(data),
+ )
+ return data
+
+
+def save_debug_rollout_data(path_template, data, *, rollout_id: int, evaluation: bool) -> None:
+ if path_template is None:
+ return
+
+ path = Path(path_template.format(rollout_id=("eval_" if evaluation else "") + str(rollout_id)))
+ logger.info(f"Save debug rollout data to {path}")
+ path.parent.mkdir(parents=True, exist_ok=True)
+
+ if evaluation:
+ samples = [sample.to_dict() for info in data.values() for sample in info["samples"]]
+ else:
+ samples = [sample.to_dict() for sample in data]
+
+ torch.save({"rollout_id": rollout_id, "samples": samples}, path)
diff --git a/slime/observability/rollout_metrics.py b/slime/observability/rollout_metrics.py
new file mode 100644
index 0000000000..e7e9ca2bb6
--- /dev/null
+++ b/slime/observability/rollout_metrics.py
@@ -0,0 +1,271 @@
+import logging
+from typing import Any
+
+import numpy as np
+import torch
+
+from slime.observability import logging_utils
+from slime.observability.metric_utils import (
+ compute_pass_rate,
+ compute_rollout_step,
+ compute_statistics,
+ dict_add_prefix,
+ has_repetition,
+)
+from slime.utils.misc import group_by, load_function
+from slime.utils.types import Sample
+
+logger = logging.getLogger(__name__)
+
+_SGLANG_REQUEST_PERF_FIELDS = (
+ ("request/e2e_latency", "e2e_latency"),
+ ("request/queue_time", "queue_time"),
+ ("decode/throughput", "decode_throughput"),
+)
+_SGLANG_PREFILL_PERF_FIELDS = (
+ ("prefill/bootstrap_queue_duration", "pd_prefill_bootstrap_queue_duration"),
+ ("prefill/bootstrap_duration", "pd_prefill_bootstrap_duration"),
+ ("prefill/alloc_wait_duration", "pd_prefill_alloc_wait_duration"),
+ ("prefill/forward_duration", "pd_prefill_forward_duration"),
+ ("prefill/transfer_queue_duration", "pd_prefill_transfer_queue_duration"),
+ ("prefill/transfer_speed_gb_s", "pd_transfer_speed_gb_s"),
+ ("prefill/transfer_total_mb", "pd_transfer_total_mb"),
+ ("prefill/retry_count", "pd_prefill_retry_count"),
+)
+_SGLANG_DECODE_PERF_FIELDS = (
+ ("decode/prealloc_duration", "pd_decode_prealloc_duration"),
+ ("decode/bootstrap_duration", "pd_decode_bootstrap_duration"),
+ ("decode/alloc_wait_duration", "pd_decode_alloc_wait_duration"),
+ ("decode/transfer_duration", "pd_decode_transfer_duration"),
+ ("decode/forward_duration", "pd_decode_forward_duration"),
+)
+
+
+def compute_metrics_from_samples(args, samples):
+ response_lengths = [sample.effective_response_length for sample in samples]
+
+ log_dict = {}
+ log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/")
+ log_dict |= _compute_zero_std_metrics(args, samples)
+ log_dict |= _compute_spec_metrics(args, samples)
+ log_dict |= _compute_prefix_cache_metrics(args, samples)
+ log_dict |= _compute_reward_cat_metrics(args, samples)
+ log_dict |= _compute_top_p_kept_vocab_metrics(args, samples)
+ log_dict["repetition_frac"] = np.mean([int(has_repetition(s.response)) for s in samples]).item()
+ log_dict["truncated_ratio"] = np.mean([int(s.status == Sample.Status.TRUNCATED) for s in samples]).item()
+ return log_dict
+
+
+def compute_perf_metrics_from_samples(args, samples, rollout_time):
+ non_generation_time = [sample.non_generation_time for sample in samples]
+
+ log_dict = {}
+ log_dict["rollout_time"] = rollout_time
+ if max(non_generation_time) > 0:
+ log_dict |= dict_add_prefix(compute_statistics(non_generation_time), "non_generation_time/")
+
+ def token_perf(response_lengths, non_generation_time, key=""):
+ max_response_length = max(response_lengths)
+ if args.rollout_num_gpus:
+ log_dict[f"{key}tokens_per_gpu_per_sec"] = sum(response_lengths) / rollout_time / args.rollout_num_gpus
+ log_dict[f"longest_{key}sample_tokens_per_sec"] = max_response_length / rollout_time
+
+ if max(non_generation_time) == 0:
+ return
+
+ non_generation_time = [
+ t for t, length in zip(non_generation_time, response_lengths, strict=True) if length == max_response_length
+ ]
+ mean_non_generation_time = sum(non_generation_time) / len(non_generation_time)
+
+ log_dict[f"longest_{key}sample_non_generation_time"] = mean_non_generation_time
+ log_dict[f"longest_{key}sample_tokens_per_sec_without_non_generation"] = max_response_length / (
+ rollout_time - mean_non_generation_time
+ )
+
+ token_perf([sample.response_length for sample in samples], non_generation_time, key="")
+ token_perf([sample.effective_response_length for sample in samples], non_generation_time, key="effective_")
+ log_dict |= _compute_sglang_request_perf_metrics(samples)
+
+ return log_dict
+
+
+def _compute_sglang_request_perf_metrics(all_samples: list[Sample]):
+ attrs_by_request = list(_iter_sglang_generate_attrs(all_samples))
+ if not attrs_by_request:
+ return {}
+
+ values_by_metric: dict[str, list[float]] = {}
+ profiled_request_count = 0
+
+ def add_value(metric_key: str, source_key: str, attrs: dict) -> bool:
+ value = attrs.get(source_key)
+ if not isinstance(value, (int, float)) or isinstance(value, bool) or not np.isfinite(value):
+ return False
+ values_by_metric.setdefault(metric_key, []).append(float(value))
+ return True
+
+ for attrs in attrs_by_request:
+ request_has_perf = False
+
+ for metric_key, source_key in _SGLANG_REQUEST_PERF_FIELDS:
+ request_has_perf |= add_value(metric_key, source_key, attrs)
+
+ for metric_key, source_key in _SGLANG_PREFILL_PERF_FIELDS:
+ request_has_perf |= add_value(metric_key, source_key, attrs)
+
+ for metric_key, source_key in _SGLANG_DECODE_PERF_FIELDS:
+ request_has_perf |= add_value(metric_key, source_key, attrs)
+
+ if request_has_perf:
+ profiled_request_count += 1
+
+ metrics: dict[str, float] = {}
+ for key, values in values_by_metric.items():
+ if not values:
+ continue
+ metrics |= dict_add_prefix(compute_statistics(values), f"{key}/")
+
+ return metrics
+
+
+def _iter_sglang_generate_attrs(all_samples: list[Sample]):
+ for sample in all_samples:
+ trace = getattr(sample, "trace", None)
+ if not isinstance(trace, dict):
+ continue
+ for event in trace.get("events") or []:
+ if event.get("type") != "span_end" or event.get("name") != "sglang_generate":
+ continue
+ attrs = event.get("attrs")
+ if isinstance(attrs, dict):
+ yield attrs
+
+
+def _compute_zero_std_metrics(args, all_samples: list[Sample]):
+ # only compute in GRPO-like algorithms where one prompt has multiple responses
+ if args.advantage_estimator == "ppo":
+ return {}
+
+ def _is_zero_std(samples: list[Sample]):
+ rewards = [sample.get_reward_value(args) for sample in samples]
+ return len(rewards) == 0 or all(rewards[0] == r for r in rewards)
+
+ all_sample_groups = group_by(all_samples, lambda s: s.group_index)
+ interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)]
+
+ interesting_rewards = [str(round(g[0].get_reward_value(args), 1)) for g in interesting_sample_groups]
+
+ return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()}
+
+
+def _compute_top_p_kept_vocab_metrics(args, all_samples: list[Sample]):
+ total_kept = 0
+ total_tokens = 0
+ for sample in all_samples:
+ offsets = sample.rollout_top_p_token_offsets
+ if offsets is None or sample.response_length == 0:
+ continue
+ offsets = torch.as_tensor(offsets, dtype=torch.int64)
+ if offsets.numel() == 0:
+ continue
+ assert (
+ offsets.numel() == sample.response_length + 1
+ ), f"top-p token offsets length {offsets.numel()} != response length + 1 {sample.response_length + 1}"
+ if sample.remove_sample:
+ continue
+ if sample.loss_mask is None:
+ total_kept += int(offsets[-1] - offsets[0])
+ total_tokens += sample.response_length
+ continue
+ loss_mask = torch.as_tensor(sample.loss_mask, dtype=torch.bool, device=offsets.device)
+ assert (
+ loss_mask.numel() == sample.response_length
+ ), f"loss mask length {loss_mask.numel()} != response length {sample.response_length}"
+ total_kept += int(torch.diff(offsets)[loss_mask].sum())
+ total_tokens += int(loss_mask.sum())
+ if total_tokens == 0:
+ return {}
+ return {"top_p_kept_vocab_per_token": total_kept / total_tokens}
+
+
+def _compute_spec_metrics(args, all_samples: list[Sample]):
+ if getattr(args, "sglang_speculative_algorithm", None) is None:
+ return {}
+ num_samples = len(all_samples)
+ metrics = {}
+ metrics["spec_accept_rate"] = sum(sample.spec_info.spec_accept_rate for sample in all_samples) / num_samples
+ metrics["spec_accept_length"] = sum(sample.spec_info.spec_accept_length for sample in all_samples) / num_samples
+ return metrics
+
+
+def _compute_prefix_cache_metrics(args, all_samples: list[Sample]):
+ num_samples = len(all_samples)
+ metrics = {}
+ total_cached_tokens = sum(sample.prefix_cache_info.cached_tokens for sample in all_samples)
+ total_prompt_tokens = sum(sample.prefix_cache_info.total_prompt_tokens for sample in all_samples)
+
+ metrics["prefix_cache_hit_rate"] = total_cached_tokens / total_prompt_tokens if total_prompt_tokens > 0 else 0.0
+ metrics["avg_cached_tokens_per_sample"] = total_cached_tokens / num_samples
+ return metrics
+
+
+def _compute_reward_cat_metrics(args, all_samples: list[Sample]):
+ reward_cat_key = args.log_reward_category
+ if reward_cat_key is None:
+ return {}
+
+ samples_of_reward_cat = group_by(all_samples, lambda s: s.reward[reward_cat_key])
+
+ return {f"error_cat/{reward_cat}": len(s) / len(all_samples) for reward_cat, s in samples_of_reward_cat.items()}
+
+
+def log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None):
+ if args.custom_eval_rollout_log_function_path is not None:
+ custom_log_func = load_function(args.custom_eval_rollout_log_function_path)
+ if custom_log_func(rollout_id, args, data, extra_metrics):
+ return
+
+ log_dict = extra_metrics or {}
+ for key in data.keys():
+ rewards = data[key]["rewards"]
+ log_dict[f"eval/{key}"] = sum(rewards) / len(rewards)
+ if (samples := data[key].get("samples")) is not None:
+ log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), f"eval/{key}/")
+ if "truncated" in data[key]:
+ truncated = data[key]["truncated"]
+ log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated)
+ if args.log_passrate:
+ log_dict |= dict_add_prefix(
+ compute_pass_rate(
+ flat_rewards=rewards,
+ group_size=args.n_samples_per_eval_prompt,
+ ),
+ f"eval/{key}-",
+ )
+
+ logger.info(f"eval {rollout_id}: {log_dict}")
+
+ step = compute_rollout_step(args, rollout_id)
+ log_dict["eval/step"] = step
+ logging_utils.log(args, log_dict, step_key="eval/step")
+
+ return log_dict
+
+
+def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time):
+ if args.custom_rollout_log_function_path is not None:
+ custom_log_func = load_function(args.custom_rollout_log_function_path)
+ if custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time):
+ return
+
+ if args.load_debug_rollout_data:
+ return
+
+ log_dict = {**(rollout_extra_metrics or {})}
+ log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), "rollout/")
+ log_dict |= dict_add_prefix(compute_perf_metrics_from_samples(args, samples, rollout_time), "perf/")
+ logger.info(f"perf {rollout_id}: {log_dict}")
+ step = compute_rollout_step(args, rollout_id)
+ log_dict["rollout/step"] = step
+ logging_utils.log(args, log_dict, step_key="rollout/step")
diff --git a/slime/utils/tensorboard_utils.py b/slime/observability/tensorboard_utils.py
similarity index 96%
rename from slime/utils/tensorboard_utils.py
rename to slime/observability/tensorboard_utils.py
index 3c384ec12a..250864636c 100644
--- a/slime/utils/tensorboard_utils.py
+++ b/slime/observability/tensorboard_utils.py
@@ -22,7 +22,7 @@ class _TensorboardAdapter(metaclass=SingletonMeta):
# tb.log({"Loss": 0.1}, step=1)
# In other files:
- # from tensorboard_utils import _TensorboardAdapter
+ # from slime.observability.tensorboard_utils import _TensorboardAdapter
# tb = _TensorboardAdapter(args) # No parameters needed to get existing instance
# tb.log({"Accuracy": 0.9}, step=1)
"""
diff --git a/slime/utils/timer.py b/slime/observability/timer.py
similarity index 98%
rename from slime/utils/timer.py
rename to slime/observability/timer.py
index ec1bdf767d..b0a89a5389 100644
--- a/slime/utils/timer.py
+++ b/slime/observability/timer.py
@@ -5,7 +5,7 @@
import torch.distributed
-from .misc import SingletonMeta
+from slime.utils.misc import SingletonMeta
__all__ = ["Timer", "timer"]
diff --git a/slime/utils/trace_utils.py b/slime/observability/trace_utils.py
similarity index 100%
rename from slime/utils/trace_utils.py
rename to slime/observability/trace_utils.py
diff --git a/slime/backends/megatron_utils/train_dump_utils.py b/slime/observability/train_data_utils.py
similarity index 100%
rename from slime/backends/megatron_utils/train_dump_utils.py
rename to slime/observability/train_data_utils.py
diff --git a/slime/observability/train_metric_utils.py b/slime/observability/train_metric_utils.py
new file mode 100644
index 0000000000..6a5a0f6874
--- /dev/null
+++ b/slime/observability/train_metric_utils.py
@@ -0,0 +1,405 @@
+import logging
+from argparse import Namespace
+from copy import deepcopy
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from slime.observability import logging_utils
+from slime.observability.metric_utils import compute_pass_rate, compute_rollout_step
+from slime.observability.timer import Timer
+from slime.utils.flops_utils import calculate_fwd_flops
+from slime.utils.types import RolloutBatch
+
+logger = logging.getLogger(__name__)
+
+
+def reduce_train_step_metrics(
+ losses_reduced: list[dict],
+ *,
+ calculate_per_token_loss: bool,
+ step_global_batch_size: int,
+ cp_size: int,
+ dp_with_cp_group,
+) -> dict[str, float]:
+ """Aggregate per-mb log dicts into the dict ``train_one_step`` reports.
+
+ Pipeline (1:1 with what the train loop used to do inline):
+ 1. Sum each metric's per-mb ``values`` tensor locally on this rank.
+ 2. All-reduce across the DP*CP group (``dp_with_cp_group``).
+ 3. Apply the per-mode divisor / cp_factor:
+ - per-token-loss: divisor = ``values[0]`` = all-reduced ``num_tokens``,
+ CP-inflated by ``cp_size`` because every CP rank computes the same
+ num_tokens off the FULL (not chunked) masks; the
+ ``cp_factor = cp_size`` multiplier cancels that inflation, leaving
+ the genuine per-token average.
+ - per-rollout-mean: divisor = constant ``step_global_batch_size`` from
+ the rollout side, never all-reduced, so no CP inflation to cancel
+ and ``cp_factor = 1``.
+
+ Tests pass a mock ``dp_with_cp_group`` and monkeypatch ``dist.all_reduce``
+ to a no-op, then pre-aggregate virtual ranks themselves — this exercises
+ the same call shape as production while staying single-process.
+ """
+ keys = losses_reduced[0]["keys"]
+ values = None
+ for item in losses_reduced:
+ values = item["values"] if values is None else values + item["values"]
+ assert len(keys) + 1 == values.numel()
+ dist.all_reduce(values, group=dp_with_cp_group)
+ values = values.tolist()
+
+ if calculate_per_token_loss:
+ num_samples_or_tokens = values[0]
+ cp_factor = cp_size
+ else:
+ num_samples_or_tokens = step_global_batch_size
+ cp_factor = 1
+ return {key: value * cp_factor / num_samples_or_tokens for key, value in zip(keys, values[1:], strict=False)}
+
+
+def rollout_log_metric_contribution(
+ per_rank_reducer_sum: float,
+ *,
+ cp_size: int,
+ num_rollouts_in_rollout: int,
+ dp_size: int,
+) -> tuple[float, float]:
+ """``(sum, count)`` tuple for a per-rollout-mean metric.
+
+ Sum across DP*CP ranks of ``count`` lands on ``num_rollouts_in_rollout``
+ (``dp_size`` here is the no-CP DP width; the gather covers ``dp_size *
+ cp_size`` ranks, and each rank emits the same ``count``, so the totals
+ cancel out the ``cp_size`` in the sum). Result: ``Σsum / Σcount =
+ sum_DP_full / num_rollouts`` — the same number ``train_one_step`` reports
+ for the same samples (when ``num_steps_per_rollout == 1``).
+
+ Pair with :func:`gather_and_reduce_log_dict` to do the full end-to-end
+ in tests.
+ """
+ sum_value = cp_size * per_rank_reducer_sum
+ count = num_rollouts_in_rollout / dp_size
+ return sum_value, count
+
+
+def gather_and_reduce_log_dict(
+ log_dict: dict,
+ *,
+ dp_size: int,
+ dp_src_rank: int,
+ dp_group,
+) -> dict | None:
+ """Gather per-rank log dicts and reduce each metric on ``dp_src_rank``.
+
+ ``(sum, count)`` tuples reduce to ``Σsum / Σcount``; plain values reduce
+ to a mean across ranks. The helper stays free of reporting side effects so
+ CPU multi-process tests can exercise it with real ``torch.distributed``.
+ """
+ if dist.get_rank() == dp_src_rank:
+ gathered = [None] * dp_size
+ dist.gather_object(log_dict, gathered, dst=dp_src_rank, group=dp_group)
+ reduced: dict = {}
+ for key in log_dict:
+ values = [item[key] for item in gathered]
+ first = values[0]
+ if isinstance(first, tuple) and len(first) == 2:
+ total_sum = sum(value[0] for value in values)
+ total_count = sum(value[1] for value in values)
+ reduced[key] = total_sum / total_count if total_count else 0.0
+ else:
+ reduced[key] = sum(values) / dp_size
+ return reduced
+ dist.gather_object(log_dict, None, dst=dp_src_rank, group=dp_group)
+ return None
+
+
+def gather_log_data(
+ metric_name: str,
+ args: Namespace,
+ rollout_id: int,
+ log_dict: dict[str, "float | tuple[float, float]"],
+) -> dict[str, float] | None:
+ """Gather per-rank metrics and report them through the configured trackers."""
+ from megatron.core import mpu
+
+ reduced = gather_and_reduce_log_dict(
+ log_dict,
+ dp_size=mpu.get_data_parallel_world_size(with_context_parallel=True),
+ dp_src_rank=mpu.get_data_parallel_src_rank(with_context_parallel=True),
+ dp_group=mpu.get_data_parallel_group_gloo(with_context_parallel=True),
+ )
+ if reduced is None:
+ return None
+ reduced_log_dict = {f"{metric_name}/{key}": value for key, value in reduced.items()}
+ logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}")
+ step = compute_rollout_step(args, rollout_id)
+ reduced_log_dict["rollout/step"] = step
+ logging_utils.log(args, reduced_log_dict, step_key="rollout/step")
+ return reduced_log_dict
+
+
+def log_rollout_data(
+ rollout_id: int,
+ args: Namespace,
+ rollout_data: RolloutBatch,
+) -> None:
+ """Summarize and report Megatron-side rollout fields."""
+ from megatron.core import mpu
+
+ from slime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean
+
+ if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
+ cp_size = mpu.get_context_parallel_world_size()
+ log_dict = {}
+ response_lengths = rollout_data["response_lengths"]
+ loss_masks = rollout_data["loss_masks"]
+ total_lengths = rollout_data["total_lengths"]
+ rollout_mask_sums = rollout_data.get("rollout_mask_sums", None)
+ dp_world = mpu.get_data_parallel_world_size(with_context_parallel=False)
+ num_rollouts_in_rollout = sum(rollout_data["global_batch_sizes"])
+
+ ignored_keys = {
+ "tokens",
+ "multimodal_train_inputs",
+ "loss_masks",
+ "sample_indices",
+ "rollout_ids",
+ "rollout_mask_sums",
+ "rollout_top_p_token_ids",
+ "rollout_top_p_token_offsets",
+ "rollout_routed_experts",
+ "global_batch_sizes",
+ "num_microbatches",
+ "micro_batch_indices",
+ "source_names",
+ "local_raw_reward",
+ }
+ per_rollout_mean_keys = {
+ "log_probs",
+ "ref_log_probs",
+ "rollout_log_probs",
+ "returns",
+ "advantages",
+ "values",
+ "teacher_log_probs",
+ "opd_reverse_kl",
+ }
+
+ for key, value in rollout_data.items():
+ if key in ignored_keys:
+ continue
+ if isinstance(value, (list, tuple)):
+ count = len(value)
+ if isinstance(value[0], torch.Tensor):
+ tensor = torch.cat(value).clone().detach()
+ if key in per_rollout_mean_keys:
+ sum_of_sample_mean = get_sum_of_sample_mean(
+ total_lengths,
+ response_lengths,
+ loss_masks,
+ rollout_mask_sums,
+ )
+ sum_value, count = rollout_log_metric_contribution(
+ sum_of_sample_mean(tensor).item(),
+ cp_size=cp_size,
+ num_rollouts_in_rollout=num_rollouts_in_rollout,
+ dp_size=dp_world,
+ )
+ log_dict[key] = (sum_value, count)
+ continue
+ per_rank_sum = tensor.mean() * cp_size * count
+ sum_value = per_rank_sum.item()
+ else:
+ sum_value = sum(value)
+ log_dict[key] = (sum_value, count)
+ elif isinstance(value, torch.Tensor):
+ log_dict[key] = (value.float().mean().item(), 1)
+ else:
+ raise ValueError(f"Unsupported type: {type(value)} for key: {key}")
+
+ reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict)
+ if args.ci_test and reduced_log_dict is not None:
+ if (
+ rollout_id == 0
+ and not getattr(args, "ci_disable_kl_checker", False)
+ and not getattr(args, "use_rollout_routing_replay", False)
+ and "rollout/log_probs" in reduced_log_dict
+ and "rollout/ref_log_probs" in reduced_log_dict
+ ):
+ assert abs(reduced_log_dict["rollout/log_probs"] - reduced_log_dict["rollout/ref_log_probs"]) < 1e-8
+ if "rollout/log_probs" in reduced_log_dict:
+ assert -1 < reduced_log_dict["rollout/log_probs"] < 0
+ if "rollout/entropy" in reduced_log_dict:
+ assert 0 < reduced_log_dict["rollout/entropy"] < 1
+
+ if args.log_multi_turn:
+ log_multi_turn_data(rollout_id, args, rollout_data)
+ if args.log_passrate:
+ log_passrate(rollout_id, args, rollout_data)
+
+ if args.log_correct_samples and mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
+ response_lengths = rollout_data["response_lengths"]
+ loss_masks = rollout_data["loss_masks"]
+ total_lengths = rollout_data["total_lengths"]
+
+ def quantile(total_value, n_quantiles, data) -> dict:
+ import math
+
+ assert n_quantiles > 1, f"n_quantiles({n_quantiles}) must be greater than 1."
+ quantiles = [(i + 1) / n_quantiles for i in range(n_quantiles)]
+ cut_points = [total_value * quantile for quantile in quantiles]
+ cut_points[-1] = total_value
+
+ count = [0] * n_quantiles
+ for value in data:
+ for i, point in enumerate(cut_points):
+ if value <= point:
+ count[i] += 1
+ break
+
+ total = sum(count) + 1e-9
+ percentile = [value / total for value in count]
+ return {
+ f"p{min(math.ceil(quantile * 100), 100)}": value
+ for quantile, value in zip(quantiles, percentile, strict=True)
+ }
+
+ raw_rewards = rollout_data["local_raw_reward"]
+ correct_response_lengths = []
+ correct_total_lengths = []
+ correct_loss_masks = []
+ correct_entropy = []
+ for i, raw_reward in enumerate(raw_rewards):
+ if raw_reward == 1:
+ correct_response_lengths.append(response_lengths[i])
+ correct_total_lengths.append(total_lengths[i])
+ correct_loss_masks.append(loss_masks[i])
+ correct_entropy.append(-rollout_data["log_probs"][i])
+ num_correct_responses = len(correct_total_lengths)
+ rollout_data["correct_response_lengths"] = correct_response_lengths
+ correct_response_length_percentile = quantile(
+ args.rollout_max_response_len,
+ 4,
+ rollout_data["correct_response_lengths"],
+ )
+ for percentile, value in correct_response_length_percentile.items():
+ rollout_data[f"correct_length/{percentile}"] = [value] * num_correct_responses
+ if correct_entropy:
+ sum_of_sample_mean = get_sum_of_sample_mean(
+ correct_total_lengths,
+ correct_response_lengths,
+ correct_loss_masks,
+ sample_denoms=None,
+ )
+ correct_entropy_value = sum_of_sample_mean(torch.cat(correct_entropy, dim=0))
+ rollout_data["correct_entropy"] = [correct_entropy_value.item()] * num_correct_responses
+ else:
+ rollout_data["correct_entropy"] = [0] * num_correct_responses
+
+
+def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
+ """Report multi-turn response-length and round-count metrics."""
+ from megatron.core import mpu
+
+ if mpu.get_tensor_model_parallel_rank() != 0 or not mpu.is_pipeline_last_stage():
+ return
+
+ log_dict = {}
+ for key, value in rollout_data.items():
+ if key == "loss_masks" and value:
+ device = value[0].device
+ raw_response_lengths = torch.tensor(
+ [item.shape[0] for item in value],
+ dtype=torch.float32,
+ device=device,
+ )
+ log_dict["raw_response_length/response_length_mean"] = raw_response_lengths.mean().item()
+ log_dict["raw_response_length/response_length_max"] = raw_response_lengths.max().item()
+ log_dict["raw_response_length/response_length_min"] = raw_response_lengths.min().item()
+ log_dict["raw_response_length/response_length_clip_ratio"] = (
+ (raw_response_lengths >= args.rollout_max_response_len).float().mean().item()
+ )
+
+ wo_obs_response_lengths = torch.tensor(
+ [item.sum().item() for item in value],
+ dtype=torch.float32,
+ device=device,
+ )
+ log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item()
+ log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item()
+ log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item()
+ if key == "round_number":
+ round_number_array = np.array(value)
+ log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array)
+ log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array)
+ log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array)
+ gather_log_data("multi_turn", args, rollout_id, log_dict)
+
+
+def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
+ """Compute and report pass@k metrics from grouped ``raw_reward`` values."""
+ from megatron.core import mpu
+
+ if mpu.get_tensor_model_parallel_rank() != 0 or not mpu.is_pipeline_last_stage():
+ return
+
+ log_dict = {}
+ for key, value in rollout_data.items():
+ if key == "raw_reward":
+ log_dict |= compute_pass_rate(
+ flat_rewards=value,
+ group_size=args.n_samples_per_prompt,
+ num_groups=args.rollout_batch_size,
+ )
+ gather_log_data("passrate", args, rollout_id, log_dict)
+
+
+def log_perf_data(
+ rollout_id: int,
+ args: Namespace,
+ extra_metrics: dict | None = None,
+) -> None:
+ from megatron.core import mpu
+
+ timer_instance = Timer()
+ log_dict_raw = deepcopy(timer_instance.log_dict())
+ timer_instance.reset()
+
+ if not (
+ mpu.get_tensor_model_parallel_rank() == 0
+ and mpu.is_pipeline_last_stage()
+ and mpu.get_data_parallel_rank(with_context_parallel=True) == 0
+ ):
+ return
+
+ log_dict = {f"perf/{key}_time": val for key, val in log_dict_raw.items()}
+ if extra_metrics:
+ log_dict.update(extra_metrics)
+
+ if "perf/actor_train_time" in log_dict:
+ total_fwd_flops = (
+ calculate_fwd_flops(seqlens=timer_instance.seq_lens, args=args) / dist.get_world_size() / 1e12
+ )
+
+ if "perf/log_probs_time" in log_dict:
+ log_dict["perf/log_probs_tflops"] = total_fwd_flops / log_dict["perf/log_probs_time"]
+
+ if "perf/ref_log_probs_time" in log_dict:
+ log_dict["perf/ref_log_probs_tflops"] = total_fwd_flops / log_dict["perf/ref_log_probs_time"]
+
+ if log_dict["perf/actor_train_time"] > 0:
+ log_dict["perf/actor_train_tflops"] = 3 * total_fwd_flops / log_dict["perf/actor_train_time"]
+ log_dict["perf/actor_train_tok_per_s"] = sum(timer_instance.seq_lens) / log_dict["perf/actor_train_time"]
+
+ if "perf/train_wait_time" in log_dict and "perf/train_time" in log_dict:
+ total_time = log_dict["perf/train_wait_time"] + log_dict["perf/train_time"]
+ if total_time > 0:
+ log_dict["perf/step_time"] = total_time
+ log_dict["perf/wait_time_ratio"] = log_dict["perf/train_wait_time"] / total_time
+
+ logger.info(f"perf {rollout_id}: {log_dict}")
+
+ step = compute_rollout_step(args, rollout_id)
+ log_dict["rollout/step"] = step
+ logging_utils.log(args, log_dict, step_key="rollout/step")
diff --git a/slime/utils/wandb_utils.py b/slime/observability/wandb_utils.py
similarity index 100%
rename from slime/utils/wandb_utils.py
rename to slime/observability/wandb_utils.py
diff --git a/slime/ray/placement_group.py b/slime/ray/placement_group.py
index eb652c3af9..2617b1f62a 100644
--- a/slime/ray/placement_group.py
+++ b/slime/ray/placement_group.py
@@ -187,7 +187,7 @@ def create_training_models(args, pgs, rollout_manager, actor_cls=None):
actor_model, actor_start_rollout_ids = create_actor_model(args, pgs, rollout_manager, actor_cls=actor_cls)
critic_model = None
- if args.use_critic:
+ if args.use_critic and args.num_rollout != 0:
from slime.utils.arguments import parse_megatron_role_args
critic_args = (
@@ -208,7 +208,7 @@ def create_training_models(args, pgs, rollout_manager, actor_cls=None):
critic_start_rollout_ids = critic_model.create(rollout_manager=rollout_manager)
# TODO how to decide rollout start id when critic is involved? For now we just require user to specify it via args.
- if args.use_critic:
+ if critic_model is not None:
start_rollout_ids = critic_start_rollout_ids
else:
start_rollout_ids = actor_start_rollout_ids
diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py
index abd8e35d70..246e828fb1 100644
--- a/slime/ray/rollout.py
+++ b/slime/ray/rollout.py
@@ -5,10 +5,8 @@
import os
import random
import time
-from pathlib import Path
from typing import Any
-import numpy as np
import ray
import torch
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
@@ -17,19 +15,25 @@
from slime.backends.sglang_utils.external import start_external_rollout_servers
from slime.backends.sglang_utils.sglang_config import ModelConfig, ServerGroupConfig, SglangConfig
from slime.backends.sglang_utils.sglang_engine import SGLangEngine
+from slime.observability import logging_utils
+from slime.observability.logging_utils import configure_logger, init_tracking
+from slime.observability.rollout_data_utils import (
+ load_debug_rollout_data,
+ save_debug_rollout_data,
+ tensorize_rollout_data_for_training,
+ validate_rollout_id_annotated,
+ validate_rollout_routed_experts_for_replay,
+)
+from slime.observability.rollout_metrics import log_eval_rollout_data, log_rollout_data
from slime.rollout.base_types import call_rollout_fn
from slime.rollout.sample_hooks import set_current_rollout_id
-from slime.utils import logging_utils
from slime.utils.data import get_source
from slime.utils.dp_schedule import build_dp_schedule
from slime.utils.health_monitor import RolloutHealthMonitor
from slime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client
-from slime.utils.logging_utils import configure_logger, init_tracking
-from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step, compute_statistics, dict_add_prefix
-from slime.utils.misc import Box, group_by, load_function
+from slime.utils.misc import Box, load_function
from slime.utils.types import Sample
-from ..utils.metric_utils import has_repetition
from .rollout_validation import validate_server_group_gpu_indices
from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock, add_default_ray_env_vars
@@ -38,108 +42,6 @@
logger = logging.getLogger(__name__)
-_ROLLOUT_DATA_TENSOR_DTYPES = {
- "tokens": torch.long,
- "loss_masks": torch.int,
- "rollout_log_probs": torch.float32,
- "rollout_top_p_token_ids": torch.int32,
- "rollout_top_p_token_offsets": torch.int32,
- "teacher_log_probs": torch.float32,
- "rollout_routed_experts": None,
-}
-
-_SGLANG_REQUEST_PERF_FIELDS = (
- ("request/e2e_latency", "e2e_latency"),
- ("request/queue_time", "queue_time"),
- ("decode/throughput", "decode_throughput"),
-)
-_SGLANG_PREFILL_PERF_FIELDS = (
- ("prefill/bootstrap_queue_duration", "pd_prefill_bootstrap_queue_duration"),
- ("prefill/bootstrap_duration", "pd_prefill_bootstrap_duration"),
- ("prefill/alloc_wait_duration", "pd_prefill_alloc_wait_duration"),
- ("prefill/forward_duration", "pd_prefill_forward_duration"),
- ("prefill/transfer_queue_duration", "pd_prefill_transfer_queue_duration"),
- ("prefill/transfer_speed_gb_s", "pd_transfer_speed_gb_s"),
- ("prefill/transfer_total_mb", "pd_transfer_total_mb"),
- ("prefill/retry_count", "pd_prefill_retry_count"),
-)
-_SGLANG_DECODE_PERF_FIELDS = (
- ("decode/prealloc_duration", "pd_decode_prealloc_duration"),
- ("decode/bootstrap_duration", "pd_decode_bootstrap_duration"),
- ("decode/alloc_wait_duration", "pd_decode_alloc_wait_duration"),
- ("decode/transfer_duration", "pd_decode_transfer_duration"),
- ("decode/forward_duration", "pd_decode_forward_duration"),
-)
-
-
-def _cpu_tensor(value, dtype: torch.dtype | None = None) -> torch.Tensor:
- if isinstance(value, np.ndarray) and not value.flags.writeable:
- value = value.copy()
- tensor = torch.as_tensor(value, dtype=dtype) if dtype is not None else torch.as_tensor(value)
- return tensor.detach().cpu().contiguous()
-
-
-def _tensorize_rollout_data_for_training(rollout_data: dict[str, Any]) -> None:
- for key, dtype in _ROLLOUT_DATA_TENSOR_DTYPES.items():
- if key in rollout_data:
- rollout_data[key] = [_cpu_tensor(value, dtype=dtype) for value in rollout_data[key]]
-
- if "multimodal_train_inputs" in rollout_data:
- rollout_data["multimodal_train_inputs"] = [
- (
- {
- key: _cpu_tensor(value) if isinstance(value, (np.ndarray, torch.Tensor)) else value
- for key, value in mm_dict.items()
- }
- if mm_dict is not None
- else None
- )
- for mm_dict in rollout_data["multimodal_train_inputs"]
- ]
-
- if "rollout_mask_sums" in rollout_data:
- rollout_data["rollout_mask_sums"] = _cpu_tensor(
- rollout_data["rollout_mask_sums"],
- dtype=torch.float32,
- )
-
-
-def _validate_rollout_routed_experts_for_replay(
- routed_experts: list[torch.Tensor],
- args,
-) -> None:
- """Reject incomplete PP routing captures before R3 consumes them."""
- if not routed_experts:
- raise ValueError("R3 is enabled but no rollout routed-experts tensors were returned.")
-
- num_layers = int(args.num_layers)
- topk = int(args.moe_router_topk)
- moe_layer_freq = getattr(args, "moe_layer_freq", None)
- if isinstance(moe_layer_freq, (list, tuple)):
- moe_layers = [layer_id for layer_id, freq in enumerate(moe_layer_freq[:num_layers]) if int(freq) != 0]
- else:
- moe_layers = list(range(num_layers))
-
- for sample_idx, experts in enumerate(routed_experts):
- experts = torch.as_tensor(experts)
- if experts.ndim != 3 or tuple(experts.shape[1:]) != (num_layers, topk):
- raise ValueError(
- "Invalid rollout routed-experts shape for R3: "
- f"sample={sample_idx}, got={tuple(experts.shape)}, "
- f"expected=(*, {num_layers}, {topk})."
- )
- if experts.shape[0] == 0:
- raise ValueError(f"R3 sample {sample_idx} has no routed-experts rows.")
- if topk > 1:
- missing_layers = [layer_id for layer_id in moe_layers if not torch.count_nonzero(experts[:, layer_id, :])]
- if missing_layers:
- raise ValueError(
- "R3 routed-experts capture is all zero for MoE layers "
- f"{missing_layers} in sample {sample_idx}. This usually means "
- "SGLang pipeline stages did not aggregate their disjoint routing "
- "captures; refusing to replay expert 0 everywhere."
- )
-
@dataclasses.dataclass
class ServerGroup:
@@ -595,8 +497,13 @@ def generate(self, rollout_id):
if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2:
self._try_ci_fault_injection()
data, metrics = self._get_rollout_data(rollout_id=rollout_id)
- self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=False)
- _log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time)
+ save_debug_rollout_data(
+ self.args.save_debug_rollout_data,
+ data,
+ rollout_id=rollout_id,
+ evaluation=False,
+ )
+ log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time)
if self.args.debug_rollout_only:
# if debug rollout only, we don't convert samples to train data and directly return
return
@@ -612,8 +519,13 @@ def eval(self, rollout_id):
result = call_rollout_fn(self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True)
data = result.data
- self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=True)
- _log_eval_rollout_data(rollout_id, self.args, data, result.metrics)
+ save_debug_rollout_data(
+ self.args.save_debug_rollout_data,
+ data,
+ rollout_id=rollout_id,
+ evaluation=True,
+ )
+ log_eval_rollout_data(rollout_id, self.args, data, result.metrics)
def save(self, rollout_id):
self.data_source.save(rollout_id)
@@ -670,18 +582,11 @@ def check_weights(self, action: str):
def _get_rollout_data(self, rollout_id):
if self.args.load_debug_rollout_data:
- data = torch.load(
- self.args.load_debug_rollout_data.format(rollout_id=rollout_id),
- weights_only=False,
- )["samples"]
- data = [Sample.from_dict(sample) for sample in data]
- if (ratio := self.args.load_debug_rollout_data_subsample) is not None:
- original_num_rows = len(data)
- rough_subsample_num_rows = int(original_num_rows * ratio)
- data = data[: rough_subsample_num_rows // 2] + data[-rough_subsample_num_rows // 2 :]
- logger.info(
- f"Subsample loaded debug rollout data using {ratio=} and change num rows {original_num_rows} -> {len(data)}"
- )
+ data = load_debug_rollout_data(
+ self.args.load_debug_rollout_data,
+ rollout_id=rollout_id,
+ subsample_ratio=self.args.load_debug_rollout_data_subsample,
+ )
metrics = None
else:
data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False)
@@ -693,32 +598,13 @@ def _get_rollout_data(self, rollout_id):
# subagent paths that split one rollout into N training samples must
# set the same rollout_id on every sibling so the loss reducer counts
# the rollout once instead of N times.
- _validate_rollout_id_annotated(data)
+ validate_rollout_id_annotated(data)
# flatten the data if it is a list of lists
while isinstance(data[0], list):
data = list(itertools.chain.from_iterable(data))
return data, metrics
- def _save_debug_rollout_data(self, data, rollout_id, evaluation: bool):
- # TODO to be refactored (originally Buffer._set_data)
- if (path_template := self.args.save_debug_rollout_data) is not None:
- path = Path(path_template.format(rollout_id=("eval_" if evaluation else "") + str(rollout_id)))
- logger.info(f"Save debug rollout data to {path}")
- path.parent.mkdir(parents=True, exist_ok=True)
-
- # TODO may improve the format
- if evaluation:
- dump_data = dict(
- samples=[sample.to_dict() for dataset_name, info in data.items() for sample in info["samples"]]
- )
- else:
- dump_data = dict(
- samples=[sample.to_dict() for sample in data],
- )
-
- torch.save(dict(rollout_id=rollout_id, **dump_data), path)
-
def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]):
if self.custom_reward_post_process_func is not None:
return self.custom_reward_post_process_func(self.args, samples)
@@ -848,7 +734,7 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
if samples[0].rollout_routed_experts is not None:
routed_experts = [torch.as_tensor(sample.rollout_routed_experts) for sample in samples]
if getattr(self.args, "use_rollout_routing_replay", False):
- _validate_rollout_routed_experts_for_replay(routed_experts, self.args)
+ validate_rollout_routed_experts_for_replay(routed_experts, self.args)
train_data["rollout_routed_experts"] = routed_experts
if samples[0].train_metadata is not None:
@@ -927,7 +813,7 @@ def _split_train_data_by_dp(self, data):
rollout_data["global_batch_sizes"] = global_batch_sizes
rollout_data["num_microbatches"] = num_microbatches
rollout_data["micro_batch_indices"] = micro_batch_indices[r]
- _tensorize_rollout_data_for_training(rollout_data)
+ tensorize_rollout_data_for_training(rollout_data)
transport = getattr(self.args, "rollout_data_transport", "object-store")
if transport == "nixl":
rollout_data_refs.append(Box(ray.put(rollout_data, _tensor_transport="nixl")))
@@ -938,38 +824,6 @@ def _split_train_data_by_dp(self, data):
return rollout_data_refs
-def _validate_rollout_id_annotated(node, depth=0):
- """Walk the rollout function's nested output and validate ``rollout_id`` only
- when a compact / subagent pattern is detected.
-
- "Compact" = the rollout function wraps multiple training samples from one
- rollout execution into a ``list[Sample]``. In slime's convention the
- default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout)
- so its leaf ``list[Sample]`` lands at depth 1 and we skip validation,
- preserving backward compatibility. A compact rollout adds a third level:
- ``list[list[list[Sample]]]`` (prompt × rollout × samples-from-one-rollout),
- so the leaf ``list[Sample]`` lands at depth ≥ 2. At that point we require
- every sibling to carry a non-None ``rollout_id`` and to share the same
- value, so the loss reducer counts the rollout once instead of N times.
- """
- if isinstance(node, Sample):
- return
- assert isinstance(node, list), f"unexpected rollout output node type: {type(node).__name__}"
- if node and isinstance(node[0], Sample):
- if depth >= 2 and len(node) > 1:
- rids = [s.rollout_id for s in node]
- missing = [i for i, r in enumerate(rids) if r is None]
- assert not missing, (
- f"Compact rollout returned {len(node)} samples but rollout_id is unset on "
- f"positions {missing}. Set Sample.rollout_id on every sibling so the loss "
- "reducer can aggregate them as one rollout instead of N."
- )
- assert len(set(rids)) == 1, f"Sibling samples from one compact rollout must share rollout_id; got {rids}."
- return
- for item in node:
- _validate_rollout_id_annotated(item, depth + 1)
-
-
def _allocate_rollout_engine_addr_and_ports_normal(
*,
args,
@@ -1096,7 +950,8 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool
router_args.disable_circuit_breaker = True
# We will not use the health check from router.
- router_args.disable_health_check = True
+ if hasattr(router_args, "disable_health_check"):
+ router_args.disable_health_check = True
logger.info(f"Launch router with args: {router_args}")
@@ -1296,233 +1151,3 @@ def _resolve_sglang_config(args) -> SglangConfig:
)
]
)
-
-
-def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None):
- if args.custom_eval_rollout_log_function_path is not None:
- custom_log_func = load_function(args.custom_eval_rollout_log_function_path)
- if custom_log_func(rollout_id, args, data, extra_metrics):
- return
-
- log_dict = extra_metrics or {}
- for key in data.keys():
- rewards = data[key]["rewards"]
- log_dict[f"eval/{key}"] = sum(rewards) / len(rewards)
- if (samples := data[key].get("samples")) is not None:
- log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), f"eval/{key}/")
- if "truncated" in data[key]:
- truncated = data[key]["truncated"]
- log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated)
- if args.log_passrate:
- log_dict |= dict_add_prefix(
- compute_pass_rate(
- flat_rewards=rewards,
- group_size=args.n_samples_per_eval_prompt,
- ),
- f"eval/{key}-",
- )
-
- logger.info(f"eval {rollout_id}: {log_dict}")
-
- step = compute_rollout_step(args, rollout_id)
- log_dict["eval/step"] = step
- logging_utils.log(args, log_dict, step_key="eval/step")
-
- return log_dict
-
-
-def _log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time):
- if args.custom_rollout_log_function_path is not None:
- custom_log_func = load_function(args.custom_rollout_log_function_path)
- if custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time):
- return
-
- if args.load_debug_rollout_data:
- return
-
- log_dict = {**(rollout_extra_metrics or {})}
- log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), "rollout/")
- log_dict |= dict_add_prefix(compute_perf_metrics_from_samples(args, samples, rollout_time), "perf/")
- logger.info(f"perf {rollout_id}: {log_dict}")
- step = compute_rollout_step(args, rollout_id)
- log_dict["rollout/step"] = step
- logging_utils.log(args, log_dict, step_key="rollout/step")
-
-
-def compute_metrics_from_samples(args, samples):
- response_lengths = [sample.effective_response_length for sample in samples]
-
- log_dict = {}
- log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/")
- log_dict |= _compute_zero_std_metrics(args, samples)
- log_dict |= _compute_spec_metrics(args, samples)
- log_dict |= _compute_prefix_cache_metrics(args, samples)
- log_dict |= _compute_reward_cat_metrics(args, samples)
- log_dict |= _compute_top_p_kept_vocab_metrics(args, samples)
- log_dict["repetition_frac"] = np.mean([int(has_repetition(s.response)) for s in samples]).item()
- log_dict["truncated_ratio"] = np.mean([int(s.status == Sample.Status.TRUNCATED) for s in samples]).item()
- return log_dict
-
-
-def compute_perf_metrics_from_samples(args, samples, rollout_time):
- non_generation_time = [sample.non_generation_time for sample in samples]
-
- log_dict = {}
- log_dict["rollout_time"] = rollout_time
- if max(non_generation_time) > 0:
- log_dict |= dict_add_prefix(compute_statistics(non_generation_time), "non_generation_time/")
-
- def token_perf(response_lengths, non_generation_time, key=""):
- max_response_length = max(response_lengths)
- if args.rollout_num_gpus:
- log_dict[f"{key}tokens_per_gpu_per_sec"] = sum(response_lengths) / rollout_time / args.rollout_num_gpus
- log_dict[f"longest_{key}sample_tokens_per_sec"] = max_response_length / rollout_time
-
- if max(non_generation_time) == 0:
- return
-
- non_generation_time = [
- t for t, length in zip(non_generation_time, response_lengths, strict=True) if length == max_response_length
- ]
- mean_non_generation_time = sum(non_generation_time) / len(non_generation_time)
-
- log_dict[f"longest_{key}sample_non_generation_time"] = mean_non_generation_time
- log_dict[f"longest_{key}sample_tokens_per_sec_without_non_generation"] = max_response_length / (
- rollout_time - mean_non_generation_time
- )
-
- token_perf([sample.response_length for sample in samples], non_generation_time, key="")
- token_perf([sample.effective_response_length for sample in samples], non_generation_time, key="effective_")
- log_dict |= _compute_sglang_request_perf_metrics(samples)
-
- return log_dict
-
-
-def _compute_sglang_request_perf_metrics(all_samples: list[Sample]):
- attrs_by_request = list(_iter_sglang_generate_attrs(all_samples))
- if not attrs_by_request:
- return {}
-
- values_by_metric: dict[str, list[float]] = {}
- profiled_request_count = 0
-
- def add_value(metric_key: str, source_key: str, attrs: dict) -> bool:
- value = attrs.get(source_key)
- if not isinstance(value, (int, float)) or isinstance(value, bool) or not np.isfinite(value):
- return False
- values_by_metric.setdefault(metric_key, []).append(float(value))
- return True
-
- for attrs in attrs_by_request:
- request_has_perf = False
-
- for metric_key, source_key in _SGLANG_REQUEST_PERF_FIELDS:
- request_has_perf |= add_value(metric_key, source_key, attrs)
-
- for metric_key, source_key in _SGLANG_PREFILL_PERF_FIELDS:
- request_has_perf |= add_value(metric_key, source_key, attrs)
-
- for metric_key, source_key in _SGLANG_DECODE_PERF_FIELDS:
- request_has_perf |= add_value(metric_key, source_key, attrs)
-
- if request_has_perf:
- profiled_request_count += 1
-
- metrics: dict[str, float] = {}
- for key, values in values_by_metric.items():
- if not values:
- continue
- metrics |= dict_add_prefix(compute_statistics(values), f"{key}/")
-
- return metrics
-
-
-def _iter_sglang_generate_attrs(all_samples: list[Sample]):
- for sample in all_samples:
- trace = getattr(sample, "trace", None)
- if not isinstance(trace, dict):
- continue
- for event in trace.get("events") or []:
- if event.get("type") != "span_end" or event.get("name") != "sglang_generate":
- continue
- attrs = event.get("attrs")
- if isinstance(attrs, dict):
- yield attrs
-
-
-def _compute_zero_std_metrics(args, all_samples: list[Sample]):
- # only compute in GRPO-like algorithms where one prompt has multiple responses
- if args.advantage_estimator == "ppo":
- return {}
-
- def _is_zero_std(samples: list[Sample]):
- rewards = [sample.get_reward_value(args) for sample in samples]
- return len(rewards) == 0 or all(rewards[0] == r for r in rewards)
-
- all_sample_groups = group_by(all_samples, lambda s: s.group_index)
- interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)]
-
- interesting_rewards = [str(round(g[0].get_reward_value(args), 1)) for g in interesting_sample_groups]
-
- return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()}
-
-
-def _compute_top_p_kept_vocab_metrics(args, all_samples: list[Sample]):
- total_kept = 0
- total_tokens = 0
- for sample in all_samples:
- offsets = sample.rollout_top_p_token_offsets
- if offsets is None or sample.response_length == 0:
- continue
- offsets = torch.as_tensor(offsets, dtype=torch.int64)
- if offsets.numel() == 0:
- continue
- assert (
- offsets.numel() == sample.response_length + 1
- ), f"top-p token offsets length {offsets.numel()} != response length + 1 {sample.response_length + 1}"
- if sample.remove_sample:
- continue
- if sample.loss_mask is None:
- total_kept += int(offsets[-1] - offsets[0])
- total_tokens += sample.response_length
- continue
- loss_mask = torch.as_tensor(sample.loss_mask, dtype=torch.bool, device=offsets.device)
- assert (
- loss_mask.numel() == sample.response_length
- ), f"loss mask length {loss_mask.numel()} != response length {sample.response_length}"
- total_kept += int(torch.diff(offsets)[loss_mask].sum())
- total_tokens += int(loss_mask.sum())
- if total_tokens == 0:
- return {}
- return {"top_p_kept_vocab_per_token": total_kept / total_tokens}
-
-
-def _compute_spec_metrics(args, all_samples: list[Sample]):
- if getattr(args, "sglang_speculative_algorithm", None) is None:
- return {}
- num_samples = len(all_samples)
- metrics = {}
- metrics["spec_accept_rate"] = sum(sample.spec_info.spec_accept_rate for sample in all_samples) / num_samples
- metrics["spec_accept_length"] = sum(sample.spec_info.spec_accept_length for sample in all_samples) / num_samples
- return metrics
-
-
-def _compute_prefix_cache_metrics(args, all_samples: list[Sample]):
- num_samples = len(all_samples)
- metrics = {}
- total_cached_tokens = sum(sample.prefix_cache_info.cached_tokens for sample in all_samples)
- total_prompt_tokens = sum(sample.prefix_cache_info.total_prompt_tokens for sample in all_samples)
-
- metrics["prefix_cache_hit_rate"] = total_cached_tokens / total_prompt_tokens if total_prompt_tokens > 0 else 0.0
- metrics["avg_cached_tokens_per_sample"] = total_cached_tokens / num_samples
- return metrics
-
-
-def _compute_reward_cat_metrics(args, all_samples: list[Sample]):
- reward_cat_key = args.log_reward_category
- if reward_cat_key is None:
- return {}
-
- samples_of_reward_cat = group_by(all_samples, lambda s: s.reward[reward_cat_key])
-
- return {f"error_cat/{reward_cat}": len(s) / len(all_samples) for reward_cat, s in samples_of_reward_cat.items()}
diff --git a/slime/ray/train_actor.py b/slime/ray/train_actor.py
index a8ba6ddc64..42fb0ee0e3 100644
--- a/slime/ray/train_actor.py
+++ b/slime/ray/train_actor.py
@@ -9,20 +9,17 @@
import torch.distributed as dist
import slime.utils.eval_config
+from slime.observability.logging_utils import configure_logger
from slime.ray.ray_actor import RayActor
+from slime.utils import accelerator
from slime.utils.distributed_utils import init_gloo_group
-from slime.utils.logging_utils import configure_logger
from slime.utils.memory_utils import clear_memory, print_memory
logger = logging.getLogger(__name__)
def get_local_gpu_id():
- cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None)
- if cvd is None:
- return ray.get_gpu_ids()[0]
- else:
- return cvd.split(",").index(str(ray.get_gpu_ids()[0]))
+ return accelerator.resolve_visible_device_id(ray.get_gpu_ids()[0])
class TrainRayActor(RayActor):
@@ -56,9 +53,14 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False):
torch.serialization.add_safe_globals([slime.utils.eval_config.EvalDatasetConfig])
local_rank = int(os.environ.get("LOCAL_RANK", 0))
- torch.cuda.set_device(f"cuda:{local_rank}")
+ accelerator.set_device(local_rank)
+ if accelerator.set_allocator_expandable_segments():
+ logger.info(
+ f"[Rank {self._rank}] Enabled {accelerator.device_type().upper()} memory allocator "
+ "expandable_segments for train actor"
+ )
- backend = args.distributed_backend
+ backend = accelerator.process_group_backend(args.distributed_backend)
dist.init_process_group(
backend=backend,
diff --git a/slime/ray/utils.py b/slime/ray/utils.py
index decdbdec2e..2575a96f04 100644
--- a/slime/ray/utils.py
+++ b/slime/ray/utils.py
@@ -2,8 +2,9 @@
import os
import ray
-import torch
+
from slime.ray.ray_actor import RayActor
+from slime.utils import accelerator
# Refer to
# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/nvidia_gpu.py#L95-L96
@@ -15,6 +16,7 @@
# https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/intel_gpu.py#L97-L98
NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [
"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
+ "RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES",
"RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES",
"RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES",
"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES",
@@ -47,9 +49,9 @@ def ray_noset_visible_devices(env_vars=os.environ):
def get_physical_gpu_id():
- device = torch.cuda.current_device()
- props = torch.cuda.get_device_properties(device)
- return str(props.uuid)
+ device = accelerator.current_device()
+ props = accelerator.get_device_properties(device)
+ return str(getattr(props, "uuid", device))
@ray.remote
diff --git a/slime/rollout/_fanout_test_helpers.py b/slime/rollout/_fanout_test_helpers.py
index e065dd13dd..d989972c2d 100644
--- a/slime/rollout/_fanout_test_helpers.py
+++ b/slime/rollout/_fanout_test_helpers.py
@@ -12,7 +12,7 @@
- ``compact_generate``: fans one input sample out to N siblings
sharing the same ``rollout_id``. That's the contract the rest of the
framework (per-rollout step splitter, per-rollout-mean reducer,
- ``_validate_rollout_id_annotated`` validator) is built around.
+ ``validate_rollout_id_annotated`` validator) is built around.
- ``grpo_normalize_by_group_index``: replaces the default
``_post_process_rewards`` reshape-by-shape logic with a proper
diff --git a/slime/rollout/on_policy_distillation.py b/slime/rollout/on_policy_distillation.py
index 9190974345..eb52a0821d 100644
--- a/slime/rollout/on_policy_distillation.py
+++ b/slime/rollout/on_policy_distillation.py
@@ -10,7 +10,7 @@ async def reward_func(args, sample, **kwargs):
# "text": sample.prompt + sample.response,
"input_ids": sample.tokens,
"sampling_params": {
- "temperature": 0,
+ "temperature": args.rollout_temperature,
"max_new_tokens": 0,
"skip_special_tokens": False,
},
diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py
index 7615f20e7a..2cbcb89662 100644
--- a/slime/rollout/sglang_rollout.py
+++ b/slime/rollout/sglang_rollout.py
@@ -13,6 +13,7 @@
from tqdm import tqdm
from slime.backends.sglang_utils.server_control import abort_servers_until_idle
+from slime.observability.trace_utils import build_sglang_meta_trace_attrs, trace_function, trace_span
from slime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput
from slime.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter, should_drop_dynamic_filter_output
from slime.rollout.sample_hooks import apply_rollout_sample_hooks
@@ -27,7 +28,6 @@
load_processor,
load_tokenizer,
)
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_function, trace_span
from slime.utils.types import Sample
from .rm_hub import async_rm, batched_async_rm
diff --git a/slime/rollout/sglang_streaming_rollout.py b/slime/rollout/sglang_streaming_rollout.py
index 195da071bd..bcc6d176fb 100644
--- a/slime/rollout/sglang_streaming_rollout.py
+++ b/slime/rollout/sglang_streaming_rollout.py
@@ -29,10 +29,10 @@
from argparse import Namespace
from typing import Any
+from slime.observability.trace_utils import build_sglang_meta_trace_attrs, trace_span
from slime.rollout.sglang_rollout import GenerateState, _prepare_prompt_ids
from slime.utils import http_utils
from slime.utils.processing_utils import encode_image_for_rollout_engine
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span
from slime.utils.types import Sample
__all__ = ["generate_streaming"]
diff --git a/slime/utils/accelerator/__init__.py b/slime/utils/accelerator/__init__.py
new file mode 100644
index 0000000000..4ece4d01c0
--- /dev/null
+++ b/slime/utils/accelerator/__init__.py
@@ -0,0 +1,394 @@
+"""Runtime-selectable, backend-neutral accelerator API for Slime.
+
+The module-level functions are compatibility shims for the historical
+``slime.utils.accelerator`` API. New code can use :func:`get_accelerator`
+when it needs capability inspection or dependency injection.
+"""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+import sys
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from .base import Accelerator
+from .cuda import CUDAAccelerator
+from .musa import MUSAAccelerator
+from .musa import is_musa_available as _is_musa_available
+
+logger = logging.getLogger(__name__)
+
+_MUSA_PATCH_IMPORTED = False
+_MUSA_BOOTSTRAP_CHECKED = False
+_ACCELERATOR: Accelerator | None = None
+_SELECTION_LOCK = threading.RLock()
+
+
+@dataclass(frozen=True)
+class _BackendRegistration:
+ factory: Callable[[], Accelerator]
+ is_available: Callable[[], bool]
+ priority: int
+ communication_backends: tuple[str, ...]
+
+
+_REGISTRY: dict[str, _BackendRegistration] = {}
+
+
+def register_accelerator(
+ name: str,
+ factory: Callable[[], Accelerator],
+ is_available: Callable[[], bool] | None = None,
+ priority: int = 0,
+ communication_backends: tuple[str, ...] = (),
+) -> None:
+ """Register a lazily constructed backend without changing Slime core."""
+ normalized = name.strip().lower()
+ if not normalized or normalized in {"auto", "none"}:
+ raise ValueError("Accelerator name must be a non-empty backend name")
+ with _SELECTION_LOCK:
+ _REGISTRY[normalized] = _BackendRegistration(
+ factory=factory,
+ is_available=is_available or (lambda: factory().is_available()),
+ priority=priority,
+ communication_backends=tuple(name.lower() for name in communication_backends),
+ )
+
+
+def _append_musa_patch_path() -> None:
+ patch_path = os.environ.get("MUSA_PATCH_PATH")
+ if patch_path and patch_path not in sys.path:
+ sys.path.append(patch_path)
+
+
+def _import_musa_patch() -> bool:
+ _append_musa_patch_path()
+ try:
+ importlib.import_module("musa_patch")
+ except ModuleNotFoundError as exc:
+ if exc.name == "musa_patch":
+ return False
+ raise RuntimeError(f"musa_patch failed because dependency {exc.name!r} is missing") from exc
+ except Exception as exc:
+ raise RuntimeError(f"musa_patch initialization failed: {exc}") from exc
+ return True
+
+
+def is_musa_available() -> bool:
+ return _is_musa_available()
+
+
+def is_musa_environment() -> bool:
+ return (
+ is_musa_available()
+ or os.environ.get("SLIME_ACCELERATOR", "").lower() == "musa"
+ or "MUSA_VISIBLE_DEVICES" in os.environ
+ or bool(os.environ.get("MUSA_PATCH_PATH"))
+ )
+
+
+def _try_import_musa_patch() -> bool:
+ global _MUSA_PATCH_IMPORTED
+ if _MUSA_PATCH_IMPORTED:
+ return True
+ if not is_musa_environment():
+ return False
+ _MUSA_PATCH_IMPORTED = _import_musa_patch()
+ if not _MUSA_PATCH_IMPORTED and is_musa_environment():
+ logger.warning("musa_patch is not importable; continuing without it")
+ return _MUSA_PATCH_IMPORTED
+
+
+def _musa_requested() -> bool:
+ configured = os.environ.get("SLIME_ACCELERATOR", "").lower()
+ if configured and configured != "auto":
+ return configured == "musa"
+ return "MUSA_VISIBLE_DEVICES" in os.environ or bool(os.environ.get("MUSA_PATCH_PATH"))
+
+
+def _bootstrap_musa_patch_if_needed() -> bool:
+ """Bootstrap the patch for an already chosen MUSA backend at most once."""
+ global _MUSA_BOOTSTRAP_CHECKED
+ if _MUSA_BOOTSTRAP_CHECKED:
+ return _MUSA_PATCH_IMPORTED
+ _MUSA_BOOTSTRAP_CHECKED = True
+ return _try_import_musa_patch()
+
+
+def _cuda_available() -> bool:
+ try:
+ return bool(torch.cuda.is_available() and torch.cuda.device_count() > 0)
+ except (ImportError, RuntimeError):
+ return False
+
+
+def _register_builtin_backends() -> None:
+ if "cuda" not in _REGISTRY:
+ register_accelerator("cuda", CUDAAccelerator, _cuda_available, priority=100, communication_backends=("nccl",))
+ if "musa" not in _REGISTRY:
+ register_accelerator(
+ "musa", MUSAAccelerator, is_musa_available, priority=200, communication_backends=("mccl",)
+ )
+
+
+def _requested_name() -> str | None:
+ value = os.environ.get("SLIME_ACCELERATOR")
+ if value and value.lower() != "auto":
+ return value.strip().lower()
+ if _musa_requested():
+ return "musa"
+ return None
+
+
+def _make_selected(name: str, explicit: bool) -> Accelerator:
+ _register_builtin_backends()
+ entry = _REGISTRY.get(name)
+ if entry is None:
+ available = ", ".join(sorted(_REGISTRY))
+ raise ValueError(f"Unknown accelerator {name!r}; registered backends: {available}")
+ if name == "musa":
+ # musa_patch may expose torch.musa, so bootstrap after MUSA has been
+ # chosen but before validating and constructing its backend.
+ _bootstrap_musa_patch_if_needed()
+ if explicit and not entry.is_available():
+ if name == "musa":
+ detail = (
+ "torch.musa is unavailable; install a MUSA-enabled PyTorch runtime and set MUSA_PATCH_PATH if required"
+ )
+ elif name == "cuda":
+ detail = "torch.cuda.is_available() is false or no CUDA device is visible"
+ else:
+ detail = "the backend availability check returned false"
+ raise RuntimeError(f"Requested accelerator {name!r} is unavailable: {detail}")
+ backend = entry.factory()
+ if not backend.is_available():
+ raise RuntimeError(f"Accelerator backend {name!r} was selected but is unavailable at runtime")
+ return backend
+
+
+def get_accelerator() -> Accelerator:
+ global _ACCELERATOR
+ if _ACCELERATOR is not None:
+ return _ACCELERATOR
+ with _SELECTION_LOCK:
+ if _ACCELERATOR is not None:
+ return _ACCELERATOR
+ _register_builtin_backends()
+ requested = _requested_name()
+ if requested is not None:
+ _ACCELERATOR = _make_selected(requested, explicit=True)
+ logger.info("Selected accelerator %s (explicit)", _ACCELERATOR.name)
+ return _ACCELERATOR
+
+ # Highest priority wins; names break priority ties deterministically.
+ candidates = sorted(_REGISTRY.items(), key=lambda item: (-item[1].priority, item[0]))
+ for name, registration in candidates:
+ if registration.is_available():
+ _ACCELERATOR = _make_selected(name, explicit=False)
+ logger.info("Selected accelerator %s (auto)", _ACCELERATOR.name)
+ return _ACCELERATOR
+ registered = ", ".join(sorted(_REGISTRY))
+ raise RuntimeError(
+ "No usable accelerator was detected. "
+ f"Registered backends: {registered}. "
+ "Set SLIME_ACCELERATOR explicitly or install a supported accelerator runtime."
+ )
+
+
+def initialize_accelerator() -> Accelerator | None:
+ """Finalize runtime selection when a backend is requested or available.
+
+ Explicit requests retain ``get_accelerator``'s fail-fast behavior. An
+ environment without accelerator hardware remains importable for CPU-only
+ tooling and documentation.
+ """
+ if _ACCELERATOR is not None:
+ return _ACCELERATOR
+ with _SELECTION_LOCK:
+ _register_builtin_backends()
+ if _requested_name() is not None or any(entry.is_available() for entry in _REGISTRY.values()):
+ return get_accelerator()
+ return None
+
+
+def set_accelerator(accelerator: Accelerator) -> None:
+ global _ACCELERATOR
+ if not isinstance(accelerator, Accelerator):
+ raise TypeError(f"Expected Accelerator, got {type(accelerator).__name__}")
+ if not accelerator.is_available():
+ raise RuntimeError(f"Cannot install unavailable accelerator backend {accelerator.name!r}")
+ with _SELECTION_LOCK:
+ _ACCELERATOR = accelerator
+
+
+def reset_accelerator() -> None:
+ """Reset the singleton; intended for tests and process initialization."""
+ global _ACCELERATOR
+ with _SELECTION_LOCK:
+ _ACCELERATOR = None
+
+
+def _backend() -> Accelerator:
+ return get_accelerator()
+
+
+def device_type() -> str:
+ return _backend().device_type
+
+
+def accelerator_module() -> Any:
+ return _backend().accelerator_module()
+
+
+def device(index: int | str | torch.device | None = None) -> torch.device:
+ return _backend().device(index)
+
+
+def device_name(index: int | str | torch.device | None = None) -> str:
+ return _backend().device_name(index)
+
+
+def set_device(index: int | str | torch.device) -> None:
+ return _backend().set_device(index)
+
+
+def current_device() -> int | str:
+ return _backend().current_device()
+
+
+def device_count() -> int:
+ return _backend().device_count()
+
+
+def synchronize(device_arg: int | str | torch.device | None = None) -> None:
+ return _backend().synchronize(device_arg)
+
+
+def current_stream(device_arg: int | str | torch.device | None = None) -> Any:
+ return _backend().current_stream(device_arg)
+
+
+def default_stream(device_arg: int | str | torch.device | None = None) -> Any:
+ return _backend().default_stream(device_arg)
+
+
+def stream(stream_arg: Any):
+ return _backend().stream(stream_arg)
+
+
+def new_stream(*args, **kwargs) -> Any:
+ stream_type = _backend().Stream
+ if stream_type is None:
+ raise NotImplementedError(f"Accelerator {_backend().name!r} does not support streams")
+ return stream_type(*args, **kwargs)
+
+
+def new_event(*args, **kwargs) -> Any:
+ event_type = _backend().Event
+ if event_type is None:
+ raise NotImplementedError(f"Accelerator {_backend().name!r} does not support events")
+ return event_type(*args, **kwargs)
+
+
+def empty_cache() -> None:
+ return _backend().empty_cache()
+
+
+def ipc_collect() -> None:
+ return _backend().ipc_collect()
+
+
+def set_allocator_expandable_segments() -> bool:
+ return _backend().set_allocator_expandable_segments()
+
+
+def mem_get_info(device_arg: int | str | torch.device | None = None) -> tuple[int, int]:
+ return _backend().mem_get_info(device_arg)
+
+
+def memory_allocated(device_arg: int | str | torch.device | None = None) -> int:
+ return _backend().memory_allocated(device_arg)
+
+
+def memory_reserved(device_arg: int | str | torch.device | None = None) -> int:
+ return _backend().memory_reserved(device_arg)
+
+
+def get_device_properties(device_arg: int | str | torch.device | None = None) -> Any:
+ return _backend().get_device_properties(device_arg)
+
+
+def memory_module() -> Any:
+ return _backend().memory_module()
+
+
+def attach_oom_observer(callback) -> bool:
+ return _backend().attach_oom_observer(callback)
+
+
+def supports(capability: str) -> bool:
+ return _backend().supports(capability)
+
+
+def autocast(*args, **kwargs):
+ return _backend().autocast(*args, **kwargs)
+
+
+def manual_seed(seed: int) -> None:
+ return _backend().manual_seed(seed)
+
+
+def manual_seed_all(seed: int) -> None:
+ return _backend().manual_seed_all(seed)
+
+
+def get_rng_state(device_arg: int | str | torch.device | None = None) -> torch.Tensor:
+ return _backend().get_rng_state(device_arg)
+
+
+def set_rng_state(state: torch.Tensor, device_arg: int | str | torch.device | None = None) -> None:
+ return _backend().set_rng_state(state, device_arg)
+
+
+def initial_seed() -> int:
+ return _backend().initial_seed()
+
+
+def distributed_device_id(index: int | str | torch.device | None = None) -> torch.device | None:
+ return _backend().distributed_device_id(index)
+
+
+def post_import_torch() -> None:
+ return _backend().post_import_torch()
+
+
+def is_accelerator_backend(backend: str) -> bool:
+ """Return whether a distributed backend belongs to a registered device accelerator."""
+ _register_builtin_backends()
+ normalized = backend.lower()
+ return any(
+ name in normalized for registration in _REGISTRY.values() for name in registration.communication_backends
+ )
+
+
+def process_group_backend(default: str = "nccl") -> str:
+ return _backend().communication_backend(default)
+
+
+def weight_update_backend(default: str = "nccl") -> str:
+ return _backend().weight_update_backend(default)
+
+
+def visible_devices_env_key() -> str:
+ return _backend().visible_devices_env
+
+
+def resolve_visible_device_id(physical_device_id: int | float | str) -> int:
+ return _backend().resolve_visible_device_id(physical_device_id)
diff --git a/slime/utils/accelerator/base.py b/slime/utils/accelerator/base.py
new file mode 100644
index 0000000000..9d760c619e
--- /dev/null
+++ b/slime/utils/accelerator/base.py
@@ -0,0 +1,179 @@
+"""Small, backend-neutral accelerator contract used by Slime.
+
+The contract intentionally contains only operations that Slime uses in its
+runtime. Vendor modules are supplied by concrete implementations and are
+never imported by this module.
+"""
+
+from __future__ import annotations
+
+import abc
+import os
+from typing import Any
+
+import torch
+
+
+class Accelerator(abc.ABC):
+ """Common device/runtime surface exposed to Slime code."""
+
+ name: str
+ device_type: str
+ communication_backend_name: str
+
+ @abc.abstractmethod
+ def is_available(self) -> bool:
+ """Return whether this backend can actually execute on this host."""
+
+ @abc.abstractmethod
+ def device(self, index: int | str | torch.device | None = None) -> torch.device:
+ """Return a :class:`torch.device` for a local device index."""
+
+ @abc.abstractmethod
+ def device_name(self, index: int | str | torch.device | None = None) -> str:
+ """Return the canonical device string used by PyTorch APIs."""
+
+ @abc.abstractmethod
+ def set_device(self, index: int | str | torch.device) -> None:
+ """Select the current local device."""
+
+ @abc.abstractmethod
+ def current_device(self) -> int | str:
+ """Return the current local device index, or ``cpu`` for CPU."""
+
+ @abc.abstractmethod
+ def device_count(self) -> int:
+ """Return the number of visible devices."""
+
+ @abc.abstractmethod
+ def synchronize(self, device: int | str | torch.device | None = None) -> None:
+ """Synchronize work on one device or the current device."""
+
+ @abc.abstractmethod
+ def current_stream(self, device: int | str | torch.device | None = None) -> Any:
+ """Return the current stream, or ``None`` when streams are unsupported."""
+
+ def default_stream(self, device: int | str | torch.device | None = None) -> Any:
+ """Return the default stream, or ``None`` when streams are unsupported."""
+ return None
+
+ def stream(self, stream: Any):
+ """Return a context manager for a stream when the backend supports it."""
+ raise NotImplementedError(f"Accelerator {self.name!r} does not support streams")
+
+ @property
+ def Stream(self) -> Any:
+ return None
+
+ @property
+ def Event(self) -> Any:
+ return None
+
+ @abc.abstractmethod
+ def empty_cache(self) -> None:
+ """Release allocator-held, currently unused memory."""
+
+ def ipc_collect(self) -> None:
+ """Collect inter-process allocator state when supported."""
+ return None
+
+ def set_allocator_expandable_segments(self) -> bool:
+ """Configure expandable allocator segments when supported."""
+ return False
+
+ @abc.abstractmethod
+ def mem_get_info(self, device: int | str | torch.device | None = None) -> tuple[int, int]:
+ """Return ``(free_bytes, total_bytes)`` for the selected device."""
+
+ @abc.abstractmethod
+ def memory_allocated(self, device: int | str | torch.device | None = None) -> int:
+ """Return currently allocated device memory in bytes."""
+
+ @abc.abstractmethod
+ def memory_reserved(self, device: int | str | torch.device | None = None) -> int:
+ """Return allocator-reserved device memory in bytes."""
+
+ def get_device_properties(self, device: int | str | torch.device | None = None) -> Any:
+ return None
+
+ def memory_module(self) -> Any:
+ """Return the backend memory namespace, if it exposes one."""
+ return None
+
+ def attach_oom_observer(self, callback) -> bool:
+ """Attach an OOM callback; return ``False`` when unsupported."""
+ return False
+
+ def supports(self, capability: str) -> bool:
+ """Return whether a named optional capability is implemented."""
+ return False
+
+ def autocast(self, *args, **kwargs):
+ """Return an autocast context for this backend."""
+ raise NotImplementedError(f"Accelerator {self.name!r} does not support autocast")
+
+ def manual_seed(self, seed: int) -> None:
+ """Seed the current device generator when supported."""
+ return None
+
+ def manual_seed_all(self, seed: int) -> None:
+ """Seed all device generators when supported."""
+ return None
+
+ def get_rng_state(self, device: int | str | torch.device | None = None) -> torch.Tensor:
+ """Return the current generator state."""
+ return torch.get_rng_state()
+
+ def set_rng_state(self, state: torch.Tensor, device: int | str | torch.device | None = None) -> None:
+ """Restore the current generator state."""
+ torch.set_rng_state(state)
+
+ def initial_seed(self) -> int:
+ return int(torch.initial_seed())
+
+ def distributed_device_id(self, index: int | str | torch.device | None = None) -> torch.device | None:
+ """Return the device id accepted by ``dist.init_process_group``."""
+ return self.device(index)
+
+ def post_import_torch(self) -> None:
+ """Apply an optional backend hook after third-party torch imports."""
+ return None
+
+ def communication_backend(self, default: str = "nccl") -> str:
+ """Map a logical default backend to this accelerator's transport."""
+ return self.communication_backend_name if default == "nccl" else default
+
+ def weight_update_backend(self, default: str = "nccl") -> str:
+ return self.communication_backend(default)
+
+ @property
+ def visible_devices_env(self) -> str:
+ return "CUDA_VISIBLE_DEVICES"
+
+ def resolve_visible_device_id(self, physical_device_id: int | float | str) -> int:
+ """Map a physical id to a local id under this backend's visibility env."""
+ raw_value = str(physical_device_id).strip()
+ visible = os.environ.get(self.visible_devices_env)
+ if not visible:
+ return int(float(raw_value))
+
+ ids = [item.strip() for item in visible.split(",") if item.strip()]
+ if raw_value in ids:
+ return ids.index(raw_value)
+
+ try:
+ value = int(float(raw_value))
+ except ValueError:
+ value = None
+ if value is not None and str(value) in ids:
+ return ids.index(str(value))
+ if value is not None and 0 <= value < len(ids):
+ return value
+ raise RuntimeError(
+ f"Device id {raw_value} is not valid under {self.visible_devices_env}={visible}. "
+ f"Expected one of {ids} (physical) or 0..{len(ids) - 1} (local)."
+ )
+
+ def accelerator_module(self) -> Any:
+ """Return the torch backend namespace, or ``None`` for CPU."""
+ return None
diff --git a/slime/utils/accelerator/cuda.py b/slime/utils/accelerator/cuda.py
new file mode 100644
index 0000000000..c9fca2afe0
--- /dev/null
+++ b/slime/utils/accelerator/cuda.py
@@ -0,0 +1,40 @@
+"""CUDA/ROCm accelerator implementation."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from .torch_accelerator import TorchAccelerator
+
+
+class CUDAAccelerator(TorchAccelerator):
+ name = "cuda"
+ device_type = "cuda"
+ communication_backend_name = "nccl"
+
+ def _module(self) -> Any:
+ return torch.cuda
+
+ def is_available(self) -> bool:
+ return bool(torch.cuda.is_available() and torch.cuda.device_count() > 0)
+
+ def attach_oom_observer(self, callback) -> bool:
+ attach = getattr(torch._C, "_cuda_attach_out_of_memory_observer", None)
+ if attach is None:
+ return False
+ attach(callback)
+ return True
+
+ def supports(self, capability: str) -> bool:
+ if capability == "nvml_affinity":
+ return torch.version.hip is None
+ if capability == "bf16":
+ checker = getattr(torch.cuda, "is_bf16_supported", None)
+ return bool(checker and checker())
+ if capability in {"cuda_int4_extension", "sglang_fp8_utils", "strict_fp32_logits", "triton_kernels"}:
+ return True
+ if capability == "requires_cpu_initialization":
+ return torch.version.hip is not None
+ return super().supports(capability)
diff --git a/slime/utils/accelerator/musa.py b/slime/utils/accelerator/musa.py
new file mode 100644
index 0000000000..62b398dfaa
--- /dev/null
+++ b/slime/utils/accelerator/musa.py
@@ -0,0 +1,88 @@
+"""MUSA accelerator implementation.
+
+Importing this module never imports ``torch_musa``. A MUSA runtime or the
+optional ``musa_patch`` bootstrap may attach ``torch.musa`` before selection.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+import torch
+
+from .torch_accelerator import TorchAccelerator
+
+
+def musa_module() -> Any:
+ return getattr(torch, "musa", None)
+
+
+def is_musa_available() -> bool:
+ module = musa_module()
+ checker = getattr(module, "is_available", None)
+ return bool(module is not None and checker is not None and checker())
+
+
+class MUSAAccelerator(TorchAccelerator):
+ name = "musa"
+ device_type = "musa"
+ communication_backend_name = "mccl"
+
+ @property
+ def visible_devices_env(self) -> str:
+ return "MUSA_VISIBLE_DEVICES"
+
+ def _module(self) -> Any:
+ module = musa_module()
+ if module is None:
+ raise RuntimeError("MUSA backend requires a runtime that exposes torch.musa")
+ return module
+
+ def is_available(self) -> bool:
+ return is_musa_available()
+
+ def weight_update_backend(self, default: str = "nccl") -> str:
+ return "cpu:gloo,musa:mccl" if default == "nccl" else default
+
+ def distributed_device_id(self, index: int | str | torch.device | None = None) -> None:
+ return None
+
+ def post_import_torch(self) -> None:
+ try:
+ module = importlib.import_module("musa_patch")
+ except ModuleNotFoundError as exc:
+ if exc.name == "musa_patch":
+ return
+ raise RuntimeError(f"musa_patch failed because dependency {exc.name!r} is missing") from exc
+ callback = getattr(module, "patch_after_import_torch", None)
+ if callback is not None:
+ callback()
+
+ def attach_oom_observer(self, callback) -> bool:
+ musa_c = getattr(self._module(), "_MUSAC", None)
+ attach = getattr(musa_c, "_musa_attach_out_of_memory_observer", None)
+ if attach is None:
+ return False
+ attach(callback)
+ return True
+
+ def autocast(self, *args, **kwargs):
+ amp = getattr(self._module(), "amp", None)
+ autocast = getattr(amp, "autocast", None)
+ if autocast is None:
+ raise NotImplementedError("MUSA runtime does not expose torch.musa.amp.autocast")
+ return autocast(*args, **kwargs)
+
+ def supports(self, capability: str) -> bool:
+ if capability in {"nvml_affinity", "sglang_fp8_utils", "strict_fp32_logits"}:
+ return False
+ if capability == "requires_cpu_initialization":
+ return True
+ if capability == "amp":
+ amp = getattr(self._module(), "amp", None)
+ return callable(getattr(amp, "autocast", None))
+ if capability == "bf16":
+ checker = getattr(self._module(), "is_bf16_supported", None)
+ return bool(checker and checker())
+ return super().supports(capability)
diff --git a/slime/utils/accelerator/torch_accelerator.py b/slime/utils/accelerator/torch_accelerator.py
new file mode 100644
index 0000000000..d961ebe3d8
--- /dev/null
+++ b/slime/utils/accelerator/torch_accelerator.py
@@ -0,0 +1,163 @@
+"""Shared adapter for PyTorch accelerator namespaces."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import torch
+
+from .base import Accelerator
+
+logger = logging.getLogger(__name__)
+
+
+class TorchAccelerator(Accelerator):
+ """Delegate common CUDA-like APIs to a vendor torch namespace."""
+
+ def _module(self) -> Any:
+ raise NotImplementedError
+
+ def accelerator_module(self) -> Any:
+ return self._module()
+
+ def is_available(self) -> bool:
+ module = self._module()
+ checker = getattr(module, "is_available", None)
+ return bool(module is not None and checker is not None and checker())
+
+ def device(self, index: int | str | torch.device | None = None) -> torch.device:
+ return torch.device(self.device_name(index))
+
+ def device_name(self, index: int | str | torch.device | None = None) -> str:
+ if isinstance(index, torch.device):
+ return str(index)
+ if isinstance(index, str):
+ return index if ":" in index else f"{self.device_type}:{index}"
+ if index is None:
+ index = self.current_device()
+ return f"{self.device_type}:{index}"
+
+ def set_device(self, index: int | str | torch.device) -> None:
+ self._module().set_device(index)
+
+ def current_device(self) -> int:
+ return int(self._module().current_device())
+
+ def device_count(self) -> int:
+ return int(self._module().device_count())
+
+ def synchronize(self, device: int | str | torch.device | None = None) -> None:
+ if device is None:
+ self._module().synchronize()
+ else:
+ self._module().synchronize(device)
+
+ def current_stream(self, device: int | str | torch.device | None = None) -> Any:
+ if device is None:
+ return self._module().current_stream()
+ return self._module().current_stream(device)
+
+ def default_stream(self, device: int | str | torch.device | None = None) -> Any:
+ default_stream = getattr(self._module(), "default_stream", None)
+ if default_stream is None:
+ raise NotImplementedError(f"Accelerator {self.name!r} does not expose a default stream")
+ if device is None:
+ return default_stream()
+ return default_stream(device)
+
+ def stream(self, stream: Any):
+ stream_context = getattr(self._module(), "stream", None)
+ if stream_context is None:
+ raise NotImplementedError(f"Accelerator {self.name!r} does not expose stream contexts")
+ return stream_context(stream)
+
+ @property
+ def Stream(self) -> Any:
+ return getattr(self._module(), "Stream", None)
+
+ @property
+ def Event(self) -> Any:
+ return getattr(self._module(), "Event", None)
+
+ def empty_cache(self) -> None:
+ self._module().empty_cache()
+
+ def ipc_collect(self) -> None:
+ collect = getattr(self._module(), "ipc_collect", None)
+ if collect is not None:
+ collect()
+
+ def set_allocator_expandable_segments(self) -> bool:
+ value = os.getenv("SLIME_ENABLE_EXPANDABLE_SEGMENTS", "0")
+ if value not in {"0", "1"}:
+ raise ValueError(f"SLIME_ENABLE_EXPANDABLE_SEGMENTS must be 0 or 1, got {value!r}")
+ if value == "0":
+ return False
+
+ memory = self.memory_module()
+ setter = getattr(memory, "_set_allocator_settings", None)
+ if setter is None:
+ logger.warning(
+ "%s memory allocator settings API is unavailable; skip expandable_segments:True",
+ self.name.upper(),
+ )
+ return False
+ setter("expandable_segments:True")
+ return True
+
+ def mem_get_info(self, device: int | str | torch.device | None = None) -> tuple[int, int]:
+ if device is None:
+ device = self.current_device()
+ free, total = self._module().mem_get_info(device)
+ return int(free), int(total)
+
+ def memory_allocated(self, device: int | str | torch.device | None = None) -> int:
+ return int(self._module().memory_allocated(device))
+
+ def memory_reserved(self, device: int | str | torch.device | None = None) -> int:
+ return int(self._module().memory_reserved(device))
+
+ def get_device_properties(self, device: int | str | torch.device | None = None) -> Any:
+ if device is None:
+ device = self.current_device()
+ return self._module().get_device_properties(device)
+
+ def memory_module(self) -> Any:
+ return getattr(self._module(), "memory", None)
+
+ def autocast(self, *args, **kwargs):
+ return torch.autocast(self.device_type, *args, **kwargs)
+
+ def manual_seed(self, seed: int) -> None:
+ self._module().manual_seed(seed)
+
+ def manual_seed_all(self, seed: int) -> None:
+ self._module().manual_seed_all(seed)
+
+ def get_rng_state(self, device: int | str | torch.device | None = None) -> torch.Tensor:
+ if device is None:
+ return self._module().get_rng_state()
+ return self._module().get_rng_state(device)
+
+ def set_rng_state(self, state: torch.Tensor, device: int | str | torch.device | None = None) -> None:
+ if device is None:
+ self._module().set_rng_state(state)
+ else:
+ self._module().set_rng_state(state, device)
+
+ def initial_seed(self) -> int:
+ return int(self._module().initial_seed())
+
+ def supports(self, capability: str) -> bool:
+ module = self._module()
+ if capability == "device_memory":
+ return all(hasattr(module, name) for name in ("empty_cache", "mem_get_info", "memory_allocated"))
+ if capability == "events":
+ return hasattr(module, "Event")
+ if capability == "rng":
+ return all(hasattr(module, name) for name in ("get_rng_state", "set_rng_state", "manual_seed"))
+ if capability == "streams":
+ return all(hasattr(module, name) for name in ("Stream", "current_stream", "stream"))
+ return capability in {"amp", "fp16"}
diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py
index 2ef7dad59b..33381cf0a3 100644
--- a/slime/utils/arguments.py
+++ b/slime/utils/arguments.py
@@ -10,8 +10,8 @@
from slime.backends.sglang_utils.arguments import sglang_parse_args
from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args
from slime.backends.sglang_utils.external import apply_external_engine_info_to_args
+from slime.observability.logging_utils import configure_logger
from slime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list
-from slime.utils.logging_utils import configure_logger
logger = logging.getLogger(__name__)
@@ -343,7 +343,7 @@ def add_rollout_arguments(parser):
"--rollout-temperature",
type=float,
default=1.0,
- help="the temperature for the inference engine during rollout.",
+ help="the temperature for the inference engine during rollout. Must be > 0.",
)
parser.add_argument(
"--rollout-top-p", type=float, default=1.0, help="the top-p for the inference engine during rollout."
@@ -1767,6 +1767,11 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
def slime_validate_args(args):
args.eval_datasets = _resolve_eval_datasets(args)
+ if args.rollout_temperature <= 0:
+ raise ValueError(
+ "--rollout-temperature must be > 0; temperature 0 is greedy decoding and is not a valid RL policy."
+ )
+
if args.kl_coef != 0 or args.use_kl_loss:
if not os.path.exists(args.ref_load):
raise FileNotFoundError(f"ref_load {args.ref_load} does not exist, please check the path.")
diff --git a/slime/utils/data.py b/slime/utils/data.py
index adac4b95a6..275cd3efa8 100644
--- a/slime/utils/data.py
+++ b/slime/utils/data.py
@@ -13,10 +13,9 @@
except ImportError:
pq = None
+from slime.observability.timer import Timer
from slime.utils.types import MultimodalTypes, Sample
-from .timer import Timer
-
__all__ = ["Dataset", "get_source"]
logger = logging.getLogger(__name__)
diff --git a/slime/utils/memory_utils.py b/slime/utils/memory_utils.py
index d4b2d89321..4bad36b51f 100644
--- a/slime/utils/memory_utils.py
+++ b/slime/utils/memory_utils.py
@@ -5,28 +5,30 @@
import torch
import torch.distributed as dist
+from slime.utils import accelerator
+
logger = logging.getLogger(__name__)
def clear_memory(clear_host_memory: bool = False):
- torch.cuda.synchronize()
+ accelerator.synchronize()
gc.collect()
- torch.cuda.empty_cache()
+ accelerator.empty_cache()
if clear_host_memory:
torch._C._host_emptyCache()
def available_memory():
- device = torch.cuda.current_device()
- free, total = torch.cuda.mem_get_info(device)
+ device = accelerator.current_device()
+ free, total = accelerator.mem_get_info(device)
vm = psutil.virtual_memory()
return {
"gpu": str(device),
"total_GB": _byte_to_gb(total),
"free_GB": _byte_to_gb(free),
"used_GB": _byte_to_gb(total - free),
- "allocated_GB": _byte_to_gb(torch.cuda.memory_allocated(device)),
- "reserved_GB": _byte_to_gb(torch.cuda.memory_reserved(device)),
+ "allocated_GB": _byte_to_gb(accelerator.memory_allocated(device)),
+ "reserved_GB": _byte_to_gb(accelerator.memory_reserved(device)),
"host_total_GB": _byte_to_gb(vm.total),
"host_available_GB": _byte_to_gb(vm.available),
"host_used_GB": _byte_to_gb(vm.used),
diff --git a/slime/utils/reloadable_process_group.py b/slime/utils/reloadable_process_group.py
index 3e79281d4d..bef0dc5e7b 100644
--- a/slime/utils/reloadable_process_group.py
+++ b/slime/utils/reloadable_process_group.py
@@ -9,6 +9,7 @@
import torch.distributed as dist
from torch.distributed.distributed_c10d import PrefixStore, _get_default_group, _get_default_store
+from slime.utils import accelerator
from slime.utils.distributed_utils import get_gloo_group, init_gloo_group, set_gloo_group
from slime.utils.memory_utils import available_memory, clear_memory, print_memory
@@ -26,11 +27,11 @@ class _DefaultProcessGroupState:
rank: int
world_size: int
generation: int = 0
- nccl_world_destroyed: bool = False
+ accelerator_world_destroyed: bool = False
def register_default_process_group(timeout: timedelta) -> None:
- """Register the NCCL WORLD group so it can be destroyed and rebuilt.
+ """Register the accelerator WORLD group so it can be destroyed and rebuilt.
Keeping a reference to the rendezvous store is intentional. It keeps the
rank-0 TCPStore alive after ``destroy_process_group()`` and lets every
@@ -58,8 +59,8 @@ def register_default_process_group(timeout: timedelta) -> None:
)
-def _uses_nccl(backend: str) -> bool:
- return "nccl" in backend.lower()
+def _uses_accelerator_backend(backend: str) -> bool:
+ return accelerator.is_accelerator_backend(backend)
def _new_default_process_group(state: _DefaultProcessGroupState, backend: str) -> None:
@@ -74,9 +75,9 @@ def _new_default_process_group(state: _DefaultProcessGroupState, backend: str) -
)
-def _destroy_default_nccl_process_group() -> None:
+def _destroy_default_accelerator_process_group() -> None:
state = default_process_group_states.get(os.getpid())
- if state is None or state.nccl_world_destroyed or not _uses_nccl(state.backend):
+ if state is None or state.accelerator_world_destroyed or not _uses_accelerator_backend(state.backend):
return
# Pure PP=4 exposed a teardown ordering deadlock here. Pipeline ranks own
@@ -96,7 +97,7 @@ def _destroy_default_nccl_process_group() -> None:
_new_default_process_group(state, backend="gloo")
set_gloo_group(_get_default_group())
- state.nccl_world_destroyed = True
+ state.accelerator_world_destroyed = True
logger.info(
"Destroyed default %s WORLD process group and initialized a temporary Gloo WORLD (generation %s)",
state.backend,
@@ -106,18 +107,18 @@ def _destroy_default_nccl_process_group() -> None:
def _reload_default_process_group() -> None:
state = default_process_group_states.get(os.getpid())
- if state is None or not state.nccl_world_destroyed:
+ if state is None or not state.accelerator_world_destroyed:
return
- # WORLD uses Gloo while the NCCL WORLD is destroyed, so this barrier does
- # not allocate CUDA or recreate an NCCL communicator before all ranks are ready.
+ # WORLD uses Gloo while the accelerator WORLD is destroyed, so this barrier
+ # does not recreate an accelerator communicator before all ranks are ready.
dist.barrier()
dist.destroy_process_group()
set_gloo_group(None)
_new_default_process_group(state, backend=state.backend)
init_gloo_group()
- state.nccl_world_destroyed = False
+ state.accelerator_world_destroyed = False
logger.info(
"Reloaded default WORLD process group with backend %s (generation %s)",
state.backend,
@@ -154,9 +155,17 @@ def monkey_patch_torch_dist():
dist.old_new_group = old_new_group
def new_group(*args, **kwargs):
- group = old_new_group(*args, **kwargs)
explicit_backend = args[2] if len(args) >= 3 else kwargs.get("backend")
backend = str(explicit_backend) if explicit_backend is not None else str(dist.get_backend())
+ normalized_backend = accelerator.process_group_backend(backend) if backend == "nccl" else backend
+ if normalized_backend != backend:
+ if len(args) >= 3:
+ args = (*args[:2], normalized_backend, *args[3:])
+ else:
+ kwargs = {**kwargs, "backend": normalized_backend}
+ backend = normalized_backend
+
+ group = old_new_group(*args, **kwargs)
# Before WORLD is registered, preserve the historical behavior of
# leaving CPU groups and singleton groups untouched. Afterwards every
@@ -457,16 +466,16 @@ def bound_device_id(self, dev):
def destroy_process_groups():
- """Destroy registered subgroups and replace NCCL WORLD with a temporary Gloo WORLD."""
+ """Destroy registered subgroups and replace accelerator WORLD with a temporary Gloo WORLD."""
state = default_process_group_states.get(os.getpid())
- if state is not None and not state.nccl_world_destroyed and _uses_nccl(state.backend):
- _destroy_default_nccl_process_group()
+ if state is not None and not state.accelerator_world_destroyed and _uses_accelerator_backend(state.backend):
+ _destroy_default_accelerator_process_group()
else:
ReloadableProcessGroup.destroy_process_groups()
def reload_process_groups():
- """Restore NCCL WORLD and recreate all registered subgroups."""
+ """Restore accelerator WORLD and recreate all registered subgroups."""
_reload_default_process_group()
ReloadableProcessGroup.reload_process_groups()
diff --git a/slime/utils/routing_replay.py b/slime/utils/routing_replay.py
index c26a9e192e..b9759ea72c 100644
--- a/slime/utils/routing_replay.py
+++ b/slime/utils/routing_replay.py
@@ -1,6 +1,9 @@
import os
import torch
+from slime.utils import accelerator
+
+
ROUTING_REPLAY = None
ORDERED_TOPK_CAPTURE_ROUTER = None
@@ -118,7 +121,7 @@ def pop_forward(self):
if hasattr(top_indices, "materialize_for_routing_replay"):
return top_indices.materialize_for_routing_replay("forward")
return top_indices.to(
- torch.cuda.current_device(),
+ accelerator.current_device(),
dtype=torch.int32,
non_blocking=top_indices.is_pinned(),
)
@@ -129,7 +132,7 @@ def pop_backward(self):
if hasattr(top_indices, "materialize_for_routing_replay"):
return top_indices.materialize_for_routing_replay("backward")
return top_indices.to(
- torch.cuda.current_device(),
+ accelerator.current_device(),
dtype=torch.int32,
non_blocking=top_indices.is_pinned(),
)
diff --git a/slime/utils/tensor_backper.py b/slime/utils/tensor_backper.py
index 2fc2a6359d..7d876a4be5 100644
--- a/slime/utils/tensor_backper.py
+++ b/slime/utils/tensor_backper.py
@@ -4,6 +4,8 @@
import torch
+from slime.utils import accelerator
+
_SourceGetter = Callable[[], Iterable[tuple[str, torch.Tensor]]]
@@ -58,7 +60,7 @@ def backup(self, tag: str) -> None:
if name not in backup_dict:
backup_dict[name] = torch.empty_like(param, device=torch.device("cpu"), pin_memory=True)
backup_dict[name].copy_(param.detach(), non_blocking=True)
- torch.cuda.synchronize()
+ accelerator.synchronize()
@torch.no_grad()
def copy(self, *, src_tag: str, dst_tag: str):
@@ -71,7 +73,7 @@ def restore(self, tag: str) -> None:
for name, param in self._source_getter():
assert name in backup_dict
param.copy_(backup_dict[name], non_blocking=True)
- torch.cuda.synchronize()
+ accelerator.synchronize()
class _TensorBackuperNoop(TensorBackuper):
@@ -94,12 +96,12 @@ def get(self, tag: str):
def backup(self, tag: str) -> None:
assert tag == self._single_tag
self._backup_hash_dict = _compute_hash_dict(dict(self._source_getter()))
- torch.cuda.synchronize()
+ accelerator.synchronize()
def restore(self, tag: str) -> None:
assert tag == self._single_tag
assert _compute_hash_dict(dict(self._source_getter())) == self._backup_hash_dict
- torch.cuda.synchronize()
+ accelerator.synchronize()
def _compute_hash_dict(tensors: dict[str, torch.Tensor]):
diff --git a/slime/utils/train_metric_utils.py b/slime/utils/train_metric_utils.py
deleted file mode 100644
index 8ecefcfbf3..0000000000
--- a/slime/utils/train_metric_utils.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import logging
-from argparse import Namespace
-from collections.abc import Callable
-from copy import deepcopy
-
-from slime.utils import logging_utils
-from slime.utils.metric_utils import compute_rollout_step
-from slime.utils.timer import Timer
-
-logger = logging.getLogger(__name__)
-
-
-def log_perf_data_raw(
- rollout_id: int,
- args: Namespace,
- is_primary_rank: bool,
- compute_total_fwd_flops: Callable,
- extra_metrics: dict | None = None,
-) -> None:
- timer_instance = Timer()
- log_dict_raw = deepcopy(timer_instance.log_dict())
- timer_instance.reset()
-
- if not is_primary_rank:
- return
-
- log_dict = {f"perf/{key}_time": val for key, val in log_dict_raw.items()}
- if extra_metrics:
- log_dict.update(extra_metrics)
-
- if ("perf/actor_train_time" in log_dict) and (compute_total_fwd_flops is not None):
- total_fwd_flops = compute_total_fwd_flops(seq_lens=timer_instance.seq_lens)
-
- if "perf/log_probs_time" in log_dict:
- log_dict["perf/log_probs_tflops"] = total_fwd_flops / log_dict["perf/log_probs_time"]
-
- if "perf/ref_log_probs_time" in log_dict:
- log_dict["perf/ref_log_probs_tflops"] = total_fwd_flops / log_dict["perf/ref_log_probs_time"]
-
- if log_dict["perf/actor_train_time"] > 0:
- log_dict["perf/actor_train_tflops"] = 3 * total_fwd_flops / log_dict["perf/actor_train_time"]
- log_dict["perf/actor_train_tok_per_s"] = sum(timer_instance.seq_lens) / log_dict["perf/actor_train_time"]
-
- if "perf/train_wait_time" in log_dict and "perf/train_time" in log_dict:
- total_time = log_dict["perf/train_wait_time"] + log_dict["perf/train_time"]
- if total_time > 0:
- log_dict["perf/step_time"] = total_time
- log_dict["perf/wait_time_ratio"] = log_dict["perf/train_wait_time"] / total_time
-
- logger.info(f"perf {rollout_id}: {log_dict}")
-
- step = compute_rollout_step(args, rollout_id)
- log_dict["rollout/step"] = step
- logging_utils.log(args, log_dict, step_key="rollout/step")
diff --git a/slime_plugins/models/flash_dot_product_attention.py b/slime_plugins/models/flash_dot_product_attention.py
index 06aef6dc3a..fd89bdc464 100644
--- a/slime_plugins/models/flash_dot_product_attention.py
+++ b/slime_plugins/models/flash_dot_product_attention.py
@@ -18,6 +18,7 @@
from megatron.core.utils import divide
from torch import Tensor
+from slime.utils import accelerator
from slime_plugins.models.learnable_softmax_attention import learnable_softmax_flash_attn_varlen
@@ -75,7 +76,7 @@ def __init__(
elif config.softmax_type == "off-by-one":
self.softmax_offset = torch.zeros(
num_heads_per_partition,
- device=torch.cuda.current_device(),
+ device=accelerator.current_device(),
dtype=config.params_dtype,
)
elif config.softmax_type == "learnable":
@@ -84,7 +85,7 @@ def __init__(
torch.nn.Parameter(
torch.empty(
num_heads_per_partition,
- device=torch.cuda.current_device(),
+ device=accelerator.current_device(),
dtype=config.params_dtype,
)
),
diff --git a/slime_plugins/models/qwen3_5.py b/slime_plugins/models/qwen3_5.py
index a53fb15dc9..d7551f9a94 100644
--- a/slime_plugins/models/qwen3_5.py
+++ b/slime_plugins/models/qwen3_5.py
@@ -9,6 +9,8 @@
from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
from transformers.activations import ACT2FN
+from slime.utils import accelerator
+
try:
from fla.modules import FusedRMSNormGated, ShortConvolution
except ImportError:
@@ -77,7 +79,7 @@ def __init__(self, config, layer_idx: int, args=None):
self.head_v_dim,
eps=self.layer_norm_epsilon,
activation=self.activation,
- device=torch.cuda.current_device(),
+ device=accelerator.current_device(),
dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(),
)
diff --git a/slime_plugins/models/qwen3_5_vl.py b/slime_plugins/models/qwen3_5_vl.py
index 2403091bc4..cdc1ad7622 100644
--- a/slime_plugins/models/qwen3_5_vl.py
+++ b/slime_plugins/models/qwen3_5_vl.py
@@ -9,6 +9,8 @@
from megatron.core.transformer.module import MegatronModule
from transformers import AutoConfig
+from slime.utils import accelerator
+
from .qwen3_5 import get_qwen3_5_spec
from .qwen3_5_vl_utils import build_packed_mrope_position_ids, gather_packed_input_ids, get_packed_cp_local_indices
@@ -56,8 +58,8 @@ def _load_vision_model(hf_config, dtype: torch.dtype, use_cpu_initialization: bo
vision_model_cls = Qwen3_5VisionModel
- device = torch.device("cpu") if use_cpu_initialization else torch.device("cuda", torch.cuda.current_device())
- with device:
+ device = torch.device("cpu") if use_cpu_initialization else accelerator.current_device()
+ with torch.device(device):
vision_model = vision_model_cls._from_config(hf_config.vision_config)
vision_model.to(dtype=dtype)
diff --git a/slime_plugins/models/qwen3_next.py b/slime_plugins/models/qwen3_next.py
index f73d57bcf7..cc2937228e 100644
--- a/slime_plugins/models/qwen3_next.py
+++ b/slime_plugins/models/qwen3_next.py
@@ -9,6 +9,8 @@
from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
from transformers.activations import ACT2FN
+from slime.utils import accelerator
+
from .hf_attention import _load_hf_config
try:
@@ -70,7 +72,7 @@ def __init__(self, config, layer_idx: int, args=None):
self.head_v_dim,
eps=self.layer_norm_epsilon,
activation=self.activation,
- device=torch.cuda.current_device(),
+ device=accelerator.current_device(),
dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(),
)
diff --git a/slime_plugins/rollout_buffer/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py
index 74b4b4c46c..ac1fb6998f 100644
--- a/slime_plugins/rollout_buffer/rollout_buffer_example.py
+++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py
@@ -123,7 +123,7 @@ def log_raw_info(args, all_meta_info, rollout_id):
wandb.log(log_dict)
if args.use_tensorboard:
- from slime.utils.tensorboard_utils import _TensorboardAdapter
+ from slime.observability.tensorboard_utils import _TensorboardAdapter
tb = _TensorboardAdapter(args)
tb.log(data=log_dict, step=step)
diff --git a/tests/_cp_dist_helpers.py b/tests/_cp_dist_helpers.py
index f5fa39ab2a..351dadfc28 100644
--- a/tests/_cp_dist_helpers.py
+++ b/tests/_cp_dist_helpers.py
@@ -41,7 +41,6 @@
import sys
import types
-
# --- Stub ``megatron.core.mpu`` (must run before cp_utils is imported) ---
#
# Both this module and any test file that imports it should *import this
diff --git a/slime/utils/compare_glm52_layerwise.py b/tests/glm52_layerwise_comparator.py
similarity index 99%
rename from slime/utils/compare_glm52_layerwise.py
rename to tests/glm52_layerwise_comparator.py
index 7d59a763a5..93227f1208 100644
--- a/slime/utils/compare_glm52_layerwise.py
+++ b/tests/glm52_layerwise_comparator.py
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
-"""Compare matching Megatron and SGLang decoder-layer outputs."""
+"""Compare matching Megatron and SGLang decoder-layer outputs for tests."""
from __future__ import annotations
diff --git a/tests/utils/test_trace_utils.py b/tests/observability/test_trace_utils.py
similarity index 93%
rename from tests/utils/test_trace_utils.py
rename to tests/observability/test_trace_utils.py
index e8337d6cbe..dc19f1f36d 100644
--- a/tests/utils/test_trace_utils.py
+++ b/tests/observability/test_trace_utils.py
@@ -5,9 +5,11 @@
import pytest
import torch
-from slime.utils.trace_utils import TRACE_CHILDREN_KEY, build_sglang_meta_trace_attrs, trace_span
+from slime.observability.trace_utils import TRACE_CHILDREN_KEY, build_sglang_meta_trace_attrs, trace_span
from slime.utils.types import Sample
+NUM_GPUS = 0
+
def _load_trace_timeline_viewer_module():
module_path = Path(__file__).resolve().parents[2] / "tools" / "trace_timeline_viewer.py"
@@ -91,3 +93,7 @@ def test_trace_timeline_viewer_omits_virtual_pd_lanes_without_pd_attrs(tmp_path:
}
assert "[P]" not in item["name"]
assert "[D]" not in item["name"]
+
+
+if __name__ == "__main__":
+ raise SystemExit(pytest.main([__file__]))
diff --git a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py
index a8380feecc..d4192e10a7 100644
--- a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py
+++ b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py
@@ -133,7 +133,7 @@ def invoke_rollout_data_postprocess(fn):
"custom_rollout_log",
"CUSTOM_ROLLOUT_LOG_FUNCTION_PATH",
"plugin_contracts.test_plugin_runtime_hook_contracts.reference_custom_rollout_log",
- "slime/ray/rollout.py",
+ "slime/observability/rollout_metrics.py",
"custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time)",
("rollout_id", "args", "samples", "rollout_extra_metrics", "rollout_time"),
invoke_custom_rollout_log,
@@ -142,7 +142,7 @@ def invoke_rollout_data_postprocess(fn):
"custom_eval_rollout_log",
"CUSTOM_EVAL_ROLLOUT_LOG_FUNCTION_PATH",
"plugin_contracts.test_plugin_runtime_hook_contracts.reference_custom_eval_rollout_log",
- "slime/ray/rollout.py",
+ "slime/observability/rollout_metrics.py",
"custom_log_func(rollout_id, args, data, extra_metrics)",
("rollout_id", "args", "data", "extra_metrics"),
invoke_custom_eval_rollout_log,
diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py
new file mode 100644
index 0000000000..301737550e
--- /dev/null
+++ b/tests/test_accelerator.py
@@ -0,0 +1,203 @@
+from types import SimpleNamespace
+
+import pytest
+
+from slime.utils import accelerator
+
+NUM_GPUS = 0
+
+
+class FakeAccelerator(accelerator.Accelerator):
+ name = "fake"
+ device_type = "fake"
+ communication_backend_name = "fake"
+
+ def is_available(self):
+ return True
+
+ def device(self, index=None):
+ return accelerator.torch.device("cpu")
+
+ def device_name(self, index=None):
+ return "cpu"
+
+ def set_device(self, index):
+ return None
+
+ def current_device(self):
+ return "cpu"
+
+ def device_count(self):
+ return 0
+
+ def synchronize(self, device=None):
+ return None
+
+ def current_stream(self, device=None):
+ return None
+
+ def empty_cache(self):
+ return None
+
+ def mem_get_info(self, device=None):
+ return 0, 0
+
+ def memory_allocated(self, device=None):
+ return 0
+
+ def memory_reserved(self, device=None):
+ return 0
+
+
+@pytest.fixture(autouse=True)
+def reset_accelerator_selection(monkeypatch):
+ registry = accelerator._REGISTRY.copy()
+ selected = accelerator._ACCELERATOR
+ patch_imported = accelerator._MUSA_PATCH_IMPORTED
+ bootstrap_checked = accelerator._MUSA_BOOTSTRAP_CHECKED
+ for name in ("SLIME_ACCELERATOR", "MUSA_VISIBLE_DEVICES", "MUSA_PATCH_PATH", "CUDA_VISIBLE_DEVICES"):
+ monkeypatch.delenv(name, raising=False)
+ accelerator._REGISTRY.clear()
+ accelerator.reset_accelerator()
+ accelerator._MUSA_PATCH_IMPORTED = False
+ accelerator._MUSA_BOOTSTRAP_CHECKED = False
+ yield
+ accelerator._REGISTRY.clear()
+ accelerator._REGISTRY.update(registry)
+ accelerator._ACCELERATOR = selected
+ accelerator._MUSA_PATCH_IMPORTED = patch_imported
+ accelerator._MUSA_BOOTSTRAP_CHECKED = bootstrap_checked
+
+
+@pytest.mark.unit
+def test_cuda_selection_does_not_bootstrap_musa(monkeypatch):
+ monkeypatch.setenv("SLIME_ACCELERATOR", "cuda")
+ monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0")
+ monkeypatch.setattr(accelerator, "_cuda_available", lambda: True)
+ monkeypatch.setattr(accelerator.CUDAAccelerator, "is_available", lambda self: True)
+ monkeypatch.setattr(
+ accelerator,
+ "_import_musa_patch",
+ lambda: pytest.fail("CUDA selection must not import musa_patch"),
+ )
+
+ assert accelerator.get_accelerator().name == "cuda"
+ assert accelerator.process_group_backend() == "nccl"
+ assert accelerator.visible_devices_env_key() == "CUDA_VISIBLE_DEVICES"
+
+
+@pytest.mark.unit
+def test_selected_musa_bootstraps_patch_once(monkeypatch):
+ imports = []
+ fake_musa = SimpleNamespace(is_available=lambda: True)
+
+ def import_musa_patch():
+ imports.append("musa_patch")
+ monkeypatch.setattr(accelerator.torch, "musa", fake_musa, raising=False)
+ return True
+
+ monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0")
+ monkeypatch.setattr(accelerator, "_import_musa_patch", import_musa_patch)
+
+ assert imports == []
+ assert accelerator.initialize_accelerator().name == "musa"
+ assert accelerator.initialize_accelerator().name == "musa"
+ assert imports == ["musa_patch"]
+
+
+@pytest.mark.unit
+def test_cpu_only_initialization_does_not_require_an_accelerator(monkeypatch):
+ monkeypatch.setattr(accelerator, "is_musa_available", lambda: False)
+ monkeypatch.setattr(accelerator, "_cuda_available", lambda: False)
+
+ assert accelerator.initialize_accelerator() is None
+
+
+@pytest.mark.unit
+def test_musa_backend_maps_devices_and_process_groups(monkeypatch):
+ monkeypatch.setattr(accelerator.MUSAAccelerator, "is_available", lambda self: True)
+ monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "2,5")
+ accelerator.set_accelerator(accelerator.MUSAAccelerator())
+
+ assert accelerator.visible_devices_env_key() == "MUSA_VISIBLE_DEVICES"
+ assert accelerator.resolve_visible_device_id("5") == 1
+ assert accelerator.process_group_backend() == "mccl"
+ assert accelerator.weight_update_backend() == "cpu:gloo,musa:mccl"
+ assert accelerator.process_group_backend("gloo") == "gloo"
+
+
+@pytest.mark.unit
+def test_cuda_visible_device_mapping(monkeypatch):
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,6")
+ accelerator.set_accelerator(FakeAccelerator())
+
+ assert accelerator.resolve_visible_device_id(4) == 0
+ assert accelerator.resolve_visible_device_id(1) == 1
+ with pytest.raises(RuntimeError, match="CUDA_VISIBLE_DEVICES=4,6"):
+ accelerator.resolve_visible_device_id(7)
+
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-aaaa,GPU-bbbb")
+ assert accelerator.resolve_visible_device_id("GPU-bbbb") == 1
+
+
+@pytest.mark.unit
+def test_registered_backend_can_be_selected(monkeypatch):
+ class RegisteredAccelerator(FakeAccelerator):
+ name = "registered"
+
+ monkeypatch.setattr(accelerator, "is_musa_available", lambda: False)
+ monkeypatch.setattr(accelerator, "_cuda_available", lambda: False)
+ accelerator.register_accelerator("registered", RegisteredAccelerator, lambda: True, priority=300)
+
+ assert accelerator.get_accelerator().name == "registered"
+
+
+@pytest.mark.unit
+def test_cuda_backend_uses_torch_cuda_namespace(monkeypatch):
+ monkeypatch.setattr(accelerator.torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(accelerator.torch.cuda, "device_count", lambda: 2)
+ monkeypatch.setattr(accelerator.torch.cuda, "current_device", lambda: 1)
+ monkeypatch.setattr(accelerator.torch.cuda, "memory_allocated", lambda device=None: 123)
+ backend = accelerator.CUDAAccelerator()
+
+ assert backend.is_available()
+ assert backend.device_name() == "cuda:1"
+ assert backend.memory_allocated() == 123
+
+
+@pytest.mark.unit
+def test_routing_replay_uses_selected_backend_current_device(monkeypatch):
+ from slime.utils import routing_replay
+
+ transfers = []
+
+ class FakeTopIndices:
+ def is_pinned(self):
+ return False
+
+ def to(self, device, *, dtype, non_blocking):
+ transfers.append((device, dtype, non_blocking))
+ return self
+
+ accelerator.set_accelerator(FakeAccelerator())
+ monkeypatch.setattr(routing_replay.RoutingReplay, "all_routing_replays", [])
+ replay = routing_replay.RoutingReplay()
+ replay.top_indices_list.append(FakeTopIndices())
+
+ replay.pop_forward()
+ replay.pop_backward()
+ assert transfers == [
+ ("cpu", accelerator.torch.int32, False),
+ ("cpu", accelerator.torch.int32, False),
+ ]
+
+
+@pytest.mark.unit
+def test_musa_availability_handles_missing_torch_namespace(monkeypatch):
+ monkeypatch.delattr(accelerator.torch, "musa", raising=False)
+
+ assert accelerator.is_musa_available() is False
+
+
+if __name__ == "__main__":
+ raise SystemExit(pytest.main([__file__]))
diff --git a/tests/test_empty_colocated_weight_bucket.py b/tests/test_empty_colocated_weight_bucket.py
index 7e5db7e0db..4b5217cd2c 100644
--- a/tests/test_empty_colocated_weight_bucket.py
+++ b/tests/test_empty_colocated_weight_bucket.py
@@ -67,6 +67,10 @@ def _install_fake_deps(monkeypatch):
update_weight_pkg.__path__ = [str(REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight")]
slime_utils_pkg = types.ModuleType("slime.utils")
slime_utils_pkg.__path__ = [str(REPO_ROOT / "slime" / "utils")]
+ accelerator_mod = types.ModuleType("slime.utils.accelerator")
+ accelerator_mod.device = lambda: "cuda:0"
+ accelerator_mod.current_device = lambda: "cuda:0"
+ accelerator_mod.ipc_collect = lambda: None
dist_mod = types.ModuleType("torch.distributed")
@@ -133,6 +137,7 @@ def gather_object(obj, object_gather_list, dst, group):
monkeypatch.setitem(sys.modules, "slime.backends.megatron_utils", megatron_utils_pkg)
monkeypatch.setitem(sys.modules, "slime.backends.megatron_utils.update_weight", update_weight_pkg)
monkeypatch.setitem(sys.modules, "slime.utils", slime_utils_pkg)
+ monkeypatch.setitem(sys.modules, "slime.utils.accelerator", accelerator_mod)
monkeypatch.setitem(sys.modules, "torch", torch_mod)
monkeypatch.setitem(sys.modules, "torch.distributed", dist_mod)
monkeypatch.setitem(sys.modules, "ray", ray_mod)
diff --git a/tests/test_glm52_6layer_deterministic_e2e.py b/tests/test_glm52_6layer_deterministic_e2e.py
index 1f94cc0732..241ec4f1c8 100644
--- a/tests/test_glm52_6layer_deterministic_e2e.py
+++ b/tests/test_glm52_6layer_deterministic_e2e.py
@@ -365,8 +365,7 @@ def run_gate(*, layerwise_zero: bool = False, rollout_max_response_len: int = 40
_run(
[
sys.executable,
- "-m",
- "slime.utils.compare_glm52_layerwise",
+ str(REPO_ROOT / "tests/glm52_layerwise_comparator.py"),
"--megatron-dir",
megatron_layerwise_dump,
"--sglang-dir",
diff --git a/tests/test_glm52_layerwise_comparison.py b/tests/test_glm52_layerwise_comparison.py
index e1d528ebb2..9228646461 100644
--- a/tests/test_glm52_layerwise_comparison.py
+++ b/tests/test_glm52_layerwise_comparison.py
@@ -1,7 +1,7 @@
import pytest
import torch
-from slime.utils.compare_glm52_layerwise import (
+from glm52_layerwise_comparator import (
TrainSequence,
_sglang_layer_token_rows,
compare_layer_outputs,
diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py
index ba3304d79c..20c133ed95 100644
--- a/tests/test_megatron_argument_validation.py
+++ b/tests/test_megatron_argument_validation.py
@@ -43,7 +43,7 @@ def load_slime_arguments_module(monkeypatch):
router_launch_mod = types.ModuleType("sglang_router.launch_router")
sglang_arguments_mod = types.ModuleType("slime.backends.sglang_utils.arguments")
sglang_external_mod = types.ModuleType("slime.backends.sglang_utils.external")
- logging_utils_mod = types.ModuleType("slime.utils.logging_utils")
+ logging_utils_mod = types.ModuleType("slime.observability.logging_utils")
router_launch_mod.RouterArgs = object
sglang_arguments_mod.sglang_parse_args = lambda *args, **kwargs: None
@@ -55,7 +55,7 @@ def load_slime_arguments_module(monkeypatch):
monkeypatch.setitem(sys.modules, "sglang_router.launch_router", router_launch_mod)
monkeypatch.setitem(sys.modules, "slime.backends.sglang_utils.arguments", sglang_arguments_mod)
monkeypatch.setitem(sys.modules, "slime.backends.sglang_utils.external", sglang_external_mod)
- monkeypatch.setitem(sys.modules, "slime.utils.logging_utils", logging_utils_mod)
+ monkeypatch.setitem(sys.modules, "slime.observability.logging_utils", logging_utils_mod)
module_path = Path(__file__).resolve().parents[1] / "slime" / "utils" / "arguments.py"
module_name = "test_slime_argument_validation_module"
@@ -255,6 +255,7 @@ def make_slime_validate_args(**overrides):
update_weight_disk_dir=None,
update_weight_local_checkpoint_dir=None,
update_weight_mode="full",
+ rollout_temperature=1.0,
)
values.update(overrides)
return types.SimpleNamespace(**values)
@@ -299,6 +300,16 @@ def test_slime_validate_args_rejects_equal_debug_data_paths(monkeypatch):
module.slime_validate_args(args)
+@pytest.mark.unit
+@pytest.mark.parametrize("temperature", [0.0, -0.1])
+def test_slime_validate_args_rejects_non_positive_rollout_temperature(monkeypatch, temperature):
+ module = load_slime_arguments_module(monkeypatch)
+ args = make_slime_validate_args(rollout_temperature=temperature)
+
+ with pytest.raises(ValueError, match="--rollout-temperature must be > 0"):
+ module.slime_validate_args(args)
+
+
@pytest.mark.unit
def test_slime_validate_args_preserves_zero_rollout_gpus_under_colocate(monkeypatch):
module = load_slime_arguments_module(monkeypatch)
diff --git a/tests/test_metric_report.py b/tests/test_metric_report.py
index d908597c5c..a83f7fda1b 100644
--- a/tests/test_metric_report.py
+++ b/tests/test_metric_report.py
@@ -1,8 +1,7 @@
"""Single-process metric-report invariance tests.
Pins train-side / rollout-side report formulas implemented in
-``slime.backends.megatron_utils.cp_utils.reduce_train_step_metrics`` and
-``rollout_log_metric_contribution``: the reported number for a given set
+``slime.observability.train_metric_utils``: the reported number for a given set
of samples must be the same regardless of
- how samples are distributed across micro-batches / DP ranks
@@ -28,11 +27,12 @@
from slime.backends.megatron_utils.cp_utils import ( # noqa: E402
get_logits_and_tokens_offset_with_cp,
get_sum_of_sample_mean,
+)
+from slime.observability.train_metric_utils import ( # noqa: E402
reduce_train_step_metrics,
rollout_log_metric_contribution,
)
-
NUM_GPUS = 0
@@ -143,7 +143,7 @@ def _simulate_rollout_report(samples_per_rank):
per-token metric branch.
Each "rank" applies the reducer once over its full sample subset, then
- ``rollout_log_metric_contribution`` (the same helper data.py uses) emits
+ ``rollout_log_metric_contribution`` (the same helper the reporter uses) emits
the ``(per_rank_sum, count)`` tuple. We aggregate via
``Σsum / Σcount`` — the same shape ``gather_log_data`` uses.
"""
diff --git a/tests/test_metric_report_dist.py b/tests/test_metric_report_dist.py
index 614720d5ed..466ad3b729 100644
--- a/tests/test_metric_report_dist.py
+++ b/tests/test_metric_report_dist.py
@@ -1,4 +1,4 @@
-"""Multi-process distributed tests for the cp_utils report helpers.
+"""Multi-process distributed tests for the train metric report helpers.
Spawn ``dp_size * cp_size`` workers with real ``torch.distributed`` (gloo
backend) and exercise the actual production helpers end-to-end. The
@@ -42,7 +42,6 @@
stub_megatron_in_worker,
)
-
NUM_GPUS = 0
@@ -68,7 +67,8 @@ def _train_step_distributed_worker(
# Import AFTER the megatron stub override so cp_utils still binds
# against the pre-installed stub (which we've now pinned for this
# worker's CP rank).
- from slime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean, reduce_train_step_metrics
+ from slime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean
+ from slime.observability.train_metric_utils import reduce_train_step_metrics
all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS
all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS
@@ -197,11 +197,8 @@ def _rollout_log_distributed_worker(
dp_group = init_worker_process_group(rank, world_size, master_port)
try:
- from slime.backends.megatron_utils.cp_utils import (
- gather_and_reduce_log_dict,
- get_sum_of_sample_mean,
- rollout_log_metric_contribution,
- )
+ from slime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean
+ from slime.observability.train_metric_utils import gather_and_reduce_log_dict, rollout_log_metric_contribution
all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS
all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS
diff --git a/tests/test_ppo_kl_metric.py b/tests/test_ppo_kl_metric.py
new file mode 100644
index 0000000000..64d85aa479
--- /dev/null
+++ b/tests/test_ppo_kl_metric.py
@@ -0,0 +1,63 @@
+import sys
+import types
+from argparse import Namespace
+
+import torch
+
+from slime.utils.ppo_utils import compute_approx_kl
+
+NUM_GPUS = 0
+
+
+def test_ppo_estimator_does_not_corrupt_logged_kl(monkeypatch):
+ previous_loss = sys.modules.pop("slime.backends.megatron_utils.loss", None)
+ previous_cp_utils = sys.modules.pop("slime.backends.megatron_utils.cp_utils", None)
+
+ mpu_stub = types.SimpleNamespace(
+ get_context_parallel_world_size=lambda: 1,
+ get_context_parallel_rank=lambda: 0,
+ is_pipeline_last_stage=lambda: True,
+ )
+ megatron_mod = types.ModuleType("megatron")
+ core_mod = types.ModuleType("megatron.core")
+ core_mod.mpu = mpu_stub
+ monkeypatch.setitem(sys.modules, "megatron", megatron_mod)
+ monkeypatch.setitem(sys.modules, "megatron.core", core_mod)
+
+ try:
+ from slime.backends.megatron_utils.loss import compute_advantages_and_returns
+
+ log_probs = [torch.tensor([0.5, 0.7, 0.9])]
+ ref_log_probs = [torch.tensor([0.4, 0.5, 0.6])]
+ expected_kl = compute_approx_kl(log_probs[0], ref_log_probs[0], kl_loss_type="k1")
+ rollout_data = {
+ "log_probs": log_probs,
+ "ref_log_probs": ref_log_probs,
+ "rewards": [1.0],
+ "values": [torch.zeros(3)],
+ "response_lengths": [3],
+ "total_lengths": [5],
+ "loss_masks": [torch.ones(3)],
+ }
+ args = Namespace(
+ advantage_estimator="ppo",
+ kl_coef=0.05,
+ kl_loss_type="k1",
+ use_rollout_logprobs=False,
+ custom_advantage_function_path=None,
+ normalize_advantages=False,
+ use_opd=False,
+ gamma=1.0,
+ lambd=1.0,
+ )
+ compute_advantages_and_returns(args, rollout_data)
+ torch.testing.assert_close(rollout_data["kl"][0], expected_kl)
+ finally:
+ if previous_loss is None:
+ sys.modules.pop("slime.backends.megatron_utils.loss", None)
+ else:
+ sys.modules["slime.backends.megatron_utils.loss"] = previous_loss
+ if previous_cp_utils is None:
+ sys.modules.pop("slime.backends.megatron_utils.cp_utils", None)
+ else:
+ sys.modules["slime.backends.megatron_utils.cp_utils"] = previous_cp_utils
diff --git a/tests/test_qwen2.5_0.5B_fanout_short.py b/tests/test_qwen2.5_0.5B_fanout_short.py
index 289a6e725c..5a90954b53 100644
--- a/tests/test_qwen2.5_0.5B_fanout_short.py
+++ b/tests/test_qwen2.5_0.5B_fanout_short.py
@@ -10,7 +10,7 @@
this test, **no e2e training run had ever exercised the full chain**:
custom_generate returns list[Sample] sharing rollout_id
- → _validate_rollout_id_annotated at depth ≥ 2 passes
+ → validate_rollout_id_annotated at depth ≥ 2 passes
→ _split_train_data_by_dp groups by rollout_id and trims to N steps
using ``rollout_batch_size * n_samples_per_prompt / global_batch_size``
(NOT total sample count, which would inflate steps once N>1)
diff --git a/tests/test_reloadable_process_group_world.py b/tests/test_reloadable_process_group_world.py
index 38822a2748..289b0bebf0 100644
--- a/tests/test_reloadable_process_group_world.py
+++ b/tests/test_reloadable_process_group_world.py
@@ -26,10 +26,10 @@ def _run_pp_group_reload_worker(rank: int, world_size: int, rendezvous_path: str
distributed_utils.init_gloo_group()
rpg.register_default_process_group(timeout=timeout)
- # Exercise the NCCL lifecycle with Gloo so this remains a CPU test. The
- # relevant contract is the global ordering of WORLD and subgroup teardown,
- # not the backend implementation.
- rpg._uses_nccl = lambda _backend: True
+ # Exercise the accelerator lifecycle with Gloo so this remains a CPU test.
+ # The contract under test is WORLD/subgroup teardown ordering, not a vendor
+ # communication backend.
+ rpg._uses_accelerator_backend = lambda _backend: True
group_specs = [
([0], "TP_0"),
@@ -65,6 +65,49 @@ def _run_pp_group_reload_worker(rank: int, world_size: int, rendezvous_path: str
dist.destroy_process_group()
+def _run_backend_normalization_worker(_rank: int) -> None:
+ calls = []
+ mapped_backends = []
+
+ def old_new_group(*args, **kwargs):
+ calls.append((args, kwargs))
+ return f"group-{len(calls)}"
+
+ def process_group_backend(backend):
+ mapped_backends.append(backend)
+ return "mccl" if backend == "nccl" else backend
+
+ rpg.old_new_group_dict.clear()
+ rpg.default_process_group_states.clear()
+ rpg.dist.new_group = old_new_group
+ rpg.dist.get_backend = lambda: "gloo"
+ rpg.accelerator.process_group_backend = process_group_backend
+ rpg.monkey_patch_torch_dist()
+
+ gloo_group = rpg.dist.new_group(ranks=[0], backend="gloo")
+ assert gloo_group == "group-1"
+ assert calls[-1][1]["backend"] == "gloo"
+ assert mapped_backends == []
+
+ mccl_group = rpg.dist.new_group([0], None, "nccl")
+ assert mccl_group == "group-2"
+ assert calls[-1][0][2] == "mccl"
+ assert mapped_backends == ["nccl"]
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("backend", ["nccl", "mccl", "cpu:gloo,musa:mccl"])
+def test_accelerator_backend_detection(backend):
+ assert rpg._uses_accelerator_backend(backend)
+
+
+@pytest.mark.unit
+def test_new_group_normalizes_only_logical_nccl_backend():
+ # monkey_patch_torch_dist replaces process-wide torch.distributed symbols,
+ # so isolate this behavior check in a spawned process.
+ mp.spawn(_run_backend_normalization_worker, nprocs=1, join=True)
+
+
@pytest.mark.unit
def test_register_default_process_group_captures_rendezvous_state(monkeypatch):
timeout = timedelta(minutes=7)
@@ -83,7 +126,7 @@ def test_register_default_process_group_captures_rendezvous_state(monkeypatch):
assert state.store == "rendezvous-store"
assert state.rank == 3
assert state.world_size == 8
- assert not state.nccl_world_destroyed
+ assert not state.accelerator_world_destroyed
@pytest.mark.unit
@@ -127,7 +170,7 @@ def init_process_group(**kwargs):
rpg.destroy_process_groups()
- assert state.nccl_world_destroyed
+ assert state.accelerator_world_destroyed
assert state.generation == 1
assert events == [
("barrier", "canonical-gloo"),
@@ -150,7 +193,7 @@ def init_process_group(**kwargs):
events.clear()
rpg.reload_process_groups()
- assert not state.nccl_world_destroyed
+ assert not state.accelerator_world_destroyed
assert state.generation == 2
assert events == [
("barrier", "WORLD"),
diff --git a/tests/test_rollout_data_utils.py b/tests/test_rollout_data_utils.py
new file mode 100644
index 0000000000..82a72de105
--- /dev/null
+++ b/tests/test_rollout_data_utils.py
@@ -0,0 +1,102 @@
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from slime.observability.rollout_data_utils import (
+ load_debug_rollout_data,
+ save_debug_rollout_data,
+ tensorize_rollout_data_for_training,
+ validate_rollout_routed_experts_for_replay,
+)
+from slime.utils.types import Sample
+
+NUM_GPUS = 0
+
+
+def _args():
+ return SimpleNamespace(
+ num_layers=6,
+ moe_router_topk=2,
+ moe_layer_freq=[0, 0, 0, 1, 1, 1],
+ )
+
+
+def test_r3_validation_accepts_dense_zeros_and_complete_moe_routes():
+ routes = torch.zeros((4, 6, 2), dtype=torch.uint8)
+ routes[:, 3:, 1] = 7
+ validate_rollout_routed_experts_for_replay([routes], _args())
+
+
+def test_r3_validation_rejects_missing_pipeline_layers():
+ routes = torch.zeros((4, 6, 2), dtype=torch.uint8)
+ routes[:, 3, 1] = 7
+
+ with pytest.raises(ValueError, match=r"all zero.*\[4, 5\]"):
+ validate_rollout_routed_experts_for_replay([routes], _args())
+
+
+def test_r3_validation_rejects_wrong_shape():
+ routes = torch.zeros((4, 5, 2), dtype=torch.uint8)
+
+ with pytest.raises(ValueError, match="Invalid rollout routed-experts shape"):
+ validate_rollout_routed_experts_for_replay([routes], _args())
+
+
+def test_tensorize_rollout_data_for_training_normalizes_cpu_tensors():
+ readonly_tokens = np.array([1, 2, 3])
+ readonly_tokens.flags.writeable = False
+ rollout_data = {
+ "tokens": [readonly_tokens],
+ "loss_masks": [[1, 0]],
+ "multimodal_train_inputs": [
+ {
+ "pixel_values": torch.tensor([1.0], requires_grad=True),
+ "metadata": "unchanged",
+ }
+ ],
+ "rollout_mask_sums": [2],
+ }
+
+ tensorize_rollout_data_for_training(rollout_data)
+
+ assert rollout_data["tokens"][0].dtype == torch.long
+ assert rollout_data["loss_masks"][0].dtype == torch.int
+ assert rollout_data["multimodal_train_inputs"][0]["metadata"] == "unchanged"
+ assert not rollout_data["multimodal_train_inputs"][0]["pixel_values"].requires_grad
+ assert rollout_data["rollout_mask_sums"].dtype == torch.float32
+
+
+def test_save_and_load_debug_rollout_data_round_trip(tmp_path):
+ path_template = str(tmp_path / "rollout_{rollout_id}.pt")
+ samples = [
+ Sample(index=1, rollout_id=3, prompt="question", response="answer", response_length=1),
+ ]
+
+ save_debug_rollout_data(path_template, samples, rollout_id=3, evaluation=False)
+ loaded = load_debug_rollout_data(path_template, rollout_id=3)
+
+ assert len(loaded) == 1
+ assert loaded[0].index == 1
+ assert loaded[0].rollout_id == 3
+ assert loaded[0].prompt == "question"
+ assert loaded[0].response == "answer"
+
+
+def test_save_debug_eval_rollout_data_flattens_datasets(tmp_path):
+ path_template = str(tmp_path / "rollout_{rollout_id}.pt")
+ data = {
+ "math": {"samples": [Sample(index=1)]},
+ "code": {"samples": [Sample(index=2)]},
+ }
+
+ save_debug_rollout_data(path_template, data, rollout_id=4, evaluation=True)
+
+ saved = torch.load(tmp_path / "rollout_eval_4.pt", weights_only=False)
+ assert saved["rollout_id"] == 4
+ assert [sample["index"] for sample in saved["samples"]] == [1, 2]
+
+
+if __name__ == "__main__":
+ raise SystemExit(pytest.main([__file__]))
diff --git a/tests/test_rollout_metrics.py b/tests/test_rollout_metrics.py
index 02fc92bce4..60faef145f 100644
--- a/tests/test_rollout_metrics.py
+++ b/tests/test_rollout_metrics.py
@@ -5,7 +5,7 @@
import pytest
import torch
-from slime.ray.rollout import _compute_top_p_kept_vocab_metrics
+from slime.observability.rollout_metrics import _compute_top_p_kept_vocab_metrics
from slime.utils.misc import decode_int32_meta_array
from slime.utils.types import Sample
@@ -215,3 +215,7 @@ def test_append_response_tokens_rejects_non_trainable_log_probs():
with pytest.raises(ValueError, match="non-trainable response tokens should not pass rollout log probabilities"):
sample.append_response_tokens(tokens=[10], log_probs=[-0.1], trainable=False)
+
+
+if __name__ == "__main__":
+ raise SystemExit(pytest.main([__file__]))
diff --git a/tests/test_rollout_routing_replay_validation.py b/tests/test_rollout_routing_replay_validation.py
deleted file mode 100644
index 55c7a4fa6b..0000000000
--- a/tests/test_rollout_routing_replay_validation.py
+++ /dev/null
@@ -1,41 +0,0 @@
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from slime.ray.rollout import _validate_rollout_routed_experts_for_replay
-
-NUM_GPUS = 0
-
-
-def _args():
- return SimpleNamespace(
- num_layers=6,
- moe_router_topk=2,
- moe_layer_freq=[0, 0, 0, 1, 1, 1],
- )
-
-
-def test_r3_validation_accepts_dense_zeros_and_complete_moe_routes():
- routes = torch.zeros((4, 6, 2), dtype=torch.uint8)
- routes[:, 3:, 1] = 7
- _validate_rollout_routed_experts_for_replay([routes], _args())
-
-
-def test_r3_validation_rejects_missing_pipeline_layers():
- routes = torch.zeros((4, 6, 2), dtype=torch.uint8)
- routes[:, 3, 1] = 7
-
- with pytest.raises(ValueError, match=r"all zero.*\[4, 5\]"):
- _validate_rollout_routed_experts_for_replay([routes], _args())
-
-
-def test_r3_validation_rejects_wrong_shape():
- routes = torch.zeros((4, 5, 2), dtype=torch.uint8)
-
- with pytest.raises(ValueError, match="Invalid rollout routed-experts shape"):
- _validate_rollout_routed_experts_for_replay([routes], _args())
-
-
-if __name__ == "__main__":
- raise SystemExit(pytest.main([__file__]))
diff --git a/tests/test_train_dump.py b/tests/test_train_data_utils.py
similarity index 99%
rename from tests/test_train_dump.py
rename to tests/test_train_data_utils.py
index 694673736a..6b94ff29ac 100644
--- a/tests/test_train_dump.py
+++ b/tests/test_train_data_utils.py
@@ -5,7 +5,7 @@
import torch
from _cp_dist_helpers import cp_chunk_response_tensor, free_port, init_worker_process_group, stub_megatron_in_worker
-from slime.backends.megatron_utils.train_dump_utils import (
+from slime.observability.train_data_utils import (
_build_dump_payload,
restore_context_parallel_fields_to_cpu,
save_debug_train_data,
diff --git a/tools/convert_hf_to_fp8.py b/tools/convert_hf_to_fp8.py
index e98e750e22..bcc4306b9d 100644
--- a/tools/convert_hf_to_fp8.py
+++ b/tools/convert_hf_to_fp8.py
@@ -28,6 +28,8 @@
import torch.nn.functional as F
from tqdm import tqdm
+from slime.utils import accelerator
+
FP8_INFO = torch.finfo(torch.float8_e4m3fn)
FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min
@@ -117,11 +119,13 @@ def process_file(input_path, output_path, filename, strategy, block_size, result
if not filename.endswith(".safetensors"):
return
- print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}")
+ print(f"Processing {filename}, memory usage: {accelerator.memory_allocated()}")
weights = {}
q_weights = {}
- with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device="cuda") as f:
+ with safetensors.safe_open(
+ os.path.join(input_path, filename), framework="pt", device=accelerator.device_name()
+ ) as f:
for k in f.keys():
weights[k] = f.get_tensor(k)
@@ -243,7 +247,7 @@ def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers=
json.dump(index_dict, open(os.path.join(output_path, "model.safetensors.index.json"), "w"), indent=2)
gc.collect()
- torch.cuda.empty_cache()
+ accelerator.empty_cache()
if __name__ == "__main__":
diff --git a/tools/convert_hf_to_int4_direct.py b/tools/convert_hf_to_int4_direct.py
index 613f65595d..7eaffc7272 100644
--- a/tools/convert_hf_to_int4_direct.py
+++ b/tools/convert_hf_to_int4_direct.py
@@ -21,6 +21,8 @@
import torch
from tqdm import tqdm
+from slime.utils import accelerator
+
try:
import fake_int4_quant_cuda
except ImportError:
@@ -166,11 +168,13 @@ def add_result(self, filename, q_weights):
def process_file(input_path, output_path, filename, group_size, is_symmetric, ignore_rules, result_collector):
- print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}")
+ print(f"Processing {filename}, memory usage: {accelerator.memory_allocated()}")
weights = {}
q_weights = {}
- with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device="cuda") as f:
+ with safetensors.safe_open(
+ os.path.join(input_path, filename), framework="pt", device=accelerator.device_name()
+ ) as f:
for k in f.keys():
weights[k] = f.get_tensor(k)
@@ -180,15 +184,15 @@ def process_file(input_path, output_path, filename, group_size, is_symmetric, ig
)
if is_ignored or not name.endswith(".weight") or weight.dim() < 2:
- print(f"Ignoring {name}, memory usage: {torch.cuda.memory_allocated()}")
+ print(f"Ignoring {name}, memory usage: {accelerator.memory_allocated()}")
q_weights[name] = weight
continue
- print(f"Packing {name}, memory usage: {torch.cuda.memory_allocated()}")
+ print(f"Packing {name}, memory usage: {accelerator.memory_allocated()}")
qw, s, zp = pack_layer(weight, group_size, is_symmetric)
qweight_name = name.replace(".weight", ".weight_packed")
scale_name = name.replace(".weight", ".weight_scale")
- weight_shape = torch.tensor(weight.shape, dtype=torch.int32, device="cuda")
+ weight_shape = torch.tensor(weight.shape, dtype=torch.int32, device=accelerator.device())
weight_shape_name = name.replace(".weight", ".weight_shape")
if zp is not None:
zp_name = name.replace(".weight", ".weight_zero_point")
@@ -273,7 +277,7 @@ def convert_int4(input_path, output_path, group_size, is_symmetric, ignore_rules
json.dump(index_dict, open(os.path.join(output_path, "model.safetensors.index.json"), "w"), indent=2)
gc.collect()
- torch.cuda.empty_cache()
+ accelerator.empty_cache()
return output_path
diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py
index 4ef5034b10..8cd72b26c3 100644
--- a/tools/convert_hf_to_torch_dist.py
+++ b/tools/convert_hf_to_torch_dist.py
@@ -13,7 +13,8 @@
from slime.backends.megatron_utils.hf_to_megatron import load_hf_weights
from slime.backends.megatron_utils.initialize import init
from slime.backends.megatron_utils.model_provider import get_model_provider_func
-from slime.utils.logging_utils import configure_logger
+from slime.observability.logging_utils import configure_logger
+from slime.utils import accelerator
from slime.utils.memory_utils import print_memory
@@ -27,7 +28,6 @@ def add_convertion_args(parser):
help="Path to a custom model provider function.",
)
parser.add_argument("--allgather-cp", action="store_true", default=False)
- parser.add_argument("--use-gated-attention", action="store_true", default=False)
try:
parser.add_argument("--padded-vocab-size", type=int, default=None)
except Exception:
@@ -81,6 +81,7 @@ def ceildiv(a, b):
def main():
if torch.version.hip:
import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module
+
from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync
filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync
@@ -93,17 +94,17 @@ def main():
local_rank = int(os.getenv("LOCAL_RANK") or os.getenv("SLURM_LOCALID") or 0)
global_rank = int(os.getenv("RANK") or os.getenv("SLURM_PROCID") or 0)
- torch.cuda.set_device(local_rank)
+ accelerator.set_device(local_rank)
os.environ.setdefault("WORLD_SIZE", str(world_size))
os.environ.setdefault("RANK", str(global_rank))
os.environ.setdefault("LOCAL_RANK", str(local_rank))
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "12355")
dist.init_process_group(
- backend="nccl",
+ backend=accelerator.process_group_backend(),
world_size=world_size,
rank=global_rank,
- device_id=torch.device(f"cuda:{local_rank}"),
+ device_id=accelerator.distributed_device_id(local_rank),
)
args = get_args()
init(args)
@@ -123,9 +124,9 @@ def main():
model[0] = model[0].cpu()
print_memory("after loading model")
- torch.cuda.synchronize()
+ accelerator.synchronize()
gc.collect()
- torch.cuda.empty_cache()
+ accelerator.empty_cache()
save_checkpoint(1, model, None, None, 0)
diff --git a/tools/convert_to_hf.py b/tools/convert_to_hf.py
index 9f16850f0c..201e35bf9c 100644
--- a/tools/convert_to_hf.py
+++ b/tools/convert_to_hf.py
@@ -5,6 +5,7 @@
import slime.backends.megatron_utils as megatron_utils
from slime.backends.megatron_utils import update_weight_utils
+from slime.utils import accelerator
from slime.utils.arguments import parse_args
@@ -57,7 +58,7 @@ def main(args):
param = param_
break
else:
- param = torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device())
+ param = torch.empty(info.shape, dtype=info.dtype, device=accelerator.current_device())
if pp_size > 1:
if info.src_rank in dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()):
diff --git a/tools/fp8_cast_bf16.py b/tools/fp8_cast_bf16.py
index c227c300f2..edc8485d90 100644
--- a/tools/fp8_cast_bf16.py
+++ b/tools/fp8_cast_bf16.py
@@ -10,6 +10,8 @@
from safetensors.torch import load_file, save_file
from tqdm import tqdm
+from slime.utils import accelerator
+
@triton.jit
def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
@@ -60,7 +62,7 @@ def get_tensor(tensor_name):
file_name = weight_map[tensor_name]
if file_name not in loaded_files:
file_path = os.path.join(fp8_path, file_name)
- loaded_files[file_name] = load_file(file_path, device="cuda")
+ loaded_files[file_name] = load_file(file_path, device=accelerator.device_name())
return loaded_files[file_name][tensor_name]
safetensor_files = list(glob(os.path.join(fp8_path, "*.safetensors")))
@@ -68,7 +70,7 @@ def get_tensor(tensor_name):
for safetensor_file in tqdm(safetensor_files):
print(f"Handling file: {safetensor_file}")
file_name = os.path.basename(safetensor_file)
- current_state_dict = load_file(safetensor_file, device="cuda")
+ current_state_dict = load_file(safetensor_file, device=accelerator.device_name())
loaded_files[file_name] = current_state_dict
new_state_dict = {}
@@ -95,7 +97,7 @@ def get_tensor(tensor_name):
if len(loaded_files) > 2:
oldest_file = next(iter(loaded_files))
del loaded_files[oldest_file]
- torch.cuda.empty_cache()
+ accelerator.empty_cache()
# Update model index
new_model_index_file = os.path.join(bf16_path, "model.safetensors.index.json")
diff --git a/train.py b/train.py
index 9cb2968866..051e5ec35d 100644
--- a/train.py
+++ b/train.py
@@ -1,8 +1,8 @@
import ray
+from slime.observability.logging_utils import configure_logger, finish_tracking, init_tracking
from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
from slime.utils.arguments import parse_args
-from slime.utils.logging_utils import configure_logger, finish_tracking, init_tracking
from slime.utils.misc import should_run_periodic_action
diff --git a/train_async.py b/train_async.py
index 9141b612fe..1f74a9a5a2 100644
--- a/train_async.py
+++ b/train_async.py
@@ -1,8 +1,8 @@
import ray
+from slime.observability.logging_utils import configure_logger, finish_tracking, init_tracking
from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
from slime.utils.arguments import parse_args
-from slime.utils.logging_utils import configure_logger, finish_tracking, init_tracking
from slime.utils.misc import should_run_periodic_action