scripts: complete slime-exact port of most scripts except for gpt-oss 20B support

docs/en/advanced/low-precision.md

@@ -0,0 +1,143 @@
+# Low Precision Training and Rollout
+
+Low precision in vime is primarily used to make rollout faster and more memory-efficient while keeping training numerically stable. For large MoE RL jobs, the recommended production path is:
+
+> **BF16 training in Megatron + FP8 rollout/inference in vLLM**
+
+Megatron keeps the trainable checkpoint in BF16/torch_dist format. vLLM serves an FP8 Hugging Face checkpoint for rollout. During weight updates, vime uses the quantization config in `--hf-checkpoint` to quantize updated BF16 weights before sending them to vLLM.
+
+## Feature Maturity
+
+| Feature | Status | Recommended Use |
+|---|---|---|
+| BF16 training + FP8 rollout/inference | Stable | Default path for large MoE RL recipes. Keeps training stable while reducing rollout memory and bandwidth. |
+| FP8 KV cache in vLLM rollout | Stable when supported by your vLLM version/GPU stack | Increase KV cache capacity for long-context or agentic rollout by passing `--vllm-kv-cache-dtype fp8_e4m3`. |
+| INT4 rollout / INT4 QAT | Beta | Use when rollout memory/throughput pressure is high and the model path has been validated. |
+| FP8 training + FP8 rollout | Experimental | Useful for research on training/inference mismatch and throughput, but still has optimizer and checkpointing caveats. |
+
+## BF16 Training with FP8 Rollout
+
+This is the main production path in vime.
+
+You can run FP8 rollout by setting `--hf-checkpoint` to a blockwise-quantized Hugging Face checkpoint. Convert a BF16 checkpoint with:
+
+```bash
+python tools/convert_hf_to_fp8.py \
+    --model-dir $BF16_MODEL \
+    --save-dir $FP8_MODEL \
+    --strategy block --block-size 128 128 \
+    --max-workers 4
+```
+
+Make sure the converted checkpoint's `config.json` contains the correct `quantization_config`. vime uses that config during weight updates, so the training side can remain BF16 while rollout receives FP8 weights.
+
+Example:
+
+```bash
+# Megatron training checkpoint remains BF16 / torch_dist.
+--ref-load /path/to/model_torch_dist
+
+# vLLM rollout checkpoint is FP8 Hugging Face.
+--hf-checkpoint /path/to/model-fp8-hf
+```
+
+## FP8 KV Cache for Rollout
+
+For long-context, multi-turn, or agentic workloads, KV cache capacity is often the bottleneck. Because vLLM arguments are passed through by adding `--vllm-`, you can enable FP8 KV cache directly:
+
+```bash
+--vllm-kv-cache-dtype fp8_e4m3
+```
+
+This is a rollout-side setting. It does not change Megatron training precision; it increases effective vLLM KV cache capacity and can allow longer contexts or higher concurrency, subject to the accuracy/performance behavior of your vLLM version and GPU stack.
+
+## FP8 Training with FP8 Rollout
+
+vime also supports experimental FP8 training paths. We observed that FP8 training plus FP8 inference can improve inference throughput and reduce training/inference mismatch in some settings. More details are available in [this blog](https://lmsys.org/blog/2025-11-25-fp8-rl/).
+
+### Quick Start
+
+1. Convert your Hugging Face model weights to FP8 format using `tools/convert_hf_to_fp8.py`.
+
+2. Add the FP8 training flags:
+
+```bash
+--fp8-format e4m3
+--fp8-recipe blockwise
+# --fp8-param-gather # optional; currently incompatible with CPU Adam
+```
+
+3. Ensure `NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1` is set. vime sets this to `1` by default for Ray actors.
+
+4. Start an FP8 training example:
+
+```bash
+# Qwen3-4B FP8 training
+bash scripts/low_precision/run-qwen3-4b-fp8.sh
+
+# Qwen3-30B-A3B FP8 training (2 nodes)
+bash scripts/low_precision/run-qwen3-30b-a3b-fp8.sh
+```
+
+### Implementation Notes
+
+1. If an FP8 recipe is enabled, TransformerEngine layers are built in an FP8 context.
+2. During training, weights and activations are quantized online to NVFP8 format, and cuBLAS FP8 GEMM is used for forward and backward GEMMs.
+3. During RL weight updates, Megatron dequantizes FP8 weights to BF16, then vime quantizes the BF16 weights to FP8 and sends them to vLLM.
+4. Checkpoints saved from the training engine are dequantized back to BF16 and saved as `torch_dist`.
+
+Only `Linear` and `GroupLinear` layers in TransformerEngine use FP8. `embedding` and `lm_head` remain in their original precision. If `--fp8-param-gather` is not enabled, TransformerEngine weights remain stored in BF16 and are cast to FP8 only during `GEMM` or `GroupGEMM`.
+
+### Known Caveat
+
+`--fp8-param-gather` can save memory, but currently requires TransformerEngine `FusedAdam`, which conflicts with the CPU Adam offload path commonly used for large Megatron-LM RL jobs.
+
+## INT4 QAT Training
+
+INT4 STE (Straight-Through Estimator) training and INT4 inference can further reduce rollout memory and improve throughput. Treat this path as beta unless you have validated the target model and reward setup.
+
+### Quick Start
+
+1. Convert Hugging Face weights to INT4:
+
+```bash
+python tools/convert_hf_to_int4_direct.py \
+  --model-dir /path/to/your/original/models \
+  --save-dir /path/to/your/save/models
+```
+
+If you only need INT4 rollout, set `--hf-checkpoint` to the converted INT4 checkpoint.
+
+2. Enable INT4 fake QAT:
+
+```json
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\",
+    \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\"
+  }
+}"
+```
+
+`OPEN_TRAINING_INT4_GROUP_SIZE` should usually be:
+
+- `128` for `moonlight-16B-A3B`, `qwen3-30B-A3B`, and `qwen3-235B-A22B-int4`;
+- `32` for `kimi-k2-Thinking-int4`.
+
+3. Launch an example:
+
+```bash
+# Moonlight-16B-A3B INT4 training
+bash scripts/low_precision/run-moonlight-16B-A3B-int4.sh
+
+# Qwen3-30B-A3B INT4 training
+bash scripts/low_precision/run-qwen3-30B-A3B-int4.sh
+
+# Qwen3-235B-A22B INT4 training (8 nodes)
+bash scripts/low_precision/run-qwen3-235B-A22B-int4.sh
+
+# Kimi-k2-Thinking INT4 training (32 nodes)
+bash scripts/low_precision/run-kimi-k2-Thinking-int4.sh
+```
+
+For multi-node environments, start the Ray service according to your cluster configuration.

docs/en/advanced/pd-disaggregation.md

@@ -10,7 +10,7 @@ Use PD Disaggregation when:
 - decode dominates rollout time;
 - prefix-cache locality matters for multi-turn sessions;
 - prefill and decode need different TP, memory, or runtime settings;
-- you want an vLLM serving topology that is closer to production serving rather than a single uniform inference group.
+- you want a vLLM serving topology that is closer to production serving rather than a single uniform inference group.
 
 For short single-turn tasks, the default regular vLLM engine layout is usually simpler.
 
@@ -30,7 +30,7 @@ This is the lightweight path used by simple scripts. It is convenient when you o
 
 ### Advanced Path: `--vllm-config`
 
-For production rollout topologies, use [vLLM Config](vllm-config.md). It lets you configure prefill and decode groups independently, and can also express EPD-style layouts, heterogeneous engine groups, multi-model serving, and per-group vLLM overrides.
+For production rollout topologies, use [vLLM Config](vllm-config.md). It lets you configure prefill and decode groups independently, and can also express EPD-style layouts, heterogeneous server groups, multi-model serving, and per-group vLLM overrides.
 
 Example:
 

docs/en/developer_guide/profiling.md

@@ -205,7 +205,7 @@ launch_train_for_profiling() {
 
   # Clean up old Ray / vLLM processes (comment out if not needed)
   ray stop --force || true
-  pkill -9 -f '[v]llm serve|VLL[M]::' || true
+  pkill -9 vllm || true
   sleep 2
 
   ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats

docs/en/examples/deepseek-r1.md

@@ -0,0 +1,206 @@
+# DeepSeek R1 with 128xH100
+
+This is an example of doing DeepSeek R1 RL training using 128xH100 GPUs.
+
+We will use bf16 for training, and an fp8 format with 128x128 blockwise quantization for inference. The maximum response length is 32k, and dynamic sampling will be used to filter data during training.
+
+Regarding parallelism, for vLLM we will enable expert parallelism (`--vllm-enable-expert-parallel`) and data parallelism (`--vllm-data-parallel-size 8`). DeepEP is disabled by default. For the Megatron part, we will use TP8, PP4, EP32, and CP4.
+
+⚠️ To save GPU memory, we will use CPU Adam. Each node (8xH100) will occupy 1.4\~1.5TB of host memory. If a single machine's host memory is insufficient, this can be resolved by adding more GPUs to expand the parallelism.
+
+## Environment Setup
+
+For instructions on setting up the environment and downloading data, please refer to [Example: Qwen3-4B](qwen3-4B.md).
+
+To prepare the DeepSeek R1 checkpoint, first you will need to download DeepSeek-R1 to a directory accessible by all machines (hereinafter referred to as `$BASE_DIR`):
+
+```bash
+hf download deepseek-ai/DeepSeek-R1 --local-dir $BASE_DIR/DeepSeek-R1
+```
+
+The Hugging Face checkpoint for DeepSeek-R1 is in a block-quantized fp8 format. To convert it into a torch_dist format that Megatron can load, you first need to convert it to a bf16 Hugging Face checkpoint:
+
+```bash
+cd vime/
+python tools/fp8_cast_bf16.py --input-fp8-hf-path $BASE_DIR/DeepSeek-R1 --output-bf16-hf-path $BASE_DIR/DeepSeek-R1-bf16/
+```
+
+Next, we need to convert the bf16 version of DeepSeek-R1 into the torch_dist format. Specifically, execute the following on 4 separate nodes:
+
+```bash
+cd vime/
+source scripts/models/deepseek-v3.sh
+PYTHONPATH=/root/Megatron-LM/ torchrun \
+   --nproc-per-node 8 \
+   --master-addr ${MASTER_ADDR} --master-port 12345 \
+   --nnodes=4 --node-rank ${NODE_RANK} \
+   tools/convert_hf_to_torch_dist.py \
+   ${MODEL_ARGS[@]} \
+   --tensor-model-parallel-size 1 \
+   --pipeline-model-parallel-size 8 \
+   --expert-tensor-parallel-size 1 \
+   --expert-model-parallel-size 4 \
+   --decoder-first-pipeline-num-layers 7 \
+   --decoder-last-pipeline-num-layers 6 \
+   --hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ \
+   --save $BASE_DIR/DeepSeek-R1_torch_dist/
+```
+
+Here, `MASTER_ADDR` is the IP of node0, and `NODE_RANK` indicates the node's index, both configured similarly to a multi-node `torchrun` setup.
+
+## Executing the Training
+
+On node0, run:
+
+```bash
+cd vime/
+bash scripts/run-deepseek-r1.sh
+```
+
+On other nodes, you need to join the Ray cluster with the following command:
+
+```bash
+ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats"
+```
+
+Alternatively, if you have a list of all node IPs, for example, an MPI hostfile (where each line is `ip slot=8`), you can add the following commands after the `ray start --head` command in `scripts/run-deepseek-r1.sh`. This allows you to execute the training entirely from node0:
+
+```bash
+for WORKER_IP in $(awk '{print $1}' $BASE_DIR/mpi_hostfile); do
+  if [[ "$WORKER_IP" == "$MASTER_ADDR" ]]; then
+    continue
+  fi
+  echo "Starting Ray worker on ${WORKER_IP}"
+  ssh root@"${WORKER_IP}" \
+    "pkill -9 vllm ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" &
+done
+wait
+```
+
+### Parameter Introduction
+
+```bash
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/models/deepseek-v3.sh"
+```
+
+This reads the model's config from [scripts/models/deepseek-v3.sh](https://github.com/vllm-project/vime/blob/main/scripts/models/deepseek-v3.sh). These configs are all Megatron parameters. When training with Megatron, it cannot read the model config from the checkpoint, so we need to configure it ourselves. We provide some examples in [scripts/models](https://github.com/vllm-project/vime/tree/main/scripts/models/).
+
+#### CKPT\_ARGS
+
+```bash
+CKPT_ARGS=(
+   # HF ckpt required by vllm, we also read the tokenizer from here
+   --hf-checkpoint $BASE_DIR/DeepSeek-R1/
+   #--hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/
+   --ref-load $BASE_DIR/DeepSeek-R1_torch_dist/
+   # Actor's load directory, if empty, it will read from `ref_load`
+   --load $BASE_DIR/DeepSeek-R1_vime/
+   --save $BASE_DIR/DeepSeek-R1_vime/
+   --save-interval 20
+)
+```
+
+vime will perform online quantization during training based on the quantization configuration in `hf_checkpoint`. For instance, in the current example, we are using the fp8 checkpoint of DeepSeek R1. This means that when updating parameters, we will first perform blockwise quantization on the parameters before passing them to vllm.
+
+#### PERF\_ARGS
+
+A set of Megatron parallelism parameters. Only `--use-dynamic-batch-size` and `--max-tokens-per-gpu` are added by vime.
+
+For the Megatron part, we have configured TP8, PP4, CP4, and EP32. Since DeepSeek-R1 has 61 layers, which is not divisible by 4, we have specifically configured the last pipeline stage to have 13 layers.
+
+`max_tokens_per_gpu` refers to the maximum number of tokens each GPU can process. When `use_dynamic_batch_size` is enabled, it will pack data of varying lengths within a batch as close to `max_tokens_per_gpu`. If a single data item exceeds `max_tokens_per_gpu`, it will form its own batch without truncation. When context parallelism (CP) is enabled, it allows CP GPUs to share a total length of `CP * max_tokens_per_gpu` tokens.
+
+When `dynamic_batch_size` is enabled, the traditional `micro_batch_size` is ignored.
+
+⚠️ vime always trains the model using data packing and strictly guarantees per-sample or per-token loss. This means enabling dynamic batch size will not affect the loss calculation. It is recommended to enable it.
+
+```bash
+PERF_ARGS=(
+   --tensor-model-parallel-size 8
+   --sequence-parallel
+   --pipeline-model-parallel-size 4
+   --context-parallel-size 4
+   --expert-model-parallel-size 32
+   --expert-tensor-parallel-size 1
+   --decoder-last-pipeline-num-layers 13
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 16384
+)
+```
+
+#### GRPO\_ARGS
+
+Currently, these are some GRPO-related parameters in vime:
+
+```bash
+GRPO_ARGS=(
+   --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
+)
+```
+
+If you wish to train without loading the reference model, you need to remove `--use-kl-loss` and set `--kl-coef 0.00` (the default value is 0).
+
+#### OPTIMIZER\_ARGS
+
+We have configured CPU Adam with the following parameters to save GPU memory.
+
+```bash
+OPTIMIZER_ARGS=(
+   ...
+
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+```
+
+#### VLLM\_ARGS
+
+These are the parameters required by vllm. Here, `--rollout-num-gpus-per-engine` basically corresponds to vllm's `tp_size`. Other vllm parameters are passed to vime by adding a `--vllm-` prefix. To fully leverage vLLM's large EP inference capabilities, we enable `--vllm-enable-expert-parallel` for expert parallelism and `--vllm-data-parallel-size 8` for data-parallel attention. DeepEP is available but disabled by default (see commented flags in the script).
+
+The final `--vllm-server-concurrency` is a parameter specific to vime. It is used to prevent the vllm server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128.
+
+```bash
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 64
+   --vllm-gpu-memory-utilization 0.7
+   --vllm-enable-expert-parallel
+
+   # dp attention
+   --vllm-data-parallel-size 8
+
+    # enable deepep for vllm
+
+    # mtp
+
+    # make every dp rank has 128 concurrency
+    --vllm-server-concurrency 1024
+    --vllm-speculative-config '{"method":"eagle","num_speculative_tokens":4}'
+)
+```
+
+#### MISC\_ARGS
+
+Some additional Megatron configurations. Note that Megatron's deepep is configured here.
+
+```bash
+MISC_ARGS=(
+   ...
+
+   # use deepep for megatron
+   --moe-enable-deepep
+   --moe-token-dispatcher-type flex
+)
+```