Skip to content

FortPercent/Slime_R3

Repository files navigation

R3 / FAST_R3

Routing replay has two modes:

  • Inline R3: default production path. SGLang returns routed_experts in the rollout response and slime consumes it directly.
  • FAST_R3: optional Mooncake sidecar path. SGLang writes routed expert tensors to Mooncake, returns None for committed samples, and slime fetches missing samples from Mooncake before replay.

The current rootfix makes FAST_R3 tensor-native and rank-local:

  • Mooncake tensor API: batch_put_tensor(keys, tensors) / batch_get_tensor(keys).
  • Routing expert ids are stored as uint8 payload bytes and cast back to int32 on read. The writer range-checks [0, 255]; overflow keeps inline payload instead of truncating.
  • SGLang commits finished requests in one scheduler-iteration batch before stream_output(). Inline payload is cleared only for samples whose Mooncake commit succeeds.
  • CPU hot paths use NumPy-backed sample hashing, pre-encoded key fields, and single-pass reader chunk grouping. These keep the chunk protocol unchanged and only reduce Python-side overhead.
  • Chunk blob protocol: tensor chunks with validated R3CV or layer-sharded R3LS headers.
  • Hex string keys derive from namespace, writer request id, sample tokens, token chunk range, and layer shard range. Token-only keys are rejected.
  • Rank-local train fetch: each training rank fetches only its local sample shard and local pipeline layer range.
  • TP broadcast fetch: with SLIME_R3_DISTRIBUTED_FETCH_MODE=tp_broadcast, one tensor-parallel leader fetches Mooncake payloads and broadcasts uint8 tensors inside the TP group.
  • Layer shards: SLIME_R3_LAYER_RANGES=0:7,7:14,14:21,21:27 for the Moonlight 27-layer, PP=4 setup.
  • Mooncake local buffer rootfix: this SGLang checkout passes DEFAULT_LOCAL_BUFFER_SIZE directly to MooncakeStore.setup(), so MOONCAKE_LOCAL_BUFFER_SIZE is ignored unless R3 patches that constant before store construction.
  • Mooncake bind-address rootfix: R3 validates MOONCAKE_LOCAL_HOSTNAME by binding it locally and derives a peer-routable source IP when the platform provides a non-bindable VIP.
  • Fail-fast correctness checks for missing chunks, bad headers, bad sizes, and missing routing after fill.

Quickstart

For the shortest command sequence, use QUICKSTART.md. The steps below include the full setup, run, and validation flow.

1. Sync Artifacts

Copy the current artifact tree to the run environment. The wrappers expect R3_DIR to point at this directory.

export R3_DIR=/gfs/space/chatrl/users/wxe/agl_fully_async/r3_artifacts
mkdir -p "${R3_DIR}"
rsync -a --delete patched_files/ "${R3_DIR}/"

Apply the integration edits to the Slime/SGLang checkout. The patcher is idempotent and should be run before launching a validation job.

cd "${R3_DIR}"
SLIME_ROOT=/gfs/space/chatrl/users/wxe/slime_per_sample \
SGLANG_ROOT=/sgl-workspace/sglang \
python3 apply_fast_r3_edits.py

2. Run Local Artifact Checks

cd "${R3_DIR}"
python3 -m py_compile routing_replay.py sglang_r3_writer.py apply_fast_r3_edits.py \
  export_r3_log_metrics_to_tb.py plot_r3_train_infer_deviation.py \
  plot_tb_train_infer_deviation.py profile_fa_r3_log.py \
  test_r3_chunk_protocol_static.py test_r3_chunk_tensor_fake_store.py
python3 test_r3_chunk_protocol_static.py
python3 test_r3_chunk_tensor_fake_store.py

test_r3_chunk_tensor_fake_store.py needs torch. If the local laptop does not have torch, run it on the training image.

3. Choose A Run Wrapper

Fully-async 4-GPU smoke:

RUN_NAME=fa_r3_smoke_$(date +%Y%m%d_%H%M%S) \
USE_TENSORBOARD=1 \
SLIME_R3_CHECKSUM=1 \
bash "${R3_DIR}/run_agl_fully_async_r3_moonlight_4gpu_smoke.sh"

Fully-async 2-node profile/regression:

RUN_NAME=fa_r3_2node_$(date +%Y%m%d_%H%M%S) \
NUM_ROLLOUT=4 \
ROLLOUT_BATCH_SIZE=4 \
GLOBAL_BATCH_SIZE=4 \
USE_TENSORBOARD=1 \
SLIME_R3_CHECKSUM=1 \
bash "${R3_DIR}/run_agl_fully_async_r3_moonlight_2node_profile.sh"

GLM-4.7 SWE-bench 2-node colocate:

RUN_NAME=glm47_colocate_r3_$(date +%Y%m%d_%H%M%S) \
SLIME_R3_CHECKSUM=1 \
bash "${R3_DIR}/run_agl_glm47_swebench_2node_colocate_r3.sh"

The wrappers set the normal FAST_R3 environment. Override only when you need to change topology:

export FAST_R3=1
export SLIME_R3_DISTRIBUTED_FETCH_MODE=tp_broadcast
export SLIME_R3_LAYER_RANGES=0:7,7:14,14:21,21:27
export SLIME_FAST_ROUTING_REPLAY=1
export SLIME_R3_WRITE_MOONCAKE=1
export MOONCAKE_LOCAL_BUFFER_SIZE=1073741824

