From 6b8e3563777ad510e8683e2f6c18ce5ba38ae3d2 Mon Sep 17 00:00:00 2001 From: kaixih Date: Thu, 9 Jul 2026 05:55:51 +0000 Subject: [PATCH 1/8] [OPD] Stabilize two-node GB200 HybridEP workflow --- .../qwen3_5_35b_selfdistill/phase2_gb200.sh | 46 +++- miles/backends/megatron_utils/model.py | 8 + miles/backends/training_utils/data.py | 22 +- miles/backends/training_utils/replay_data.py | 19 +- miles/ray/actor_group.py | 17 +- workflow_gb200.md | 198 ++++++++++++++++++ 6 files changed, 297 insertions(+), 13 deletions(-) create mode 100644 workflow_gb200.md diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh index 512abe575c3..94c6350fcb2 100755 --- a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh @@ -42,6 +42,18 @@ EXAMPLE_DIR=${EXAMPLE_DIR:-$(cd "$(dirname "$0")" && pwd)} MILES_DIR=${MILES_DIR:-/workspace/miles} RAY_ADDRESS=${RAY_ADDRESS:-http://127.0.0.1:8265} OPD_KL_COEF=${OPD_KL_COEF:-0.2} +MOE_FLEX_DISPATCHER_BACKEND=${MOE_FLEX_DISPATCHER_BACKEND:-hybridep} +NUM_ROLLOUT=${NUM_ROLLOUT:-12} +ROLLOUT_BATCH_SIZE=${ROLLOUT_BATCH_SIZE:-32} +N_SAMPLES_PER_PROMPT=${N_SAMPLES_PER_PROMPT:-8} +OVER_SAMPLING_BATCH_SIZE=${OVER_SAMPLING_BATCH_SIZE:-32} +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-256} +ROLLOUT_MAX_RESPONSE_LEN=${ROLLOUT_MAX_RESPONSE_LEN:-24576} +MAX_TOKENS_PER_GPU=${MAX_TOKENS_PER_GPU:-16384} +SAVE_INTERVAL=${SAVE_INTERVAL:-5} +EVAL_INTERVAL=${EVAL_INTERVAL:-5} +SKIP_EVAL_BEFORE_TRAIN=${SKIP_EVAL_BEFORE_TRAIN:-0} +read -r -a SGLANG_CUDA_GRAPH_BS_ARRAY <<< "${SGLANG_CUDA_GRAPH_BS:-1 2 4}" mkdir -p "${OUTPUT_DIR}" EVAL_CONFIG="${OUTPUT_DIR}/eval_dapo_heldout.yaml" @@ -63,6 +75,7 @@ MODEL_ARGS=( --untie-embeddings-and-output-weights --vocab-size 248320 --rotary-base 10000000 --moe-ffn-hidden-size 512 --moe-shared-expert-intermediate-size 512 --moe-router-score-function softmax --moe-token-dispatcher-type flex + --moe-flex-dispatcher-backend ${MOE_FLEX_DISPATCHER_BACKEND} --moe-router-topk 8 --moe-layer-freq "[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]" --num-experts 256 --moe-grouped-gemm --moe-token-drop-policy probs --moe-router-dtype fp32 @@ -71,21 +84,29 @@ MODEL_ARGS=( ) CKPT_ARGS=( --hf-checkpoint ${MODEL_DIR}/Qwen3.5-35B-A3B - --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist - --load ${OUTPUT_DIR} --save ${OUTPUT_DIR} --save-interval 5 + --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist --load ${OUTPUT_DIR} ) +if [ "${SAVE_INTERVAL}" != "none" ]; then + CKPT_ARGS+=(--save ${OUTPUT_DIR} --save-interval ${SAVE_INTERVAL}) +fi OPD_ARGS=( --use-opd --opd-type megatron --opd-teacher-load ${TEACHER_LOAD} --opd-kl-coef ${OPD_KL_COEF} ) ROLLOUT_ARGS=( --prompt-data ${DATA_DIR}/dapo_train.jsonl --input-key prompt --label-key label --apply-chat-template --rollout-shuffle - --num-rollout 12 --rollout-batch-size 32 --n-samples-per-prompt 8 - --rollout-max-response-len 24576 --rollout-temperature 1 --num-steps-per-rollout 1 - --over-sampling-batch-size 32 --global-batch-size 256 --balance-data + --num-rollout ${NUM_ROLLOUT} --rollout-batch-size ${ROLLOUT_BATCH_SIZE} + --n-samples-per-prompt ${N_SAMPLES_PER_PROMPT} + --rollout-max-response-len ${ROLLOUT_MAX_RESPONSE_LEN} + --rollout-temperature 1 --num-steps-per-rollout 1 + --over-sampling-batch-size ${OVER_SAMPLING_BATCH_SIZE} + --global-batch-size ${GLOBAL_BATCH_SIZE} --balance-data ) RM_ARGS=( --custom-rm-path ${RM_FUNC} ) -EVAL_ARGS=( --eval-interval 5 --eval-config ${EVAL_CONFIG} ) +EVAL_ARGS=( --eval-interval ${EVAL_INTERVAL} --eval-config ${EVAL_CONFIG} ) +if [ "${SKIP_EVAL_BEFORE_TRAIN}" = "1" ]; then + EVAL_ARGS+=(--skip-eval-before-train) +fi GRPO_ARGS=( --advantage-estimator grpo --kl-loss-type low_var_kl --entropy-coef 0.00 --eps-clip 0.2 --eps-clip-high 0.28 --use-tis @@ -99,7 +120,7 @@ PERF_ARGS=( --tensor-model-parallel-size 2 --sequence-parallel --pipeline-model-parallel-size 1 --context-parallel-size 2 --expert-model-parallel-size 8 --expert-tensor-parallel-size 1 --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 - --use-dynamic-batch-size --max-tokens-per-gpu 16384 --log-probs-chunk-size 4096 + --use-dynamic-batch-size --max-tokens-per-gpu ${MAX_TOKENS_PER_GPU} --log-probs-chunk-size 4096 ) # Blackwell (GB200 sm100 / B300 sm103) backends, per scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py: # - moe-runner-backend flashinfer_cutlass: the default triton fused-MoE mis-shards @@ -109,7 +130,7 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine 8 --sglang-mem-fraction-static 0.7 --sglang-ep-size 8 --sglang-watchdog-timeout 1800 --sglang-enable-metrics --sglang-moe-runner-backend flashinfer_cutlass --sglang-attention-backend trtllm_mha - --sglang-cuda-graph-bs 1 2 4 8 16 32 --use-rollout-routing-replay + --sglang-cuda-graph-bs "${SGLANG_CUDA_GRAPH_BS_ARRAY[@]}" --use-rollout-routing-replay --sglang-mamba-scheduler-strategy extra_buffer ) MISC_ARGS=( @@ -126,7 +147,14 @@ WANDB_ARGS=( --use-wandb --wandb-project miles-opd --wandb-group qwen3.5-35b-opd # during cross-node ncclCommInitRank ("unhandled cuda error" / CUDA 999), which # kills the EP8 sglang engine spanning both nodes. MNNVL stays enabled (pod env) # for the NVLink fabric itself; only NVLS multicast is disabled. -RUNTIME_ENV_JSON="{\"env_vars\": {\"PYTHONPATH\": \"${MILES_DIR}:/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"WANDB_API_KEY\": \"${WANDB_API_KEY}\", \"PROMETHEUS_PORT\": \"9090\", \"NCCL_NVLS_ENABLE\": \"0\"}}" +TRACE_RUNTIME_ENV="" +if [ -n "${TMS_CUMEM_TRACE_LIB:-}" ]; then + TRACE_RUNTIME_ENV=", \"TMS_CUMEM_TRACE_LIB\": \"${TMS_CUMEM_TRACE_LIB}\", \"CUMEM_TRACE_ALL\": \"${CUMEM_TRACE_ALL:-}\"" +fi +if [ -n "${TORCH_SHM_UNLINK_COMPAT_LIB:-}" ]; then + TRACE_RUNTIME_ENV="${TRACE_RUNTIME_ENV}, \"TORCH_SHM_UNLINK_COMPAT_LIB\": \"${TORCH_SHM_UNLINK_COMPAT_LIB}\"" +fi +RUNTIME_ENV_JSON="{\"env_vars\": {\"PYTHONPATH\": \"${MILES_DIR}:/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"WANDB_API_KEY\": \"${WANDB_API_KEY}\", \"PROMETHEUS_PORT\": \"9090\", \"NCCL_NVLS_ENABLE\": \"0\"${TRACE_RUNTIME_ENV}}}" cd "${MILES_DIR}" ray job submit --address="${RAY_ADDRESS}" --submission-id qwen3.5-opd-${MODE} --no-wait \ diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 44fdaeedcd8..503b14f657d 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -269,6 +269,10 @@ def forward_step( args.data_pad_size_multiplier, args.qkv_format, allgather_cp=args.allgather_cp, + pad_to_ep_max=( + getattr(args, "moe_token_dispatcher_type", None) == "flex" + and getattr(args, "moe_flex_dispatcher_backend", None) == "hybridep" + ), ) unconcat_tokens = batch["unconcat_tokens"] tokens = batch["tokens"] @@ -421,6 +425,10 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p args.data_pad_size_multiplier, args.qkv_format, allgather_cp=args.allgather_cp, + pad_to_ep_max=( + getattr(args, "moe_token_dispatcher_type", None) == "flex" + and getattr(args, "moe_flex_dispatcher_backend", None) == "hybridep" + ), ) from miles.utils.replay_base import all_replay_managers diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index a43e0aecfbf..693b9a436b9 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -107,6 +107,7 @@ def get_batch( qkv_format: str = "thd", get_position_ids: bool = False, allgather_cp: bool = False, + pad_to_ep_max: bool = False, ) -> dict[str, torch.Tensor | list[torch.Tensor] | None]: """ Generate a CP-ready micro-batch with packed sequence parameters. @@ -121,6 +122,8 @@ def get_batch( data_iterator: Iterator providing micro-batch data. keys: List of keys to fetch from the iterator. pad_multiplier: Multiplier for padding size calculation (default: 128). + pad_to_ep_max: Pad THD token rows to the largest aligned row count in + the expert-parallel group. Required by HybridEP for variable batches. Returns a dict including: - "tokens": torch.LongTensor of shape [1, T_padded] on the current CUDA device @@ -195,8 +198,23 @@ def get_batch( tokens = torch.cat(tokens) - # Always pad to reduce memory fragmentation and maybe make the computation faster - pad = (pad_size - tokens.size(0) % pad_size) % pad_size + # HybridEP assumes every rank contributes the same number of token rows + # to metadata all-gather. Dynamic batches violate that assumption unless + # the local packed stream is padded to an EP-wide maximum. + local_target = (tokens.size(0) + pad_size - 1) // pad_size * pad_size + if pad_to_ep_max: + assert parallel_state.ep.group is not None + target = torch.tensor( + [local_target], + dtype=torch.int64, + device=tokens.device, + ) + dist.all_reduce(target, op=dist.ReduceOp.MAX, group=parallel_state.ep.group) + target_num_tokens = int(target.item()) + else: + target_num_tokens = local_target + + pad = target_num_tokens - tokens.size(0) if pad != 0: tokens = F.pad(tokens, (0, pad), value=pad_token_id) cu_seqlens.append(cu_seqlens[-1] + pad) diff --git a/miles/backends/training_utils/replay_data.py b/miles/backends/training_utils/replay_data.py index 1a35551ebf6..fc8a13e9c57 100644 --- a/miles/backends/training_utils/replay_data.py +++ b/miles/backends/training_utils/replay_data.py @@ -1,6 +1,7 @@ from typing import Protocol import torch +import torch.distributed as dist from .cp_utils import slice_with_cp from .parallel import get_parallel_state @@ -115,7 +116,23 @@ def pad_func(data, pad): replay_data[i] = torch.where(r != -1, r + offset, r) offset += r.shape[0] replay_data = torch.cat(replay_data, dim=0) - pad = (pad_size - replay_data.size(0) % pad_size) % pad_size + local_target = (replay_data.size(0) + pad_size - 1) // pad_size * pad_size + if ( + getattr(args, "moe_token_dispatcher_type", None) == "flex" + and getattr(args, "moe_flex_dispatcher_backend", None) == "hybridep" + ): + assert parallel_state.ep.group is not None + target = torch.tensor( + [local_target], + dtype=torch.int64, + device=torch.cuda.current_device(), + ) + dist.all_reduce(target, op=dist.ReduceOp.MAX, group=parallel_state.ep.group) + target_num_tokens = int(target.item()) + else: + target_num_tokens = local_target + + pad = target_num_tokens - replay_data.size(0) if pad != 0: replay_data = pad_func(replay_data, pad) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index f729d948aed..4766766d523 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -74,7 +74,22 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): ) assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." - env_vars["LD_PRELOAD"] = dynlib_path + preload_libs = [] + shm_compat_lib = os.environ.get("TORCH_SHM_UNLINK_COMPAT_LIB") + if shm_compat_lib: + assert os.path.exists( + shm_compat_lib + ), f"torch shared-memory cleanup library {shm_compat_lib} does not exist." + preload_libs.append(shm_compat_lib) + + trace_lib = os.environ.get("TMS_CUMEM_TRACE_LIB") + if trace_lib: + assert os.path.exists(trace_lib), f"cuMemCreate trace library {trace_lib} does not exist." + preload_libs.append(trace_lib) + env_vars["CUMEM_TRACE_ALL"] = os.environ.get("CUMEM_TRACE_ALL", "") + + preload_libs.append(dynlib_path) + env_vars["LD_PRELOAD"] = ":".join(preload_libs) env_vars["TMS_INIT_ENABLE"] = "1" env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" diff --git a/workflow_gb200.md b/workflow_gb200.md new file mode 100644 index 00000000000..eac89f66787 --- /dev/null +++ b/workflow_gb200.md @@ -0,0 +1,198 @@ +# Two-Node GB200 Pure-OPD Workflow + +This document records the working two-node GB200 workflow for the Qwen3.5-35B-A3B pure-OPD experiment. It describes the exact checkout, mounts, runtime shims, launch sequence, and the reason for each local change. + +## 1. Checkout and Branch + +Use the dedicated checkout and branch: + +```bash +cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main +git switch opd-gb200-main +``` + +`opd-gb200-main` contains the four commits from PR #1488 cherry-picked onto the current `origin/main`, followed by the GB200 fixes documented below. Do not confuse it with `/home/scratch.kaixih_ent/repo/miles`, which is a different checkout. + +The local `lab/` directory is intentionally not part of the commit. It contains experiment launchers, temporary repros, and credentials used only in the DLCluster workspace. + +## 2. Storage and Container Mounts + +The same shared storage has three path names: + +```text +DLCluster login: /home/scratch.kaixih_ent +GB compute node: /mnt/cifs/home/scratch.kaixih_ent +Miles container: /scratch +``` + +Launch the ARM Miles image with these bind mounts: + +```text +/mnt/cifs/home/scratch.kaixih_ent/repo/miles-opd-gb200-main -> /workspace/miles +/mnt/cifs/home/scratch.kaixih_ent -> /scratch +``` + +The launcher uses: + +```bash +MILES_DIR=/workspace/miles +PYTHONPATH=/workspace/miles:/root/Megatron-LM/ +``` + +Therefore Python imports Miles from the mounted checkout, including its committed changes. Megatron, PyTorch, torch-memory-saver, and SGLang still come from `radixark/miles:latest`; their installed source files are not edited. + +## 3. Build the Runtime Compatibility Shims + +The temporary shim sources are local lab artifacts and are intentionally not committed: + +```text +lab/opd_gb200/tms_cumem_trace.c +lab/opd_gb200/torch_shm_unlink_compat.c +``` + +Compile them on an ARM GB200 compute node, inside the same Miles image. Do not compile them on the x86 login node. + +From an active two-node allocation shell: + +```bash +cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main +export ALLOC_JOB_ID="${SLURM_JOB_ID}" +bash lab/opd_gb200/00_build_tms_compat.sh +``` + +The build script runs the equivalent of: + +```bash +mkdir -p /scratch/repro/opd-self-distill/runtime/aarch64 + +cuda_header=$(find /usr/local/cuda* -path '*/targets/*/include/cuda.h' -print -quit) +cuda_include=$(dirname "${cuda_header}") + +cc -shared -fPIC -O2 -I"${cuda_include}" \ + /workspace/miles/lab/opd_gb200/tms_cumem_trace.c \ + -o /scratch/repro/opd-self-distill/runtime/aarch64/libtms_cumem_compat.so \ + -ldl + +cc -shared -fPIC -O2 \ + /workspace/miles/lab/opd_gb200/torch_shm_unlink_compat.c \ + -o /scratch/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so \ + -ldl +``` + +The persistent host equivalents are: + +```text +/home/scratch.kaixih_ent/repro/opd-self-distill/runtime/aarch64/libtms_cumem_compat.so +/home/scratch.kaixih_ent/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so +``` + +Verify both outputs before launching: + +```bash +file /home/scratch.kaixih_ent/repro/opd-self-distill/runtime/aarch64/*.so +``` + +They must be AArch64 shared objects. `phase2_gb200.sh` passes their container paths into the Ray runtime environment, and `miles/ray/actor_group.py` prepends them to `LD_PRELOAD` before the original torch-memory-saver library. + +## 4. Launch the Two-Node Run + +The allocation must contain two `gb200nvl72` nodes with four GPUs per node. The local launcher starts one container per node, creates a Ray head on task 0, joins task 1 as a worker, and submits the Phase-2 payload from `/workspace/miles`. + +Run the current accuracy-focused recipe: + +```bash +cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main + +RUN_ID="opd-gb200-$(date +%m%d-%H%M)" \ +NUM_ROLLOUT=12 \ +SKIP_EVAL_BEFORE_TRAIN=1 \ +EVAL_INTERVAL=2 \ +bash lab/opd_gb200/04_run_long_8k_hybridep.sh +``` + +The effective settings are: + +```text +2 nodes x 4 GB200 GPUs +TP=2, CP=2, EP=8 +Flex dispatcher with HybridEP +global batch size=256 +max tokens per GPU=8192 +CUDA graph capture sizes=1 2 4 +checkpoint saving disabled +``` + +CUDA graph sizes larger than 4 currently use eager decode. This is a conservative mitigation for a separate replay hang after online weight handoff; it is not a final SGLang fix. + +## 5. What Changed and Why + +### Miles model-input alignment + +Files: + +```text +miles/backends/megatron_utils/model.py +miles/backends/training_utils/data.py +``` + +Dynamic THD batches can produce different token-row counts on different EP ranks. HybridEP expects rank-consistent metadata and buffer capacity. A mismatch caused `HYBRID-EP ALLGATHER TIMEOUT` and corrupted updates with extremely large gradient norms and importance-sampling ratios. + +For Flex + HybridEP only, Miles now all-reduces the maximum aligned token-row count across the EP group and pads every rank's model input to that common target. + +### Miles routing-replay alignment + +File: + +```text +miles/backends/training_utils/replay_data.py +``` + +Model inputs and routing replay must use the same EP-wide row target. Local-only replay padding produced errors such as `replay n_tokens ... does not match scores n_tokens ...` under TP2 plus sequence parallelism. Replay rows are now padded to the same EP-wide target before TP/SP slicing; padding remains invalid with value `-1` and stays outside the loss. + +### GB200 recipe and runtime environment + +File: + +```text +examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh +``` + +The script now selects HybridEP explicitly and makes rollout count, batch sizes, response length, token budget, evaluation cadence, checkpoint cadence, and SGLang CUDA graph sizes configurable. It also forwards the two shim paths into Ray actors. + +### Ray actor preload composition + +File: + +```text +miles/ray/actor_group.py +``` + +The actor environment composes this preload order: + +```text +libtorch_shm_unlink_compat.so +libtms_cumem_compat.so +the original torch-memory-saver preload library +``` + +`libtms_cumem_compat.so` bypasses torch-memory-saver only for small allocations that are not aligned to the GB200 CUDA VMM granularity, routing them to native `cudaMalloc`. + +`libtorch_shm_unlink_compat.so` makes cleanup idempotent when PyTorch calls `shm_unlink()` for an already-removed `/torch_*` CUDA IPC refcount file. Only `ENOENT` for `/torch_*` is converted to success. + +These two shims are dependency workarounds. The permanent fixes belong in torch-memory-saver/container packaging and PyTorch respectively, not in OPD mathematics. + +## 6. Validation Signals + +A healthy run must have all of the following: + +- no `HYBRID-EP ALLGATHER TIMEOUT`; +- no routing replay row mismatch; +- no fatal CUDA IPC cleanup error; +- `OIS`, `ESS`, and `TIS` remain near 1; +- finite, ordinary-scale gradient norms; +- successful online weight synchronization and the next rollout; +- response length and reverse-KL move toward the teacher behavior. + +The validated run changed mean response length from about 18.8k to 6.6k tokens after the first update, reduced truncation from 0.367 to 0.043, and reduced per-token OPD reverse-KL from 0.0448 to 0.0130 while keeping the training diagnostics healthy. + +No Megatron, PyTorch, torch-memory-saver, or SGLang installed source file is modified by this workflow. The active code consists of the mounted Miles checkout, the two external runtime shims, and recipe flags. From ffb7f320220f8b8f1a90613344f351ffe8b09848 Mon Sep 17 00:00:00 2001 From: kaixih Date: Sat, 11 Jul 2026 22:09:20 +0000 Subject: [PATCH 2/8] [OPD] Fix stacked branch pre-commit --- .../qwen3_5_35b_selfdistill/make_split.py | 6 +++--- .../on_policy_distillation/qwen3_5_35b_selfdistill/rm.py | 2 +- miles_plugins/mbridge/qwen3_5.py | 4 +--- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py index d15366d2643..8a91a271486 100644 --- a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py @@ -42,10 +42,10 @@ def main(): rows.append(line) random.Random(args.seed).shuffle(rows) # fixed seed: reproducible, unbiased split - eval_rows, train_rows = rows[-args.eval_n:], rows[:-args.eval_n] + eval_rows, train_rows = rows[-args.eval_n :], rows[: -args.eval_n] - ek = {prompt_text(json.loads(l)) for l in eval_rows} - tk = {prompt_text(json.loads(l)) for l in train_rows} + ek = {prompt_text(json.loads(line)) for line in eval_rows} + tk = {prompt_text(json.loads(line)) for line in train_rows} assert ek.isdisjoint(tk), "LEAK: eval prompt found in train split" with open(os.path.join(args.out_dir, "dapo_train.jsonl"), "w") as f: diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py index 9a68f09b1fe..88b8cbd572e 100644 --- a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py @@ -28,8 +28,8 @@ ``--custom-rm-path examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func``. """ -from miles.rollout.rm_hub.math_utils import grade_answer_verl from miles.rollout.rm_hub.math_dapo_utils import compute_score as _dapo_score +from miles.rollout.rm_hub.math_utils import grade_answer_verl def _is_eval_sample(sample) -> bool: diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index 9f467eedf0b..67d1dcb35f7 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -211,9 +211,7 @@ def _experts_fused(self) -> bool: index = getattr(io, "index", None) if io is not None else None if not index: return True - fused = any( - "model.language_model.layers." in k and k.endswith("mlp.experts.gate_up_proj") for k in index - ) + fused = any("model.language_model.layers." in k and k.endswith("mlp.experts.gate_up_proj") for k in index) self._experts_fused_cached = fused return fused From 87b48d5422198573162ebb0de898233f9e121e81 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Mon, 13 Jul 2026 14:35:56 -0700 Subject: [PATCH 3/8] Document validated GB200 CUDA graph workflow --- workflow_gb200.md | 290 +++++++++++++++++++++------------------------- 1 file changed, 135 insertions(+), 155 deletions(-) diff --git a/workflow_gb200.md b/workflow_gb200.md index eac89f66787..6c046308f6a 100644 --- a/workflow_gb200.md +++ b/workflow_gb200.md @@ -1,58 +1,64 @@ # Two-Node GB200 Pure-OPD Workflow -This document records the working two-node GB200 workflow for the Qwen3.5-35B-A3B pure-OPD experiment. It describes the exact checkout, mounts, runtime shims, launch sequence, and the reason for each local change. +Validated on 2026-07-13 with two independent two-node runs. Each run completed +two real rollouts and two online training updates with overlap scheduling, +CUDA graph padding, and power-of-two capture enabled. -## 1. Checkout and Branch +## Required Stack -Use the dedicated checkout and branch: +Use these pieces together: -```bash -cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main -git switch opd-gb200-main -``` +1. Miles PR #1634 from branch `agent/gb200-hybridep-fixes`. +2. The narrow PyTorch `shm_unlink(ENOENT)` compatibility shim. +3. A real installation of torch-memory-saver PR #82 at commit + `c96bf60e093b4bec2b045cb5b8a08601d0ae8a79`. +4. SGLang PR #27140 behavior: invalidate stale CUDA graphs and recapture them + after torch-memory-saver resumes graph-owned memory. +5. Synchronize asynchronous replay before torch-memory-saver unmaps graph + memory. +6. SGLang PR #31072: for hybrid linear-attention models, publish the + overlap-scheduler read-done event after CUDA graph replay, not before it. +7. Capture exactly `1 2 4 8 16 32 64 128 256 512`, with graph padding and + overlap scheduling enabled. -`opd-gb200-main` contains the four commits from PR #1488 cherry-picked onto the current `origin/main`, followed by the GB200 fixes documented below. Do not confuse it with `/home/scratch.kaixih_ent/repo/miles`, which is a different checkout. +SGLang PRs #30895 and #30974 are not required for this hang. The second +validation completed with both disabled. They may remain useful as padded-row +hygiene, but they are not part of the minimum validated stack. -The local `lab/` directory is intentionally not part of the commit. It contains experiment launchers, temporary repros, and credentials used only in the DLCluster workspace. +Do not use the historical `libtms_cumem_compat.so` shim. Keep +`TMS_CUMEM_TRACE_LIB` empty. The local `lab/` directory contains launchers and +temporary patches and is intentionally not committed. -## 2. Storage and Container Mounts +## 1. Get Two Nodes -The same shared storage has three path names: +Run this from the visible tmux shell on `dl3`: -```text -DLCluster login: /home/scratch.kaixih_ent -GB compute node: /mnt/cifs/home/scratch.kaixih_ent -Miles container: /scratch +```bash +salloc_node gb200nvl72 2 ``` -Launch the ARM Miles image with these bind mounts: - -```text -/mnt/cifs/home/scratch.kaixih_ent/repo/miles-opd-gb200-main -> /workspace/miles -/mnt/cifs/home/scratch.kaixih_ent -> /scratch -``` +The launcher uses one Slurm task and four GPUs per node. -The launcher uses: +## 2. Check the Checkouts and Mounts -```bash -MILES_DIR=/workspace/miles -PYTHONPATH=/workspace/miles:/root/Megatron-LM/ +```text +Miles PR #1634: /home/scratch.kaixih_ent/repo/miles-hybridep-pr +Launcher/lab: /home/scratch.kaixih_ent/repo/miles-opd-gb200-main +TMS PR #82: /home/scratch.kaixih_ent/repo/torch_memory_saver ``` -Therefore Python imports Miles from the mounted checkout, including its committed changes. Megatron, PyTorch, torch-memory-saver, and SGLang still come from `radixark/miles:latest`; their installed source files are not edited. - -## 3. Build the Runtime Compatibility Shims - -The temporary shim sources are local lab artifacts and are intentionally not committed: +The compute-node mounts are: ```text -lab/opd_gb200/tms_cumem_trace.c -lab/opd_gb200/torch_shm_unlink_compat.c +/mnt/cifs/home/scratch.kaixih_ent/repo/miles-hybridep-pr -> /workspace/miles +/mnt/cifs/home/scratch.kaixih_ent/repo/miles-opd-gb200-main/lab -> /workspace/lab +/mnt/cifs/home/scratch.kaixih_ent/repo/torch_memory_saver -> /workspace/torch_memory_saver +/mnt/cifs/home/scratch.kaixih_ent -> /scratch ``` -Compile them on an ARM GB200 compute node, inside the same Miles image. Do not compile them on the x86 login node. +## 3. Build the Shim and Install TMS -From an active two-node allocation shell: +Build the AArch64 `shm_unlink` shim inside the Miles container on a GB200 node: ```bash cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main @@ -60,139 +66,113 @@ export ALLOC_JOB_ID="${SLURM_JOB_ID}" bash lab/opd_gb200/00_build_tms_compat.sh ``` -The build script runs the equivalent of: - -```bash -mkdir -p /scratch/repro/opd-self-distill/runtime/aarch64 - -cuda_header=$(find /usr/local/cuda* -path '*/targets/*/include/cuda.h' -print -quit) -cuda_include=$(dirname "${cuda_header}") - -cc -shared -fPIC -O2 -I"${cuda_include}" \ - /workspace/miles/lab/opd_gb200/tms_cumem_trace.c \ - -o /scratch/repro/opd-self-distill/runtime/aarch64/libtms_cumem_compat.so \ - -ldl - -cc -shared -fPIC -O2 \ - /workspace/miles/lab/opd_gb200/torch_shm_unlink_compat.c \ - -o /scratch/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so \ - -ldl -``` - -The persistent host equivalents are: +Only this runtime shim is loaded: ```text -/home/scratch.kaixih_ent/repro/opd-self-distill/runtime/aarch64/libtms_cumem_compat.so -/home/scratch.kaixih_ent/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so +/scratch/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so ``` -Verify both outputs before launching: +`node_entrypoint.sh` copies the mounted TMS source to node-local storage, runs +`make reinstall`, and verifies the installed PR #82 library. Both nodes must +print: -```bash -file /home/scratch.kaixih_ent/repro/opd-self-distill/runtime/aarch64/*.so +```text +TMS_PR82_VERIFIED=1 +TMS_SOURCE_COMMIT=c96bf60e093b4bec2b045cb5b8a08601d0ae8a79 ``` -They must be AArch64 shared objects. `phase2_gb200.sh` passes their container paths into the Ray runtime environment, and `miles/ray/actor_group.py` prepends them to `LD_PRELOAD` before the original torch-memory-saver library. - -## 4. Launch the Two-Node Run +## 4. Launch a Two-Update Validation -The allocation must contain two `gb200nvl72` nodes with four GPUs per node. The local launcher starts one container per node, creates a Ray head on task 0, joins task 1 as a worker, and submits the Phase-2 payload from `/workspace/miles`. - -Run the current accuracy-focused recipe: +Until SGLang PR #31072 is merged, the launcher applies +`lab/opd_gb200/patches/sglang_hybrid_post_replay_war.patch` when +`SGLANG_HYBRID_POST_REPLAY_WAR=1`. ```bash cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main -RUN_ID="opd-gb200-$(date +%m%d-%H%M)" \ -NUM_ROLLOUT=12 \ +RUN_ID="opd-p2-hybridwar-$(date +%m%d-%H%M)" \ +HOST_MILES=/home/scratch.kaixih_ent/repo/miles-hybridep-pr \ +NUM_ROLLOUT=2 \ SKIP_EVAL_BEFORE_TRAIN=1 \ -EVAL_INTERVAL=2 \ -bash lab/opd_gb200/04_run_long_8k_hybridep.sh -``` - -The effective settings are: - -```text -2 nodes x 4 GB200 GPUs -TP=2, CP=2, EP=8 -Flex dispatcher with HybridEP -global batch size=256 -max tokens per GPU=8192 -CUDA graph capture sizes=1 2 4 -checkpoint saving disabled -``` - -CUDA graph sizes larger than 4 currently use eager decode. This is a conservative mitigation for a separate replay hang after online weight handoff; it is not a final SGLang fix. - -## 5. What Changed and Why - -### Miles model-input alignment - -Files: +EVAL_INTERVAL=5 \ +SAVE_INTERVAL=none \ +ROLLOUT_BATCH_SIZE=32 \ +N_SAMPLES_PER_PROMPT=8 \ +OVER_SAMPLING_BATCH_SIZE=32 \ +GLOBAL_BATCH_SIZE=256 \ +MAX_TOKENS_PER_GPU=8192 \ +MOE_FLEX_DISPATCHER_BACKEND=hybridep \ +SGLANG_CUDA_GRAPH_BS="1 2 4 8 16 32 64 128 256 512" \ +SGLANG_RECAPTURE_AFTER_WEIGHT_UPDATE=pr27140 \ +SGLANG_SYNC_BEFORE_MEMORY_RELEASE=1 \ +SGLANG_HYBRID_POST_REPLAY_WAR=1 \ +SGLANG_PR30895=0 \ +SGLANG_PR30974=0 \ +SGLANG_DISABLE_CUDA_GRAPH_PADDING=0 \ +SGLANG_NEUTRALIZE_MOE_CUDA_GRAPH_PADDING=0 \ +FLASHINFER_DISTRIBUTED_AUTOTUNE_SYNC=0 \ +TMS_CUMEM_TRACE_LIB="" \ +TORCH_SHM_UNLINK_COMPAT_LIB=/scratch/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so \ +bash lab/opd_gb200/01_launch_phase2_pure.sh +``` + +For a B200-comparable curve, set `NUM_ROLLOUT=12` and +`SKIP_EVAL_BEFORE_TRAIN=0`; keep the runtime stack and capture list unchanged. + +## 5. Validate the Result + +A successful run must show: + +- `TMS_PR82_VERIFIED=1` on both nodes; +- the #27140 recapture and pre-release synchronization patches verified; +- `SGLANG_HYBRID_POST_REPLAY_WAR_VERIFIED=1` on both nodes; +- `disable_overlap_schedule=False` and `disable_cuda_graph_padding=False`; +- all ten power-of-two graph buckets recaptured after weight updates; +- second-rollout padded batches reporting `cuda graph: True`; +- no HybridEP timeout, CUDA IPC error, replay error, or non-finite metric; +- OIS, ESS, and TIS near 1 and a finite ordinary-scale gradient norm; +- Ray reporting `Job 'qwen3.5-opd-pure' succeeded`. + +## Validation Evidence + +Both runs below used the new overlap-event ordering and completed two real +rollout/update cycles: + +| Run | Old padding PRs | Response length | Reverse-KL | Gradient norm | +| --- | --- | --- | --- | --- | +| `opd-p2-hybridwar-r1c-0713-1937` | #30895/#30974 on | 18881.8 -> 6469.1 | 0.04488 -> 0.01315 | 0.3104 -> 0.1387 | +| `opd-p2-hybridwar-no56-r1c-0713-2041` | #30895/#30974 off | 18931.8 -> 6551.2 | 0.04466 -> 0.01359 | 0.3112 -> 0.1376 | + +For the A/B run with #30895/#30974 disabled, update 1 finished with +`OIS=1.0000001`, `ESS=0.9999997`, and `TIS=1.0000067`. Its second rollout +crossed padded request counts including 78, 75, 73, 70, and 69 while remaining +on CUDA graph replay. + +The event-order unit test also passes in the Miles container: ```text -miles/backends/megatron_utils/model.py -miles/backends/training_utils/data.py -``` - -Dynamic THD batches can produce different token-row counts on different EP ranks. HybridEP expects rank-consistent metadata and buffer capacity. A mismatch caused `HYBRID-EP ALLGATHER TIMEOUT` and corrupted updates with extremely large gradient norms and importance-sampling ratios. - -For Flex + HybridEP only, Miles now all-reduces the maximum aligned token-row count across the EP group and pads every rank's model input to that common target. - -### Miles routing-replay alignment - -File: - -```text -miles/backends/training_utils/replay_data.py -``` - -Model inputs and routing replay must use the same EP-wide row target. Local-only replay padding produced errors such as `replay n_tokens ... does not match scores n_tokens ...` under TP2 plus sequence parallelism. Replay rows are now padded to the same EP-wide target before TP/SP slicing; padding remains invalid with value `-1` and stays outside the loss. - -### GB200 recipe and runtime environment - -File: - -```text -examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh -``` - -The script now selects HybridEP explicitly and makes rollout count, batch sizes, response length, token budget, evaluation cadence, checkpoint cadence, and SGLang CUDA graph sizes configurable. It also forwards the two shim paths into Ray actors. - -### Ray actor preload composition - -File: - -```text -miles/ray/actor_group.py -``` - -The actor environment composes this preload order: - -```text -libtorch_shm_unlink_compat.so -libtms_cumem_compat.so -the original torch-memory-saver preload library -``` - -`libtms_cumem_compat.so` bypasses torch-memory-saver only for small allocations that are not aligned to the GB200 CUDA VMM granularity, routing them to native `cudaMalloc`. - -`libtorch_shm_unlink_compat.so` makes cleanup idempotent when PyTorch calls `shm_unlink()` for an already-removed `/torch_*` CUDA IPC refcount file. Only `ENOENT` for `/torch_*` is converted to success. - -These two shims are dependency workarounds. The permanent fixes belong in torch-memory-saver/container packaging and PyTorch respectively, not in OPD mathematics. - -## 6. Validation Signals - -A healthy run must have all of the following: - -- no `HYBRID-EP ALLGATHER TIMEOUT`; -- no routing replay row mismatch; -- no fatal CUDA IPC cleanup error; -- `OIS`, `ESS`, and `TIS` remain near 1; -- finite, ordinary-scale gradient norms; -- successful online weight synchronization and the next rollout; -- response length and reverse-KL move toward the teacher behavior. - -The validated run changed mean response length from about 18.8k to 6.6k tokens after the first update, reduced truncation from 0.367 to 0.043, and reduced per-token OPD reverse-KL from 0.0448 to 0.0130 while keeping the training diagnostics healthy. - -No Megatron, PyTorch, torch-memory-saver, or SGLang installed source file is modified by this workflow. The active code consists of the mounted Miles checkout, the two external runtime shims, and recipe flags. +plain attention: load -> read-done event -> replay +hybrid attention: load -> replay -> read-done event +2 passed +``` + +## Why Each Change Exists + +- **Miles #1634:** keeps model-input and routing-replay token rows aligned + across HybridEP ranks. +- **PyTorch unlink shim:** makes cleanup idempotent only for an already-removed + `/torch_*` CUDA-IPC refcount file. +- **TMS #82:** fixes the GB200 CUDA VMM allocation-size failure. +- **SGLang #27140:** prevents replaying graph objects whose memory was released + and then restored by TMS. +- **Pre-release synchronization:** finishes asynchronous replay before TMS + unmaps its backing memory. +- **Hybrid overlap-event fix:** the existing WAR fast path publishes read-done + after `load_batch`, which is safe for plain attention. Hybrid/Mamba decode + keeps reading shared request/state buffers during graph replay, so the + scheduler could mutate the next iteration's buffers too early. Publishing + the event after replay closes that race while preserving the fast path for + plain attention. + +At teardown, W&B may emit an `atexit` `BrokenPipeError`. Treat it as teardown +noise only when both updates completed and Ray explicitly reported success. From 95393cdd4ebe35e7905f33129fd1026f3902ddbe Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Mon, 13 Jul 2026 15:01:19 -0700 Subject: [PATCH 4/8] Document GB200 fixes by upstream PR --- workflow_gb200.md | 52 +++++++++++++++++++++++++++++++---------------- 1 file changed, 35 insertions(+), 17 deletions(-) diff --git a/workflow_gb200.md b/workflow_gb200.md index 6c046308f6a..c8068f98555 100644 --- a/workflow_gb200.md +++ b/workflow_gb200.md @@ -8,22 +8,31 @@ CUDA graph padding, and power-of-two capture enabled. Use these pieces together: -1. Miles PR #1634 from branch `agent/gb200-hybridep-fixes`. -2. The narrow PyTorch `shm_unlink(ENOENT)` compatibility shim. -3. A real installation of torch-memory-saver PR #82 at commit +1. [Miles PR #1634](https://github.com/radixark/miles/pull/1634) from + branch `agent/gb200-hybridep-fixes`. +2. The narrow PyTorch `shm_unlink(ENOENT)` compatibility shim. This is the + only remaining local workaround without an upstream PR. +3. A real installation of + [torch-memory-saver PR #82](https://github.com/fzyzcjy/torch_memory_saver/pull/82) + at commit `c96bf60e093b4bec2b045cb5b8a08601d0ae8a79`. -4. SGLang PR #27140 behavior: invalidate stale CUDA graphs and recapture them - after torch-memory-saver resumes graph-owned memory. -5. Synchronize asynchronous replay before torch-memory-saver unmaps graph - memory. -6. SGLang PR #31072: for hybrid linear-attention models, publish the - overlap-scheduler read-done event after CUDA graph replay, not before it. +4. [SGLang PR #27140](https://github.com/sgl-project/sglang/pull/27140): + invalidate stale CUDA graphs and recapture them after torch-memory-saver + resumes graph-owned memory. +5. [SGLang PR #31073](https://github.com/sgl-project/sglang/pull/31073): + synchronize asynchronous device work before torch-memory-saver unmaps its + backing memory. +6. [SGLang PR #31072](https://github.com/sgl-project/sglang/pull/31072): for + hybrid linear-attention models, publish the overlap-scheduler read-done + event after CUDA graph replay, not before it. 7. Capture exactly `1 2 4 8 16 32 64 128 256 512`, with graph padding and overlap scheduling enabled. SGLang PRs #30895 and #30974 are not required for this hang. The second validation completed with both disabled. They may remain useful as padded-row -hygiene, but they are not part of the minimum validated stack. +hygiene, but they are not part of the minimum validated stack. The other +padding, replay-debug, and FlashInfer patches under `lab/opd_gb200/patches/` +are disabled experiments and are not part of this workflow. Do not use the historical `libtms_cumem_compat.so` shim. Keep `TMS_CUMEM_TRACE_LIB` empty. The local `lab/` directory contains launchers and @@ -83,9 +92,17 @@ TMS_SOURCE_COMMIT=c96bf60e093b4bec2b045cb5b8a08601d0ae8a79 ## 4. Launch a Two-Update Validation -Until SGLang PR #31072 is merged, the launcher applies -`lab/opd_gb200/patches/sglang_hybrid_post_replay_war.patch` when -`SGLANG_HYBRID_POST_REPLAY_WAR=1`. +Until the SGLang PRs are present in the container image, the launcher applies +these exact backports: + +| Upstream PR | Launcher patch | Switch | +| --- | --- | --- | +| [#27140](https://github.com/sgl-project/sglang/pull/27140) | `sglang_pr27140_forward_port.patch` and `sglang_pr27140_graph_reset.patch` | `SGLANG_RECAPTURE_AFTER_WEIGHT_UPDATE=pr27140` | +| [#31073](https://github.com/sgl-project/sglang/pull/31073) | `sglang_sync_before_tms_pause.patch` | `SGLANG_SYNC_BEFORE_MEMORY_RELEASE=1` | +| [#31072](https://github.com/sgl-project/sglang/pull/31072) | `sglang_hybrid_post_replay_war.patch` | `SGLANG_HYBRID_POST_REPLAY_WAR=1` | + +Do not add an unnamed source patch to this workflow. Create or reference its +upstream PR first. ```bash cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main @@ -124,7 +141,8 @@ For a B200-comparable curve, set `NUM_ROLLOUT=12` and A successful run must show: - `TMS_PR82_VERIFIED=1` on both nodes; -- the #27140 recapture and pre-release synchronization patches verified; +- the #27140 recapture and #31073 pre-release synchronization backports + verified; - `SGLANG_HYBRID_POST_REPLAY_WAR_VERIFIED=1` on both nodes; - `disable_overlap_schedule=False` and `disable_cuda_graph_padding=False`; - all ten power-of-two graph buckets recaptured after weight updates; @@ -165,9 +183,9 @@ hybrid attention: load -> replay -> read-done event - **TMS #82:** fixes the GB200 CUDA VMM allocation-size failure. - **SGLang #27140:** prevents replaying graph objects whose memory was released and then restored by TMS. -- **Pre-release synchronization:** finishes asynchronous replay before TMS - unmaps its backing memory. -- **Hybrid overlap-event fix:** the existing WAR fast path publishes read-done +- **SGLang #31073:** finishes asynchronous device work before TMS unmaps its + backing memory. +- **SGLang #31072:** the existing WAR fast path publishes read-done after `load_batch`, which is safe for plain attention. Hybrid/Mamba decode keeps reading shared request/state buffers during graph replay, so the scheduler could mutate the next iteration's buffers too early. Publishing From 0c7701e4fa972d851b1c289a68698475c3c7c426 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Mon, 13 Jul 2026 17:18:37 -0700 Subject: [PATCH 5/8] Revert unrelated make_split formatting --- .../qwen3_5_35b_selfdistill/make_split.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py index 8a91a271486..d15366d2643 100644 --- a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py @@ -42,10 +42,10 @@ def main(): rows.append(line) random.Random(args.seed).shuffle(rows) # fixed seed: reproducible, unbiased split - eval_rows, train_rows = rows[-args.eval_n :], rows[: -args.eval_n] + eval_rows, train_rows = rows[-args.eval_n:], rows[:-args.eval_n] - ek = {prompt_text(json.loads(line)) for line in eval_rows} - tk = {prompt_text(json.loads(line)) for line in train_rows} + ek = {prompt_text(json.loads(l)) for l in eval_rows} + tk = {prompt_text(json.loads(l)) for l in train_rows} assert ek.isdisjoint(tk), "LEAK: eval prompt found in train split" with open(os.path.join(args.out_dir, "dapo_train.jsonl"), "w") as f: From c0469398c28b01b6ce17a8b8efc80464ba9cf504 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Mon, 13 Jul 2026 17:24:48 -0700 Subject: [PATCH 6/8] Narrow GB200 HybridEP PR scope --- .../qwen3_5_35b_selfdistill/phase2_gb200.sh | 47 +---- .../qwen3_5_35b_selfdistill/rm.py | 2 +- miles/backends/training_utils/replay_data.py | 2 +- miles/ray/actor_group.py | 17 +- miles_plugins/mbridge/qwen3_5.py | 4 +- workflow_gb200.md | 196 ------------------ 6 files changed, 16 insertions(+), 252 deletions(-) delete mode 100644 workflow_gb200.md diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh index 94c6350fcb2..46ca7c53bb7 100755 --- a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh @@ -42,18 +42,6 @@ EXAMPLE_DIR=${EXAMPLE_DIR:-$(cd "$(dirname "$0")" && pwd)} MILES_DIR=${MILES_DIR:-/workspace/miles} RAY_ADDRESS=${RAY_ADDRESS:-http://127.0.0.1:8265} OPD_KL_COEF=${OPD_KL_COEF:-0.2} -MOE_FLEX_DISPATCHER_BACKEND=${MOE_FLEX_DISPATCHER_BACKEND:-hybridep} -NUM_ROLLOUT=${NUM_ROLLOUT:-12} -ROLLOUT_BATCH_SIZE=${ROLLOUT_BATCH_SIZE:-32} -N_SAMPLES_PER_PROMPT=${N_SAMPLES_PER_PROMPT:-8} -OVER_SAMPLING_BATCH_SIZE=${OVER_SAMPLING_BATCH_SIZE:-32} -GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-256} -ROLLOUT_MAX_RESPONSE_LEN=${ROLLOUT_MAX_RESPONSE_LEN:-24576} -MAX_TOKENS_PER_GPU=${MAX_TOKENS_PER_GPU:-16384} -SAVE_INTERVAL=${SAVE_INTERVAL:-5} -EVAL_INTERVAL=${EVAL_INTERVAL:-5} -SKIP_EVAL_BEFORE_TRAIN=${SKIP_EVAL_BEFORE_TRAIN:-0} -read -r -a SGLANG_CUDA_GRAPH_BS_ARRAY <<< "${SGLANG_CUDA_GRAPH_BS:-1 2 4}" mkdir -p "${OUTPUT_DIR}" EVAL_CONFIG="${OUTPUT_DIR}/eval_dapo_heldout.yaml" @@ -75,7 +63,7 @@ MODEL_ARGS=( --untie-embeddings-and-output-weights --vocab-size 248320 --rotary-base 10000000 --moe-ffn-hidden-size 512 --moe-shared-expert-intermediate-size 512 --moe-router-score-function softmax --moe-token-dispatcher-type flex - --moe-flex-dispatcher-backend ${MOE_FLEX_DISPATCHER_BACKEND} + --moe-flex-dispatcher-backend hybridep --moe-router-topk 8 --moe-layer-freq "[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]" --num-experts 256 --moe-grouped-gemm --moe-token-drop-policy probs --moe-router-dtype fp32 @@ -84,29 +72,21 @@ MODEL_ARGS=( ) CKPT_ARGS=( --hf-checkpoint ${MODEL_DIR}/Qwen3.5-35B-A3B - --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist --load ${OUTPUT_DIR} + --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist + --load ${OUTPUT_DIR} --save ${OUTPUT_DIR} --save-interval 5 ) -if [ "${SAVE_INTERVAL}" != "none" ]; then - CKPT_ARGS+=(--save ${OUTPUT_DIR} --save-interval ${SAVE_INTERVAL}) -fi OPD_ARGS=( --use-opd --opd-type megatron --opd-teacher-load ${TEACHER_LOAD} --opd-kl-coef ${OPD_KL_COEF} ) ROLLOUT_ARGS=( --prompt-data ${DATA_DIR}/dapo_train.jsonl --input-key prompt --label-key label --apply-chat-template --rollout-shuffle - --num-rollout ${NUM_ROLLOUT} --rollout-batch-size ${ROLLOUT_BATCH_SIZE} - --n-samples-per-prompt ${N_SAMPLES_PER_PROMPT} - --rollout-max-response-len ${ROLLOUT_MAX_RESPONSE_LEN} - --rollout-temperature 1 --num-steps-per-rollout 1 - --over-sampling-batch-size ${OVER_SAMPLING_BATCH_SIZE} - --global-batch-size ${GLOBAL_BATCH_SIZE} --balance-data + --num-rollout 12 --rollout-batch-size 32 --n-samples-per-prompt 8 + --rollout-max-response-len 24576 --rollout-temperature 1 --num-steps-per-rollout 1 + --over-sampling-batch-size 32 --global-batch-size 256 --balance-data ) RM_ARGS=( --custom-rm-path ${RM_FUNC} ) -EVAL_ARGS=( --eval-interval ${EVAL_INTERVAL} --eval-config ${EVAL_CONFIG} ) -if [ "${SKIP_EVAL_BEFORE_TRAIN}" = "1" ]; then - EVAL_ARGS+=(--skip-eval-before-train) -fi +EVAL_ARGS=( --eval-interval 5 --eval-config ${EVAL_CONFIG} ) GRPO_ARGS=( --advantage-estimator grpo --kl-loss-type low_var_kl --entropy-coef 0.00 --eps-clip 0.2 --eps-clip-high 0.28 --use-tis @@ -120,7 +100,7 @@ PERF_ARGS=( --tensor-model-parallel-size 2 --sequence-parallel --pipeline-model-parallel-size 1 --context-parallel-size 2 --expert-model-parallel-size 8 --expert-tensor-parallel-size 1 --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 - --use-dynamic-batch-size --max-tokens-per-gpu ${MAX_TOKENS_PER_GPU} --log-probs-chunk-size 4096 + --use-dynamic-batch-size --max-tokens-per-gpu 8192 --log-probs-chunk-size 4096 ) # Blackwell (GB200 sm100 / B300 sm103) backends, per scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py: # - moe-runner-backend flashinfer_cutlass: the default triton fused-MoE mis-shards @@ -130,7 +110,7 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine 8 --sglang-mem-fraction-static 0.7 --sglang-ep-size 8 --sglang-watchdog-timeout 1800 --sglang-enable-metrics --sglang-moe-runner-backend flashinfer_cutlass --sglang-attention-backend trtllm_mha - --sglang-cuda-graph-bs "${SGLANG_CUDA_GRAPH_BS_ARRAY[@]}" --use-rollout-routing-replay + --sglang-cuda-graph-bs 1 2 4 8 16 32 --use-rollout-routing-replay --sglang-mamba-scheduler-strategy extra_buffer ) MISC_ARGS=( @@ -147,14 +127,7 @@ WANDB_ARGS=( --use-wandb --wandb-project miles-opd --wandb-group qwen3.5-35b-opd # during cross-node ncclCommInitRank ("unhandled cuda error" / CUDA 999), which # kills the EP8 sglang engine spanning both nodes. MNNVL stays enabled (pod env) # for the NVLink fabric itself; only NVLS multicast is disabled. -TRACE_RUNTIME_ENV="" -if [ -n "${TMS_CUMEM_TRACE_LIB:-}" ]; then - TRACE_RUNTIME_ENV=", \"TMS_CUMEM_TRACE_LIB\": \"${TMS_CUMEM_TRACE_LIB}\", \"CUMEM_TRACE_ALL\": \"${CUMEM_TRACE_ALL:-}\"" -fi -if [ -n "${TORCH_SHM_UNLINK_COMPAT_LIB:-}" ]; then - TRACE_RUNTIME_ENV="${TRACE_RUNTIME_ENV}, \"TORCH_SHM_UNLINK_COMPAT_LIB\": \"${TORCH_SHM_UNLINK_COMPAT_LIB}\"" -fi -RUNTIME_ENV_JSON="{\"env_vars\": {\"PYTHONPATH\": \"${MILES_DIR}:/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"WANDB_API_KEY\": \"${WANDB_API_KEY}\", \"PROMETHEUS_PORT\": \"9090\", \"NCCL_NVLS_ENABLE\": \"0\"${TRACE_RUNTIME_ENV}}}" +RUNTIME_ENV_JSON="{\"env_vars\": {\"PYTHONPATH\": \"${MILES_DIR}:/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"WANDB_API_KEY\": \"${WANDB_API_KEY}\", \"PROMETHEUS_PORT\": \"9090\", \"NCCL_NVLS_ENABLE\": \"0\"}}" cd "${MILES_DIR}" ray job submit --address="${RAY_ADDRESS}" --submission-id qwen3.5-opd-${MODE} --no-wait \ diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py index 88b8cbd572e..9a68f09b1fe 100644 --- a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py @@ -28,8 +28,8 @@ ``--custom-rm-path examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func``. """ -from miles.rollout.rm_hub.math_dapo_utils import compute_score as _dapo_score from miles.rollout.rm_hub.math_utils import grade_answer_verl +from miles.rollout.rm_hub.math_dapo_utils import compute_score as _dapo_score def _is_eval_sample(sample) -> bool: diff --git a/miles/backends/training_utils/replay_data.py b/miles/backends/training_utils/replay_data.py index fc8a13e9c57..d26a514035e 100644 --- a/miles/backends/training_utils/replay_data.py +++ b/miles/backends/training_utils/replay_data.py @@ -125,7 +125,7 @@ def pad_func(data, pad): target = torch.tensor( [local_target], dtype=torch.int64, - device=torch.cuda.current_device(), + device=replay_data.device, ) dist.all_reduce(target, op=dist.ReduceOp.MAX, group=parallel_state.ep.group) target_num_tokens = int(target.item()) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 4766766d523..f729d948aed 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -74,22 +74,7 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): ) assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." - preload_libs = [] - shm_compat_lib = os.environ.get("TORCH_SHM_UNLINK_COMPAT_LIB") - if shm_compat_lib: - assert os.path.exists( - shm_compat_lib - ), f"torch shared-memory cleanup library {shm_compat_lib} does not exist." - preload_libs.append(shm_compat_lib) - - trace_lib = os.environ.get("TMS_CUMEM_TRACE_LIB") - if trace_lib: - assert os.path.exists(trace_lib), f"cuMemCreate trace library {trace_lib} does not exist." - preload_libs.append(trace_lib) - env_vars["CUMEM_TRACE_ALL"] = os.environ.get("CUMEM_TRACE_ALL", "") - - preload_libs.append(dynlib_path) - env_vars["LD_PRELOAD"] = ":".join(preload_libs) + env_vars["LD_PRELOAD"] = dynlib_path env_vars["TMS_INIT_ENABLE"] = "1" env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index 67d1dcb35f7..9f467eedf0b 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -211,7 +211,9 @@ def _experts_fused(self) -> bool: index = getattr(io, "index", None) if io is not None else None if not index: return True - fused = any("model.language_model.layers." in k and k.endswith("mlp.experts.gate_up_proj") for k in index) + fused = any( + "model.language_model.layers." in k and k.endswith("mlp.experts.gate_up_proj") for k in index + ) self._experts_fused_cached = fused return fused diff --git a/workflow_gb200.md b/workflow_gb200.md deleted file mode 100644 index c8068f98555..00000000000 --- a/workflow_gb200.md +++ /dev/null @@ -1,196 +0,0 @@ -# Two-Node GB200 Pure-OPD Workflow - -Validated on 2026-07-13 with two independent two-node runs. Each run completed -two real rollouts and two online training updates with overlap scheduling, -CUDA graph padding, and power-of-two capture enabled. - -## Required Stack - -Use these pieces together: - -1. [Miles PR #1634](https://github.com/radixark/miles/pull/1634) from - branch `agent/gb200-hybridep-fixes`. -2. The narrow PyTorch `shm_unlink(ENOENT)` compatibility shim. This is the - only remaining local workaround without an upstream PR. -3. A real installation of - [torch-memory-saver PR #82](https://github.com/fzyzcjy/torch_memory_saver/pull/82) - at commit - `c96bf60e093b4bec2b045cb5b8a08601d0ae8a79`. -4. [SGLang PR #27140](https://github.com/sgl-project/sglang/pull/27140): - invalidate stale CUDA graphs and recapture them after torch-memory-saver - resumes graph-owned memory. -5. [SGLang PR #31073](https://github.com/sgl-project/sglang/pull/31073): - synchronize asynchronous device work before torch-memory-saver unmaps its - backing memory. -6. [SGLang PR #31072](https://github.com/sgl-project/sglang/pull/31072): for - hybrid linear-attention models, publish the overlap-scheduler read-done - event after CUDA graph replay, not before it. -7. Capture exactly `1 2 4 8 16 32 64 128 256 512`, with graph padding and - overlap scheduling enabled. - -SGLang PRs #30895 and #30974 are not required for this hang. The second -validation completed with both disabled. They may remain useful as padded-row -hygiene, but they are not part of the minimum validated stack. The other -padding, replay-debug, and FlashInfer patches under `lab/opd_gb200/patches/` -are disabled experiments and are not part of this workflow. - -Do not use the historical `libtms_cumem_compat.so` shim. Keep -`TMS_CUMEM_TRACE_LIB` empty. The local `lab/` directory contains launchers and -temporary patches and is intentionally not committed. - -## 1. Get Two Nodes - -Run this from the visible tmux shell on `dl3`: - -```bash -salloc_node gb200nvl72 2 -``` - -The launcher uses one Slurm task and four GPUs per node. - -## 2. Check the Checkouts and Mounts - -```text -Miles PR #1634: /home/scratch.kaixih_ent/repo/miles-hybridep-pr -Launcher/lab: /home/scratch.kaixih_ent/repo/miles-opd-gb200-main -TMS PR #82: /home/scratch.kaixih_ent/repo/torch_memory_saver -``` - -The compute-node mounts are: - -```text -/mnt/cifs/home/scratch.kaixih_ent/repo/miles-hybridep-pr -> /workspace/miles -/mnt/cifs/home/scratch.kaixih_ent/repo/miles-opd-gb200-main/lab -> /workspace/lab -/mnt/cifs/home/scratch.kaixih_ent/repo/torch_memory_saver -> /workspace/torch_memory_saver -/mnt/cifs/home/scratch.kaixih_ent -> /scratch -``` - -## 3. Build the Shim and Install TMS - -Build the AArch64 `shm_unlink` shim inside the Miles container on a GB200 node: - -```bash -cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main -export ALLOC_JOB_ID="${SLURM_JOB_ID}" -bash lab/opd_gb200/00_build_tms_compat.sh -``` - -Only this runtime shim is loaded: - -```text -/scratch/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so -``` - -`node_entrypoint.sh` copies the mounted TMS source to node-local storage, runs -`make reinstall`, and verifies the installed PR #82 library. Both nodes must -print: - -```text -TMS_PR82_VERIFIED=1 -TMS_SOURCE_COMMIT=c96bf60e093b4bec2b045cb5b8a08601d0ae8a79 -``` - -## 4. Launch a Two-Update Validation - -Until the SGLang PRs are present in the container image, the launcher applies -these exact backports: - -| Upstream PR | Launcher patch | Switch | -| --- | --- | --- | -| [#27140](https://github.com/sgl-project/sglang/pull/27140) | `sglang_pr27140_forward_port.patch` and `sglang_pr27140_graph_reset.patch` | `SGLANG_RECAPTURE_AFTER_WEIGHT_UPDATE=pr27140` | -| [#31073](https://github.com/sgl-project/sglang/pull/31073) | `sglang_sync_before_tms_pause.patch` | `SGLANG_SYNC_BEFORE_MEMORY_RELEASE=1` | -| [#31072](https://github.com/sgl-project/sglang/pull/31072) | `sglang_hybrid_post_replay_war.patch` | `SGLANG_HYBRID_POST_REPLAY_WAR=1` | - -Do not add an unnamed source patch to this workflow. Create or reference its -upstream PR first. - -```bash -cd /home/scratch.kaixih_ent/repo/miles-opd-gb200-main - -RUN_ID="opd-p2-hybridwar-$(date +%m%d-%H%M)" \ -HOST_MILES=/home/scratch.kaixih_ent/repo/miles-hybridep-pr \ -NUM_ROLLOUT=2 \ -SKIP_EVAL_BEFORE_TRAIN=1 \ -EVAL_INTERVAL=5 \ -SAVE_INTERVAL=none \ -ROLLOUT_BATCH_SIZE=32 \ -N_SAMPLES_PER_PROMPT=8 \ -OVER_SAMPLING_BATCH_SIZE=32 \ -GLOBAL_BATCH_SIZE=256 \ -MAX_TOKENS_PER_GPU=8192 \ -MOE_FLEX_DISPATCHER_BACKEND=hybridep \ -SGLANG_CUDA_GRAPH_BS="1 2 4 8 16 32 64 128 256 512" \ -SGLANG_RECAPTURE_AFTER_WEIGHT_UPDATE=pr27140 \ -SGLANG_SYNC_BEFORE_MEMORY_RELEASE=1 \ -SGLANG_HYBRID_POST_REPLAY_WAR=1 \ -SGLANG_PR30895=0 \ -SGLANG_PR30974=0 \ -SGLANG_DISABLE_CUDA_GRAPH_PADDING=0 \ -SGLANG_NEUTRALIZE_MOE_CUDA_GRAPH_PADDING=0 \ -FLASHINFER_DISTRIBUTED_AUTOTUNE_SYNC=0 \ -TMS_CUMEM_TRACE_LIB="" \ -TORCH_SHM_UNLINK_COMPAT_LIB=/scratch/repro/opd-self-distill/runtime/aarch64/libtorch_shm_unlink_compat.so \ -bash lab/opd_gb200/01_launch_phase2_pure.sh -``` - -For a B200-comparable curve, set `NUM_ROLLOUT=12` and -`SKIP_EVAL_BEFORE_TRAIN=0`; keep the runtime stack and capture list unchanged. - -## 5. Validate the Result - -A successful run must show: - -- `TMS_PR82_VERIFIED=1` on both nodes; -- the #27140 recapture and #31073 pre-release synchronization backports - verified; -- `SGLANG_HYBRID_POST_REPLAY_WAR_VERIFIED=1` on both nodes; -- `disable_overlap_schedule=False` and `disable_cuda_graph_padding=False`; -- all ten power-of-two graph buckets recaptured after weight updates; -- second-rollout padded batches reporting `cuda graph: True`; -- no HybridEP timeout, CUDA IPC error, replay error, or non-finite metric; -- OIS, ESS, and TIS near 1 and a finite ordinary-scale gradient norm; -- Ray reporting `Job 'qwen3.5-opd-pure' succeeded`. - -## Validation Evidence - -Both runs below used the new overlap-event ordering and completed two real -rollout/update cycles: - -| Run | Old padding PRs | Response length | Reverse-KL | Gradient norm | -| --- | --- | --- | --- | --- | -| `opd-p2-hybridwar-r1c-0713-1937` | #30895/#30974 on | 18881.8 -> 6469.1 | 0.04488 -> 0.01315 | 0.3104 -> 0.1387 | -| `opd-p2-hybridwar-no56-r1c-0713-2041` | #30895/#30974 off | 18931.8 -> 6551.2 | 0.04466 -> 0.01359 | 0.3112 -> 0.1376 | - -For the A/B run with #30895/#30974 disabled, update 1 finished with -`OIS=1.0000001`, `ESS=0.9999997`, and `TIS=1.0000067`. Its second rollout -crossed padded request counts including 78, 75, 73, 70, and 69 while remaining -on CUDA graph replay. - -The event-order unit test also passes in the Miles container: - -```text -plain attention: load -> read-done event -> replay -hybrid attention: load -> replay -> read-done event -2 passed -``` - -## Why Each Change Exists - -- **Miles #1634:** keeps model-input and routing-replay token rows aligned - across HybridEP ranks. -- **PyTorch unlink shim:** makes cleanup idempotent only for an already-removed - `/torch_*` CUDA-IPC refcount file. -- **TMS #82:** fixes the GB200 CUDA VMM allocation-size failure. -- **SGLang #27140:** prevents replaying graph objects whose memory was released - and then restored by TMS. -- **SGLang #31073:** finishes asynchronous device work before TMS unmaps its - backing memory. -- **SGLang #31072:** the existing WAR fast path publishes read-done - after `load_batch`, which is safe for plain attention. Hybrid/Mamba decode - keeps reading shared request/state buffers during graph replay, so the - scheduler could mutate the next iteration's buffers too early. Publishing - the event after replay closes that race while preserving the fast path for - plain attention. - -At teardown, W&B may emit an `atexit` `BrokenPipeError`. Treat it as teardown -noise only when both updates completed and Ray explicitly reported success. From 15d8abedd70459f819c4f345fadc99ab383d1f64 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Mon, 13 Jul 2026 18:53:23 -0700 Subject: [PATCH 7/8] [OPD] Place HybridEP replay row collective on CUDA --- miles/backends/training_utils/replay_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/miles/backends/training_utils/replay_data.py b/miles/backends/training_utils/replay_data.py index d26a514035e..fc8a13e9c57 100644 --- a/miles/backends/training_utils/replay_data.py +++ b/miles/backends/training_utils/replay_data.py @@ -125,7 +125,7 @@ def pad_func(data, pad): target = torch.tensor( [local_target], dtype=torch.int64, - device=replay_data.device, + device=torch.cuda.current_device(), ) dist.all_reduce(target, op=dist.ReduceOp.MAX, group=parallel_state.ep.group) target_num_tokens = int(target.item()) From dff2ef8418c03606acec56f1615f88322af6e926 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Mon, 13 Jul 2026 23:13:23 -0700 Subject: [PATCH 8/8] [OPD] Disable SGLang overlap scheduling on GB200 --- .../qwen3_5_35b_selfdistill/phase2_gb200.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh index 46ca7c53bb7..ea47c962c13 100755 --- a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh @@ -111,6 +111,9 @@ SGLANG_ARGS=( --sglang-watchdog-timeout 1800 --sglang-enable-metrics --sglang-moe-runner-backend flashinfer_cutlass --sglang-attention-backend trtllm_mha --sglang-cuda-graph-bs 1 2 4 8 16 32 --use-rollout-routing-replay + # Hybrid linear-attention CUDA graphs keep reading shared recurrent state + # during replay. Disable overlap so the next batch cannot overwrite it early. + --sglang-disable-overlap-schedule --sglang-mamba-scheduler-strategy extra_buffer ) MISC_ARGS=(