Use rank_fetch only when TP broadcast is not applicable. Do not use the old rank0 scatter path; it has been removed from the supported path.

Keep inline R3 as default by leaving FAST_R3=0 or unset.

4. Validate A Completed Run

Set the run directory and validate payload-level correctness:

RUN_DIR=/gfs/space/chatrl/users/wxe/agl_runs/<run-name>
LOG_FILE="${RUN_DIR}/node0.log" SLIME_R3_CHECKSUM=1 \
  bash "${R3_DIR}/validate_fa_r3_agent_smoke.sh" "${RUN_DIR}"
bash "${R3_DIR}/validate_r3_run.sh" "${RUN_DIR}/node0.log"
python3 "${R3_DIR}/profile_fa_r3_log.py" "${RUN_DIR}/node0.log" "${RUN_DIR}/node1.log"

The checksum pass condition is:

reader_without_writer=0
duplicate_reads=0
writer_conflict_keys=0
matched_unique > 0

This proves every reader payload checksum has a matching writer payload checksum for the same key, chunk, token count, layer range, byte count, and crc32.

5. TensorBoard And Deviation Plots

For runs with TensorBoard events:

python3 "${R3_DIR}/plot_tb_train_infer_deviation.py" "${RUN_DIR}" \
  --out "${RUN_DIR}/tb_train_infer_deviation.png" \
  --csv "${RUN_DIR}/tb_train_infer_deviation.csv"

For log-only runs:

python3 "${R3_DIR}/plot_r3_train_infer_deviation.py" \
  "${RUN_DIR}/node0.log" \
  --labels "$(basename "${RUN_DIR}")" \
  --out "${RUN_DIR}/train_infer_deviation.png" \
  --csv "${RUN_DIR}/train_infer_deviation.csv"

Maintenance Surfaces

R3 itself should be maintained as the canonical design and artifact repo, but it has three coupled integration surfaces that must stay versioned together:

  • Slime: routing_replay.py, actor-side replay fill, rollout batch payload preservation, and train/rollout timing instrumentation.
  • SGLang: routed-experts capture, tensor-native Mooncake writer, scheduler batch flush hook, and writer profiling.
  • AGL / launchers: two-pod wrappers, runtime env, Mooncake master/metadata startup, pod role mapping, and reproducible experiment commands.

Do not treat a FAST_R3 result as reproducible unless the report records the R3 artifact revision plus the Slime, SGLang, and AGL commits/branches used for that run.

Validation Signals

FAST_R3 is healthy when logs show:

[r3-writer] mooncake tensor API roundtrip OK
[r3-writer] init OK: tensor-native Mooncake store ready
[r3-writer] override MOONCAKE_LOCAL_HOSTNAME: ... (src=bindable-probe)
[r3-writer/BATCH] batches=... samples=... chunks=... tensor_bytes=...
[r3-writer/CHECK] key=... chunk=... layers=... bytes=... crc32=...
[R3-fast] filling ... samples from Mooncake ... layers=x:y/27
[R3-fast/distributed] rank ... tp-broadcast leader prefetching ... unique seqs
[R3-fast/distributed] tp-broadcast tensor payload: unique=... bytes=... wire_dtype=uint8
[R3-fast] chunk keys ... layers=x:y/27, stored_shards=[...]
[R3-fast/STATS] fetched_samples=... payload_bytes=...
[R3-fast/CHECK] key=... chunk=... layers=... bytes=... crc32=...

It is not healthy if any of these appear:

FAST_R3 missing ... chunks
R3 routing data missing
bad R3 chunk header
R3 chunk size mismatch
reader checksum without writer
writer emitted multiple payload crc32 values for identical R3 key
Failed to allocate buffer
Traceback
RuntimeError
AssertionError

When FAST_R3 Pays Off

FAST_R3 is useful when inline routing payload becomes a scaling pressure:

  • multiple rollout batches resident at once, such as rollout pipeline overlap;
  • multiple consumers reading the same routing data, such as actor + ref + RM;
  • constrained Ray plasma or large concurrent rollout buffers;
  • larger payloads from more samples, longer responses, more MoE layers, higher top-k, or larger cluster scale.

For the current single-batch, single-consumer, large-plasma SWE-bench smoke shape, inline R3 remains the lower-latency path. FAST_R3 is the optional path for the scaling cases above.

Measured on the current two-pod Moonlight SWE-bench shape:

  • Old FAST_R3 data preprocess was 70.35s-74.39s by perf metric.
  • Current FAST_R3 data preprocess is 0.1878s-0.2286s by perf metric.
  • Best current FAST_R3 step time is 706.51s; latest batch/CPU run is 713.42s.
  • Inline step time is 698.85s in the same validation setup.
  • Root cause: R3 fetch is now under 0.5s; the single-batch step is dominated by train wait and train compute, so this shape does not expose an E2E FAST_R3 win over inline.

Remote SWE-bench Smoke

Head:

ssh -p <port> <head-ssh-user>@<ssh-proxy-host>

Worker:

ssh -p <port> <worker-ssh-user>@<ssh-proxy-host>

Main wrapper:

/gfs/space/chatrl/users/wxe/agl_fully_async/run_gemini_r3_moonlight.sh

Status and measured results are tracked in SWEbench_FAST_R3_status.md.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors